rb_tree.cr

# Okasaki style Functional Red Black Tree
# https://www.cs.tufts.edu/comp/150FP/archive/chris-okasaki/redblack99.pdf
#
# leaves and root are black
# BST
# No red node has a red child
# Every path from the root to a leaf has the same number of black nodes
#
# Based off Tenderlove's ruby version: https://gist.github.com/tenderlove/c843996b39e50ea6eb129871f07e589d
#
# Conversion to Crystal was pretty simple except for the tree balance method, which as Tenderlove 
# demonstrated is much simpler when using Ruby 3's new pattern matching features.
#
module RBTree
  class Leaf
    include Enumerable(Int32)

    def color
      :black
    end

    def leaf?
      true
    end

    def key
      0
    end

    def left
      LEAF
    end

    def right
      LEAF
    end

    def red?
      color == :red
    end

    def key?(_arg)
      false
    end

    def num_black
      1
    end

    def balanced?
      true
    end

    def each(&blk)
      yield self
    end

    def deconstruct
      [:black, nil, nil, nil]
    end

    def insert(val)
      RBTree.insert self, val
    end
  end

  class Node < Leaf
    getter :color, :key, :left, :right

    @color : Symbol
    @key : Int32
    @left : Leaf
    @right : Leaf

    def initialize(color, key, left, right)
      @color = color
      @key = key
      @left = left
      @right = right
    end

    def leaf?
      false
    end

    def deconstruct
      [color, key, left, right]
    end

    def key?(k)
      return true if k == key
      k < key ? left.key?(k) : right.key?(k)
    end

    def insert(val)
      RBTree.insert self, val
    end

    def each(&blk)
      yield self
      left.each(&blk)
      right.each(&blk)
    end
  end

  private def self.balance(color, key, left, right)
    if color == :black
      x = y = z = 0
      a = b = c = d = LEAF

      if left.red? && left.left && left.left.red?
        z = key
        y = left.key
        x = left.left.key
        a = left.left.left
        b = left.left.right
        c = left.right
        d = right
      elsif left.red? && left.right && left.right.red?
        z = key
        x = left.key
        y = left.left.key
        a = left.left
        b = left.right.left
        c = left.right.right
        d = right
      elsif right.red? && right.left && right.left.red?
        x = key
        z = right.key
        y = right.left.key
        a = left
        b = right.left.left
        c = right.left.right
        d = right.right
      elsif right.red? && right.left && right.left.red?
        x = key
        y = right.key
        z = right.right.key
        a = left
        b = right.left
        c = right.right.left
        d = right.right.right
      else
        return Node.new(color, key, left, right)
      end

      Node.new(:red, y, Node.new(:black, x, a, b), Node.new(:black, z, c, d))
    end

    Node.new(color, key, left, right)
  end

  private def self.insert_aux(tree, x)
    if tree.leaf?
      Node.new :red, x, LEAF, LEAF
    else
      if x < tree.key
        balance(tree.color, tree.key, insert_aux(tree.left, x), tree.right)
      elsif x > tree.key
        balance(tree.color, tree.key, tree.left, insert_aux(tree.right, x))
      else
        tree
      end
    end
  end

  def self.insert(tree, x)
    root = insert_aux(tree, x)
    if root.color == :red
      Node.new(:black, root.key, root.left, root.right)
    else
      root
    end
  end

  LEAF = Leaf.new

  def self.new
    LEAF
  end
end

tree = RBTree.new
tree = tree.insert 5
tree = tree.insert 4
tree = tree.insert 2
tree = tree.insert 7
tree = tree.insert 8
tree = tree.insert 9
tree = tree.insert 1
tree = tree.insert 3
5000.times { tree = tree.insert rand(123456) }
p tree.key?(8)  # true
p tree.key?(10) # maybe false