module Network.Wai.Handler.Warp.MultiMap ( MMap , Some(..) , empty , singleton , insert , search , searchWith , isEmpty , valid , pruneWith , fromList , toList , fromSortedList , toSortedList , merge ) where import Control.Applicative ((<$>)) import Data.List (foldl') ---------------------------------------------------------------- -- | One ore more list to implement multimap. data Some a = One !a | Tom !a !(Some a) -- Two or more deriving (Eq,Show) -- This is slow but assuming rarely used. snoc :: Some a -> a -> Some a snoc (One x) y = Tom x (One y) snoc (Tom x xs) y = Tom x (snoc xs y) top :: Some a -> a top (One x) = x top (Tom x _) = x ---------------------------------------------------------------- -- | Red black tree as multimap. data MMap k v = Leaf -- color is Black | Node Color !(MMap k v) !k !(Some v) !(MMap k v) deriving (Show) data Color = B -- ^ Black | R -- ^ Red deriving (Eq, Show) ---------------------------------------------------------------- instance (Eq k, Eq v) => Eq (MMap k v) where t1 == t2 = toSortedList t1 == toSortedList t2 ---------------------------------------------------------------- -- | O(log N) search :: Ord k => k -> MMap k v -> Maybe v search _ Leaf = Nothing search xk (Node _ l k v r) = case compare xk k of LT -> search xk l GT -> search xk r EQ -> Just $ top v -- | O(log N) searchWith :: Ord k => k -> (Some v -> Maybe v) -> MMap k v -> Maybe v searchWith _ _ Leaf = Nothing searchWith xk f (Node _ l k v r) = case compare xk k of LT -> searchWith xk f l GT -> searchWith xk f r EQ -> f v ---------------------------------------------------------------- -- | O(1) isEmpty :: (Eq k, Eq v) => MMap k v -> Bool isEmpty Leaf = True isEmpty _ = False -- | O(1) empty :: MMap k v empty = Leaf ---------------------------------------------------------------- -- | O(1) singleton :: Ord k => k -> v -> MMap k v singleton k v = Node B Leaf k (One v) Leaf ---------------------------------------------------------------- -- | O(log N) insert :: Ord k => k -> v -> MMap k v -> MMap k v insert kx kv t = turnB (insert' kx kv t) insert' :: Ord k => k -> v -> MMap k v -> MMap k v insert' xk xv Leaf = Node R Leaf xk (One xv) Leaf insert' xk xv (Node B l k v r) = case compare xk k of LT -> balanceL' (insert' xk xv l) k v r GT -> balanceR' l k v (insert' xk xv r) EQ -> Node B l k (snoc v xv) r insert' xk xv (Node R l k v r) = case compare xk k of LT -> Node R (insert' xk xv l) k v r GT -> Node R l k v (insert' xk xv r) EQ -> Node R l k (snoc v xv) r balanceL' :: MMap k v -> k -> Some v -> MMap k v -> MMap k v balanceL' (Node R (Node R a xk xv b) yk yv c) zk zv d = Node R (Node B a xk xv b) yk yv (Node B c zk zv d) balanceL' (Node R a xk xv (Node R b yk yv c)) zk zv d = Node R (Node B a xk xv b) yk yv (Node B c zk zv d) balanceL' l k v r = Node B l k v r balanceR' :: MMap k v -> k -> Some v -> MMap k v -> MMap k v balanceR' a xk xv (Node R b yk yv (Node R c zk zv d)) = Node R (Node B a xk xv b) yk yv (Node B c zk zv d) balanceR' a xk xv (Node R (Node R b yk yv c) zk zv d) = Node R (Node B a xk xv b) yk yv (Node B c zk zv d) balanceR' l xk xv r = Node B l xk xv r turnB :: MMap k v -> MMap k v turnB Leaf = error "turnB" turnB (Node _ l k v r) = Node B l k v r ---------------------------------------------------------------- -- | O(N log N) fromList :: Ord k => [(k,v)] -> MMap k v fromList = foldl' (\t (k,v) -> insert k v t) empty -- | O(N) toList :: MMap k v -> [(k,v)] toList t = inorder t [] where inorder Leaf xs = xs inorder (Node _ l k v r) xs = inorder l (pairs k v ++ inorder r xs) pairs k (One v) = [(k,v)] pairs k (Tom v vs) = (k,v) : pairs k vs ---------------------------------------------------------------- -- | O(N) -- "Constructing Red-Black Trees" by Ralf Hinze fromSortedList :: Ord k => [(k,Some v)] -> MMap k v fromSortedList = linkAll . foldr add [] data Digit k v = Uno k (Some v) (MMap k v) | Due k (Some v) (MMap k v) k (Some v) (MMap k v) deriving (Eq,Show) incr :: Digit k v -> [Digit k v] -> [Digit k v] incr (Uno k v t) [] = [Uno k v t] incr (Uno k1 v1 t1) (Uno k2 v2 t2 : ps) = Due k1 v1 t1 k2 v2 t2 : ps 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 incr _ _ = error "incr" add :: (k,Some v) -> [Digit k v] -> [Digit k v] add (k,v) ps = incr (Uno k v Leaf) ps linkAll :: [Digit k v] -> MMap k v linkAll = foldl' link Leaf link :: MMap k v -> Digit k v -> MMap k v link l (Uno k v t) = Node B l k v t --link l (Due k1 v1 t1 k2 v2 t2) = Node B (Node R l k1 v1 t1) k2 v2 t2 link l (Due k1 v1 t1 k2 v2 t2) = Node B l k1 v1 (Node R t1 k2 v2 t2) ---------------------------------------------------------------- -- | O(N) toSortedList :: MMap k v -> [(k,Some v)] toSortedList t = inorder t [] where inorder Leaf xs = xs inorder (Node _ l k v r) xs = inorder l ((k,v) : inorder r xs) ---------------------------------------------------------------- -- | O(N) pruneWith :: Ord k => MMap k v -> (k -> Some v -> IO [(k, Some v)]) -> IO (MMap k v) pruneWith t run = fromSortedList <$> inorder t [] where inorder Leaf xs = return xs inorder (Node _ l k v r) xs = do ys <- run k v zs <- inorder r xs inorder l (ys ++ zs) ---------------------------------------------------------------- -- O(N log N) where N is the size of the second argument merge :: Ord k => MMap k v -> MMap k v -> MMap k v merge base m = foldl' ins base xs where ins t (k,v) = insert k v t xs = toList m ---------------------------------------------------------------- -- for testing valid :: Ord k => MMap k v -> Bool valid t = isBalanced t && isOrdered t isBalanced :: MMap k v -> Bool isBalanced t = isBlackSame t && isRedSeparate t isBlackSame :: MMap k v -> Bool isBlackSame t = all (n==) ns where n:ns = blacks t blacks :: MMap k v -> [Int] blacks = blacks' 0 where blacks' n Leaf = [n+1] blacks' n (Node R l _ _ r) = blacks' n l ++ blacks' n r blacks' n (Node B l _ _ r) = blacks' n' l ++ blacks' n' r where n' = n + 1 isRedSeparate :: MMap k v -> Bool isRedSeparate = reds B reds :: Color -> MMap k v -> Bool reds _ Leaf = True reds R (Node R _ _ _ _) = False reds _ (Node c l _ _ r) = reds c l && reds c r isOrdered :: Ord k => MMap k v -> Bool isOrdered t = ordered $ toSortedList t where ordered [] = True ordered [_] = True ordered (x:y:xys) = fst x <= fst y && ordered (y:xys)