Fold testsuite.git into ghc.git (re #8545)
[ghc.git] / testsuite / tests / polykinds / RedBlack.hs
1 -- From http://www.reddit.com/r/haskell/comments/ti5il/redblack_trees_in_haskell_using_gadts_existential/
2
3 {-# LANGUAGE GADTs #-}
4 {-# LANGUAGE RankNTypes #-}
5 {-# LANGUAGE StandaloneDeriving #-}
6 {-# LANGUAGE DataKinds#-}
7 {-# LANGUAGE KindSignatures#-}
8 module RedBlackTree where
9
10 data Nat = Zero | Succ Nat deriving (Eq, Ord, Show)
11 type One = Succ Zero
12
13 data RedBlack = Black | Red deriving (Eq, Ord, Show)
14
15 -- red-black trees are rooted at a black node
16 data RedBlackTree a = forall n. T ( Node Black n a )
17 deriving instance Show a => Show (RedBlackTree a)
18
19 -- all paths from a node to a leaf have exactly n black nodes
20 data Node :: RedBlack -> Nat -> * -> * where
21 -- all leafs are black
22 Leaf :: Node Black One a
23 -- internal black nodes can have children of either color
24 B :: Node cL n a -> a -> Node cR n a -> Node Black (Succ n) a
25 -- internal red nodes can only have black children
26 R :: Node Black n a -> a -> Node Black n a -> Node Red n a
27 deriving instance Show a => Show (Node c n a)
28
29 -- one-hole context for red-black trees
30 data Context :: Nat -> RedBlack -> Nat -> * -> * where
31 -- if we're at the root, the hole is a black node
32 Root :: Context n Black n a
33 -- we can go left or right from a red node hole, creating a hole for a black node
34 BC :: Bool -> a -> Node Black n a -> Context m Red n a -> Context m Black n a
35 -- we can go left or right from a black node hole, creating a hole for either
36 EC :: Bool -> a -> Node cY n a -> Context m Black (Succ n) a -> Context m cX n a
37 deriving instance Show a => Show (Context m c n a)
38
39 data Zipper m a = forall c n. Zipper (Node c n a) (Context m c n a)
40 deriving instance Show a => Show (Zipper m a)
41
42 -- create a zipper
43 unZip :: Node Black n a -> Zipper n a
44 unZip = flip Zipper Root
45
46 -- destroy a zipper
47 zipUp :: Zipper m a -> Node Black m a
48 zipUp (Zipper x Root) = x
49 zipUp (Zipper x (BC goLeft a y c)) = zipUp $ Zipper (if goLeft then R x a y else R y a x) c
50 zipUp (Zipper x (EC goLeft a y c)) = zipUp $ Zipper (if goLeft then B x a y else B y a x) c
51
52 -- locate the node that should contain a in the red-black tree
53 zipTo :: Ord a => a -> Zipper n a -> Zipper n a
54 zipTo _ z@(Zipper Leaf _) = z
55 zipTo a z@(Zipper (R l a' r) c) = case compare a a' of
56 EQ -> z
57 LT -> zipTo a $ Zipper l (BC True a' r c)
58 GT -> zipTo a $ Zipper r (BC False a' l c)
59 zipTo a z@(Zipper (B l a' r) c) = case compare a a' of
60 EQ -> z
61 LT -> zipTo a $ Zipper l (EC True a' r c)
62 GT -> zipTo a $ Zipper r (EC False a' l c)
63
64 -- create a red-black tree
65 empty :: RedBlackTree a
66 empty = T Leaf
67
68 -- insert a node into a red-black tree
69 -- (see http://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Insertion)
70 insert :: Ord a => a -> RedBlackTree a -> RedBlackTree a
71 insert a t@(T root) = case zipTo a (unZip root) of
72 -- find matching leaf and replace with red node (pointing to two leaves)
73 Zipper Leaf c -> insertAt (R Leaf a Leaf) c
74 -- if it's already in the tree, there's no need to modify it
75 _ -> t
76
77 insertAt :: Node Red n a -> Context m c n a -> RedBlackTree a
78 -- 1) new node is root => paint it black and done
79 insertAt (R l a r) Root = T $ B l a r
80 -- 2) new node's parent is black => done
81 insertAt x (EC b a y c) = T . zipUp $ Zipper x (EC b a y c)
82 -- 3) uncle is red => paint parent/uncle black, g'parent red. recurse on g'parent
83 insertAt x (BC pb pa py (EC gb ga (R ul ua ur) gc)) = insertAt g gc
84 where p = if pb then B x pa py else B py pa x
85 u = B ul ua ur
86 g = if gb then R p ga u else R u ga p
87 -- 4) node is between parent and g'parent => inner rotation
88 insertAt (R l a r) (BC False pa py pc@(EC True _ _ _)) = insertAt (R py pa l) (BC True a r pc)
89 insertAt (R l a r) (BC True pa py pc@(EC False _ _ _)) = insertAt (R r pa py) (BC False a l pc)
90
91 -- 5) otherwise => outer rotation
92 insertAt x (BC True pa py (EC True ga gy@Leaf gc)) =
93 T . zipUp $ Zipper (B x pa $ R py ga gy) gc
94 insertAt x (BC True pa py (EC True ga gy@(B _ _ _) gc)) =
95 T . zipUp $ Zipper (B x pa $ R py ga gy) gc
96 -- XXX: GHC seems unable to infer that gy is Black so I have to do both cases
97 -- explicitly, rather than
98 -- insertAt x (BC False pa py (EC False ga gy gc)) =
99 -- T . zipUp $ Zipper (B (R gy ga py) pa x) gc
100 --
101 -- insertAt x (BC True pa py (EC True ga gy gc)) =
102 -- T . zipUp $ Zipper (B x pa $ R py ga gy) gc
103 {-
104 BC :: Bool -> a -> Node Black n a -> Context m Red n a -> Context m Black n a
105 EC :: Bool -> a -> Node cY n a -> Context m Black (Succ n) a -> Context m cX n a
106
107
108 BC True pa py (EC True ga gy gc) :: Context m c n a
109 Hence c~Black
110 EC True ga gy gc :: Context m Red n a
111 gy :: Node cY n a
112 -}
113
114 insertAt x (BC False pa py (EC False ga gy@Leaf gc)) =
115 T . zipUp $ Zipper (B (R gy ga py) pa x) gc
116 insertAt x (BC False pa py (EC False ga gy@(B _ _ _) gc)) =
117 T . zipUp $ Zipper (B (R gy ga py) pa x) gc
118
119 -- can't derive, since we abstract over n, so we have to manually
120 -- check for identical structure
121 instance Eq a => Eq (RedBlackTree a) where
122 T Leaf == T Leaf = True
123 T (B l@(B _ _ _) a r@(B _ _ _)) == T (B l'@(B _ _ _) a' r'@(B _ _ _)) =
124 a == a' && T l == T l' && T r == T r'
125 T (B (R ll la lr) a r@(B _ _ _)) == T (B (R ll' la' lr') a' r'@(B _ _ _)) =
126 a == a' && la == la' &&
127 T ll == T ll' && T lr == T lr' && T r == T r'
128 T (B l@(B _ _ _) a r@(R rl ra rr)) == T (B l'@(B _ _ _) a' r'@(R rl' ra' rr')) =
129 a == a' && ra == ra' &&
130 T l == T l' && T rl == T rl' && T rr == T rr'
131 T (B (R ll la lr) a r@(R rl ra rr)) == T (B (R ll' la' lr') a' r'@(R rl' ra' rr')) =
132 a == a' && la == la' && ra == ra' &&
133 T ll == T ll' && T lr == T lr' && T rl == T rl' && T rr == T rr'
134 _ == _ = False
135
136 -- can't derive, since B abstracts over child node colors, so
137 -- manually check for identical structure
138 instance (Eq a) => Eq (Node c n a) where
139 Leaf == Leaf = True
140 R l a r == R l' a' r' = a == a' && l == l' && r == r'
141 b@(B _ _ _) == b'@(B _ _ _) = T b == T b'
142 _ == _ = False