red-black-tree.arr

include string-dict

data ControlFlow<a, b>: next(v :: a) | brk(v :: b) end
fun loop<a, b>(func :: (a -> ControlFlow<a, b>), init :: a) -> b:
  cases(ControlFlow) func(init):
    | brk(v) => v
    | next(v) => loop(func, v)
  end
end

data NodeColor:
  | black
  | red
end

data RedBlackTree<k, v>:
  | rb-tree-empty with:
    method contains-key(self, key :: k) -> Boolean:
      false
    end,
    method l-balance(self) -> RedBlackTree<k, v>:
      self
    end,
    method r-balance(self) -> RedBlackTree<k, v>:
      self
    end
  | rb-tree-node(
      color :: NodeColor, 
      key :: k,
      value :: v,
      left :: RedBlackTree<k, v>,
      right :: RedBlackTree<k, v>
    ) with:
    method contains-key(self, key :: k) -> Boolean:
      (self.key == key) or if key < self.key: self.left else: self.right end.contains-key(key)
    end,
    method l-balance(self) -> RedBlackTree<k, v>:
      ask:
        | (self.color == black) and
          is-rb-tree-node(self.left) and is-rb-tree-node(self.left.left) and 
          (self.left.color == red) and (self.left.left.color == red) then:
          #|
                        bz      
                        /\               ry
                      ry  d             /  \
                      /\        =>    bx    bz
                    rx  c             /\    /\
                    /\               a  b  c  d
                   a  b
          |#
          rb-tree-node(
            red, self.left.key, self.left.value,
            rb-tree-node(
              black, self.left.left.key, self.left.left.value, 
              self.left.left.left, 
              self.left.left.right
             ),
            rb-tree-node(
              black, self.key, self.value, 
              self.left.right, 
              self.right
            )
          )
        | is-rb-tree-node(self.left) and is-rb-tree-node(self.left.right) and
          (self.left.color == red) and (self.left.right.color == red) then:
          #|
                        bz      
                        /\                ry
                      rx  d              /  \
                      /\        =>      bx   bz
                     a  ry             /\    /\
                        /\            a  b  c  d
                       b  c
          |#
          rb-tree-node(
            red, self.left.right.key, self.left.right.value,
            rb-tree-node(
              black, self.left.key, self.left.value, 
              self.left.left, 
              self.left.right.left
            ),
            rb-tree-node(
              black, self.key, self.value, 
              self.left.right.right, 
              self.right
            )
          )
        | otherwise: self
      end
    end,
    method r-balance(self) -> RedBlackTree<k, v>:
      ask:
        | (self.color == black) and
          is-rb-tree-node(self.right) and is-rb-tree-node(self.right.left) and 
          (self.right.color == red) and (self.right.left.color == red) then:
          #|
                       bx      
                       /\                 ry
                      a  rz              /  \
                         /\      =>    bx    bz
                       ry  d           /\    /\
                       /\             a  b  c  d
                      b  c
          |#
          rb-tree-node(
            red, self.right.left.key, self.right.left.value,
            rb-tree-node(
              black, self.key, self.value, 
              self.left, 
              self.right.left.left
            ),
            rb-tree-node(
              black, self.right.key, self.right.value, 
              self.right.left.right, 
              self.right.right
            )
          )
        | is-rb-tree-node(self.right) and is-rb-tree-node(self.right.right) and 
          (self.right.color == red) and (self.right.right.color == red) then:
          #|
                       bx      
                       /\                   ry
                      a  ry                /  \
                         /\       =>     bx    bz
                        b  rz            /\    /\
                           /\           a  b  c  d
                          c  d
          |#
          rb-tree-node(
            red, self.right.key, self.right.value,
            rb-tree-node(
              black, self.key, self.value, 
              self.left, 
              self.right.left
            ),
            rb-tree-node(
              black, self.right.right.key, self.right.right.value, 
              self.right.right.left, 
              self.right.right.right
            )
          )
        | otherwise: self
      end
    end
