@@ -1,13 +1,299 @@
11
2−empty = []
3−singleton k v = [(k,v)]
4−insert k v dict = (k,v):dict
2+module Map (empty,singleton,insert,lookup,remove,member,fold) where
53
6−fst pair = let (a,b) = pair in a
4+import Data.Maybe (isJust)
75
8−member key = exists (\pair -> fst pair == key)
6+data NColor = Red | Black
97
10−delete key = filter (\pair -> fst pair == key)
11−intersect d1 d2 = filter (\pair -> member (fst pair) d2) d1
12−difference d1 d2 = filter (\pair -> not (member (fst pair) d2)) d1
13−union d1 d2 = d1 ++ difference d2 d1
8+data RBTree k v = Node NColor k v (RBTree k v) (RBTree k v) | Empty
9+
10+raise = console.log
11+
12+empty = Empty
13+
14+-- Helpers for checking invariants
15+
16+-- Check that the tree has an equal number of black nodes on each path
17+equal_pathLen t =
18+ let path_numBlacks t =
19+ case t of
20+ { Empty -> 1
21+ ; Node col _ _ l r ->
22+ let { bl = path_numBlacks l ; br = path_numBlacks r } in
23+ if bl /= br || bl == 0-1 || br == 0-1
24+ then 0-1
25+ else bl + (if col == Red then 0 else 1)
26+ }
27+ in 0-1 /= path_numBlacks t
28+
29+rootBlack t =
30+ case t of
31+ { Empty -> True
32+ ; Node Black _ _ _ _ -> True
33+ ; _ -> False }
34+
35+redBlack_children t =
36+ case t of
37+ { Node Red _ _ (Node Red _ _ _ _) _ -> False
38+ ; Node Red _ _ _ (Node Red _ _ _ _) -> False
39+ ; Empty -> True
40+ ; Node _ _ _ l r -> redBlack_children l && redBlack_children r
41+ }
42+
43+findExtreme f t =
44+ case t of
45+ { Empty -> Nothing
46+ ; Node c k _ l r ->
47+ case findExtreme f (f (l,r)) of
48+ { Nothing -> Just k
49+ ; Just k' -> Just k' }
50+ }
51+
52+findminRbt t = findExtreme fst t
53+findmaxRbt t = findExtreme snd t
54+
55+-- "Option LT than"
56+-- Returns True if either xo or yo is Nothing
57+-- Otherwise returns the result of comparing the values using f
58+
59+optionRelation f u xo yo =
60+ case (xo,yo) of
61+ { (Nothing,_) -> u
62+ ; (_,Nothing) -> u
63+ ; (Just x, Just y) -> f x y }
64+
65+olt xo yo = optionRelation (< ) True xo yo
66+olte xo yo = optionRelation (<=) True xo yo
67+
68+ordered t =
69+ case t of
70+ { Empty -> True
71+ ; Node c k v l r ->
72+ let (lmax,rmin) = (findmaxRbt l, findminRbt r) in
73+ olte lmax (Just k) && olte (Just k) rmin && ordered l && ordered r
74+ }
75+
76+-- Check that there aren't any right red nodes in the tree *)
77+leftLeaning t =
78+ case t of
79+ { Empty -> True
80+ ; Node _ _ _ (Node Black _ _ _ _) (Node Red _ _ _ _) -> False
81+ ; Node _ _ _ Empty (Node Red _ _ _ _) -> False
82+ ; Node _ _ _ l r -> (leftLeaning l) && (leftLeaning r)
83+ }
84+
85+invariants_hold t =
86+ ordered t && rootBlack t && redBlack_children t &&
87+ equal_pathLen t && leftLeaning t
88+
89+--** End invariant helpers *****
90+
91+
92+min t =
93+ case t of
94+ { Node _ k v Empty _ -> (k,v)
95+ ; Node _ _ _ l _ -> min l
96+ ; Empty -> console.log "(min Empty) is not defined"
97+ }
98+
99+max t =
100+ case t of
101+ { Node _ k v _ Empty -> (k,v)
102+ ; Node _ _ _ _ r -> max r
103+ ; Empty -> console.log "(max Empty) is not defined"
104+ }
105+
106+lookup k t =
107+ case t of
108+ { Empty -> Nothing
109+ ; Node _ k' v l r ->
110+ case compare k k' of
111+ { LT -> lookup k l
112+ ; EQ -> Just v
113+ ; GT -> lookup k r }
114+ }
115+
116+-- Does t contain k?
117+member k t = isJust $ lookup k t
118+
119+rotateLeft t =
120+ case t of
121+ { Node cy ky vy a (Node cz kz vz b c) -> Node cy kz vz (Node Red ky vy a b) c
122+ ; _ -> raise "rotateLeft of a node without enough children" }
123+
124+-- rotateRight -- the reverse, and
125+-- makes Y have Z's color, and makes Z Red.
126+rotateRight t =
127+ case t of
128+ { Node cz kz vz (Node cy ky vy a b) c -> Node cz ky vy a (Node Red kz vz b c)
129+ ; _ -> raise "rotateRight of a node without enough children" }
130+
131+rotateLeftIfNeeded t =
132+ case t of
133+ { Node _ _ _ _ (Node Red _ _ _ _) -> rotateLeft t
134+ ; _ -> t }
135+
136+rotateRightIfNeeded t =
137+ case t of
138+ { Node _ _ _ (Node Red _ _ (Node Red _ _ _ _) _) _ -> rotateRight t
139+ ; _ -> t }
140+
141+otherColor c = case c of { Red -> Black ; Black -> Red }
142+
143+color_flip t =
144+ case t of
145+ { Node c1 bk bv (Node c2 ak av la ra) (Node c3 ck cv lc rc) ->
146+ Node (otherColor c1) bk bv
147+ (Node (otherColor c2) ak av la ra)
148+ (Node (otherColor c3) ck cv lc rc)
149+ ; _ -> raise "color_flip called on a Empty or Node with a Empty child" }
150+
151+color_flipIfNeeded t =
152+ case t of
153+ { Node _ _ _ (Node Red _ _ _ _) (Node Red _ _ _ _) -> color_flip t
154+ ; _ -> t }
155+
156+fixUp t = color_flipIfNeeded (rotateRightIfNeeded (rotateLeftIfNeeded t))
157+
158+
159+ensureBlackRoot t =
160+ case t of
161+ { Node Red k v l r -> Node Black k v l r
162+ ; _ -> t }
163+
164+-- Invariant: t is a valid left-leaning rb tree *)
165+insert k v t =
166+ let ins t =
167+ case t of
168+ { Empty -> Node Red k v Empty Empty
169+ ; Node c k' v' l r ->
170+ let h = case compare k k' of
171+ { LT -> Node c k' v' (ins l) r
172+ ; EQ -> Node c k' v l r -- replace
173+ ; GT -> Node c k' v' l (ins r) }
174+ in fixUp h }
175+ in if not (invariants_hold t) then
176+ raise "invariants broken before insert"
177+ else (let new_t = ensureBlackRoot (ins t) in
178+ if not (invariants_hold new_t) then
179+ raise "invariants broken after insert"
180+ else new_t)
181+
182+singleton k v = insert k v Empty
183+
184+
185+isRed t =
186+ case t of
187+ { Node Red _ _ _ _ -> True
188+ ; _ -> False }
189+
190+isRedLeft t =
191+ case t of
192+ { Node _ _ _ (Node Red _ _ _ _) _ -> True
193+ ; _ -> False }
194+
195+isRedLeftLeft t =
196+ case t of
197+ { Node _ _ _ (Node _ _ _ (Node Red _ _ _ _) _) _ -> True
198+ ; _ -> False }
199+
200+isRedRight t =
201+ case t of
202+ { Node _ _ _ _ (Node Red _ _ _ _) -> True
203+ ; _ -> False }
204+
205+isRedRightLeft t =
206+ case t of
207+ { Node _ _ _ _ (Node _ _ _ (Node Red _ _ _ _) _) -> True
208+ ; _ -> False }
209+
210+
211+moveRedLeft t =
212+ let t' = color_flip t in
213+ case t' of
214+ { Node c k v l r ->
215+ case r of
216+ { Node _ _ _ (Node Red _ _ _ _) _ ->
217+ color_flip (rotateLeft (Node c k v l (rotateRight r)))
218+ ; _ -> t' }
219+ ; _ -> t' }
220+
221+moveRedRight t =
222+ let t' = color_flip t in
223+ if isRedLeftLeft t' then color_flip (rotateRight t') else t'
224+
225+moveRedLeftIfNeeded t =
226+ if not (isRedLeft t) && not (isRedLeftLeft t)
227+ then moveRedLeft t
228+ else t
229+
230+moveRedRightIfNeeded t =
231+ if not (isRedRight t) && not (isRedRightLeft t)
232+ then moveRedRight t
233+ else t
234+
235+deleteMin t =
236+ let del t =
237+ case t of
238+ { Node _ _ _ Empty _ -> Empty
239+ ; _ -> let t' = moveRedLeftIfNeeded t in
240+ case t' of
241+ { Node c k v l r -> fixUp (Node c k v (del l) r)
242+ ; Empty -> Empty }
243+ }
244+ in ensureBlackRoot (del t)
245+
246+deleteMax t =
247+ let del t =
248+ let t' = if isRedLeft t then rotateRight t else t in
249+ case t' of
250+ { Node _ _ _ _ Empty -> Empty
251+ ; _ -> let t'' = moveRedRightIfNeeded t' in
252+ case t'' of
253+ { Node c k v l r -> fixUp (Node c k v l (del r))
254+ ; Empty -> Empty } }
255+ in ensureBlackRoot (del t)
256+
257+remove k t =
258+ let {
259+ eq_and_noRightNode t = case t of { Node _ k' _ _ Empty -> k == k' ; _ -> False }
260+ ; eq t = case t of { Node _ k' _ _ _ -> k == k' ; _ -> False }
261+ ; delLT t =
262+ let t' = moveRedLeftIfNeeded t in
263+ case t' of
264+ { Node c k' v l r -> fixUp (Node c k' v (del l) r)
265+ ; Empty -> raise "delLT on Empty" }
266+ ; delEQ t =
267+ case t of -- Replace with successor
268+ { Node c _ _ l r ->
269+ let (k',v') = min r in
270+ fixUp (Node c k' v' l (deleteMin r))
271+ ; Empty -> raise "delEQ called on a Empty" }
272+ ; delGT t =
273+ case t of
274+ { Node c k' v l r -> fixUp (Node c k' v l (del r))
275+ ; Empty -> raise "delGT called on a Empty" }
276+ ; del t =
277+ case t of
278+ { Empty -> Empty
279+ ; Node _ k' _ _ _ ->
280+ if k < k' then delLT t
281+ else (let t' = if isRedLeft t then rotateRight t else t in
282+ if eq_and_noRightNode t' then Empty
283+ else (let t = moveRedRightIfNeeded t in
284+ if eq t then delEQ t else delGT t)) }
285+ }
286+ in if not (invariants_hold t) then
287+ raise "invariants broken before remove"
288+ else (let t' = ensureBlackRoot (del t) in
289+ if invariants_hold t' then t' else
290+ raise "invariants broken after remove")
291+
292+fold f acc t =
293+ case t of
294+ { Empty -> acc
295+ ; Node _ k v l r -> fold f (f k v (fold f acc l)) r
296+ }
297+
298+{--
299+--}