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