sharing:
  method insert(self, key :: k, val :: v) -> RedBlackTree<k, v>:
    fun ins(s :: RedBlackTree<k, v>) -> RedBlackTree<k, v>:
      cases(RedBlackTree) s:
        | rb-tree-empty => rb-tree-node(red, key, val, rb-tree-empty, rb-tree-empty)
        | rb-tree-node(color, yk, y, left, right) =>
          ask:
            | key < yk then: 
              rb-tree-node(color, yk, y, ins(left), right).l-balance()
            | yk < key then:
              rb-tree-node(color, yk, y, left, ins(right)).r-balance()
            | otherwise: s
          end
      end
    end
    res = ins(self)
    rb-tree-node(black, res.key, res.value, res.left, res.right)
  end,
  method rank(self) -> Number:
    for loop({rmax; l} from { 0; [list: {0; self}] }):
      cases(List) l:
        | empty => brk(rmax)
        | link(x, next-l) =>
          {r; n} = x
          cases(RedBlackTree) n:
            | rb-tree-empty => next({ num-max(r, rmax); next-l })
            | rb-tree-node(_, _, _, left, right) =>
              next({ num-max(r, rmax); next-l.push({r + 1; left}).push({r + 1; right}) })
          end
      end
    end
  end,
  method is-leaf(self) -> Boolean:
    is-rb-tree-node(self) and is-rb-tree-empty(self.left) and is-rb-tree-empty(self.right)
  end,
  method each(
      self, 
      foo :: ({Option<{Number; Number}>; {Number; Number}; NodeColor; k; v} -> Nothing)
      ) -> Nothing:
    fun traverse(
        t :: RedBlackTree<k, v>,
        n :: Number,
        prev-rng :: Option<{Number; Number}>,
        rng :: {Number; Number},
        f :: ({Number; Option<{Number; Number}>; {Number; Number}; NodeColor; k; v} -> Nothing)
        ) -> Nothing:
      ask:
        | is-rb-tree-empty(t) then: nothing
        | t.is-leaf() then:
          block:
            f({n; prev-rng; rng; t.color; t.key; t.value})
            nothing
          end
        | otherwise:
          block:
            m-rng = (rng.{0} + rng.{1}) / 2
            l-rng = {rng.{0}; m-rng}
            r-rng = {m-rng; rng.{1}}
            traverse(t.left, n + 1, some(rng), l-rng, f)
            f({n; prev-rng; rng; t.color; t.key; t.value})
            traverse(t.right, n + 1, some(rng), r-rng, f)
            nothing
          end
      end
    end
    traverse(
      self, 0, none, {0; 1},
      lam({n; prev-rng; rng; color; key; value}):
        prev-coords = prev-rng
          .and-then({(it):
            x0 = ((it.{0} + it.{1}) / 2)
            y0 = n
            {x0; y0}
          })
        
        x1 = ((rng.{0} + rng.{1}) / 2)
        y1 = n + 1
        coords = {x1; y1}
        foo({prev-coords; coords; color; key; value})
      end
    )
  end,
  method draw(self) -> Image block:
    height-multiplier = 48
    radius = height-multiplier / 3
    width = num-expt(2, num-max(self.rank() - 1, 3)) * 2 * radius
    
    fun make-txt(color :: NodeColor, key :: k, value :: v) -> Image:
      txt = to-string(key) + ": " + to-string(value)
      text(txt, 10, if color == black: "white" else: "black" end)
    end

    fun draw-node(
        back :: Image, 
        color :: NodeColor, 
        coords :: {Number; Number}, 
        key :: k, 
        value :: v) -> Image:
      y = coords.{1} * height-multiplier
      x = coords.{0} * width
       
      place-image(
        circle(
          radius, 
          "solid", 
          if color == black: "black" else: "pink" end
        ), 
        x, 
        y, 
        back
      ) ^ place-image(make-txt(color, key, value), x, y, _)
    end
    
    var back = empty-scene(width, (self.rank() + 1) * height-multiplier)
    
    # drawing lines
    self.each(
      lam({prev-coords; coords; _; _; _}):
        cases(Option) prev-coords:
          | none => nothing
          | some({x0; y0}) =>
            block:
              {shadow y0; shadow y1} = {y0 * height-multiplier; coords.{1} * height-multiplier}
              shadow x0 = x0 * width
              x1 = coords.{0} * width
              back := scene-line(back, x0, y0, x1, y1, "black")
              nothing
            end
        end
      end
    )
    
    # drawing nodes
    self.each(
      lam({_; coords; color; key; value}):
        back := draw-node(back, color, coords, key, value)
      end
    )
    back
  end
