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