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  fun update (set as (compare, tree, n), key, data) =
53      let val addone = ref true
54          fun ins LEAF = RED(key,data NONE,LEAF,LEAF)
55            | ins (BLACK(k,x,left,right)) =
56              (case compare(key, k) of
57                   LESS    => lbalance k x (ins left) right
58                 | GREATER => rbalance k x left (ins right)
59                 | EQUAL   => (addone := false; BLACK(key, data (SOME x), left, right)))
60            | ins (RED(k, x,left,right)) =
61              (case compare(key, k) of
62                   LESS    => RED(k, x, (ins left), right)
63                 | GREATER => RED(k, x, left, (ins right))
64                 | EQUAL   => (addone := false; RED(key, data (SOME x), left, right)))
65      in  ( compare
66          , case ins tree of
67                RED x => BLACK x
68              | tree  => tree
69          , if !addone then n+1 else n) end
70
71  local fun K x _ = x in
72    fun insert (set, key, data) = update (set, key, K data)
73  end
74
75  fun insertList (m, xs) =
76    List.foldl (fn ((i, v), m) => insert (m, i, v)) m xs
77
78  fun fromList compare xs =
79    insertList (mkDict compare, xs)
80
81  fun push LEAF stack = stack
82    | push tree stack = tree :: stack
83
84  fun pushNode left k x right stack =
85      left :: (BLACK(k, x, LEAF, LEAF) :: (push right stack))
86
87  fun getMin []             some none = none
88    | getMin (tree :: rest) some none =
89      case tree of
90          LEAF                 => getMin rest some none
91        | RED  (k, x, LEAF, b) => some k x (push b rest)
92        | BLACK(k, x, LEAF, b) => some k x (push b rest)
93        | RED  (k, x, a, b)    => getMin(pushNode a k x b rest) some none
94        | BLACK(k, x, a, b)    => getMin(pushNode a k x b rest) some none
95
96  fun getMax []             some none = none
97    | getMax (tree :: rest) some none =
98      case tree of
99          LEAF                 => getMax rest some none
100        | RED  (k, x, a, LEAF) => some k x (push a rest)
101        | BLACK(k, x, a, LEAF) => some k x (push a rest)
102        | RED  (k, x, a, b)    => getMax(pushNode b k x a rest) some none
103        | BLACK(k, x, a, b)    => getMax(pushNode b k x a rest) some none
104
105  fun fold get f e (compare, tree, n) =
106      let fun loop stack acc =
107              get stack (fn k =>fn x =>fn stack => loop stack (f(k,x,acc))) acc
108      in  loop [tree] e end
109
110  fun foldl f = fold getMin f
111
112  fun foldr f = fold getMax f
113
114  fun listItems set = foldr (fn(k,x,res) => (k,x)::res) [] set
115
116  fun appAll get f (compare, tree, n) =
117      let fun loop stack = get stack (fn k => fn x => (f(k,x); loop)) ()
118      in  loop [tree] end
119
120  fun app f = appAll getMin f
121
122  fun revapp f = appAll getMax f
123
124
125  exception RedBlackMapError
126
127  (* remove a la Stefan M. Kahrs *)
128  fun redden (BLACK arg) = RED arg
129    | redden _ = raise RedBlackMapError
130
131  fun balleft y yd (RED(x,xd,a,b)) c                =
132      RED(y, yd, BLACK(x, xd, a, b), c)
133    | balleft x xd bl (BLACK(y, yd, a, b))          =
134      rbalance x xd bl (RED(y, yd, a, b))
135    | balleft x xd bl (RED(z,zd,BLACK(y,yd,a,b),c)) =
136      RED(y, yd, BLACK(x, xd, bl, a), rbalance z zd b (redden c))
137    | balleft _ _ _ _ = raise RedBlackMapError
138
139  fun balright x xd a             (RED(y, yd ,b,c)) =
140      RED(x, xd, a, BLACK(y, yd, b, c))
141    | balright y yd (BLACK(x,xd,a,b))          br =
142      lbalance y yd (RED(x,xd,a,b)) br
143    | balright z zd (RED(x,xd,a,BLACK(y,yd,b,c))) br =
144      RED(y, yd, lbalance x xd (redden a) b, BLACK(z, zd, c, br))
145    | balright _ _ _ _ = raise RedBlackMapError
146
147
148  (* [append left right] constructs a new tree t.
149  PRECONDITIONS: RB left /\ RB right
150              /\ !e in left => !x in right e < x
151  POSTCONDITION: not (RB t)
152  *)
153  fun append LEAF right                    = right
154    | append left LEAF                     = left
155    | append (RED(x,xd,a,b)) (RED(y,yd,c,d))     =
156      (case append b c of
157           RED(z, zd, b, c) => RED(z, zd, RED(x, xd, a, b), RED(y, yd, c, d))
158         | bc           => RED(x, xd, a, RED(y, yd, bc, d)))
159    | append a (RED(x,xd,b,c))                = RED(x, xd, append a b, c)
160    | append (RED(x,xd,a,b)) c                = RED(x, xd, a, append b c)
161    | append (BLACK(x,xd,a,b)) (BLACK(y,yd,c,d)) =
162      (case append b c of
163           RED(z, zd, b, c) => RED(z, zd, BLACK(x,xd,a,b), BLACK(y,yd,c,d))
164         | bc           => balleft x xd a (BLACK(y, yd, bc, d)))
165
166  fun remove ((compare, tree, n), key) =
167      let fun delShared k x a b =
168              case compare(key, k) of
169                  EQUAL   => (x, append a b)
170                | LESS    =>
171                  let val (res, a') = del a
172                  in  (res, case a of
173                                BLACK _ => balleft k x a' b
174                              | _       => RED(k, x, a', b)) end
175                | GREATER =>
176                  let val (res, b') = del b
177                  in  (res, case b of
178                                BLACK _ => balright k x a b'
179                              | _       => RED(k, x, a, b')) end
180          and del LEAF                = raise NotFound
181            | del (RED(k, x, a, b))   = delShared k x a b
182            | del (BLACK(k, x, a, b)) = delShared k x a b
183
184          val (res, tree) = case del tree of
185                                (res, RED arg) => (res, BLACK arg)
186                              | x              => x
187      in  ((compare, tree, n-1), res) end
188
189  fun map f (compare, tree, n) =
190      let fun loop LEAF = LEAF
191            | loop (RED(k,x,a,b)) =
192              let val a = loop a
193                  val x = f(k,x)
194              in  RED(k,x,a, loop b) end
195            | loop (BLACK(k,x,a,b)) =
196              let val a = loop a
197                  val x = f(k,x)
198              in  BLACK(k,x,a, loop b) end
199      in  (compare, loop tree, n) end
200
201  fun transform f (compare, tree, n) =
202      let fun loop LEAF = LEAF
203            | loop (RED(k,x,a,b)) =
204              let val a = loop a
205              in  RED(k, f x, a, loop b) end
206            | loop (BLACK(k,x,a,b)) =
207              let val a = loop a
208              in  BLACK(k, f x, a, loop b) end
209      in  (compare, loop tree, n) end
210end
211(*
212val t1 = Redblackset.addList(Redblackset.empty Int.compare, [43,25,13,14]);
213val t2 = Redblackset.addList(Redblackset.empty Int.compare, [43,1,2,3]);
214val t3 = Redblackset.addList(Redblackset.empty Int.compare, [1,3]);
215*)
216