where:
  t0 = rb-tree-empty.insert(0, "a").insert(1, "a").insert(2, "a")
  t0.rank() is 2
  
  t1 = rb-tree-empty.insert(2, "a").insert(1, "a").insert(0, "a")
  t1.rank() is 2
  
  t2 = rb-tree-empty.insert(2, "a").insert(0, "a").insert(1, "a")
  t2.rank() is 2
  
  t3 = rb-tree-empty.insert(1, "a").insert(0, "a").insert(2, "a")
  t3.rank() is 2
  
  t4 = rb-tree-empty
    .insert(0, "a") 
    .insert(1, "a")
    .insert(97, "a")
    .insert(3, "a")
    .insert(95, "a")
    .insert(4, "a")
    .insert(5, "a")
    .insert(94, "a")
    .insert(6, "a")
    .insert(99, "a")
    .insert(98, "a")
    .insert(96, "a")
    .insert(2, "a")
    .insert(93, "a")
    .insert(92, "a")
    .insert(10002, "a")
    .insert(10093, "a")
    .insert(10092, "a")
    .insert(5002, "a")
    .insert(5093, "a")
    .insert(5092, "a")
  num-abs(t4.right.rank() - t4.left.rank()) <= 2 is true
  
end

fun draw-balances():
  fun make-unbalanced-node-case-1() -> RedBlackTree<Number, String>:
    rb-tree-node(
      black, 10, "z",
      rb-tree-node(
        red, 8, "y", 
        rb-tree-node(red, 6, "x", rb-tree-empty,rb-tree-empty), 
        rb-tree-empty
      ), 
      rb-tree-empty
    )
  end
  
  fun make-unbalanced-node-case-2() -> RedBlackTree<Number, String>:
    rb-tree-node(
      black, 20, "z",
      rb-tree-node(
        red, 16, "x", 
        rb-tree-empty, 
        rb-tree-node(red, 18, "y", rb-tree-empty, rb-tree-empty)
        ), 
      rb-tree-empty
    )
  end
  
  fun make-unbalanced-node-case-3() -> RedBlackTree<Number, String>:
    rb-tree-node(
      black, 10, "z",
      rb-tree-empty,
      rb-tree-node(
        red, 12, "y", 
        rb-tree-node(red, 11, "x", rb-tree-empty,rb-tree-empty), 
        rb-tree-empty
      )
    )
  end
  
  fun make-unbalanced-node-case-4() -> RedBlackTree<Number, String>:
    rb-tree-node(
      black, 10, "z",
      rb-tree-empty,
      rb-tree-node(
        red, 12, "y", 
        rb-tree-empty, 
        rb-tree-node(red, 14, "x", rb-tree-empty,rb-tree-empty)
      )
    )
  end
  
  fun draw-t-l(t):
    t.draw()
      ^ above(_, text("becomes", 24, "blue"))
      ^ above(_, t.l-balance().draw())
  end

  fun draw-t-r(t):
    t.draw()
      ^ above(_, text("becomes", 24, "blue"))
      ^ above(_, t.r-balance().draw())
  end
  
  u1 = make-unbalanced-node-case-1()
  u2 = make-unbalanced-node-case-2()
  u3 = make-unbalanced-node-case-3()
  u4 = make-unbalanced-node-case-4()
  
  text("lefts:", 24, "blue")
    ^ above(_, beside(draw-t-l(u1), draw-t-l(u2)))
    ^ above(_, text("rights:", 24, "blue"))
    ^ above(_, beside(draw-t-r(u3), draw-t-r(u4)))
end

fun draw-random-tree(size :: Number) -> Image:
  letters = [array: "a", "b", "c", "d", "e", "f", "g", "h"]
  l = for loop({i; lst} from {0; empty}):
    if i == size:
      brk(lst)
    else:
      letter-id = num-random(letters.length())
      next({i + 1; link({i; letters.get-now(letter-id)}, lst)})
    end
  end
  
  
  for loop({t; lst} from {rb-tree-empty; shuffle(l)}):
    cases(List) lst:
      | empty => brk(t)
      | link(x, xs) => next({t.insert(x.{0}, x.{1}); xs})
    end
  end.draw()
end