@@ -0,0 +1,236 @@
1+module Network.Wai.Handler.Warp.MultiMap (
2+ MMap
3+ , Some(..)
4+ , empty
5+ , singleton
6+ , insert
7+ , search
8+ , searchWith
9+ , isEmpty
10+ , valid
11+ , pruneWith
12+ , fromList
13+ , toList
14+ , fromSortedList
15+ , toSortedList
16+ , merge
17+ ) where
18+
19+import Control.Applicative ((<$>))
20+import Data.List (foldl')
21+
22+----------------------------------------------------------------
23+
24+-- | One ore more list to implement multimap.
25+data Some a = One !a
26+ | Tom !a !(Some a) -- Two or more
27+ deriving (Eq,Show)
28+
29+-- This is slow but assuming rarely used.
30+snoc :: Some a -> a -> Some a
31+snoc (One x) y = Tom x (One y)
32+snoc (Tom x xs) y = Tom x (snoc xs y)
33+
34+top :: Some a -> a
35+top (One x) = x
36+top (Tom x _) = x
37+
38+----------------------------------------------------------------
39+
40+-- | Red black tree as multimap.
41+data MMap k v = Leaf -- color is Black
42+ | Node Color !(MMap k v) !k !(Some v) !(MMap k v)
43+ deriving (Show)
44+
45+data Color = B -- ^ Black
46+ | R -- ^ Red
47+ deriving (Eq, Show)
48+
49+----------------------------------------------------------------
50+
51+instance (Eq k, Eq v) => Eq (MMap k v) where
52+ t1 == t2 = toSortedList t1 == toSortedList t2
53+
54+----------------------------------------------------------------
55+
56+-- | O(log N)
57+search :: Ord k => k -> MMap k v -> Maybe v
58+search _ Leaf = Nothing
59+search xk (Node _ l k v r) = case compare xk k of
60+ LT -> search xk l
61+ GT -> search xk r
62+ EQ -> Just $ top v
63+
64+-- | O(log N)
65+searchWith :: Ord k => k -> (Some v -> Maybe v) -> MMap k v -> Maybe v
66+searchWith _ _ Leaf = Nothing
67+searchWith xk f (Node _ l k v r) = case compare xk k of
68+ LT -> searchWith xk f l
69+ GT -> searchWith xk f r
70+ EQ -> f v
71+
72+----------------------------------------------------------------
73+
74+-- | O(1)
75+isEmpty :: (Eq k, Eq v) => MMap k v -> Bool
76+isEmpty Leaf = True
77+isEmpty _ = False
78+
79+-- | O(1)
80+empty :: MMap k v
81+empty = Leaf
82+
83+----------------------------------------------------------------
84+
85+-- | O(1)
86+singleton :: Ord k => k -> v -> MMap k v
87+singleton k v = Node B Leaf k (One v) Leaf
88+
89+----------------------------------------------------------------
90+
91+-- | O(log N)
92+insert :: Ord k => k -> v -> MMap k v -> MMap k v
93+insert kx kv t = turnB (insert' kx kv t)
94+
95+insert' :: Ord k => k -> v -> MMap k v -> MMap k v
96+insert' xk xv Leaf = Node R Leaf xk (One xv) Leaf
97+insert' xk xv (Node B l k v r) = case compare xk k of
98+ LT -> balanceL' (insert' xk xv l) k v r
99+ GT -> balanceR' l k v (insert' xk xv r)
100+ EQ -> Node B l k (snoc v xv) r
101+insert' xk xv (Node R l k v r) = case compare xk k of
102+ LT -> Node R (insert' xk xv l) k v r
103+ GT -> Node R l k v (insert' xk xv r)
104+ EQ -> Node R l k (snoc v xv) r
105+
106+balanceL' :: MMap k v -> k -> Some v -> MMap k v -> MMap k v
107+balanceL' (Node R (Node R a xk xv b) yk yv c) zk zv d =
108+ Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
109+balanceL' (Node R a xk xv (Node R b yk yv c)) zk zv d =
110+ Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
111+balanceL' l k v r = Node B l k v r
112+
113+balanceR' :: MMap k v -> k -> Some v -> MMap k v -> MMap k v
114+balanceR' a xk xv (Node R b yk yv (Node R c zk zv d)) =
115+ Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
116+balanceR' a xk xv (Node R (Node R b yk yv c) zk zv d) =
117+ Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
118+balanceR' l xk xv r = Node B l xk xv r
119+
120+turnB :: MMap k v -> MMap k v
121+turnB Leaf = error "turnB"
122+turnB (Node _ l k v r) = Node B l k v r
123+
124+----------------------------------------------------------------
125+
126+-- | O(N log N)
127+fromList :: Ord k => [(k,v)] -> MMap k v
128+fromList = foldl' (\t (k,v) -> insert k v t) empty
129+
130+-- | O(N)
131+toList :: MMap k v -> [(k,v)]
132+toList t = inorder t []
133+ where
134+ inorder Leaf xs = xs
135+ inorder (Node _ l k v r) xs = inorder l (pairs k v ++ inorder r xs)
136+ pairs k (One v) = [(k,v)]
137+ pairs k (Tom v vs) = (k,v) : pairs k vs
138+
139+----------------------------------------------------------------
140+
141+-- | O(N)
142+-- "Constructing Red-Black Trees" by Ralf Hinze
143+fromSortedList :: Ord k => [(k,Some v)] -> MMap k v
144+fromSortedList = linkAll . foldr add []
145+
146+data Digit k v = Uno k (Some v) (MMap k v)
147+ | Due k (Some v) (MMap k v) k (Some v) (MMap k v)
148+ deriving (Eq,Show)
149+
150+incr :: Digit k v -> [Digit k v] -> [Digit k v]
151+incr (Uno k v t) [] = [Uno k v t]
152+incr (Uno k1 v1 t1) (Uno k2 v2 t2 : ps) = Due k1 v1 t1 k2 v2 t2 : ps
153+incr (Uno k1 v1 t1) (Due k2 v2 t2 k3 v3 t3 : ps) = Uno k1 v1 t1 : incr (Uno k2 v2 (Node B t2 k3 v3 t3)) ps
154+incr _ _ = error "incr"
155+
156+add :: (k,Some v) -> [Digit k v] -> [Digit k v]
157+add (k,v) ps = incr (Uno k v Leaf) ps
158+
159+linkAll :: [Digit k v] -> MMap k v
160+linkAll = foldl' link Leaf
161+
162+link :: MMap k v -> Digit k v -> MMap k v
163+link l (Uno k v t) = Node B l k v t
164+--link l (Due k1 v1 t1 k2 v2 t2) = Node B (Node R l k1 v1 t1) k2 v2 t2
165+link l (Due k1 v1 t1 k2 v2 t2) = Node B l k1 v1 (Node R t1 k2 v2 t2)
166+
167+----------------------------------------------------------------
168+
169+-- | O(N)
170+toSortedList :: MMap k v -> [(k,Some v)]
171+toSortedList t = inorder t []
172+ where
173+ inorder Leaf xs = xs
174+ inorder (Node _ l k v r) xs = inorder l ((k,v) : inorder r xs)
175+
176+----------------------------------------------------------------
177+
178+-- | O(N)
179+pruneWith :: Ord k =>
180+ MMap k v
181+ -> (k -> Some v -> IO [(k, Some v)])
182+ -> IO (MMap k v)
183+pruneWith t run = fromSortedList <$> inorder t []
184+ where
185+ inorder Leaf xs = return xs
186+ inorder (Node _ l k v r) xs = do
187+ ys <- run k v
188+ zs <- inorder r xs
189+ inorder l (ys ++ zs)
190+
191+----------------------------------------------------------------
192+
193+-- O(N log N) where N is the size of the second argument
194+merge :: Ord k => MMap k v -> MMap k v -> MMap k v
195+merge base m = foldl' ins base xs
196+ where
197+ ins t (k,v) = insert k v t
198+ xs = toList m
199+
200+----------------------------------------------------------------
201+-- for testing
202+
203+valid :: Ord k => MMap k v -> Bool
204+valid t = isBalanced t && isOrdered t
205+
206+isBalanced :: MMap k v -> Bool
207+isBalanced t = isBlackSame t && isRedSeparate t
208+
209+isBlackSame :: MMap k v -> Bool
210+isBlackSame t = all (n==) ns
211+ where
212+ n:ns = blacks t
213+
214+blacks :: MMap k v -> [Int]
215+blacks = blacks' 0
216+ where
217+ blacks' n Leaf = [n+1]
218+ blacks' n (Node R l _ _ r) = blacks' n l ++ blacks' n r
219+ blacks' n (Node B l _ _ r) = blacks' n' l ++ blacks' n' r
220+ where
221+ n' = n + 1
222+
223+isRedSeparate :: MMap k v -> Bool
224+isRedSeparate = reds B
225+
226+reds :: Color -> MMap k v -> Bool
227+reds _ Leaf = True
228+reds R (Node R _ _ _ _) = False
229+reds _ (Node c l _ _ r) = reds c l && reds c r
230+
231+isOrdered :: Ord k => MMap k v -> Bool
232+isOrdered t = ordered $ toSortedList t
233+ where
234+ ordered [] = True
235+ ordered [_] = True
236+ ordered (x:y:xys) = fst x <= fst y && ordered (y:xys)