diff --git a/Control/Monad/Lookup.hs b/Control/Monad/Lookup.hs new file mode 100644 --- /dev/null +++ b/Control/Monad/Lookup.hs @@ -0,0 +1,18 @@ +module Control.Monad.Lookup where + +import Control.Monad + +newtype Lookup r a = Lookup {runLookup :: r -> (a -> r) -> r} + +instance Functor (Lookup r) where + fmap f m = Lookup $ \ no yes -> runLookup m no (yes . f) + +instance Monad (Lookup r) where + return a = Lookup $ \ _ yes -> yes a + m >>= k = Lookup $ \ no yes -> + runLookup m no (\ a -> runLookup (k a) no yes) + fail _ = mzero + +instance MonadPlus (Lookup r) where + mzero = Lookup $ \ no _ -> no + m `mplus` k = Lookup $ \ no yes -> runLookup m (runLookup k no yes) yes diff --git a/Data/TrieMap.hs b/Data/TrieMap.hs --- a/Data/TrieMap.hs +++ b/Data/TrieMap.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UnboxedTuples, ImplicitParams, RecordWildCards, FlexibleContexts #-} module Data.TrieMap ( -- * Map type @@ -75,19 +75,31 @@ foldrWithKey, foldlWithKey, -- * Conversion + keysSet, + -- ** Lists elems, keys, - keysSet, assocs, - -- ** Lists fromList, fromListWith, fromListWithKey, + -- ** Vectors + elemsVector, + keysVector, + assocsVector, + fromVector, + fromVectorWith, + fromVectorWithKey, -- ** Ordered lists fromAscList, fromAscListWith, fromAscListWithKey, fromDistinctAscList, + -- ** Ordered vectors + fromAscVector, + fromAscVectorWith, + fromAscVectorWithKey, + fromDistinctAscVector, -- * Filter filter, filterWithKey, @@ -125,22 +137,27 @@ maxViewWithKey ) where +import Control.Monad import Control.Monad.Ends +import Control.Monad.Lookup import Data.TrieMap.Class import Data.TrieMap.Class.Instances() -import Data.TrieMap.TrieKey +import Data.TrieMap.TrieKey hiding (union, isect, diff, mapMaybe, mapEither) +import qualified Data.TrieMap.TrieKey.Projection as Proj +import qualified Data.TrieMap.TrieKey.SetOp as Set import Data.TrieMap.Representation import Data.TrieMap.Representation.Instances () -import Data.TrieMap.Sized -import Data.TrieMap.Utils -import Control.Applicative hiding (empty) -import Control.Monad import qualified Data.Foldable as F import Data.Maybe hiding (mapMaybe) -import Data.Monoid(Monoid(..)) +import Data.Vector.Build +import qualified Data.Vector.Generic as G +import Data.Vector.Fusion.Util (unId) +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Monadic as S + import GHC.Exts (build) import Prelude hiding (lookup, foldr, null, map, filter, reverse) @@ -179,14 +196,14 @@ -- | /O(1)/. Is the map empty? {-# INLINE null #-} null :: TKey k => TMap k a -> Bool -null (TMap m) = nullM m +null (TMap m) = isNull m -- | Lookup the value at a key in the map. -- -- The function will return the corresponding value as @('Just' value)@, or 'Nothing' if the key isn't in the map. {-# INLINE lookup #-} lookup :: TKey k => k -> TMap k a -> Maybe a -lookup k (TMap m) = option (lookupM (toRep k) m) Nothing (Just . getValue) +lookup k (TMap m) = runLookup (lookupMC (toRep k) m) Nothing (Just . getValue) -- | The expression @('findWithDefault' def k map)@ returns the value at key @k@ or returns default value @def@ -- when the key is not in the map. @@ -342,7 +359,7 @@ -- | Map each key\/element pair to an action, evaluate these actions from left to right, and collect the results. {-# INLINE traverseWithKey #-} traverseWithKey :: (TKey k, Applicative f) => (k -> a -> f b) -> TMap k a -> f (TMap k b) -traverseWithKey f (TMap m) = TMap <$> traverseM (\ (Assoc k a) -> Assoc k <$> f k a) m +traverseWithKey f (TMap m) = TMap <$> traverse (\ (Assoc k a) -> Assoc k <$> f k a) m -- | Map a function over all values in the map. -- @@ -357,7 +374,7 @@ -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")] {-# INLINEABLE mapWithKey #-} mapWithKey :: TKey k => (k -> a -> b) -> TMap k a -> TMap k b -mapWithKey f (TMap m) = TMap (fmapM (\ (Assoc k a) -> Assoc k (f k a)) m) +mapWithKey f (TMap m) = TMap (fmap (\ (Assoc k a) -> Assoc k (f k a)) m) -- | -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. @@ -445,7 +462,7 @@ -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] {-# INLINEABLE unionMaybeWithKey #-} unionMaybeWithKey :: TKey k => (k -> a -> a -> Maybe a) -> TMap k a -> TMap k a -> TMap k a -unionMaybeWithKey f (TMap m1) (TMap m2) = TMap (unionM f' m1 m2) where +unionMaybeWithKey f (TMap m1) (TMap m2) = TMap (Set.union f' m1 m2) where f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b -- | 'symmetricDifference' is equivalent to @'unionMaybeWith' (\ _ _ -> Nothing)@. @@ -488,7 +505,7 @@ -- @'mapMaybe' 'id' ('intersectionWithKey' f m1 m2)@. {-# INLINEABLE intersectionMaybeWithKey #-} intersectionMaybeWithKey :: TKey k => (k -> a -> b -> Maybe c) -> TMap k a -> TMap k b -> TMap k c -intersectionMaybeWithKey f (TMap m1) (TMap m2) = TMap (isectM f' m1 m2) where +intersectionMaybeWithKey f (TMap m1) (TMap m2) = TMap (Set.isect f' m1 m2) where f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b -- | Difference of two maps. @@ -529,7 +546,7 @@ -- > == singleton 3 "3:b|B" {-# INLINEABLE differenceWithKey #-} differenceWithKey :: TKey k => (k -> a -> b -> Maybe a) -> TMap k a -> TMap k b -> TMap k a -differenceWithKey f (TMap m1) (TMap m2) = TMap (diffM f' m1 m2) where +differenceWithKey f (TMap m1) (TMap m2) = TMap (Set.diff f' m1 m2) where f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b -- | Retrieves the value associated with minimal key of the @@ -632,10 +649,10 @@ deleteFindMin :: TKey k => TMap k a -> ((k, a), TMap k a) deleteFindMin m = fromMaybe (error "Cannot return the minimal element of an empty map") (minViewWithKey m) --- | Delete and find the minimal element. +-- | Delete and find the maximal element. -- --- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) --- > deleteFindMin Error: can not return the minimal element of an empty map +-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList[(3,"b"),(5,"a")]) +-- > deleteFindMax Error: can not return the maximal element of an empty map {-# INLINEABLE deleteFindMax #-} deleteFindMax :: TKey k => TMap k a -> ((k, a), TMap k a) deleteFindMax m = fromMaybe (error "Cannot return the maximal element of an empty map") (maxViewWithKey m) @@ -671,6 +688,13 @@ elems :: TKey k => TMap k a -> [a] elems m = build (\ c n -> foldrWithKey (\ _ a -> c a) n m) +{-# INLINE elemsVector #-} +-- | +-- Return all elements of the map in the ascending order of their keys. +-- Does not currently fuse. +elemsVector :: (TKey k, G.Vector v a) => TMap k a -> v a +elemsVector (TMap m) = toVectorMapN (sizeM m) (\ (Assoc _ a) -> a) m + -- | Return all keys of the map in ascending order. -- -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] @@ -679,6 +703,11 @@ keys :: TKey k => TMap k a -> [k] keys m = build (\ c n -> foldrWithKey (\ k _ -> c k) n m) +-- | Return all keys of the map in ascending order. +-- Does not currently fuse. +keysVector :: (TKey k, G.Vector v k) => TMap k a -> v k +keysVector (TMap m) = toVectorMapN (sizeM m) (\ (Assoc k _) -> k) m + -- | Return all key\/value pairs in the map in ascending key order. -- -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] @@ -687,6 +716,12 @@ assocs :: TKey k => TMap k a -> [(k, a)] assocs m = build (\ c n -> foldrWithKey (curry c) n m) +{-# INLINE assocsVector #-} +-- | Return all key\/value pairs in the map in ascending key order. +-- Does not currently fuse. +assocsVector :: (TKey k, G.Vector v (k, a)) => TMap k a -> v (k, a) +assocsVector (TMap m) = toVectorMapN (sizeM m) (\ (Assoc k a) -> (k, a)) m + -- | Map values and separate the 'Left' and 'Right' results. -- -- > let f a = if a < "c" then Left a else Right a @@ -709,7 +744,7 @@ -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) {-# INLINEABLE mapEitherWithKey #-} mapEitherWithKey :: TKey k => (k -> a -> Either b c) -> TMap k a -> (TMap k b, TMap k c) -mapEitherWithKey f (TMap m) = case mapEitherM f' m of +mapEitherWithKey f (TMap m) = case Proj.mapEither f' m of (# mL, mR #) -> (TMap mL, TMap mR) where f' (Assoc k a) = case f k a of Left b -> (# Just (Assoc k b), Nothing #) @@ -729,7 +764,7 @@ -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" {-# INLINEABLE mapMaybeWithKey #-} mapMaybeWithKey :: TKey k => (k -> a -> Maybe b) -> TMap k a -> TMap k b -mapMaybeWithKey f (TMap m) = TMap (mapMaybeM (\ (Assoc k a) -> Assoc k <$> f k a) m) +mapMaybeWithKey f (TMap m) = TMap (Proj.mapMaybe (\ (Assoc k a) -> Assoc k <$> f k a) m) -- | Partition the map according to a predicate. The first -- map contains all elements that satisfy the predicate, the second all @@ -760,14 +795,14 @@ -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty {-# INLINE filter #-} filter :: TKey k => (a -> Bool) -> TMap k a -> TMap k a -filter = filterWithKey . const +filter p = mapMaybeWithKey (\ _ a -> mfilter p (Just a)) -- | Filter all keys\/values that satisfy the predicate. -- -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" {-# INLINE filterWithKey #-} filterWithKey :: TKey k => (k -> a -> Bool) -> TMap k a -> TMap k a -filterWithKey p = mapMaybeWithKey (\ k a -> if p k a then Just a else Nothing) +filterWithKey p = mapMaybeWithKey (\ k a -> mfilter (p k) (Just a)) -- | The expression (@'split' k map@) is a pair @(map1,map2)@ where -- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. @@ -821,9 +856,25 @@ -} {-# INLINEABLE isSubmapOfBy #-} isSubmapOfBy :: TKey k => (a -> b -> Bool) -> TMap k a -> TMap k b -> Bool -isSubmapOfBy (<=) (TMap m1) (TMap m2) = isSubmapM (<<=) m1 m2 where - Assoc _ a <<= Assoc _ b = a <= b +isSubmapOfBy (<=) (TMap m1) (TMap m2) = let ?le = \ (Assoc _ a) (Assoc _ b) -> a <= b in m1 <=? m2 +{-# INLINE fromFoldStream #-} +fromFoldStream :: (Repr k, TrieKey (Rep k), Monad m) => + FromList z (Rep k) (Assoc k a) -> Stream m (k, a) -> m (TMap k a) +fromFoldStream Foldl{..} (Stream suc s0 _) = run s0 where + run s = do + step <- suc s + case step of + Done -> return empty + Skip s' -> run s' + Yield (k, a) s' -> run' (begin (toRep k) (Assoc k a)) s' + run' stack s = do + step <- suc s + case step of + Done -> return (TMap (done stack)) + Skip s' -> run' stack s' + Yield (k, a) s' -> run' (snoc stack (toRep k) (Assoc k a)) s' + -- | Build a map from a list of key\/value pairs. See also 'fromAscList'. -- If the list contains more than one value for the same key, the last value -- for the key is retained. @@ -835,6 +886,11 @@ fromList :: TKey k => [(k, a)] -> TMap k a fromList = fromListWith const +{-# INLINE fromVector #-} +-- | Equivalent to @'fromList' ('G.toList' xs)@. +fromVector :: (TKey k, G.Vector v (k, a)) => v (k, a) -> TMap k a +fromVector = fromVectorWith const + -- | Build a map from an ascending list in linear time. -- /The precondition (input list is ascending) is not checked./ -- @@ -844,6 +900,11 @@ fromAscList :: TKey k => [(k, a)] -> TMap k a fromAscList = fromAscListWith const +{-# INLINE fromAscVector #-} +-- | Equivalent to @'fromAscList' ('G.toList' xs)@. +fromAscVector :: (TKey k, G.Vector v (k, a)) => v (k, a) -> TMap k a +fromAscVector = fromAscVectorWith const + -- | Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. -- -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] @@ -852,6 +913,11 @@ fromListWith :: TKey k => (a -> a -> a) -> [(k, a)] -> TMap k a fromListWith = fromListWithKey . const +{-# INLINE fromVectorWith #-} +-- | Equivalent to @'fromListWith' f ('G.toList' xs)@. +fromVectorWith :: (TKey k, G.Vector v (k, a)) => (a -> a -> a) -> v (k, a) -> TMap k a +fromVectorWith = fromVectorWithKey . const + -- | Build a map from an ascending list in linear time with a combining function for equal keys. -- /The precondition (input list is ascending) is not checked./ -- @@ -860,15 +926,26 @@ fromAscListWith :: TKey k => (a -> a -> a) -> [(k, a)] -> TMap k a fromAscListWith = fromAscListWithKey . const +{-# INLINE fromAscVectorWith #-} +-- | Equivalent to @'fromAscListWith' f ('G.toList' xs)@. +fromAscVectorWith :: (TKey k, G.Vector v (k, a)) => (a -> a -> a) -> v (k, a) -> TMap k a +fromAscVectorWith = fromAscVectorWithKey . const + -- | Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. -- -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] -- > fromListWith (++) [] == empty {-# INLINEABLE fromListWithKey #-} fromListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a -fromListWithKey f xs = TMap (fromListM f' [(toRep k, Assoc k a) | (k, a) <- xs]) +fromListWithKey f xs = unId $ fromFoldStream (uFold f') (S.fromList xs) where f' (Assoc k a) (Assoc _ b) = Assoc k (f k a b) +{-# INLINE fromVectorWithKey #-} +-- | Equivalent to @'fromListWithKey' f ('G.toList' xs)@. +fromVectorWithKey :: (TKey k, G.Vector v (k, a)) => (k -> a -> a -> a) -> v (k, a) -> TMap k a +fromVectorWithKey f xs = unId $ fromFoldStream (uFold f') (G.stream xs) + where f' (Assoc k a) (Assoc _ b) = Assoc k (f k a b) + -- | Build a map from an ascending list in linear time. -- /The precondition (input list is ascending) is not checked./ -- @@ -876,17 +953,28 @@ -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] {-# INLINEABLE fromAscListWithKey #-} fromAscListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a -fromAscListWithKey f xs = TMap (fromAscListM f' [(toRep k, Assoc k a) | (k, a) <- xs]) +fromAscListWithKey f xs = unId $ fromFoldStream (aFold f') (S.fromList xs) where f' (Assoc k a) (Assoc _ b) = Assoc k (f k a b) +{-# INLINE fromAscVectorWithKey #-} +-- | Equivalent to @'fromAscListWithKey' f ('G.toList' xs)@. +fromAscVectorWithKey :: (TKey k, G.Vector v (k, a)) => (k -> a -> a -> a) -> v (k, a) -> TMap k a +fromAscVectorWithKey f xs = unId $ fromFoldStream (aFold f') (G.stream xs) + where f' (Assoc k a) (Assoc _ b) = Assoc k (f k a b) + -- | Build a map from an ascending list of distinct elements in linear time. -- /The precondition is not checked./ -- -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] {-# INLINEABLE fromDistinctAscList #-} fromDistinctAscList :: TKey k => [(k, a)] -> TMap k a -fromDistinctAscList xs = TMap (fromDistAscListM [(toRep k, Assoc k a) | (k, a) <- xs]) +fromDistinctAscList xs = unId $ fromFoldStream daFold (S.fromList xs) +{-# INLINE fromDistinctAscVector #-} +-- | Equivalent to @'fromDistinctAscList' ('G.toList' xs)@. +fromDistinctAscVector :: (TKey k, G.Vector v (k, a)) => v (k, a) -> TMap k a +fromDistinctAscVector xs = unId $ fromFoldStream daFold (G.stream xs) + -- | /O(1)/. The number of elements in the map. -- -- > size empty == 0 @@ -918,7 +1006,7 @@ -- > keysSet empty == Data.TrieSet.empty {-# INLINE keysSet #-} keysSet :: TKey k => TMap k a -> TSet k -keysSet (TMap m) = TSet (fmapM (\ (Assoc k _) -> Elem k) m) +keysSet (TMap m) = TSet (fmap (\ (Assoc k _) -> Elem k) m) -- | /O(1)/. The key marking the position of the \"hole\" in the map. {-# INLINE key #-} @@ -967,11 +1055,8 @@ -- @'elemAt' i m == let (v, loc) = 'index' i m in ('key' loc, v)@ {-# INLINEABLE index #-} index :: TKey k => Int -> TMap k a -> (a, TLocation k a) -index i m - | i < 0 || i >= size m - = error "TrieMap.index: index out of range" -index i (TMap m) = case indexM i m of - (# _, Assoc k a, hole #) -> (a, TLoc k hole) +index i (TMap m) = case indexM m (unbox i) of + (# _, Assoc k a, hole #) -> (a, TLoc k hole) {-# INLINE extract #-} extract :: (TKey k, Functor m, MonadPlus m) => TMap k a -> m (a, TLocation k a) diff --git a/Data/TrieMap/Class.hs b/Data/TrieMap/Class.hs --- a/Data/TrieMap/Class.hs +++ b/Data/TrieMap/Class.hs @@ -4,12 +4,7 @@ import Data.TrieMap.TrieKey import Data.TrieMap.Representation.Class -import Data.TrieMap.Sized -import Data.Functor -import Data.Foldable -import Data.Traversable - import Prelude hiding (foldr, foldl, foldl1, foldr1) -- | A map from keys @k@ to values @a@, backed by a trie. @@ -26,18 +21,14 @@ instance (Repr k, TrieKey (Rep k)) => TKey k instance TKey k => Functor (TMap k) where - fmap f (TMap m) = TMap (fmapM (fmap f) m) + fmap f (TMap m) = TMap (fmap (fmap f) m) instance TKey k => Foldable (TMap k) where foldMap f (TMap m) = foldMap (foldMap f) m foldr f z (TMap m) = foldr (flip $ foldr f) z m foldl f z (TMap m) = foldl (foldl f) z m - foldr1 f (TMap m) = getElem (foldr1 f' m') where - f' (Elem a) (Elem b) = Elem (f a b) - m' = fmapM (\ (Assoc _ a) -> Elem a) m - foldl1 f (TMap m) = getElem (foldl1 f' m') where - f' (Elem a) (Elem b) = Elem (f a b) - m' = fmapM (\ (Assoc _ a) -> Elem a) m + foldr1 f (TMap m) = foldr1 f (getValue <$> m) + foldl1 f (TMap m) = foldl1 f (getValue <$> m) instance TKey k => Traversable (TMap k) where - traverse f (TMap m) = TMap <$> traverseM (traverse f) m + traverse f (TMap m) = TMap <$> traverse (traverse f) m diff --git a/Data/TrieMap/Key.hs b/Data/TrieMap/Key.hs --- a/Data/TrieMap/Key.hs +++ b/Data/TrieMap/Key.hs @@ -1,30 +1,55 @@ -{-# LANGUAGE TypeFamilies, MagicHash, CPP, FlexibleInstances #-} +{-# LANGUAGE TypeFamilies, CPP, FlexibleInstances, FlexibleContexts, NamedFieldPuns, RecordWildCards, UndecidableInstances #-} +{-# LANGUAGE MultiParamTypeClasses, UnboxedTuples #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.Key () where -import Data.Functor -import Data.Foldable - import Data.TrieMap.Class import Data.TrieMap.TrieKey -import Data.TrieMap.Sized import Data.TrieMap.Representation.Class import Data.TrieMap.Modifiers import Prelude hiding (foldr, foldl, foldr1, foldl1) -keyMap :: (TKey k, Sized a) => TrieMap (Rep k) a -> TrieMap (Key k) a +type RepMap k = TrieMap (Rep k) + +keyMap :: (Repr k, TrieKey (Rep k), Sized a) => TrieMap (Rep k) a -> TrieMap (Key k) a keyMap m = KeyMap (sizeM m) m #define KMAP(m) KeyMap{tMap = m} +#define CONTEXT(cl) (Repr k, TrieKey (Rep k), cl (RepMap k)) -instance TKey k => Foldable (TrieMap (Key k)) where +instance CONTEXT(Foldable) => Foldable (TrieMap (Key k)) where foldMap f KMAP(m) = foldMap f m foldr f z KMAP(m) = foldr f z m foldl f z KMAP(m) = foldl f z m - foldr1 f KMAP(m) = foldr1 f m - foldl1 f KMAP(m) = foldl1 f m +instance CONTEXT(Functor) => Functor (TrieMap (Key k)) where + fmap f KeyMap{..} = KeyMap{sz, tMap = f <$> tMap} + +instance CONTEXT(Traversable) => Traversable (TrieMap (Key k)) where + traverse f KeyMap{..} = KeyMap sz <$> traverse f tMap + +instance CONTEXT(Subset) => Subset (TrieMap (Key k)) where + KMAP(m1) <=? KMAP(m2) = m1 <=? m2 + +instance (Repr k, TrieKey (Rep k), Buildable (RepMap k) (Rep k)) => Buildable (TrieMap (Key k)) (Key k) where + type UStack (TrieMap (Key k)) = UMStack (Rep k) + uFold = fmap keyMap . mapFoldlKeys keyRep . uFold + type AStack (TrieMap (Key k)) = AMStack (Rep k) + aFold = fmap keyMap . mapFoldlKeys keyRep . aFold + type DAStack (TrieMap (Key k)) = DAMStack (Rep k) + daFold = keyMap <$> mapFoldlKeys keyRep daFold + +#define SETOP(op) op f KMAP(m1) KMAP(m2) = keyMap (op f m1 m2) +instance CONTEXT(SetOp) => SetOp (TrieMap (Key k)) where + SETOP(union) + SETOP(isect) + SETOP(diff) + +instance CONTEXT(Project) => Project (TrieMap (Key k)) where + mapMaybe f KMAP(m) = keyMap $ mapMaybe f m + mapEither f KMAP(m) = both keyMap (mapEither f) m + -- | @'TrieMap' ('Key' k) a@ is a wrapper around a @TrieMap (Rep k) a@. instance TKey k => TrieKey (Key k) where data TrieMap (Key k) a = KeyMap {sz :: !Int, tMap :: !(TrieMap (Rep k) a)} @@ -34,15 +59,7 @@ singletonM (Key k) a = KeyMap (getSize a) (singletonM (toRep k) a) getSimpleM KMAP(m) = getSimpleM m sizeM = sz - lookupM (Key k) KMAP(m) = lookupM (toRep k) m - traverseM f KMAP(m) = keyMap <$> traverseM f m - fmapM f KMAP(m) = keyMap (fmapM f m) - mapMaybeM f KMAP(m) = keyMap (mapMaybeM f m) - mapEitherM f KMAP(m) = both keyMap keyMap (mapEitherM f) m - unionM f KMAP(m1) KMAP(m2) = keyMap (unionM f m1 m2) - isectM f KMAP(m1) KMAP(m2) = keyMap (isectM f m1 m2) - diffM f KMAP(m1) KMAP(m2) = keyMap (diffM f m1 m2) - isSubmapM (<=) KMAP(m1) KMAP(m2) = isSubmapM (<=) m1 m2 + lookupMC (Key k) KMAP(m) = lookupMC (toRep k) m singleHoleM (Key k) = KeyHole (singleHoleM (toRep k)) beforeM (KeyHole hole) = keyMap (beforeM hole) @@ -50,12 +67,13 @@ afterM (KeyHole hole) = keyMap (afterM hole) afterWithM a (KeyHole hole) = keyMap (afterWithM a hole) searchMC (Key k) KMAP(m) = mapSearch KeyHole (searchMC (toRep k) m) - indexM i KMAP(m) = onThird KeyHole (indexM i) m + indexM KMAP(m) i = case indexM m i of + (# i', a, hole #) -> (# i', a, KeyHole hole #) extractHoleM KMAP(m) = fmap KeyHole <$> extractHoleM m assignM v (KeyHole hole) = keyMap (assignM v hole) clearM (KeyHole hole) = keyMap (clearM hole) insertWithM f (Key k) a KMAP(m) = keyMap (insertWithM f (toRep k) a m) - fromListM f xs = keyMap (fromListM f [(toRep k, a) | (Key k, a) <- xs]) - fromAscListM f xs = keyMap (fromAscListM f [(toRep k, a) | (Key k, a) <- xs]) - fromDistAscListM xs = keyMap (fromDistAscListM [(toRep k, a) | (Key k, a) <- xs]) + +keyRep :: (Repr k, TrieKey (Rep k)) => Key k -> Rep k +keyRep (Key k) = toRep k diff --git a/Data/TrieMap/Modifiers.hs b/Data/TrieMap/Modifiers.hs --- a/Data/TrieMap/Modifiers.hs +++ b/Data/TrieMap/Modifiers.hs @@ -1,9 +1,16 @@ {-# LANGUAGE FlexibleContexts, UndecidableInstances, TypeFamilies #-} + +-- | +-- Operators for use in 'Repr' instances for types. module Data.TrieMap.Modifiers where import Data.TrieMap.Representation.Class +-- | Denotes that maps on this type should be implemented with traditional binary search trees. newtype Ordered a = Ord {unOrd :: a} deriving (Eq, Ord) + +-- | Denotes that maps on this type should be treated as reversed. For instance, @'Rep' 'Int'@ might be +-- implemented as @'Either' ('Rev' Word) Word@, to handle negative numbers properly. newtype Rev k = Rev {getRev :: k} deriving (Eq) instance Ord k => Ord (Rev k) where compare (Rev a) (Rev b) = compare b a @@ -11,7 +18,6 @@ Rev a <= Rev b = b <= a (>) = flip (<) (>=) = flip (<=) - instance Functor Ordered where fmap f (Ord a) = Ord (f a) @@ -19,6 +25,17 @@ instance Functor Rev where fmap f (Rev a) = Rev (f a) +-- | Indicates that the map for this type should be bootstrapped from its @TKey@ instance. +-- This modifier is necessary to define a @TKey@ instance for a recursively defined type. +-- For example: +-- +-- > data Tree = Node Char [Tree] +-- > +-- > instance 'Repr' Tree where +-- > type 'Rep' Tree = ('Rep' 'Char', ['Key' Tree]) +-- > 'toRep' (Node label children) = ('toRep' label, 'map' 'Key' children) +-- > ... +-- newtype Key k = Key {getKey :: k} instance (Repr k, Eq (Rep k)) => Eq (Key k) where diff --git a/Data/TrieMap/OrdMap.hs b/Data/TrieMap/OrdMap.hs --- a/Data/TrieMap/OrdMap.hs +++ b/Data/TrieMap/OrdMap.hs @@ -1,18 +1,16 @@ {-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, PatternGuards, MagicHash, CPP, TupleSections, NamedFieldPuns, FlexibleInstances #-} +{-# LANGUAGE RecordWildCards, ImplicitParams, GeneralizedNewtypeDeriving, StandaloneDeriving, MultiParamTypeClasses #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.OrdMap () where +import Control.Monad.Lookup + import Data.TrieMap.TrieKey import Data.TrieMap.Sized import Data.TrieMap.Modifiers -import Control.Applicative (Applicative(..), (<$>)) -import Control.Monad hiding (join) - -import Data.Foldable -import Data.Monoid - import Prelude hiding (lookup, foldr, foldl, foldr1, foldl1, map) +import GHC.Exts #define DELTA 5 #define RATIO 2 @@ -30,6 +28,15 @@ #define TIP SNode{node=Tip} #define BIN(args) SNode{node=Bin args} +-- Morally reprehensible exploitation of generalized newtype deriving. +class ImmoralCast a b where + immoralCast :: SNode k a -> SNode k b + +instance ImmoralCast a a where + immoralCast = id + +deriving instance ImmoralCast a (Elem a) + instance Sized a => Sized (Node k a) where getSize# m = unbox $ case m of Tip -> 0 @@ -48,6 +55,38 @@ tip :: SNode k a tip = SNode 0 0 Tip +instance Ord k => Subset (TrieMap (Ordered k)) where + OrdMap m1 <=? OrdMap m2 = m1 <=? m2 + +instance Functor (TrieMap (Ordered k)) where + fmap f (OrdMap m) = OrdMap (f <$> m) + +instance Foldable (TrieMap (Ordered k)) where + foldMap f (OrdMap m) = foldMap f m + foldr f z (OrdMap m) = foldr f z m + foldl f z (OrdMap m) = foldl f z m + +instance Traversable (TrieMap (Ordered k)) where + traverse f (OrdMap m) = OrdMap <$> traverse f m + +instance Ord k => Buildable (TrieMap (Ordered k)) (Ordered k) where + type UStack (TrieMap (Ordered k)) = TrieMap (Ordered k) + uFold = defaultUFold emptyM singletonM insertWithM + type AStack (TrieMap (Ordered k)) = Distinct (Ordered k) (Stack k) + aFold = combineFold daFold + type DAStack (TrieMap (Ordered k)) = Stack k + daFold = OrdMap <$> mapFoldlKeys unOrd fromDistAscList + +#define SETOP(op) op f (OrdMap m1) (OrdMap m2) = OrdMap (op f m1 m2) +instance Ord k => SetOp (TrieMap (Ordered k)) where + SETOP(union) + SETOP(isect) + SETOP(diff) + +instance Ord k => Project (TrieMap (Ordered k)) where + mapMaybe f (OrdMap m) = OrdMap $ mapMaybe f m + mapEither f (OrdMap m) = both OrdMap (mapEither f) m + -- | @'TrieMap' ('Ordered' k) a@ is based on "Data.Map". instance Ord k => TrieKey (Ordered k) where newtype TrieMap (Ordered k) a = OrdMap (SNode k a) @@ -56,23 +95,13 @@ | Full k !(Path k a) !(SNode k a) !(SNode k a) emptyM = OrdMap tip singletonM (Ord k) a = OrdMap (singleton k a) - lookupM (Ord k) (OrdMap m) = lookup k m + lookupMC (Ord k) (OrdMap m) = lookupC k m getSimpleM (OrdMap m) = case m of TIP -> Null BIN(_ a TIP TIP) -> Singleton a _ -> NonSimple sizeM (OrdMap m) = sz m - traverseM f (OrdMap m) = OrdMap <$> traverse f m - fmapM f (OrdMap m) = OrdMap (map f m) - mapMaybeM f (OrdMap m) = OrdMap (mapMaybe f m) - mapEitherM f (OrdMap m) = both OrdMap OrdMap (mapEither f) m - isSubmapM (<=) (OrdMap m1) (OrdMap m2) = isSubmap (<=) m1 m2 - fromAscListM f xs = OrdMap $ fromAscList f [(k, a) | (Ord k, a) <- xs] - fromDistAscListM xs = OrdMap $ fromDistinctAscList [(k, a) | (Ord k, a) <- xs] - unionM f (OrdMap m1) (OrdMap m2) = OrdMap $ hedgeUnion f (const LT) (const GT) m1 m2 - isectM f (OrdMap m1) (OrdMap m2) = OrdMap $ isect f m1 m2 - diffM f (OrdMap m1) (OrdMap m2) = OrdMap $ hedgeDiff f (const LT) (const GT) m1 m2 singleHoleM (Ord k) = Empty k Root beforeM (Empty _ path) = OrdMap $ before tip path @@ -84,14 +113,14 @@ afterWithM a (Empty k path) = OrdMap $ after (singleton k a) path afterWithM a (Full k path _ r) = OrdMap $ after (insertMin k a r) path searchMC (Ord k) (OrdMap m) = search k m - indexM i (OrdMap m) = indexT Root i m where - indexT path i BIN(kx x l r) - | i < sl = indexT (LeftBin kx x path r) i l - | i < sx = (# i - sl, x, Full kx path l r #) - | otherwise = indexT (RightBin kx x l path) (i - sx) r - where !sl = getSize l - !sx = getSize x + sl - indexT _ _ _ = indexFail () + indexM (OrdMap m) i = indexT Root i m where + indexT path !i SNode{sz, node = Bin kx x l r} + | i <# sl = indexT (LeftBin kx x path r) i l + | i <# sx = (# i -# sl, x, Full kx path l r #) + | otherwise = indexT (RightBin kx x l path) (i -# sx) r + where !sl = getSize# l + !sx = unbox $ sz - getSize r + indexT _ _ _ = indexFail () extractHoleM (OrdMap m) = extractHole Root m where extractHole path BIN(kx x l r) = extractHole (LeftBin kx x path r) l `mplus` @@ -105,29 +134,47 @@ assignM x (Full k path l r) = OrdMap $ rebuild (join k x l r) path unifierM (Ord k') (Ord k) a = case compare k' k of - EQ -> Nothing - LT -> Just $ Empty k' (LeftBin k a Root tip) - GT -> Just $ Empty k' (RightBin k a tip Root) + EQ -> mzero + LT -> return $ Empty k' (LeftBin k a Root tip) + GT -> return $ Empty k' (RightBin k a tip Root) + unifyM (Ord k1) a1 (Ord k2) a2 = case compare k1 k2 of + EQ -> mzero + LT -> return $ OrdMap $ bin k1 a1 tip (singleton k2 a2) + GT -> return $ OrdMap $ bin k1 a1 (singleton k2 a2) tip + + {-# INLINE insertWithM #-} + insertWithM f (Ord k) a (OrdMap m) = OrdMap (insertWith f k a m) +insertWith :: (Ord k, Sized a) => (a -> a) -> k -> a -> SNode k a -> SNode k a +insertWith f k a = k `seq` ins where + ins BIN(kx x l r) = case compare k kx of + EQ -> bin kx (f x) l r + LT -> balance kx x (ins l) r + GT -> balance kx x l (ins r) + ins TIP = singleton k a + rebuild :: Sized a => SNode k a -> Path k a -> SNode k a rebuild t Root = t rebuild t (LeftBin kx x path r) = rebuild (balance kx x t r) path rebuild t (RightBin kx x l path) = rebuild (balance kx x l t) path -lookup :: Ord k => k -> SNode k a -> Lookup a -lookup k = look where +lookupC :: Ord k => k -> SNode k a -> Lookup r a +lookupC k = look where look BIN(kx x l r) = case compare k kx of - LT -> lookup k l - EQ -> some x - GT -> lookup k r - look _ = none + LT -> look l + EQ -> return x + GT -> look r + look _ = mzero singleton :: Sized a => k -> a -> SNode k a singleton k a = bin k a tip tip -traverse :: (Applicative f, Sized b) => (a -> f b) -> SNode k a -> f (SNode k b) -traverse _ TIP = pure tip -traverse f BIN(k a l r) = balance k <$> f a <*> traverse f l <*> traverse f r +instance Traversable (SNode k) where + traverse f = trav where + trav TIP = pure tip + trav SNode{node = Bin k a l r, ..} = + let done a' l' r' = SNode sz count (Bin k a' l' r') in + done <$> f a <*> trav l <*> trav r instance Foldable (SNode k) where foldMap _ TIP = mempty @@ -137,76 +184,61 @@ foldr f z BIN(_ a l r) = foldr f (a `f` foldr f z r) l foldl _ z TIP = z foldl f z BIN(_ a l r) = foldl f (foldl f z l `f` a) r - - foldr1 _ TIP = foldr1Empty - foldr1 f BIN(_ a l TIP) = foldr f a l - foldr1 f BIN(_ a l r) = foldr f (a `f` foldr1 f r) l - - foldl1 _ TIP = foldl1Empty - foldl1 f BIN(_ a TIP r) = foldl f a r - foldl1 f BIN(_ a l r) = foldl f (foldl1 f l `f` a) r -instance Foldable (TrieMap (Ordered k)) where - foldMap f (OrdMap m) = foldMap f m - foldr f z (OrdMap m) = foldr f z m - foldl f z (OrdMap m) = foldl f z m - foldl1 f (OrdMap m) = foldl1 f m - foldr1 f (OrdMap m) = foldr1 f m - -map :: (Ord k, Sized b) => (a -> b) -> SNode k a -> SNode k b -map f BIN(k a l r) = join k (f a) (map f l) (map f r) -map _ _ = tip - -mapMaybe :: (Ord k, Sized b) => (a -> Maybe b) -> SNode k a -> SNode k b -mapMaybe f BIN(k a l r) = joinMaybe k (f a) (mapMaybe f l) (mapMaybe f r) -mapMaybe _ _ = tip +instance Functor (SNode k) where + fmap f = map where + map SNode{node = Bin k a l r, ..} = SNode {node = Bin k (f a) (map l) (map r), ..} + map _ = tip -mapEither :: (Ord k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> - SNode k a -> (# SNode k b, SNode k c #) -mapEither f BIN(k a l r) = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #) - where !(# aL, aR #) = f a; !(# lL, lR #) = mapEither f l; !(# rL, rR #) = mapEither f r -mapEither _ _ = (# tip, tip #) +instance Ord k => Project (SNode k) where + mapMaybe f = mMaybe where + mMaybe BIN(k a l r) = joinMaybe k (f a) (mMaybe l) (mMaybe r) + mMaybe _ = tip + mapEither f = mEither where + mEither BIN(k a l r) = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #) + where !(# aL, aR #) = f a + !(# lL, lR #) = mEither l + !(# rL, rR #) = mEither r + mEither _ = (# tip, tip #) -splitLookup :: (Ord k, Sized a) => k -> SNode k a -> (SNode k a -> Maybe a -> SNode k a -> r) -> r +splitLookup :: Ord k => k -> SNode k (Elem a) -> (SNode k (Elem a) -> Maybe (Elem a) -> SNode k (Elem a) -> r) -> r splitLookup k t cont = search k t (split Nothing) (split . Just) where split v (Empty _ path) = cont (before tip path) v (after tip path) split v (Full _ path l r) = cont (before l path) v (after r path) -isSubmap :: (Ord k, Sized a, Sized b) => LEq a b -> LEq (SNode k a) (SNode k b) -isSubmap _ TIP _ = True -isSubmap _ _ TIP = False -isSubmap (<=) BIN(kx x l r) t = splitLookup kx t result - where result _ Nothing _ = False - result tl (Just y) tr = x <= y && isSubmap (<=) l tl && isSubmap (<=) r tr +instance Ord k => Subset (SNode k) where + t1 <=? t2 = immoralCast t1 `subMap` immoralCast t2 where + TIP `subMap` _ = True + _ `subMap` TIP = False + BIN(kx x l r) `subMap` t = splitLookup kx t result + where result _ Nothing _ = False + result tl (Just y) tr = x <=? y && l `subMap` tl && r `subMap` tr -fromAscList :: (Eq k, Sized a) => (a -> a -> a) -> [(k, a)] -> SNode k a -fromAscList f xs = fromDistinctAscList (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 xx zz) xs - | otherwise = (kz,zz):combineEq' x xs +fromDistAscList :: (Eq k, Sized a) => Foldl (Stack k) k a (SNode k a) +fromDistAscList = Foldl{zero = tip, ..} where + incr !t (Yes t' stk) = No (incr (t' `glue` t) stk) + incr !t (No stk) = Yes t stk + incr !t End = Yes t End + + begin k a = Yes (singleton k a) End + + snoc stk k a = incr (singleton k a) stk + + roll !t End = t + roll !t (No stk) = roll t stk + roll !t (Yes t' stk) = roll (t' `glue` t) stk + + done = roll tip -fromDistinctAscList :: Sized a => [(k, a)] -> SNode k a -fromDistinctAscList 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 k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton 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 +data Stack k a = No (Stack k a) | Yes !(SNode k a) (Stack k a) | End - 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 k x l r) zs +instance Ord k => SetOp (SNode k) where + union f = hedgeUnion f (const LT) (const GT) + diff f = hedgeDiff f (const LT) (const GT) + isect f m1 m2 = immoralCast m1 `intersection` m2 where + t1@BIN(_ _ _ _) `intersection` BIN(k2 x2 l2 r2) = splitLookup k2 t1 result where + result tl found tr = joinMaybe k2 (found >>= \ (Elem x1') -> f x1' x2) (tl `intersection` l2) (tr `intersection` r2) + _ `intersection` _ = tip hedgeUnion :: (Ord k, Sized a) => (a -> a -> Maybe a) @@ -234,7 +266,7 @@ LT -> join kx x (filterGt cmp l) r GT -> filterGt cmp r EQ -> r - + filterLt :: (Ord k, Sized a) => (k -> Ordering) -> SNode k a -> SNode k a filterLt _ TIP = tip filterLt cmp BIN(kx x l r) @@ -256,15 +288,10 @@ trimLookupLo lo cmphi t@BIN(kx x l r) = case compare lo kx of LT -> case cmphi kx of - GT -> (option (lookup lo t) Nothing (\ a -> Just (lo, a)), t) + GT -> (runLookup (lookupC lo t) Nothing (\ a -> Just (lo, a)), t) _ -> trimLookupLo lo cmphi l GT -> trimLookupLo lo cmphi r EQ -> (Just (kx,x),trim (compare lo) cmphi r) - -isect :: (Ord k, Sized a, Sized b, Sized c) => (a -> b -> Maybe c) -> SNode k a -> SNode k b -> SNode k c -isect f t1@BIN(_ _ _ _) BIN(k2 x2 l2 r2) = splitLookup k2 t1 result where - result tl found tr = joinMaybe k2 (found >>= \ x1' -> f x1' x2) (isect f tl l2) (isect f tr r2) -isect _ _ _ = tip hedgeDiff :: (Ord k, Sized a) => (a -> b -> Maybe a) diff --git a/Data/TrieMap/ProdMap.hs b/Data/TrieMap/ProdMap.hs --- a/Data/TrieMap/ProdMap.hs +++ b/Data/TrieMap/ProdMap.hs @@ -1,24 +1,45 @@ -{-# LANGUAGE UnboxedTuples, TupleSections, PatternGuards, TypeFamilies, FlexibleInstances #-} - +{-# LANGUAGE TupleSections, TypeFamilies, FlexibleInstances, RecordWildCards, CPP, FlexibleContexts, UnboxedTuples #-} +{-# LANGUAGE MultiParamTypeClasses #-} module Data.TrieMap.ProdMap () where -import Data.TrieMap.Sized import Data.TrieMap.TrieKey -import Control.Monad -import Data.Functor -import Data.Foldable hiding (foldlM, foldrM) +import Prelude hiding (foldl, foldl1, foldr, foldr1) -import Data.Sequence ((|>)) -import qualified Data.Sequence as Seq +#define CONTEXT(cl) (TrieKey k1, TrieKey k2, cl (TrieMap k1), cl (TrieMap k2)) -import Prelude hiding (foldl, foldl1, foldr, foldr1) +instance CONTEXT(Functor) => Functor (TrieMap (k1, k2)) where + fmap f (PMap m) = PMap (fmap (fmap f) m) -instance (TrieKey k1, TrieKey k2) => Foldable (TrieMap (k1, k2)) where +instance CONTEXT(Foldable) => Foldable (TrieMap (k1, k2)) where foldMap f (PMap m) = foldMap (foldMap f) m foldr f z (PMap m) = foldr (flip $ foldr f) z m foldl f z (PMap m) = foldl (foldl f) z m +instance CONTEXT(Traversable) => Traversable (TrieMap (k1, k2)) where + traverse f (PMap m) = PMap <$> traverse (traverse f) m + +instance CONTEXT(Subset) => Subset (TrieMap (k1, k2)) where + PMap m1 <=? PMap m2 = m1 <<=? m2 + +instance (TrieKey k1, TrieKey k2) => Buildable (TrieMap (k1, k2)) (k1, k2) where + type UStack (TrieMap (k1, k2)) = TrieMap (k1, k2) + uFold = defaultUFold emptyM singletonM insertWithM + type AStack (TrieMap (k1, k2)) = Stack k1 k2 (DAMStack k1) (AMStack k2) + aFold f = prodFold daFold (aFold f) + type DAStack (TrieMap (k1, k2)) = Stack k1 k2 (DAMStack k1) (DAMStack k2) + daFold = prodFold daFold daFold + +#define SETOP(op,opM) op f (PMap m1) (PMap m2) = PMap ((op) ((opM) f) m1 m2) +instance CONTEXT(SetOp) => SetOp (TrieMap (k1, k2)) where + SETOP(union,unionM) + SETOP(isect,isectM) + SETOP(diff,diffM) + +instance CONTEXT(Project) => Project (TrieMap (k1, k2)) where + mapMaybe f (PMap m) = PMap $ mapMaybe (mapMaybeM f) m + mapEither f (PMap m) = both' PMap PMap (mapEither (mapEitherM f)) m + -- | @'TrieMap' (k1, k2) a@ is implemented as a @'TrieMap' k1 ('TrieMap' k2 a)@. instance (TrieKey k1, TrieKey k2) => TrieKey (k1, k2) where newtype TrieMap (k1, k2) a = PMap (TrieMap k1 (TrieMap k2 a)) @@ -28,22 +49,10 @@ singletonM (k1, k2) = PMap . singletonM k1 . singletonM k2 getSimpleM (PMap m) = getSimpleM m >>= getSimpleM sizeM (PMap m) = sizeM m - lookupM (k1, k2) (PMap m) = lookupM k1 m >>= lookupM k2 - traverseM f (PMap m) = PMap <$> traverseM (traverseM f) m - fmapM f (PMap m) = PMap (fmapM (fmapM f) m) - mapMaybeM f (PMap m) = PMap (mapMaybeM (mapMaybeM' f) m) - mapEitherM f (PMap m) = both PMap PMap (mapEitherM (mapEitherM' f)) m - isSubmapM (<=) (PMap m1) (PMap m2) = isSubmapM (isSubmapM (<=)) m1 m2 - unionM f (PMap m1) (PMap m2) = PMap (unionM (unionM' f) m1 m2) - isectM f (PMap m1) (PMap m2) = PMap (isectM (isectM' f) m1 m2) - diffM f (PMap m1) (PMap m2) = PMap (diffM (diffM' f) m1 m2) + lookupMC (k1, k2) (PMap m) = lookupMC k1 m >>= lookupMC k2 insertWithM f (k1, k2) a (PMap m) = PMap (insertWithM f' k1 (singletonM k2 a) m) where f' = insertWithM f k2 a - fromAscListM f xs = PMap (fromDistAscListM - [(a, fromAscListM f ys) | (a, Elem ys) <- breakFst xs]) - fromDistAscListM xs = PMap (fromDistAscListM - [(a, fromDistAscListM ys) | (a, Elem ys) <- breakFst xs]) - + singleHoleM (k1, k2) = PHole (singleHoleM k1) (singleHoleM k2) beforeM (PHole hole1 hole2) = PMap (beforeMM (gNull beforeM hole2) hole1) beforeWithM a (PHole hole1 hole2) = PMap (beforeWithM (beforeWithM a hole2) hole1) @@ -52,8 +61,9 @@ searchMC (k1, k2) (PMap m) f g = searchMC k1 m f' g' where f' hole1 = f (PHole hole1 (singleHoleM k2)) g' m' hole1 = mapSearch (PHole hole1) (searchMC k2 m') f g - indexM i (PMap m) = onThird (PHole hole1) (indexM i') m' - where !(# i', m', hole1 #) = indexM i m + indexM (PMap m) i = case indexM m i of + (# i', m', hole1 #) -> case indexM m' i' of + (# i'', a, hole2 #) -> (# i'', a, PHole hole1 hole2 #) extractHoleM (PMap m) = do (m', hole1) <- extractHoleM m (v, hole2) <- extractHoleM m' @@ -62,17 +72,33 @@ clearM (PHole hole1 hole2) = PMap (fillHoleM (clearM' hole2) hole1) assignM a (PHole hole1 hole2) = PMap (assignM (assignM a hole2) hole1) - unifierM (k1', k2') (k1, k2) a = case unifierM k1' k1 (singletonM k2 a) of - Just hole1 -> Just (PHole hole1 (singleHoleM k2')) - Nothing -> PHole (singleHoleM k1) <$> unifierM k2' k2 a + unifierM (k1', k2') (k1, k2) a = + (fmap (`PHole` singleHoleM k2') $ unifierM k1' k1 (singletonM k2 a)) + `mplus` (PHole (singleHoleM k1) <$> unifierM k2' k2 a) + unifyM (k11, k12) a1 (k21, k22) a2 = + let unify1 = unifyM k11 (singletonM k12 a1) k21 (singletonM k22 a2) + unify2 = singletonM k11 <$> unifyM k12 a1 k22 a2 + in PMap <$> (unify1 `mplus` unify2) gNull :: TrieKey k => (x -> TrieMap k a) -> x -> Maybe (TrieMap k a) -gNull = (guardNullM .) +gNull = (guardNull .) -breakFst :: Eq k1 => [((k1, k2), a)] -> [(k1, Elem [(k2, a)])] -breakFst [] = [] -breakFst (((a, b),v):xs) = breakFst' a (Seq.singleton (b, v)) xs where - breakFst' a vs (((a', b'), v'):xs) - | a == a' = breakFst' a' (vs |> (b', v')) xs - | otherwise = (a, Elem $ toList vs):breakFst' a' (Seq.singleton (b', v')) xs - breakFst' a vs [] = [(a, Elem $ toList vs)] +prodFold :: Eq k1 => FromList z1 k1 (TrieMap k2 a) -> FromList z2 k2 a -> + FromList (Stack k1 k2 z1 z2) (k1, k2) a +prodFold Foldl{snoc = snoc1, begin = begin1, zero = zero1, done = done1} + Foldl{snoc = snoc2, begin = begin2, done = done2} + = Foldl{zero = PMap zero1, ..} + where snoc (First k1 stk2) (k1', k2') a + | k1' == k1 = First k1 (snoc2 stk2 k2' a) + snoc (Stack k1 stk1 stk2) (k1', k2') a + | k1' == k1 = Stack k1 stk1 (snoc2 stk2 k2' a) + snoc stk (k1, k2) a = Stack k1 (collapse stk) (begin2 k2 a) + + collapse (First k1 stk2) = begin1 k1 (done2 stk2) + collapse (Stack k1 stk1 stk2) = snoc1 stk1 k1 (done2 stk2) + + begin (k1, k2) a = First k1 (begin2 k2 a) + + done = PMap . done1 . collapse + +data Stack k1 k2 z1 z2 a = First k1 (z2 a) | Stack k1 (z1 (TrieMap k2 a)) (z2 a) diff --git a/Data/TrieMap/RadixTrie.hs b/Data/TrieMap/RadixTrie.hs --- a/Data/TrieMap/RadixTrie.hs +++ b/Data/TrieMap/RadixTrie.hs @@ -1,17 +1,13 @@ -{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, MagicHash, FlexibleInstances #-} - +{-# LANGUAGE BangPatterns, FlexibleContexts, TypeFamilies, FlexibleInstances, MultiParamTypeClasses, TypeSynonymInstances, CPP #-} +{-# LANGUAGE UnboxedTuples #-} module Data.TrieMap.RadixTrie () where -import Data.TrieMap.TrieKey -import Data.TrieMap.Sized +import Control.Monad.Unpack -import Data.Functor -import Data.Foldable (Foldable(..)) -import Control.Monad +import Data.TrieMap.TrieKey import Data.Vector (Vector) -import qualified Data.Vector.Storable as S -import Data.Traversable +import qualified Data.Vector.Primitive as P import Data.Word import Data.TrieMap.RadixTrie.Edge @@ -19,11 +15,49 @@ import Prelude hiding (length, and, zip, zipWith, foldr, foldl) -instance TrieKey k => Foldable (TrieMap (Vector k)) where +#define VINSTANCE(cl) (TrieKey k, cl (TrieMap k)) => cl (TrieMap (Vector k)) + +instance VINSTANCE(Functor) where + fmap f (Radix m) = Radix (fmap f <$> m) + +instance VINSTANCE(Foldable) where foldMap f (Radix m) = foldMap (foldMap f) m foldr f z (Radix m) = foldl (foldr f) z m foldl f z (Radix m) = foldl (foldl f) z m +instance VINSTANCE(Traversable) where + traverse _ (Radix Nothing) = pure emptyM + traverse f (Radix (Just m)) = Radix . Just <$> traverse f m + +instance VINSTANCE(Subset) where + Radix m1 <=? Radix m2 = m1 <<=? m2 + +instance TrieKey k => Buildable (TrieMap (Vector k)) (Vector k) where + type UStack (TrieMap (Vector k)) = Edge Vector k + {-# INLINE uFold #-} + uFold f = Foldl{ + zero = emptyM, + begin = singletonEdge, + snoc = \ e ks a -> insertEdge (f a) ks a e, + done = Radix . Just} + type AStack (TrieMap (Vector k)) = Stack Vector k + {-# INLINE aFold #-} + aFold f = Radix <$> fromAscListEdge f + type DAStack (TrieMap (Vector k)) = Stack Vector k + {-# INLINE daFold #-} + daFold = aFold const + +#define SETOP(rad,op,opE) op f (rad m1) (rad m2) = rad (op (opE f) m1 m2) + +instance VINSTANCE(SetOp) where + SETOP(Radix,union,unionEdge) + SETOP(Radix,isect,isectEdge) + SETOP(Radix,diff,diffEdge) + +instance VINSTANCE(Project) where + mapMaybe f (Radix m) = Radix (mapMaybe (mapMaybeEdge f) m) + mapEither f (Radix m) = both' Radix Radix (mapEither (mapEitherEdge f)) m + -- | @'TrieMap' ('Vector' k) a@ is a traditional radix trie. instance TrieKey k => TrieKey (Vector k) where newtype TrieMap (Vector k) a = Radix (MEdge Vector k a) @@ -34,24 +68,17 @@ getSimpleM (Radix Nothing) = Null getSimpleM (Radix (Just e)) = getSimpleEdge e sizeM (Radix m) = getSize m - lookupM ks (Radix m) = liftMaybe m >>= lookupEdge ks - - fmapM f (Radix m) = Radix (mapEdge f <$> m) - mapMaybeM f (Radix m) = Radix (m >>= mapMaybeEdge f) - mapEitherM f (Radix e) = both Radix Radix (mapEitherMaybe (mapEitherEdge f)) e - traverseM f (Radix m) = Radix <$> traverse (traverseEdge f) m - - unionM f (Radix m1) (Radix m2) = Radix (unionMaybe (unionEdge f) m1 m2) - isectM f (Radix m1) (Radix m2) = Radix (isectMaybe (isectEdge f) m1 m2) - diffM f (Radix m1) (Radix m2) = Radix (diffMaybe (diffEdge f) m1 m2) - - isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubEdge (<=)) m1 m2 + lookupMC ks (Radix (Just e)) = lookupEdge ks e + lookupMC _ _ = mzero singleHoleM ks = Hole (singleLoc ks) {-# INLINE searchMC #-} - searchMC ks (Radix (Just e)) = mapSearch Hole (searchEdgeC ks e) - searchMC ks _ = \ f _ -> f (singleHoleM ks) - indexM i (Radix (Just e)) = onThird Hole (indexEdge i e) root + searchMC ks (Radix m) nomatch match = case m of + Just e -> searchEdgeC ks e nomatch' match' + Nothing -> nomatch' $~ singleLoc ks + where nomatch' = unpack (nomatch . Hole); match' a = unpack (match a . Hole) + indexM (Radix (Just e)) i = case indexEdge e i of + (# i', a, loc #) -> (# i', a, Hole loc #) indexM _ _ = indexFail () clearM (Hole loc) = Radix (clearEdge loc) @@ -67,50 +94,73 @@ afterWithM a (Hole loc) = Radix (afterEdge (Just a) loc) insertWithM f ks v (Radix e) = Radix (Just (maybe (singletonEdge ks v) (insertEdge f ks v) e)) - fromListM _ [] = emptyM - fromListM f ((k, a):xs) = Radix (Just (roll (singletonEdge k a) xs)) where - roll !e [] = e - roll !e ((ks, a):xs) = roll (insertEdge (f a) ks a e) xs - -type WordVec = S.Vector Word -instance Foldable (TrieMap (S.Vector Word)) where +type WordVec = P.Vector Word + +#define PINSTANCE(cl) cl (TrieMap (P.Vector Word)) + +instance PINSTANCE(Functor) where + fmap f (WRadix m) = WRadix (fmap f <$> m) + +instance PINSTANCE(Foldable) where foldMap f (WRadix m) = foldMap (foldMap f) m foldr f z (WRadix m) = foldl (foldr f) z m foldl f z (WRadix m) = foldl (foldl f) z m --- | @'TrieMap' ('S.Vector' Word) a@ is a traditional radix trie specialized for word arrays. -instance TrieKey (S.Vector Word) where - newtype TrieMap WordVec a = WRadix (MEdge S.Vector Word a) - newtype Hole WordVec a = WHole (EdgeLoc S.Vector Word a) +instance PINSTANCE(Traversable) where + traverse _ (WRadix Nothing) = pure emptyM + traverse f (WRadix (Just m)) = WRadix . Just <$> traverse f m + +instance PINSTANCE(Subset) where + WRadix m1 <=? WRadix m2 = m1 <<=? m2 + +instance PINSTANCE(SetOp) where + SETOP(WRadix,union,unionEdge) + SETOP(WRadix,isect,isectEdge) + SETOP(WRadix,diff,diffEdge) + +instance Buildable (TrieMap WordVec) WordVec where + type UStack (TrieMap WordVec) = Edge P.Vector Word + {-# INLINE uFold #-} + uFold f = Foldl{ + zero = emptyM, + begin = singletonEdge, + snoc = \ e ks a -> insertEdge (f a) ks a e, + done = WRadix . Just} + type AStack (TrieMap WordVec) = Stack P.Vector Word + {-# INLINE aFold #-} + aFold f = WRadix <$> fromAscListEdge f + type DAStack (TrieMap WordVec) = Stack P.Vector Word + {-# INLINE daFold #-} + daFold = aFold const + +instance PINSTANCE(Project) where + mapMaybe f (WRadix m) = WRadix (mapMaybe (mapMaybeEdge f) m) + mapEither f (WRadix m) = both' WRadix WRadix (mapEither (mapEitherEdge f)) m + +-- | @'TrieMap' ('P.Vector' Word) a@ is a traditional radix trie specialized for word arrays. +instance TrieKey (P.Vector Word) where + newtype TrieMap WordVec a = WRadix (MEdge P.Vector Word a) + newtype Hole WordVec a = WHole (EdgeLoc P.Vector Word a) emptyM = WRadix Nothing singletonM ks a = WRadix (Just (singletonEdge ks a)) getSimpleM (WRadix Nothing) = Null getSimpleM (WRadix (Just e)) = getSimpleEdge e sizeM (WRadix m) = getSize m - lookupM ks (WRadix m) = liftMaybe m >>= lookupEdge ks - - fmapM f (WRadix m) = WRadix (mapEdge f <$> m) - mapMaybeM f (WRadix m) = WRadix (m >>= mapMaybeEdge f) - mapEitherM f (WRadix e) = both WRadix WRadix (mapEitherMaybe (mapEitherEdge f)) e - traverseM f (WRadix m) = WRadix <$> traverse (traverseEdge f) m - - unionM f (WRadix m1) (WRadix m2) = WRadix (unionMaybe (unionEdge f) m1 m2) - isectM f (WRadix m1) (WRadix m2) = WRadix (isectMaybe (isectEdge f) m1 m2) - diffM f (WRadix m1) (WRadix m2) = WRadix (diffMaybe (diffEdge f) m1 m2) - - isSubmapM (<=) (WRadix m1) (WRadix m2) = subMaybe (isSubEdge (<=)) m1 m2 + lookupMC ks (WRadix (Just e)) = lookupEdge ks e + lookupMC _ _ = mzero singleHoleM ks = WHole (singleLoc ks) {-# INLINE searchMC #-} - searchMC ks (WRadix (Just e)) f g = searchEdgeC ks e f' g' where - f' loc = f (WHole loc) - g' a loc = g a (WHole loc) - searchMC ks _ f _ = f (singleHoleM ks) - indexM i (WRadix (Just e)) = onThird WHole (indexEdge i e) root - indexM _ (WRadix Nothing) = indexFail () - + searchMC ks (WRadix m) nomatch match = case m of + Just e -> searchEdgeC ks e nomatch' match' + Nothing -> nomatch' $~ singleLoc ks + where nomatch' = unpack (nomatch . WHole); match' a = unpack (match a . WHole) + indexM (WRadix (Just e)) i = case indexEdge e i of + (# i', a, loc #) -> (# i', a, WHole loc #) + indexM _ _ = indexFail () + clearM (WHole loc) = WRadix (clearEdge loc) {-# INLINE assignM #-} assignM a (WHole loc) = WRadix (Just (assignEdge a loc)) @@ -126,8 +176,3 @@ afterWithM a (WHole loc) = WRadix (afterEdge (Just a) loc) insertWithM f ks v (WRadix e) = WRadix (Just (maybe (singletonEdge ks v) (insertEdge f ks v) e)) - {-# INLINE fromListM #-} - fromListM _ [] = emptyM - fromListM f ((k, a):xs) = WRadix (Just (roll (singletonEdge k a) xs)) where - roll !e [] = e - roll !e ((ks, a):xs) = roll (insertEdge (f a) ks a e) xs diff --git a/Data/TrieMap/RadixTrie/Edge.hs b/Data/TrieMap/RadixTrie/Edge.hs --- a/Data/TrieMap/RadixTrie/Edge.hs +++ b/Data/TrieMap/RadixTrie/Edge.hs @@ -1,84 +1,127 @@ -{-# LANGUAGE MagicHash, BangPatterns, UnboxedTuples, PatternGuards, CPP, ViewPatterns #-} -{-# OPTIONS -funbox-strict-fields #-} -module Data.TrieMap.RadixTrie.Edge where +{-# LANGUAGE MagicHash, BangPatterns, UnboxedTuples, PatternGuards, CPP, ViewPatterns, NamedFieldPuns, ScopedTypeVariables #-} +{-# LANGUAGE RecordWildCards, TypeOperators, FlexibleContexts #-} +{-# OPTIONS -funbox-strict-fields -O -fspec-constr -fliberate-case -fstatic-argument-transformation #-} +module Data.TrieMap.RadixTrie.Edge ( searchEdgeC, + afterEdge, + assignEdge, + beforeEdge, + clearEdge, + diffEdge, + extractEdgeLoc, + indexEdge, + insertEdge, + isectEdge, + lookupEdge, + mapEitherEdge, + mapMaybeEdge, + unionEdge, + fromAscListEdge) where -import Data.TrieMap.Sized +import Control.Monad.Lookup +import Control.Monad.Ends +import Control.Monad.Unpack + import Data.TrieMap.TrieKey import Data.TrieMap.WordMap () import Data.TrieMap.RadixTrie.Label import Data.TrieMap.RadixTrie.Slice -import Control.Applicative -import Control.Monad - -import Data.Foldable -import Data.Monoid import Data.Word import Data.Vector.Generic (length) import qualified Data.Vector (Vector) -import qualified Data.Vector.Storable (Vector) -import Prelude hiding (length, foldr, foldl, zip, take) +import qualified Data.Vector.Primitive (Vector) +import Prelude hiding (length, foldr, foldl, zip, take, map) +import GHC.Exts + #define V(f) f (Data.Vector.Vector) (k) -#define U(f) f (Data.Vector.Storable.Vector) (Word) -#define EDGE(args) (eView -> Edge args) +#define U(f) f (Data.Vector.Primitive.Vector) (Word) +#define EDGE(args) (!(eView -> Edge args)) #define LOC(args) !(locView -> Loc args) +#define DEEP(args) !(pView -> Deep args) -{-# SPECIALIZE lookupEdge :: - TrieKey k => V() -> V(Edge) a -> Lookup a, - U() -> U(Edge) a -> Lookup a #-} -lookupEdge :: (Eq k, Label v k) => v k -> Edge v k a -> Lookup a -lookupEdge = lookupE where - lookupE !ks !EDGE(_ ls v ts) = if kLen < lLen then none else matchSlice matcher matches ks ls where - !kLen = length ks - !lLen = length ls - matcher k l z - | k == l = z - | otherwise = none - matches _ _ - | kLen == lLen = liftMaybe v - | (_, k, ks') <- splitSlice lLen ks - = lookupM k ts >>= lookupE ks' +instance Label v k => Functor (Edge v k) where + {-# SPECIALIZE instance TrieKey k => Functor (V(Edge)) #-} + {-# SPECIALIZE instance Functor (U(Edge)) #-} + fmap f = map where + map EDGE(sz ks v ts) = edge' sz ks (f <$> v) (map <$> ts) -{-# INLINE searchEdgeC #-} -searchEdgeC :: (Eq k, Label v k) => v k -> Edge v k a -> (EdgeLoc v k a -> r) -> (a -> EdgeLoc v k a -> r) -> r -searchEdgeC ks0 e f g = searchE ks0 e root where - searchE !ks !e@EDGE(_ !ls !v ts) path = matcher 0 where +instance Label v k => Foldable (Edge v k) where + {-# SPECIALIZE instance TrieKey k => Foldable (V(Edge)) #-} + {-# SPECIALIZE instance Foldable (U(Edge)) #-} + foldMap f = fold where + foldBranch = foldMap fold + fold e = case eView e of + Edge _ _ Nothing ts -> foldBranch ts + Edge _ _ (Just a) ts -> f a `mappend` foldBranch ts + + foldr f = flip fold where + foldBranch = foldr fold + fold e z = case eView e of + Edge _ _ Nothing ts -> foldBranch z ts + Edge _ _ (Just a) ts -> a `f` foldBranch z ts + + foldl f = fold where + foldBranch = foldl fold + fold z e = case eView e of + Edge _ _ Nothing ts -> foldBranch z ts + Edge _ _ (Just a) ts -> foldBranch (z `f` a) ts + +instance Label v k => Traversable (Edge v k) where + {-# SPECIALIZE instance TrieKey k => Traversable (V(Edge)) #-} + {-# SPECIALIZE instance Traversable (U(Edge)) #-} + traverse f = trav where + travBranch = traverse trav + trav e = case eView e of + Edge sz ks Nothing ts -> edge' sz ks Nothing <$> travBranch ts + Edge sz ks (Just a) ts -> edge' sz ks . Just <$> f a <*> travBranch ts + +{-# SPECIALIZE lookupEdge :: + TrieKey k => V() -> V(Edge) a -> Lookup r a, + U() -> U(Edge) a -> Lookup r a #-} +lookupEdge :: (Eq k, Label v k) => v k -> Edge v k a -> Lookup r a +lookupEdge ks e = Lookup $ \ no yes -> let + lookupE !ks !EDGE(_ ls !v ts) = if kLen < lLen then no else matchSlice matcher matches ks ls where !kLen = length ks !lLen = length ls - !len = min kLen lLen - {-# INLINE kk #-} - kk = ks !$ lLen - matcher !i - | i < len = let k = ks !$ i; l = ls !$ i in case unifierM k l (dropEdge (i+1) e) of - Nothing -> matcher (i+1) - Just tHole -> f (loc (dropSlice (i+1) ks) emptyM (deep path (takeSlice i ls) Nothing tHole)) - matcher _ - | kLen < lLen - = let lPre = takeSlice kLen ls; l = ls !$ kLen; e' = dropEdge (kLen + 1) e in - f (loc lPre (singletonM l e') path) - | kLen == lLen - = maybe f g v (loc ls ts path) - | otherwise = let - ks' = dropSlice (lLen + 1) ks - f' tHole = f (loc ks' emptyM (deep path ls v tHole)) - g' e' tHole = searchE ks' e' (deep path ls v tHole) - in searchMC kk ts f' g' + matcher k l z + | k == l = z + | otherwise = no + matches _ _ + | kLen == lLen = maybe no yes v + | (_, k, ks') <- splitSlice lLen ks + = runLookup (lookupMC k ts) no (lookupE ks') + in lookupE ks e -{-# SPECIALIZE mapEdge :: - (TrieKey k, Sized b) => (a -> b) -> V(Edge) a -> V(Edge) b, - Sized b => (a -> b) -> U(Edge) a -> U(Edge) b #-} -mapEdge :: (Label v k, Sized b) => (a -> b) -> Edge v k a -> Edge v k b -mapEdge f = mapE where - mapE !EDGE(_ ks v ts) = edge ks (f <$> v) (fmapM mapE ts) +{-# SPECIALIZE INLINE searchEdgeC :: + TrieKey k => V() -> V(Edge) a -> (V(EdgeLoc) a :~> r) -> (a -> V(EdgeLoc) a :~> r) -> r, + U() -> U(Edge) a -> (U(EdgeLoc) a :~> r) -> (a -> U(EdgeLoc) a :~> r) -> r #-} +searchEdgeC :: (Eq k, Label v k, Unpackable (EdgeLoc v k a)) => + v k -> Edge v k a -> (EdgeLoc v k a :~> r) -> (a -> EdgeLoc v k a :~> r) -> r +searchEdgeC ks0 e nomatch match = searchE ks0 e root where + searchE !ks e@EDGE(_ !ls !v ts) path = iMatchSlice matcher matches ks ls where + matcher i k l z = + runLookup (unifierM k l (dropEdge (i+1) e)) z + (\ tHole -> nomatch $~ loc (dropSlice (i+1) ks) emptyM (deep path (takeSlice i ls) Nothing tHole)) + matches kLen lLen = case compare kLen lLen of + LT -> let lPre = takeSlice kLen ls; l = ls !$ kLen; e' = dropEdge (kLen + 1) e in + nomatch $~ loc lPre (singletonM l e') path + EQ -> maybe nomatch match v $~ loc ls ts path + GT -> let + {-# INLINE kk #-} + kk = ks !$ lLen + ks' = dropSlice (lLen + 1) ks + nomatch' tHole = nomatch $~ loc ks' emptyM (deep path ls v tHole) + match' e' tHole = searchE ks' e' (deep path ls v tHole) + in searchMC kk ts nomatch' match' {-# SPECIALIZE mapMaybeEdge :: (TrieKey k, Sized b) => (a -> Maybe b) -> V(Edge) a -> V(MEdge) b, Sized b => (a -> Maybe b) -> U(Edge) a -> U(MEdge) b #-} mapMaybeEdge :: (Label v k, Sized b) => (a -> Maybe b) -> Edge v k a -> MEdge v k b mapMaybeEdge f = mapMaybeE where - mapMaybeE EDGE(_ ks v ts) = cEdge ks (v >>= f) (mapMaybeM mapMaybeE ts) + mapMaybeE !EDGE(_ ks !v ts) = let !v' = v >>= f in cEdge ks v' (mapMaybe mapMaybeE ts) {-# SPECIALIZE mapEitherEdge :: (TrieKey k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> V(Edge) a -> (# V(MEdge) b, V(MEdge) c #), @@ -86,35 +129,9 @@ mapEitherEdge :: (Label v k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> Edge v k a -> (# MEdge v k b, MEdge v k c #) mapEitherEdge f = mapEitherE where - mapEitherE !EDGE(_ ks v ts) = (# cEdge ks vL tsL, cEdge ks vR tsR #) - where !(# vL, vR #) = mapEitherMaybe f v - !(# tsL, tsR #) = mapEitherM mapEitherE ts - -{-# SPECIALIZE traverseEdge :: - (TrieKey k, Applicative f, Sized b) => (a -> f b) -> V(Edge) a -> f (V(Edge) b), - (Applicative f, Sized b) => (a -> f b) -> U(Edge) a -> f (U(Edge) b) #-} -traverseEdge :: (Label v k, Applicative f, Sized b) => - (a -> f b) -> Edge v k a -> f (Edge v k b) -traverseEdge f = traverseE where - traverseE e = case eView e of - Edge _ ks Nothing ts -> edge ks Nothing <$> traverseM traverseE ts - Edge _ ks (Just v) ts -> edge ks . Just <$> f v <*> traverseM traverseE ts - -instance Label v k => Foldable (EView v k) where - {-# INLINE foldr #-} - {-# INLINE foldl #-} - {-# INLINE foldMap #-} - foldMap f (Edge _ _ Nothing ts) = foldMap (foldMap f) ts - foldMap f (Edge _ _ (Just v) ts) = f v `mappend` foldMap (foldMap f) ts - foldr f z (Edge _ _ v ts) = foldr f (foldr (flip $ foldr f) z ts) v - foldl f z (Edge _ _ v ts) = foldl (foldl f) (foldl f z v) ts - -instance Label v k => Foldable (Edge v k) where - {-# SPECIALIZE instance TrieKey k => Foldable (V(Edge)) #-} - {-# SPECIALIZE instance Foldable (U(Edge)) #-} - foldMap f e = foldMap f (eView e) - foldr f z e = foldr f z (eView e) - foldl f z e = foldl f z (eView e) + mapEitherE EDGE(_ ks v ts) = (# cEdge ks vL tsL, cEdge ks vR tsR #) + where !(# vL, vR #) = mapEither f v + !(# tsL, tsR #) = mapEither mapEitherE ts {-# INLINE assignEdge #-} assignEdge :: (Label v k, Sized a) => a -> EdgeLoc v k a -> Edge v k a @@ -124,20 +141,17 @@ (TrieKey k, Sized a) => V(Edge) a -> V(Path) a -> V(Edge) a, Sized a => U(Edge) a -> U(Path) a -> U(Edge) a #-} assign :: (Label v k, Sized a) => Edge v k a -> Path v k a -> Edge v k a -assign !e path = case pView path of - Root -> e - Deep path ks v tHole - -> assign (edge ks v (assignM e tHole)) path +assign e DEEP(path ks v tHole) = assign (edge ks v (assignM e tHole)) path +assign e _ = e {-# SPECIALIZE clearEdge :: (TrieKey k, Sized a) => V(EdgeLoc) a -> V(MEdge) a, Sized a => U(EdgeLoc) a -> U(MEdge) a #-} clearEdge :: (Label v k, Sized a) => EdgeLoc v k a -> MEdge v k a clearEdge LOC(ks ts path) = rebuild (cEdge ks Nothing ts) path where - rebuild !e path = case pView path of - Root -> e - Deep path ks v tHole - -> rebuild (cEdge ks v (fillHoleM e tHole)) path + rebuild Nothing DEEP(path ks v tHole) = rebuild (cEdge ks v (clearM tHole)) path + rebuild Nothing _ = Nothing + rebuild (Just e) path = Just $ assign e path {-# SPECIALIZE unionEdge :: (TrieKey k, Sized a) => (a -> a -> Maybe a) -> V(Edge) a -> V(Edge) a -> V(MEdge) a, @@ -145,14 +159,13 @@ unionEdge :: (Label v k, Sized a) => (a -> a -> Maybe a) -> Edge v k a -> Edge v k a -> MEdge v k a unionEdge f = unionE where - unionE !eK@EDGE(_ ks0 vK tsK) !eL@EDGE(_ ls0 vL tsL) = iMatchSlice matcher matches ks0 ls0 where - matcher i k l z = case unifyM k eK' l eL' of - Nothing -> z - Just ts -> Just (edge (takeSlice i ks0) Nothing ts) + unionE !eK@EDGE(_ ks0 !vK tsK) !eL@EDGE(_ ls0 !vL tsL) = iMatchSlice matcher matches ks0 ls0 where + matcher !i k l z = runLookup (unifyM k eK' l eL') z $ Just . edge (takeSlice i ks0) Nothing where eK' = dropEdge (i+1) eK eL' = dropEdge (i+1) eL + matches kLen lLen = case compare kLen lLen of - EQ -> cEdge ks0 (unionMaybe f vK vL) $ unionM unionE tsK tsL + EQ -> cEdge ks0 (union f vK vL) $ union unionE tsK tsL LT -> searchMC l tsK nomatch match where eL' = dropEdge (kLen + 1) eL; l = ls0 !$ kLen nomatch holeKT = cEdge ks0 vK $ assignM eL' holeKT @@ -171,15 +184,11 @@ isectE !eK@EDGE(_ ks0 vK tsK) !eL@EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where matcher k l z = guard (k == l) >> z matches kLen lLen = case compare kLen lLen of - EQ -> compact $ edge ks0 (isectMaybe f vK vL) $ isectM isectE tsK tsL - LT -> let l = ls0 !$ kLen in do - eK' <- toMaybe $ lookupM l tsK - let eL' = dropEdge (kLen + 1) eL - unDropEdge (kLen + 1) <$> eK' `isectE` eL' - GT -> let k = ks0 !$ lLen in do - eL' <- toMaybe $ lookupM k tsL - let eK' = dropEdge (lLen + 1) eK - unDropEdge (lLen + 1) <$> eK' `isectE` eL' + EQ -> cEdge ks0 (isect f vK vL) $ isect isectE tsK tsL + LT -> let l = ls0 !$ kLen in runLookup (lookupMC l tsK) Nothing $ \ eK' -> + let eL' = dropEdge (kLen + 1) eL in unDropEdge (kLen + 1) <$> eK' `isectE` eL' + GT -> let k = ks0 !$ lLen in runLookup (lookupMC k tsL) Nothing $ \ eL' -> + let eK' = dropEdge (lLen + 1) eK in unDropEdge (lLen + 1) <$> eK' `isectE` eL' {-# SPECIALIZE diffEdge :: (TrieKey k, Sized a) => (a -> b -> Maybe a) -> V(Edge) a -> V(Edge) b -> V(MEdge) a, @@ -187,79 +196,149 @@ diffEdge :: (Eq k, Label v k, Sized a) => (a -> b -> Maybe a) -> Edge v k a -> Edge v k b -> MEdge v k a diffEdge f = diffE where - diffE !eK@EDGE(_ ks0 vK tsK) !eL@EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where + diffE !eK@EDGE(_ ks0 !vK tsK) !eL@EDGE(_ ls0 !vL tsL) = matchSlice matcher matches ks0 ls0 where matcher k l z | k == l = z | otherwise = Just eK matches kLen lLen = case compare kLen lLen of - EQ -> cEdge ks0 (diffMaybe f vK vL) $ diffM diffE tsK tsL + EQ -> cEdge ks0 (diff f vK vL) $ diff diffE tsK tsL LT -> searchMC l tsK nomatch match where l = ls0 !$ kLen; eL' = dropEdge (kLen + 1) eL nomatch _ = Just eK match eK' holeKT = cEdge ks0 vK $ fillHoleM (eK' `diffE` eL') holeKT GT -> let k = ks0 !$ lLen; eK' = dropEdge (lLen + 1) eK in - option (lookupM k tsL) (Just eK) (\ eL' -> fmap (unDropEdge (lLen + 1)) (eK' `diffE` eL')) + runLookup (lookupMC k tsL) (Just eK) (\ eL' -> fmap (unDropEdge (lLen + 1)) (eK' `diffE` eL')) -{-# SPECIALIZE isSubEdge :: - TrieKey k => LEq a b -> LEq (V(Edge) a) (V(Edge) b), - LEq a b -> LEq (U(Edge) a) (U(Edge) b) #-} -isSubEdge :: (Eq k, Label v k) => LEq a b -> LEq (Edge v k a) (Edge v k b) -isSubEdge (<=) = isSubE where - isSubE !eK@EDGE(_ ks0 vK tsK) !EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where +instance (Eq k, Label v k) => Subset (Edge v k) where + {-# SPECIALIZE instance (Eq k, TrieKey k) => Subset (V(Edge)) #-} + {-# SPECIALIZE instance Subset (U(Edge)) #-} + eK@EDGE(_ ks0 vK tsK) <=? EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where matcher k l z = k == l && z matches kLen lLen = case compare kLen lLen of LT -> False - EQ -> subMaybe (<=) vK vL && isSubmapM isSubE tsK tsL - GT -> let k = ks0 !$ lLen in option (lookupM k tsL) False (isSubE (dropEdge (lLen + 1) eK)) + EQ -> vK <=? vL && tsK <<=? tsL + GT -> let k = ks0 !$ lLen in runLookup (lookupMC k tsL) False (dropEdge (lLen + 1) eK <=?) {-# SPECIALIZE beforeEdge :: (TrieKey k, Sized a) => Maybe a -> V(EdgeLoc) a -> V(MEdge) a, Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} -beforeEdge :: (Label v k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a -beforeEdge v LOC(ks ts path) = buildBefore (cEdge ks v ts) path where - buildBefore !e path = case pView path of - Root -> e - Deep path ks v tHole -> buildBefore (cEdge ks v $ beforeMM e tHole) path - {-# SPECIALIZE afterEdge :: (TrieKey k, Sized a) => Maybe a -> V(EdgeLoc) a -> V(MEdge) a, Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} -afterEdge :: (Label v k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a -afterEdge v LOC(ks ts path) = buildAfter (cEdge ks v ts) path where - buildAfter !e path = case pView path of - Root -> e - Deep path ks v tHole - -> buildAfter (cEdge ks v $ afterMM e tHole) path +beforeEdge, afterEdge :: (Label v k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a +beforeEdge v LOC(ks ts path) = case cEdge ks v ts of + Nothing -> before path + Just e -> Just $ beforeWith e path + where before DEEP(path ks v tHole) = case cEdge ks v (beforeM tHole) of + Nothing -> before path + Just e -> Just $ beforeWith e path + before _ = Nothing + beforeWith e DEEP(path ks v tHole) + = beforeWith (edge ks v (beforeWithM e tHole)) path + beforeWith e _ = e +afterEdge v LOC(ks ts path) = case cEdge ks v ts of + Nothing -> after path + Just e -> Just $ afterWith e path + where after DEEP(path ks _ tHole) = case cEdge ks Nothing (afterM tHole) of + Nothing -> after path + Just e -> Just $ afterWith e path + after _ = Nothing + afterWith e DEEP(path ks _ tHole) + = afterWith (edge ks Nothing (afterWithM e tHole)) path + afterWith e _ = e + {-# SPECIALIZE extractEdgeLoc :: - (TrieKey k, Functor m, MonadPlus m) => V(Edge) a -> V(Path) a -> m (a, V(EdgeLoc) a), - (Functor m, MonadPlus m) => U(Edge) a -> U(Path) a -> m (a, U(EdgeLoc) a) #-} + TrieKey k => V(Edge) a -> V(Path) a -> First (a, V(EdgeLoc) a), + TrieKey k => V(Edge) a -> V(Path) a -> Last (a, V(EdgeLoc) a), + U(Edge) a -> U(Path) a -> First (a, U(EdgeLoc) a), + U(Edge) a -> U(Path) a -> Last (a, U(EdgeLoc) a)#-} extractEdgeLoc :: (Label v k, Functor m, MonadPlus m) => Edge v k a -> Path v k a -> m (a, EdgeLoc v k a) -extractEdgeLoc !EDGE(_ ks v ts) path = case v of +extractEdgeLoc EDGE(_ ks v ts) path = case v of Nothing -> extractTS Just a -> return (a, loc ks ts path) `mplus` extractTS - where extractTS = do (e', tHole) <- extractHoleM ts - extractEdgeLoc e' (deep path ks v tHole) + where extractTS = do (e', tHole) <- extractHoleM ts + extractEdgeLoc e' (deep path ks v tHole) -{-# SPECIALIZE INLINE indexEdge :: - (TrieKey k, Sized a) => Int -> V(Edge) a -> V(Path) a -> (# Int, a, V(EdgeLoc) a #), - Sized a => Int -> U(Edge) a -> U(Path) a -> (# Int, a, U(EdgeLoc) a #) #-} -indexEdge :: (Label v k, Sized a) => Int -> Edge v k a -> Path v k a -> (# Int, a, EdgeLoc v k a #) -indexEdge = indexE where - indexE !i e path = case eView e of - Edge _ ks v@(Just a) ts - | i < sv -> (# i, a, loc ks ts path #) - | (# i', e', tHole #) <- indexM (i - sv) ts - -> indexE i' e' (deep path ks v tHole) - where !sv = getSize a - Edge _ ks Nothing ts - -> indexE i' e' (deep path ks Nothing tHole) - where !(# i', e', tHole #) = indexM i ts +{-# SPECIALIZE indexEdge :: + (TrieKey k, Sized a) => V(Edge) a -> Int# -> (# Int#, a, V(EdgeLoc) a #), + Sized a => U(Edge) a -> Int# -> (# Int#, a, U(EdgeLoc) a #) #-} +indexEdge :: (Label v k, Sized a) => Edge v k a -> Int# -> (# Int#, a, EdgeLoc v k a #) +indexEdge e i = let + indexE i !e path = case eView e of + Edge sE ks v@(Just a) ts + | i <# sv -> (# i, a, loc ks ts path #) + | otherwise -> case indexM ts (i -# sv) of + (# i', e', tHole #) -> indexE i' e' (deep path ks v tHole) + where !sv = unbox $ sE - sizeM ts + Edge _ ks Nothing ts -> case indexM ts i of + (# i', e', tHole #) -> indexE i' e' (deep path ks Nothing tHole) + in indexE i e root {-# SPECIALIZE insertEdge :: (TrieKey k, Sized a) => (a -> a) -> V() -> a -> V(Edge) a -> V(Edge) a, Sized a => (a -> a) -> U() -> a -> U(Edge) a -> U(Edge) a #-} insertEdge :: (Label v k, Sized a) => (a -> a) -> v k -> a -> Edge v k a -> Edge v k a -insertEdge f ks v e = searchEdgeC ks e nomatch match where - nomatch = assignEdge v - match = assignEdge . f +insertEdge f ks0 a e = insertE ks0 e where + !sza = getSize a + insertE !ks eL@EDGE(szL ls !v ts) = iMatchSlice matcher matches ks ls where + !szV = szL - sizeM ts + matcher !i k l z = runLookup (unifyM k eK' l eL') z (edge (takeSlice i ls) Nothing) + where eK' = edge' sza (dropSlice (i+1) ks) (Just a) emptyM + eL' = dropEdge (i+1) eL + matches kLen lLen = case compare kLen lLen of + LT -> (edge' (sza + szL) ks (Just a) (singletonM l eL')) + where l = ls !$ kLen; eL' = dropEdge (kLen+1) eL + EQ -> (edge ls (Just (maybe a f v)) ts) + GT -> edge' sz' ls v ts' where + ks' = dropSlice (lLen + 1) ks + k = ks !$ lLen + ts' = insertWithM (insertE ks') k (edge' sza ks' (Just a) emptyM) ts + sz' = sizeM ts' + szV + +{-# SPECIALIZE fromAscListEdge :: + (TrieKey k, Sized a) => (a -> a -> a) -> Foldl (V(Stack)) (V()) a (V(MEdge) a), + Sized a => (a -> a -> a) -> Foldl (U(Stack)) (U()) a (U(MEdge) a) #-} +fromAscListEdge :: forall v k a .(Label v k, Sized a) => (a -> a -> a) -> + Foldl (Stack v k) (v k) a (MEdge v k a) +fromAscListEdge f = case inline daFold of + Foldl{snoc = snocB, begin = beginB, done = doneB} + -> Foldl{..} where + begin ks a = stack ks (Just a) Nothing Nothing + zero = Nothing + + snoc stk ks vK = snoc' ks stk where + snoc' !ks !stk = case sView stk of + Stack ls !vL !brL !lStack -> iMatchSlice matcher matches ks ls where + matcher i k l z + | k == l = z + | otherwise = let + ksPre = takeSlice i ks + ksSuf = dropSlice (i+1) ks + ls' = dropSlice (i+1) ls + eL = roll (stack ls' vL brL lStack) + in stack ksPre Nothing (Just (beginB l eL)) (Just (k, begin ksSuf vK)) + matches kLen lLen + | kLen > lLen = let + ksPre = takeSlice lLen ks + k = ks !$ lLen + ksSuf = dropSlice (lLen + 1) ks + in case lStack of + Just (lChar, lStack) + | k == lChar -> stack ksPre vL brL (Just (lChar, snoc' ksSuf lStack)) + | otherwise -> stack ksPre vL (Just $ snocBranch brL lChar lStack) + (Just (k, begin ksSuf vK)) + Nothing -> stack ksPre vL brL (Just (k, begin ksSuf vK)) + | otherwise = stack ks (Just (maybe vK (f vK) vL)) brL lStack + + + snocBranch Nothing k stack = beginB k (roll stack) + snocBranch (Just s) k stack = snocB s k (roll stack) + + roll stack = case sView stack of + Stack ks (Just vK) _ Nothing -> singletonEdge ks vK + Stack ks vK brK (Just (kChar, stack')) -> + edge ks vK $ inline doneB $ snocBranch brK kChar stack' + _ -> error "Error: bad stack" + + done = Just . roll diff --git a/Data/TrieMap/RadixTrie/Label.hs b/Data/TrieMap/RadixTrie/Label.hs --- a/Data/TrieMap/RadixTrie/Label.hs +++ b/Data/TrieMap/RadixTrie/Label.hs @@ -2,34 +2,39 @@ {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.RadixTrie.Label where +import Control.Monad.Unpack +import Control.Monad.Trans.Reader + import Data.TrieMap.TrieKey -import Data.TrieMap.Sized import Data.TrieMap.RadixTrie.Slice import Data.TrieMap.WordMap import Data.Word import Data.Vector.Generic import qualified Data.Vector as V -import qualified Data.Vector.Storable as S +import qualified Data.Vector.Primitive as P import Prelude hiding (length) #define V(ty) (ty (V.Vector) (k)) -#define U(ty) (ty (S.Vector) Word) +#define U(ty) (ty (P.Vector) Word) -class (Vector v k, TrieKey k) => Label v k where +class (Unpackable (v k), Vector v k, TrieKey k) => Label v k where data Edge v k :: * -> * data Path v k :: * -> * data EdgeLoc v k :: * -> * + data Stack v k :: * -> * edge :: Sized a => v k -> Maybe a -> Branch v k a -> Edge v k a edge' :: Int -> v k -> Maybe a -> Branch v k a -> Edge v k a root :: Path v k a deep :: Path v k a -> v k -> Maybe a -> BHole v k a -> Path v k a loc :: v k -> Branch v k a -> Path v k a -> EdgeLoc v k a + stack :: v k -> Maybe a -> Maybe (DAMStack k (Edge v k a)) -> Maybe (k, Stack v k a) -> Stack v k a eView :: Edge v k a -> EView v k a pView :: Path v k a -> PView v k a locView :: EdgeLoc v k a -> LocView v k a + sView :: Stack v k a -> StackView v k a type BHole v k a = Hole k (Edge v k a) @@ -39,25 +44,33 @@ data LocView v k a = Loc !( v k) (Branch v k a) (Path v k a) data PView v k a = Root | Deep (Path v k a) (v k) (Maybe a) (BHole v k a) +data StackView v k a = Stack (v k) (Maybe a) (Maybe (DAMStack k (Edge v k a))) (Maybe (k, Stack v k a)) type MEdge v k a = Maybe (Edge v k a) instance Sized (EView v k a) where getSize# (Edge sz _ _ _) = unbox sz instance Label v k => Sized (Edge v k a) where - {-# SPECIALIZE instance TrieKey k => Sized (Edge V.Vector k a) #-} + {-# SPECIALIZE instance TrieKey k => Sized (V(Edge) a) #-} + {-# SPECIALIZE instance Sized (U(Edge) a) #-} getSize# e = getSize# (eView e) instance TrieKey k => Label V.Vector k where data Edge V.Vector k a = - VEdge Int !(V()) (V(Branch) a) - | VEdgeX Int !(V()) a (V(Branch) a) + VEdge !Int !(V()) (V(Branch) a) + | VEdgeX !Int !(V()) a (V(Branch) a) data Path V.Vector k a = VRoot | VDeep (V(Path) a) !(V()) (V(BHole) a) | VDeepX (V(Path) a) !(V()) a (V(BHole) a) data EdgeLoc V.Vector k a = VLoc !(V()) (V(Branch) a) (V(Path) a) + data Stack V.Vector k a = + VStackAZ !(V()) a (DAMStack k (V(Edge) a)) k (V(Stack) a) + | VStackA !(V()) a k (V(Stack) a) + | VStackZ !(V()) (DAMStack k (V(Edge) a)) k (V(Stack) a) + | VTip !(V()) a + {-# INLINE edge #-} edge !ks Nothing ts = VEdge (sizeM ts) ks ts edge !ks (Just a) ts = VEdgeX (sizeM ts + getSize a) ks a ts edge' s !ks Nothing ts = VEdge s ks ts @@ -75,21 +88,61 @@ pView (VDeep path ks tHole) = Deep path ks Nothing tHole pView (VDeepX path ks v tHole) = Deep path ks (Just v) tHole locView (VLoc ks ts path) = Loc ks ts path + + {-# INLINE stack #-} + stack !ks (Just a) (Just z) (Just (k, stack)) = + VStackAZ ks a z k stack + stack !ks (Just a) Nothing (Just (k, stack)) = + VStackA ks a k stack + stack !ks Nothing (Just z) (Just (k, stack)) = + VStackZ ks z k stack + stack !ks (Just a) Nothing Nothing = VTip ks a + stack _ _ _ _ = error "Error: bad stack" + {-# INLINE sView #-} + sView (VTip ks v) = Stack ks (Just v) Nothing Nothing + sView (VStackAZ ks a z k stack) = Stack ks (Just a) (Just z) (Just (k, stack)) + sView (VStackA ks a k stack) = Stack ks (Just a) Nothing (Just (k, stack)) + sView (VStackZ ks z k stack) = Stack ks Nothing (Just z) (Just (k, stack)) -instance Label S.Vector Word where - data Edge S.Vector Word a = - SEdge !Int !(U()) !(SNode (U(Edge) a)) +instance TrieKey k => Unpackable (V(EdgeLoc) a) where + newtype UnpackedReaderT (EdgeLoc V.Vector k a) m r = + VLocRT {runVLocRT :: UnpackedReaderT (V.Vector k) (ReaderT (V(Branch) a) (ReaderT (V(Path) a) m)) r} + runUnpackedReaderT func (VLoc ks ts path) = + runVLocRT func `runUnpackedReaderT` ks `runReaderT` ts `runReaderT` path + unpackedReaderT func = VLocRT $ unpackedReaderT $ \ ks -> ReaderT $ \ ts -> ReaderT $ \ path -> func (VLoc ks ts path) + +instance Label P.Vector Word where + data Edge P.Vector Word a = + SEdge !(U()) !(SNode (U(Edge) a)) | SEdgeX !Int !(U()) a !(SNode (U(Edge) a)) - data Path S.Vector Word a = + data Path P.Vector Word a = SRoot | SDeep (U(Path) a) !(U()) !(WHole (U(Edge) a)) | SDeepX (U(Path) a) !(U()) a !(WHole (U(Edge) a)) - data EdgeLoc S.Vector Word a = + data EdgeLoc P.Vector Word a = SLoc !(U()) !(SNode (U(Edge) a)) (U(Path) a) + data Stack P.Vector Word a = + PStackAZ !(U()) a !(WordStack (U(Edge) a)) !Word (U(Stack) a) + | PStackA !(U()) a !Word (U(Stack) a) + | PStackZ !(U()) !(WordStack (U(Edge) a)) !Word (U(Stack) a) + | PTip !(U()) a - edge !ks Nothing ts = SEdge (sizeM ts) ks (getWordMap ts) + {-# INLINE stack #-} + stack !ks a z stack = case (a, z, stack) of + (Just a, Just z, Just (k, stack)) -> PStackAZ ks a z k stack + (Just a, Nothing, Just (k, stack)) -> PStackA ks a k stack + (Nothing, Just z, Just (k, stack)) -> PStackZ ks z k stack + (Just a, Nothing, Nothing) -> PTip ks a + _ -> error "Error: bad stack" + {-# INLINE sView #-} + sView (PStackAZ ks a z k stack) = Stack ks (Just a) (Just z) (Just (k, stack)) + sView (PStackA ks a k stack) = Stack ks (Just a) Nothing (Just (k, stack)) + sView (PStackZ ks z k stack) = Stack ks Nothing (Just z) (Just (k, stack)) + sView (PTip ks a) = Stack ks (Just a) Nothing Nothing + + edge !ks Nothing ts = SEdge ks (getWordMap ts) edge !ks (Just v) ts = SEdgeX (getSize v + sizeM ts) ks v (getWordMap ts) - edge' sz !ks Nothing ts = SEdge sz ks (getWordMap ts) + edge' _ !ks Nothing ts = SEdge ks (getWordMap ts) edge' sz !ks (Just v) ts = SEdgeX sz ks v (getWordMap ts) root = SRoot @@ -98,18 +151,26 @@ loc ks ts path = SLoc ks (getWordMap ts) path - eView (SEdge s ks ts) = Edge s ks Nothing (WordMap ts) + eView (SEdge ks ts) = Edge (getSize ts) ks Nothing (WordMap ts) eView (SEdgeX s ks v ts) = Edge s ks (Just v) (WordMap ts) pView SRoot = Root pView (SDeep path ks tHole) = Deep path ks Nothing (Hole tHole) pView (SDeepX path ks v tHole) = Deep path ks (Just v) (Hole tHole) locView (SLoc ks ts path) = Loc ks (WordMap ts) path +instance Unpackable (U(EdgeLoc) a) where + newtype UnpackedReaderT (U(EdgeLoc) a) m r = + ULocRT {runULocRT :: UnpackedReaderT (U()) (UnpackedReaderT (SNode (U(Edge) a)) (ReaderT (U(Path) a) m)) r} + runUnpackedReaderT func (SLoc ks ts path) = + runULocRT func `runUnpackedReaderT` ks `runUnpackedReaderT` ts `runReaderT` path + unpackedReaderT func = ULocRT $ unpackedReaderT $ \ ks -> unpackedReaderT $ \ ts -> ReaderT $ \ path -> + func (SLoc ks ts path) + {-# SPECIALIZE singletonEdge :: (TrieKey k, Sized a) => V() -> a -> V(Edge) a, Sized a => U() -> a -> U(Edge) a #-} singletonEdge :: (Label v k, Sized a) => v k -> a -> Edge v k a -singletonEdge ks a = edge ks (Just a) emptyM +singletonEdge !ks a = edge' (getSize a) ks (Just a) emptyM {-# SPECIALIZE singleLoc :: TrieKey k => V() -> V(EdgeLoc) a, @@ -122,31 +183,50 @@ U(Edge) a -> Simple a #-} getSimpleEdge :: Label v k => Edge v k a -> Simple a getSimpleEdge !(eView -> Edge _ _ v ts) - | nullM ts = maybe Null Singleton v + | isNull ts = maybe Null Singleton v | otherwise = NonSimple -{-# SPECIALIZE INLINE dropEdge :: +{-# SPECIALIZE dropEdge :: TrieKey k => Int -> V(Edge) a -> V(Edge) a, Int -> U(Edge) a -> U(Edge) a #-} -{-# SPECIALIZE INLINE unDropEdge :: +{-# SPECIALIZE unDropEdge :: TrieKey k => Int -> V(Edge) a -> V(Edge) a, Int -> U(Edge) a -> U(Edge) a #-} dropEdge, unDropEdge :: Label v k => Int -> Edge v k a -> Edge v k a dropEdge !n !(eView -> Edge sz# ks v ts) = edge' sz# (dropSlice n ks) v ts unDropEdge !n !(eView -> Edge sz# ks v ts) = edge' sz# (unDropSlice n ks) v ts -{-# SPECIALIZE compact :: - TrieKey k => V(Edge) a -> V(MEdge) a, - U(Edge) a -> U(MEdge) a #-} -compact :: Label v k => Edge v k a -> MEdge v k a -compact !e@(eView -> Edge _ ks Nothing ts) = case getSimpleM ts of - Null -> Nothing - Singleton e' -> Just (unDropEdge (length ks + 1) e') - NonSimple -> Just e -compact e = Just e - {-# SPECIALIZE cEdge :: (TrieKey k, Sized a) => V() -> Maybe a -> V(Branch) a -> V(MEdge) a, Sized a => U() -> Maybe a -> U(Branch) a -> U(MEdge) a #-} cEdge :: (Label v k, Sized a) => v k -> Maybe a -> Branch v k a -> MEdge v k a -cEdge ks v ts = compact (edge ks v ts) +cEdge !ks v ts = case v of + Nothing -> case getSimpleM ts of + Null -> Nothing + Singleton e' -> Just (unDropEdge (length ks + 1) e') + NonSimple -> Just (edge ks Nothing ts) + _ -> Just (edge ks v ts) + +-- data StackView v k a z = Stack (v k) a (TrieMap Word (Hang a z)) + +data HangView a z = + Branch !Int (Maybe a) (Maybe z) +data Hang a z = H !Int z | HT !Int a z | T !Int a + +branch :: Int -> Maybe a -> Maybe z -> Hang a z +branch !i Nothing (Just z) = H i z +branch !i (Just a) Nothing = T i a +branch !i (Just a) (Just z) = HT i a z +branch _ _ _ = error "Error: bad branch" + +bView :: Hang a z -> HangView a z +bView (H i z) = Branch i Nothing (Just z) +bView (HT i a z) = Branch i (Just a) (Just z) +bView (T i a) = Branch i (Just a) Nothing + +instance Sized (Hang a z) where + getSize# _ = 1# + +{-# RULES + "sView/stack" forall ks a z branch . sView (stack ks a z branch) = Stack ks a z branch + #-} diff --git a/Data/TrieMap/RadixTrie/Slice.hs b/Data/TrieMap/RadixTrie/Slice.hs --- a/Data/TrieMap/RadixTrie/Slice.hs +++ b/Data/TrieMap/RadixTrie/Slice.hs @@ -4,7 +4,6 @@ import Control.Exception (assert) import Data.Vector.Generic -import qualified Data.Vector as V import Prelude hiding (length, zip, foldr) @@ -24,11 +23,17 @@ {-# INLINE matchSlice #-} matchSlice :: (Vector v a, Vector v b) => (a -> b -> z -> z) -> (Int -> Int -> z) -> v a -> v b -> z -matchSlice f z !xs !ys = foldr (\ (a, b) -> f a b) (z (length xs) (length ys)) (V.zip (convert xs) (convert ys)) +matchSlice f = iMatchSlice (const f) {-# INLINE iMatchSlice #-} iMatchSlice :: (Vector v a, Vector v b) => (Int -> a -> b -> z -> z) -> (Int -> Int -> z) -> v a -> v b -> z -iMatchSlice f z !xs !ys = ifoldr (\ i (a, b) -> f i a b) (z (length xs) (length ys)) (V.zip (convert xs) (convert ys)) +iMatchSlice f z !xs !ys = matcher 0 where + !xLen = length xs + !yLen = length ys + !len = min xLen yLen + matcher i + | i < len = f i (xs !$ i) (ys !$ i) (matcher (i+1)) + | otherwise = z xLen yLen {-# INLINE (!$) #-} (!$) :: Vector v a => v a -> Int -> a diff --git a/Data/TrieMap/Representation/Class.hs b/Data/TrieMap/Representation/Class.hs --- a/Data/TrieMap/Representation/Class.hs +++ b/Data/TrieMap/Representation/Class.hs @@ -19,10 +19,14 @@ toRep :: a -> Rep a toRepList :: [a] -> RepList a +-- | A default implementation of @'RepList' a@. type DRepList a = Vector (Rep a) + +-- | A default implementation of 'toRepList'. dToRepList :: Repr a => [a] -> DRepList a dToRepList = fromList . Prelude.map toRep +-- | Uses the 'RepList' instance of @a@. (This allows for efficient and automatic implementations of e.g. @Rep String@.) instance Repr a => Repr [a] where type Rep [a] = RepList a type RepList [a] = Vector (RepList a) diff --git a/Data/TrieMap/Representation/Instances.hs b/Data/TrieMap/Representation/Instances.hs --- a/Data/TrieMap/Representation/Instances.hs +++ b/Data/TrieMap/Representation/Instances.hs @@ -7,7 +7,7 @@ import Data.Bits import Data.TrieMap.Modifiers import qualified Data.Vector as V -import qualified Data.Vector.Storable as S +import qualified Data.Vector.Primitive as P import qualified Data.Set as S import qualified Data.Map as M import qualified Data.Sequence as Seq @@ -45,14 +45,14 @@ genRepr ''Ratio instance Repr Integer where - type Rep Integer = Either (Rev (Word, S.Vector Word)) (Word, S.Vector Word) + type Rep Integer = Either (Rev (Word, P.Vector Word)) (Word, P.Vector Word) toRep x - | x < 0 = let bs = unroll (-x); n = fromIntegral (S.length bs) in Left (Rev (n, bs)) - | otherwise = let bs = unroll x; n = fromIntegral (S.length bs) in Right (n, bs) + | x < 0 = let bs = unroll (-x); n = fromIntegral (P.length bs) in Left (Rev (n, bs)) + | otherwise = let bs = unroll x; n = fromIntegral (P.length bs) in Right (n, bs) DefList(Integer) -unroll :: Integer -> S.Vector Word -unroll x = S.reverse (S.unfoldr split x) +unroll :: Integer -> P.Vector Word +unroll x = P.reverse (P.unfoldr split x) where wSize = bitSize (0 :: Word) split 0 = Nothing split x = Just (fromIntegral x :: Word, shiftR x wSize) diff --git a/Data/TrieMap/Representation/Instances/ByteString.hs b/Data/TrieMap/Representation/Instances/ByteString.hs --- a/Data/TrieMap/Representation/Instances/ByteString.hs +++ b/Data/TrieMap/Representation/Instances/ByteString.hs @@ -1,27 +1,87 @@ -{-# LANGUAGE UndecidableInstances, TypeFamilies #-} +{-# LANGUAGE UndecidableInstances, TypeFamilies, BangPatterns, CPP #-} module Data.TrieMap.Representation.Instances.ByteString () where +#include "MachDeps.h" + import Data.TrieMap.Representation.Class -import Data.TrieMap.Representation.Instances.Vectors () +import Data.TrieMap.Utils +import Control.Monad +import Data.Primitive.ByteArray + +import Foreign.Ptr +import Foreign.Storable +import Foreign.ForeignPtr +import Foreign.Marshal.Array + +import Data.Bits import Data.Word -import Data.ByteString.Internal (ByteString(..)) +import Data.ByteString.Internal import qualified Data.ByteString as B -import qualified Data.ByteString.Lazy as L -import Data.Vector.Storable +import Data.Vector.Primitive +import Prelude + -- | @'Rep' 'ByteString' = 'Rep' ('Vector' 'Word8')@ instance Repr ByteString where - type Rep ByteString = Rep (Vector Word8) - toRep (PS fp off len) = toRep (unsafeFromForeignPtr fp off len) + type Rep ByteString = (Vector Word, Word) + toRep !bs = (bsToRep bs, fromIntegral (B.length bs)) type RepList ByteString = DRepList ByteString toRepList = dToRepList --- | @'Rep' 'L.ByteString' = 'Rep' ('Vector' 'Word8')@ -instance Repr L.ByteString where - type Rep L.ByteString = Rep (Vector Word8) - toRep = toRep . B.concat . L.toChunks - type RepList L.ByteString = DRepList L.ByteString - toRepList = dToRepList +bsToRep :: ByteString -> Vector Word +bsToRep (PS fp off n) = if n <= 0 then empty else inlinePerformIO $ withForeignPtr fp $ \ p0 -> + let !src = p0 `advancePtr` off :: Ptr Word8 in do + !dest <- newByteArray (n' * bytesPerWord) + let go !i = if ii < n' then (readWordAt src i >>= out >> go ii) else readLastWordAt n i src >>= out + where !ii = i + 1 + out = writeByteArray dest i + go 0 + unsafeFreeze (MVector 0 n' dest) + where n' = (n + (bytesPerWord - 1)) `quoPow` bytesPerWord + +bytesPerWord :: Int +bytesPerWord = sizeOf (0 :: Word) + +readWordAt :: Ptr Word8 -> Int -> IO Word +readWordAt ptr off = +#if WORD_SIZE_IN_BITS == 32 + accum 3 $ accum 2 $ accum 1 $ accum 0 $ return 0 +#else + accum 7 $ accum 6 $ accum 5 $ accum 4 $ accum 3 $ accum 2 $ accum 1 $ accum 0 $ return 0 +#endif + where !off' = off * bytesPerWord + accum x w = let s = 8 * (bytesPerWord - 1 - x) in + liftM2 (.|.) w $ liftM (\ w -> fromIntegral w .<<. s) $ peekElemOff ptr (x + off') + +readLastWordAt :: Int -> Int -> Ptr Word8 -> IO Word +readLastWordAt !n !off !ptr = + let w0 = accum 0 (return 0) + w1 = accum 1 w0 + w2 = accum 2 w1 + w3 = accum 3 w2 +#if WORD_SIZE_IN_BITS > 32 + w4 = accum 4 w3 + w5 = accum 5 w4 + w6 = accum 6 w5 + w7 = accum 7 w6 +#endif + in case n `remPow` bytesPerWord of + 1 -> w0 + 2 -> w1 + 3 -> w2 +#if WORD_SIZE_IN_BITS > 32 + 4 -> w3 + 5 -> w4 + 6 -> w5 + 7 -> w6 + _ -> w7 +#else + _ -> w3 +#endif + where !off' = off * bytesPerWord + {-# INLINE accum #-} + accum x w = let s = 8 * (bytesPerWord - 1 - x) in + liftM2 (.|.) w $ liftM (\ w -> fromIntegral w .<<. s) $ peekElemOff ptr (x + off') diff --git a/Data/TrieMap/Representation/Instances/Prim.hs b/Data/TrieMap/Representation/Instances/Prim.hs --- a/Data/TrieMap/Representation/Instances/Prim.hs +++ b/Data/TrieMap/Representation/Instances/Prim.hs @@ -9,7 +9,7 @@ import Data.Int import Data.Char import Data.Bits -import Data.Vector.Storable +import Data.Vector.Primitive import qualified Data.Vector.Unboxed as U import Prelude hiding (map) diff --git a/Data/TrieMap/Representation/Instances/Vectors.hs b/Data/TrieMap/Representation/Instances/Vectors.hs --- a/Data/TrieMap/Representation/Instances/Vectors.hs +++ b/Data/TrieMap/Representation/Instances/Vectors.hs @@ -3,14 +3,13 @@ module Data.TrieMap.Representation.Instances.Vectors (i2w) where import Control.Monad.Primitive +import Data.Primitive.Types import Data.Word import Data.Int import Data.Bits import Foreign.Storable (Storable) -import Foreign.Ptr -import Foreign.ForeignPtr import Data.Vector.Generic (convert, stream, unstream) import qualified Data.Vector.Generic as G @@ -39,79 +38,79 @@ toRep = V.map toRep DefList(V.Vector a) -instance Repr (S.Vector Word) where - type Rep (S.Vector Word) = S.Vector Word +instance Repr (P.Vector Word) where + type Rep (P.Vector Word) = P.Vector Word toRep = id - DefList(S.Vector Word) + DefList(P.Vector Word) -{-# INLINE unsafeCastStorable #-} -unsafeCastStorable :: (Storable a, Storable b) => (Int -> Int) -> S.Vector a -> S.Vector b -unsafeCastStorable f xs = unsafeInlineST $ do - S.MVector ptr n fp <- S.unsafeThaw xs +{-# INLINE unsafeCastPrim #-} +unsafeCastPrim :: (Prim a, Prim b) => (Int -> Int) -> P.Vector a -> P.Vector b +unsafeCastPrim f xs = unsafeInlineST $ do + P.MVector off n arr <- P.unsafeThaw xs let n' = f n - S.unsafeFreeze (S.MVector (castPtr ptr) n' (castForeignPtr fp)) + P.unsafeFreeze (P.MVector off n' arr) wordSize :: Int wordSize = bitSize (0 :: Word) #define VEC_WORD_INST(vec,wTy) \ instance Repr (vec wTy) where { \ - type Rep (vec wTy) = Rep (S.Vector wTy); \ + type Rep (vec wTy) = Rep (P.Vector wTy); \ toRep xs = toHangingVector xs;\ DefList(vec wTy)} #define HANGINSTANCE(wTy) \ - instance Repr (S.Vector wTy) where { \ - type Rep (S.Vector wTy) = (S.Vector Word, Word);\ + instance Repr (P.Vector wTy) where { \ + type Rep (P.Vector wTy) = (P.Vector Word, Word);\ {-# INLINE toRep #-}; \ toRep xs = toHangingVector xs; \ - DefList(S.Vector wTy) }; \ - VEC_WORD_INST(P.Vector,wTy); \ + DefList(P.Vector wTy) }; \ + VEC_WORD_INST(S.Vector,wTy); \ VEC_WORD_INST(U.Vector,wTy) {-# INLINE toHangingVector #-} -toHangingVector :: (G.Vector v w, Bits w, Integral w, Storable w) => v w -> (S.Vector Word, Word) -toHangingVector xs = let !ys = unstream (packStream (stream xs)) in (S.unsafeInit ys, S.unsafeLast ys) +toHangingVector :: (G.Vector v w, Bits w, Integral w, Storable w) => v w -> (P.Vector Word, Word) +toHangingVector xs = let !ys = unstream (packStream (stream xs)) in (P.unsafeInit ys, P.unsafeLast ys) --- | @'Rep' ('S.Vector' 'Word8') = 'S.Vector' 'Word'@, by packing multiple 'Word8's into each 'Word' for space efficiency. +-- | @'Rep' ('P.Vector' 'Word8') = 'P.Vector' 'Word'@, by packing multiple 'Word8's into each 'Word' for space efficiency. HANGINSTANCE(Word8) --- | @'Rep' ('S.Vector' 'Word16') = 'S.Vector' 'Word'@, by packing multiple 'Word16's into each 'Word' for space efficiency. +-- | @'Rep' ('P.Vector' 'Word16') = 'P.Vector' 'Word'@, by packing multiple 'Word16's into each 'Word' for space efficiency. HANGINSTANCE(Word16) #if WORD_SIZE_IN_BITS == 32 -instance Repr (S.Vector Word32) where - type Rep (S.Vector Word32) = S.Vector Word - toRep xs = unsafeCastStorable id xs - DefList (S.Vector Word32) -instance Repr (U.Vector Word32) where - type Rep (U.Vector Word32) = S.Vector Word - toRep xs = unsafeCastStorable id (convert xs) - DefList (U.Vector Word32) instance Repr (P.Vector Word32) where - type Rep (P.Vector Word32) = S.Vector Word - toRep xs = unsafeCastStorable id (convert xs) + type Rep (P.Vector Word32) = P.Vector Word + toRep xs = unsafeCastPrim id xs DefList (P.Vector Word32) +instance Repr (U.Vector Word32) where + type Rep (U.Vector Word32) = P.Vector Word + toRep xs = unsafeCastPrim id (convert xs) + DefList (U.Vector Word32) +instance Repr (S.Vector Word32) where + type Rep (S.Vector Word32) = P.Vector Word + toRep xs = unsafeCastPrim id (convert xs) + DefList (S.Vector Word32) #elif WORD_SIZE_IN_BITS > 32 HANGINSTANCE(Word32) #endif #if WORD_SIZE_IN_BITS == 32 --- | @'Rep' ('S.Vector' 'Word64') = 'S.Vector' 'Word'@, by viewing each 'Word64' as two 'Word's. +-- | @'Rep' ('P.Vector' 'Word64') = 'P.Vector' 'Word'@, by viewing each 'Word64' as two 'Word's. #else --- | @'Rep' ('S.Vector' 'Word64') = 'S.Vector' 'Word'@ +-- | @'Rep' ('P.Vector' 'Word64') = 'P.Vector' 'Word'@ #endif -instance Repr (S.Vector Word64) where - type Rep (S.Vector Word64) = S.Vector Word - toRep xs = unsafeCastStorable (ratio *) xs +instance Repr (P.Vector Word64) where + type Rep (P.Vector Word64) = P.Vector Word + toRep xs = unsafeCastPrim (ratio *) xs where !wordBits = bitSize (0 :: Word); ratio = quoPow 64 wordBits - DefList(S.Vector Word64) + DefList(P.Vector Word64) #define VEC_INT_INST(vec,iTy,wTy) \ instance Repr (vec iTy) where { \ - type Rep (vec iTy) = Rep (S.Vector wTy); \ - toRep xs = (toRep :: S.Vector wTy -> Rep (S.Vector wTy)) (convert (G.map (i2w :: iTy -> wTy) xs)); \ + type Rep (vec iTy) = Rep (P.Vector wTy); \ + toRep xs = (toRep :: P.Vector wTy -> Rep (P.Vector wTy)) (convert (G.map (i2w :: iTy -> wTy) xs)); \ DefList(vec iTy)} #define VEC_INT_INSTANCES(iTy,wTy) \ - VEC_INT_INST(S.Vector,iTy,wTy); \ VEC_INT_INST(P.Vector,iTy,wTy); \ + VEC_INT_INST(S.Vector,iTy,wTy); \ VEC_INT_INST(U.Vector,iTy,wTy) VEC_INT_INSTANCES(Int8, Word8) @@ -122,16 +121,16 @@ #define VEC_ENUM_INST(ty, vec) \ instance Repr (vec ty) where { \ - type Rep (vec ty) = S.Vector Word; \ + type Rep (vec ty) = P.Vector Word; \ {-# INLINE toRep #-}; \ toRep xs = convert (G.map (fromIntegral . fromEnum) xs);\ DefList(vec ty)} #define VEC_ENUM_INSTANCES(ty) \ - VEC_ENUM_INST(ty,S.Vector); \ VEC_ENUM_INST(ty,P.Vector); \ + VEC_ENUM_INST(ty,S.Vector); \ VEC_ENUM_INST(ty,U.Vector) --- | @'Rep' ('S.Vector' 'Char') = 'S.Vector' 'Word'@ +-- | @'Rep' ('P.Vector' 'Char') = 'P.Vector' 'Word'@ VEC_ENUM_INSTANCES(Char) -- | We embed IntN into WordN, but we have to be careful about overflow. @@ -164,20 +163,15 @@ Skip s' -> return $ Skip (PackState w i s') Yield ww s' -> return $ Skip (PackState ((w .<<. bitSize (0 :: w)) .|. fromIntegral ww) (i-1) s') -instance Repr (S.Vector Bool) where - type Rep (S.Vector Bool) = (S.Vector Word, Word) - toRep = boolVecToRep - DefList(S.Vector Bool) - instance Repr (U.Vector Bool) where - type Rep (U.Vector Bool) = (S.Vector Word, Word) + type Rep (U.Vector Bool) = (P.Vector Word, Word) {-# INLINE toRep #-} toRep xs = boolVecToRep xs DefList(U.Vector Bool) {-# INLINE boolVecToRep #-} -boolVecToRep :: G.Vector v Bool => v Bool -> (S.Vector Word, Word) -boolVecToRep xs = let !ys = unstream (packBoolStream (stream xs)) in (S.unsafeInit ys, S.unsafeLast ys) +boolVecToRep :: G.Vector v Bool => v Bool -> (P.Vector Word, Word) +boolVecToRep xs = let !ys = unstream (packBoolStream (stream xs)) in (P.unsafeInit ys, P.unsafeLast ys) {-# INLINE packBoolStream #-} packBoolStream :: Monad m => Stream m Bool -> Stream m Word diff --git a/Data/TrieMap/ReverseMap.hs b/Data/TrieMap/ReverseMap.hs --- a/Data/TrieMap/ReverseMap.hs +++ b/Data/TrieMap/ReverseMap.hs @@ -1,18 +1,16 @@ -{-# LANGUAGE TypeFamilies, MagicHash, UnboxedTuples, GeneralizedNewtypeDeriving, FlexibleInstances #-} +{-# LANGUAGE TypeFamilies, FlexibleContexts, GeneralizedNewtypeDeriving, FlexibleInstances, NamedFieldPuns, RecordWildCards #-} +{-# LANGUAGE MultiParamTypeClasses, CPP, UnboxedTuples, MagicHash #-} module Data.TrieMap.ReverseMap () where -import Control.Applicative -import Control.Monad import Control.Monad.Ends -import Data.Foldable import qualified Data.Monoid as M import Data.TrieMap.TrieKey import Data.TrieMap.Modifiers -import Data.TrieMap.Sized import Prelude hiding (foldr, foldl, foldr1, foldl1) +import GHC.Exts newtype DualPlus m a = DualPlus {runDualPlus :: m a} deriving (Functor, Monad) newtype Dual f a = Dual {runDual :: f a} deriving (Functor) @@ -25,13 +23,41 @@ mzero = DualPlus mzero DualPlus m `mplus` DualPlus k = DualPlus (k `mplus` m) -instance TrieKey k => Foldable (TrieMap (Rev k)) where +#define INSTANCE(cl) (TrieKey k, cl (TrieMap k)) => cl (TrieMap (Rev k)) + +instance INSTANCE(Functor) where + fmap f (RevMap m) = RevMap (f <$> m) + +instance INSTANCE(Foldable) where foldMap f (RevMap m) = M.getDual (foldMap (M.Dual . f) m) foldr f z (RevMap m) = foldl (flip f) z m foldl f z (RevMap m) = foldr (flip f) z m - foldr1 f (RevMap m) = foldl1 (flip f) m - foldl1 f (RevMap m) = foldr1 (flip f) m +instance INSTANCE(Traversable) where + traverse f (RevMap m) = RevMap <$> runDual (traverse (Dual . f) m) + +instance INSTANCE(Subset) where + RevMap m1 <=? RevMap m2 = m1 <=? m2 + +instance TrieKey k => Buildable (TrieMap (Rev k)) (Rev k) where + type UStack (TrieMap (Rev k)) = UMStack k + uFold = fmap RevMap . mapFoldlKeys getRev . uFold + type AStack (TrieMap (Rev k)) = RevFold (AMStack k) k + aFold = fmap RevMap . mapFoldlKeys getRev . reverseFold . aFold + type DAStack (TrieMap (Rev k)) = RevFold (DAMStack k) k + daFold = RevMap <$> mapFoldlKeys getRev (reverseFold daFold) + +#define SETOP(op) op f (RevMap m1) (RevMap m2) = RevMap (op f m1 m2) + +instance INSTANCE(SetOp) where + SETOP(union) + SETOP(diff) + SETOP(isect) + +instance INSTANCE(Project) where + mapMaybe f (RevMap m) = RevMap $ mapMaybe f m + mapEither f (RevMap m) = both RevMap (mapEither f) m + -- | @'TrieMap' ('Rev' k) a@ is a wrapper around a @'TrieMap' k a@ that reverses the order of the operations. instance TrieKey k => TrieKey (Rev k) where newtype TrieMap (Rev k) a = RevMap (TrieMap k a) @@ -39,30 +65,20 @@ emptyM = RevMap emptyM singletonM (Rev k) a = RevMap (singletonM k a) - lookupM (Rev k) (RevMap m) = lookupM k m + lookupMC (Rev k) (RevMap m) = lookupMC k m sizeM (RevMap m) = sizeM m getSimpleM (RevMap m) = getSimpleM m - fmapM f (RevMap m) = RevMap (fmapM f m) - traverseM f (RevMap m) = RevMap <$> runDual (traverseM (Dual . f) m) - - mapMaybeM f (RevMap m) = RevMap (mapMaybeM f m) - mapEitherM f (RevMap m) = both RevMap RevMap (mapEitherM f) m - unionM f (RevMap m1) (RevMap m2) = RevMap (unionM f m1 m2) - isectM f (RevMap m1) (RevMap m2) = RevMap (isectM f m1 m2) - diffM f (RevMap m1) (RevMap m2) = RevMap (diffM f m1 m2) - isSubmapM (<=) (RevMap m1) (RevMap m2) = isSubmapM (<=) m1 m2 - singleHoleM (Rev k) = RHole (singleHoleM k) beforeM (RHole hole) = RevMap (afterM hole) beforeWithM a (RHole hole) = RevMap (afterWithM a hole) afterM (RHole hole) = RevMap (beforeM hole) afterWithM a (RHole hole) = RevMap (beforeWithM a hole) searchMC (Rev k) (RevMap m) = mapSearch RHole (searchMC k m) - indexM i (RevMap m) = case indexM (revIndex i m) m of - (# i', a, hole #) -> (# revIndex i' a, a, RHole hole #) - where revIndex :: Sized a => Int -> a -> Int - revIndex i a = getSize a - 1 - i + indexM (RevMap m) i = case indexM m (revIndex i m) of + (# i', a, hole #) -> (# revIndex i' a, a, RHole hole #) + where revIndex :: Sized a => Int# -> a -> Int# + revIndex i a = getSize# a -# 1# -# i extractHoleM (RevMap m) = fmap RHole <$> runDualPlus (extractHoleM m) firstHoleM (RevMap m) = First (fmap RHole <$> getLast (lastHoleM m)) @@ -72,8 +88,19 @@ clearM (RHole m) = RevMap (clearM m) insertWithM f (Rev k) a (RevMap m) = RevMap (insertWithM f k a m) - fromListM f xs = RevMap (fromListM f [(k, a) | (Rev k, a) <- xs]) - fromAscListM f xs = RevMap (fromAscListM (flip f) [(k, a) | (Rev k, a) <- reverse xs]) - fromDistAscListM xs = RevMap (fromDistAscListM [(k, a) | (Rev k, a) <- reverse xs]) unifierM (Rev k') (Rev k) a = RHole <$> unifierM k' k a + +{-# INLINE reverseFold #-} +reverseFold :: FromList z k a -> FromList (RevFold z k) k a +reverseFold Foldl{snoc = snoc0, begin = begin0, zero, done = done0} + = Foldl {..} where + snoc g k a = RevFold $ \ m -> case m of + Nothing -> runRevFold g (Just $ begin0 k a) + Just m -> runRevFold g (Just $ snoc0 m k a) + + begin = snoc (RevFold $ maybe zero done0) + + done g = runRevFold g Nothing + +newtype RevFold z k a = RevFold {runRevFold :: Maybe (z a) -> TrieMap k a} diff --git a/Data/TrieMap/Sized.hs b/Data/TrieMap/Sized.hs --- a/Data/TrieMap/Sized.hs +++ b/Data/TrieMap/Sized.hs @@ -1,23 +1,30 @@ -{-# LANGUAGE MagicHash, DeriveFunctor, DeriveFoldable, DeriveTraversable #-} +{-# LANGUAGE MagicHash, DeriveFunctor, DeriveFoldable, DeriveTraversable, ImplicitParams #-} module Data.TrieMap.Sized where +import Data.TrieMap.TrieKey.Subset import Data.Foldable import Data.Traversable import GHC.Exts class Sized a where - getSize# :: a -> Int# + getSize# :: a -> Int# data Assoc k a = Assoc {getK :: k, getValue :: a} deriving (Functor, Foldable, Traversable) newtype Elem a = Elem {getElem :: a} deriving (Functor, Foldable, Traversable) +instance Subset Elem where + Elem a <=? Elem b = ?le a b + +instance Subset (Assoc k) where + Assoc _ a <=? Assoc _ b = ?le a b + instance Sized (Elem a) where - getSize# _ = 1# + getSize# _ = 1# instance Sized (Assoc k a) where - getSize# _ = 1# + getSize# _ = 1# instance Sized a => Sized (Maybe a) where getSize# (Just a) = getSize# a diff --git a/Data/TrieMap/TrieKey.hs b/Data/TrieMap/TrieKey.hs --- a/Data/TrieMap/TrieKey.hs +++ b/Data/TrieMap/TrieKey.hs @@ -1,45 +1,50 @@ -{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash, FlexibleContexts, TupleSections, Rank2Types #-} +{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash, FlexibleContexts, TupleSections, Rank2Types, ExistentialQuantification #-} +{-# LANGUAGE NamedFieldPuns, RecordWildCards, ImplicitParams, TypeOperators #-} -module Data.TrieMap.TrieKey where +module Data.TrieMap.TrieKey ( + module Data.TrieMap.TrieKey, + module Data.Foldable, + module Data.Traversable, + module Control.Applicative, + module Data.TrieMap.Sized, + module Data.TrieMap.Utils, + module Data.TrieMap.TrieKey.Subset, + module Data.TrieMap.TrieKey.Buildable, + module Data.TrieMap.TrieKey.SetOp, + module Data.TrieMap.TrieKey.Projection, + module Data.TrieMap.TrieKey.Search, + MonadPlus(..), + Monoid(..), + guard) where import Data.TrieMap.Sized import Data.TrieMap.Utils +import Data.TrieMap.TrieKey.Subset +import Data.TrieMap.TrieKey.Buildable +import Data.TrieMap.TrieKey.SetOp +import Data.TrieMap.TrieKey.Projection +import Data.TrieMap.TrieKey.Search -import Control.Applicative (Applicative) +import Control.Applicative hiding (empty) import Control.Monad +import Control.Monad.Lookup import Control.Monad.Ends -import Data.Foldable hiding (foldrM, foldlM) -import qualified Data.List as L +import Data.Monoid (Monoid(..)) +import Data.Foldable +import Data.Traversable import Prelude hiding (foldr, foldl) import GHC.Exts -type LEq a b = a -> b -> Bool -type SearchCont h a r = (h -> r) -> (a -> h -> r) -> r -type Lookup a = Maybe a +type FromList stack k a = Foldl stack k a (TrieMap k a) +type UMStack k = UStack (TrieMap k) +type AMStack k = AStack (TrieMap k) +type DAMStack k = DAStack (TrieMap k) data Simple a = Null | Singleton a | NonSimple -class (Functor f, Monad f) => Option f where - none :: f a - some :: a -> f a - option :: f a -> r -> (a -> r) -> r - -instance Option Maybe where - none = Nothing - some = Just - option m a f = maybe a f m - -{-# INLINE [0] liftMaybe #-} -liftMaybe :: Option f => Maybe a -> f a -liftMaybe = maybe none some - -{-# INLINE [0] toMaybe #-} -toMaybe :: Option f => f a -> Maybe a -toMaybe x = option x Nothing Just - instance Monad Simple where return = Singleton Null >>= _ = Null @@ -64,66 +69,59 @@ -- | A @TrieKey k@ instance implies that @k@ is a standardized representation for which a -- generalized trie structure can be derived. -class (Ord k, Foldable (TrieMap k)) => TrieKey k where - data TrieMap k :: * -> * - emptyM :: TrieMap k a - singletonM :: Sized a => k -> a -> TrieMap k a - getSimpleM :: TrieMap k a -> Simple a - sizeM# :: Sized a => TrieMap k a -> Int# - sizeM :: Sized a => TrieMap k a -> Int - lookupM :: k -> TrieMap k a -> Lookup a - fmapM :: Sized b => (a -> b) -> TrieMap k a -> TrieMap k b - traverseM :: (Applicative f, Sized b) => - (a -> f b) -> TrieMap k a -> f (TrieMap k b) - mapMaybeM :: Sized b => (a -> Maybe b) -> TrieMap k a -> TrieMap k b - mapEitherM :: (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> TrieMap k a -> (# TrieMap k b, TrieMap k c #) - unionM :: Sized a => (a -> a -> Maybe a) -> TrieMap k a -> TrieMap k a -> TrieMap k a - isectM :: (Sized a, Sized b, Sized c) => - (a -> b -> Maybe c) -> TrieMap k a -> TrieMap k b -> TrieMap k c - diffM :: Sized a => (a -> b -> Maybe a) -> TrieMap k a -> TrieMap k b -> TrieMap k a - isSubmapM :: (Sized a, Sized b) => LEq a b -> LEq (TrieMap k a) (TrieMap k b) - - fromListM, fromAscListM :: Sized a => (a -> a -> a) -> [(k, a)] -> TrieMap k a - fromDistAscListM :: Sized a => [(k, a)] -> TrieMap k a - insertWithM :: (TrieKey k, Sized a) => (a -> a) -> k -> a -> TrieMap k a -> TrieMap k a - - data Hole k :: * -> * - singleHoleM :: k -> Hole k a - beforeM, afterM :: Sized a => Hole k a -> TrieMap k a - beforeWithM, afterWithM :: Sized a => a -> Hole k a -> TrieMap k a - searchMC :: k -> TrieMap k a -> SearchCont (Hole k a) a r - indexM :: Sized a => Int -> TrieMap k a -> (# Int, a, Hole k a #) - indexM# :: Sized a => Int# -> TrieMap k a -> (# Int#, a, Hole k a #) +class (Ord k, + Buildable (TrieMap k) k, + Subset (TrieMap k), + Traversable (TrieMap k), + SetOp (TrieMap k), + Project (TrieMap k)) => TrieKey k where + data TrieMap k :: * -> * + emptyM :: TrieMap k a + singletonM :: Sized a => k -> a -> TrieMap k a + getSimpleM :: TrieMap k a -> Simple a + sizeM# :: Sized a => TrieMap k a -> Int# + sizeM :: Sized a => TrieMap k a -> Int + lookupMC :: k -> TrieMap k a -> Lookup r a + + insertWithM :: (TrieKey k, Sized a) => (a -> a) -> k -> a -> TrieMap k a -> TrieMap k a + + data Hole k :: * -> * + singleHoleM :: k -> Hole k a + beforeM, afterM :: Sized a => Hole k a -> TrieMap k a + beforeWithM, afterWithM :: Sized a => a -> Hole k a -> TrieMap k a + searchMC :: k -> TrieMap k a -> SearchCont (Hole k a) a r + indexM :: Sized a => TrieMap k a -> Int# -> (# Int#, a, Hole k a #) - -- By combining rewrite rules and these NOINLINE pragmas, we automatically derive - -- specializations of functions for every instance of TrieKey. - extractHoleM :: (Functor m, MonadPlus m) => Sized a => TrieMap k a -> m (a, Hole k a) - {-# NOINLINE firstHoleM #-} - {-# NOINLINE lastHoleM #-} - {-# NOINLINE sizeM# #-} - {-# NOINLINE indexM# #-} - sizeM# m = unbox (inline sizeM m) - indexM# i# m = case inline indexM (I# i#) m of - (# I# i'#, a, hole #) -> (# i'#, a, hole #) - firstHoleM :: Sized a => TrieMap k a -> First (a, Hole k a) - firstHoleM m = inline extractHoleM m - lastHoleM :: Sized a => TrieMap k a -> Last (a, Hole k a) - lastHoleM m = inline extractHoleM m - - insertWithM f k a m = inline searchMC k m (assignM a) (assignM . f) - - assignM :: Sized a => a -> Hole k a -> TrieMap k a - clearM :: Sized a => Hole k a -> TrieMap k a - unifierM :: Sized a => k -> k -> a -> Maybe (Hole k a) - - fromListM f = L.foldl' (\ m (k, a) -> insertWithM (f a) k a m) emptyM - fromAscListM = fromListM - fromDistAscListM = fromAscListM const - unifierM k' k a = searchMC k' (singletonM k a) Just (\ _ _ -> Nothing) + -- By combining rewrite rules and these NOINLINE pragmas, we automatically derive + -- specializations of functions for every instance of TrieKey. + extractHoleM :: (Functor m, MonadPlus m) => Sized a => TrieMap k a -> m (a, Hole k a) + {-# NOINLINE firstHoleM #-} + {-# NOINLINE lastHoleM #-} + {-# NOINLINE sizeM# #-} + sizeM# m = unbox (inline sizeM m) + firstHoleM :: Sized a => TrieMap k a -> First (a, Hole k a) + firstHoleM m = inline extractHoleM m + lastHoleM :: Sized a => TrieMap k a -> Last (a, Hole k a) + lastHoleM m = inline extractHoleM m + + insertWithM f k a m = inline searchMC k m (assignM a) (assignM . f) + + assignM :: Sized a => a -> Hole k a -> TrieMap k a + clearM :: Sized a => Hole k a -> TrieMap k a + unifierM :: Sized a => k -> k -> a -> Lookup r (Hole k a) + unifyM :: Sized a => k -> a -> k -> a -> Lookup r (TrieMap k a) + + unifierM k' k a = Lookup $ \ no yes -> searchMC k' (singletonM k a) yes (\ _ _ -> no) + unifyM k1 a1 k2 a2 = assignM a1 <$> unifierM k1 k2 a2 instance (TrieKey k, Sized a) => Sized (TrieMap k a) where getSize# = sizeM# +instance TrieKey k => Nullable (TrieMap k) where + isNull m = case getSimpleM m of + Null -> True + _ -> False + foldl1Empty :: a foldl1Empty = error "Error: cannot call foldl1 on an empty map" @@ -134,41 +132,20 @@ fillHoleM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a fillHoleM = maybe clearM assignM -{-# INLINE mapSearch #-} -mapSearch :: (hole -> hole') -> SearchCont hole a r -> SearchCont hole' a r -mapSearch f run nomatch match = run nomatch' match' where - nomatch' hole = nomatch (f hole) - match' a hole = match a (f hole) +{-# INLINE lookupM #-} +lookupM :: TrieKey k => k -> TrieMap k a -> Maybe a +lookupM k m = runLookup (lookupMC k m) Nothing Just -{-# INLINE unifyM #-} -unifyM :: (TrieKey k, Sized a) => k -> a -> k -> a -> Maybe (TrieMap k a) -unifyM k1 a1 k2 a2 = case unifierM k1 k2 a2 of - Nothing -> Nothing - Just hole -> Just $ inline assignM a1 hole +{-# INLINE mappendM #-} +mappendM :: Monoid m => Maybe m -> Maybe m -> m +Nothing `mappendM` Nothing = mempty +Nothing `mappendM` Just m = m +Just m `mappendM` Nothing = m +Just m1 `mappendM` Just m2 = m1 `mappend` m2 insertWithM' :: (TrieKey k, Sized a) => (a -> a) -> k -> a -> Maybe (TrieMap k a) -> TrieMap k a insertWithM' f k a = maybe (singletonM k a) (insertWithM f k a) -mapMaybeM' :: (TrieKey k, Sized b) => (a -> Maybe b) -> TrieMap k a -> Maybe (TrieMap k b) -mapMaybeM' = guardNullM .: mapMaybeM - -mapEitherM' :: (TrieKey k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> TrieMap k a -> - (# Maybe (TrieMap k b), Maybe (TrieMap k c) #) -mapEitherM' = both guardNullM guardNullM . mapEitherM - -mapEitherM'' :: (TrieKey k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> Maybe (TrieMap k a) -> - (# Maybe (TrieMap k b), Maybe (TrieMap k c) #) -mapEitherM'' = mapEitherMaybe . mapEitherM' - -unionM' :: (TrieKey k, Sized a) => (a -> a -> Maybe a) -> TrieMap k a -> TrieMap k a -> Maybe (TrieMap k a) -unionM' f m1 m2 = guardNullM (unionM f m1 m2) - -isectM' :: (TrieKey k, Sized a, Sized b, Sized c) => (a -> b -> Maybe c) -> TrieMap k a -> TrieMap k b -> Maybe (TrieMap k c) -isectM' f m1 m2 = guardNullM (isectM f m1 m2) - -diffM' :: (TrieKey k, Sized a) => (a -> b -> Maybe a) -> TrieMap k a -> TrieMap k b -> Maybe (TrieMap k a) -diffM' f m1 m2 = guardNullM (diffM f m1 m2) - {-# INLINE beforeMM #-} beforeMM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a beforeMM = maybe beforeM beforeWithM @@ -178,7 +155,7 @@ afterMM = maybe afterM afterWithM clearM' :: (TrieKey k, Sized a) => Hole k a -> Maybe (TrieMap k a) -clearM' hole = guardNullM (clearM hole) +clearM' hole = guardNull (clearM hole) {-# INLINE alterM #-} alterM :: (TrieKey k, Sized a) => (Maybe a -> Maybe a) -> k -> TrieMap k a -> TrieMap k a @@ -193,63 +170,16 @@ searchMC' k Nothing f _ = f (singleHoleM k) searchMC' k (Just m) f g = searchMC k m f g -nullM :: TrieKey k => TrieMap k a -> Bool -nullM m = case getSimpleM m of - Null -> True - _ -> False - -guardNullM :: TrieKey k => TrieMap k a -> Maybe (TrieMap k a) -guardNullM m - | nullM m = Nothing - | otherwise = Just m - -sides :: (b -> d) -> (a -> (# b, c, b #)) -> a -> (# d, c, d #) -sides g f a = case f a of - (# x, y, z #) -> (# g x, y, g z #) - -both :: (b -> b') -> (c -> c') -> (a -> (# b, c #)) -> a -> (# b', c' #) -both g1 g2 f a = case f a of - (# x, y #) -> (# g1 x, g2 y #) - elemsM :: TrieKey k => TrieMap k a -> [a] elemsM m = build (\ f z -> foldr f z m) -mapEitherMaybe :: (a -> (# Maybe b, Maybe c #)) -> Maybe a -> (# Maybe b, Maybe c #) -mapEitherMaybe f (Just a) = f a -mapEitherMaybe _ _ = (# Nothing, Nothing #) - -{-# INLINE unionMaybe #-} -unionMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a -unionMaybe f (Just x) (Just y) = f x y -unionMaybe _ Nothing y = y -unionMaybe _ x Nothing = x - -isectMaybe :: (a -> b -> Maybe c) -> Maybe a -> Maybe b -> Maybe c -isectMaybe f (Just x) (Just y) = f x y -isectMaybe _ _ _ = Nothing - -diffMaybe :: (a -> b -> Maybe a) -> Maybe a -> Maybe b -> Maybe a -diffMaybe _ Nothing _ = Nothing -diffMaybe _ (Just x) Nothing = Just x -diffMaybe f (Just x) (Just y) = f x y - -subMaybe :: (a -> b -> Bool) -> Maybe a -> Maybe b -> Bool -subMaybe _ Nothing _ = True -subMaybe (<=) (Just a) (Just b) = a <= b -subMaybe _ _ _ = False - -indexFail :: a -> (# Int, b, c #) -indexFail _ = (# error err, error err, error err #) where - err = "Error: not a valid index" +indexFail :: a +indexFail = error "Error: index out of bounds" {-# RULES "extractHoleM/First" [0] extractHoleM = firstHoleM; "extractHoleM/Last" [0] extractHoleM = lastHoleM; "sizeM" [0] forall m . sizeM m = I# (sizeM# m); - "indexM" [0] forall i m . indexM i m = case indexM# (unbox i) m of { - (# i'#, a, m #) -> (# I# i'#, a, m #)}; "getSimpleM/emptyM" getSimpleM emptyM = Null; "getSimpleM/singletonM" forall k a . getSimpleM (singletonM k a) = Singleton a; - "toMaybe" forall f . toMaybe f = f; - "liftMaybe" forall m . liftMaybe m = m; #-} diff --git a/Data/TrieMap/TrieKey/Buildable.hs b/Data/TrieMap/TrieKey/Buildable.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/TrieKey/Buildable.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE TypeFamilies, NamedFieldPuns, RecordWildCards, FunctionalDependencies, BangPatterns, MultiParamTypeClasses, ViewPatterns #-} +module Data.TrieMap.TrieKey.Buildable ( + Buildable(..), + Foldl(..), + mapFoldlKeys, + runFoldl, + defaultUFold, + Distinct, + combineFold) where + +import Data.TrieMap.Sized + +class Buildable f k | f -> k where + type UStack f :: * -> * + uFold :: Sized a => (a -> a -> a) -> Foldl (UStack f) k a (f a) + type AStack f :: * -> * + aFold :: Sized a => (a -> a -> a) -> Foldl (AStack f) k a (f a) + type DAStack f :: * -> * + daFold :: Sized a => Foldl (DAStack f) k a (f a) + +data Foldl stack k a result = + Foldl {snoc :: stack a -> k -> a -> stack a, + begin :: k -> a -> stack a, + zero :: result, + done :: stack a -> result} + +instance Functor (Foldl stack k a) where + fmap f Foldl{..} = Foldl{zero = f zero, done = f . done, ..} + +{-# INLINE runFoldl #-} +runFoldl :: Foldl stack k a result -> [(k, a)] -> result +runFoldl Foldl{..} = fL where + fL [] = zero + fL ((k, a):xs) = fL' (begin k a) xs + + fL' !s ((k, a):xs) = fL' (snoc s k a) xs + fL' s [] = done s + +{-# INLINE mapFoldlKeys #-} +mapFoldlKeys :: (k -> k') -> Foldl stack k' a result -> Foldl stack k a result +mapFoldlKeys f Foldl{..} = Foldl{snoc = \ z k a -> snoc z (f k) a, begin = begin . f, ..} + +{-# INLINE defaultUFold #-} +defaultUFold :: f a -> (k -> a -> f a) -> ((a -> a) -> k -> a -> f a -> f a) -> + (a -> a -> a) -> Foldl f k a (f a) +defaultUFold empty single insert f = Foldl{ + zero = empty, + begin = single, + snoc = \ m k a -> insert (f a) k a m, + done = id} + +data Distinct k z a = Begin k a | Dist k a (z a) + +{-# INLINE combineFold #-} +combineFold :: Eq k => Foldl stack k a result -> (a -> a -> a) -> Foldl (Distinct k stack) k a result +combineFold Foldl{..} f = Foldl{snoc = snoc', begin = Begin, zero, done = done'} where + snoc' (Begin k a) k' a' + | k == k' = Begin k (f a' a) + snoc' (Dist k a stk) k' a' + | k == k' = Dist k (f a' a) stk + snoc' stk k a = Dist k a (collapse stk) + + done' = done . collapse + + collapse (Begin k a) = begin k a + collapse (Dist k a stk) = snoc stk k a diff --git a/Data/TrieMap/TrieKey/Projection.hs b/Data/TrieMap/TrieKey/Projection.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/TrieKey/Projection.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE LiberalTypeSynonyms, UnboxedTuples, ScopedTypeVariables, Rank2Types #-} +module Data.TrieMap.TrieKey.Projection (MapMaybe, MapEither, Project(..), mapMaybeM, mapEitherM, both, both') where + +import Data.TrieMap.Sized +import Data.TrieMap.TrieKey.Subset + +type MapMaybe f a b = f a -> Maybe (f b) +type MapEither f a b c = f a -> (# Maybe (f b), Maybe (f c) #) +type Id a = a + +class Project f where + mapMaybe :: Sized b => MapMaybe Id a b -> f a -> f b + mapEither :: (Sized b, Sized c) => MapEither Id a b c -> f a -> (# f b, f c #) + + mapEither f a = (# mapMaybe f1 a, mapMaybe f2 a #) where + f1 a = case f a of + (# b, _ #) -> b + f2 a = case f a of + (# _, c #) -> c + mapMaybe (f :: MapMaybe Id a b) a = case mapEither g a of + (# fb, _ #) -> fb + where g :: MapEither Id a b (Elem a) + g a = (# f a, Nothing #) + +instance Project Maybe where + mapMaybe f m = m >>= f + mapEither _ Nothing = (# Nothing, Nothing #) + mapEither f (Just a) = f a + +mapMaybeM :: (Sized b, Project f, Nullable f) => MapMaybe Id a b -> MapMaybe f a b +mapMaybeM f a = guardNull (mapMaybe f a) + +mapEitherM :: (Sized b, Sized c, Project f, Nullable f) => MapEither Id a b c -> MapEither f a b c +mapEitherM f a = case mapEither f a of + (# b, c #) -> (# guardNull b, guardNull c #) + +both :: (Sized b, Sized c) => (forall x . Sized x => f x -> f' x) -> (a -> (# f b, f c #)) -> a -> (# f' b, f' c #) +both g f a = case f a of + (# x, y #) -> (# g x, g y #) + +both' :: (b -> b') -> (c -> c') -> (a -> (# b, c #)) -> a -> (# b', c' #) +both' g1 g2 f a = case f a of + (# x, y #) -> (# g1 x, g2 y #) diff --git a/Data/TrieMap/TrieKey/Search.hs b/Data/TrieMap/TrieKey/Search.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/TrieKey/Search.hs @@ -0,0 +1,10 @@ +module Data.TrieMap.TrieKey.Search where + +type SearchCont h a r = (h -> r) -> (a -> h -> r) -> r + +{-# INLINE mapSearch #-} +mapSearch :: (hole -> hole') -> SearchCont hole a r -> SearchCont hole' a r +mapSearch f run nomatch match = run nomatch' match' where + nomatch' hole = nomatch (f hole) + match' a hole = match a (f hole) + diff --git a/Data/TrieMap/TrieKey/SetOp.hs b/Data/TrieMap/TrieKey/SetOp.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/TrieKey/SetOp.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE LiberalTypeSynonyms, ImplicitParams, TypeOperators, CPP #-} +module Data.TrieMap.TrieKey.SetOp ( + IsectM, UnionM, DiffM, + isectM, unionM, diffM, + Isect, Union, Diff, + SetOp(..)) where + +import Data.TrieMap.Sized +import Data.TrieMap.TrieKey.Subset + +type IsectM f a b c = f a -> f b -> Maybe (f c) +type UnionM f a = f a -> f a -> Maybe (f a) +type DiffM f a b = f a -> f b -> Maybe (f a) + +type Isect f a b c = f a -> f b -> f c +type Union f a = f a -> f a -> f a +type Diff f a b = f a -> f b -> f a + +type Id a = a + +class SetOp f where + isect :: Sized c => IsectM Id a b c -> Isect f a b c + union :: Sized a => UnionM Id a -> Union f a + diff :: Sized a => DiffM Id a b -> Diff f a b + +instance SetOp Maybe where + {-# INLINE isect #-} + {-# INLINE union #-} + {-# INLINE diff #-} + isect f (Just a) (Just b) = f a b + isect _ _ _ = Nothing + union f (Just a) (Just b) = f a b + union _ (Just a) Nothing = Just a + union _ Nothing (Just b) = Just b + union _ Nothing Nothing = Nothing + diff f (Just a) (Just b) = f a b + diff _ (Just a) Nothing = Just a + diff _ Nothing _ = Nothing + +{-# INLINE isectM #-} +isectM :: (Nullable f, SetOp f, Sized c) => IsectM Id a b c -> IsectM f a b c +isectM f a b = guardNull (isect f a b) + +{-# INLINE diffM #-} +diffM :: (Nullable f, SetOp f, Sized a) => DiffM Id a b -> DiffM f a b +diffM f a b = guardNull (diff f a b) + +{-# INLINE unionM #-} +unionM :: (Nullable f, SetOp f, Sized a) => UnionM Id a -> UnionM f a +unionM f a b = guardNull (union f a b) diff --git a/Data/TrieMap/TrieKey/Subset.hs b/Data/TrieMap/TrieKey/Subset.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/TrieKey/Subset.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE ImplicitParams #-} + +module Data.TrieMap.TrieKey.Subset where + +type LEq a b = a -> b -> Bool +class Subset f where + (<=?) :: (?le :: LEq a b) => LEq (f a) (f b) + +(<<=?) :: (Subset f, Subset g, ?le :: LEq a b) => LEq (f (g a)) (f (g b)) +f <<=? g = let ?le = (<=?) in f <=? g + +instance Subset Maybe where + Nothing <=? _ = True + Just a <=? Just b = a Bool + +{-# INLINE guardNull #-} +guardNull :: Nullable f => f a -> Maybe (f a) +guardNull a + | isNull a = Nothing + | otherwise = Just a + +instance Nullable Maybe where + isNull Nothing = True + isNull Just{} = False + +( LEq a b +( 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 ^ Nothing = MapL m1 +Nothing ^ Just m2 = MapR 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 +mapLR :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a +mapLR 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) +singletonL k a = MapL (singletonM k a) singletonR :: (TrieKey k1, TrieKey k2, Sized a) => k2 -> a -> TrieMap (Either k1 k2) a -singletonR k a = K2 (singletonM k a) +singletonR k a = MapR (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) +{-# INLINE uView #-} 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 (MapL m1) = UView (Just m1) Nothing +uView (MapR 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 (HoleXR hole1 m2) = Hole1 hole1 (Just m2) hView (Hole0X hole2) = Hole2 Nothing hole2 -hView (Hole1X m1 hole2) = Hole2 (Just m1) hole2 +hView (HoleLX 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 +hole1 hole1 (Just m2) = HoleXR 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 +hole2 (Just m1) hole2 = HoleLX m1 hole2 #define UVIEW uView -> UView +#define CONTEXT(cl) (TrieKey k1, TrieKey k2, cl (TrieMap k1), cl (TrieMap k2)) -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 CONTEXT(Functor) => Functor (TrieMap (Either k1 k2)) where + fmap _ Empty = Empty + fmap f (MapL m1) = MapL (f <$> m1) + fmap f (MapR m2) = MapR (f <$> m2) + fmap f (Union s m1 m2) = Union s (f <$> m1) (f <$> 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 +instance CONTEXT(Foldable) => Foldable (TrieMap (Either k1 k2)) where + foldMap f (UVIEW m1 m2) = fmap (foldMap f) m1 `mappendM` fmap (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 CONTEXT(Traversable) => Traversable (TrieMap (Either k1 k2)) where + traverse _ Empty = pure Empty + traverse f (MapL m1) = MapL <$> traverse f m1 + traverse f (MapR m2) = MapR <$> traverse f m2 + traverse f (Union s m1 m2) = Union s <$> traverse f m1 <*> traverse f m2 + +instance CONTEXT(Subset) => Subset (TrieMap (Either k1 k2)) where + (UVIEW m11 m12) <=? (UVIEW m21 m22) + = m11 <<=? m21 && m12 <<=? m22 + +instance (TrieKey k1, TrieKey k2) => Buildable (TrieMap (Either k1 k2)) (Either k1 k2) where + type UStack (TrieMap (Either k1 k2)) = TrieMap (Either k1 k2) + uFold = defaultUFold emptyM singletonM insertWithM + type AStack (TrieMap (Either k1 k2)) = Stack (AMStack k1) (AMStack k2) + aFold f = unionFold (aFold f) (aFold f) + type DAStack (TrieMap (Either k1 k2)) = Stack (DAMStack k1) (DAMStack k2) + daFold = unionFold daFold daFold + +{-# INLINE runUView #-} +runUView :: TrieMap (Either k1 k2) a -> (Maybe (TrieMap k1 a) -> Maybe (TrieMap k2 a) -> r) -> r +runUView Empty f = inline f Nothing Nothing +runUView (MapL mL) f = inline f (Just mL) Nothing +runUView (MapR mR) f = inline f Nothing (Just mR) +runUView (Union _ mL mR) f = inline f (Just mL) (Just mR) + +instance CONTEXT(SetOp) => SetOp (TrieMap (Either k1 k2)) where + union f m1 m2 + | Empty <- m1 = m2 + | otherwise = runUView m1 (runUView m2 .: run) + where {-# INLINE run #-} + run m1L m1R m2L m2R + | Empty <- m2 = m1 + | otherwise = union (unionM f) m1L m2L ^ union (unionM f) m1R m2R + isect f m1 m2 = runUView m1 (runUView m2 .: run) where + run m1L m1R m2L m2R = isect (isectM f) m1L m2L ^ isect (isectM f) m1R m2R + diff _ m1 Empty = m1 + diff f m1 m2 = runUView m2 (runUView m1 .: run) where + run m2L m2R m1L m1R = diff (diffM f) m1L m2L ^ diff (diffM f) m1R m2R + +instance CONTEXT(Project) => Project (TrieMap (Either k1 k2)) where + mapMaybe f (UVIEW m1 m2) = mapMaybe (mapMaybeM f) m1 ^ mapMaybe (mapMaybeM f) m2 + mapEither f (UVIEW m1 m2) = (# m11 ^ m21, m12 ^ m22 #) + where !(# m11, m12 #) = mapEither (mapEitherM f) m1 + !(# m21, m22 #) = mapEither (mapEitherM f) m2 + -- | @'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) + | MapL (TrieMap k1 a) + | MapR (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) + | HoleXR (Hole k1 a) (TrieMap k2 a) | Hole0X (Hole k2 a) - | Hole1X (TrieMap k1 a) (Hole k2 a) + | HoleLX (TrieMap k1 a) (Hole k2 a) emptyM = Empty singletonM = either singletonL singletonR @@ -104,78 +133,50 @@ mSimple = maybe mzero getSimpleM sizeM Empty = 0 - sizeM (K1 m1) = sizeM m1 - sizeM (K2 m2) = sizeM m2 + sizeM (MapL m1) = sizeM m1 + sizeM (MapR 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 + lookupMC (Left k) (UVIEW (Just m1) _) = lookupMC k m1 + lookupMC (Right k) (UVIEW _ (Just m2)) = lookupMC k m2 + lookupMC _ _ = mzero 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) + Hole1 h1 __ -> guardNull (beforeM h1) ^ Nothing + Hole2 m1 h2 -> m1 ^ guardNull (beforeM h2) beforeWithM a hole = case hView hole of - Hole1 h1 __ -> K1 (beforeWithM a h1) + Hole1 h1 __ -> MapL (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) + Hole1 h1 m2 -> guardNull (afterM h1) ^ m2 + Hole2 __ h2 -> Nothing ^ guardNull (afterM h2) afterWithM a hole = case hView hole of Hole1 h1 m2 -> Just (afterWithM a h1) ^ m2 - Hole2 __ h2 -> K2 (afterWithM a h2) + Hole2 __ h2 -> MapR (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 () + indexM m i = case m of + MapL m1 -> case indexM m1 i of + (# i', a, hole1 #) -> (# i', a, HoleX0 hole1 #) + MapR m2 -> case indexM m2 i of + (# i', a, hole2 #) -> (# i', a, Hole0X hole2 #) + Union _ m1 m2 + | i <# s1, (# i', a, hole1 #) <- indexM m1 i + -> (# i', a, HoleXR hole1 m2 #) + | (# i', a, hole2 #) <- indexM m2 (i -# s1) + -> (# i', a, HoleLX m1 hole2 #) + where !s1 = sizeM# m1 + _ -> indexFail () extractHoleM (UVIEW !m1 !m2) = holes1 `mplus` holes2 where holes1 = holes extractHoleM (`hole1` m2) m1 @@ -189,19 +190,41 @@ 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 (Left k') (Right k) a = return $ HoleXR (singleHoleM k') (singletonM k a) + unifierM (Right k') (Left k) a = return $ HoleLX (singletonM k a) (singleHoleM k') unifierM (Right k') (Right k) a = Hole0X <$> unifierM k' k a + + unifyM (Left k1) a1 (Left k2) a2 = MapL <$> unifyM k1 a1 k2 a2 + unifyM (Left k1) a1 (Right k2) a2 = return $ singletonM k1 a1 `mapLR` singletonM k2 a2 + unifyM (Right k2) a2 (Left k1) a1 = return $ singletonM k1 a1 `mapLR` singletonM k2 a2 + unifyM (Right k1) a1 (Right k2) a2 = MapR <$> unifyM k1 a1 k2 a2 {-# 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) +{-# INLINE unionFold #-} +unionFold :: (TrieKey k1, TrieKey k2, Sized a) => + FromList z1 k1 a -> FromList z2 k2 a -> FromList (Stack z1 z2) (Either k1 k2) a +unionFold Foldl{snoc = snocL, begin = beginL, done = doneL} + Foldl{snoc = snocR, begin = beginR, done = doneR} + = Foldl{zero = Empty, ..} + where snoc (JustL s1) (Left k) a = JustL (snocL s1 k a) + snoc (JustL s1) (Right k) a = Both s1 (beginR k a) + snoc (JustR s2) (Left k) a = Both (beginL k a) s2 + snoc (JustR s2) (Right k) a = JustR (snocR s2 k a) + snoc (Both s1 s2) (Left k) a = Both (snocL s1 k a) s2 + snoc (Both s1 s2) (Right k) a = Both s1 (snocR s2 k a) + + begin (Left k) a = JustL (beginL k a) + begin (Right k) a = JustR (beginR k a) + + done (JustL sL) = MapL (doneL sL) + done (JustR sR) = MapR (doneR sR) + done (Both sL sR) = doneL sL `mapLR` doneR sR -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) +data Stack s1 s2 a = + JustL (s1 a) + | JustR (s2 a) + | Both (s1 a) (s2 a) diff --git a/Data/TrieMap/UnitMap.hs b/Data/TrieMap/UnitMap.hs --- a/Data/TrieMap/UnitMap.hs +++ b/Data/TrieMap/UnitMap.hs @@ -1,26 +1,52 @@ -{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash, FlexibleInstances #-} - +{-# LANGUAGE TypeFamilies, FlexibleInstances, CPP, MultiParamTypeClasses, UnboxedTuples #-} module Data.TrieMap.UnitMap () where +import Data.Maybe (fromMaybe) import Data.TrieMap.TrieKey -import Data.TrieMap.Sized -import Data.Functor -import Control.Monad - -import Data.Foldable -import Data.Traversable -import Data.Maybe - import Prelude hiding (foldr, foldl, foldr1, foldl1) +instance Functor (TrieMap ()) where + fmap f (Unit m) = Unit (f <$> m) + instance Foldable (TrieMap ()) where foldMap f (Unit m) = foldMap f m foldr f z (Unit m) = foldr f z m foldl f z (Unit m) = foldl f z m - foldr1 f (Unit m) = foldr1 f m - foldl1 f (Unit m) = foldl1 f m +instance Traversable (TrieMap ()) where + traverse f (Unit (Just a)) = Unit . Just <$> f a + traverse _ _ = pure (Unit Nothing) + +instance Subset (TrieMap ()) where + Unit m1 <=? Unit m2 = m1 <=? m2 + +instance Buildable (TrieMap ()) () where + type UStack (TrieMap ()) = Elem + uFold f = Foldl{ + zero = emptyM, + begin = const Elem, + snoc = \ (Elem a) _ a' -> Elem (f a' a), + done = \ (Elem a) -> single a} + type AStack (TrieMap ()) = Elem + aFold = uFold + type DAStack (TrieMap ()) = TrieMap () + daFold = Foldl{ + zero = emptyM, + begin = const single, + snoc = error "Error: duplicate keys", + done = id} + +#define SETOP(op) op f (Unit m1) (Unit m2) = Unit (op f m1 m2) +instance SetOp (TrieMap ()) where + SETOP(union) + SETOP(isect) + SETOP(diff) + +instance Project (TrieMap ()) where + mapMaybe f (Unit m) = Unit (mapMaybe f m) + mapEither f (Unit m) = both Unit (mapEither f) m + -- | @'TrieMap' () a@ is implemented as @'Maybe' a@. instance TrieKey () where newtype TrieMap () a = Unit (Maybe a) @@ -30,19 +56,10 @@ singletonM _ = single getSimpleM (Unit m) = maybe Null Singleton m sizeM (Unit m) = getSize m - lookupM _ (Unit m) = liftMaybe m - traverseM f (Unit m) = Unit <$> traverse f m - fmapM f (Unit m) = Unit (f <$> m) - mapMaybeM f (Unit m) = Unit (m >>= f) - mapEitherM f (Unit a) = both Unit Unit (mapEitherMaybe f) a - unionM f (Unit m1) (Unit m2) = Unit (unionMaybe f m1 m2) - isectM f (Unit m1) (Unit m2) = Unit (isectMaybe f m1 m2) - diffM f (Unit m1) (Unit m2) = Unit (diffMaybe f m1 m2) - isSubmapM (<=) (Unit m1) (Unit m2) = subMaybe (<=) m1 m2 + lookupMC _ (Unit (Just a)) = return a + lookupMC _ _ = mzero insertWithM f _ a (Unit m) = Unit (Just (maybe a f m)) - fromListM _ [] = emptyM - fromListM f ((_, v):xs) = single (foldl (\ v' -> f v' . snd) v xs) singleHoleM _ = Hole beforeM _ = emptyM @@ -53,10 +70,11 @@ searchMC _ (Unit (Just v)) _ g = g v Hole searchMC _ _ f _ = f Hole - indexM i (Unit (Just v)) = (# i, v, Hole #) - indexM _ _ = indexFail () + indexM (Unit v) i = + (# i, fromMaybe indexFail v, Hole #) - unifierM _ _ _ = Nothing + unifierM _ _ _ = mzero + unifyM _ _ _ _ = mzero extractHoleM (Unit (Just v)) = return (v, Hole) extractHoleM _ = mzero diff --git a/Data/TrieMap/Utils.hs b/Data/TrieMap/Utils.hs --- a/Data/TrieMap/Utils.hs +++ b/Data/TrieMap/Utils.hs @@ -1,6 +1,8 @@ -{-# LANGUAGE Rank2Types, BangPatterns, MagicHash #-} +{-# LANGUAGE Rank2Types, BangPatterns, MagicHash, TypeOperators #-} module Data.TrieMap.Utils where +import Control.Monad.Unpack + import Data.Bits import qualified Data.Foldable @@ -9,6 +11,10 @@ import GHC.Exts +{-# INLINE mapInput #-} +mapInput :: (Unpackable a, Unpackable b) => (a -> b) -> (b :~> c) -> (a :~> c) +mapInput f func = unpack $ \ a -> func $~ f a + {-# INLINE toVectorN #-} toVectorN :: Vector v a => (forall b . (a -> b -> b) -> b -> f -> b) -> (f -> Int) -> f -> v a toVectorN fold size xs = create $ do @@ -46,5 +52,6 @@ {-# RULES "or 0" forall w# . or# w# 0## = w#; "0 or" forall w# . or# 0## w# = w#; + "shiftL 0" forall w# . uncheckedShiftL# w# 0# = w#; "plusAddr 0" forall a# . plusAddr# a# 0# = a#; #-} diff --git a/Data/TrieMap/WordMap.hs b/Data/TrieMap/WordMap.hs --- a/Data/TrieMap/WordMap.hs +++ b/Data/TrieMap/WordMap.hs @@ -1,19 +1,17 @@ -{-# LANGUAGE UnboxedTuples, BangPatterns, TypeFamilies, PatternGuards, MagicHash, CPP, NamedFieldPuns, FlexibleInstances #-} -{-# OPTIONS -funbox-strict-fields #-} -module Data.TrieMap.WordMap (SNode, WHole, TrieMap(WordMap), Hole(Hole), getWordMap, getHole) where +{-# LANGUAGE UnboxedTuples, BangPatterns, TypeFamilies, PatternGuards, MagicHash, CPP, NamedFieldPuns, FlexibleInstances, RecordWildCards #-} +{-# LANGUAGE TemplateHaskell, TypeOperators, MultiParamTypeClasses #-} +{-# OPTIONS -funbox-strict-fields -O -fspec-constr -fliberate-case -fstatic-argument-transformation #-} +module Data.TrieMap.WordMap (SNode, WHole, TrieMap(WordMap), Hole(Hole), WordStack, getWordMap, getHole) where import Data.TrieMap.TrieKey import Data.TrieMap.Sized import Control.Exception (assert) -import Control.Applicative (Applicative(..), (<$>)) -import Control.Monad hiding (join) +import Control.Monad.Lookup +import Control.Monad.Unpack import Data.Bits -import Data.Foldable import Data.Maybe hiding (mapMaybe) -import Data.Monoid -import Data.TrieMap.Utils import GHC.Exts @@ -38,7 +36,6 @@ data SNode a = SNode {sz :: !Size, node :: (Node a)} {-# ANN type SNode ForceSpecConstr #-} data Node a = Nil | Tip !Key a | Bin !Prefix !Mask !(SNode a) !(SNode a) -{-# ANN type Node ForceSpecConstr #-} instance Sized (SNode a) where getSize# SNode{sz} = unbox sz @@ -54,13 +51,57 @@ sNode !n = SNode (getSize n) n data WHole a = WHole !Key (Path a) +{-# ANN type WHole ForceSpecConstr #-} +$(noUnpackInstance ''Path) +$(noUnpackInstance ''Node) +$(unpackInstance ''WHole) +$(unpackInstance ''SNode) + {-# INLINE hole #-} hole :: Key -> Path a -> Hole Word a hole k path = Hole (WHole k path) #define HOLE(args) (Hole (WHole args)) +instance Subset (TrieMap Word) where + WordMap m1 <=? WordMap m2 = m1 <=? m2 + +instance Functor (TrieMap Word) where + fmap f (WordMap m) = WordMap (f <$> m) + +instance Foldable (TrieMap Word) where + foldMap f (WordMap m) = foldMap f m + foldr f z (WordMap m) = foldr f z m + foldl f z (WordMap m) = foldl f z m + foldr1 f (WordMap m) = foldr1 f m + foldl1 f (WordMap m) = foldl1 f m + +instance Traversable (TrieMap Word) where + traverse f (WordMap m) = WordMap <$> traverse f m + +instance Buildable (TrieMap Word) Word where + type UStack (TrieMap Word) = SNode + {-# INLINE uFold #-} + uFold = fmap WordMap . defaultUFold nil singleton (\ f k a -> insertWithC f k (getSize a) a) + type AStack (TrieMap Word) = WordStack + {-# INLINE aFold #-} + aFold = fmap WordMap . fromAscList + type DAStack (TrieMap Word) = WordStack + {-# INLINE daFold #-} + daFold = aFold const + +#define SETOP(op) op f (WordMap m1) (WordMap m2) = WordMap (op f m1 m2) + +instance SetOp (TrieMap Word) where + SETOP(union) + SETOP(isect) + SETOP(diff) + +instance Project (TrieMap Word) where + mapMaybe f (WordMap m) = WordMap $ mapMaybe f m + mapEither f (WordMap m) = both WordMap (mapEither f) m + -- | @'TrieMap' 'Word' a@ is based on "Data.IntMap". instance TrieKey Word where newtype TrieMap Word a = WordMap {getWordMap :: SNode a} @@ -72,92 +113,126 @@ Tip _ a -> Singleton a _ -> NonSimple sizeM (WordMap t) = getSize t - lookupM k (WordMap m) = lookup k m - traverseM f (WordMap m) = WordMap <$> traverse f m - fmapM f (WordMap m) = WordMap (map f m) - mapMaybeM f (WordMap m) = WordMap (mapMaybe f m) - mapEitherM f (WordMap m) = both WordMap WordMap (mapEither f) m - unionM f (WordMap m1) (WordMap m2) = WordMap (unionWith f m1 m2) - isectM f (WordMap m1) (WordMap m2) = WordMap (intersectionWith f m1 m2) - diffM f (WordMap m1) (WordMap m2) = WordMap (differenceWith f m1 m2) - isSubmapM (<=) (WordMap m1) (WordMap m2) = isSubmapOfBy (<=) m1 m2 + lookupMC k (WordMap m) = lookupC k m singleHoleM k = hole k Root - beforeM HOLE(_ path) = WordMap (before nil path) - beforeWithM a HOLE(k path) = WordMap (before (singleton k a) path) - afterM HOLE(_ path) = WordMap (after nil path) - afterWithM a HOLE(k path) = WordMap (after (singleton k a) path) + beforeM HOLE(_ path) = WordMap (before path) + beforeWithM a HOLE(k path) = WordMap (beforeWith (singleton k a) path) + afterM HOLE(_ path) = WordMap (after path) + afterWithM a HOLE(k path) = WordMap (afterWith (singleton k a) path) {-# INLINE searchMC #-} - searchMC !k (WordMap t) = mapSearch (hole k) (searchC k t) - indexM i (WordMap m) = indexT i m Root where - indexT !i TIP(kx x) path = (# i, x, hole kx path #) - indexT !i BIN(p m l r) path - | i < sl = indexT i l (LeftBin p m path r) - | otherwise = indexT (i - sl) r (RightBin p m l path) - where !sl = getSize l - indexT _ NIL _ = indexFail () + searchMC !k (WordMap t) notfound found = searchC k t (unpack (notfound . Hole)) (\ a -> unpack (found a . Hole)) + {-# INLINE indexM #-} + indexM (WordMap m) i = index i m extractHoleM (WordMap m) = extractHole Root m where extractHole _ (SNode _ Nil) = mzero extractHole path TIP(kx x) = return (x, hole kx path) extractHole path BIN(p m l r) = extractHole (LeftBin p m path r) l `mplus` extractHole (RightBin p m l path) r - clearM HOLE(_ path) = WordMap (assign nil path) + {-# INLINE clearM #-} + clearM HOLE(_ path) = case clear path of + (# sz#, node #) -> WordMap SNode{sz = I# sz#, node} {-# INLINE assignM #-} - assignM v HOLE(kx path) = WordMap (assign (singleton kx v) path) + assignM v HOLE(kx path) = case assign (singleton kx v) path of + (# sz#, node #) -> WordMap SNode{sz = I# sz#, node} + {-# INLINE unifyM #-} + unifyM k1 a1 k2 a2 = WordMap <$> unify k1 a1 k2 a2 + {-# INLINE unifierM #-} unifierM k' k a = Hole <$> unifier k' k a -{-# INLINE searchC #-} -searchC :: Key -> SNode a -> SearchCont (Path a) a r + {-# INLINE insertWithM #-} + insertWithM f k a (WordMap m) = WordMap (insertWithC f k (getSize a) a m) + +insertWithC :: Sized a => (a -> a) -> Key -> Int -> a -> SNode a -> SNode a +insertWithC f !k !szA a !t = ins' t where + {-# INLINE tip #-} + tip = SNode {sz = szA, node = Tip k a} + + {-# INLINE out #-} + out SNode{sz = I# sz#, node} = (# sz#, node #) + {-# INLINE ins' #-} + ins' t = case ins t of + (# sz#, node #) -> SNode{sz = I# sz#, node} + ins !t = case t of + BIN(p m l r) + | nomatch k p m -> out $ join k tip p t + | mask0 k m -> out $ bin' p m (ins' l) r + | otherwise -> out $ bin' p m l (ins' r) + TIP(kx x) + | k == kx -> out $ singleton kx (f x) + | otherwise -> out $ join k tip kx t + NIL -> out tip + +index :: Int# -> SNode a -> (# Int#, a, Hole Word a #) +index i !t = indexT i t Root where + indexT i TIP(kx x) path = (# i, x, hole kx path #) + indexT i BIN(p m l r) path + | i <# sl = indexT i l (LeftBin p m path r) + | otherwise = indexT (i -# sl) r (RightBin p m l path) + where !sl = getSize# l + indexT _ NIL _ = indexFail () + +searchC :: Key -> SNode a -> (WHole a :~> r) -> (a -> WHole a :~> r) -> r searchC !k t notfound found = seek Root t where seek path t@BIN(p m l r) - | nomatch k p m = notfound (branchHole k p path t) - | zero k m + | nomatch k p m = notfound $~ WHole k (branchHole k p path t) + | mask0 k m = seek (LeftBin p m path r) l | otherwise = seek (RightBin p m l path) r seek path t@TIP(ky y) - | k == ky = found y path - | otherwise = notfound (branchHole k ky path t) - seek path NIL = notfound path + | k == ky = found y $~ WHole k path + | otherwise = notfound $~ WHole k (branchHole k ky path t) + seek path NIL = notfound $~ WHole k path -before, after :: SNode a -> Path a -> SNode a -before !t Root = t -before !t (LeftBin _ _ path _) = before t path -before !t (RightBin p m l path) = before (bin p m l t) path -after !t Root = t -after !t (RightBin _ _ _ path) = after t path -after !t (LeftBin p m path r) = after (bin p m t r) path +before, after :: Path a -> SNode a +beforeWith, afterWith :: SNode a -> Path a -> SNode a -assign :: Sized a => SNode a -> Path a -> SNode a -assign NIL Root = nil -assign NIL (LeftBin _ _ path r) = assign' r path -assign NIL (RightBin _ _ l path) = assign' l path -assign t Root = t -assign t (LeftBin p m path r) = assign' (bin' p m t r) path -assign t (RightBin p m l path) = assign' (bin' p m l t) path +before Root = nil +before (LeftBin _ _ path _) = before path +before (RightBin _ _ l path) = beforeWith l path -assign' :: Sized a => SNode a -> Path a -> SNode a -assign' !t Root = t -assign' !t (LeftBin p m path r) = assign' (bin' p m t r) path -assign' !t (RightBin p m l path) = assign' (bin' p m l t) path +beforeWith !t Root = t +beforeWith !t (LeftBin _ _ path _) = beforeWith t path +beforeWith !t (RightBin p m l path) = beforeWith (bin' p m l t) path +after Root = nil +after (RightBin _ _ _ path) = after path +after (LeftBin _ _ path r) = afterWith r path + +afterWith !t Root = t +afterWith !t (RightBin _ _ _ path) = afterWith t path +afterWith !t (LeftBin p m path r) = afterWith (bin' p m t r) path + +clear :: Path a -> (# Int#, Node a #) +assign :: SNode a -> Path a -> (# Int#, Node a #) +clear Root = (# 0#, Nil #) +clear (LeftBin _ _ path r) = assign r path +clear (RightBin _ _ l path) = assign l path + +assign SNode{sz = I# sz#, node} Root = (# sz#, node #) +assign !t (LeftBin p m path r) = assign (bin' p m t r) path +assign !t (RightBin p m l path) = assign (bin' p m l t) path + branchHole :: Key -> Prefix -> Path a -> SNode a -> Path a branchHole !k !p path t - | zero k m = LeftBin p' m path t + | mask0 k m = LeftBin p' m path t | otherwise = RightBin p' m t path where m = branchMask k p p' = mask k m -lookup :: Key -> SNode a -> Lookup a -lookup !k = look where - look BIN(_ m l r) = look (if zeroN k m then l else r) +{-# INLINE lookupC #-} +lookupC :: Key -> SNode a -> Lookup r a +lookupC !k !t = Lookup $ \ no yes -> let + look BIN(_ m l r) = if zeroN k m then look l else look r look TIP(kx x) - | k == kx = some x - look _ = none + | k == kx = yes x + look _ = no + in look t singleton :: Sized a => Key -> a -> SNode a singleton k a = sNode (Tip k a) @@ -165,133 +240,126 @@ singletonMaybe :: Sized a => Key -> Maybe a -> SNode a singletonMaybe k = maybe nil (singleton k) -traverse :: (Applicative f, Sized b) => (a -> f b) -> SNode a -> f (SNode b) -traverse f = trav where - trav NIL = pure nil - trav TIP(kx x) = singleton kx <$> f x - trav BIN(p m l r) = bin' p m <$> trav l <*> trav r +instance Functor SNode where + fmap f = map where + map SNode{sz, node} = SNode sz $ case node of + Nil -> Nil + Tip k x -> Tip k (f x) + Bin p m l r -> Bin p m (map l) (map r) instance Foldable SNode where - foldMap _ NIL = mempty - foldMap f TIP(_ x) = f x - foldMap f BIN(_ _ l r) = foldMap f l `mappend` foldMap f r - - foldr f z BIN(_ _ l r) = foldr f (foldr f z r) l - foldr f z TIP(_ x) = f x z - foldr _ z NIL = z - - foldl f z BIN(_ _ l r) = foldl f (foldl f z l) r - foldl f z TIP(_ x) = f z x - foldl _ z NIL = z + foldMap f = fold where + fold NIL = mempty + fold TIP(_ x) = f x + fold BIN(_ _ l r) = fold l `mappend` fold r - foldr1 _ NIL = foldr1Empty - foldr1 _ TIP(_ x) = x - foldr1 f BIN(_ _ l r) = foldr f (foldr1 f r) l + foldr f = flip fold where + fold BIN(_ _ l r) z = fold l (fold r z) + fold TIP(_ x) z = f x z + fold NIL z = z - foldl1 _ NIL = foldl1Empty - foldl1 _ TIP(_ x) = x - foldl1 f BIN(_ _ l r) = foldl f (foldl1 f l) r + foldl f = fold where + fold z BIN(_ _ l r) = fold (fold z l) r + fold z TIP(_ x) = f z x + fold z NIL = z -instance Foldable (TrieMap Word) where - foldMap f (WordMap m) = foldMap f m - foldr f z (WordMap m) = foldr f z m - foldl f z (WordMap m) = foldl f z m - foldr1 f (WordMap m) = foldr1 f m - foldl1 f (WordMap m) = foldl1 f m +instance Traversable SNode where + traverse f = trav where + trav NIL = pure nil + trav SNode{sz, node = Tip kx x} + = SNode sz . Tip kx <$> f x + trav SNode{sz, node = Bin p m l r} + = SNode sz .: Bin p m <$> trav l <*> trav r -map :: Sized b => (a -> b) -> SNode a -> SNode b -map f BIN(p m l r) = bin' p m (map f l) (map f r) -map f TIP(kx x) = singleton kx (f x) -map _ _ = nil +instance Subset SNode where + (<=?) = subMap where + t1@BIN(p1 m1 l1 r1) `subMap` BIN(p2 m2 l2 r2) + | shorter m1 m2 = False + | shorter m2 m1 = match p1 p2 m2 && (if mask0 p1 m2 then t1 `subMap` l2 + else t1 `subMap` r2) + | otherwise = (p1==p2) && l1 `subMap` l2 && r1 `subMap` r2 + BIN({}) `subMap` _ = False + TIP(k x) `subMap` t2 = runLookup (lookupC k t2) False (x (a -> Maybe b) -> SNode a -> SNode b -mapMaybe f BIN(p m l r) = bin p m (mapMaybe f l) (mapMaybe f r) -mapMaybe f TIP(kx x) = singletonMaybe kx (f x) -mapMaybe _ _ = nil +instance SetOp SNode where + union f = (\/) where + n1@(SNode _ t1) \/ n2@(SNode _ t2) = case (t1, t2) of + (Nil, _) -> n2 + (_, Nil) -> n1 + (Tip k x, _) -> alter (maybe (Just x) (f x)) k n2 + (_, Tip k x) -> alter (maybe (Just x) (`f` x)) k n1 + (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) + | shorter m1 m2 -> union1 + | shorter m2 m1 -> union2 + | p1 == p2 -> bin p1 m1 (l1 \/ l2) (r1 \/ r2) + | otherwise -> join p1 n1 p2 n2 + where + union1 | nomatch p2 p1 m1 = join p1 n1 p2 n2 + | mask0 p2 m1 = bin p1 m1 (l1 \/ n2) r1 + | otherwise = bin p1 m1 l1 (r1 \/ n2) -mapEither :: (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> - SNode a -> (# SNode b, SNode c #) -mapEither f BIN(p m l r) = both (bin p m lL) (bin p m lR) (mapEither f) r - where !(# lL, lR #) = mapEither f l -mapEither f TIP(kx x) = both (singletonMaybe kx) (singletonMaybe kx) f x -mapEither _ _ = (# nil, nil #) + union2 | nomatch p1 p2 m2 = join p1 n1 p2 n2 + | mask0 p1 m2 = bin p2 m2 (n1 \/ l2) r2 + | otherwise = bin p2 m2 l2 (n1 \/ r2) + isect f = (/\) where + n1@(SNode _ t1) /\ n2@(SNode _ t2) = case (t1, t2) of + (Nil, _) -> nil + (Tip{}, Nil) -> nil + (Bin{}, Nil) -> nil + (Tip k x, _) -> runLookup (lookupC k n2) nil (singletonMaybe k . f x) + (_, Tip k y) -> runLookup (lookupC k n1) nil (singletonMaybe k . flip f y) + (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) + | shorter m1 m2 -> intersection1 + | shorter m2 m1 -> intersection2 + | p1 == p2 -> bin p1 m1 (l1 /\ l2) (r1 /\ r2) + | otherwise -> nil + where + intersection1 | nomatch p2 p1 m1 = nil + | mask0 p2 m1 = l1 /\ n2 + | otherwise = r1 /\ n2 -unionWith :: Sized a => (a -> a -> Maybe a) -> SNode a -> SNode a -> SNode a -unionWith f n1@(SNode _ t1) n2@(SNode _ t2) = case (t1, t2) of - (Nil, _) -> n2 - (_, Nil) -> n1 - (Tip k x, _) -> alter (maybe (Just x) (f x)) k n2 - (_, Tip k x) -> alter (maybe (Just x) (`f` x)) k n1 - (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) - | shorter m1 m2 -> union1 - | shorter m2 m1 -> union2 - | p1 == p2 -> bin p1 m1 (unionWith f l1 l2) (unionWith f r1 r2) - | otherwise -> join p1 n1 p2 n2 - where - union1 | nomatch p2 p1 m1 = join p1 n1 p2 n2 - | zero p2 m1 = bin p1 m1 (unionWith f l1 n2) r1 - | otherwise = bin p1 m1 l1 (unionWith f r1 n2) + intersection2 | nomatch p1 p2 m2 = nil + | mask0 p1 m2 = n1 /\ l2 + | otherwise = n1 /\ r2 + diff f = (\\) where + n1@(SNode _ t1) \\ n2@(SNode _ t2) = case (t1, t2) of + (Nil, _) -> nil + (_, Nil) -> n1 + (Tip k x, _) -> runLookup (lookupC k n2) n1 (singletonMaybe k . f x) + (_, Tip k y) -> alter (>>= flip f y) k n1 + (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) + | shorter m1 m2 -> difference1 + | shorter m2 m1 -> difference2 + | p1 == p2 -> bin p1 m1 (l1 \\ l2) (r1 \\ r2) + | otherwise -> n1 + where + difference1 | nomatch p2 p1 m1 = n1 + | mask0 p2 m1 = bin p1 m1 (l1 \\ n2) r1 + | otherwise = bin p1 m1 l1 (r1 \\ n2) - union2 | nomatch p1 p2 m2 = join p1 n1 p2 n2 - | zero p1 m2 = bin p2 m2 (unionWith f n1 l2) r2 - | otherwise = bin p2 m2 l2 (unionWith f n1 r2) + difference2 | nomatch p1 p2 m2 = n1 + | mask0 p1 m2 = n1 \\ l2 + | otherwise = n1 \\ r2 +instance Project SNode where + mapMaybe f = mMaybe where + mMaybe BIN(p m l r) = bin p m (mMaybe l) (mMaybe r) + mMaybe TIP(kx x) = singletonMaybe kx (f x) + mMaybe NIL = nil + mapEither f = mEither where + mEither BIN(p m l r) = (# bin p m l1 r1, bin p m l2 r2 #) + where !(# l1, l2 #) = mEither l + !(# r1, r2 #) = mEither r + mEither TIP(kx x) = both (singletonMaybe kx) f x + mEither NIL = (# nil, nil #) + {-# INLINE alter #-} alter :: Sized a => (Maybe a -> Maybe a) -> Key -> SNode a -> SNode a alter f k t = getWordMap $ alterM f k (WordMap t) -intersectionWith :: Sized c => (a -> b -> Maybe c) -> SNode a -> SNode b -> SNode c -intersectionWith f n1@(SNode _ t1) n2@(SNode _ t2) = case (t1, t2) of - (Nil, _) -> nil - (_, Nil) -> nil - (Tip k x, _) -> option (lookup k n2) nil (singletonMaybe k . f x) - (_, Tip k y) -> option (lookup k n1) nil (singletonMaybe k . flip f y) - (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) - | shorter m1 m2 -> intersection1 - | shorter m2 m1 -> intersection2 - | p1 == p2 -> bin p1 m1 (intersectionWith f l1 l2) (intersectionWith f r1 r2) - | otherwise -> nil - where - intersection1 | nomatch p2 p1 m1 = nil - | zero p2 m1 = intersectionWith f l1 n2 - | otherwise = intersectionWith f r1 n2 - - intersection2 | nomatch p1 p2 m2 = nil - | zero p1 m2 = intersectionWith f n1 l2 - | otherwise = intersectionWith f n1 r2 - -differenceWith :: Sized a => (a -> b -> Maybe a) -> SNode a -> SNode b -> SNode a -differenceWith f n1@(SNode _ t1) n2@(SNode _ t2) = case (t1, t2) of - (Nil, _) -> nil - (_, Nil) -> n1 - (Tip k x, _) -> option (lookup k n2) n1 (singletonMaybe k . f x) - (_, Tip k y) -> alter (>>= flip f y) k n1 - (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) - | shorter m1 m2 -> difference1 - | shorter m2 m1 -> difference2 - | p1 == p2 -> bin p1 m1 (differenceWith f l1 l2) (differenceWith f r1 r2) - | otherwise -> n1 - where - difference1 | nomatch p2 p1 m1 = n1 - | zero p2 m1 = bin p1 m1 (differenceWith f l1 n2) r1 - | otherwise = bin p1 m1 l1 (differenceWith f r1 n2) - - difference2 | nomatch p1 p2 m2 = n1 - | zero p1 m2 = differenceWith f n1 l2 - | otherwise = differenceWith f n1 r2 - -isSubmapOfBy :: LEq a b -> LEq (SNode a) (SNode b) -isSubmapOfBy (<=) t1@BIN(p1 m1 l1 r1) BIN(p2 m2 l2 r2) - | shorter m1 m2 = False - | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy (<=) t1 l2 - else isSubmapOfBy (<=) t1 r2) - | otherwise = (p1==p2) && isSubmapOfBy (<=) l1 l2 && isSubmapOfBy (<=) r1 r2 -isSubmapOfBy _ BIN(_ _ _ _) _ = False -isSubmapOfBy (<=) TIP(k x) t2 = option (lookup k t2) False (x <=) -isSubmapOfBy _ NIL _ = True - -zero :: Key -> Mask -> Bool -zero i m +mask0 :: Key -> Mask -> Bool +mask0 i m = i .&. m == 0 nomatch,match :: Key -> Prefix -> Mask -> Bool @@ -306,7 +374,7 @@ mask :: Nat -> Nat -> Prefix mask i m - = i .&. compl ((m-1) .|. m) + = i .&. compl ((m-1) `xor` m) shorter :: Mask -> Mask -> Bool shorter m1 m2 @@ -329,11 +397,12 @@ {-# INLINE join #-} join :: Prefix -> SNode a -> Prefix -> SNode a -> SNode a join p1 t1 p2 t2 - | zero p1 m = bin' p m t1 t2 - | otherwise = bin' p m t2 t1 + | mask0 p1 m = SNode{sz = sz', node = Bin p m t1 t2} + | otherwise = SNode{sz = sz', node = Bin p m t2 t1} where m = branchMask p1 p2 p = mask p1 m + sz' = sz t1 + sz t2 nil :: SNode a nil = SNode 0 Nil @@ -349,8 +418,43 @@ where nonempty NIL = False nonempty _ = True +{-# INLINE unify #-} +unify :: Sized a => Key -> a -> Key -> a -> Lookup r (SNode a) +unify k1 a1 k2 a2 = Lookup $ \ no yes -> + if k1 == k2 then no else yes (join k1 (singleton k1 a1) k2 (singleton k2 a2)) + {-# INLINE unifier #-} -unifier :: Sized a => Key -> Key -> a -> Maybe (WHole a) -unifier k' k a - | k' == k = Nothing - | otherwise = Just (WHole k' $ branchHole k' k Root (singleton k a)) +unifier :: Sized a => Key -> Key -> a -> Lookup r (WHole a) +unifier k' k a = Lookup $ \ no yes -> + if k == k' then no else yes (WHole k' $ branchHole k' k Root (singleton k a)) + +{-# INLINE fromAscList #-} +fromAscList :: Sized a => (a -> a -> a) -> Foldl WordStack Key a (SNode a) +fromAscList f = Foldl{zero = nil, ..} where + begin kx vx = WordStack kx vx Nada + + snoc (WordStack kx vx stk) kz vz + | kx == kz = WordStack kx (f vz vx) stk + | otherwise = WordStack kz vz $ reduce (branchMask kx kz) kx (singleton kx vx) stk + +-- reduce :: Mask -> Prefix -> SNode a -> Stack a -> Stack a + reduce !m !px !tx (Push py ty stk') + | shorter m mxy = reduce m pxy (bin' pxy mxy ty tx) stk' + where mxy = branchMask px py; pxy = mask px mxy + reduce _ px tx stk = Push px tx stk + + done (WordStack kx vx stk) = case finish kx (singleton kx vx) stk of + (# sz#, node #) -> SNode {sz = I# sz#, node} + + finish !px !tx (Push py ty stk) = finish p (join' py ty px tx) stk + where m = branchMask px py; p = mask px m + finish _ SNode{sz, node} Nada = (# unbox sz, node #) + + join' p1 t1 p2 t2 + = SNode{sz = sz t1 + sz t2, node = Bin p m t1 t2} + where + m = branchMask p1 p2 + p = mask p1 m + +data WordStack a = WordStack !Key a (Stack a) +data Stack a = Push !Prefix !(SNode a) !(Stack a) | Nada diff --git a/Data/TrieSet.hs b/Data/TrieSet.hs --- a/Data/TrieSet.hs +++ b/Data/TrieSet.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UnboxedTuples, ImplicitParams, RecordWildCards, FlexibleContexts #-} module Data.TrieSet ( -- * Set type TSet, @@ -41,6 +41,10 @@ deleteFindMax, minView, maxView, + -- * Index + elemAt, + deleteAt, + lookupIndex, -- * Conversion -- ** Map mapSet, @@ -48,22 +52,36 @@ elems, toList, fromList, + -- ** Vector + toVector, + fromVector, -- ** Ordered lists toAscList, fromAscList, - fromDistinctAscList) - where + fromDistinctAscList, + -- ** Ordered vectors + fromAscVector, + fromDistinctAscVector) + where +import Control.Monad +import Control.Monad.Ends +import Control.Monad.Lookup + import Data.TrieMap.Class import Data.TrieMap.Class.Instances () -import Data.TrieMap.TrieKey +import Data.TrieMap.TrieKey hiding (foldr, foldl, toList, union, diff, isect) +import qualified Data.TrieMap.TrieKey.SetOp as Set import Data.TrieMap.Representation.Class -import Data.TrieMap.Sized -import Data.TrieMap.Utils -import Control.Monad.Ends +import Data.Vector.Build +import qualified Data.Vector.Generic as G +import Data.Vector.Fusion.Util (unId) +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Monadic as S -import Data.Maybe +import Data.Maybe(fromJust) + import qualified Data.Foldable as F import Data.Monoid (Monoid (..)) @@ -102,29 +120,29 @@ -- | The union of two 'TSet's, preferring the first set when -- equal elements are encountered. union :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `union` TSet s2 = TSet (unionM (const . Just) s1 s2) +TSet s1 `union` TSet s2 = TSet (Set.union (const . Just) s1 s2) -- | The symmetric difference of two 'TSet's. symmetricDifference :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `symmetricDifference` TSet s2 = TSet (unionM (\ _ _ -> Nothing) s1 s2) +TSet s1 `symmetricDifference` TSet s2 = TSet (Set.union (\ _ _ -> Nothing) s1 s2) -- | Difference of two 'TSet's. difference :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `difference` TSet s2 = TSet (diffM (\ _ _ -> Nothing) s1 s2) +TSet s1 `difference` TSet s2 = TSet (Set.diff (\ _ _ -> Nothing) s1 s2) -- | Intersection of two 'TSet's. Elements of the result come from the first set. intersection :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `intersection` TSet s2 = TSet (isectM (const . Just) s1 s2) +TSet s1 `intersection` TSet s2 = TSet (Set.isect (const . Just) s1 s2) -- | Filter all elements that satisfy the predicate. filter :: TKey a => (a -> Bool) -> TSet a -> TSet a -filter p (TSet s) = TSet (mapMaybeM (\ (Elem a) -> if p a then Just (Elem a) else Nothing) s) +filter p (TSet s) = TSet (mapMaybe (\ (Elem a) -> if p a then return (Elem a) else mzero) s) -- | Partition the set into two sets, one with all elements that satisfy -- the predicate and one with all elements that don't satisfy the predicate. -- See also 'split'. partition :: TKey a => (a -> Bool) -> TSet a -> (TSet a, TSet a) -partition p (TSet s) = case mapEitherM f s of +partition p (TSet s) = case mapEither f s of (# s1, s2 #) -> (TSet s1, TSet s2) where f e@(Elem a) | p a = (# Just e, Nothing #) @@ -153,7 +171,7 @@ map f s = fromList [f x | x <- elems s] -- | --- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is monotonic. +-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly monotonic. -- /The precondition is not checked./ -- Semi-formally, we have: -- @@ -161,7 +179,7 @@ -- > ==> mapMonotonic f s == map f s -- > where ls = toList s mapMonotonic :: (TKey a, TKey b) => (a -> b) -> TSet a -> TSet b -mapMonotonic f s = fromAscList [f x | x <- toAscList s] +mapMonotonic f s = fromDistinctAscList [f x | x <- toAscList s] -- | Post-order fold. foldr :: TKey a => (a -> b -> b) -> b -> TSet a -> b @@ -226,23 +244,66 @@ toAscList :: TKey a => TSet a -> [a] toAscList s = build (\ c n -> foldr c n s) +{-# INLINE fromFoldStream #-} +fromFoldStream :: (Monad m, Repr a, TrieKey (Rep a)) => FromList z (Rep a) (Elem a) -> Stream m a -> m (TSet a) +fromFoldStream Foldl{..} (Stream suc s0 _) = run s0 where + run s = do + step <- suc s + case step of + Done -> return empty + Skip s' -> run s' + Yield x s' -> run' (begin (toRep x) (Elem x)) s' + run' stack s = do + step <- suc s + case step of + Done -> return (TSet (done stack)) + Skip s' -> run' stack s' + Yield x s' -> run' (snoc stack (toRep x) (Elem x)) s' + +{-# INLINE fromList #-} -- | Create a set from a list of elements. fromList :: TKey a => [a] -> TSet a -fromList xs = TSet (fromListM const [(toRep x, Elem x) | x <- xs]) +fromList xs = unId (fromFoldStream (uFold const) (S.fromList xs)) +{-# INLINE fromVector #-} +-- | Create a set from a vector of elements. +fromVector :: (TKey a, G.Vector v a) => v a -> TSet a +fromVector xs = unId (fromFoldStream (uFold const) (G.stream xs)) + +{-# INLINE fromAscList #-} -- | Build a set from an ascending list in linear time. -- /The precondition (input list is ascending) is not checked./ fromAscList :: TKey a => [a] -> TSet a -fromAscList xs = TSet (fromAscListM const [(toRep x, Elem x) | x <- xs]) +fromAscList xs = unId (fromFoldStream (aFold const) (S.fromList xs)) +{-# INLINE fromAscVector #-} +-- | Build a set from an ascending vector in linear time. +-- /The precondition (input vector is ascending) is not checked./ +fromAscVector :: (TKey a, G.Vector v a) => v a -> TSet a +fromAscVector xs = unId (fromFoldStream (aFold const) (G.stream xs)) + +{-# INLINE fromDistinctAscList #-} -- | /O(n)/. Build a set from an ascending list of distinct elements in linear time. -- /The precondition (input list is strictly ascending) is not checked./ fromDistinctAscList :: TKey a => [a] -> TSet a -fromDistinctAscList xs = TSet (fromDistAscListM [(toRep x, Elem x) | x <- xs]) +fromDistinctAscList xs = unId (fromFoldStream daFold (S.fromList xs)) +{-# INLINE fromDistinctAscVector #-} +-- | /O(n)/. Build a set from an ascending vector of distinct elements in linear time. +-- /The precondition (input vector is strictly ascending) is not checked./ +fromDistinctAscVector :: (TKey a, G.Vector v a) => v a -> TSet a +fromDistinctAscVector xs = unId (fromFoldStream daFold (G.stream xs)) + +{-# INLINE toVector #-} +-- | /O(n)/. Construct a vector from the elements of this set. Does not currently fuse. +toVector :: (TKey a, G.Vector v a) => TSet a -> v a +toVector (TSet s) = toVectorMapN (sizeM s) getElem s +-- If we want this to fuse, our best bet is probably a method to iterate a hole to the next key...or something. +-- This seems difficult, but perhaps not impossible. + -- | /O(1)/. Is this the empty set? null :: TKey a => TSet a -> Bool -null (TSet s) = nullM s +null (TSet s) = isNull s -- | /O(1)/. The number of elements in the set. size :: TKey a => TSet a -> Int @@ -250,7 +311,7 @@ -- | Is the element in the set? member :: TKey a => a -> TSet a -> Bool -member a (TSet s) = option (lookupM (toRep a) s) False (const True) +member a (TSet s) = runLookup (lookupMC (toRep a) s) False (const True) -- | Is the element not in the set? notMember :: TKey a => a -> TSet a -> Bool @@ -258,7 +319,7 @@ -- | Is this a subset? @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@. isSubsetOf :: TKey a => TSet a -> TSet a -> Bool -TSet s1 `isSubsetOf` TSet s2 = isSubmapM (\ _ _ -> True) s1 s2 +TSet s1 `isSubsetOf` TSet s2 = let ?le = \ _ _ -> True in s1 <=? s2 -- | Is this a proper subset? (ie. a subset but not equal). isProperSubsetOf :: TKey a => TSet a -> TSet a -> Bool @@ -271,4 +332,18 @@ {-# INLINE [1] mapSet #-} -- | Generate a 'TMap' by mapping on the elements of a 'TSet'. mapSet :: TKey a => (a -> b) -> TSet a -> TMap a b -mapSet f (TSet s) = TMap (fmapM (\ (Elem a) -> Assoc a (f a)) s) +mapSet f (TSet s) = TMap (fmap (\ (Elem a) -> Assoc a (f a)) s) + +-- | Returns the element at the specified index. Throws an error if an invalid index is specified. +elemAt :: TKey a => Int -> TSet a -> a +elemAt i (TSet s) = case indexM s (unbox i) of + (# _, Elem a, _ #) -> a + +-- | Deletes the element at the specified index. Throws an error if an invalid index is specified. +deleteAt :: TKey a => Int -> TSet a -> TSet a +deleteAt i (TSet s) = case indexM s (unbox i) of + (# _, _, hole #) -> TSet (clearM hole) + +-- | If the specified element is in the set, returns 'Just' the index of the element, otherwise returns 'Nothing'. +lookupIndex :: TKey a => a -> TSet a -> Maybe Int +lookupIndex a (TSet s) = searchMC (toRep a) s (\ _ -> Nothing) (\ _ hole -> Just $ sizeM (beforeM hole)) diff --git a/Data/Vector/Build.hs b/Data/Vector/Build.hs new file mode 100644 --- /dev/null +++ b/Data/Vector/Build.hs @@ -0,0 +1,24 @@ +{-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-} +module Data.Vector.Build where + +import Control.Monad.Primitive +import Data.Vector.Generic (Vector, create) +import Data.Vector.Generic.Mutable + +import Data.Foldable + +import Prelude hiding (foldr) +import GHC.Exts + +{-# INLINE toMVectorMapN #-} +toMVectorMapN :: (Foldable f, PrimMonad m, MVector v b) => + Int -> (a -> b) -> f a -> m (v (PrimState m) b) +toMVectorMapN !n f xs = do + !mv <- new n + let writer a k i# = unsafeWrite mv (I# i#) (f a) >> k (i# +# 1#) + foldr writer (\ _ -> return ()) xs 0# + return mv + +{-# INLINE toVectorMapN #-} +toVectorMapN :: (Foldable f, Vector v b) => Int -> (a -> b) -> f a -> v b +toVectorMapN !n f xs = create (toMVectorMapN n f xs) diff --git a/Tests.hs b/Tests.hs --- a/Tests.hs +++ b/Tests.hs @@ -6,7 +6,8 @@ import Control.Applicative import qualified Data.TrieMap as T import qualified Data.Map as M -import Data.List (foldl') +import Data.Ord +import Data.List (foldl', sortBy) import Data.TrieMap.Representation import Test.QuickCheck import Prelude hiding (null, lookup) @@ -200,6 +201,31 @@ Just (m', tm') -> verify m' tm' ops verify _ _ [] = True +newtype SortedAssoc k a = SortedAssoc [(k, a)] deriving (Show) +newtype SortedDistinctAssoc k a = SDA [(k, a)] deriving (Show) + +instance (Ord k, Arbitrary k, Arbitrary a) => Arbitrary (SortedAssoc k a) where + arbitrary = do + xs <- arbitrary + return (SortedAssoc (sortBy (comparing fst) xs)) + shrink (SortedAssoc xs) = do + xs' <- shrink xs + return (SortedAssoc (sortBy (comparing fst) xs')) + +instance (Ord k, Arbitrary k, Arbitrary a) => Arbitrary (SortedDistinctAssoc k a) where + arbitrary = do + SortedAssoc xs <- arbitrary + return (SDA $ sNub fst xs) + shrink (SDA xs) = do + SortedAssoc xs' <- shrink (SortedAssoc xs) + return (SDA $ sNub fst xs') + +fromAscListTest :: [(Key, Val)] -> [(Key, Val)] +fromAscListTest ((k1, v1):xs@((k2,v2):xs')) + | k1 == k2 = fromAscListTest ((k1, v2 ++ v1):xs') + | otherwise = (k1, v1) : fromAscListTest xs +fromAscListTest xs = xs + concretes :: [Property] concretes = [ printTestCase "extending by a single 0 makes a difference" @@ -208,10 +234,32 @@ (let input = [(BS.pack [0], "a"), (BS.pack [0,0,0,0,0], "a")] in T.assocs (T.fromList input) == input), printTestCase "comparisons are correct" (let input = [(BS.pack [0], "a"), (BS.pack [0,0,0,0,maxBound], "a")] in T.assocs (T.fromList input) == input), - printTestCase "genOptRepr is consistent with equality" (\ a b -> ((a :: Key') == b) == (toRep a == toRep b)), printTestCase "deleteAt works for OrdMap" - (let input = [(1.4 :: Double, 'a'), (-4.0, 'b')] in T.assocs (T.deleteAt 0 (T.fromList input)) == [(1.4, 'a')]) + (let input = [(1.4 :: Double, 'a'), (-4.0, 'b')] in T.assocs (T.deleteAt 0 (T.fromList input)) == [(1.4, 'a')]), + printTestCase "genOptRepr is consistent with equality" (\ a b -> ((a :: Key') == b) == (toRep a == toRep b)) + ,printTestCase "after works for RadixTrie" + (let input = [("abcd", 'a'), ("abcdef", 'b')]; m = T.fromList input in + T.assocs (T.after (snd (T.search "abcde" m))) == [("abcdef", 'b')]) + , + (printTestCase "fromDistinctAscList" + (\ (SDA sinput) -> expect (sinput :: [(Key, Val)]) (T.assocs (T.fromDistinctAscList sinput)))) + , + printTestCase "fromAscList" + (\ (SortedAssoc sinput) -> expect (fromAscListTest sinput) (T.assocs (T.fromAscListWith (++) sinput))) ] + +expect :: (Eq a, Show a) => a -> a -> Property +expect expected result = printTestCase ("Expected:\t" ++ show expected ++ "\nActual:\t\t" ++ show result) (expected == result) + +sNub :: Ord b => (a -> b) -> [a] -> [a] +sNub f xs = nubber xs'' + where xs' = [(x, f x) | x <- xs] + xs'' = sortBy (comparing snd) xs' + nubber ((x1, y1):xs@((_, y2):xs')) + | y1 == y2 = nubber ((x1, y1):xs') + | otherwise = x1:nubber xs + nubber [(x, _)] = [x] + nubber [] = [] $(genRepr ''Key) $(genOptRepr ''Key') diff --git a/TrieMap.cabal b/TrieMap.cabal --- a/TrieMap.cabal +++ b/TrieMap.cabal @@ -1,22 +1,15 @@ name: TrieMap -version: 3.0.1 +version: 4.0.0 cabal-version: >= 1.6 tested-with: GHC category: Algorithms synopsis: Automatic type inference of generalized tries with Template Haskell. description: Provides a efficient and compact implementation of generalized tries, and Template Haskell tools to generate - the necessary translation code. This is meant as a drop-in replacement for Data.Map. - + the necessary translation code. This is meant as a drop-in replacement for Data.Map, and can be used anywhere + @Data.Map@ can be used. + . The most recent release combines zipper-based ideas from recently proposed changes to Data.Map, as well as heavily optimized ByteString and Vector instances based on the vector package. - - Since version 2, unit tests and benchmarks have been taken much more seriously, and major optimizations - have been made. - - Compared to Data.Map and Data.Set, on e.g. @ByteString@s, TrieMaps support 6-12x faster @union@, - @intersection@, and @difference@ operations, 2x faster @lookup@, but 2x slower @toList@, and 4x slower @filter@. - Other operations are closely tied. TrieMaps tend to use somewhat more memory, and frequently perform better - with increased heap space and allocation area. license: BSD3 license-file: LICENSE author: Louis Wasserman @@ -29,10 +22,13 @@ location: git@github.com:lowasser/TrieMap.git Library{ -build-Depends: base < 5.0.0.0, containers, template-haskell, bytestring, th-expand-syns, vector, primitive +build-Depends: base < 5.0.0.0, containers, template-haskell >= 2.5.0.0, bytestring >= 0.9.1.0, th-expand-syns, + vector >= 0.6, primitive >= 0.3, unpack-funcs >= 0.1.2, transformers >= 0.2.0.0 ghc-options: -Wall -fno-warn-name-shadowing -fno-warn-orphans -O2 -fno-spec-constr-count -fno-spec-constr-threshold -fno-liberate-case-threshold -fmax-worker-args=100 +if impl(ghc >= 7.0.0) + ghc-options: -fllvm -optlo-O3 -optlo-std-compile-opts exposed-modules: Data.TrieMap, Data.TrieSet, @@ -41,9 +37,16 @@ Data.TrieMap.Modifiers other-modules: Control.Monad.Ends, + Control.Monad.Lookup, + Data.Vector.Build, Data.TrieMap.TrieKey, Data.TrieMap.Utils, Data.TrieMap.Sized, + Data.TrieMap.TrieKey.Search, + Data.TrieMap.TrieKey.Subset, + Data.TrieMap.TrieKey.Buildable, + Data.TrieMap.TrieKey.SetOp, + Data.TrieMap.TrieKey.Projection, Data.TrieMap.Representation.Class, Data.TrieMap.Representation.TH, Data.TrieMap.Representation.TH.Utils,