{-# LANGUAGE UnboxedTuples, TypeFamilies, PatternGuards #-} module Data.TrieMap.OrdMap () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized import Data.TrieMap.Modifiers import Control.Applicative (Applicative(..), Alternative(..), (<$>)) import Control.Monad hiding (join) import Prelude hiding (lookup) type OrdMap k = TrieMap (Ordered k) instance Ord k => TrieKey (Ordered k) where data TrieMap (Ordered k) a = Tip | Bin {-# UNPACK #-} !Int k a !(OrdMap k a) !(OrdMap k a) emptyM = Tip singletonM s (Ord k) = singleton s k nullM Tip = True nullM _ = False sizeM _ = size lookupM (Ord k) = lookup k alterM s f (Ord k) = alter s f k alterLookupM s f (Ord k) = alterLookup s f k traverseWithKeyM s f = traverseWithKey s (f . Ord) foldWithKeyM f = foldrWithKey (f . Ord) foldlWithKeyM f = foldlWithKey (f . Ord) mapMaybeM s f = mapMaybe s (f . Ord) mapEitherM s1 s2 f = mapEither s1 s2 (f . Ord) extractM s f = extract s (f . Ord) splitLookupM s f (Ord k) = splitLookup s f k isSubmapM = isSubmap fromAscListM s f xs = fromAscList s (f . Ord) [(k, a) | (Ord k, a) <- xs] fromDistAscListM s xs = fromDistinctAscList s [(k, a) | (Ord k, a) <- xs] unionM _ _ Tip m2 = m2 unionM _ _ m1 Tip = m1 unionM s f m1 m2 = hedgeUnionWithKey s (f . Ord) (const LT) (const GT) m1 m2 isectM s f = isect s (f . Ord) diffM _ _ Tip _ = Tip diffM _ _ m1 Tip = m1 diffM s f m1 m2 = hedgeDiffWithKey s (f . Ord) (const LT) (const GT) m1 m2 lookup :: Ord k => k -> OrdMap k a -> Maybe a lookup k (Bin _ k' v l r) = case compare k k' of LT -> lookup k l EQ -> Just v GT -> lookup k r lookup _ _ = Nothing alter :: Ord k => Sized a -> (Maybe a -> Maybe a) -> k -> OrdMap k a -> OrdMap k a alter s f k Tip = case f Nothing of Nothing -> Tip Just x -> singleton s k x alter s f k (Bin _ kx x l r) = case compare k kx of LT -> balance s kx x (alter s f k l) r EQ -> case f (Just x) of Nothing -> glue s l r Just x' -> balance s k x' l r GT -> balance s kx x l (alter s f k r) alterLookup :: Ord k => Sized a -> (Maybe a -> (# z, Maybe a #)) -> k -> OrdMap k a -> (# z, OrdMap k a #) alterLookup s f k Tip = onUnboxed (maybe Tip (singleton s k)) f Nothing alterLookup s f k (Bin _ kx x l r) = case compare k kx of LT -> onUnboxed (\ l' -> balance s kx x l' r) (alterLookup s f k) l EQ -> onUnboxed (maybe (glue s l r) (\ x' -> balance s k x' l r)) f (Just x) GT -> onUnboxed (balance s kx x l) (alterLookup s f k) r singleton :: Sized a -> k -> a -> OrdMap k a singleton s k a = Bin (s a) k a Tip Tip traverseWithKey :: Applicative f => Sized b -> (k -> a -> f b) -> OrdMap k a -> f (OrdMap k b) traverseWithKey _ _ Tip = pure Tip traverseWithKey s f (Bin _ k a l r) = balance s k <$> f k a <*> traverseWithKey s f l <*> traverseWithKey s f r foldrWithKey :: (k -> a -> b -> b) -> OrdMap k a -> b -> b foldrWithKey _ Tip = id foldrWithKey f (Bin _ k a l r) = foldrWithKey f l . f k a . foldrWithKey f r foldlWithKey :: (k -> b -> a -> b) -> OrdMap k a -> b -> b foldlWithKey _ Tip = id foldlWithKey f (Bin _ k a l r) = foldlWithKey f r . flip (f k) a . foldlWithKey f l mapMaybe :: Ord k => Sized b -> (k -> a -> Maybe b) -> OrdMap k a -> OrdMap k b mapMaybe _ _ Tip = Tip mapMaybe s f (Bin _ k a l r) = joinMaybe s k (f k a) (mapMaybe s f l) (mapMaybe s f r) mapEither :: Ord k => Sized b -> Sized c -> EitherMap k a b c -> OrdMap k a -> (# OrdMap k b, OrdMap k c #) mapEither _ _ _ Tip = (# Tip, Tip #) mapEither s1 s2 f (Bin _ k a l r) | (# aL, aR #) <- f k a, (# lL, lR #) <- mapEither s1 s2 f l, (# rL, rR #) <- mapEither s1 s2 f r = (# joinMaybe s1 k aL lL rL, joinMaybe s2 k aR lR rR #) splitLookup :: Ord k => Sized a -> SplitMap a x -> k -> OrdMap k a -> (# OrdMap k a, Maybe x, OrdMap k a #) splitLookup s f k m = case m of Tip -> (# Tip, Nothing, Tip #) Bin _ kx x l r -> case compare k kx of LT -> case splitLookup s f k l of (# lL, ans, lR #) -> (# lL, ans, join s kx x lR r #) EQ -> case f x of (# xL, ans, xR #) -> (# maybe l (\ xL -> insertMax s kx xL l) xL, ans, maybe r (\ xR -> insertMin s kx xR r) xR #) GT -> case splitLookup s f k r of (# rL, ans, rR #) -> (# join s kx x l rL, ans, rR #) isSubmap :: Ord k => LEq a b -> LEq (OrdMap k a) (OrdMap k b) isSubmap _ Tip _ = True isSubmap _ _ Tip = False isSubmap (<=) (Bin _ kx x l r) t = case splitLookup (const 1) (\ x -> (# Nothing, Just x, Nothing #)) kx t of (# lt, found, gt #) -> case found of Nothing -> False Just y -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt fromAscList :: Eq k => Sized a -> (k -> a -> a -> a) -> [(k, a)] -> OrdMap k a fromAscList s f xs = fromDistinctAscList s (combineEq xs) where combineEq (x:xs) = combineEq' x xs combineEq [] = [] combineEq' z [] = [z] combineEq' (kz, zz) (x@(kx, xx):xs) | kz == kx = combineEq' (kx, f kx xx zz) xs | otherwise = (kz,zz):combineEq' x xs fromDistinctAscList :: Sized a -> [(k, a)] -> OrdMap k a fromDistinctAscList s xs = build const (length xs) xs where -- 1) use continutations so that we use heap space instead of stack space. -- 2) special case for n==5 to build bushier trees. build c 0 xs' = c Tip xs' build c 5 xs' = case xs' of ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx) -> c (bin s k4 x4 (bin s k2 x2 (singleton s k1 x1) (singleton s k3 x3)) (singleton s k5 x5)) xx _ -> error "fromDistinctAscList build" build c n xs' = seq nr $ build (buildR nr c) nl xs' where nl = n `div` 2 nr = n - nl - 1 buildR n c l ((k,x):ys) = build (buildB l k x c) n ys buildR _ _ _ [] = error "fromDistinctAscList buildR []" buildB l k x c r zs = c (bin s k x l r) zs hedgeUnionWithKey :: Ord k => Sized a -> (k -> a -> a -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k a -> OrdMap k a hedgeUnionWithKey _ _ _ _ t1 Tip = t1 hedgeUnionWithKey s _ cmplo cmphi Tip (Bin _ kx x l r) = join s kx x (filterGt s cmplo l) (filterLt s cmphi r) hedgeUnionWithKey s f cmplo cmphi (Bin _ kx x l r) t2 = joinMaybe s kx newx (hedgeUnionWithKey s f cmplo cmpkx l lt) (hedgeUnionWithKey s f cmpkx cmphi r gt) where cmpkx k = compare kx k lt = trim cmplo cmpkx t2 (found,gt) = trimLookupLo kx cmphi t2 newx = case found of Nothing -> Just x Just (_,y) -> f kx x y filterGt :: Ord k => Sized a -> (k -> Ordering) -> OrdMap k a -> OrdMap k a filterGt _ _ Tip = Tip filterGt s cmp (Bin _ kx x l r) = case cmp kx of LT -> join s kx x (filterGt s cmp l) r GT -> filterGt s cmp r EQ -> r filterLt :: Ord k => Sized a -> (k -> Ordering) -> OrdMap k a -> OrdMap k a filterLt _ _ Tip = Tip filterLt s cmp (Bin _ kx x l r) = case cmp kx of LT -> filterLt s cmp l GT -> join s kx x l (filterLt s cmp r) EQ -> l trim :: (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k a trim _ _ Tip = Tip trim cmplo cmphi t@(Bin _ kx _ l r) = case cmplo kx of LT -> case cmphi kx of GT -> t _ -> trim cmplo cmphi l _ -> trim cmplo cmphi r trimLookupLo :: Ord k => k -> (k -> Ordering) -> OrdMap k a -> (Maybe (k,a), OrdMap k a) trimLookupLo _ _ Tip = (Nothing,Tip) trimLookupLo lo cmphi t@(Bin _ kx x l r) = case compare lo kx of LT -> case cmphi kx of GT -> (((,) lo) <$> lookup lo t, t) _ -> trimLookupLo lo cmphi l GT -> trimLookupLo lo cmphi r EQ -> (Just (kx,x),trim (compare lo) cmphi r) isect :: Ord k => Sized c -> IsectFunc k a b c -> OrdMap k a -> OrdMap k b -> OrdMap k c isect s f t1@Bin{} (Bin _ k2 x2 l2 r2) | (# lt, found, gt #) <- splitLookup (const 1) (\ x -> (# Nothing, Just x, Nothing #)) k2 t1 = let tl = isect s f lt l2 tr = isect s f gt r2 in joinMaybe s k2 (found >>= \ x1' -> f k2 x1' x2) tl tr isect _ _ _ _ = Tip hedgeDiffWithKey :: Ord k => Sized a -> (k -> a -> b -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k b -> OrdMap k a hedgeDiffWithKey _ _ _ _ Tip _ = Tip hedgeDiffWithKey s _ cmplo cmphi (Bin _ kx x l r) Tip = join s kx x (filterGt s cmplo l) (filterLt s cmphi r) hedgeDiffWithKey s f cmplo cmphi t (Bin _ kx x l r) = case found of Nothing -> merge s tl tr Just (ky,y) -> case f ky y x of Nothing -> merge s tl tr Just z -> join s ky z tl tr where cmpkx k = compare kx k lt = trim cmplo cmpkx t (found,gt) = trimLookupLo kx cmphi t tl = hedgeDiffWithKey s f cmplo cmpkx lt l tr = hedgeDiffWithKey s f cmpkx cmphi gt r joinMaybe :: Ord k => Sized a -> k -> Maybe a -> OrdMap k a -> OrdMap k a -> OrdMap k a joinMaybe s kx = maybe (merge s) (join s kx) join :: Ord k => Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a join s kx x Tip r = insertMin s kx x r join s kx x l Tip = insertMax s kx x l join s kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) | delta*sizeL <= sizeR = balance s kz z (join s kx x l lz) rz | delta*sizeR <= sizeL = balance s ky y ly (join s kx x ry r) | otherwise = bin s kx x l r -- insertMin and insertMax don't perform potentially expensive comparisons. insertMax,insertMin :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a insertMax s kx x t = case t of Tip -> singleton s kx x Bin _ ky y l r -> balance s ky y l (insertMax s kx x r) insertMin s kx x t = case t of Tip -> singleton s kx x Bin _ ky y l r -> balance s ky y (insertMin s kx x l) r {-------------------------------------------------------------------- [merge l r]: merges two trees. --------------------------------------------------------------------} merge :: Sized a -> OrdMap k a -> OrdMap k a -> OrdMap k a merge _ Tip r = r merge _ l Tip = l merge s l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) | delta*sizeL <= sizeR = balance s ky y (merge s l ly) ry | delta*sizeR <= sizeL = balance s kx x lx (merge s rx r) | otherwise = glue s l r {-------------------------------------------------------------------- [glue l r]: glues two trees together. Assumes that [l] and [r] are already balanced with respect to each other. --------------------------------------------------------------------} glue :: Sized a -> OrdMap k a -> OrdMap k a -> OrdMap k a glue _ Tip r = r glue _ l Tip = l glue s l r | size l > size r = let (f,l') = deleteFindMax s (\ k a -> (balance s k a, Nothing)) l in f l' r | otherwise = let (f,r') = deleteFindMin s (\ k a -> (balance s k a, Nothing)) r in f l r' extract :: Alternative t => Sized a -> (k -> a -> t (z, Maybe a)) -> OrdMap k a -> t (z, OrdMap k a) extract s f t = case t of Bin _ k x l r -> fmap (\ l' -> balance s k x l' r) <$> extract s f l <|> fmap (maybe (glue s l r) (\ x' -> balance s k x' l r)) <$> f k x <|> fmap (balance s k x l) <$> extract s f r Tip -> empty deleteFindMin :: Sized a -> (k -> a -> (x, Maybe a)) -> OrdMap k a -> (x, OrdMap k a) deleteFindMin s f t = case t of Bin _ k x Tip r -> let (ans, x') = f k x in (ans, maybe r (\ y' -> bin s k y' Tip r) x') Bin _ k x l r -> let (km,l') = deleteFindMin s f l in (km,balance s k x l' r) Tip -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) deleteFindMax :: Sized a -> (k -> a -> (x, Maybe a)) -> OrdMap k a -> (x, OrdMap k a) deleteFindMax s f t = case t of Bin _ k x l Tip -> let (ans, x') = f k x in (ans, maybe l (\ y -> bin s k y l Tip) x') Bin _ k x l r -> let (km,r') = deleteFindMax s f r in (km,balance s k x l r') Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) delta,ratio :: Int delta = 5 ratio = 2 size :: OrdMap k a -> Int size Tip = 0 size (Bin s _ _ _ _) = s balance :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a balance s k x l r | sizeL + sizeR <= 1 = Bin sizeX k x l r | sizeR >= delta*sizeL = rotateL s k x l r | sizeL >= delta*sizeR = rotateR s k x l r | otherwise = Bin sizeX k x l r where sizeL = size l sizeR = size r sizeX = sizeL + sizeR + s x -- rotate rotateL :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a rotateL s k x l r@(Bin _ _ _ ly ry) | size ly < ratio*size ry = singleL s k x l r | otherwise = doubleL s k x l r rotateL _ _ _ _ Tip = error "rotateL Tip" rotateR :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a rotateR s k x l@(Bin _ _ _ ly ry) r | size ry < ratio*size ly = singleR s k x l r | otherwise = doubleR s k x l r rotateR _ _ _ Tip _ = error "rotateR Tip" -- basic rotations singleL, singleR :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a singleL s k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin s k2 x2 (bin s k1 x1 t1 t2) t3 singleL s k1 x1 t1 Tip = bin s k1 x1 t1 Tip singleR s k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin s k2 x2 t1 (bin s k1 x1 t2 t3) singleR s k1 x1 Tip t2 = bin s k1 x1 Tip t2 doubleL, doubleR :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a doubleL s k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin s k3 x3 (bin s k1 x1 t1 t2) (bin s k2 x2 t3 t4) doubleL s k1 x1 t1 t2 = singleL s k1 x1 t1 t2 doubleR s k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin s k3 x3 (bin s k2 x2 t1 t2) (bin s k1 x1 t3 t4) doubleR s k1 x1 t1 t2 = singleR s k1 x1 t1 t2 bin :: Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a bin s k x l r = Bin (size l + size r + s x) k x l r