1(* Redblackmap --                                                   *)
2(*    applicative maps implemented by Okasaki-style Red-Black trees *)
3(* Ken Friis Larsen <ken@friislarsen.net>                           *)
4structure Redblackmap :>  Redblackmap =
5struct
6
7  datatype ('key, 'a) tree =
8           LEAF
9         | RED   of 'key * 'a * ('key, 'a) tree * ('key, 'a) tree
10         | BLACK of 'key * 'a * ('key, 'a) tree * ('key, 'a) tree
11
12  type ('key, 'a) dict  = ('key * 'key -> order) * ('key, 'a) tree * int
13
14  exception NotFound
15
16  fun mkDict compare = (compare, LEAF, 0)
17
18  fun numItems (_, _, n) = n
19
20  fun isEmpty (_, _, n) = (n = 0)
21
22  fun findKey ((compare, tree, n), key) =
23      let fun loopShared k x left right =
24              case compare(key, k) of
25                  EQUAL   => (k, x)
26                | LESS    => loop left
27                | GREATER => loop right
28          and loop LEAF                       = raise NotFound
29            | loop (RED(k, x, left, right))   = loopShared k x left right
30            | loop (BLACK(k, x, left, right)) = loopShared k x left right
31      in  loop tree end
32
33  fun find x = #2(findKey x)
34
35  fun peek (set, key) = SOME(find(set, key))
36                        handle NotFound => NONE
37
38  fun lbalance z zd (RED(y,yd,RED(x,xd,a,b),c)) d =
39      RED(y,yd,BLACK(x,xd,a,b),BLACK(z,zd,c,d))
40    | lbalance z zd (RED(x,xd,a,RED(y,yd,b,c))) d =
41      RED(y,yd,BLACK(x,xd,a,b),BLACK(z,zd,c,d))
42    | lbalance k x left right = BLACK(k, x, left, right)
43
44  fun rbalance x xd a (RED(y,yd,b,RED(z,zd,c,d))) =
45      RED(y,yd,BLACK(x,xd,a,b),BLACK(z,zd,c,d))
46    | rbalance x xd a (RED(z,zd,RED(y,yd,b,c),d)) =
47      RED(y,yd,BLACK(x,xd,a,b),BLACK(z,zd,c,d))
48    | rbalance k x left right = BLACK(k, x, left, right)
49
50  exception GETOUT
51
52  local open Uref in
53  fun update (set as (compare, tree, n), key, data) =
54      let val addone = Uref.new true
55          fun ins LEAF = RED(key,data NONE,LEAF,LEAF)
56            | ins (BLACK(k,x,left,right)) =
57              (case compare(key, k) of
58                   LESS    => lbalance k x (ins left) right
59                 | GREATER => rbalance k x left (ins right)
60                 | EQUAL   => (addone := false; BLACK(key, data (SOME x), left, right)))
61            | ins (RED(k, x,left,right)) =
62              (case compare(key, k) of
63                   LESS    => RED(k, x, (ins left), right)
64                 | GREATER => RED(k, x, left, (ins right))
65                 | EQUAL   => (addone := false; RED(key, data (SOME x), left, right)))
66      in  ( compare
67          , case ins tree of
68                RED x => BLACK x
69              | tree  => tree
70          , if !addone then n+1 else n) end
71  end
72
73  local fun K x _ = x in
74    fun insert (set, key, data) = update (set, key, K data)
75  end
76
77  fun insertList (m, xs) =
78    List.foldl (fn ((i, v), m) => insert (m, i, v)) m xs
79
80  fun fromList compare xs =
81    insertList (mkDict compare, xs)
82
83  fun push LEAF stack = stack
84    | push tree stack = tree :: stack
85
86  fun pushNode left k x right stack =
87      left :: (BLACK(k, x, LEAF, LEAF) :: (push right stack))
88
89  fun getMin []             some none = none
90    | getMin (tree :: rest) some none =
91      case tree of
92          LEAF                 => getMin rest some none
93        | RED  (k, x, LEAF, b) => some k x (push b rest)
94        | BLACK(k, x, LEAF, b) => some k x (push b rest)
95        | RED  (k, x, a, b)    => getMin(pushNode a k x b rest) some none
96        | BLACK(k, x, a, b)    => getMin(pushNode a k x b rest) some none
97
98  fun getMax []             some none = none
99    | getMax (tree :: rest) some none =
100      case tree of
101          LEAF                 => getMax rest some none
102        | RED  (k, x, a, LEAF) => some k x (push a rest)
103        | BLACK(k, x, a, LEAF) => some k x (push a rest)
104        | RED  (k, x, a, b)    => getMax(pushNode b k x a rest) some none
105        | BLACK(k, x, a, b)    => getMax(pushNode b k x a rest) some none
106
107  fun fold get f e (compare, tree, n) =
108      let fun loop stack acc =
109              get stack (fn k =>fn x =>fn stack => loop stack (f(k,x,acc))) acc
110      in  loop [tree] e end
111
112  fun foldl f = fold getMin f
113
114  fun foldr f = fold getMax f
115
116  fun listItems set = foldr (fn(k,x,res) => (k,x)::res) [] set
117
118  fun appAll get f (compare, tree, n) =
119      let fun loop stack = get stack (fn k => fn x => (f(k,x); loop)) ()
120      in  loop [tree] end
121
122  fun app f = appAll getMin f
123
124  fun revapp f = appAll getMax f
125
126
127  exception RedBlackMapError
128
129  (* remove a la Stefan M. Kahrs *)
130  fun redden (BLACK arg) = RED arg
131    | redden _ = raise RedBlackMapError
132
133  fun balleft y yd (RED(x,xd,a,b)) c                =
134      RED(y, yd, BLACK(x, xd, a, b), c)
135    | balleft x xd bl (BLACK(y, yd, a, b))          =
136      rbalance x xd bl (RED(y, yd, a, b))
137    | balleft x xd bl (RED(z,zd,BLACK(y,yd,a,b),c)) =
138      RED(y, yd, BLACK(x, xd, bl, a), rbalance z zd b (redden c))
139    | balleft _ _ _ _ = raise RedBlackMapError
140
141  fun balright x xd a             (RED(y, yd ,b,c)) =
142      RED(x, xd, a, BLACK(y, yd, b, c))
143    | balright y yd (BLACK(x,xd,a,b))          br =
144      lbalance y yd (RED(x,xd,a,b)) br
145    | balright z zd (RED(x,xd,a,BLACK(y,yd,b,c))) br =
146      RED(y, yd, lbalance x xd (redden a) b, BLACK(z, zd, c, br))
147    | balright _ _ _ _ = raise RedBlackMapError
148
149
150  (* [append left right] constructs a new tree t.
151  PRECONDITIONS: RB left /\ RB right
152              /\ !e in left => !x in right e < x
153  POSTCONDITION: not (RB t)
154  *)
155  fun append LEAF right                    = right
156    | append left LEAF                     = left
157    | append (RED(x,xd,a,b)) (RED(y,yd,c,d))     =
158      (case append b c of
159           RED(z, zd, b, c) => RED(z, zd, RED(x, xd, a, b), RED(y, yd, c, d))
160         | bc           => RED(x, xd, a, RED(y, yd, bc, d)))
161    | append a (RED(x,xd,b,c))                = RED(x, xd, append a b, c)
162    | append (RED(x,xd,a,b)) c                = RED(x, xd, a, append b c)
163    | append (BLACK(x,xd,a,b)) (BLACK(y,yd,c,d)) =
164      (case append b c of
165           RED(z, zd, b, c) => RED(z, zd, BLACK(x,xd,a,b), BLACK(y,yd,c,d))
166         | bc           => balleft x xd a (BLACK(y, yd, bc, d)))
167
168  fun remove ((compare, tree, n), key) =
169      let fun delShared k x a b =
170              case compare(key, k) of
171                  EQUAL   => (x, append a b)
172                | LESS    =>
173                  let val (res, a') = del a
174                  in  (res, case a of
175                                BLACK _ => balleft k x a' b
176                              | _       => RED(k, x, a', b)) end
177                | GREATER =>
178                  let val (res, b') = del b
179                  in  (res, case b of
180                                BLACK _ => balright k x a b'
181                              | _       => RED(k, x, a, b')) end
182          and del LEAF                = raise NotFound
183            | del (RED(k, x, a, b))   = delShared k x a b
184            | del (BLACK(k, x, a, b)) = delShared k x a b
185
186          val (res, tree) = case del tree of
187                                (res, RED arg) => (res, BLACK arg)
188                              | x              => x
189      in  ((compare, tree, n-1), res) end
190
191  fun map f (compare, tree, n) =
192      let fun loop LEAF = LEAF
193            | loop (RED(k,x,a,b)) =
194              let val a = loop a
195                  val x = f(k,x)
196              in  RED(k,x,a, loop b) end
197            | loop (BLACK(k,x,a,b)) =
198              let val a = loop a
199                  val x = f(k,x)
200              in  BLACK(k,x,a, loop b) end
201      in  (compare, loop tree, n) end
202
203  fun transform f (compare, tree, n) =
204      let fun loop LEAF = LEAF
205            | loop (RED(k,x,a,b)) =
206              let val a = loop a
207              in  RED(k, f x, a, loop b) end
208            | loop (BLACK(k,x,a,b)) =
209              let val a = loop a
210              in  BLACK(k, f x, a, loop b) end
211      in  (compare, loop tree, n) end
212end
213(*
214val t1 = Redblackset.addList(Redblackset.empty Int.compare, [43,25,13,14]);
215val t2 = Redblackset.addList(Redblackset.empty Int.compare, [43,1,2,3]);
216val t3 = Redblackset.addList(Redblackset.empty Int.compare, [1,3]);
217*)
218