summaryrefslogtreecommitdiffstats
path: root/map.sml
blob: d74cfd55de517a76b264ea2716e43a06e86640e0 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
(* 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