<commit_before>#
# A red/black binary search tree
#

import strutils, optional_t

type
    Color = enum ## The color of a node
        red, black

    Node[T] = ref object ## A node within a red/black tree

        # The (potential) connections in this node
        left, right, parent: Node[T]

        # The color of this node
        color: Color

        # The  value of this node
        value: T

    RedBlackTree*[T] = object ## A Red/Black tree

        compare: proc (a, b: T): int

        # The root of the tree
        root: Node[T]


proc newRBTree*[T]( compare: proc (a, b: T): int = cmp ): RedBlackTree[T] =
  # Creates a new Red/Black tree
  return RedBlackTree[T]( root: nil, compare: compare )

proc `$` [T]( self: Node[T] ): string =
    ## Converts a node to a string
    if self == nil:
        result = "()"
    else:
        result = "(" &
            (if self.color == red: "R" else: "B") & " " &
            $(self.value)
        if self.left != nil or self.right != nil:
            result.add(" " & $(self.left) & " " & $(self.right))
        result.add(")")

proc `$`* [T]( self: RedBlackTree[T] ): string =
    ## Returns a tree as a string
    return "RedBlackTree" & `$`[T](self.root)

proc find[T]( tree: RedBlackTree[T], value: T ): Node[T] =
    ## Find a value in the tree and returns the containing node. Or nil
    var examine = tree.root
    while examine != nil:
        if examine.value == value:
            return examine
        elif tree.compare(value, examine.value) <= 0:
            examine = examine.left
        else:
            examine = examine.right
    return nil



proc newNode[T]( value: T, parent: Node[T] = nil ): Node[T] {.inline.} =
    return Node[T](
        left: nil, right: nil, parent: parent, color: red, value: value )

proc insert[T](
    self: var Node[T], compare: proc(a, b: T): int, value: T
): Node[T] =
    ## Does a basic binary search tree insert, returning the new node
    if compare(value, self.value) <= 0:
        if self.left == nil:
            result = newNode( value, self )
            self.left = result
        else:
            result = insert(self.left, compare, value)
    else:
        if self.right == nil:
            result = newNode( value, self )
            self.right = result
        else:
            result = insert(self.right, compare, value)

proc grandparent[T](node: var Node[T]): Node[T] {.inline.} =
    ## Returns the grandparent of a node; the parent of a parent
    if node != nil and node.parent != nil:
        return node.parent.parent
    else:
        return nil

proc uncle[T](node: var Node[T]): Node[T] {.inline.} =
    ## Returns the uncle (the parent's sibling) of a node
    let grandparent = grandparent(node)
    if grandparent == nil:
        return nil # No grandparent means no uncle
    elif node.parent == grandparent.left:
        return grandparent.right
    else:
        return grandparent.left

proc rotateLeft[T]( tree: var RedBlackTree[T], node: var Node[T] ) {.inline.} =
    var parent = node.parent
    if parent == nil:
        return

    # It looks like nim is trying to garbage collect this node before it's
    # actually gone for some reason. This is a hack to get around that
    var saved = node

    var grandparent = parent.parent
    var child = saved.left

    # Move the child over
    parent.right = child
    if child != nil:
        child.parent = parent

    # Move the parent around
    saved.left = parent
    parent.parent = saved

    # Move the node itself
    saved.parent = grandparent

    # Update the grandparent, swapping the root of the tree if needed
    if grandparent == nil:
        tree.root = saved
    elif grandparent.left == parent:
        grandparent.left = saved
    else:
        grandparent.right = saved

proc rotateRight[T]( tree: var RedBlackTree[T], node: var Node[T] ) {.inline.} =
    var parent = node.parent
    if parent == nil:
        return

    # It looks like nim is trying to garbage collect this node before it's
    # actually gone for some reason. This is a hack to get around that
    var saved = node

    var grandparent = parent.parent
    var child = saved.right

    # Move the child over
    parent.left = child
    if child != nil:
        child.parent = parent

    # Move the parent around
    saved.right = parent
    parent.parent = saved

    # Move the node itself
    saved.parent = grandparent

    # Update the grandparent, swapping the root of the tree if needed
    if grandparent == nil:
        tree.root = saved
    elif grandparent.left == parent:
        grandparent.left = saved
    else:
        grandparent.right = saved

template isOn[T]( self: Node[T], side: expr ): bool =
    ## Whether this node is the right child of its parent
    self.parent != nil and self.parent.`side` == self

proc isRed[T]( node: Node[T] ): bool {.inline.} =
    ## Safely checks whether a value is a red node or not
    return node != nil and node.color == red

proc isBlack[T]( node: Node[T] ): bool {.inline.} =
    ## Safely checks whether a node is black
    return node == nil or node.color == black

proc isAllBlack[T]( node: Node[T] ): bool {.inline.} =
    ## Safely checks if a node and its children are all black
    return node.isBlack and node.left.isBlack and node.right.isBlack

