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