@@ -9,9 +9,7 @@ module Network.Wai.Handler.Warp.MultiMap (
99 , search
1010 , searchWith
1111 , isEmpty
12− , valid
1312 , pruneWith
14− , fromList
1513 , toList
1614 , fromSortedList
1715 , toSortedList
@@ -21,7 +19,8 @@ module Network.Wai.Handler.Warp.MultiMap (
2119#if __GLASGOW_HASKELL__ < 709
2220import Control.Applicative ((<$>))
2321#endif
24−import Data.List (foldl')
22+import Data.IntMap (IntMap)
23+import qualified Data.IntMap as I
2524
2625----------------------------------------------------------------
2726
@@ -30,211 +29,92 @@ data Some a = One !a
3029 | Tom !a !(Some a) -- Two or more
3130 deriving (Eq,Show)
3231
33−-- This is slow but assuming rarely used.
34−snoc :: Some a -> a -> Some a
35−snoc (One x) y = Tom x (One y)
36−snoc (Tom x xs) y = Tom x (snoc xs y)
32+{-
33+cons :: Some a -> a -> Some a
34+cons s x = Tom x s
35+-}
3736
3837top :: Some a -> a
3938top (One x) = x
4039top (Tom x _) = x
4140
42−----------------------------------------------------------------
43−
44−-- | Red black tree as multimap.
45−data MMap k v = Leaf -- color is Black
46− | Node Color !(MMap k v) !k !(Some v) !(MMap k v)
47− deriving (Show)
48−
49−data Color = B -- ^ Black
50− | R -- ^ Red
51− deriving (Eq, Show)
41+union :: Some a -> Some a -> Some a
42+union (One x) s = Tom x s
43+union (Tom x xs) s = union xs (Tom x s)
5244
5345----------------------------------------------------------------
5446
55−instance (Eq k, Eq v) => Eq (MMap k v) where
56− t1 == t2 = toSortedList t1 == toSortedList t2
47+type MMap v = IntMap (Some v)
5748
5849----------------------------------------------------------------
5950
60−-- | O(log N)
61−search :: Ord k => k -> MMap k v -> Maybe v
62−search _ Leaf = Nothing
63−search xk (Node _ l k v r) = case compare xk k of
64− LT -> search xk l
65− GT -> search xk r
66− EQ -> Just $ top v
67−
68−-- | O(log N)
69−searchWith :: Ord k => k -> (Some v -> Maybe v) -> MMap k v -> Maybe v
70−searchWith _ _ Leaf = Nothing
71−searchWith xk f (Node _ l k v r) = case compare xk k of
72− LT -> searchWith xk f l
73− GT -> searchWith xk f r
74− EQ -> f v
51+-- | O(log n)
52+search :: Int -> MMap v -> Maybe v
53+search k m = case I.lookup k m of
54+ Nothing -> Nothing
55+ Just s -> Just $! top s
56+
57+-- | O(log n)
58+searchWith :: Int -> (Some v -> Maybe v) -> MMap v -> Maybe v
59+searchWith k f m = case I.lookup k m of
60+ Nothing -> Nothing
61+ Just s -> f s
7562
7663----------------------------------------------------------------
7764
7865-- | O(1)
79−isEmpty :: MMap k v -> Bool
80−isEmpty Leaf = True
81−isEmpty _ = False
66+isEmpty :: MMap v -> Bool
67+isEmpty = I.null
8268
8369-- | O(1)
84−empty :: MMap k v
85−empty = Leaf
70+empty :: MMap v
71+empty = I.empty
8672
8773----------------------------------------------------------------
8874
8975-- | O(1)
90−singleton :: Ord k => k -> v -> MMap k v
91−singleton k v = Node B Leaf k (One v) Leaf
76+singleton :: Int -> v -> MMap v
77+singleton k v = I.singleton k (One v)
9278
9379----------------------------------------------------------------
9480
95−-- | O(log N)
96−insert :: Ord k => k -> v -> MMap k v -> MMap k v
97−insert kx kv t = turnB (insert' kx kv t)
98−
99−insert' :: Ord k => k -> v -> MMap k v -> MMap k v
100−insert' xk xv Leaf = Node R Leaf xk (One xv) Leaf
101−insert' xk xv (Node B l k v r) = case compare xk k of
102− LT -> balanceL' (insert' xk xv l) k v r
103− GT -> balanceR' l k v (insert' xk xv r)
104− EQ -> Node B l k (snoc v xv) r
105−insert' xk xv (Node R l k v r) = case compare xk k of
106− LT -> Node R (insert' xk xv l) k v r
107− GT -> Node R l k v (insert' xk xv r)
108− EQ -> Node R l k (snoc v xv) r
109−
110−balanceL' :: MMap k v -> k -> Some v -> MMap k v -> MMap k v
111−balanceL' (Node R (Node R a xk xv b) yk yv c) zk zv d =
112− Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
113−balanceL' (Node R a xk xv (Node R b yk yv c)) zk zv d =
114− Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
115−balanceL' l k v r = Node B l k v r
116−
117−balanceR' :: MMap k v -> k -> Some v -> MMap k v -> MMap k v
118−balanceR' a xk xv (Node R b yk yv (Node R c zk zv d)) =
119− Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
120−balanceR' a xk xv (Node R (Node R b yk yv c) zk zv d) =
121− Node R (Node B a xk xv b) yk yv (Node B c zk zv d)
122−balanceR' l xk xv r = Node B l xk xv r
123−
124−turnB :: MMap k v -> MMap k v
125−turnB Leaf = error "turnB"
126−turnB (Node _ l k v r) = Node B l k v r
81+-- | O(log n)
82+insert :: Int -> v -> MMap v -> MMap v
83+insert k v m = I.insertWith union k (One v) m
12784
12885----------------------------------------------------------------
12986
130−-- | O(N log N)
131−fromList :: Ord k => [(k,v)] -> MMap k v
132−fromList = foldl' (\t (k,v) -> insert k v t) empty
133−
134−-- | O(N)
135−toList :: MMap k v -> [(k,v)]
136−toList t = inorder t []
87+-- | O(n)
88+toList :: MMap v -> [v]
89+toList m = concatMap f $ I.toAscList m
13790 where
138− inorder Leaf xs = xs
139− inorder (Node _ l k v r) xs = inorder l (pairs k v ++ inorder r xs)
140− pairs k (One v) = [(k,v)]
141− pairs k (Tom v vs) = (k,v) : pairs k vs
91+ f (_,s) = go s []
92+ where
93+ go (One x) acc = x : acc
94+ go (Tom x xs) acc = go xs (x : acc)
14295
14396----------------------------------------------------------------
14497
145−-- | O(N)
146−-- "Constructing Red-Black Trees" by Ralf Hinze
147−fromSortedList :: Ord k => [(k,Some v)] -> MMap k v
148−fromSortedList = linkAll . foldr add []
149−
150−data Digit k v = Uno k (Some v) (MMap k v)
151− | Due k (Some v) (MMap k v) k (Some v) (MMap k v)
152− deriving (Eq,Show)
153−
154−incr :: Digit k v -> [Digit k v] -> [Digit k v]
155−incr (Uno k v t) [] = [Uno k v t]
156−incr (Uno k1 v1 t1) (Uno k2 v2 t2 : ps) = Due k1 v1 t1 k2 v2 t2 : ps
157−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
158−incr _ _ = error "incr"
159−
160−add :: (k,Some v) -> [Digit k v] -> [Digit k v]
161−add (k,v) ps = incr (Uno k v Leaf) ps
162−
163−linkAll :: [Digit k v] -> MMap k v
164−linkAll = foldl' link Leaf
165−
166−link :: MMap k v -> Digit k v -> MMap k v
167−link l (Uno k v t) = Node B l k v t
168−--link l (Due k1 v1 t1 k2 v2 t2) = Node B (Node R l k1 v1 t1) k2 v2 t2
169−link l (Due k1 v1 t1 k2 v2 t2) = Node B l k1 v1 (Node R t1 k2 v2 t2)
98+-- | O(n)
99+fromSortedList :: [(Int,Some v)] -> MMap v
100+fromSortedList = I.fromAscList
170101
171102----------------------------------------------------------------
172103
173−-- | O(N)
174−toSortedList :: MMap k v -> [(k,Some v)]
175−toSortedList t = inorder t []
176− where
177− inorder Leaf xs = xs
178− inorder (Node _ l k v r) xs = inorder l ((k,v) : inorder r xs)
104+-- | O(n)
105+toSortedList :: MMap v -> [(Int,Some v)]
106+toSortedList = I.toAscList
179107
180108----------------------------------------------------------------
181109
182−-- | O(N)
183−pruneWith :: Ord k =>
184− MMap k v
185− -> (k -> Some v -> IO [(k, Some v)])
186− -> IO (MMap k v)
187−pruneWith t run = fromSortedList <$> inorder t []
188− where
189− inorder Leaf xs = return xs
190− inorder (Node _ l k v r) xs = do
191− ys <- run k v
192− zs <- inorder r xs
193− inorder l (ys ++ zs)
110+-- | O(n)
111+pruneWith :: MMap v
112+ -> ((Int,Some v) -> IO [(Int, Some v)])
113+ -> IO (MMap v)
114+pruneWith m action = fromSortedList . concat <$> mapM action (toSortedList m)
194115
195116----------------------------------------------------------------
196117
197−-- O(N log N) where N is the size of the second argument
198−merge :: Ord k => MMap k v -> MMap k v -> MMap k v
199−merge base m = foldl' ins base xs
200− where
201− ins t (k,v) = insert k v t
202− xs = toList m
203−
204−----------------------------------------------------------------
205−-- for testing
206−
207−valid :: Ord k => MMap k v -> Bool
208−valid t = isBalanced t && isOrdered t
209−
210−isBalanced :: MMap k v -> Bool
211−isBalanced t = isBlackSame t && isRedSeparate t
212−
213−isBlackSame :: MMap k v -> Bool
214−isBlackSame t = all (n==) ns
215− where
216− n:ns = blacks t
217−
218−blacks :: MMap k v -> [Int]
219−blacks = blacks' 0
220− where
221− blacks' n Leaf = [n+1]
222− blacks' n (Node R l _ _ r) = blacks' n l ++ blacks' n r
223− blacks' n (Node B l _ _ r) = blacks' n' l ++ blacks' n' r
224− where
225− n' = n + 1
226−
227−isRedSeparate :: MMap k v -> Bool
228−isRedSeparate = reds B
229−
230−reds :: Color -> MMap k v -> Bool
231−reds _ Leaf = True
232−reds R (Node R _ _ _ _) = False
233−reds _ (Node c l _ _ r) = reds c l && reds c r
234−
235−isOrdered :: Ord k => MMap k v -> Bool
236−isOrdered t = ordered $ toSortedList t
237− where
238− ordered [] = True
239− ordered [_] = True
240− ordered (x:y:xys) = fst x <= fst y && ordered (y:xys)
118+-- O(n + m) where N is the size of the second argument
119+merge :: MMap v -> MMap v -> MMap v
120+merge m1 m2 = I.unionWith union m1 m2