9873463ec03762f6e5541bf603bed209cc78a2c7
[ghc.git] / testsuite / tests / indexed-types / should_compile / red-black-delete.hs
1 {-# LANGUAGE InstanceSigs,GADTs, DataKinds, KindSignatures, MultiParamTypeClasses, FlexibleInstances, TypeFamilies #-}
2
3 -- Implementation of deletion for red black trees by Matt Might
4 -- Editing to preserve the red/black tree invariants by Stephanie Weirich,
5 -- Dan Licata and John Hughes, at WG2.8 August 2014
6
7 -- Original available from:
8 -- http://matt.might.net/articles/red-black-delete/code/RedBlackSet.hs
9 -- Slides:
10 -- http://matt.might.net/papers/might2014redblack-talk.pdf
11
12 module MightRedBlackGADT where
13
14 import Prelude hiding (max)
15 -- import Test.QuickCheck hiding (elements)
16 import Data.List(nub,sort)
17 import Control.Monad(liftM)
18 import Data.Type.Equality
19 import Data.Maybe(isJust)
20
21 --
22 data Color = Red | Black | DoubleBlack | NegativeBlack
23
24 -- Singleton type, connects the type level color to data
25 -- not necessary in a full-spectrum dependently-typed language
26 data SColor (c :: Color) where
27 R :: SColor Red -- red
28 B :: SColor Black -- black
29 BB :: SColor DoubleBlack -- double black
30 NB :: SColor NegativeBlack -- negative black
31
32
33 -- natural numbers for tracking the black height
34 data Nat = Z | S Nat
35
36 -- Well-formed Red/Black trees
37 -- n statically tracks the black height of the tree
38 -- c statically tracks the color of the root node
39 data CT (n :: Nat) (c :: Color) (a :: *) where
40 E :: CT Z Black a
41 T :: Valid c c1 c2 =>
42 SColor c -> (CT n c1 a) -> a -> (CT n c2 a) -> CT (Incr c n) c a
43
44 -- this type class forces the children of red nodes to be black
45 -- and also excludes double-black and negative-black nodes from
46 -- a valid red-black tree.
47 class Valid (c :: Color) (c1 :: Color) (c2 :: Color) where
48 instance Valid Red Black Black
49 instance Valid Black c1 c2
50
51 -- incrementing the black height, based on the color
52 type family Incr (c :: Color) (n :: Nat) :: Nat
53 type instance Incr Black n = S n
54 type instance Incr DoubleBlack n = S (S n)
55 type instance Incr Red n = n
56 type instance Incr NegativeBlack (S n) = n
57
58 -- the top level type of red-black trees
59 -- the data constructor forces the root to be black and
60 -- hides the black height.
61 data RBSet a where
62 Root :: (CT n Black a) -> RBSet a
63
64 -- We can't automatically derive show and equality
65 -- methods for GADTs.
66 instance Show (SColor c) where
67 show R = "R"
68 show B = "B"
69 show BB = "BB"
70 show NB = "NB"
71
72 instance Show a => Show (CT n c a) where
73 show E = "E"
74 show (T c l x r) =
75 "(T " ++ show c ++ " " ++ show l ++ " "
76 ++ show x ++ " " ++ show r ++ ")"
77 instance Show a => Show (RBSet a) where
78 show (Root x) = show x
79
80
81 showT :: CT n c a -> String
82 showT E = "E"
83 showT (T c l x r) =
84 "(T " ++ show c ++ " " ++ showT l ++ " "
85 ++ "..." ++ " " ++ showT r ++ ")"
86
87
88 -- the test equality class gives us a proof that type level
89 -- colors are equal.
90 -- testEquality :: SColor c1 -> SColor c2 -> Maybe (c1 ~ c2)
91 instance TestEquality SColor where
92 testEquality R R = Just Refl
93 testEquality B B = Just Refl
94 testEquality BB BB = Just Refl
95 testEquality NB NB = Just Refl
96 testEquality _ _ = Nothing
97
98 -- comparing two Red/Black trees for equality.
99 -- unfortunately, because we have two instances, we can't
100 -- use the testEquality class.
101 (%==%) :: Eq a => CT n1 c1 a -> CT n2 c2 a -> Bool
102 E %==% E = True
103 T c1 a1 x1 b1 %==% T c2 a2 x2 b2 =
104 isJust (testEquality c1 c2) && a1 %==% a2 && x1 == x2 && b1 %==% b2
105
106 instance Eq a => Eq (RBSet a) where
107 (Root t1) == (Root t2) = t1 %==% t2
108
109
110 -- Public operations --
111
112 empty :: RBSet a
113 empty = Root E
114
115 member :: (Ord a) => a -> RBSet a -> Bool
116 member x (Root t) = aux x t where
117 aux :: Ord a => a -> CT n c a -> Bool
118 aux x E = False
119 aux x (T _ l y r) | x < y = aux x l
120 | x > y = aux x r
121 | otherwise = True
122
123 elements :: Ord a => RBSet a -> [a]
124 elements (Root t) = aux t [] where
125 aux :: Ord a => CT n c a -> [a] -> [a]
126 aux E acc = acc
127 aux (T _ a x b) acc = aux a (x : aux b acc)
128
129 -- INSERTION --
130
131 insert :: (Ord a) => a -> RBSet a -> RBSet a
132 insert x (Root s) = blacken (ins x s)
133 where ins :: Ord a => a -> CT n c a -> IR n a
134 ins x E = IN R E x E
135 ins x s@(T color a y b) | x < y = balanceL color (ins x a) y b
136 | x > y = balanceR color a y (ins x b)
137 | otherwise = (IN color a y b)
138
139
140 blacken :: IR n a -> RBSet a
141 blacken (IN _ a x b) = Root (T B a x b)
142
143 -- an intermediate data structure that temporarily violates the
144 -- red-black tree invariants during insertion. (IR stands for infra-red).
145 -- This tree must be non-empty, but is allowed to have two red nodes in
146 -- a row.
147 data IR n a where
148 IN :: SColor c -> CT n c1 a -> a -> CT n c2 a -> IR (Incr c n) a
149
150 -- `balance` rotates away coloring conflicts
151 -- we separate it into two cases based on whether the infraction could
152 -- be on the left or the right
153
154 balanceL :: SColor c -> IR n a -> a -> CT n c1 a -> IR (Incr c n) a
155 balanceL B (IN R (T R a x b) y c) z d = IN R (T B a x b) y (T B c z d)
156 balanceL B (IN R a x (T R b y c)) z d = IN R (T B a x b) y (T B c z d)
157 balanceL c (IN B a x b) z d = IN c (T B a x b) z d
158 balanceL c (IN R a@(T B _ _ _) x b@(T B _ _ _)) z d = IN c (T R a x b) z d
159 balanceL c (IN R a@E x b@E) z d = IN c (T R a x b) z d
160
161
162 balanceR :: SColor c -> CT n c1 a -> a -> IR n a -> IR (Incr c n) a
163 balanceR B a x (IN R (T R b y c) z d) = IN R (T B a x b) y (T B c z d)
164 balanceR B a x (IN R b y (T R c z d)) = IN R (T B a x b) y (T B c z d)
165 balanceR c a x (IN B b z d) = IN c a x (T B b z d)
166 balanceR c a x (IN R b@(T B _ _ _) z d@(T B _ _ _)) = IN c a x (T R b z d)
167 balanceR c a x (IN R b@E z d@E) = IN c a x (T R b z d)
168
169
170 -- DELETION --
171
172 -- like insert, delete relies on a helper function that may (slightly)
173 -- violate the red-black tree invariant
174
175 delete :: (Ord a) => a -> RBSet a -> RBSet a
176 delete x (Root s) = blacken (del x s)
177 where blacken :: DT n a -> RBSet a
178 blacken (DT B a x b) = Root (T B a x b)
179 blacken (DT BB a x b) = Root (T B a x b)
180 blacken DE = Root E
181 blacken DEE = Root E
182 -- note: del always produces a black or
183 -- double black tree. These last two cases are unreachable
184 -- but it may be difficult to prove it.
185 blacken (DT R a x b) = Root (T B a x b)
186 blacken (DT NB a x b) = Root (T B a x b)
187
188 del :: Ord a => a -> CT n c a -> DT n a
189 del x E = DE
190 del x s@(T color a y b) | x < y = bubbleL color (del x a) y b
191 | x > y = bubbleR color a y (del x b)
192 | otherwise = remove s
193
194
195 -- deletion tree, the result of del
196 -- could have a double black node
197 -- (but only at the top or as a single leaf)
198 data DT n a where
199 DT :: SColor c -> CT n c1 a -> a -> CT n c2 a -> DT (Incr c n) a
200 DE :: DT Z a
201 DEE :: DT (S Z) a
202
203 -- Note: we could unify this with the IR data structure by
204 -- adding the DE and DEE data constructors. That would allow us to
205 -- reuse the same balancing function (and blacken function)
206 -- for both insertion and deletion.
207 -- However, if an Agda version did that it might get into trouble,
208 -- trying to show that insertion never produces a leaf.
209
210
211 instance Show a => Show (DT n a) where
212 show DE = "DE"
213 show DEE = "DEE"
214 show (DT c l x r) =
215 "(DT " ++ show c ++ " " ++ show l ++ " "
216 ++ "..." ++ " " ++ show r ++ ")"
217
218 -- Note: eliding a to make error easier below.
219 showDT :: DT n a -> String
220 showDT DE = "DE"
221 showDT DEE = "DEE"
222 showDT (DT c l x r) =
223 "(DT " ++ show c ++ " " ++ showT l ++ " "
224 ++ "..." ++ " " ++ showT r ++ ")"
225
226
227 -- maximum element from a red-black tree
228 max :: CT n c a -> a
229 max E = error "no largest element"
230 max (T _ _ x E) = x
231 max (T _ _ x r) = max r
232
233 -- Remove the top node: it might leave behind a double black node
234 -- if the removed node is black (note that the height is preserved)
235 remove :: CT n c a -> DT n a
236 remove (T R E _ E) = DE
237 remove (T B E _ E) = DEE
238 remove (T B E _ (T R a x b)) = DT B a x b
239 remove (T B (T R a x b) _ E) = DT B a x b
240 remove (T color l y r) = bubbleL color (removeMax l) (max l) r
241
242 -- remove the maximum element of the tree
243 removeMax :: CT n c a -> DT n a
244 removeMax E = error "no maximum to remove"
245 removeMax s@(T _ _ _ E) = remove s
246 removeMax s@(T color l x r) = bubbleR color l x (removeMax r)
247
248 -- make a node more red (i.e. decreasing its black height)
249 -- note that the color of the resulting node is not part of the
250 -- result height because the result of this operation doesn't
251 -- necessarily satisfy the RBT invariants
252 redden :: CT (S n) c a -> DT n a
253 redden (T B a x y) = DT R a x y
254 redden (T BB a x y) = DT B a x y
255 redden (T R a x y) = DT NB a x y
256
257 -- assert that a tree is red. Dynamically fail otherwise
258 -- this is a cheat that Haskell let's us get away with, but Agda
259 -- would not
260 assertRed :: DT n a -> CT n Red a
261 assertRed (DT R (T B aa ax ay) x (T B ac az ad)) =
262 (T R (T B aa ax ay) x (T B ac az ad))
263 assertRed t = error ("not a red tree " ++ showDT t)
264
265 dbalanceL :: SColor c -> DT n a -> a -> CT n c1 a -> DT (Incr c n) a
266 -- reshuffle (same as insert)
267 dbalanceL B (DT R (T R a x b) y c) z d = DT R (T B a x b) y (T B c z d)
268 dbalanceL B (DT R a x (T R b y c)) z d = DT R (T B a x b) y (T B c z d)
269 -- double-black (new for delete)
270 dbalanceL BB (DT R (T R a x b) y c) z d = DT B (T B a x b) y (T B c z d)
271 dbalanceL BB (DT R a x (T R b y c)) z d = DT B (T B a x b) y (T B c z d)
272 dbalanceL BB (DT NB a@(T B _ _ _) x (T B b y c)) z d =
273 case (dbalanceL B (redden a) x b) of
274 r@(DT R _ _ _) -> DT B (assertRed r) y (T B c z d)
275 (DT B a1 x1 y1) -> DT B (T B a1 x1 y1) y (T B c z d)
276 -- fall through cases (same as insert)
277 dbalanceL c (DT B a x b) z d = DT c (T B a x b) z d
278 dbalanceL c (DT R a@(T B _ _ _) x b@(T B _ _ _)) z d = DT c (T R a x b) z d
279 dbalanceL c (DT R a@E x b@E) z d = DT c (T R a x b) z d
280 -- more fall through cases (new for delete)
281 dbalanceL c DE x b = (DT c E x b)
282 dbalanceL c a x b = error ("no case for " ++ showDT a)
283
284
285 dbalanceR :: SColor c -> CT n c1 a -> a -> DT n a -> DT (Incr c n) a
286 dbalanceR B a x (DT R (T R b y c) z d) = DT R (T B a x b) y (T B c z d)
287 dbalanceR B a x (DT R b y (T R c z d)) = DT R (T B a x b) y (T B c z d)
288 dbalanceR BB a x (DT R (T R b y c) z d) = DT B (T B a x b) y (T B c z d)
289 dbalanceR BB a x (DT R b y (T R c z d)) = DT B (T B a x b) y (T B c z d)
290 dbalanceR BB a x (DT NB (T B b y c) z d@(T B _ _ _)) =
291 case (dbalanceR B c z (redden d)) of
292 r@(DT R _ _ _) -> DT B (T B a x b) y (assertRed r)
293 (DT B a1 x1 y1) -> DT B (T B a x b) y (T B a1 x1 y1)
294
295 dbalanceR c a x (DT B b z d) = DT c a x (T B b z d)
296 dbalanceR c a x (DT R b@(T B _ _ _) z d@(T B _ _ _)) = DT c a x (T R b z d)
297 dbalanceR c a x (DT R b@E z d@E) = DT c a x (T R b z d)
298 dbalanceR c a x DE = DT c a x E
299 dbalanceR c a x b = error ("no case for " ++ showDT b)
300
301
302 bubbleL :: SColor c -> DT n a -> a -> CT n c1 a -> DT (Incr c n) a
303 -- don't want to prove generally that
304 -- (Incr (Blacker c) 'Z ~ Incr c ('S 'Z)) so expand by cases.
305 -- (Note: Do we need all three of these cases?)
306 bubbleL B DEE x r = dbalanceR BB E x (redden r)
307 bubbleL R DEE x r = dbalanceR B E x (redden r)
308 bubbleL NB DEE x r = dbalanceR R E x (redden r)
309 -- same thing here
310 bubbleL B l@(DT BB a y b) x r = dbalanceR BB (T B a y b) x (redden r)
311 bubbleL R l@(DT BB a y b) x r = dbalanceR B (T B a y b) x (redden r)
312 bubbleL NB l@(DT BB a y b) x r = dbalanceR R (T B a y b) x (redden r)
313 -- fall through case
314 bubbleL color l x r = dbalanceL color l x r
315
316
317 bubbleR :: SColor c -> CT n c1 a -> a -> DT n a -> DT (Incr c n) a
318 bubbleR B r x DEE = dbalanceL BB (redden r) x E
319 bubbleR R r x DEE = dbalanceL B (redden r) x E
320 bubbleR NB r x DEE = dbalanceL R (redden r) x E
321 -- same thing here
322 bubbleR B r x l@(DT BB a y b) = dbalanceL BB (redden r) x (T B a y b)
323 bubbleR R r x l@(DT BB a y b) = dbalanceL B (redden r) x (T B a y b)
324 bubbleR NB r x l@(DT BB a y b) = dbalanceL R (redden r) x (T B a y b)
325 -- fall through case
326 bubbleR color l x r = dbalanceR color l x r
327