{-# LANGUAGE UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns, MagicHash, CPP, BangPatterns, FlexibleInstances #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.UnionMap () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized import Data.TrieMap.UnitMap () import Control.Applicative import Control.Monad import Data.Monoid import Data.Foldable (Foldable(..)) import Prelude hiding (foldr, foldr1, foldl, foldl1, (^)) (&) :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a m1 & m2 = guardNullM m1 ^ guardNullM m2 {-# INLINE (^) #-} (^) :: (TrieKey k1, TrieKey k2, Sized a) => Maybe (TrieMap k1 a) -> Maybe (TrieMap k2 a) -> TrieMap (Either k1 k2) a Nothing ^ Nothing = Empty Just m1 ^ Nothing = K1 m1 Nothing ^ Just m2 = K2 m2 Just m1 ^ Just m2 = Union (sizeM m1 + sizeM m2) m1 m2 union :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a union m1 m2 = Union (sizeM m1 + getSize m2) m1 m2 singletonL :: (TrieKey k1, TrieKey k2, Sized a) => k1 -> a -> TrieMap (Either k1 k2) a singletonL k a = K1 (singletonM k a) singletonR :: (TrieKey k1, TrieKey k2, Sized a) => k2 -> a -> TrieMap (Either k1 k2) a singletonR k a = K2 (singletonM k a) data UView k1 k2 a = UView (Maybe (TrieMap k1 a)) (Maybe (TrieMap k2 a)) data HView k1 k2 a = Hole1 (Hole k1 a) (Maybe (TrieMap k2 a)) | Hole2 (Maybe (TrieMap k1 a)) (Hole k2 a) uView :: TrieMap (Either k1 k2) a -> UView k1 k2 a uView Empty = UView Nothing Nothing uView (K1 m1) = UView (Just m1) Nothing uView (K2 m2) = UView Nothing (Just m2) uView (Union _ m1 m2) = UView (Just m1) (Just m2) hView :: Hole (Either k1 k2) a -> HView k1 k2 a hView (HoleX0 hole1) = Hole1 hole1 Nothing hView (HoleX2 hole1 m2) = Hole1 hole1 (Just m2) hView (Hole0X hole2) = Hole2 Nothing hole2 hView (Hole1X m1 hole2) = Hole2 (Just m1) hole2 hole1 :: Hole k1 a -> Maybe (TrieMap k2 a) -> Hole (Either k1 k2) a hole1 hole1 Nothing = HoleX0 hole1 hole1 hole1 (Just m2) = HoleX2 hole1 m2 hole2 :: Maybe (TrieMap k1 a) -> Hole k2 a -> Hole (Either k1 k2) a hole2 Nothing hole2 = Hole0X hole2 hole2 (Just m1) hole2 = Hole1X m1 hole2 #define UVIEW uView -> UView instance (TrieKey k1, TrieKey k2) => Foldable (UView k1 k2) where {-# INLINE foldr #-} {-# INLINE foldl #-} {-# INLINE foldMap #-} foldMap f (UView m1 m2) = foldMap (foldMap f) m1 `mappend` foldMap (foldMap f) m2 foldr f z (UView m1 m2) = foldl (foldr f) (foldl (foldr f) z m2) m1 foldl f z (UView m1 m2) = foldl (foldl f) (foldl (foldl f) z m1) m2 instance (TrieKey k1, TrieKey k2) => Foldable (TrieMap (Either k1 k2)) where foldMap f m = foldMap f (uView m) foldr f z m = foldr f z (uView m) foldl f z m = foldl f z (uView m) foldl1 _ Empty = foldl1Empty foldl1 f (K1 m1) = foldl1 f m1 foldl1 f (K2 m2) = foldl1 f m2 foldl1 f (Union _ m1 m2) = foldl f (foldl1 f m1) m2 foldr1 _ Empty = foldr1Empty foldr1 f (K1 m1) = foldr1 f m1 foldr1 f (K2 m2) = foldr1 f m2 foldr1 f (Union _ m1 m2) = foldr f (foldr1 f m2) m1 -- | @'TrieMap' ('Either' k1 k2) a@ is essentially a @(TrieMap k1 a, TrieMap k2 a)@, but -- specialized for the cases where one or both maps are empty. instance (TrieKey k1, TrieKey k2) => TrieKey (Either k1 k2) where {-# SPECIALIZE instance TrieKey (Either () ()) #-} data TrieMap (Either k1 k2) a = Empty | K1 (TrieMap k1 a) | K2 (TrieMap k2 a) | Union !Int (TrieMap k1 a) (TrieMap k2 a) data Hole (Either k1 k2) a = HoleX0 (Hole k1 a) | HoleX2 (Hole k1 a) (TrieMap k2 a) | Hole0X (Hole k2 a) | Hole1X (TrieMap k1 a) (Hole k2 a) emptyM = Empty singletonM = either singletonL singletonR getSimpleM (UVIEW m1 m2) = mSimple m1 `mplus` mSimple m2 where mSimple :: TrieKey k => Maybe (TrieMap k a) -> Simple a mSimple = maybe mzero getSimpleM sizeM Empty = 0 sizeM (K1 m1) = sizeM m1 sizeM (K2 m2) = sizeM m2 sizeM (Union s _ _) = s lookupM (Left k) (UVIEW m1 _) = liftMaybe m1 >>= lookupM k lookupM (Right k) (UVIEW _ m2) = liftMaybe m2 >>= lookupM k traverseM f (Union _ m1 m2) = union <$> traverseM f m1 <*> traverseM f m2 traverseM f (K1 m1) = K1 <$> traverseM f m1 traverseM f (K2 m2) = K2 <$> traverseM f m2 traverseM _ _ = pure Empty fmapM f (Union _ m1 m2) = fmapM f m1 `union` fmapM f m2 fmapM f (K1 m1) = K1 (fmapM f m1) fmapM f (K2 m2) = K2 (fmapM f m2) fmapM _ _ = Empty mapMaybeM f (UVIEW m1 m2) = (m1 >>= mapMaybeM' f) ^ (m2 >>= mapMaybeM' f) mapEitherM f (UVIEW m1 m2) = (# m1L ^ m2L, m1R ^ m2R #) where !(# m1L, m1R #) = mapEitherM'' f m1 !(# m2L, m2R #) = mapEitherM'' f m2 unionM _ Empty m2 = m2 unionM f m1@(UVIEW m11 m12) m2@(UVIEW m21 m22) | Empty <- m2 = m1 | otherwise = unionMaybe (unionM' f) m11 m21 ^ unionMaybe (unionM' f) m12 m22 isectM f (UVIEW m11 m12) (UVIEW m21 m22) = isectMaybe (isectM' f) m11 m21 ^ isectMaybe (isectM' f) m12 m22 diffM f m1@(UVIEW m11 m12) m2@(UVIEW m21 m22) | Empty <- m2 = m1 | otherwise = diffMaybe (diffM' f) m11 m21 ^ diffMaybe (diffM' f) m12 m22 isSubmapM (<=) (UVIEW m11 m12) (UVIEW m21 m22) = subMaybe (isSubmapM (<=)) m11 m21 && subMaybe (isSubmapM (<=)) m12 m22 insertWithM f (Left k) a (UVIEW m1 m2) = Just (insertWithM' f k a m1) ^ m2 insertWithM f (Right k) a (UVIEW m1 m2) = m1 ^ Just (insertWithM' f k a m2) fromListM f = onPair (&) (fromListM f) (fromListM f) . partEithers fromAscListM f = onPair (&) (fromAscListM f) (fromAscListM f) . partEithers fromDistAscListM = onPair (&) fromDistAscListM fromDistAscListM . partEithers singleHoleM = either (HoleX0 . singleHoleM) (Hole0X . singleHoleM) beforeM hole = case hView hole of Hole1 h1 __ -> guardNullM (beforeM h1) ^ Nothing Hole2 m1 h2 -> m1 ^ guardNullM (beforeM h2) beforeWithM a hole = case hView hole of Hole1 h1 __ -> K1 (beforeWithM a h1) Hole2 m1 h2 -> m1 ^ Just (beforeWithM a h2) afterM hole = case hView hole of Hole1 h1 m2 -> guardNullM (afterM h1) ^ m2 Hole2 __ h2 -> Nothing ^ guardNullM (afterM h2) afterWithM a hole = case hView hole of Hole1 h1 m2 -> Just (afterWithM a h1) ^ m2 Hole2 __ h2 -> K2 (afterWithM a h2) searchMC (Left k) (UVIEW m1 m2) = mapSearch (`hole1` m2) (searchMC' k m1) searchMC (Right k) (UVIEW m1 m2) = mapSearch (hole2 m1) (searchMC' k m2) indexM i (K1 m1) = onThird HoleX0 (indexM i) m1 indexM i (K2 m2) = onThird Hole0X (indexM i) m2 indexM i (Union _ m1 m2) | i < s1 = onThird (`HoleX2` m2) (indexM i) m1 | otherwise = onThird (Hole1X m1) (indexM (i - s1)) m2 where !s1 = sizeM m1 indexM _ _ = indexFail () extractHoleM (UVIEW !m1 !m2) = holes1 `mplus` holes2 where holes1 = holes extractHoleM (`hole1` m2) m1 holes2 = holes extractHoleM (hole2 m1) m2 clearM hole = case hView hole of Hole1 h1 m2 -> clearM' h1 ^ m2 Hole2 m1 h2 -> m1 ^ clearM' h2 assignM v hole = case hView hole of Hole1 h1 m2 -> Just (assignM v h1) ^ m2 Hole2 m1 h2 -> m1 ^ Just (assignM v h2) unifierM (Left k') (Left k) a = HoleX0 <$> unifierM k' k a unifierM (Left k') (Right k) a = Just $ HoleX2 (singleHoleM k') (singletonM k a) unifierM (Right k') (Left k) a = Just $ Hole1X (singletonM k a) (singleHoleM k') unifierM (Right k') (Right k) a = Hole0X <$> unifierM k' k a {-# INLINE holes #-} holes :: (Functor m, Functor f, MonadPlus m) => (a -> m (f b)) -> (b -> c) -> Maybe a -> m (f c) holes k f (Just a) = fmap f <$> k a holes _ _ Nothing = mzero onPair :: (c -> d -> e) -> (a -> c) -> (b -> d) -> (a, b) -> e onPair f g h (a, b) = f (g a) (h b) partEithers :: [(Either a b, x)] -> ([(a, x)], [(b, x)]) partEithers = foldr part ([], []) where part (Left x, z) (xs, ys) = ((x,z):xs, ys) part (Right y, z) (xs, ys) = (xs, (y, z):ys)