module Data.TrieMap.OrdMap (Ordered (..)) where
import Data.TrieMap.TrieKey
import Data.TrieMap.Sized
import Data.TrieMap.Applicative
import Control.Applicative (Applicative(..), (<$>))
import Control.Arrow
import Control.Monad hiding (join)
import Data.Monoid
import Data.Maybe
import Data.Traversable
import Prelude hiding (lookup)
newtype Ordered a = Ord {unOrd :: a} deriving (Eq, Ord)
data OrdMap k a ix = Tip
| Bin !Int k (a ix) !(OrdMap k a ix) !(OrdMap k a ix)
type instance TrieMap (Ordered k) = OrdMap k
instance Ord k => TrieKey (Ordered k) (OrdMap k) where
emptyM = Tip
nullM Tip = True
nullM _ = False
sizeM _ = size
lookupM (Ord k) = lookup k
lookupIxM _ (Ord k) = lookupIx 0 k
assocAtM _ i m = fromJust (do (i', k, a) <- assocAt 0 i m
return (i', Ord k, a))
updateAtM s f = updateAt s (\ i -> f i . Ord)
alterM s f (Ord k) = alter s f k
traverseWithKeyM s f = traverseWithKey s (f . Ord)
foldWithKeyM f = foldrWithKey (f . Ord)
foldlWithKeyM f = foldlWithKey (f . Ord)
mapEitherM s1 s2 f = mapEither s1 s2 (f . Ord)
extractMinM s Tip = mzero
extractMinM s m = return (first (first Ord) $ deleteFindMin s m)
extractMaxM s Tip = mzero
extractMaxM s m = return (first (first Ord) $ deleteFindMax s m)
alterMinM s f = updateMin s (f . Ord)
alterMaxM s f = updateMax 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 s f m1 m2 = case (m1, m2) of
(Tip, _) -> m2
(_, Tip) -> m1
_ -> hedgeUnionWithKey s (f . Ord) (const LT) (const GT) m1 m2
isectM s f = isect s (f . Ord)
diffM s f m1 m2 = case (m1, m2) of
(Tip, _) -> Tip
(_, Tip) -> m1
_ -> hedgeDiffWithKey s (f . Ord) (const LT) (const GT) m1 m2
lookup :: Ord k => k -> OrdMap k a ix -> Maybe (a ix)
lookup k Tip = Nothing
lookup k (Bin _ k' v l r) = case compare k k' of
LT -> lookup k l
EQ -> Just v
GT -> lookup k r
lookupIx :: Ord k => Int -> k -> OrdMap k a ix -> Maybe (Int, a ix)
lookupIx i _ _ | i `seq` False = undefined
lookupIx _ _ Tip = Nothing
lookupIx i k (Bin sz k' v l r) = case compare k k' of
LT -> lookupIx i k l
EQ -> Just (size l, v)
GT -> lookupIx (i + sz size r) k r
assocAt :: Int -> Int -> OrdMap k a ix -> Maybe (Int, k, a ix)
assocAt i0 i _ | i0 `seq` i `seq` False = Nothing
assocAt _ _ Tip = Nothing
assocAt i0 i (Bin sz k a l r)
| i < sL = assocAt i0 i l
| i < sK = Just (i0 + sL, k, a)
| otherwise = assocAt (i0 + sK) (i sK) r
where sL = size l
sK = sz size r
updateAt :: Sized a -> (Int -> k -> a ix -> Maybe (a ix)) -> Int -> OrdMap k a ix -> OrdMap k a ix
updateAt _ _ i _ | i `seq` False = undefined
updateAt _ _ _ Tip = Tip
updateAt s f i (Bin sz k a l r)
| i < sL = balance s k a (updateAt s f i l) r
| i < sK = case f sK k a of
Nothing -> glue s l r
Just a' -> bin s k a' l r
| otherwise = balance s k a l (updateAt s (f . (+ sK)) (i sK) r)
where sL = size l
sK = sz size r
alter :: Ord k => Sized a -> (Maybe (a ix) -> Maybe (a ix)) -> k -> OrdMap k a ix -> OrdMap k a ix
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)
singleton :: Sized a -> k -> a ix -> OrdMap k a ix
singleton s k a = Bin (s a) k a Tip Tip
traverseWithKey :: Applicative f => Sized b -> (k -> a ix -> f (b ix)) -> OrdMap k a ix -> f (OrdMap k b ix)
traverseWithKey s f 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 ix -> b -> b) -> OrdMap k a ix -> b -> b
foldrWithKey f Tip = id
foldrWithKey f (Bin _ k a l r) = foldrWithKey f l . f k a . foldrWithKey f r
foldlWithKey :: (k -> b -> a ix -> b) -> OrdMap k a ix -> b -> b
foldlWithKey f Tip = id
foldlWithKey f (Bin _ k a l r) = foldlWithKey f r . flip (f k) a . foldlWithKey f l
mapEither :: Ord k => Sized b -> Sized c -> EitherMap k (a ix) (b ix) (c ix) ->
OrdMap k a ix -> (OrdMap k b ix, OrdMap k c ix)
mapEither s1 s2 f m = case m of
Tip -> (Tip, Tip)
Bin _ k a l r -> case (f k a, mapEither s1 s2 f l, mapEither s1 s2 f r) of
((aL, aR), (lL, lR), (rL, rR)) ->
(joinMaybe s1 k aL lL rL, joinMaybe s2 k aR lR rR)
updateMin :: Ord k => Sized a -> (k -> a ix -> Maybe (a ix)) -> OrdMap k a ix -> OrdMap k a ix
updateMin s f m = case m of
Tip -> Tip
Bin _ k a Tip r -> case f k a of
Nothing -> r
Just a' -> insertMin s k a' r
Bin _ k a l r -> balance s k a (updateMin s f l) r
updateMax :: Ord k => Sized a -> (k -> a ix -> Maybe (a ix)) -> OrdMap k a ix -> OrdMap k a ix
updateMax s f m = case m of
Tip -> Tip
Bin _ k a l Tip -> case f k a of
Nothing -> l
Just a' -> insertMax s k a' l
Bin _ k a l r -> balance s k a l (updateMax s f r)
splitLookup :: Ord k => Sized a -> SplitMap (a ix) x -> k -> OrdMap k a ix -> (OrdMap k a ix, Maybe x, OrdMap k a ix)
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 ix) (b ix) -> LEq (OrdMap k a ix) (OrdMap k b ix)
isSubmap (<=) Tip _ = True
isSubmap (<=) _ Tip = False
isSubmap (<=) (Bin _ kx x l r) t = case found of
Nothing -> False
Just y -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt
where (lt, found, gt) = splitLookup (const 1) (\ x -> (Nothing, Just x, Nothing)) kx t
fromAscList :: Eq k => Sized a -> (k -> a ix -> a ix -> a ix) -> [(k, a ix)] -> OrdMap k a ix
fromAscList s f xs = fromDistinctAscList s (combineEq xs) where
combineEq (x:xs) = combineEq' x xs
combineEq [] = []
combineEq' z [] = [z]
combineEq' z@(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 ix)] -> OrdMap k a ix
fromDistinctAscList s xs = build const (length xs) xs
where
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 ix -> a ix -> Maybe (a ix))
-> (k -> Ordering) -> (k -> Ordering)
-> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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 ix -> OrdMap k a ix
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 ix -> OrdMap k a ix
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 ix -> OrdMap k a ix
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 ix -> (Maybe (k,a ix), OrdMap k a ix)
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 ix) (b ix) (c ix) -> OrdMap k a ix -> OrdMap k b ix -> OrdMap k c ix
isect s f Tip _ = Tip
isect s f _ Tip = Tip
isect s f t1@(Bin _ k1 x1 l1 r1) t2@(Bin _ k2 x2 l2 r2) =
let (lt, found, gt) = splitLookup (const 1) (\ x -> (Nothing, Just x, Nothing)) k2 t1
tl = isect s f lt l2
tr = isect s f gt r2
in joinMaybe s k2 (found >>= \ x1' -> f k2 x1' x2) tl tr
hedgeDiffWithKey :: Ord k
=> Sized a -> (k -> a ix -> b ix -> Maybe (a ix))
-> (k -> Ordering) -> (k -> Ordering)
-> OrdMap k a ix -> OrdMap k b ix -> OrdMap k a ix
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 ix) -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
joinMaybe s kx = maybe (merge s) (join s kx)
join :: Ord k => Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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
insertMax,insertMin :: Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix
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 :: Sized a -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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 :: Sized a -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
glue _ Tip r = r
glue _ l Tip = l
glue s l r
| size l > size r = let ((km,m),l') = deleteFindMax s l in balance s km m l' r
| otherwise = let ((km,m),r') = deleteFindMin s r in balance s km m l r'
deleteFindMin :: Sized a -> OrdMap k a ix -> ((k, a ix), OrdMap k a ix)
deleteFindMin s t
= case t of
Bin _ k x Tip r -> ((k,x),r)
Bin _ k x l r -> let (km,l') = deleteFindMin s 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 -> OrdMap k a ix -> ((k, a ix), OrdMap k a ix)
deleteFindMax s t
= case t of
Bin _ k x l Tip -> ((k,x),l)
Bin _ k x l r -> let (km,r') = deleteFindMax s 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 ix -> Int
size Tip = 0
size (Bin s _ _ _ _) = s
balance :: Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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
rotateL :: Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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 ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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"
singleL, singleR :: Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
singleL s k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin s k2 x2 (bin s k1 x1 t1 t2) t3
singleL _ _ _ _ Tip = error "singleL Tip"
singleR s k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin s k2 x2 t1 (bin s k1 x1 t2 t3)
singleR _ _ _ Tip _ = error "singleR Tip"
doubleL, doubleR :: Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
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 _ _ _ _ _ = error "doubleL"
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 _ _ _ _ _ = error "doubleR"
bin :: Sized a -> k -> a ix -> OrdMap k a ix -> OrdMap k a ix -> OrdMap k a ix
bin s k x l r
= Bin (size l + size r + s x) k x l r