template far[T]( node: Node[T], direction: expr ): Node[T] =
    ## Walks every child in a specific direction down to the bottom
    var result = node<commit_msg>Pull rotate code into a template<commit_after>#
# A red/black binary search tree
#

import strutils, optional_t

type
    Color = enum ## The color of a node
        red, black

    Node[T] = ref object ## A node within a red/black tree

        # The (potential) connections in this node
        left, right, parent: Node[T]

        # The color of this node
        color: Color

        # The  value of this node
        value: T

    RedBlackTree*[T] = object ## A Red/Black tree

        compare: proc (a, b: T): int

        # The root of the tree
        root: Node[T]


proc newRBTree*[T]( compare: proc (a, b: T): int = cmp ): RedBlackTree[T] =
  # Creates a new Red/Black tree
  return RedBlackTree[T]( root: nil, compare: compare )

proc `$` [T]( self: Node[T] ): string =
    ## Converts a node to a string
    if self == nil:
        result = "()"
    else:
        result = "(" &
            (if self.color == red: "R" else: "B") & " " &
            $(self.value)
        if self.left != nil or self.right != nil:
            result.add(" " & $(self.left) & " " & $(self.right))
        result.add(")")

proc `$`* [T]( self: RedBlackTree[T] ): string =
    ## Returns a tree as a string
    return "RedBlackTree" & `$`[T](self.root)

proc find[T]( tree: RedBlackTree[T], value: T ): Node[T] =
    ## Find a value in the tree and returns the containing node. Or nil
    var examine = tree.root
    while examine != nil:
        if examine.value == value:
            return examine
        elif tree.compare(value, examine.value) <= 0:
            examine = examine.left
        else:
            examine = examine.right
    return nil



proc newNode[T]( value: T, parent: Node[T] = nil ): Node[T] {.inline.} =
    return Node[T](
        left: nil, right: nil, parent: parent, color: red, value: value )

proc insert[T](
    self: var Node[T], compare: proc(a, b: T): int, value: T
): Node[T] =
    ## Does a basic binary search tree insert, returning the new node
    if compare(value, self.value) <= 0:
        if self.left == nil:
            result = newNode( value, self )
            self.left = result
        else:
            result = insert(self.left, compare, value)
    else:
        if self.right == nil:
            result = newNode( value, self )
            self.right = result
        else:
            result = insert(self.right, compare, value)

proc grandparent[T](node: var Node[T]): Node[T] {.inline.} =
    ## Returns the grandparent of a node; the parent of a parent
    if node != nil and node.parent != nil:
        return node.parent.parent
    else:
        return nil

proc uncle[T](node: var Node[T]): Node[T] {.inline.} =
    ## Returns the uncle (the parent's sibling) of a node
    let grandparent = grandparent(node)
    if grandparent == nil:
        return nil # No grandparent means no uncle
    elif node.parent == grandparent.left:
        return grandparent.right
    else:
        return grandparent.left


template rotate( direction: expr, opposite: expr ) {.immediate.} =
    ## Rotates a node into its parents position

    var parent = node.parent
    if parent == nil:
        return

    # It looks like nim is trying to garbage collect this node before it's
    # actually gone for some reason. This is a hack to get around that
    var saved = node

    var grandparent = parent.parent
    var child = saved.`direction`

    # Move the child over
    parent.`opposite` = child
    if child != nil:
        child.parent = parent

    # Move the parent around
    saved.`direction` = parent
    parent.parent = saved

    # Move the node itself
    saved.parent = grandparent

    # Update the grandparent, swapping the root of the tree if needed
    if grandparent == nil:
        tree.root = saved
    elif grandparent.left == parent:
        grandparent.left = saved
    else:
        grandparent.right = saved

proc rotateLeft[T]( tree: var RedBlackTree[T], node: var Node[T] ) {.inline.} =
    ## Rotates a node to the left
    rotate(left, right)

proc rotateRight[T]( tree: var RedBlackTree[T], node: var Node[T] ) {.inline.} =
    ## Rotates a node to the right
    rotate(right, left)


template isOn[T]( self: Node[T], side: expr ): bool =
    ## Whether this node is the right child of its parent
    self.parent != nil and self.parent.`side` == self

proc isRed[T]( node: Node[T] ): bool {.inline.} =
    ## Safely checks whether a value is a red node or not
    return node != nil and node.color == red

proc isBlack[T]( node: Node[T] ): bool {.inline.} =
    ## Safely checks whether a node is black
    return node == nil or node.color == black

proc isAllBlack[T]( node: Node[T] ): bool {.inline.} =
    ## Safely checks if a node and its children are all black
    return node.isBlack and node.left.isBlack and node.right.isBlack

template far[T]( node: Node[T], direction: expr ): Node[T] =
    ## Walks every child in a specific direction down to the bottom
    var result = node
{'fraction_non_alphanumeric': 0.11223021582733812, 'fraction_numerical': 0.00038369304556354914, 'mean_word_length': 2.3556485355648538, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1483597', 'n_tokens_mistral': 3140, 'n_tokens_neox': 3073, 'n_words': 1205}