(* map is a red-black tree implementing an immutable key-value store. *)
signature MAP =
sig
  type key
  type 'a map

  val empty : 'a map
  val insert : 'a map -> key * 'a -> 'a map
  val lookup : 'a map -> key -> 'a option
  val delete : 'a map -> key -> 'a map
end

functor Map (type k
             val cmp : k * k -> order) :> MAP where type key = k =
struct
  type key = k

  datatype color = Red | Black
  datatype 'a node = Node of color * (key * 'a) * ('a node) option * ('a node) option
  type 'a map = ('a node) option
  
  val empty = NONE

  fun rebalance (Node (Black, g, SOME (Node (Red, p, SOME (Node (Red, n, nLeft, nRight)), pRight)), u)) : 'a node =
        (*     b
              /
             r
            /
           r     *)
        Node (Red, p, SOME (Node (Black, n, nLeft, nRight)),
                      SOME (Node (Black, g, pRight, u)))
    | rebalance (Node (Black, g, u, SOME (Node (Red, p, pLeft, SOME (Node (Red, n, nLeft, nRight)))))) =
        (* b
            \
             r
              \
               r *)
        Node (Red, p, SOME (Node (Black, g, u, pLeft)),
                      SOME (Node (Black, n, nLeft, nRight)))
    | rebalance (Node (Black, g, u, SOME (Node (Red, p, SOME (Node (Red, n, nLeft, nRight)), pRight)))) =
        (* b
            \
             r
            /
           r   *)
        Node (Red, n, SOME (Node (Black, g, u, nLeft)),
                      SOME (Node (Black, p, nRight, pRight)))
    | rebalance (Node (Black, g, SOME (Node (Red, p, pLeft, SOME (Node (Red, n, nLeft, nRight)))), u)) =
        (*   b
            /
           r
            \
             r *)
       Node (Red, n, SOME (Node (Black, p, pLeft, nLeft)),
                     SOME (Node (Black, g, nRight, u)))
    | rebalance n = n
  
  fun insert m (kv as (k, _)) =
        let
          fun insertNode NONE = Node (Red, kv, NONE, NONE)
            | insertNode (SOME (Node (color, nodeKV as (nodeKey, _), left, right))) =
                case cmp (k, nodeKey) of
                    LESS => rebalance (Node (color, nodeKV, SOME (insertNode left), right))
                  | EQUAL => Node (color, kv, left, right)
                  | GREATER => rebalance (Node (color, nodeKV, left, SOME (insertNode right)))
          val Node (_, nodeKV, left, right) = insertNode m
        in
          SOME (Node (Black, nodeKV, left, right))
        end

  fun lookup m k =
        let fun loop NONE = NONE
              | loop (SOME (Node (_, (nodeKey, nodeVal), left, right))) =
                  case cmp (k, nodeKey) of
                      LESS => loop left
                    | EQUAL => SOME nodeVal
                    | GREATER => loop right
        in loop m
        end

  fun sub1 (SOME (Node (Black, x, a, b))) = Node (Red, x, a, b)
    | sub1 _ = raise Fail "invariance violation"

  (* Thanks to https://github.com/sweirich/dth/blob/master/examples/red-black/RedBlack.lhs *)
  fun delete t k =
        let
          fun del NONE = NONE
            | del (SOME (Node (_, (nodeKV as (nodeKey, _)), left, right))) =
                case cmp (k, nodeKey) of
                    LESS => SOME (delLeft nodeKV left right)
                  | EQUAL => merge left right
                  | GREATER => SOME (delRight nodeKV left right)

          and delLeft nodeKV (a as SOME (Node (Black, _, _, _))) b = balLeft nodeKV (del a) b
            | delLeft nodeKV a b = Node (Red, nodeKV, del a, b)

          and balLeft x (SOME (Node (Red, y, a, b))) c = Node (Red, x, SOME (Node (Black, y, a, b)), c)
            | balLeft x bl (SOME (Node (Black, y, a, b))) =
                rebalance (Node (Black, x, bl, SOME (Node (Red, y, a, b))))
            | balLeft x bl (SOME (Node (Red, z, SOME (Node (Black, y, a, b)), c))) =
                Node (Red, y, SOME (Node (Black, x, bl, a)), SOME (rebalance (Node (Black, z, b, SOME (sub1 c)))))
            | balLeft _ _ _ = raise Fail "unreachable"

          and delRight nodeKV a (b as SOME (Node (Black, _, _, _))) = balRight nodeKV a (del b)
            | delRight nodeKV a b = Node (Red, nodeKV, a, del b)

          and balRight x a (SOME (Node (Red, y, b, c))) = Node (Red, x, a, SOME (Node (Black, y, b, c)))
            | balRight y (SOME (Node (Black, x, a, b))) bl = rebalance (Node (Black, y, SOME (Node (Red, x, a, b)), bl))
            | balRight z (SOME (Node (Red, x, a, SOME (Node (Black, y, b, c))))) bl =
                Node (Red, y, SOME (rebalance (Node (Black, x, SOME (sub1 a), b))),
                              SOME (Node (Black, z, c, bl)))
            | balRight _ _ _ = raise Fail "unreachable"

          and merge NONE x = x
            | merge x NONE = x
            | merge (SOME (Node (Red, x, a, b))) (SOME (Node (Red, y, c, d))) =
                (case merge b c of
                     SOME (Node (Red, z, b', c')) =>
                       SOME (Node (Red, z, SOME (Node (Red, x, a, b')),
                                           SOME (Node (Red, y, c', d))))
                   | bc => SOME (Node (Red, x, a, SOME (Node (Red, y, bc, d)))))
            | merge (SOME (Node (Black, x, a, b))) (SOME (Node (Black, y, c, d))) =
                (case merge b c of
                     SOME (Node (Red, z, b', c')) =>
                       SOME (Node (Red, z, SOME (Node (Black, x, a, b')),
                                           SOME (Node (Black, y, c', d))))
                   | bc => SOME (balLeft x a (SOME (Node (Black, y, bc, d)))))
            | merge a (SOME (Node (Red, x, b, c))) = SOME (Node (Red, x, merge a b, c))
            | merge (SOME (Node (Red, x, a, b))) c = SOME (Node (Red, x, a, merge b c))
        in
          case del t of
              NONE => NONE
            | SOME (Node (_, kv, left, right)) => SOME (Node (Black, kv, left, right))
        end
end