diff --git a/Data/TrieMap.hs b/Data/TrieMap.hs --- a/Data/TrieMap.hs +++ b/Data/TrieMap.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE TypeFamilies, FlexibleContexts, UnboxedTuples #-} +{-# LANGUAGE TypeFamilies, FlexibleContexts, UnboxedTuples, RecordWildCards #-} module Data.TrieMap ( -- * Map type @@ -129,8 +129,8 @@ import Data.TrieMap.Class.Instances() import Data.TrieMap.TrieKey import Data.TrieMap.Applicative -import Data.TrieMap.Rep -import Data.TrieMap.Rep.Instances () +import Data.TrieMap.Representation +import Data.TrieMap.Representation.Instances () import Data.TrieMap.Sized import Control.Applicative hiding (empty) @@ -161,32 +161,38 @@ -- -- 1. A 'TLocation' (and the value at that position, if any) is obtained from a 'TMap' by searching or indexing. -- 2. A new 'TMap' is made from a 'TLocation' by either filling the hole with a value ('assign') or erasing it ('clear'). -newtype TLocation k a = TLoc (Hole (Rep k) (Elem a)) +data TLocation k a = TLoc k (Hole (Rep k) (Assoc k a)) +{-# INLINE empty #-} -- | /O(1)/. The empty map. empty :: TKey k => TMap k a empty = TMap emptyM -- | /O(1)/. A map with a single element. +{-# INLINE singleton #-} singleton :: TKey k => k -> a -> TMap k a -singleton k a = insert k a empty +singleton k a = TMap (singletonM (toRep k) (Assoc k a)) -- | /O(1)/. Is the map empty? +{-# INLINE null #-} null :: TKey k => TMap k a -> Bool null (TMap m) = nullM 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) = getElem <$> lookupM (toRep k) m +lookup k (TMap m) = getValue <$> lookupM (toRep k) m -- | 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. +{-# INLINE findWithDefault #-} findWithDefault :: TKey k => a -> k -> TMap k a -> a findWithDefault a = fromMaybe a .: lookup -- | Find the value at a key. Calls 'error' when the element can not be found. +{-# INLINE (!) #-} (!) :: TKey k => TMap k a -> k -> a m ! k = fromMaybe (error "Element not found") (lookup k m) @@ -299,6 +305,7 @@ -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +{-# INLINE update #-} update :: TKey k => (a -> Maybe a) -> k -> TMap k a -> TMap k a update f = updateWithKey (const f) @@ -311,6 +318,7 @@ -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +{-# INLINE updateWithKey #-} updateWithKey :: TKey k => (k -> a -> Maybe a) -> k -> TMap k a -> TMap k a updateWithKey f k m = case search k m of (Nothing, _) -> m @@ -318,17 +326,20 @@ -- | Post-order fold. The function will be applied from the lowest -- value to the highest. +{-# INLINE foldrWithKey #-} foldrWithKey :: TKey k => (k -> a -> b -> b) -> b -> TMap k a -> b -foldrWithKey f z (TMap m) = foldrWithKeyM (\ k (Elem a) -> f (fromRep k) a) m z +foldrWithKey f z (TMap m) = foldrM (\ (Assoc k a) -> f k a) m z -- | Pre-order fold. The function will be applied from the highest -- value to the lowest. +{-# INLINE foldlWithKey #-} foldlWithKey :: TKey k => (b -> k -> a -> b) -> b -> TMap k a -> b -foldlWithKey f z (TMap m) = foldlWithKeyM (\ k z (Elem a) -> f z (fromRep k) a) m z +foldlWithKey f z (TMap m) = foldlM (\ z (Assoc k a) -> f z k a) m z -- | 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 <$> traverseWithKeyM (\ k (Elem a) -> Elem <$> f (fromRep k) a) m +traverseWithKey f (TMap m) = TMap <$> traverseM (\ (Assoc k a) -> Assoc k <$> f k a) m -- | Map a function over all values in the map. -- @@ -343,7 +354,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 (mapWithKeyM (\ k (Elem a) -> Elem (f (fromRep k) a)) m) +mapWithKey f (TMap m) = TMap (fmapM (\ (Assoc k a) -> Assoc k (f k a)) m) -- | -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. @@ -432,9 +443,10 @@ {-# 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 - f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b + f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b -- | 'symmetricDifference' is equivalent to @'unionMaybeWith' (\ _ _ -> Nothing)@. +{-# INLINE symmetricDifference #-} symmetricDifference :: TKey k => TMap k a -> TMap k a -> TMap k a symmetricDifference = unionMaybeWith (\ _ _ -> Nothing) @@ -474,13 +486,14 @@ {-# 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 - f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b + f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b -- | Difference of two maps. -- Return elements of the first map not existing in the second map. -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. -- -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" +{-# INLINE difference #-} difference :: TKey k => TMap k a -> TMap k b -> TMap k a difference = differenceWith (\ _ _ -> Nothing) @@ -514,7 +527,7 @@ {-# 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 - f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b + f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b -- | Retrieves the value associated with minimal key of the -- map, and the map stripped of that element, or 'Nothing' if passed an @@ -584,10 +597,10 @@ updateMax = updateMaxWithKey . const {-# INLINE updateHelper #-} -updateHelper :: (TKey k, MonadPlus m) => (k -> a -> Maybe a) -> TMap k a -> m (Maybe (Elem a), Hole (Rep k) (Elem a)) +updateHelper :: (TKey k, MonadPlus m) => (k -> a -> Maybe a) -> TMap k a -> m (Maybe (Assoc k a), Hole (Rep k) (Assoc k a)) updateHelper f (TMap m) = do - (Elem a, loc) <- extractHoleM m - return (Elem <$> f (fromRep (keyM loc)) a, loc) + (Assoc k a, loc) <- extractHoleM m + return (Assoc k <$> f k a, loc) -- | Update the value at the minimal key. -- @@ -625,23 +638,23 @@ 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) -{-# INLINE minViewWithKey #-} -- | Retrieves the minimal (key,value) pair of the map, and -- the map stripped of that element, or 'Nothing' if passed an empty map. -- -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a") -- > minViewWithKey empty == Nothing minViewWithKey :: TKey k => TMap k a -> Maybe ((k, a), TMap k a) +{-# INLINE minViewWithKey #-} minViewWithKey m = do (a, loc) <- minLocation m return ((key loc, a), after loc) -{-# INLINE maxViewWithKey #-} -- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and -- the map stripped of that element, or 'Nothing' if passed an empty map. -- -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b") -- > maxViewWithKey empty == Nothing +{-# INLINE maxViewWithKey #-} maxViewWithKey :: TKey k => TMap k a -> Maybe ((k, a), TMap k a) maxViewWithKey m = do (a, loc) <- maxLocation m @@ -696,9 +709,9 @@ 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 (# mL, mR #) -> (TMap mL, TMap mR) - where f' k (Elem a) = case f (fromRep k) a of - Left b -> (# Just (Elem b), Nothing #) - Right c -> (# Nothing, Just (Elem c) #) + where f' (Assoc k a) = case f k a of + Left b -> (# Just (Assoc k b), Nothing #) + Right c -> (# Nothing, Just (Assoc k c) #) -- | /O(n)/. Map values and collect the 'Just' results. -- @@ -714,7 +727,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 (\ k (Elem a) -> Elem <$> f (fromRep k) a) m) +mapMaybeWithKey f (TMap m) = TMap (mapMaybeM (\ (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 @@ -763,6 +776,7 @@ -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty) -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty) +{-# INLINE split #-} split :: TKey k => k -> TMap k a -> (TMap k a, TMap k a) split k m = case splitLookup k m of (mL, _, mR) -> (mL, mR) @@ -806,7 +820,7 @@ {-# INLINEABLE isSubmapOfBy #-} isSubmapOfBy :: TKey k => (a -> b -> Bool) -> TMap k a -> TMap k b -> Bool isSubmapOfBy (<=) (TMap m1) (TMap m2) = isSubmapM (<<=) m1 m2 where - Elem a <<= Elem b = a <= b + Assoc _ a <<= Assoc _ b = a <= b -- | 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 @@ -850,8 +864,8 @@ -- > fromListWith (++) [] == empty {-# INLINEABLE fromListWithKey #-} fromListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a -fromListWithKey f xs = TMap (fromListM f' [(toRep k, Elem a) | (k, a) <- xs]) - where f' k (Elem a) (Elem b) = Elem (f (fromRep k) a b) +fromListWithKey f xs = TMap (fromListM f' [(toRep k, Assoc k a) | (k, a) <- 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./ @@ -860,8 +874,8 @@ -- > 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, Elem a) | (k, a) <- xs]) - where f' k (Elem a) (Elem b) = Elem (f (fromRep k) a b) +fromAscListWithKey f xs = TMap (fromAscListM f' [(toRep k, Assoc k a) | (k, a) <- 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./ @@ -869,13 +883,14 @@ -- > 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, Elem a) | (k, a) <- xs]) +fromDistinctAscList xs = TMap (fromDistAscListM [(toRep k, Assoc k a) | (k, a) <- xs]) -- | /O(1)/. The number of elements in the map. -- -- > size empty == 0 -- > size (singleton 1 'a') == 1 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 +{-# INLINE size #-} size :: TKey k => TMap k a -> Int size (TMap m) = getSize m @@ -904,16 +919,19 @@ keysSet m = TSet (() <$ m) -- | /O(1)/. The key marking the position of the \"hole\" in the map. +{-# INLINE key #-} key :: TKey k => TLocation k a -> k -key (TLoc hole) = fromRep (keyM hole) +key (TLoc k _) = k -- | @'before' loc@ is the submap with keys less than @'key' loc@. +{-# INLINE before #-} before :: TKey k => TLocation k a -> TMap k a -before (TLoc hole) = TMap (beforeM Nothing hole) +before (TLoc _ hole) = TMap (beforeM Nothing hole) -- | @'after' loc@ is the submap with keys greater than @'key' loc@. +{-# INLINE after #-} after :: TKey k => TLocation k a -> TMap k a -after (TLoc hole) = TMap (afterM Nothing hole) +after (TLoc _ hole) = TMap (afterM Nothing hole) -- | Search the map for the given key, returning the -- corresponding value (if any) and an updatable location for that key. @@ -927,9 +945,11 @@ -- @ -- -- @'lookup' k m == 'fst' ('search' k m)@ +{-# INLINE search #-} search :: TKey k => k -> TMap k a -> (Maybe a, TLocation k a) search k (TMap m) = case searchM (toRep k) m of - (# a, hole #) -> (getElem <$> a, TLoc hole) + (# Just (Assoc k a), hole #) -> (Just a, TLoc k hole) + (# _, hole #) -> (Nothing, TLoc k hole) -- | Return the value and an updatable location for the -- /i/th key in the map. Calls 'error' if /i/ is out of range. @@ -949,13 +969,13 @@ | i < 0 || i >= size m = error "TrieMap.index: index out of range" index i (TMap m) = case indexM (unbox i) m of - (# _, Elem a, hole #) -> (a, TLoc hole) + (# _, Assoc k a, hole #) -> (a, TLoc k hole) {-# INLINE extract #-} extract :: (TKey k, MonadPlus m) => TMap k a -> m (a, TLocation k a) extract (TMap m) = do - (Elem a, hole) <- extractHoleM m - return (a, TLoc hole) + (Assoc k a, hole) <- extractHoleM m + return (a, TLoc k hole) -- | /O(log n)/. Return the value and an updatable location for the -- least key in the map, or 'Nothing' if the map is empty. @@ -993,14 +1013,16 @@ -- at the location (replacing an existing value, if any). -- -- @'assign' v loc == 'before' loc `union` 'singleton' ('key' loc) v `union` 'after' loc@ +{-# INLINE assign #-} assign :: TKey k => a -> TLocation k a -> TMap k a -assign a (TLoc hole) = TMap (assignM (Elem a) hole) +assign a (TLoc k hole) = TMap (assignM (Just $ Assoc k a) hole) -- | Return a map obtained by erasing the location. -- -- @'clear' loc == 'before' loc `union` 'after' loc@ +{-# INLINE clear #-} clear :: TKey k => TLocation k a -> TMap k a -clear (TLoc hole) = TMap (clearM hole) +clear (TLoc _ hole) = TMap (assignM Nothing hole) {-# INLINE fillHole #-} fillHole :: TKey k => Maybe a -> TLocation k a -> TMap k a diff --git a/Data/TrieMap/Class.hs b/Data/TrieMap/Class.hs --- a/Data/TrieMap/Class.hs +++ b/Data/TrieMap/Class.hs @@ -1,21 +1,24 @@ {-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, UndecidableInstances #-} -module Data.TrieMap.Class (TMap(..), TSet (..), TKey, Rep, TrieMap, TrieKey) where +module Data.TrieMap.Class (TMap(..), TSet(..), TKey, Rep, TrieMap, TrieKey) where import Data.TrieMap.TrieKey -import Data.TrieMap.Rep +import Data.TrieMap.Representation.Class import Data.TrieMap.Sized import Control.Applicative -import Data.Foldable +import Data.Foldable hiding (foldrM, foldlM) import Data.Traversable import Prelude hiding (foldr) -newtype TMap k a = TMap {getTMap :: TrieMap (Rep k) (Elem a)} +newtype TMap k a = TMap {getTMap :: TrieMap (Rep k) (Assoc k a)} newtype TSet a = TSet (TMap a ()) +-- | @'TKey' k@ is a handy alias for @('Repr' k, 'TrieKey' ('Rep' k))@. To make a type an instance of 'TKey', +-- use the methods available in "Data.TrieMap.Representation.TH" to generate a 'Repr' instance that will +-- satisfy @'TrieKey' ('Rep' k)@. class (Repr k, TrieKey (Rep k)) => TKey k instance (Repr k, TrieKey (Rep k)) => TKey k @@ -24,7 +27,7 @@ fmap = fmapDefault instance TKey k => Foldable (TMap k) where - foldr f z (TMap m) = foldrWithKeyM (\ _ (Elem a) -> f a) m z + foldr f z (TMap m) = foldrM (\ (Assoc _ a) -> f a) m z instance TKey k => Traversable (TMap k) where - traverse f (TMap m) = TMap <$> traverseWithKeyM (\ _ (Elem a) -> Elem <$> f a) m + traverse f (TMap m) = TMap <$> traverseM (\ (Assoc k a) -> Assoc k <$> f a) m diff --git a/Data/TrieMap/Class/Instances.hs b/Data/TrieMap/Class/Instances.hs --- a/Data/TrieMap/Class/Instances.hs +++ b/Data/TrieMap/Class/Instances.hs @@ -2,14 +2,12 @@ import Data.TrieMap.Class () import Data.TrieMap.TrieKey () -import Data.TrieMap.Rep () -import Data.TrieMap.Rep.Instances () -import Data.TrieMap.Representation () +import Data.TrieMap.Representation.Instances () import Data.TrieMap.Sized () +import Data.TrieMap.ReverseMap () import Data.TrieMap.RadixTrie () import Data.TrieMap.IntMap () import Data.TrieMap.OrdMap () -import Data.TrieMap.ReverseMap () import Data.TrieMap.ProdMap () import Data.TrieMap.UnionMap () import Data.TrieMap.UnitMap() diff --git a/Data/TrieMap/IntMap.hs b/Data/TrieMap/IntMap.hs --- a/Data/TrieMap/IntMap.hs +++ b/Data/TrieMap/IntMap.hs @@ -12,62 +12,49 @@ import Data.Maybe hiding (mapMaybe) import Data.Word +import GHC.Exts + import Prelude hiding (lookup, null, foldl, foldr) #include "MachDeps.h" -#if WORD_SIZE_IN_BITS == 32 -import GHC.Prim -import GHC.Word - -complement32 (W32# w#) = W32# (not# w#) -#elif WORD_SIZE_IN_BITS > 32 -complement32 = xor (bit 32 - 1) -#else -import GHC.Prim -import GHC.IntWord32 -complement32 (W32# w#) = W32# (not32# w#) -#endif -complement32 :: Word32 -> Word32 - -{-# RULES - "complement/Word32" complement = complement32 - #-} - -type Nat = Word32 +type Nat = Word -type Prefix = Word32 -type Mask = Word32 -type Key = Word32 +type Prefix = Word +type Mask = Word +type Key = Word type Size = Int# data Path a = Root - | LeftBin !Prefix !Mask !(Path a) !(TrieMap Word32 a) - | RightBin !Prefix !Mask !(TrieMap Word32 a) !(Path a) + | LeftBin !Prefix !Mask !(Path a) !(TrieMap Word a) + | RightBin !Prefix !Mask !(TrieMap Word a) !(Path a) -instance TrieKey Word32 where - data TrieMap Word32 a = Nil +instance TrieKey Word where + (=?) = (==) + cmp = compare + + data TrieMap Word a = Nil | Tip !Size !Key a - | Bin !Size !Prefix !Mask !(TrieMap Word32 a) !(TrieMap Word32 a) - data Hole Word32 a = Hole !Key !(Path a) + | Bin !Size !Prefix !Mask !(TrieMap Word a) !(TrieMap Word a) + data Hole Word a = Hole !Key !(Path a) emptyM = Nil singletonM = singleton - nullM = null + getSimpleM Nil = Null + getSimpleM (Tip _ _ a) = Singleton a + getSimpleM _ = NonSimple sizeM = size lookupM = lookup - traverseWithKeyM = traverseWithKey - foldrWithKeyM = foldr - foldlWithKeyM = foldl - mapWithKeyM = mapWithKey + traverseM = traverse + foldrM = foldr + foldlM = foldl + fmapM = mapWithKey mapMaybeM = mapMaybe mapEitherM = mapEither - unionM = unionWithKey - isectM = intersectionWithKey - diffM = differenceWithKey --- extractM f = extract f + unionM = unionWith + isectM = intersectionWith + diffM = differenceWith isSubmapM = isSubmapOfBy singleHoleM k = Hole k Root - keyM (Hole k _) = k beforeM a (Hole k path) = before (singletonMaybe k a) path where before t Root = t before t (LeftBin _ _ path _) = before t path @@ -76,7 +63,7 @@ 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 - searchM !k = onUnboxed (Hole k) (search Root) where + searchM !k = onSnd (Hole k) (search Root) where search path t@(Bin _ p m l r) | nomatch k p m = (# Nothing, branchHole k p path t #) | zero k m @@ -88,8 +75,7 @@ | otherwise = (# Nothing, branchHole k ky path t #) search path _ = (# Nothing, path #) indexM i# t = indexT i# t Root where - indexT _ Nil _ = (# error err, error err, error err #) where - err = "Error: empty trie" + indexT _ Nil _ = indexFail () 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) @@ -101,27 +87,25 @@ extractHole path (Bin _ p m l r) = extractHole (LeftBin p m path r) l `mplus` extractHole (RightBin p m l path) r - assignM v (Hole kx path) = assign (singleton kx v) path where + assignM v (Hole kx path) = assign (singletonM' kx v) path where 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 - clearM (Hole _ path) = clear Nil path where - clear t Root = t - clear t (LeftBin p m path r) = clear (bin p m t r) path - clear t (RightBin p m l path) = clear (bin p m l t) path + {-# INLINE unifyM #-} + unifyM = unify -branchHole :: Key -> Prefix -> Path a -> TrieMap Word32 a -> Path a +branchHole :: Key -> Prefix -> Path a -> TrieMap Word a -> Path a branchHole !k !p path t | zero k m = LeftBin p' m path t | otherwise = RightBin p' m t path where m = branchMask k p p' = mask k m -natFromInt :: Word32 -> Nat +natFromInt :: Word -> Nat natFromInt = id -intFromNat :: Nat -> Word32 +intFromNat :: Nat -> Word intFromNat = id shiftRL :: Nat -> Key -> Nat @@ -135,127 +119,121 @@ shiftRL x i = shiftR x (fromIntegral i) -- #endif -size :: TrieMap Word32 a -> Int# +size :: TrieMap Word a -> Int# size Nil = 0# size (Tip sz _ _) = sz size (Bin sz _ _ _ _) = sz -null :: TrieMap Word32 a -> Bool -null Nil = True -null _ = False - -lookup :: Nat -> TrieMap Word32 a -> Maybe a -lookup k (Bin _ _ m l r) = lookup k (if zeroN k m then l else r) +lookup :: Nat -> TrieMap Word a -> Maybe a +lookup !k (Bin _ _ m l r) = lookup k (if zeroN k m then l else r) lookup k (Tip _ kx x) | k == kx = Just x lookup _ _ = Nothing -singleton :: Sized a => Key -> a -> TrieMap Word32 a +singleton :: Sized a => Key -> a -> TrieMap Word a singleton k a = Tip (getSize# a) k a -singletonMaybe :: Sized a => Key -> Maybe a -> TrieMap Word32 a +singletonMaybe :: Sized a => Key -> Maybe a -> TrieMap Word a singletonMaybe k = maybe Nil (singleton k) -traverseWithKey :: (Applicative f, Sized b) => (Key -> a -> f b) -> TrieMap Word32 a -> f (TrieMap Word32 b) -traverseWithKey f t = case t of +traverse :: (Applicative f, Sized b) => (a -> f b) -> TrieMap Word a -> f (TrieMap Word b) +traverse f t = case t of Nil -> pure Nil - Tip _ kx x -> singleton kx <$> f kx x - Bin _ p m l r -> bin p m <$> traverseWithKey f l <*> traverseWithKey f r + Tip _ kx x -> singleton kx <$> f x + Bin _ p m l r -> bin p m <$> traverse f l <*> traverse f r -foldr :: (Key -> a -> b -> b) -> TrieMap Word32 a -> b -> b +foldr :: (a -> b -> b) -> TrieMap Word a -> b -> b foldr f t = case t of Bin _ _ _ l r -> foldr f l . foldr f r - Tip _ k x -> f k x + Tip _ _ x -> f x Nil -> id -foldl :: (Key -> b -> a -> b) -> TrieMap Word32 a -> b -> b +foldl :: (b -> a -> b) -> TrieMap Word a -> b -> b foldl f t = case t of Bin _ _ _ l r -> foldl f r . foldl f l - Tip _ k x -> flip (f k) x + Tip _ _ x -> flip f x Nil -> id -mapWithKey :: Sized b => (Key -> a -> b) -> TrieMap Word32 a -> TrieMap Word32 b +mapWithKey :: Sized b => (a -> b) -> TrieMap Word a -> TrieMap Word b mapWithKey f (Bin _ p m l r) = bin p m (mapWithKey f l) (mapWithKey f r) -mapWithKey f (Tip _ kx x) = singleton kx (f kx x) +mapWithKey f (Tip _ kx x) = singleton kx (f x) mapWithKey _ _ = Nil -mapMaybe :: Sized b => (Key -> a -> Maybe b) -> TrieMap Word32 a -> TrieMap Word32 b +mapMaybe :: Sized b => (a -> Maybe b) -> TrieMap Word a -> TrieMap Word 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 kx x) +mapMaybe f (Tip _ kx x) = singletonMaybe kx (f x) mapMaybe _ _ = Nil -mapEither :: (Sized b, Sized c) => EitherMap Key a b c -> - TrieMap Word32 a -> (# TrieMap Word32 b, TrieMap Word32 c #) -mapEither f (Bin _ p m l r) - | (# lL, lR #) <- mapEither f l, - (# rL, rR #) <- mapEither f r - = (# bin p m lL rL, bin p m lR rR #) -mapEither f (Tip _ kx x) = both (singletonMaybe kx) (singletonMaybe kx) (f kx) x +mapEither :: (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> + TrieMap Word a -> (# TrieMap Word b, TrieMap Word 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 #) -unionWithKey :: Sized a => UnionFunc Key a -> TrieMap Word32 a -> TrieMap Word32 a -> TrieMap Word32 a -unionWithKey _ Nil t = t -unionWithKey _ t Nil = t -unionWithKey f (Tip _ k x) t = alterM (maybe (Just x) (f k x)) k t -unionWithKey f t (Tip _ k x) = alterM (maybe (Just x) (flip (f k) x)) k t -unionWithKey f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) +unionWith :: Sized a => (a -> a -> Maybe a) -> TrieMap Word a -> TrieMap Word a -> TrieMap Word a +unionWith _ Nil t = t +unionWith _ t Nil = t +unionWith f (Tip _ k x) t = alterM (maybe (Just x) (f x)) k t +unionWith f t (Tip _ k x) = alterM (maybe (Just x) (flip f x)) k t +unionWith f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) | shorter m1 m2 = union1 | shorter m2 m1 = union2 - | p1 == p2 = bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2) + | p1 == p2 = bin p1 m1 (unionWith f l1 l2) (unionWith f r1 r2) | otherwise = join p1 t1 p2 t2 where union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2 - | zero p2 m1 = bin p1 m1 (unionWithKey f l1 t2) r1 - | otherwise = bin p1 m1 l1 (unionWithKey f r1 t2) + | zero p2 m1 = bin p1 m1 (unionWith f l1 t2) r1 + | otherwise = bin p1 m1 l1 (unionWith f r1 t2) union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 - | zero p1 m2 = bin p2 m2 (unionWithKey f t1 l2) r2 - | otherwise = bin p2 m2 l2 (unionWithKey f t1 r2) + | zero p1 m2 = bin p2 m2 (unionWith f t1 l2) r2 + | otherwise = bin p2 m2 l2 (unionWith f t1 r2) -intersectionWithKey :: Sized c => IsectFunc Key a b c -> TrieMap Word32 a -> TrieMap Word32 b -> TrieMap Word32 c -intersectionWithKey _ Nil _ = Nil -intersectionWithKey _ _ Nil = Nil -intersectionWithKey f (Tip _ k x) t2 - = singletonMaybe k (lookup (natFromInt k) t2 >>= f k x) -intersectionWithKey f t1 (Tip _ k y) - = singletonMaybe k (lookup (natFromInt k) t1 >>= flip (f k) y) -intersectionWithKey f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) +intersectionWith :: Sized c => (a -> b -> Maybe c) -> TrieMap Word a -> TrieMap Word b -> TrieMap Word c +intersectionWith _ Nil _ = Nil +intersectionWith _ _ Nil = Nil +intersectionWith f (Tip _ k x) t2 + = singletonMaybe k (lookup (natFromInt k) t2 >>= f x) +intersectionWith f t1 (Tip _ k y) + = singletonMaybe k (lookup (natFromInt k) t1 >>= flip f y) +intersectionWith f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) | shorter m1 m2 = intersection1 | shorter m2 m1 = intersection2 - | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2) + | 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 = intersectionWithKey f l1 t2 - | otherwise = intersectionWithKey f r1 t2 + | zero p2 m1 = intersectionWith f l1 t2 + | otherwise = intersectionWith f r1 t2 intersection2 | nomatch p1 p2 m2 = Nil - | zero p1 m2 = intersectionWithKey f t1 l2 - | otherwise = intersectionWithKey f t1 r2 + | zero p1 m2 = intersectionWith f t1 l2 + | otherwise = intersectionWith f t1 r2 -differenceWithKey :: Sized a => (Key -> a -> b -> Maybe a) -> TrieMap Word32 a -> TrieMap Word32 b -> TrieMap Word32 a -differenceWithKey _ Nil _ = Nil -differenceWithKey _ t Nil = t -differenceWithKey f t1@(Tip _ k x) t2 - = maybe t1 (singletonMaybe k . f k x) (lookup (natFromInt k) t2) -differenceWithKey f t (Tip _ k y) = alterM (>>= flip (f k) y) k t -differenceWithKey f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) +differenceWith :: Sized a => (a -> b -> Maybe a) -> TrieMap Word a -> TrieMap Word b -> TrieMap Word a +differenceWith _ Nil _ = Nil +differenceWith _ t Nil = t +differenceWith f t1@(Tip _ k x) t2 + = maybe t1 (singletonMaybe k . f x) (lookup (natFromInt k) t2) +differenceWith f t (Tip _ k y) = alterM (>>= flip f y) k t +differenceWith f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) | shorter m1 m2 = difference1 | shorter m2 m1 = difference2 - | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2) + | p1 == p2 = bin p1 m1 (differenceWith f l1 l2) (differenceWith f r1 r2) | otherwise = t1 where difference1 | nomatch p2 p1 m1 = t1 - | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1 - | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2) + | zero p2 m1 = bin p1 m1 (differenceWith f l1 t2) r1 + | otherwise = bin p1 m1 l1 (differenceWith f r1 t2) difference2 | nomatch p1 p2 m2 = t1 - | zero p1 m2 = differenceWithKey f t1 l2 - | otherwise = differenceWithKey f t1 r2 + | zero p1 m2 = differenceWith f t1 l2 + | otherwise = differenceWith f t1 r2 -isSubmapOfBy :: LEq a b -> LEq (TrieMap Word32 a) (TrieMap Word32 b) +isSubmapOfBy :: LEq a b -> LEq (TrieMap Word a) (TrieMap Word 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 @@ -268,12 +246,6 @@ isSubmapOfBy _ Nil _ = True --- extract :: Alternative f => Sized a -> (Key -> a -> f (x, Maybe a)) -> TrieMap Word32 a -> f (x, TrieMap Word32 a) --- extract f (Bin _ p m l r) = --- fmap (\ l' -> bin p m l' r) <$> extract f l <|> fmap (bin p m l) <$> extract f r --- extract f (Tip _ k x) = fmap (singletonMaybe k) <$> f k x --- extract _ _ _ = empty - mask :: Key -> Mask -> Prefix mask i m = maskW (natFromInt i) (natFromInt m) @@ -318,7 +290,8 @@ x5 -> x5 `xor` shiftRL x5 1 #endif -join :: Prefix -> TrieMap Word32 a -> Prefix -> TrieMap Word32 a -> TrieMap Word32 a +{-# INLINE join #-} +join :: Prefix -> TrieMap Word a -> Prefix -> TrieMap Word a -> TrieMap Word a join p1 t1 p2 t2 | zero p1 m = bin p m t1 t2 | otherwise = bin p m t2 t1 @@ -326,7 +299,19 @@ m = branchMask p1 p2 p = mask p1 m -bin :: Prefix -> Mask -> TrieMap Word32 a -> TrieMap Word32 a -> TrieMap Word32 a +bin :: Prefix -> Mask -> TrieMap Word a -> TrieMap Word a -> TrieMap Word a bin _ _ l Nil = l bin _ _ Nil r = r bin p m l r = Bin (size l +# size r) p m l r + +{-# INLINE unify #-} +unify :: Sized a => Key -> a -> Key -> a -> Unified Word a +unify k1 _ k2 _ + | k1 == k2 = Left (Hole k1 Root) +unify k1 a1 k2 a2 = Right (if zero k1 m then outBin t1 t2 else outBin t2 t1) + where !s1# = getSize# a1 + !s2# = getSize# a2 + t1 = Tip s1# k1 a1 + t2 = Tip s2# k2 a2 + m = branchMask k1 k2 + outBin = Bin (s1# +# s2#) (mask k1 m) m diff --git a/Data/TrieMap/Key.hs b/Data/TrieMap/Key.hs --- a/Data/TrieMap/Key.hs +++ b/Data/TrieMap/Key.hs @@ -1,42 +1,52 @@ {-# LANGUAGE TypeFamilies, UnboxedTuples #-} -module Data.TrieMap.Key (Key(..)) where +module Data.TrieMap.Key () where import Control.Applicative + import Data.TrieMap.Class import Data.TrieMap.TrieKey -import Data.TrieMap.Rep +import Data.TrieMap.Representation.Class import Data.TrieMap.Modifiers +import Data.TrieMap.ProdMap() +import Data.TrieMap.UnionMap() +import Data.TrieMap.IntMap() +import Data.TrieMap.OrdMap() +import Data.TrieMap.RadixTrie() + instance TKey k => TrieKey (Key k) where + Key k1 =? Key k2 = toRep k1 =? toRep k2 + Key k1 `cmp` Key k2 = toRep k1 `cmp` toRep k2 + newtype TrieMap (Key k) a = KeyMap (TrieMap (Rep k) a) newtype Hole (Key k) a = KeyHole (Hole (Rep k) a) emptyM = KeyMap emptyM singletonM (Key k) a = KeyMap (singletonM (toRep k) a) - nullM (KeyMap m) = nullM m + getSimpleM (KeyMap m) = getSimpleM m sizeM (KeyMap m) = sizeM m lookupM (Key k) (KeyMap m) = lookupM (toRep k) m - traverseWithKeyM f (KeyMap m) = KeyMap <$> traverseWithKeyM (f . Key . fromRep) m - foldrWithKeyM f (KeyMap m) = foldrWithKeyM (f . Key . fromRep) m - foldlWithKeyM f (KeyMap m) = foldlWithKeyM (f . Key . fromRep) m - mapWithKeyM f (KeyMap m) = KeyMap (mapWithKeyM (f . Key . fromRep) m) - mapMaybeM f (KeyMap m) = KeyMap (mapMaybeM (f . Key . fromRep) m) - mapEitherM f (KeyMap m) = both KeyMap KeyMap (mapEitherM (f . Key . fromRep)) m - unionM f (KeyMap m1) (KeyMap m2) = KeyMap (unionM (f . Key . fromRep) m1 m2) - isectM f (KeyMap m1) (KeyMap m2) = KeyMap (isectM (f . Key . fromRep) m1 m2) - diffM f (KeyMap m1) (KeyMap m2) = KeyMap (diffM (f . Key . fromRep) m1 m2) + traverseM f (KeyMap m) = KeyMap <$> traverseM f m + foldrM f (KeyMap m) = foldrM f m + foldlM f (KeyMap m) = foldlM f m + fmapM f (KeyMap m) = KeyMap (fmapM f m) + mapMaybeM f (KeyMap m) = KeyMap (mapMaybeM f m) + mapEitherM f (KeyMap m) = both KeyMap KeyMap (mapEitherM f) m + unionM f (KeyMap m1) (KeyMap m2) = KeyMap (unionM f m1 m2) + isectM f (KeyMap m1) (KeyMap m2) = KeyMap (isectM f m1 m2) + diffM f (KeyMap m1) (KeyMap m2) = KeyMap (diffM f m1 m2) isSubmapM (<=) (KeyMap m1) (KeyMap m2) = isSubmapM (<=) m1 m2 singleHoleM (Key k) = KeyHole (singleHoleM (toRep k)) - keyM (KeyHole hole) = Key (fromRep (keyM hole)) beforeM a (KeyHole hole) = KeyMap (beforeM a hole) afterM a (KeyHole hole) = KeyMap (afterM a hole) - searchM (Key k) (KeyMap m) = onUnboxed KeyHole (searchM (toRep k)) m + searchM (Key k) (KeyMap m) = onSnd KeyHole (searchM (toRep k)) m indexM i (KeyMap m) = case indexM i m of (# i', v, hole #) -> (# i', v, KeyHole hole #) extractHoleM (KeyMap m) = do (v, hole) <- extractHoleM m return (v, KeyHole hole) assignM v (KeyHole hole) = KeyMap (assignM v hole) - clearM (KeyHole hole) = KeyMap (clearM hole) + + unifyM (Key k1) a1 (Key k2) a2 = either (Left . KeyHole) (Right . KeyMap) (unifyM (toRep k1) a1 (toRep k2) a2) diff --git a/Data/TrieMap/Modifiers.hs b/Data/TrieMap/Modifiers.hs --- a/Data/TrieMap/Modifiers.hs +++ b/Data/TrieMap/Modifiers.hs @@ -1,7 +1,7 @@ {-# LANGUAGE FlexibleContexts, UndecidableInstances, TypeFamilies #-} module Data.TrieMap.Modifiers where -import Data.TrieMap.Rep +import Data.TrieMap.Representation.Class newtype Ordered a = Ord {unOrd :: a} deriving (Eq, Ord) newtype Rev k = Rev {getRev :: k} deriving (Eq) @@ -16,17 +16,6 @@ newtype Key k = Key {getKey :: k} -instance (Repr k, Eq (Rep k)) => Eq (Key k) where - Key k1 == Key k2 = toRep k1 == toRep k2 - -instance (Repr k, Ord (Rep k)) => Ord (Key k) where - Key k1 `compare` Key k2 = toRep k1 `compare` toRep k2 - Key k1 <= Key k2 = toRep k1 <= toRep k2 - Key k1 < Key k2 = toRep k1 < toRep k2 - Key k1 >= Key k2 = toRep k1 >= toRep k2 - Key k1 > Key k2 = toRep k1 > toRep k2 - instance Repr k => Repr (Key k) where type Rep (Key k) = Rep k toRep (Key k) = toRep k - fromRep = Key . fromRep diff --git a/Data/TrieMap/OrdMap.hs b/Data/TrieMap/OrdMap.hs --- a/Data/TrieMap/OrdMap.hs +++ b/Data/TrieMap/OrdMap.hs @@ -7,9 +7,9 @@ import Data.TrieMap.Modifiers import Control.Applicative -import Control.Monad hiding (join) +import Control.Monad hiding (join, fmap) -import Prelude hiding (lookup) +import Prelude hiding (lookup, foldr, foldl, fmap) import GHC.Exts @@ -27,6 +27,9 @@ singletonMaybe k = maybe Tip (singleton k) instance Ord k => TrieKey (Ordered k) where + Ord k1 =? Ord k2 = k1 == k2 + Ord k1 `cmp` Ord k2 = k1 `compare` k2 + data TrieMap (Ordered k) a = Tip | Bin Int# k a !(OrdMap k a) !(OrdMap k a) data Hole (Ordered k) a = @@ -34,30 +37,29 @@ | Full k !(Path k a) !(OrdMap k a) !(OrdMap k a) emptyM = Tip singletonM (Ord k) = singleton k - nullM Tip = True - nullM _ = False - sizeM = size# lookupM (Ord k) = lookup k - traverseWithKeyM f = traverseWithKey (f . Ord) - foldrWithKeyM f = foldrWithKey (f . Ord) - foldlWithKeyM f = foldlWithKey (f . Ord) - mapWithKeyM f = mapWithKey (f . Ord) - mapMaybeM f = mapMaybe (f . Ord) - mapEitherM f = mapEither (f . Ord) + getSimpleM Tip = Null + getSimpleM (Bin _ _ a Tip Tip) = Singleton a + getSimpleM _ = NonSimple + sizeM = size# + traverseM = traverse + foldrM = foldr + foldlM = foldl + fmapM = fmap + mapMaybeM = mapMaybe + mapEitherM = mapEither isSubmapM = isSubmap - fromAscListM f xs = fromAscList (f . Ord) [(k, a) | (Ord k, a) <- xs] + fromAscListM f xs = fromAscList f [(k, a) | (Ord k, a) <- xs] fromDistAscListM xs = fromDistinctAscList [(k, a) | (Ord k, a) <- xs] unionM _ Tip m2 = m2 unionM _ m1 Tip = m1 - unionM f m1 m2 = hedgeUnionWithKey (f . Ord) (const LT) (const GT) m1 m2 - isectM f = isect (f . Ord) + unionM f m1 m2 = hedgeUnion f (const LT) (const GT) m1 m2 + isectM = isect diffM _ Tip _ = Tip diffM _ m1 Tip = m1 - diffM f m1 m2 = hedgeDiffWithKey (f . Ord) (const LT) (const GT) m1 m2 + diffM f m1 m2 = hedgeDiff f (const LT) (const GT) m1 m2 singleHoleM (Ord k) = Empty k Root - keyM (Empty k _) = Ord k - keyM (Full k _ _ _) = Ord k beforeM a (Empty k path) = before (singletonMaybe k a) path beforeM a (Full k path l _) = before t path where t = case a of @@ -76,18 +78,20 @@ | otherwise = indexT (RightBin kx x l path) (i# -# sx#) r where !sl# = size# l !sx# = getSize# x +# sl# - indexT _ _ _ = (# error err, error err, error err #) where - err = "Error: empty trie" + indexT _ _ _ = indexFail () extractHoleM = extractHole Root where extractHole path (Bin _ kx x l r) = extractHole (LeftBin kx x path r) l `mplus` return (x, Full kx path l r) `mplus` extractHole (RightBin kx x l path) r extractHole _ _ = mzero - assignM x (Empty k path) = rebuild (singleton k x) path - assignM x (Full k path l r) = rebuild (join k x l r) path - clearM (Empty _ path) = rebuild Tip path - clearM (Full _ path l r) = rebuild (merge l r) path + assignM x (Empty k path) = rebuild (maybe Tip (singleton k) x) path + assignM x (Full k path l r) = rebuild (joinMaybe k x l r) path + + unifyM (Ord k1) a1 (Ord k2) a2 = case compare k1 k2 of + EQ -> Left $ Empty k1 Root + LT -> Right $ bin k1 a1 Tip (singleton k2 a2) + GT -> Right $ bin k1 a1 (singleton k2 a2) Tip rebuild :: Sized a => OrdMap k a -> Path k a -> OrdMap k a rebuild t Root = t @@ -104,63 +108,57 @@ singleton :: Sized a => k -> a -> OrdMap k a singleton k a = Bin (getSize# a) k a Tip Tip -traverseWithKey :: (Applicative f, Sized b) => (k -> a -> f b) -> OrdMap k a -> f (OrdMap k b) -traverseWithKey _ Tip = pure Tip -traverseWithKey f (Bin _ k a l r) = balance k <$> f k a <*> traverseWithKey f l <*> traverseWithKey f r +traverse :: (Applicative f, Sized b) => (a -> f b) -> OrdMap k a -> f (OrdMap k b) +traverse _ Tip = pure Tip +traverse f (Bin _ k a l r) = balance k <$> f a <*> traverse f l <*> traverse f r -foldrWithKey :: (k -> a -> b -> b) -> OrdMap k a -> b -> b -foldrWithKey _ Tip = id -foldrWithKey f (Bin _ k a l r) = foldrWithKey f l . f k a . foldrWithKey f r +foldr :: (a -> b -> b) -> OrdMap k a -> b -> b +foldr _ Tip = id +foldr f (Bin _ _ a l r) = foldr f l . f a . foldr f r -foldlWithKey :: (k -> b -> a -> b) -> OrdMap k a -> b -> b -foldlWithKey _ Tip = id -foldlWithKey f (Bin _ k a l r) = foldlWithKey f r . flip (f k) a . foldlWithKey f l +foldl :: (b -> a -> b) -> OrdMap k a -> b -> b +foldl _ Tip = id +foldl f (Bin _ _ a l r) = foldl f r . flip f a . foldl f l -mapWithKey :: (Ord k, Sized b) => (k -> a -> b) -> OrdMap k a -> OrdMap k b -mapWithKey f (Bin _ k a l r) = join k (f k a) (mapWithKey f l) (mapWithKey f r) -mapWithKey _ _ = Tip +fmap :: (Ord k, Sized b) => (a -> b) -> OrdMap k a -> OrdMap k b +fmap f (Bin _ k a l r) = join k (f a) (fmap f l) (fmap f r) +fmap _ _ = Tip -mapMaybe :: (Ord k, Sized b) => (k -> a -> Maybe b) -> OrdMap k a -> OrdMap k b -mapMaybe f (Bin _ k a l r) = joinMaybe k (f k a) (mapMaybe f l) (mapMaybe f r) +mapMaybe :: (Ord k, Sized b) => (a -> Maybe b) -> OrdMap k a -> OrdMap k b +mapMaybe f (Bin _ k a l r) = joinMaybe k (f a) (mapMaybe f l) (mapMaybe f r) mapMaybe _ _ = Tip -mapEither :: (Ord k, Sized b, Sized c) => EitherMap k a b c -> +mapEither :: (Ord k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> OrdMap k a -> (# OrdMap k b, OrdMap k c #) -mapEither f (Bin _ k a l r) - | (# aL, aR #) <- f k a, - (# lL, lR #) <- mapEither f l, - (# rL, rR #) <- mapEither f r - = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #) +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 #) splitLookup :: (Ord k, Sized a) => SplitMap a x -> k -> OrdMap k a -> (# OrdMap k a, Maybe x, OrdMap k a #) splitLookup f k m = case m of Tip -> (# Tip, Nothing, Tip #) Bin _ kx x l r -> case compare k kx of - LT -> case splitLookup f k l of - (# lL, ans, lR #) -> (# lL, ans, join kx x lR r #) - EQ -> case f x of - (# xL, ans, xR #) -> (# maybe l (\ xL -> insertMax kx xL l) xL, ans, - maybe r (\ xR -> insertMin kx xR r) xR #) - GT -> case splitLookup f k r of - (# rL, ans, rR #) -> (# join kx x l rL, ans, rR #) + LT -> let !(# lL, ans, lR #) = splitLookup f k l in (# lL, ans, join kx x lR r #) + EQ -> let !(# xL, ans, xR #) = f x in + (# maybe l (\ xL -> insertMax kx xL l) xL, ans, maybe r (\ xR -> insertMin kx xR r) xR #) + GT -> let !(# rL, ans, rR #) = splitLookup f k r in (# join kx x l rL, ans, rR #) isSubmap :: (Ord k, Sized a, Sized b) => LEq a b -> LEq (OrdMap k a) (OrdMap k b) isSubmap _ Tip _ = True isSubmap _ _ Tip = False -isSubmap (<=) (Bin _ kx x l r) t = case splitLookup (\ x -> (# Nothing, Just (Elem x), Nothing #)) kx t of - (# lt, found, gt #) -> case found of +isSubmap (<=) (Bin _ kx x l r) t = case found of Nothing -> False Just (Elem y) -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt + where !(# lt, found, gt #) = splitLookup (\ x -> (# Nothing, Just (Elem x), Nothing #)) kx t -fromAscList :: (Eq k, Sized a) => (k -> a -> a -> a) -> [(k, a)] -> OrdMap k a +fromAscList :: (Eq k, Sized a) => (a -> a -> a) -> [(k, a)] -> OrdMap 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 kx xx zz) xs + | kz == kx = combineEq' (kx, f xx zz) xs | otherwise = (kz,zz):combineEq' x xs fromDistinctAscList :: Sized a => [(k, a)] -> OrdMap k a @@ -182,24 +180,24 @@ buildR _ _ _ [] = error "fromDistinctAscList buildR []" buildB l k x c r zs = c (bin k x l r) zs -hedgeUnionWithKey :: (Ord k, Sized a) - => (k -> a -> a -> Maybe a) +hedgeUnion :: (Ord k, Sized a) + => (a -> a -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k a -> OrdMap k a -hedgeUnionWithKey _ _ _ t1 Tip +hedgeUnion _ _ _ t1 Tip = t1 -hedgeUnionWithKey _ cmplo cmphi Tip (Bin _ kx x l r) +hedgeUnion _ cmplo cmphi Tip (Bin _ kx x l r) = join kx x (filterGt cmplo l) (filterLt cmphi r) -hedgeUnionWithKey f cmplo cmphi (Bin _ kx x l r) t2 - = joinMaybe kx newx (hedgeUnionWithKey f cmplo cmpkx l lt) - (hedgeUnionWithKey f cmpkx cmphi r gt) +hedgeUnion f cmplo cmphi (Bin _ kx x l r) t2 + = joinMaybe kx newx (hedgeUnion f cmplo cmpkx l lt) + (hedgeUnion f cmpkx cmphi r gt) where cmpkx k = compare kx k lt = trim cmplo cmpkx t2 (found,gt) = trimLookupLo kx cmphi t2 newx = case found of Nothing -> Just x - Just (_,y) -> f kx x y + Just (_,y) -> f x y filterGt :: (Ord k, Sized a) => (k -> Ordering) -> OrdMap k a -> OrdMap k a filterGt _ Tip = Tip @@ -236,35 +234,35 @@ GT -> trimLookupLo lo cmphi r EQ -> (Just (kx,x),trim (compare lo) cmphi r) -isect :: (Ord k, Sized a, Sized b, Sized c) => IsectFunc k a b c -> OrdMap k a -> OrdMap k b -> OrdMap k c +isect :: (Ord k, Sized a, Sized b, Sized c) => (a -> b -> Maybe c) -> OrdMap k a -> OrdMap k b -> OrdMap k c isect f t1@Bin{} (Bin _ k2 x2 l2 r2) - | (# found, hole #) <- search k2 Root t1 - = let tl = isect f (beforeM Nothing hole) l2 - tr = isect f (afterM Nothing hole) r2 - in joinMaybe k2 (found >>= \ x1' -> f k2 x1' x2) tl tr + = joinMaybe k2 (found >>= \ x1' -> f x1' x2) tl tr + where !(# found, hole #) = search k2 Root t1 + tl = isect f (beforeM Nothing hole) l2 + tr = isect f (afterM Nothing hole) r2 isect _ _ _ = Tip -hedgeDiffWithKey :: (Ord k, Sized a) - => (k -> a -> b -> Maybe a) +hedgeDiff :: (Ord k, Sized a) + => (a -> b -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k b -> OrdMap k a -hedgeDiffWithKey _ _ _ Tip _ +hedgeDiff _ _ _ Tip _ = Tip -hedgeDiffWithKey _ cmplo cmphi (Bin _ kx x l r) Tip +hedgeDiff _ cmplo cmphi (Bin _ kx x l r) Tip = join kx x (filterGt cmplo l) (filterLt cmphi r) -hedgeDiffWithKey f cmplo cmphi t (Bin _ kx x l r) +hedgeDiff f cmplo cmphi t (Bin _ kx x l r) = case found of Nothing -> merge tl tr Just (ky,y) -> - case f ky y x of + case f y x of Nothing -> merge tl tr Just z -> join ky z tl tr where cmpkx k = compare kx k lt = trim cmplo cmpkx t (found,gt) = trimLookupLo kx cmphi t - tl = hedgeDiffWithKey f cmplo cmpkx lt l - tr = hedgeDiffWithKey f cmpkx cmphi gt r + tl = hedgeDiff f cmplo cmpkx lt l + tr = hedgeDiff f cmpkx cmphi gt r joinMaybe :: (Ord k, Sized a) => k -> Maybe a -> OrdMap k a -> OrdMap k a -> OrdMap k a joinMaybe kx = maybe merge (join kx) @@ -310,24 +308,22 @@ glue Tip r = r glue l Tip = l glue l r - | size# l ># size# r = case deleteFindMax (\ k a -> (# balance k a, Nothing #)) l of - (# f, l' #) -> f l' r - | otherwise = case deleteFindMin (\ k a -> (# balance k a, Nothing #)) r of - (# f, r' #) -> f l r' + | size# l ># size# r = let !(# f, l' #) = deleteFindMax (\ k a -> (# balance k a, Nothing #)) l in f l' r + | otherwise = let !(# f, r' #) = deleteFindMin (\ k a -> (# balance k a, Nothing #)) r in f l r' deleteFindMin :: Sized a => (k -> a -> (# x, Maybe a #)) -> OrdMap k a -> (# x, OrdMap k a #) deleteFindMin f t = case t of - Bin _ k x Tip r -> onUnboxed (maybe r (\ y' -> bin k y' Tip r)) (f k) x - Bin _ k x l r -> onUnboxed (\ l' -> balance k x l' r) (deleteFindMin f) l - _ -> (# error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip #) + Bin _ k x Tip r -> onSnd (maybe r (\ y' -> bin k y' Tip r)) (f k) x + Bin _ k x l r -> onSnd (\ l' -> balance k x l' r) (deleteFindMin f) l + _ -> (# error "Map.deleteFindMin: can not return the minimal element of an empty fmap", Tip #) deleteFindMax :: Sized a => (k -> a -> (# x, Maybe a #)) -> OrdMap k a -> (# x, OrdMap k a #) deleteFindMax f t = case t of - Bin _ k x l Tip -> onUnboxed (maybe l (\ y -> bin k y l Tip)) (f k) x - Bin _ k x l r -> onUnboxed (balance k x l) (deleteFindMax f) r - Tip -> (# error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip #) + Bin _ k x l Tip -> onSnd (maybe l (\ y -> bin k y l Tip)) (f k) x + Bin _ k x l r -> onSnd (balance k x l) (deleteFindMax f) r + Tip -> (# error "Map.deleteFindMax: can not return the maximal element of an empty fmap", Tip #) size# :: OrdMap k a -> Int# size# Tip = 0# diff --git a/Data/TrieMap/ProdMap.hs b/Data/TrieMap/ProdMap.hs --- a/Data/TrieMap/ProdMap.hs +++ b/Data/TrieMap/ProdMap.hs @@ -4,65 +4,65 @@ import Data.TrieMap.Sized import Data.TrieMap.TrieKey -import Data.TrieMap.Applicative import Control.Applicative -import Data.Foldable +import Data.Foldable hiding (foldlM, foldrM) +import Data.Monoid import Data.Sequence ((|>)) import qualified Data.Sequence as Seq instance (TrieKey k1, TrieKey k2) => TrieKey (k1, k2) where + (k11, k12) =? (k21, k22) = k11 =? k21 && k12 =? k22 + (k11, k12) `cmp` (k21, k22) = (k11 `cmp` k21) `mappend` (k12 `cmp` k22) + newtype TrieMap (k1, k2) a = PMap (TrieMap k1 (TrieMap k2 a)) data Hole (k1, k2) a = PHole (Hole k1 (TrieMap k2 a)) (Hole k2 a) emptyM = PMap emptyM - singletonM (k1, k2) a = PMap (singletonM k1 (singletonM k2 a)) - nullM (PMap m) = nullM m + 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 - traverseWithKeyM f (PMap m) = PMap <$> traverseWithKeyM (\ a -> traverseWithKeyM (f . (a,))) m - foldrWithKeyM f (PMap m) = foldrWithKeyM (\ a -> foldrWithKeyM (f . (a,))) m - foldlWithKeyM f (PMap m) = foldlWithKeyM (\ a -> flip (foldlWithKeyM (f . (a,)))) m - mapWithKeyM f (PMap m) = PMap (mapWithKeyM (\ a -> mapWithKeyM (f . (a,))) m) - mapMaybeM f (PMap m) = PMap (mapMaybeM g m) where - g a = guardNullM . mapMaybeM (f . (a,)) - mapEitherM f (PMap m) = both PMap PMap (mapEitherM g) m where - g a m = both guardNullM guardNullM (mapEitherM (f . (a,))) m + traverseM f (PMap m) = PMap <$> traverseM (traverseM f) m + foldrM f (PMap m) = foldrM (foldrM f) m + foldlM f (PMap m) = foldlM (flip $ foldlM 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 (\ a -> guardNullM .: unionM (f . (a,))) m1 m2) - isectM f (PMap m1) (PMap m2) = PMap (isectM (\ a -> guardNullM .: isectM (f . (a,))) m1 m2) - diffM f (PMap m1) (PMap m2) = PMap (diffM (\ a -> guardNullM .: diffM (f . (a,))) m1 m2) - fromListM f xs = PMap (mapWithKeyM (\ a (Elem xs) -> fromListM (f . (a,)) xs) - (fromListM (\ _ (Elem xs) (Elem ys) -> Elem (xs ++ ys)) (breakFst xs))) + 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) fromAscListM f xs = PMap (fromDistAscListM - [(a, fromAscListM (f . (a,)) ys) | (a, Elem ys) <- breakFst xs]) + [(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) - keyM (PHole hole1 hole2) = (keyM hole1, keyM hole2) - assignM v (PHole hole1 hole2) = PMap (assignM (assignM v hole2) hole1) - clearM (PHole hole1 hole2) = PMap (fillHoleM (guardNullM (clearM hole2)) hole1) - beforeM a (PHole hole1 hole2) - = PMap (beforeM (guardNullM (beforeM a hole2)) hole1) - afterM a (PHole hole1 hole2) - = PMap (afterM (guardNullM (afterM a hole2)) hole1) - searchM (k1, k2) (PMap m) = case searchM k1 m of - (# Nothing, hole1 #) -> (# Nothing, PHole hole1 (singleHoleM k2) #) - (# Just m', hole1 #) -> onUnboxed (PHole hole1) (searchM k2) m' - indexM i (PMap m) - | (# i', m', hole1 #) <- indexM i m, - (# i'', v, hole2 #) <- indexM i' m' - = (# i'', v, PHole hole1 hole2 #) + assignM v (PHole hole1 hole2) = PMap (assignM (assignM' v hole2) hole1) + beforeM a (PHole hole1 hole2) = PMap (beforeM (beforeM' a hole2) hole1) + afterM a (PHole hole1 hole2) = PMap (afterM (afterM' a hole2) hole1) + searchM (k1, k2) (PMap m) = onSnd (PHole hole1) (searchM' k2) m' + where !(# m', hole1 #) = searchM k1 m + indexM i (PMap m) = onThird (PHole hole1) (indexM i') m' + where !(# i', m', hole1 #) = indexM i m extractHoleM (PMap m) = do (m', hole1) <- extractHoleM m (v, hole2) <- extractHoleM m' return (v, PHole hole1 hole2) + + unifyM (k11, k12) a1 (k21, k22) a2 = case unifyM k11 (singletonM k12 a1) k21 (singletonM k22 a2) of + Left hole -> case unifyM k12 a1 k22 a2 of + Left hole' -> Left (PHole hole hole') + Right m' -> Right (PMap (assignM (Just m') hole)) + Right m -> Right (PMap m) -breakFst :: Eq k1 => [((k1, k2), a)] -> [(k1, Elem [(k2, a)])] +breakFst :: TrieKey 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 + | 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)] diff --git a/Data/TrieMap/RadixTrie.hs b/Data/TrieMap/RadixTrie.hs --- a/Data/TrieMap/RadixTrie.hs +++ b/Data/TrieMap/RadixTrie.hs @@ -1,267 +1,128 @@ -{-# LANGUAGE BangPatterns, UnboxedTuples, TupleSections, TypeFamilies, PatternGuards, MagicHash #-} +{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, MagicHash, FlexibleInstances #-} module Data.TrieMap.RadixTrie () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized --- import Data.TrieMap.Applicative import Control.Applicative import Control.Monad +import Foreign.Storable + import Data.Maybe +import Data.Monoid +import Data.Ord import Data.Foldable (foldr, foldl) +import Data.Vector.Generic hiding (Vector, cmp, foldl, foldr) +import Data.Vector (Vector) +import qualified Data.Vector as V +import qualified Data.Vector.Storable as S import Data.Traversable - -import GHC.Exts - -import Prelude hiding (lookup, foldr, foldl) +import Data.Word -data Assoc k a = Empty | Assoc [k] a -data Edge k a = Edge Int# [k] (Assoc k a) (TrieMap k (Edge k a)) -type MEdge k a = Maybe (Edge k a) +import Data.TrieMap.RadixTrie.Slice +import Data.TrieMap.RadixTrie.Edge -instance Sized (Edge k a) where - getSize# (Edge sz _ _ _) = sz +import Prelude hiding (length, and, zip, zipWith, foldr, foldl) -instance Sized a => Sized (Assoc k a) where - getSize# (Assoc _ a) = getSize# a - getSize# _ = 0# +instance TrieKey k => TrieKey (Vector k) where + ks =? ls = length ks == length ls && and (zipWith (=?) ks ls) + ks `cmp` ls = V.foldr (\ (k, l) z -> (k `cmp` l) `mappend` z) (comparing length ks ls) (zip ks ls) -data Path k a = Root - | Deep (Path k a) [k] (Assoc k a) (Hole k (Edge k a)) + newtype TrieMap (Vector k) a = Radix (MEdge Vector k a) + newtype Hole (Vector k) a = Hole (EdgeLoc Vector k a) + + emptyM = Radix Nothing + singletonM ks a = Radix (Just (singletonEdge (v2S ks) a)) + getSimpleM (Radix Nothing) = Null + getSimpleM (Radix (Just e)) = getSimpleEdge e + sizeM (Radix m) = getSize# m + lookupM ks (Radix m) = m >>= lookupEdge ks -instance TrieKey k => TrieKey [k] where - newtype TrieMap [k] a = Radix (MEdge k a) - data Hole [k] a = Hole [k] [k] (TrieMap k (Edge k a)) (Path k a) + 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 - emptyM = Radix Nothing - singletonM ks a = Radix (Just (Edge (getSize# a) ks (Assoc ks a) emptyM)) - nullM (Radix m) = isNothing m - sizeM (Radix (Just e)) = getSize# e - sizeM _ = 0# - lookupM ks (Radix m) = m >>= lookup ks - traverseWithKeyM f (Radix m) = Radix <$> traverse (traverseE f) m - foldrWithKeyM f (Radix m) z = foldr (foldrE f) z m - foldlWithKeyM f (Radix m) z = foldl (foldlE f) z m - mapWithKeyM f (Radix m) = Radix (mapWithKeyE f <$> m) - mapMaybeM f (Radix m) = Radix (m >>= mapMaybeE f) - mapEitherM _ (Radix Nothing) = (# emptyM, emptyM #) - mapEitherM f (Radix (Just m)) = both Radix Radix (mapEitherE f) m - unionM f (Radix m1) (Radix m2) = Radix (unionMaybe (unionE f) m1 m2) - isectM f (Radix m1) (Radix m2) = Radix (isectMaybe (isectE f) m1 m2) - diffM f (Radix m1) (Radix m2) = Radix (diffMaybe (diffE f) m1 m2) - isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubmapE (<=)) m1 m2 + foldrM f (Radix m) z = foldr (foldrEdge f) z m + foldlM f (Radix m) z = foldl (foldlEdge f) z m - singleHoleM ks = Hole ks ks emptyM Root - keyM (Hole ks _ _ _) = ks - beforeM a (Hole ks0 ks ts path) = before (compact (edge ks v ts)) path where - v = case a of - Nothing -> Empty - Just a -> Assoc ks0 a - before t Root = Radix t - before e (Deep path ks v tHole) = - before (compact $ edge ks v $ beforeM e tHole) path - afterM a (Hole ks0 ks ts path) = after (compact (edge ks v ts)) path where - v = case a of - Nothing -> Empty - Just a -> Assoc ks0 a - after t Root = Radix t - after e (Deep path ks v tHole) = - after (compact $ edge ks v $ afterM e tHole) path + 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 - searchM ks (Radix Nothing) = (# Nothing, singleHoleM ks #) - searchM ks (Radix (Just e)) = case searchE ks e Root of - (# v, holer #) -> (# v, holer ks #) + singleHoleM ks = Hole (singleLoc (v2S ks)) + searchM ks (Radix (Just e)) = case searchEdge (v2S ks) e Root of + (a, loc) -> (# a, Hole loc #) + searchM ks _ = (# Nothing, singleHoleM ks #) + indexM i (Radix (Just e)) = case indexEdge i e Root of + (# i', a, loc #) -> (# i', a, Hole loc #) + indexM _ (Radix Nothing) = indexFail () - indexM _ (Radix Nothing) = (# error err, error err, error err #) - where err = "Error: trie map is empty" - indexM i# (Radix (Just e)) = indexE i# e Root + assignM a (Hole loc) = Radix (fillHoleEdge a loc) - extractHoleM (Radix Nothing) = mzero - extractHoleM (Radix (Just e)) = extractHoleE Root e + extractHoleM (Radix (Just e)) = do + (a, loc) <- extractEdgeLoc e Root + return (a, Hole loc) + extractHoleM _ = mzero - assignM a (Hole ks0 ks ts path) = Radix $ rebuild (compact (edge ks (Assoc ks0 a) ts)) path + beforeM a (Hole loc) = Radix (beforeEdge a loc) + afterM a (Hole loc) = Radix (afterEdge a loc) - clearM (Hole _ ks ts path) = Radix $ rebuild (compact (edge ks Empty ts)) path - -rebuild :: (TrieKey k, Sized a) => MEdge k a -> Path k a -> MEdge k a -rebuild e (Deep path ks v tHole) = - rebuild (compact (edge ks v (fillHoleM e tHole))) path -rebuild e _ = e - -cat :: [k] -> Edge k a -> Edge k a -ks `cat` Edge sz ls v ts = Edge sz (ks ++ ls) v ts - -cons :: k -> Edge k a -> Edge k a -k `cons` Edge sz ks v ts = Edge sz (k:ks) v ts - -edge :: (TrieKey k, Sized a) => [k] -> Assoc k a -> TrieMap k (Edge k a) -> Edge k a -edge ks v ts = Edge (getSize# v +# getSize# ts) ks v ts - -compact :: TrieKey k => Edge k a -> MEdge k a -compact e@(Edge _ ks Empty ts) = case assocsM ts of - [] -> Nothing - [(l, e')] -> compact (ks `cat` (l `cons` e')) - _ -> Just e -compact e = Just e - -lookup :: (Eq k, TrieKey k) => [k] -> Edge k a -> Maybe a -lookup ks (Edge _ ls v ts) = match ks ls where - match (k:ks) (l:ls) - | k == l = match ks ls - match (k:ks) [] = lookupM k ts >>= lookup ks - match [] [] = case v of - Assoc _ a -> Just a - _ -> Nothing - match _ _ = Nothing - -traverseA :: Applicative f => ([k] -> a -> f b) -> Assoc k a -> f (Assoc k b) -traverseA f (Assoc ks a) = Assoc ks <$> f ks a -traverseA _ _ = pure Empty - -traverseE :: (Applicative f, TrieKey k, Sized b) => ([k] -> a -> f b) -> Edge k a -> f (Edge k b) -traverseE f (Edge _ ks v ts) - = edge ks <$> traverseA f v <*> traverseM (traverseE f) ts - -foldrA :: ([k] -> a -> b -> b) -> Assoc k a -> b -> b -foldrA f (Assoc ks a) = f ks a -foldrA _ _ = id - -foldlA :: ([k] -> b -> a -> b) -> b -> Assoc k a -> b -foldlA f z (Assoc ks a) = f ks z a -foldlA _ z _ = z - -foldrE :: TrieKey k => ([k] -> a -> b -> b) -> Edge k a -> b -> b -foldrE f (Edge _ _ v ts) z = foldrA f v (foldr (foldrE f) z ts) - -foldlE :: TrieKey k => ([k] -> b -> a -> b) -> b -> Edge k a -> b -foldlE f z (Edge _ _ v ts) = foldl (foldlE f) (foldlA f z v) ts - -mapWithKeyA :: ([k] -> a -> b) -> Assoc k a -> Assoc k b -mapWithKeyA f (Assoc ks a) = Assoc ks (f ks a) -mapWithKeyA _ _ = Empty - -mapWithKeyE :: (TrieKey k, Sized b) => ([k] -> a -> b) -> Edge k a -> Edge k b -mapWithKeyE f (Edge _ ks v ts) = edge ks (mapWithKeyA f v) (fmapM (mapWithKeyE f) ts) - -mapMaybeA :: ([k] -> a -> Maybe b) -> Assoc k a -> Assoc k b -mapMaybeA f (Assoc ks a) = maybe Empty (Assoc ks) (f ks a) -mapMaybeA _ _ = Empty - -mapMaybeE :: (TrieKey k, Sized b) => ([k] -> a -> Maybe b) -> Edge k a -> MEdge k b -mapMaybeE f (Edge _ ks v ts) = compact (edge ks (mapMaybeA f v) - (mapMaybeM (const $ mapMaybeE f) ts)) - -mapEitherA :: ([k] -> a -> (# Maybe b, Maybe c #)) -> Assoc k a -> (# Assoc k b, Assoc k c #) -mapEitherA f (Assoc ks a) = case f ks a of - (# vL, vR #) -> (# maybe Empty (Assoc ks) vL, maybe Empty (Assoc ks) vR #) -mapEitherA _ _ = (# Empty, Empty #) - -mapEitherE :: (TrieKey k, Sized b, Sized c) => ([k] -> a -> (# Maybe b, Maybe c #)) -> Edge k a -> - (# MEdge k b, MEdge k c #) -mapEitherE f (Edge _ ks v ts) = case mapEitherA f v of - (# vL, vR #) -> case mapEitherM (\ _ -> mapEitherE f) ts of - (# tsL, tsR #) -> (# compact (edge ks vL tsL), compact (edge ks vR tsR) #) - -unionE :: (TrieKey k, Sized a) => ([k] -> a -> a -> Maybe a) -> Edge k a -> Edge k a -> MEdge k a -unionE f (Edge szK# ks0 vK tsK) (Edge szL# ls0 vL tsL) = match 0 ks0 ls0 where - match !i (k:ks) (l:ls) = case compare k l of - EQ -> match (i+1) ks ls - LT -> Just $ Edge (szK# +# szL#) (take i ks0) Empty (fromDistAscListM - [(k, Edge szK# ks vK tsK), (l, Edge szL# ls vL tsL)]) - GT -> Just $ Edge (szK# +# szL#) (take i ks0) Empty (fromDistAscListM - [(l, Edge szL# ls vL tsL), (k, Edge szK# ks vK tsK)]) - match _ [] (l:ls) = compact (edge ks0 vK (alterM g l tsK)) where - g (Just eK') = unionE f eK' (Edge szL# ls vL tsL) - g Nothing = Just (Edge szL# ls vL tsL) - match _ (k:ks) [] = compact (edge ls0 vL (alterM g k tsL)) where - g Nothing = Just (Edge szK# ks vK tsK) - g (Just eL') = unionE f (Edge szK# ks vK tsK) eL' - match _ [] [] = compact (edge ls0 (unionA f vK vL) (unionM (const $ unionE f) tsK tsL)) - -unionA :: ([k] -> a -> a -> Maybe a) -> Assoc k a -> Assoc k a -> Assoc k a -unionA f (Assoc ks v1) (Assoc _ v2) = maybe Empty (Assoc ks) (f ks v1 v2) -unionA _ Empty v = v -unionA _ v Empty = v + unifyM ks1 a1 ks2 a2 = either (Left . Hole) (Right . Radix . Just) (unifyEdge (v2S ks1) a1 (v2S ks2) a2) -isectE :: (TrieKey k, Sized c) => ([k] -> a -> b -> Maybe c) -> Edge k a -> Edge k b -> MEdge k c -isectE f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where - match (k:ks) (l:ls) - | k == l = match ks ls - match (k:ks) [] = do eL' <- lookupM k tsL - cat ls0 <$> cons k <$> isectE f (Edge szK ks vK tsK) eL' - match [] (l:ls) = do eK' <- lookupM l tsK - cat ks0 <$> cons l <$> isectE f eK' (Edge szL ls vL tsL) - match [] [] = compact (edge ks0 (isectA f vK vL) (isectM (const $ isectE f) tsK tsL)) - match _ _ = Nothing +type WordVec = S.Vector Word -isectA :: ([k] -> a -> b -> Maybe c) -> Assoc k a -> Assoc k b -> Assoc k c -isectA f (Assoc ks a) (Assoc _ b) = maybe Empty (Assoc ks) (f ks a b) -isectA _ _ _ = Empty +vZipWith :: (Storable a, Storable b) => (a -> b -> c) -> S.Vector a -> S.Vector b -> Vector c +vZipWith f xs ys = V.zipWith f (convert xs) (convert ys) -diffE :: (TrieKey k, Sized a) => ([k] -> a -> b -> Maybe a) -> Edge k a -> Edge k b -> MEdge k a -diffE f eK@(Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where - match (k:ks) (l:ls) - | k == l = match ks ls - match (k:ks) [] - | Just eL' <- lookupM k tsL - = cat ls0 . cons k <$> diffE f (Edge szK ks vK tsK) eL' - match [] (l:ls) - = compact (edge ks0 vK (alterM (>>= g) l tsK)) - where g eK' = diffE f eK' (Edge szL ls vL tsL) - match [] [] = compact (edge ks0 (diffA f vK vL) (diffM (const $ diffE f) tsK tsL)) - match _ _ = Just eK +instance TrieKey (S.Vector Word) where + ks =? ls = length ks == length ls && and (vZipWith (=?) ks ls) + ks `cmp` ls = V.foldr (\ (k, l) z -> (k `cmp` l) `mappend` z) (comparing length ks ls) (vZipWith (,) ks ls) -diffA :: ([k] -> a -> b -> Maybe a) -> Assoc k a -> Assoc k b -> Assoc k a -diffA f (Assoc ks a) (Assoc _ b) = maybe Empty (Assoc ks) (f ks a b) -diffA _ a@Assoc{} Empty = a -diffA _ Empty _ = Empty + newtype TrieMap WordVec a = WRadix (MEdge S.Vector Word a) + newtype Hole WordVec a = WHole (EdgeLoc S.Vector Word a) + + emptyM = WRadix Nothing + singletonM ks a = WRadix (Just (singletonEdge (v2S ks) a)) + getSimpleM (WRadix Nothing) = Null + getSimpleM (WRadix (Just e)) = getSimpleEdge e + sizeM (WRadix m) = getSize# m + lookupM ks (WRadix m) = m >>= lookupEdge ks -isSubmapE :: TrieKey k => LEq a b -> LEq (Edge k a) (Edge k b) -isSubmapE (<=) (Edge szK ks vK tsK) (Edge _ ls vL tsL) = match ks ls where - match (k:ks) (l:ls) - | k == l = match ks ls - match (k:ks) [] - | Just eL' <- lookupM k tsL - = isSubmapE (<=) (Edge szK ks vK tsK) eL' - match [] [] = subA (<=) vK vL && isSubmapM (isSubmapE (<=)) tsK tsL - match _ _ = False + 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 -subA :: LEq a b -> LEq (Assoc k a) (Assoc k b) -subA (<=) (Assoc _ a) (Assoc _ b) = a <= b -subA _ Assoc{} Empty = False -subA _ Empty _ = True + foldrM f (WRadix m) z = foldr (foldrEdge f) z m + foldlM f (WRadix m) z = foldl (foldlEdge f) z m -searchE :: TrieKey k => [k] -> Edge k a -> Path k a -> (# Maybe a, [k] -> Hole [k] a #) -searchE ks0 (Edge sz ls0 v ts) path = match 0 ks0 ls0 where - match !_ [] [] = (# assocToMaybe v, \ k0 -> Hole k0 ls0 ts path #) - match _ (k:ks) [] = case searchM k ts of - (# Just e', tHole #) -> searchE ks e' (Deep path ls0 v tHole) - (# Nothing, tHole #) -> (# Nothing, \ k0 -> Hole k0 ks emptyM (Deep path ls0 v tHole) #) - match i [] (l:ls) = (# Nothing, \ k0 -> Hole k0 (take i ls0) (singletonM l (Edge sz ls v ts)) path #) - match i (k:ks) (l:ls) - | k == l = match (i+1) ks ls - | (# _, kHole #) <- searchM k (singletonM l (Edge sz ls v ts)) - = (# Nothing, \ k0 -> Hole k0 ks emptyM (Deep path (take i ls0) Empty kHole) #) + 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 -assocToMaybe :: Assoc k a -> Maybe a -assocToMaybe (Assoc _ a) = Just a -assocToMaybe _ = Nothing + singleHoleM ks = WHole (singleLoc (v2S ks)) + searchM ks (WRadix (Just e)) = case searchEdge (v2S ks) e Root of + (a, loc) -> (# a, WHole loc #) + searchM ks _ = (# Nothing, singleHoleM ks #) + indexM i (WRadix (Just e)) = case indexEdge i e Root of + (# i', a, loc #) -> (# i', a, WHole loc #) + indexM _ (WRadix Nothing) = indexFail () -indexE :: (TrieKey k, Sized a) => Int# -> Edge k a -> Path k a -> (# Int#, a, Hole [k] a #) -indexE i# (Edge _ ks Empty ts) path - | (# i'#, e, tHole #) <- indexM i# ts - = indexE i'# e (Deep path ks Empty tHole) -indexE i# (Edge _ ks v@(Assoc ks0 a) ts) path - | i# <# sa# = (# i#, a, Hole ks0 ks ts path #) - | (# i'#, e, tHole #) <- indexM (i# -# sa#) ts - = indexE i'# e (Deep path ks v tHole) - where !sa# = getSize# a + assignM a (WHole loc) = WRadix (fillHoleEdge a loc) + + extractHoleM (WRadix (Just e)) = do + (a, loc) <- extractEdgeLoc e Root + return (a, WHole loc) + extractHoleM _ = mzero -extractHoleE :: (TrieKey k, MonadPlus m) => Path k a -> Edge k a -> m (a, Hole [k] a) -extractHoleE path (Edge _ ks v ts) = case v of - Empty -> tsHoles - Assoc ks0 a -> return (a, Hole ks0 ks ts path) `mplus` tsHoles - where tsHoles = do (e, tHole) <- extractHoleM ts - extractHoleE (Deep path ks v tHole) e + beforeM a (WHole loc) = WRadix (beforeEdge a loc) + afterM a (WHole loc) = WRadix (afterEdge a loc) + + unifyM ks1 a1 ks2 a2 = either (Left . WHole) (Right . WRadix . Just) (unifyEdge (v2S ks1) a1 (v2S ks2) a2) diff --git a/Data/TrieMap/RadixTrie/Edge.hs b/Data/TrieMap/RadixTrie/Edge.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/RadixTrie/Edge.hs @@ -0,0 +1,269 @@ +{-# LANGUAGE MagicHash, BangPatterns, UnboxedTuples, PatternGuards, CPP #-} +{-# OPTIONS -funbox-strict-fields #-} +module Data.TrieMap.RadixTrie.Edge where + +import Data.TrieMap.Sized +import Data.TrieMap.TrieKey +import Data.TrieMap.RadixTrie.Slice +import Data.TrieMap.IntMap () +import Data.TrieMap.Applicative () + +import Control.Applicative +import Control.Monad +import Data.Word +import Data.Traversable +import Data.Foldable (foldr, foldl) + +import Data.Vector.Generic hiding (indexM, cmp, foldr, foldl) +import qualified Data.Vector +import qualified Data.Vector.Storable +import Prelude hiding (length, foldr, foldl, zip, take) + +import GHC.Exts + +#define V(f) f (Data.Vector.Vector) (k) +#define U(f) f (Data.Vector.Storable.Vector) (Word) + +type Branch v k a = TrieMap k (Edge v k a) +data Edge v k a = + Edge Int# !(Slice v k) !(Maybe a) (Branch v k a) +data EdgeLoc v k a = Loc !(Slice v k) (Branch v k a) (Path v k a) +data Path v k a = Root + | Deep (Path v k a) !(Slice v k) !(Maybe a) (Hole k (Edge v k a)) +type MEdge v k a = Maybe (Edge v k a) + +instance Sized (Edge v k a) where + getSize# (Edge s# _ _ _) = s# + +{-# SPECIALIZE singleLoc :: U(Slice) -> U(EdgeLoc) a #-} +singleLoc :: TrieKey k => Slice v k -> EdgeLoc v k a +singleLoc ks = Loc ks emptyM Root + +{-# SPECIALIZE singletonEdge :: Sized a => U(Slice) -> a -> U(Edge) a #-} +singletonEdge :: (TrieKey k, Sized a) => Slice v k -> a -> Edge v k a +singletonEdge ks a = edge ks (Just a) emptyM + +{-# SPECIALIZE getSimpleEdge :: U(Edge) a -> Simple a #-} +getSimpleEdge :: TrieKey k => Edge v k a -> Simple a +getSimpleEdge (Edge _ _ v ts) + | nullM ts = maybe Null Singleton v + | otherwise = NonSimple + +{-# SPECIALIZE edge :: Sized a => U(Slice) -> Maybe a -> U(Branch) a -> U(Edge) a #-} +edge :: (TrieKey k, Sized a) => Slice v k -> Maybe a -> Branch v k a -> Edge v k a +edge ks v ts = Edge (getSize# v +# sizeM ts) ks v ts + +{-# INLINE compact #-} +-- TODO: figure out a way to GC dead keys +compact :: TrieKey k => Edge v k a -> MEdge v k a +compact e@(Edge _ ks Nothing ts) = case getSimpleM ts of + Null -> Nothing + Singleton e' -> Just (unDropEdge (len ks + 1) e') + _ -> Just e +compact e = Just e + +dropEdge :: Int -> Edge v k a -> Edge v k a +dropEdge n (Edge s# ks v ts) = Edge s# (dropSlice n ks) v ts + +unDropEdge :: Int -> Edge v k a -> Edge v k a +unDropEdge n (Edge s# ks v ts) = Edge s# (unDropSlice n ks) v ts + +{-# SPECIALIZE lookupEdge :: TrieKey k => V() -> V(Edge) a -> Maybe a #-} +{-# SPECIALIZE lookupEdge :: U() -> U(Edge) a -> Maybe a #-} +lookupEdge :: (TrieKey k, Vector v k) => v k -> Edge v k a -> Maybe a +lookupEdge = lookupE where + lookupE !ks (Edge _ ls v ts) = if kLen < lLen then Nothing else matchSliceV matcher matches ks ls where + !kLen = length ks + !lLen = len ls + matcher k l z + | k =? l = z + | otherwise = Nothing + matches _ _ + | kLen == lLen = v + | otherwise = do e' <- lookupM (ks `unsafeIndex` lLen) ts + lookupE (unsafeDrop (lLen + 1) ks) e' + +{-# SPECIALIZE searchEdge :: TrieKey k => V(Slice) -> V(Edge) a -> V(Path) a -> (Maybe a, V(EdgeLoc) a) #-} +{-# SPECIALIZE searchEdge :: U(Slice) -> U(Edge) a -> U(Path) a -> (Maybe a, U(EdgeLoc) a) #-} +searchEdge :: (TrieKey k, Vector v k) => Slice v k -> Edge v k a -> Path v k a -> (Maybe a, EdgeLoc v k a) +searchEdge = searchE where + searchE !ks e@(Edge _ ls v ts) path = iMatchSlice matcher matches ks ls where + matcher i k l z + | k =? l = z + | (# _, tHole #) <- searchM k (singletonM l (dropEdge (i+1) e)) + = (Nothing, Loc (dropSlice (i+1) ks) emptyM (Deep path (takeSlice i ls) Nothing tHole)) + matches kLen lLen = case compare kLen lLen of + EQ -> (v, Loc ls ts path) + LT -> let (lPre, !l, _) = splitSlice kLen ls in + (Nothing, Loc lPre (singletonM l (dropEdge (kLen + 1) e)) path) + GT -> let (_, !k, ks') = splitSlice lLen ks in case searchM k ts of + (# Nothing, tHole #) -> (Nothing, Loc ks' emptyM (Deep path ls v tHole)) + (# Just e', tHole #) -> searchE ks' e' (Deep path ls v tHole) + +{-# SPECIALIZE mapEdge :: Sized b => (a -> b) -> U(Edge) a -> U(Edge) b #-} +mapEdge :: (TrieKey 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 mapMaybeEdge :: Sized b => (a -> Maybe b) -> U(Edge) a -> U(MEdge) b #-} +mapMaybeEdge :: (TrieKey k, Sized b) => (a -> Maybe b) -> Edge v k a -> MEdge v k b +mapMaybeEdge f = mapMaybeE where + mapMaybeE (Edge _ ks v ts) = compact (edge ks (v >>= f) (mapMaybeM mapMaybeE ts)) + +{-# SPECIALIZE mapEitherEdge :: (Sized b, Sized c) => + (a -> (# Maybe b, Maybe c #)) -> U(Edge) a -> (# U(MEdge) b, U(MEdge) c #) #-} +mapEitherEdge :: (TrieKey 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) = (# compact (edge ks vL tsL), compact (edge ks vR tsR) #) + where !(# vL, vR #) = mapEitherMaybe f v + !(# tsL, tsR #) = mapEitherM mapEitherE ts + +{-# SPECIALIZE traverseEdge :: (Applicative f, Sized b) => + (a -> f b) -> U(Edge) a -> f (U(Edge) b) #-} +traverseEdge :: (TrieKey k, Applicative f, Sized b) => + (a -> f b) -> Edge v k a -> f (Edge v k b) +traverseEdge f = traverseE where + traverseE (Edge _ ks v ts) = edge ks <$> traverse f v <*> traverseM traverseE ts + +{-# SPECIALIZE foldrEdge :: (a -> b -> b) -> U(Edge) a -> b -> b #-} +foldrEdge :: TrieKey k => (a -> b -> b) -> Edge v k a -> b -> b +foldrEdge f = foldrE where + foldrE (Edge _ _ v ts) z = foldr f (foldrM foldrE ts z) v + +foldlEdge :: TrieKey k => (b -> a -> b) -> b -> Edge v k a -> b +foldlEdge f = foldlE where + foldlE z (Edge _ _ v ts) = foldlM foldlE ts (foldl f z v) + +{-# SPECIALIZE rebuild :: Sized a => U(MEdge) a -> U(Path) a -> U(MEdge) a #-} +rebuild :: (TrieKey k, Sized a) => MEdge v k a -> Path v k a -> MEdge v k a +rebuild e Root = e +rebuild e (Deep path ks v tHole) = rebuild (compact $ edge ks v $ assignM e tHole) path + +{-# SPECIALIZE fillHoleEdge :: Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} +fillHoleEdge :: (TrieKey k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a +fillHoleEdge v (Loc ks ts path) = rebuild (compact (edge ks v ts)) path + +{-# SPECIALIZE unionEdge :: (TrieKey k, Sized a) => + (a -> a -> Maybe a) -> V(Edge) a -> V(Edge) a -> V(MEdge) a #-} +{-# SPECIALIZE unionEdge :: Sized a => + (a -> a -> Maybe a) -> U(Edge) a -> U(Edge) a -> U(MEdge) a #-} +unionEdge :: (TrieKey k, Vector v k, Sized a) => + (a -> a -> Maybe a) -> Edge v k a -> Edge v k a -> MEdge v k a +unionEdge f = unionE where + eK@(Edge _ ks0 vK tsK) `unionE` eL@(Edge _ ls0 vL tsL) = iMatchSlice matcher matches ks0 ls0 where + matcher i k l z = case unifyM k eK' l eL' of + Left{} -> z + Right ts -> Just (edge (takeSlice i ks0) Nothing ts) + where eK' = dropEdge (i+1) eK + eL' = dropEdge (i+1) eL + matches kLen lLen = case compare kLen lLen of + EQ -> compact $ edge ks0 (unionMaybe f vK vL) $ unionM unionE tsK tsL + LT -> let eL' = dropEdge (kLen + 1) eL; l = ls0 !$ kLen; !(# eK', holeKT #) = searchM l tsK + in compact $ edge ks0 vK $ assignM (maybe (Just eL') (`unionE` eL') eK') holeKT + GT -> let eK' = dropEdge (lLen + 1) eK; k = ks0 !$ lLen; !(# eL', holeLT #) = searchM k tsL + in compact $ edge ls0 vL $ assignM (maybe (Just eK') (eK' `unionE`) eL') holeLT + +{-# SPECIALIZE isectEdge :: (TrieKey k, Sized c) => + (a -> b -> Maybe c) -> V(Edge) a -> V(Edge) b -> V(MEdge) c #-} +{-# SPECIALIZE isectEdge :: Sized c => + (a -> b -> Maybe c) -> U(Edge) a -> U(Edge) b -> U(MEdge) c #-} +isectEdge :: (TrieKey k, Vector v k, Sized c) => + (a -> b -> Maybe c) -> Edge v k a -> Edge v k b -> MEdge v k c +isectEdge f = isectE where + eK@(Edge _ ks0 vK tsK) `isectE` 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' <- lookupM l tsK + let eL' = dropEdge (kLen + 1) eL + unDropEdge (kLen + 1) <$> eK' `isectE` eL' + GT -> let k = ks0 !$ lLen in do + eL' <- lookupM k tsL + let eK' = dropEdge (lLen + 1) eK + 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 #-} +{-# SPECIALIZE diffEdge :: Sized a => + (a -> b -> Maybe a) -> U(Edge) a -> U(Edge) b -> U(MEdge) a #-} +diffEdge :: (TrieKey k, Vector v k, Sized a) => + (a -> b -> Maybe a) -> Edge v k a -> Edge v k b -> MEdge v k a +diffEdge f = diffE where + eK@(Edge _ ks0 vK tsK) `diffE` 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 -> compact $ edge ks0 (diffMaybe f vK vL) $ diffM diffE tsK tsL + LT -> let l = ls0 !$ kLen; eL' = dropEdge (kLen + 1) eL in case searchM l tsK of + (# Nothing, _ #) -> Just eK + (# Just eK', holeKT #) -> compact $ edge ks0 vK $ assignM (eK' `diffE` eL') holeKT + GT -> let k = ks0 !$ lLen; eK' = dropEdge (lLen + 1) eK in case lookupM k tsL of + Nothing -> Just eK + Just eL' -> fmap (unDropEdge (lLen + 1)) (eK' `diffE` eL') + +{-# SPECIALIZE isSubEdge :: TrieKey k => LEq a b -> LEq (V(Edge) a) (V(Edge) b) #-} +{-# SPECIALIZE isSubEdge :: LEq a b -> LEq (U(Edge) a) (U(Edge) b) #-} +isSubEdge :: (TrieKey k, Vector v k) => LEq a b -> LEq (Edge v k a) (Edge v k b) +isSubEdge (<=) = isSubE where + eK@(Edge _ ks0 vK tsK) `isSubE` (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 case lookupM k tsL of + Nothing -> False + Just eL' -> isSubE (dropEdge (lLen + 1) eK) eL' + +{-# SPECIALIZE beforeEdge :: Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} +beforeEdge :: (TrieKey k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a +beforeEdge v (Loc ks ts path) = buildBefore (compact (edge ks v ts)) path where + buildBefore !e Root + = e + buildBefore e (Deep path ks v tHole) + = buildBefore (compact $ edge ks v $ beforeM e tHole) path + +{-# SPECIALIZE afterEdge :: Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} +afterEdge :: (TrieKey k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a +afterEdge v (Loc ks ts path) = buildAfter (compact (edge ks v ts)) path where + buildAfter !e Root + = e + buildAfter e (Deep path ks v tHole) + = buildAfter (compact $ edge ks v $ afterM e tHole) path + +{-# SPECIALIZE extractEdgeLoc :: MonadPlus m => U(Edge) a -> U(Path) a -> m (a, U(EdgeLoc) a) #-} +extractEdgeLoc :: (TrieKey k, 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 + 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) + +{-# SPECIALIZE indexEdge :: Sized a => Int# -> U(Edge) a -> U(Path) a -> (# Int#, a, U(EdgeLoc) a #) #-} +indexEdge :: (TrieKey k, Sized a) => Int# -> Edge v k a -> Path v k a -> (# Int#, a, EdgeLoc v k a #) +indexEdge = indexE where + indexE i# (Edge _ ks v@(Just a) ts) path + | 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 + indexE i# (Edge _ ks Nothing ts) path + = indexE i'# e' (Deep path ks Nothing tHole) + where !(# i'#, e', tHole #) = indexM i# ts + +{-# SPECIALIZE unifyEdge :: (TrieKey k, Sized a) => V(Slice) -> a -> V(Slice) -> a -> Either (V(EdgeLoc) a) (V(Edge) a) #-} +{-# SPECIALIZE unifyEdge :: Sized a => U(Slice) -> a -> U(Slice) -> a -> Either (U(EdgeLoc) a) (U(Edge) a) #-} +unifyEdge :: (Vector v k, TrieKey k, Sized a) => Slice v k -> a -> Slice v k -> a -> Either (EdgeLoc v k a) (Edge v k a) +unifyEdge ks1 a1 ks2 a2 = iMatchSlice matcher matches ks1 ks2 where + matcher !i k1 k2 z = + case unifyM k1 (singletonEdge (dropSlice (i+1) ks1) a1) k2 (singletonEdge (dropSlice (i+1) ks2) a2) of + Left{} -> z + Right ts -> Right (edge (takeSlice i ks1) Nothing ts) + matches len1 len2 = case compare len1 len2 of + LT -> let (_,k2,ks2') = splitSlice len1 ks2 in + Right (edge ks1 (Just a1) (singletonM k2 (singletonEdge ks2' a2))) + GT -> let (_,k1,ks1') = splitSlice len2 ks1 in + Right (edge ks2 (Just a2) (singletonM k1 (singletonEdge ks1' a1))) + _ -> Left (singleLoc ks1) diff --git a/Data/TrieMap/RadixTrie/Slice.hs b/Data/TrieMap/RadixTrie/Slice.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/RadixTrie/Slice.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE BangPatterns #-} +{-# OPTIONS -funbox-strict-fields #-} +module Data.TrieMap.RadixTrie.Slice where + +import Control.Exception (assert) +import Data.Vector.Generic +import qualified Data.Vector as V + +import Prelude hiding (length, zip, foldr) + +data Slice v a = Slice {sliceSrc :: v a, _sliceIx :: !Int, len :: !Int} + +{-# INLINE splitSlice #-} +splitSlice :: Vector v a => Int -> Slice v a -> (Slice v a, a, Slice v a) +splitSlice !i !slice = (takeSlice i slice, slice !$ i, dropSlice (i+1) slice) + +takeSlice :: Int -> Slice v a -> Slice v a +takeSlice !n (Slice xs i _) = Slice xs i n + +dropSlice :: Int -> Slice v a -> Slice v a +dropSlice !m (Slice xs i n) = assert (n >= m) $ Slice xs (i+m) (n-m) + +unDropSlice :: Int -> Slice v a -> Slice v a +unDropSlice !m (Slice xs i n) = assert (i >= m) $ Slice xs (i-m) (n+m) + +{-# INLINE s2V #-} +s2V :: Vector v a => Slice v a -> v a +s2V (Slice xs i n) = assert (i >= 0) $ assert (i + n < length xs) $ unsafeSlice i n xs + +{-# INLINE v2S #-} +v2S :: Vector v a => v a -> Slice v a +v2S xs = Slice xs 0 (length xs) + +{-# INLINE matchSliceV #-} +matchSliceV :: (Vector v a, Vector v b) => (a -> b -> z -> z) -> (Int -> Int -> z) -> v a -> Slice v b -> z +matchSliceV f z !xs !ys = foldr (\ (a, b) -> f a b) (z (length xs) (len ys)) (V.zip (convert xs) (convert $ s2V ys)) + +{-# INLINE matchSlice #-} +matchSlice :: (Vector v a, Vector v b) => (a -> b -> z -> z) -> (Int -> Int -> z) -> Slice v a -> Slice v b -> z +matchSlice f z !xs !ys = foldr (\ (a, b) -> f a b) (z (len xs) (len ys)) (V.zip (convert $ s2V xs) (convert $ s2V ys)) + +{-# INLINE iMatchSlice #-} +iMatchSlice :: (Vector v a, Vector v b) => (Int -> a -> b -> z -> z) -> (Int -> Int -> z) -> Slice v a -> Slice v b -> z +iMatchSlice f z !xs !ys = ifoldr (\ i (a, b) -> f i a b) (z (len xs) (len ys)) (V.zip (convert $ s2V xs) (convert $ s2V ys)) + +{-# INLINE (!$) #-} +(!$) :: Vector v a => Slice v a -> Int -> a +Slice xs i n !$ j = assert (j >= 0 && j < n) $ unsafeIndex xs (i + j) diff --git a/Data/TrieMap/Rep.hs b/Data/TrieMap/Rep.hs deleted file mode 100644 --- a/Data/TrieMap/Rep.hs +++ /dev/null @@ -1,25 +0,0 @@ -{-# LANGUAGE UndecidableInstances, FlexibleContexts, TypeFamilies, KindSignatures #-} - -module Data.TrieMap.Rep where - -class Repr a where - type Rep a - toRep :: a -> Rep a - fromRep :: Rep a -> a - -class Functor (RepT f) => ReprT f where - type RepT f :: * -> * - toRepT :: f a -> RepT f a - fromRepT :: RepT f a -> f a - toRepTMap :: (a -> b) -> f a -> RepT f b - fromRepTMap :: (b -> a) -> RepT f b -> f a - - toRepT = toRepTMap id - fromRepT = fromRepTMap id - toRepTMap f = fmap f . toRepT - fromRepTMap f = fromRepT . fmap f - -{-# RULES - "toRep/fromRep" forall x . toRep (fromRep x) = x; --- "fromRep/toRep" forall x . fromRep (toRep x) = x; - #-} diff --git a/Data/TrieMap/Rep/Instances.hs b/Data/TrieMap/Rep/Instances.hs deleted file mode 100644 --- a/Data/TrieMap/Rep/Instances.hs +++ /dev/null @@ -1,188 +0,0 @@ -{-# LANGUAGE RankNTypes, FlexibleContexts, UndecidableInstances, TypeFamilies, TypeOperators, TemplateHaskell, NPlusKPatterns #-} -{-# OPTIONS -funbox-strict-fields #-} - -module Data.TrieMap.Rep.Instances() where - -import Data.TrieMap.Rep -import Data.TrieMap.Rep.TH -import Data.TrieMap.Modifiers - -import Data.Char -import Data.Int -import Data.Word -import Data.Foldable (toList) -import Data.Bits -import qualified Data.IntSet as ISet -import qualified Data.IntMap as IMap -import Data.ByteString hiding (map) -import qualified Data.ByteString as BS - -import Data.Sequence ((|>)) -import qualified Data.Sequence as Seq -import qualified Data.Foldable as Fold - -import qualified Data.Map as Map -import qualified Data.Set as Set - -import Prelude hiding (concat, take, length) - -type Pair a = (,) a -type Sum a = Either a - -instance ReprT Rev where - type RepT Rev = Rev - toRepTMap = fmap - fromRepTMap = fmap - -genRepr [t| Rev |] - -instance ReprT [] where - type RepT [] = [] - toRepTMap = map - fromRepTMap = map - -genRepr [t| [] |] - -genTupleRepr 2 -genTupleRepr 3 -genTupleRepr 4 -genTupleRepr 5 -genTupleRepr 6 -genTupleRepr 7 -genTupleRepr 8 - -instance (Repr a, Repr b) => Repr (Either a b) where - type Rep (Either a b) = Either (Rep a) (Rep b) - toRep (Left a) = Left (toRep a) - toRep (Right b) = Right (toRep b) - fromRep (Left a) = Left (fromRep a) - fromRep (Right b) = Right (fromRep b) - -instance Repr Char where - type Rep Char = Word32 - toRep = fromIntegral . ord - fromRep = chr . fromIntegral - -instance Repr () where - type Rep () = () - toRep _ = () - fromRep _ = () - -instance Repr Int where - type Rep Int = Rep Int32 - toRep = toSigned - fromRep = fromSigned - -instance Repr Word8 where - type Rep Word8 = Word32 - toRep = fromIntegral - fromRep = fromIntegral - -instance Repr Word16 where - type Rep Word16 = Word32 - toRep = fromIntegral - fromRep = fromIntegral - -instance Repr Word where - type Rep Word = Word32 - toRep = fromIntegral - fromRep = fromIntegral - -instance Repr Int8 where - type Rep Int8 = Rep Int32 - toRep = toSigned - fromRep = fromSigned - -instance Repr Int16 where - type Rep Int16 = Rep Int32 - toRep = toSigned - fromRep = fromSigned - -instance Repr Int32 where - type Rep Int32 = Sum (Rev Word32) Word32 - toRep = toSigned - fromRep = fromSigned - -instance Repr Word64 where - type Rep Word64 = Pair Word32 Word32 - toRep x = (fromIntegral (x `shiftR` 32), fromIntegral x) - fromRep (x, y) = fromIntegral x `shiftL` 32 .|. fromIntegral y - -instance Repr Int64 where - type Rep Int64 = Sum (Rev (Rep Word64)) (Rep Word64) - toRep x | x < 0 = Left (Rev (toRep' (fromIntegral (-x)))) - | otherwise = Right (toRep' (fromIntegral x)) - where toRep' = toRep :: Word64 -> Rep Word64 - fromRep (Left (Rev x)) = - fromIntegral ((fromRep :: Rep Word64 -> Word64) x) - fromRep (Right x) = fromIntegral ((fromRep :: Rep Word64 -> Word64) x) - -{-# INLINE toSigned #-} -toSigned :: Integral a => a -> Sum (Rev Word32) Word32 -toSigned x - | x < 0 = Left (Rev (fromIntegral (-x))) - | otherwise = Right (fromIntegral x) - -{-# INLINE fromSigned #-} -fromSigned :: Integral a => Sum (Rev Word32) Word32 -> a -fromSigned = either (\ (Rev x) -> - fromIntegral x) fromIntegral - -instance Repr Word32 where - type Rep Word32 = Word32 - toRep = id - fromRep = id - -instance Repr ByteString where - type Rep ByteString = ([Word32], Word32) - toRep xs = (toList64 xs, fromIntegral (length xs)) - fromRep (xs, n) = case xs of - [] -> BS.empty - (x:xs) -> fst (unfoldrN (fromIntegral n) toBlock (W (Words 3 x) xs)) - -data Words = Words !Int !Word32 -data Words' = W !Words [Word32] - -toList64 :: ByteString -> [Word32] -toList64 xs = case BS.foldl' c (Words 4 0, Seq.empty) xs of - (Words _ w32, ys) -> toList ys ++ [w32] - where (Words 0 w, xs) `c` w8 - = (Words 3 (w .|. sL w8 24), xs |> w) - (Words i' w, xs) `c` w8 - = let !i = i' - 1 in (Words i (w .|. sL w8 (8 * i)), xs) - sL :: Word8 -> Int -> Word32 - w `sL` x = fromIntegral w `shiftL` x - -toBlock :: Words' -> Maybe (Word8, Words') -toBlock (W (Words i0@(i+1) w) xs) = Just (extract w (8 * i0), (W (Words i w) xs)) - where extract :: Word32 -> Int -> Word8 - extract w x = fromIntegral (w `shiftR` x) -toBlock (W (Words 0 w) (x:xs)) = Just (fromIntegral w, (W (Words 3 x) xs)) -toBlock _ = Nothing - -instance ReprT Set.Set where - type RepT Set.Set = [] - toRepTMap f s = Fold.foldr ((:) . f) [] s - fromRepTMap f xs = Set.fromDistinctAscList [f x | x <- xs] - -genRepr [t| Set.Set |] - -instance (Repr k, Repr a) => Repr (Map.Map k a) where - type Rep (Map.Map k a) = [(Rep k, Rep a)] - toRep m = [(toRep k, toRep a) | (k, a) <- Map.assocs m] - fromRep xs = Map.fromDistinctAscList [(fromRep k, fromRep a) | (k, a) <- xs] - -instance Repr ISet.IntSet where - type Rep ISet.IntSet = Rep [Int] - toRep = toRep . ISet.toList - fromRep = ISet.fromDistinctAscList . fromRep - -instance Repr a => Repr (IMap.IntMap a) where - type Rep (IMap.IntMap a) = [(Rep Int, Rep a)] - toRep m = [(toRep i, toRep a) | (i, a) <- IMap.assocs m] - fromRep xs = IMap.fromDistinctAscList [(fromRep i, fromRep a) | (i, a) <- xs] - -instance ReprT Seq.Seq where - type RepT Seq.Seq = [] - toRepTMap f = Fold.foldr (\ a xs -> f a:xs) [] - fromRepTMap f = Fold.foldl (\ xs a -> xs |> f a) Seq.empty - -genRepr [t| Seq.Seq |] diff --git a/Data/TrieMap/Rep/TH.hs b/Data/TrieMap/Rep/TH.hs deleted file mode 100644 --- a/Data/TrieMap/Rep/TH.hs +++ /dev/null @@ -1,38 +0,0 @@ -{-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, TemplateHaskell, QuasiQuotes, UndecidableInstances #-} - -module Data.TrieMap.Rep.TH where - -import Language.Haskell.TH -import Data.TrieMap.Rep - -genRepr :: Q Type -> Q [Dec] -genRepr typ = do - t <- typ - let a = VarT (mkName "a") - toRepImpl <- [| toRepTMap toRep |] - fromRepImpl <- [| fromRepTMap fromRep |] - return [InstanceD [ClassP ''Repr [a]] - (ConT ''Repr `AppT` (t `AppT` a)) - [TySynInstD ''Rep [t `AppT` a] ((ConT ''RepT `AppT` t) `AppT` (ConT ''Rep `AppT` a)), - ValD (VarP 'toRep) - (NormalB toRepImpl) [], - ValD (VarP 'fromRep) - (NormalB fromRepImpl) []]] - -genTupleRepr :: Int -> Q [Dec] -genTupleRepr n = do - let ts = [mkName [a] | a <- take n ['a'..]] - xs <- sequence [newName [a] | a <- take n ['a'..]] - xReps <- sequence [newName (a:"Rep") | a <- take n ['a'..]] - let toR = 'toRep - let fromR = 'fromRep - let tupleT = foldl AppT (TupleT n) [VarT t | t <- ts] - return [InstanceD [ClassP ''Repr [VarT t] | t <- ts] - (ConT ''Repr `AppT` tupleT) - [TySynInstD ''Rep [tupleT] (foldl AppT (TupleT n) [ConT ''Rep `AppT` VarT t | t <- ts]), - FunD toR - [Clause [TupP [VarP x | x <- xs]] - (NormalB (TupE [VarE toR `AppE` VarE x | x <- xs])) []], - FunD fromR - [Clause [TupP [VarP xRep | xRep <- xReps]] - (NormalB (TupE [VarE fromR `AppE` VarE xRep | xRep <- xReps])) []]]] diff --git a/Data/TrieMap/Representation.hs b/Data/TrieMap/Representation.hs --- a/Data/TrieMap/Representation.hs +++ b/Data/TrieMap/Representation.hs @@ -1,42 +1,5 @@ -{-# LANGUAGE TypeFamilies, TemplateHaskell, UndecidableInstances #-} -module Data.TrieMap.Representation (Repr(..)) where +module Data.TrieMap.Representation (Repr(..), genRepr, genOptRepr, genOrdRepr) where -import Data.TrieMap.Sized -import Data.TrieMap.TrieKey -import Data.TrieMap.Class -import Data.TrieMap.Rep -import Data.TrieMap.Rep.Instances () +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.Instances () import Data.TrieMap.Representation.TH - -import Data.Complex -import Data.Tree -import Data.Ratio -import Foreign.C.Types - -instance (TKey k, Repr a) => Repr (TMap k a) where - type Rep (TMap k a) = [(Rep k, Rep a)] - toRep (TMap m) = foldrWithKeyM (\ k (Elem a) xs -> (k, toRep a):xs) m [] - fromRep xs = TMap (fromDistAscListM [(k, Elem (fromRep a)) | (k, a) <- xs]) - -genOrdRepr ''Float -genOrdRepr ''Double -genRepr ''Bool -genRepr ''Tree -genRepr ''Ratio -genRepr ''Maybe -genRepr ''Complex -genRepr ''CInt -genRepr ''CChar -genRepr ''CSChar -genRepr ''CUChar -genRepr ''CShort -genRepr ''CUShort -genRepr ''CUInt -genRepr ''CLong -genRepr ''CULong -genRepr ''CLLong -genRepr ''CULLong -genRepr ''CClock -genRepr ''CTime -genRepr ''CFloat -genRepr ''CDouble diff --git a/Data/TrieMap/Representation/Class.hs b/Data/TrieMap/Representation/Class.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Class.hs @@ -0,0 +1,16 @@ +{-# LANGUAGE TypeFamilies #-} +module Data.TrieMap.Representation.Class where + +-- | The @Repr@ type class denotes that a type can be decomposed to a representation +-- built out of pieces for which the 'TrieKey' class defines a generalized trie structure. +-- +-- It is required that, if @('Repr' a, 'Eq' a)@, and @x, y :: a@, then @x '==' y@ +-- if and only if @'toRep' x '==' 'toRep' y@. It is typically the case that +-- @'compare' x y == 'compare' ('toRep' x) ('toRep' y)@, as well, but this is not +-- strictly required. (It is, however, the case for all instances built into the package.) +-- +-- As an additional note, the 'Key' modifier is used for \"bootstrapping\" 'Repr' instances, +-- allowing a type to be used in its own 'Repr' definition when wrapped in a 'Key' modifier. +class Repr a where + type Rep a + toRep :: a -> Rep a diff --git a/Data/TrieMap/Representation/Instances.hs b/Data/TrieMap/Representation/Instances.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Instances.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE TemplateHaskell, QuasiQuotes, TypeFamilies, FlexibleInstances #-} +module Data.TrieMap.Representation.Instances () where + +import Data.Tree +import Data.Ratio +import Data.Word +import Data.Bits +import Data.TrieMap.Modifiers +import qualified Data.Vector as V +import qualified Data.Vector.Storable as S +import qualified Data.Set as S +import qualified Data.Map as M +import qualified Data.Sequence as Seq + +import Data.TrieMap.Utils +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.Instances.Prim () +import Data.TrieMap.Representation.Instances.Basic () +import Data.TrieMap.Representation.Instances.ByteString () +import Data.TrieMap.Representation.Instances.Vectors () +import Data.TrieMap.Representation.Instances.Foreign () +import Data.TrieMap.Representation.TH + +instance Repr a => Repr (S.Set a) where + type Rep (S.Set a) = V.Vector (Rep a) + toRep s = toVectorN (\ f -> S.fold (f . toRep)) S.size s + +instance (Repr k, Repr a) => Repr (M.Map k a) where + type Rep (M.Map k a) = V.Vector (Rep k, Rep a) + toRep m = toVectorN (\ f -> M.foldrWithKey (\ k a -> f (toRep k, toRep a))) + M.size m + +instance Repr a => Repr (Seq.Seq a) where + type Rep (Seq.Seq a) = V.Vector (Rep a) + toRep = toVectorF toRep Seq.length + +genRepr ''Tree +genRepr ''Ratio + +instance Repr Integer where + type Rep Integer = Either (Rev (Word, S.Vector Word)) (Word, S.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) + +unroll :: Integer -> S.Vector Word +unroll x = S.reverse (S.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/Basic.hs b/Data/TrieMap/Representation/Instances/Basic.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Instances/Basic.hs @@ -0,0 +1,39 @@ +{-# LANGUAGE TemplateHaskell, TypeFamilies #-} +module Data.TrieMap.Representation.Instances.Basic () where + +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.TH + +import Control.Monad + +import qualified Data.Vector as V + +import Language.Haskell.TH + +instance Repr a => Repr [a] where + type Rep [a] = V.Vector (Rep a) + toRep = V.map toRep . V.fromList + +$(let genTupleRepr n = do + let ts = [mkName [a] | a <- take n ['a'..]] + xs <- sequence [newName [a] | a <- take n ['a'..]] + let toR = 'toRep + let tupleT = foldl AppT (TupleT n) [VarT t | t <- ts] + return [InstanceD [ClassP ''Repr [VarT t] | t <- ts] + (ConT ''Repr `AppT` tupleT) + [TySynInstD ''Rep [tupleT] (foldl AppT (TupleT n) [ConT ''Rep `AppT` VarT t | t <- ts]), + FunD toR + [Clause [TupP [VarP x | x <- xs]] + (NormalB (TupE [VarE toR `AppE` VarE x | x <- xs])) []] {-, + FunD fromR + [Clause [TupP [VarP xRep | xRep <- xReps]] + (NormalB (TupE [VarE fromR `AppE` VarE xRep | xRep <- xReps])) []] -}]] + in liftM concat $ mapM genTupleRepr [2..10]) + +genOrdRepr ''Float +genOrdRepr ''Double +genRepr ''Maybe +genRepr ''Either +genRepr ''Bool +genRepr ''() +genRepr ''Ordering diff --git a/Data/TrieMap/Representation/Instances/ByteString.hs b/Data/TrieMap/Representation/Instances/ByteString.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Instances/ByteString.hs @@ -0,0 +1,21 @@ +{-# LANGUAGE UndecidableInstances, TypeFamilies #-} +module Data.TrieMap.Representation.Instances.ByteString () where + +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.Instances.Vectors () + +import Data.Word + +import Data.ByteString.Internal (ByteString(..)) +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L + +import Data.Vector.Storable + +instance Repr ByteString where + type Rep ByteString = (Vector Word, Word) + toRep (PS fp off len) = toRep (unsafeFromForeignPtr fp off len) + +instance Repr L.ByteString where + type Rep L.ByteString = (Vector Word, Word) + toRep = toRep . B.concat . L.toChunks diff --git a/Data/TrieMap/Representation/Instances/Foreign.hs b/Data/TrieMap/Representation/Instances/Foreign.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Instances/Foreign.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE TemplateHaskell, TypeFamilies, UndecidableInstances #-} +module Data.TrieMap.Representation.Instances.Foreign () where + +import Foreign.C.Types +import Data.TrieMap.Representation.Instances.Prim () +import Data.TrieMap.Representation.Instances.Basic () +import Data.TrieMap.Representation.TH + +genRepr ''CChar +genRepr ''CSChar +genRepr ''CUChar +genRepr ''CShort +genRepr ''CUShort +genRepr ''CInt +genRepr ''CUInt +genRepr ''CLong +genRepr ''CULong +genRepr ''CPtrdiff +genRepr ''CSize +genRepr ''CWchar +genRepr ''CSigAtomic +genRepr ''CLLong +genRepr ''CULLong +genRepr ''CClock +genRepr ''CTime +genRepr ''CFloat +genRepr ''CDouble diff --git a/Data/TrieMap/Representation/Instances/Prim.hs b/Data/TrieMap/Representation/Instances/Prim.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Instances/Prim.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE ScopedTypeVariables, BangPatterns, TypeFamilies, UndecidableInstances, CPP #-} +module Data.TrieMap.Representation.Instances.Prim (i2w) where + +#include "MachDeps.h" + +import Data.TrieMap.Representation.Class +import Data.Word +import Data.Int +import Data.Char +import Data.Bits + +instance Repr Char where + type Rep Char = Word + toRep = fromIntegral . ord + +#define WREPR(wTy) \ +instance Repr wTy where { \ + type Rep wTy = Word; \ + toRep = fromIntegral} + +WREPR(Word) +WREPR(Word8) +WREPR(Word16) +WREPR(Word32) + +#if WORD_SIZE_IN_BITS < 64 +instance Repr Word64 where + type Rep Word64 = (Rep Word32, Rep Word32) + toRep w = (toRep pre, toRep suf) + where pre = fromIntegral (w `shiftR` 32) :: Word32 + suf = fromIntegral w :: Word32 +#else +WREPR(Word64) +#endif + +-- | We embed IntN into WordN, but we have to be careful about overflow. +{-# INLINE [1] i2w #-} +i2w :: forall i w . (Integral i, Bits w, Bits i, Integral w) => i -> w +i2w !i | i < 0 = mB - fromIntegral (-i) + | otherwise = mB + fromIntegral i + where mB = bit (bitSize (0 :: i) - 1) :: w + +#define IREPR(iTy,wTy) \ +instance Repr iTy where { \ + type Rep iTy = Rep wTy; \ + toRep = toRep . (i2w :: iTy -> wTy)} + +IREPR(Int8,Word8) +IREPR(Int16,Word16) +IREPR(Int32,Word32) +IREPR(Int64,Word64) +IREPR(Int,Word) diff --git a/Data/TrieMap/Representation/Instances/Vectors.hs b/Data/TrieMap/Representation/Instances/Vectors.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/Instances/Vectors.hs @@ -0,0 +1,130 @@ +{-# LANGUAGE TypeFamilies, FlexibleInstances, CPP, BangPatterns, UndecidableInstances, TemplateHaskell #-} +module Data.TrieMap.Representation.Instances.Vectors () where + +import Control.Monad.Primitive + +import Data.Word +import Data.Int +import Data.Bits + +import Foreign.Storable (Storable) +import Foreign.Ptr +import Foreign.ForeignPtr + +import Data.Vector.Generic (convert) +import qualified Data.Vector.Generic as G +import qualified Data.Vector as V +import qualified Data.Vector.Storable as S +import qualified Data.Vector.Primitive as P +import qualified Data.Vector.Unboxed as U + +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.Instances.Prim + +import Language.Haskell.TH.Syntax + +#include "MachDeps.h" + +instance Repr a => Repr (V.Vector a) where + type Rep (V.Vector a) = V.Vector (Rep a) + toRep = V.map toRep + +instance Repr (S.Vector Word) where + type Rep (S.Vector Word) = S.Vector Word + toRep = id + +type Overhang = Word +-- When storing a vector of WordNs, we view it as a vector of Words plus an overhang. +-- We store the length of the overhang (which can be up to (WORD_SIZE_IN_BITS / N - 1)) in the top +-- N bits of the Overhang, and k leftover WordNs (however large k is) in the low kN bits of the Overhang. + +-- Just a version of 'quot' for dividing by powers of 2. +quoPow :: Int -> Int -> Int +quoPow n d = $(foldr ($) [| n `quot` d |] + [\ other -> [| if d == $(lift (bit i :: Int)) then n `shiftR` $(lift i) else $other |] + | i <- [0..6]]) + +-- Just a version of 'rem' for modding by powers of 2. +remPow :: Int -> Int -> Int +remPow n d = n .&. (d - 1) + +unsafeToPtr :: Storable a => S.Vector a -> (Ptr a, Int, ForeignPtr a) +unsafeToPtr xs = unsafeInlineST $ do + S.MVector ptr n fp <- S.unsafeThaw xs + return (ptr, n, fp) + +unsafeFromPtr :: Storable a => Ptr b -> Int -> ForeignPtr b -> S.Vector a +unsafeFromPtr ptr n fp = unsafeInlineST $ S.unsafeFreeze (S.MVector (castPtr ptr) n (castForeignPtr fp)) + +#define HANGINSTANCE(wTy) \ + instance Repr (S.Vector wTy) where \ + type Rep (S.Vector wTy) = (S.Vector Word, Overhang); \ + {-# NOINLINE toRep #-}; \ + toRep !xs0 = let { \ + !b = bitSize (0 :: wTy); \ + !wordSize = bitSize (0 :: Word); \ + !ratio = quoPow wordSize b; \ + !n' = quoPow n0 ratio; \ + !nHang = remPow n0 ratio; \ + !xHang = S.drop (n0 - nHang) xs0; \ + !overhang = (fromIntegral nHang `shiftL` (wordSize - b)) .|. \ + S.foldl' (\ hang w -> (hang `shiftL` b) .|. fromIntegral w) 0 xHang; \ + !(ptr, !n0, fp) = unsafeToPtr xs0} \ + in (unsafeFromPtr ptr n' fp, overhang) + +HANGINSTANCE(Word8) +HANGINSTANCE(Word16) +#if WORD_SIZE_IN_BITS == 32 +instance Repr (S.Vector Word32) where + type Rep (S.Vector Word32) = S.Vector Word + toRep xs = case unsafeToPtr xs of + (p, n, fp) -> unsafeFromPtr p n fp +#elif WORD_SIZE_IN_BITS > 32 +HANGINSTANCE(Word32) +#endif + +instance Repr (S.Vector Word64) where + type Rep (S.Vector Word64) = S.Vector Word + toRep xs = case unsafeToPtr xs of + (p, n, fp) -> unsafeFromPtr p (n * ratio) fp + where !wordBits = bitSize (0 :: Word); ratio = quoPow 64 wordBits + +#define VEC_WORD_INST(vec,wTy) \ + instance Repr (vec wTy) where { \ + type Rep (vec wTy) = Rep (S.Vector wTy); \ + toRep = (toRep :: S.Vector wTy -> Rep (S.Vector wTy)) . convert} +#define VEC_WORD_INSTANCES(wTy) \ + VEC_WORD_INST(U.Vector,wTy); \ + VEC_WORD_INST(P.Vector,wTy) + +VEC_WORD_INSTANCES(Word8) +VEC_WORD_INSTANCES(Word16) +VEC_WORD_INSTANCES(Word32) +VEC_WORD_INSTANCES(Word64) +VEC_WORD_INSTANCES(Word) + +#define VEC_INT_INST(vec,iTy,wTy) \ + instance Repr (vec iTy) where { \ + type Rep (vec iTy) = Rep (S.Vector wTy); \ + toRep = (toRep :: S.Vector wTy -> Rep (S.Vector wTy)) . convert . G.map (i2w :: iTy -> wTy)} +#define VEC_INT_INSTANCES(iTy,wTy) \ + VEC_INT_INST(S.Vector,iTy,wTy); \ + VEC_INT_INST(P.Vector,iTy,wTy); \ + VEC_INT_INST(U.Vector,iTy,wTy) + +VEC_INT_INSTANCES(Int8, Word8) +VEC_INT_INSTANCES(Int16, Word16) +VEC_INT_INSTANCES(Int32, Word32) +VEC_INT_INSTANCES(Int64, Word64) +VEC_INT_INSTANCES(Int, Word) + +#define VEC_ENUM_INST(ty, vec) \ + instance Repr (vec ty) where { \ + type Rep (vec ty) = S.Vector Word; \ + toRep = convert . G.map (fromIntegral . fromEnum)} +#define VEC_ENUM_INSTANCES(ty) \ + VEC_ENUM_INST(ty,S.Vector); \ + VEC_ENUM_INST(ty,P.Vector); \ + VEC_ENUM_INST(ty,U.Vector) + +VEC_ENUM_INSTANCES(Char) diff --git a/Data/TrieMap/Representation/TH.hs b/Data/TrieMap/Representation/TH.hs --- a/Data/TrieMap/Representation/TH.hs +++ b/Data/TrieMap/Representation/TH.hs @@ -1,160 +1,133 @@ -{-# LANGUAGE TemplateHaskell, QuasiQuotes, PatternGuards, DoAndIfThenElse #-} +{-# LANGUAGE BangPatterns, TypeFamilies, TemplateHaskell, PatternGuards, DoAndIfThenElse, ImplicitParams #-} -module Data.TrieMap.Representation.TH (genRepr, genOrdRepr) where +module Data.TrieMap.Representation.TH (genRepr, genOptRepr, genOrdRepr) where -import Data.TrieMap.Modifiers -import Data.TrieMap.Rep -import Data.TrieMap.Rep.Instances () -import Language.Haskell.TH +import Language.Haskell.TH.Syntax import Language.Haskell.TH.ExpandSyns -data ToRepCase = ToRepCase [Pat] Exp -data FromRepCase = FromRepCase Pat [Exp] -type ToRep = [ToRepCase] -type FromRep = [FromRepCase] +import qualified Data.Vector as V -type Representation = (Type, ToRep, FromRep) +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.TH.Utils +import Data.TrieMap.Representation.TH.Representation +import Data.TrieMap.Representation.TH.Factorized +import Data.TrieMap.Representation.TH.ReprMonad -- | Given a type with an associated 'Ord' instance, generates a representation that will cause its 'TMap' -- implementation to be essentially equivalent to "Data.Map". genOrdRepr :: Name -> Q [Dec] -genOrdRepr tycon = do - TyConI dec <- reify tycon - let theTyp = foldl AppT (ConT tycon) . map tyVarBndrType +genOrdRepr tycon = execReprMonad $ do + (cxt, ty, _) <- getDataForName tycon + outputRepr cxt ty =<< ordRepr ty + +getDataForName :: Quasi m => Name -> m (Cxt, Type, [AlgCon]) +getDataForName tycon = do + TyConI dec <- qReify tycon + let theTyp = compose tycon . map tyVarBndrVar case dec of - DataD cxt _ tyvars _ _ -> do - repr <- ordRepr (theTyp tyvars) - return (decsForRepr cxt (theTyp tyvars) repr) - NewtypeD cxt _ tyvars _ _ -> do - repr <- ordRepr (theTyp tyvars) - return (decsForRepr cxt (theTyp tyvars) repr) - _ -> fail ("Cannot generate Repr instance for " ++ pprint dec) + DataD cxt _ tyvars cons _ -> + return (cxt, theTyp tyvars, map algCon cons) + NewtypeD cxt _ tyvars con _ -> + return (cxt, theTyp tyvars, [algCon con]) + _ -> error "Error: could not get kind of type constructor" -ordRepr :: Type -> Q Representation -ordRepr t0 = do - x <- newName "x" - return (ConT ''Ordered `AppT` t0, - [ToRepCase [VarP x] (ConE 'Ord `AppE` VarE x)], - [FromRepCase (ConP 'Ord [VarP x]) - [VarE x]]) - +getDataForType :: Quasi m => Type -> m (Cxt, [AlgCon]) +getDataForType ty + | (ConT tyCon, args) <- decompose ty + = do TyConI dec <- qReify tyCon + let subAll tyvars cxt cons = let subs = zip (map tyVarBndrVar tyvars) args in + ([foldr substInPred p subs | p <- cxt], [foldr substInAlgCon (algCon con) subs | con <- cons]) + case dec of + DataD cxt _ tyvars cons _ -> + return (subAll tyvars cxt cons) + NewtypeD cxt _ tyvars con _ -> + return (subAll tyvars cxt [con]) + _ -> failure + | otherwise = failure + where failure = fail "Error: could not reify type constructor" -- | Given the name of a type constructor, automatically generates an efficient 'Repr' instance. +-- If you have several mutually dependent (or even mutually recursive) types, 'genRepr' will +-- construct instances for all of them. +-- +-- 'genRepr' guarantees that any instances it generates are consistent with the ordering that +-- would be generated by @deriving ('Ord')@ in the data declaration. That is, if 'genRepr' +-- generates an instance @Repr a@, then it is guaranteed that if @x, y :: a@, and @a@ +-- has a derived 'Ord' instance, then @compare x y == compare (toRep x) (toRep y)@. genRepr :: Name -> Q [Dec] -genRepr tycon = do - TyConI dec <- reify tycon - let theTyp = foldl AppT (ConT tycon) . map tyVarBndrType - case dec of - DataD cxt _ tyvars cons _ -> do - conReprs <- mapM conRepr cons - return (decsForRepr cxt (theTyp tyvars) (foldr1 union conReprs)) - NewtypeD cxt _ tyvars con _ -> do - theConRepr <- conRepr con - return (decsForRepr cxt (theTyp tyvars) theConRepr) - _ -> fail ("Cannot generate Repr instance for " ++ pprint dec) +genRepr tyCon = execReprMonad $ do + (_, ty, _) <- getDataForName tyCon + let ?combine = mergeWith sumRepr + genReprMain ty -tyVarBndrType :: TyVarBndr -> Type -tyVarBndrType (PlainTV tyvar) = VarT tyvar -tyVarBndrType (KindedTV tyvar _) = VarT tyvar +-- | Given the name of a type constructor, automatically generates an efficient 'Repr' instance. +-- If you have several mutually dependent (or even mutually recursive) types, 'genOptRepr' will +-- construct instances for all of them. The instance generated by 'genOptRepr' may, in some +-- cases, be more efficient than the instance generated by 'genRepr' -- in particular, +-- arguments common to several constructors may be factored out, reducing the complexity of the +-- associated 'TrieKey' instance, but leaving an ordering inconsistent with 'Ord'. +-- +-- Therefore, 'genOptRepr' guarantees that any instances it generates are consistent with the +-- ordering that would be generated by @deriving ('Eq')@ in the data declaration. That is, if +-- 'genOptRepr' generates an instance @Repr a@, then it is guaranteed that if @x, y :: a@, and +-- @a@ has a derived 'Eq' instance, then @(x == y) == (toRep x == toRep y)@. +genOptRepr :: Name -> Q [Dec] +genOptRepr tyCon = execReprMonad $ do + (_, ty, _) <- getDataForName tyCon + let ?combine = unify + genReprMain ty -decsForRepr :: Cxt -> Type -> Representation -> [Dec] -decsForRepr cxt t (tRep, toR, fromR) = [ - InstanceD cxt (ConT ''Repr `AppT` t) - [TySynInstD ''Rep [t] tRep, - FunD 'toRep - [Clause pats (NormalB e) [] | ToRepCase pats e <- toR], - FunD 'fromRep - [Clause [pat] (NormalB e) [] | FromRepCase pat [e] <- fromR]]] +mustBreakTy :: Type -> ReprMonad Bool +mustBreakTy ty = case decompose ty of + (ConT tyCon, _) -> mustBreak tyCon + _ -> return False -decompose :: Type -> (Type, [Type]) -decompose (tyfun `AppT` ty) = case decompose tyfun of - (tyfun, tys) -> (tyfun, tys ++ [ty]) -decompose ty = (ty, []) +recurseTy :: Type -> ReprMonad a -> ReprMonad a +recurseTy ty m = case decompose ty of + (ConT tyCon, _) -> recurse tyCon m + _ -> m -type ReprM = Q +genReprMain :: (?combine :: [Representation] -> Representation) => Type -> ReprMonad Type +genReprMain ty = do + breakTy <- mustBreakTy ty + if breakTy then fail "Cannot recurse here" + else do + knownInst <- getInstance ty + case knownInst of + Just known -> return known + Nothing -> do + (cxt, cons) <- getDataForType ty + conReprs <- mapM (recurseTy ty . conRepr) cons + outputRepr cxt ty (checkEnumRepr $ ?combine conReprs) -conRepr :: Con -> ReprM Representation -conRepr (RecC con args) = conRepr (NormalC con [(strict, typ) | (_, strict, typ) <- args]) -conRepr (InfixC t1 con t2) = conRepr (NormalC con [t1, t2]) -conRepr (NormalC con []) = return $ conify con unit -conRepr (NormalC con args) = do - argCons <- mapM (typeRepr . snd) args - return (conify con (foldr1 prod argCons)) -conRepr con = fail ("Cannot generate representation for existential constructor " ++ pprint con) +conRepr :: (?combine :: [Representation] -> Representation) => AlgCon -> ReprMonad Representation +conRepr (con, []) = return $ conify con unitRepr +conRepr (con, args) = do + argReprs <- mapM typeRepr args + return (conify con (foldr1 prodRepr argReprs)) -typeRepr :: Type -> ReprM Representation -typeRepr t00 = expandSyns t00 >>= \ t0 -> case decompose t0 of +typeRepr :: (?combine :: [Representation] -> Representation) => Type -> ReprMonad Representation +typeRepr t00 = liftQuasi (expandSyns t00) >>= \ t0 -> case decompose t0 of (ListT, [t]) -> do - (tRep, toR, fromR) <- typeRepr t - xs <- newName "elems" - x <- newName "el" - xsRep <- newName "elemReps" - xRep <- newName "elemRep" - return (ListT `AppT` tRep, - [ToRepCase [VarP xs] - (CompE [BindS (VarP x) (VarE xs), - NoBindS (CaseE (VarE x) [Match pat (NormalB e) [] | ToRepCase [pat] e <- toR])])], - [FromRepCase (VarP xsRep) - [CompE [BindS (VarP xRep) (VarE xsRep), - NoBindS (CaseE (VarE xRep) [Match pat (NormalB e) [] | FromRepCase pat [e] <- fromR])]]]) - (TupleT 0, _) -> return unit + tRepr <- typeRepr t + vectorizeRepr (VarE 'V.fromList) tRepr + (TupleT 0, _) -> return unitRepr (TupleT _, ts) -> do reps <- mapM typeRepr ts - let (tRep, toR, fromR) = foldr1 prod reps - return (tRep, [ToRepCase [TupP pats] e | ToRepCase pats e <- toR], [FromRepCase pat [TupE es] | FromRepCase pat es <- fromR]) + return $ mapReprInput TupP $ mergeWith prodRepr reps (ConT con, ts) - | con == ''() -> return unit + | con == ''() -> return unitRepr | con == ''Either, [tL, tR] <- ts - -> do (tRepL, lToR, lFromR) <- typeRepr tL - (tRepR, rToR, rFromR) <- typeRepr tR - return (ConT ''Either `AppT` tRepL `AppT` tRepR, - [ToRepCase [ConP 'Left pats] (ConE 'Left `AppE` e) | ToRepCase pats e <- lToR] ++ - [ToRepCase [ConP 'Right pats] (ConE 'Right `AppE` e) | ToRepCase pats e <- rToR], - [FromRepCase (ConP 'Left [pat]) [ConE 'Left `AppE` e] | FromRepCase pat [e] <- lFromR] ++ - [FromRepCase (ConP 'Right [pat]) [ConE 'Right `AppE` e] | FromRepCase pat [e] <- rFromR]) - | otherwise -> do ClassI _ instances <- reify ''Repr - let knowns = [tycon | ClassInstance{ci_tys = [ConT tycon]} <- instances] - -- TODO: recognize preexisting higher-arity instances - if con `elem` knowns && null ts then do - arg <- newName "arg" - argRep <- newName "argRep" - return (ConT ''Rep `AppT` ConT con, - [ToRepCase [VarP arg] (VarE 'toRep `AppE` VarE arg)], - [FromRepCase (VarP argRep) [VarE 'fromRep `AppE` VarE argRep]]) - else recursiveRepr t0 - _ -> recursiveRepr t0 - -recursiveRepr :: Type -> ReprM Representation -recursiveRepr t0 = do -- TODO: handle type synonyms here - x <- newName "arg" - return (ConT ''Key `AppT` t0, - [ToRepCase [VarP x] (ConE 'Key `AppE` VarE x)], - [FromRepCase (ConP 'Key [VarP x]) [VarE x]]) - -unit :: Representation -unit = (TupleT 0, [ToRepCase [] (TupE [])], [FromRepCase WildP []]) - -prod :: Representation -> Representation -> Representation -prod (t1, toRep1, fromRep1) - (t2, toRep2, fromRep2) = - (TupleT 2 `AppT` t1 `AppT` t2, - do ToRepCase pats1 out1 <- toRep1 - ToRepCase pats2 out2 <- toRep2 - return (ToRepCase (pats1 ++ pats2) (TupE [out1, out2])), - do FromRepCase pat1 out1 <- fromRep1 - FromRepCase pat2 out2 <- fromRep2 - return (FromRepCase (TupP [pat1, pat2]) (out1 ++ out2))) - -conify :: Name -> Representation -> Representation -conify conName (t, toR, fromR) = - (t, [ToRepCase [ConP conName args] e | ToRepCase args e <- toR], - [FromRepCase p [foldl AppE (ConE conName) outs] | FromRepCase p outs <- fromR]) + -> do reprL <- typeRepr tL + reprR <- typeRepr tR + return (mapReprInput (ConP leftN) reprL `sumRepr` mapReprInput (ConP rightN) reprR) + | con == ''Maybe, [t] <- ts + -> do tRepr <- typeRepr t + return (conify 'Nothing unitRepr `sumRepr` conify 'Just tRepr) + _ -> bootstrapRepr t0 -union :: Representation -> Representation -> Representation -union (t1, toRep1, fromRep1) - (t2, toRep2, fromRep2) = - (ConT ''Either `AppT` t1 `AppT` t2, - [ToRepCase pats (ConE 'Left `AppE` e) | ToRepCase pats e <- toRep1] ++ - [ToRepCase pats (ConE 'Right `AppE` e) | ToRepCase pats e <- toRep2], - [FromRepCase (ConP 'Left [pat]) es | FromRepCase pat es <- fromRep1] ++ - [FromRepCase (ConP 'Right [pat]) es | FromRepCase pat es <- fromRep2]) +bootstrapRepr :: (?combine :: [Representation] -> Representation) => Type -> ReprMonad Representation +bootstrapRepr t0 = qRecover fallback + (do _tRep <- genReprMain t0 + recursiveRepr (ConT ''Rep `AppT` t0) (VarE 'toRep)) + where fallback = keyRepr t0 diff --git a/Data/TrieMap/Representation/TH/Factorized.hs b/Data/TrieMap/Representation/TH/Factorized.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/TH/Factorized.hs @@ -0,0 +1,76 @@ +{-# LANGUAGE ParallelListComp, NamedFieldPuns, RecordWildCards #-} +module Data.TrieMap.Representation.TH.Factorized (unify) where + +import Control.Exception + +import Data.List +import Data.Maybe +import Data.Ord + +import Language.Haskell.TH +import Data.TrieMap.Representation.TH.Representation +import Data.TrieMap.Representation.TH.Utils + +data FactorCase = FCase {fInput :: [Pat], fFactor :: Exp, fOutput :: Exp} +data Factored = Factored {factorType :: Type, fRestType :: Type, fCases :: [FactorCase]} + +factorRepr, otherRepr :: Factored -> Representation +factorRepr Factored{..} = + Repr {reprType = factorType, cases = map factorCase fCases} +otherRepr Factored{..} = + Repr {reprType = fRestType, cases = map otherCase fCases} + +factorCase, otherCase :: FactorCase -> Case +factorCase FCase{..} = Case{input = fInput, output = fFactor} +otherCase FCase{..} = Case{input = fInput, output = fOutput} + +caseFactor :: Case -> FactorCase +caseFactor Case{..} = FCase{fInput = input, fFactor = output, fOutput = TupE []} + +combFCase :: Case -> FactorCase -> FactorCase +combFCase Case{..} FCase{..} = + assert (input == fInput) $ FCase{fOutput = TupE [output, fOutput], ..} + +combFactor :: Representation -> Factored -> Factored +combFactor Repr{..} Factored{fRestType = TupleT 0,..} = + Factored{factorType, fRestType = reprType, fCases = [FCase{fOutput = output,..} | (FCase{..}, Case{output}) <- zip fCases cases]} +combFactor Repr{..} Factored{..} = + Factored{factorType, fRestType = reprType `tyProd` fRestType, fCases = zipWith combFCase cases fCases} + +factors :: Representation -> [Factored] +factors repr@Repr{..} = case reprType of + TupleT 2 `AppT` _ `AppT` _ + -> let fs1 = map (combFactor (sndRepr repr)) (factors (fstRepr repr)) + fs2 = map (combFactor (fstRepr repr)) (factors (sndRepr repr)) + in baseFactor:fs1 ++ fs2 + _ -> [baseFactor] + where baseFactor = Factored { + factorType = reprType, + fRestType = TupleT 0, + fCases = map caseFactor cases} + +distinctFactors :: [Representation] -> [Type] +distinctFactors reprs = nub [factorType | repr <- reprs, Factored{factorType} <- factors repr, factorType /= TupleT 0] + +factorWith :: Type -> Representation -> Maybe Factored +factorWith fTy repr = listToMaybe [factor | factor@Factored{factorType} <- factors repr, factorType == fTy] + +factorOut :: Type -> [Representation] -> ([Factored], [Representation]) +factorOut _ [] = ([], []) +factorOut fTy (repr:reprs) = case (factorWith fTy repr, factorOut fTy reprs) of + (Nothing, (factors, others)) -> (factors, repr:others) + (Just f, (factors, others)) -> (f:factors, others) + +unify :: [Representation] -> Representation +unify reprs = case (allFactors, bestOption) of + ([], _) -> checkEnumRepr (mergeWith sumRepr reprs) + (_, ([_], _)) -> checkEnumRepr (mergeWith sumRepr reprs) + (_, (factors, [])) -> distributeMany factors + (_, (factors, others)) -> distributeMany factors `sumRepr` unify others + where allFactors = distinctFactors reprs + options = map (`factorOut` reprs) (distinctFactors reprs) + bestOption = maximumBy (comparing (length . fst)) options + +distributeMany :: [Factored] -> Representation +distributeMany factors = + foldr1 unifySumRepr (map factorRepr factors) `unifyProdRepr` unify (map otherRepr factors) diff --git a/Data/TrieMap/Representation/TH/ReprMonad.hs b/Data/TrieMap/Representation/TH/ReprMonad.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/TH/ReprMonad.hs @@ -0,0 +1,82 @@ +{-# LANGUAGE ViewPatterns, TemplateHaskell #-} +module Data.TrieMap.Representation.TH.ReprMonad ( + ReprMonad, + liftQuasi, + recurse, + getInstance, + outputInstance, + mustBreak, + execReprMonad) where + +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.TH.Utils + +import Control.Monad +import Language.Haskell.TH.Syntax +import Language.Haskell.TH.ExpandSyns + +type Instances = [(Name, ([Name], Type))] + +newtype ReprMonad a = ReprMonad {runReprMonad :: + Instances -- tycons of known instances + -> [Name] -- tycons of instances in progress (breakpoints of recursive loopies) + -> Q ([Dec], Instances, a) -- output decs, new known instances + } + +instance Monad ReprMonad where + return x = ReprMonad $ \ knowns _ -> return ([], knowns, x) + m >>= k = ReprMonad $ \ knowns breaks -> do + (outDecs, knowns', a) <- runReprMonad m knowns breaks + (outDecs', knowns'', b) <- runReprMonad (k a) knowns' breaks + return (outDecs ++ outDecs', knowns'', b) + fail err = ReprMonad $ \ _ _ -> fail err + +instance Functor ReprMonad where + fmap = liftM + +liftQuasi :: Q a -> ReprMonad a +liftQuasi q = ReprMonad $ \ knowns _ -> do + a <- q + return ([], knowns, a) + +instance Quasi ReprMonad where + qNewName = liftQuasi . qNewName + qReport b str = liftQuasi (qReport b str) + qRecover m k = ReprMonad $ \ knowns breaks -> qRecover (runReprMonad m knowns breaks) (runReprMonad k knowns breaks) + qReify = liftQuasi . qReify + qClassInstances name typs = liftQuasi (qClassInstances name typs) + qLocation = liftQuasi qLocation + qRunIO = liftQuasi . qRunIO + +insNub :: Eq a => a -> [a] -> [a] +insNub x ys0@(y:ys) + | x == y = ys0 + | otherwise = y:insNub x ys +insNub x [] = [x] + +recurse :: Name -> ReprMonad a -> ReprMonad a +recurse breakTy m = ReprMonad $ \ knowns breaks -> runReprMonad m knowns (breakTy `insNub` breaks) + +outputInstance :: Type -> Type -> [Dec] -> ReprMonad () +outputInstance ty tyRep decs = ReprMonad $ \ knowns _ -> case decompose' ty of + Just (tyCon, tyArgs) + -> return (decs, (tyCon, (tyArgs, tyRep)):knowns, ()) + _ -> return (decs, knowns, ()) + +getInstance :: Type -> ReprMonad (Maybe Type) +getInstance typ = case decompose typ of + (ConT tyCon, tyArgs) -> ReprMonad $ \ knowns _ -> case lookup tyCon knowns of + Nothing -> return ([], knowns, Nothing) + Just (tyArgs', tyRep) -> return ([], knowns, Just $ foldr substInType tyRep (zip tyArgs' tyArgs)) + _ -> return Nothing + +mustBreak :: Name -> ReprMonad Bool +mustBreak tyCon = ReprMonad $ \ knowns breaks -> return ([], knowns, tyCon `elem` breaks) + +execReprMonad :: ReprMonad a -> Q [Dec] +execReprMonad m = do + ClassI _ instances <- reify ''Repr + let instanceHeads = [(tyConName, (tyArgs, ConT ''Rep `AppT` compose tyConName tyArgs)) + | ClassInstance{ci_tys = [decompose' -> Just (tyConName, tyArgs)]} <- instances] + (decs, _, _) <- runReprMonad m instanceHeads [] + return decs diff --git a/Data/TrieMap/Representation/TH/Representation.hs b/Data/TrieMap/Representation/TH/Representation.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/TH/Representation.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE TemplateHaskell, RecordWildCards, NamedFieldPuns, PatternGuards #-} +module Data.TrieMap.Representation.TH.Representation ( + Representation(..), + Case(..), + fstRepr, + sndRepr, + prodRepr, + sumRepr, + unifyProdRepr, + unifySumRepr, + checkEnumRepr, + unitRepr, + vectorizeRepr, + mapReprInput, + conify, + ordRepr, + outputRepr, + recursiveRepr, + keyRepr) where + +import Control.Exception (assert) +import Control.Monad + +import Data.Word +import Data.Maybe +import qualified Data.Vector as V + +import Language.Haskell.TH.Syntax + +import Data.TrieMap.Modifiers +import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.TH.Utils +import Data.TrieMap.Representation.TH.ReprMonad + +data Representation = Repr {reprType :: Type, cases :: [Case]} deriving (Show) +data Case = Case {input :: [Pat], output :: Exp} deriving (Show) + +unitRepr :: Representation +unitRepr = Repr {reprType = TupleT 0, cases = [Case [] (TupE [])]} + +vectorizeRepr :: Quasi m => Exp -> Representation -> m Representation +vectorizeRepr toVecE Repr{..} = do + xs <- qNewName "xs" + eToR <- qNewName "eToR" + let mapE f xs = VarE 'V.map `AppE` f `AppE` xs + let eToRDec = FunD eToR (map caseToClause cases) + return $ Repr { + reprType = ConT ''V.Vector `AppT` reprType, + cases = [Case {input = [VarP xs], + output = mapE (LetE [eToRDec] (VarE eToR)) (toVecE `AppE` VarE xs)}]} + +fstRepr, sndRepr :: Representation -> Representation +fstRepr = mapReprOutput fstTy fstExp +sndRepr = mapReprOutput sndTy sndExp + +prodCase :: Case -> Case -> Case +prodCase Case{input = input1, output = output1} Case{input = input2, output = output2} + = Case {input = input1 ++ input2, output = TupE [output1, output2]} + +unifyProdCase :: Case -> Case -> Maybe Case +unifyProdCase Case{input = input1, output = output1} Case{input = input2, output = output2} + = do guard (input1 == input2) + return Case{input = input1, output = TupE [output1, output2]} + +mapCaseInput :: ([Pat] -> Pat) -> Case -> Case +mapCaseInput f Case{..} = Case{input = [f input],..} + +mapCaseOutput :: (Exp -> Exp) -> Case -> Case +mapCaseOutput f Case{..} = Case{output = f output,..} + +prodRepr, sumRepr, unifySumRepr, unifyProdRepr :: Representation -> Representation -> Representation +prodRepr Repr{reprType = repr1, cases = cases1} Repr{reprType = repr2, cases = cases2} + = Repr {reprType = repr1 `tyProd` repr2, cases = liftM2 prodCase cases1 cases2} + +sumRepr Repr{reprType = repr1, cases = cases1} Repr{reprType = repr2, cases = cases2} + = Repr {reprType = repr1 `tySum` repr2, + cases = map (mapCaseOutput leftExp) cases1 ++ map (mapCaseOutput rightExp) cases2} + +unifySumRepr Repr{reprType = repr1, cases = cases1} Repr{reprType = repr2, cases = cases2} + = assert (repr1 == repr2) $ Repr {reprType = repr1, cases = cases1 ++ cases2} + +unifyProdRepr Repr{reprType = repr1, cases = cases1} Repr{reprType = repr2, cases = cases2} + = Repr {reprType = repr1 `tyProd` repr2, cases = catMaybes (liftM2 unifyProdCase cases1 cases2)} + +mapReprInput :: ([Pat] -> Pat) -> Representation -> Representation +mapReprInput f Repr{..} = Repr{cases = map (mapCaseInput f) cases, ..} + +conify :: Name -> Representation -> Representation +conify con = mapReprInput (ConP con) + +mapReprOutput :: (Type -> Type) -> (Exp -> Exp) -> Representation -> Representation +mapReprOutput tyOp outOp Repr{..} = Repr{reprType = tyOp reprType, cases = map (mapCaseOutput outOp) cases} + +checkEnumRepr :: Representation -> Representation +checkEnumRepr Repr{..} + | isEnumTy reprType, length cases > 2 + = Repr {reprType = ConT ''Word, cases = [Case{input, output = LitE (IntegerL i)} | (i, Case{..}) <- zip [0..] cases]} +checkEnumRepr repr = repr + +ordRepr :: Quasi m => Type -> m Representation +ordRepr ty = do + x <- qNewName "ordK" + return Repr{reprType = ConT ''Ordered `AppT` ty, + cases = [Case {input = [VarP x], output = ConE 'Ord `AppE` VarE x}]} + +caseToClause :: Case -> Clause +caseToClause Case{..} = Clause input (NormalB output) [] + +outputRepr :: Cxt -> Type -> Representation -> ReprMonad Type +outputRepr cxt ty Repr{..} = do + outputInstance ty reprType + [InstanceD cxt (ConT ''Repr `AppT` ty) + [TySynInstD ''Rep [ty] reprType, + FunD 'toRep + (map caseToClause cases)]] + return reprType + +recursiveRepr :: Quasi m => Type -> Exp -> m Representation +recursiveRepr reprType toRepE = do + deep <- qNewName "deep" + return Repr{reprType, cases = [Case{input = [VarP deep], output = toRepE `AppE` VarE deep}]} + +keyRepr :: Quasi m => Type -> m Representation +keyRepr ty = do + shallow <- qNewName "shallow" + let keyCon = ConE 'Key + return Repr{reprType = ConT ''Key `AppT` ty, cases = [Case{input = [VarP shallow], output = keyCon `AppE` VarE shallow}]} diff --git a/Data/TrieMap/Representation/TH/Utils.hs b/Data/TrieMap/Representation/TH/Utils.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Representation/TH/Utils.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE TemplateHaskell #-} +module Data.TrieMap.Representation.TH.Utils where + +import Language.Haskell.TH +import Language.Haskell.TH.ExpandSyns + +decompose :: Type -> (Type, [Type]) +decompose (tyfun `AppT` ty) = case decompose tyfun of + (tyfun, tys) -> (tyfun, tys ++ [ty]) +decompose ty = (ty, []) + +decompose' :: Type -> Maybe (Name, [Name]) +decompose' (tyfun `AppT` VarT ty) = do + (tyfun, tys) <- decompose' tyfun + return (tyfun, tys ++ [ty]) +decompose' (ConT ty) = return (ty, []) +decompose' _ = Nothing + +compose :: Name -> [Name] -> Type +compose tyCon tyArgs = foldl AppT (ConT tyCon) (map VarT tyArgs) + +tyVarBndrVar :: TyVarBndr -> Name +tyVarBndrVar (PlainTV tyvar) = tyvar +tyVarBndrVar (KindedTV tyvar _) = tyvar + +tyVarBndrType :: TyVarBndr -> Type +tyVarBndrType = VarT . tyVarBndrVar + +tyProd, tySum :: Type -> Type -> Type +tyProd t1 t2 = TupleT 2 `AppT` t1 `AppT` t2 +tySum t1 t2 = ConT ''Either `AppT` t1 `AppT` t2 + +fstExp, sndExp :: Exp -> Exp +fstExp (TupE [e, _]) = e +fstExp e = VarE 'fst `AppE` e +sndExp (TupE [_, e]) = e +sndExp e = VarE 'snd `AppE` e + +leftN, rightN :: Name +leftN = 'Left +rightN = 'Right + +leftExp, rightExp :: Exp -> Exp +leftExp = AppE (ConE leftN) +rightExp = AppE (ConE rightN) + +fstTy, sndTy :: Type -> Type +fstTy (TupleT 2 `AppT` t1 `AppT` _) = t1 +fstTy _ = error "Error: not a pair type" +sndTy (TupleT 2 `AppT` _ `AppT` t2) = t2 +sndTy _ = error "Error: not a pair type" + +isEnumTy :: Type -> Bool +isEnumTy (ConT eith `AppT` t1 `AppT` t2) + = eith == ''Either && isEnumTy t1 && isEnumTy t2 +isEnumTy (TupleT 0) + = True +isEnumTy _ = False + +type AlgCon = (Name, [Type]) + +algCon :: Con -> AlgCon +algCon (NormalC name args) = (name, map snd args) +algCon (RecC name args) = (name, [argTy | (_, _, argTy) <- args]) +algCon (InfixC (_, ty1) name (_, ty2)) = (name, [ty1, ty2]) +algCon _ = error "Error: universally quantified constructors are not algebraic" + +substInAlgCon :: (Name, Type) -> AlgCon -> AlgCon +substInAlgCon sub (conName, args) = (conName, map (substInType sub) args) + +substInPred :: (Name, Type) -> Pred -> Pred +substInPred sub (ClassP cName tys) = ClassP cName (map (substInType sub) tys) +substInPred sub (EqualP ty1 ty2) = EqualP (substInType sub ty1) (substInType sub ty2) + +mergeWith :: (a -> a -> a) -> [a] -> a +mergeWith _ [a] = a +mergeWith _ [] = error "Error: mergeWith called with empty list" +mergeWith f xs = mergeWith f (combine xs) where + combine (x1:x2:xs) = f x1 x2:combine xs + combine xs = xs diff --git a/Data/TrieMap/ReverseMap.hs b/Data/TrieMap/ReverseMap.hs --- a/Data/TrieMap/ReverseMap.hs +++ b/Data/TrieMap/ReverseMap.hs @@ -1,58 +1,57 @@ -{-# LANGUAGE UnboxedTuples, TypeFamilies, BangPatterns, MagicHash #-} +{-# LANGUAGE TypeFamilies, MagicHash, UnboxedTuples #-} +module Data.TrieMap.ReverseMap () where -module Data.TrieMap.ReverseMap (reverse, unreverse) where +import Control.Applicative +import Data.TrieMap.Applicative import Data.TrieMap.TrieKey -import Data.TrieMap.Sized import Data.TrieMap.Modifiers -import Data.TrieMap.Applicative - -import Control.Applicative - -import Prelude hiding (reverse) -import qualified Data.List as L +import Data.TrieMap.Sized import GHC.Exts instance TrieKey k => TrieKey (Rev k) where - newtype TrieMap (Rev k) a = RMap (TrieMap k a) + newtype TrieMap (Rev k) a = RevMap (TrieMap k a) newtype Hole (Rev k) a = RHole (Hole k a) - emptyM = RMap emptyM - singletonM (Rev k) a = RMap (singletonM k a) - nullM (RMap m) = nullM m - sizeM (RMap m) = sizeM m - lookupM (Rev k) (RMap m) = lookupM k m - mapWithKeyM f (RMap m) = RMap (mapWithKeyM (f . Rev) m) - traverseWithKeyM f (RMap m) = RMap <$> runDual (traverseWithKeyM g m) - where g k a = Dual (f (Rev k) a) - mapMaybeM f (RMap m) = RMap (mapMaybeM (f . Rev) m) - mapEitherM f (RMap m) = both RMap RMap (mapEitherM (f . Rev)) m - foldrWithKeyM f (RMap m) = foldlWithKeyM (flip . f . Rev) m - foldlWithKeyM f (RMap m) = foldrWithKeyM (flip . f . Rev) m - unionM f (RMap m1) (RMap m2) = RMap (unionM (f . Rev) m1 m2) - isectM f (RMap m1) (RMap m2) = RMap (isectM (f . Rev) m1 m2) - diffM f (RMap m1) (RMap m2) = RMap (diffM (f . Rev) m1 m2) - isSubmapM (<=) (RMap m1) (RMap m2) = isSubmapM (<=) m1 m2 - fromListM f xs = RMap (fromListM (f . Rev) [(k, a) | (Rev k, a) <- xs]) - fromAscListM f xs = RMap (fromAscListM (\ k a1 a2 -> f (Rev k) a2 a1) [(k, a) | (Rev k, a) <- L.reverse xs]) - fromDistAscListM xs = RMap (fromDistAscListM [(k, a) | (Rev k, a) <- L.reverse xs]) + Rev k1 =? Rev k2 = k1 =? k2 + Rev k1 `cmp` Rev k2 = k2 `cmp` k1 + + emptyM = RevMap emptyM + singletonM (Rev k) a = RevMap (singletonM k a) + lookupM (Rev k) (RevMap m) = lookupM 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) + + foldlM f (RevMap m) = foldrM (flip f) m + foldrM f (RevMap m) = foldlM (flip 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) - keyM (RHole hole) = Rev (keyM hole) - beforeM a (RHole hole) = RMap (afterM a hole) - afterM a (RHole hole) = RMap (beforeM a hole) - searchM (Rev k) (RMap m) = onUnboxed RHole (searchM k) m - indexM i# (RMap m) = case indexM (sm# -# 1# -# i#) m of - (# i'#, v, hole #) -> (# getSize# v -# 1# -# i'#, v, RHole hole #) - where !sm# = sizeM m - extractHoleM (RMap m) = do - (v, hole) <- runDualPlus (extractHoleM m) - return (v, RHole hole) - assignM x (RHole hole) = RMap (assignM x hole) - clearM (RHole hole) = RMap (clearM hole) - -reverse :: TrieMap k a -> TrieMap (Rev k) a -reverse = RMap + beforeM a (RHole hole) = RevMap (afterM a hole) + afterM a (RHole hole) = RevMap (beforeM a hole) + searchM (Rev k) (RevMap m) = onSnd RHole (searchM k) m + indexM i# (RevMap m) = case indexM (revIndex i# m) m of + (# i'#, a, hole #) -> (# revIndex i'# a, a, RHole hole #) + extractHoleM (RevMap m) = runDualPlus $ do + (a, hole) <- extractHoleM m + return (a, RHole hole) + assignM v (RHole m) = RevMap (assignM v 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]) + + unifyM (Rev k1) a1 (Rev k2) a2 = either (Left . RHole) (Right . RevMap) (unifyM k1 a1 k2 a2) -unreverse :: TrieMap (Rev k) a -> TrieMap k a -unreverse (RMap m) = m +revIndex :: Sized a => Int# -> a -> Int# +revIndex i# a = getSize# a -# 1# -# i# diff --git a/Data/TrieMap/Sized.hs b/Data/TrieMap/Sized.hs --- a/Data/TrieMap/Sized.hs +++ b/Data/TrieMap/Sized.hs @@ -7,10 +7,19 @@ class Sized a where getSize# :: a -> Int# -newtype Elem a = Elem {getElem :: a} +data Assoc k a = Assoc {getK :: k, getValue :: a} +newtype Elem a = Elem a + instance Sized (Elem a) where getSize# _ = 1# + +instance Sized (Assoc k a) where + getSize# _ = 1# + +instance Sized a => Sized (Maybe a) where + getSize# (Just a) = getSize# a + getSize# _ = 0# getSize :: Sized a => a -> Int getSize a = I# (getSize# a) diff --git a/Data/TrieMap/TrieKey.hs b/Data/TrieMap/TrieKey.hs deleted file mode 100644 --- a/Data/TrieMap/TrieKey.hs +++ /dev/null @@ -1,135 +0,0 @@ -{-# LANGUAGE TupleSections, TypeFamilies, UnboxedTuples, MagicHash #-} - -module Data.TrieMap.TrieKey where - -import Data.TrieMap.Sized - -import Control.Applicative -import Control.Monad - -import Data.Monoid -import Data.Foldable - -import Prelude hiding (foldr, foldl) - - -import GHC.Exts - -type EitherMap k a b c = k -> a -> (# Maybe b, Maybe c #) -type SplitMap a x = a -> (# Maybe a, Maybe x, Maybe a #) -type UnionFunc k a = k -> a -> a -> Maybe a -type IsectFunc k a b c = k -> a -> b -> Maybe c -type DiffFunc k a b = k -> a -> b -> Maybe a -type LEq a b = a -> b -> Bool - -onUnboxed :: (c -> d) -> (a -> (# b, c #)) -> a -> (# b, d #) -onUnboxed g f a = case f a of - (# b, c #) -> (# b, g c #) - -instance TrieKey k => Foldable (TrieMap k) where - foldr f z m = foldrWithKeyM (const f) m z - foldl f z m = foldlWithKeyM (const f) m z - -class Ord k => TrieKey k where - data TrieMap k :: * -> * - emptyM :: TrieMap k a - singletonM :: Sized a => k -> a -> TrieMap k a - nullM :: TrieMap k a -> Bool - sizeM :: Sized a => TrieMap k a -> Int# - lookupM :: k -> TrieMap k a -> Maybe a - mapWithKeyM :: Sized b => (k -> a -> b) -> TrieMap k a -> TrieMap k b - traverseWithKeyM :: (Applicative f, Sized b) => - (k -> a -> f b) -> TrieMap k a -> f (TrieMap k b) - foldrWithKeyM :: (k -> a -> b -> b) -> TrieMap k a -> b -> b - foldlWithKeyM :: (k -> b -> a -> b) -> TrieMap k a -> b -> b - mapMaybeM :: Sized b => (k -> a -> Maybe b) -> TrieMap k a -> TrieMap k b - mapEitherM :: (Sized b, Sized c) => EitherMap k a b c -> TrieMap k a -> (# TrieMap k b, TrieMap k c #) - unionM :: Sized a => UnionFunc k a -> TrieMap k a -> TrieMap k a -> TrieMap k a - isectM :: (Sized a, Sized b, Sized c) => IsectFunc k a b c -> TrieMap k a -> TrieMap k b -> TrieMap k c - diffM :: Sized a => DiffFunc k a b -> 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 => (k -> a -> a -> a) -> [(k, a)] -> TrieMap k a - fromDistAscListM :: Sized a => [(k, a)] -> TrieMap k a - - data Hole k :: * -> * - singleHoleM :: k -> Hole k a - keyM :: Hole k a -> k - beforeM :: Sized a => Maybe a -> Hole k a -> TrieMap k a - afterM :: Sized a => Maybe a -> Hole k a -> TrieMap k a - searchM :: k -> TrieMap k a -> (# Maybe a, Hole k a #) - indexM :: Sized a => Int# -> TrieMap k a -> (# Int#, a, Hole k a #) - {-# SPECIALIZE extractHoleM :: Sized a => TrieMap k a -> First (a, Hole k a) #-} - {-# SPECIALIZE extractHoleM :: Sized a => TrieMap k a -> Last (a, Hole k a) #-} - extractHoleM :: MonadPlus m => Sized a => TrieMap k a -> m (a, Hole k a) - assignM :: Sized a => a -> Hole k a -> TrieMap k a - clearM :: Sized a => Hole k a -> TrieMap k a - - singletonM k a = assignM a (singleHoleM k) - lookupM k m = case searchM k m of - (# a, _ #) -> a - foldrWithKeyM f = appEndo . getConst . traverseWithKeyM (endofy f) where - endofy :: (k -> a -> b -> b) -> k -> a -> Const (Endo b) (Elem ()) - endofy f k a = Const (Endo (f k a)) - foldlWithKeyM f m = foldrWithKeyM (\ k a g z -> g (f k z a)) m id - fromListM f = foldr (uncurry (insertWithKeyM f)) emptyM - fromAscListM = fromListM - fromDistAscListM = fromAscListM (const const) - -instance (TrieKey k, Sized a) => Sized (TrieMap k a) where - getSize# = sizeM - -{-# INLINE alterM #-} -alterM :: (TrieKey k, Sized a) => (Maybe a -> Maybe a) -> k -> TrieMap k a -> TrieMap k a -alterM f k m = case searchM k m of - (# Nothing, hole #) -> maybe m (\ a -> assignM a hole) (f Nothing) - (# a, hole #) -> fillHoleM (f a) hole - -traverseM :: (Applicative f, TrieKey k, Sized b) => (a -> f b) -> TrieMap k a -> f (TrieMap k b) -traverseM f = traverseWithKeyM (const f) - -guardNullM :: TrieKey k => TrieMap k a -> Maybe (TrieMap k a) -guardNullM m - | nullM m = Nothing - | otherwise = Just m - -fillHoleM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a -fillHoleM Nothing hole = clearM hole -fillHoleM (Just a) hole = assignM a hole - -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 #) - -fmapM :: (TrieKey k, Sized b) => (a -> b) -> TrieMap k a -> TrieMap k b -fmapM = mapWithKeyM . const - -assocsM :: TrieKey k => TrieMap k a -> [(k, a)] -assocsM m = build (\ f z -> foldrWithKeyM (\ k a xs -> (k, a) `f` xs) m z) - -insertWithKeyM :: (TrieKey k, Sized a) => (k -> a -> a -> a) -> k -> a -> TrieMap k a -> TrieMap k a -insertWithKeyM f k a m = case searchM k m of - (# Nothing, hole #) -> assignM a hole - (# Just a', hole #) -> assignM (f k a a') hole - -unionMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a -unionMaybe _ Nothing y = y -unionMaybe _ x Nothing = x -unionMaybe f (Just x) (Just y) = f x y - -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 diff --git a/Data/TrieMap/UnionMap.hs b/Data/TrieMap/UnionMap.hs --- a/Data/TrieMap/UnionMap.hs +++ b/Data/TrieMap/UnionMap.hs @@ -1,130 +1,182 @@ -{-# LANGUAGE PatternGuards, UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns, MagicHash #-} +{-# LANGUAGE UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns, MagicHash, CPP, BangPatterns #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.UnionMap () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized +import Data.TrieMap.UnitMap () import Control.Applicative import Control.Monad +import Data.Foldable (foldr) +import Prelude hiding (foldr, (^)) import GHC.Exts (&) :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a -m1 & m2 - | nullM m1, nullM m2 = Empty - | otherwise = Union (getSize# m1 +# getSize# m2) m1 m2 +m1 & m2 = guardNullM m1 ^ guardNullM m2 +{-# INLINE (^) #-} +(^) :: (TrieKey k1, TrieKey k2, Sized a) => Maybe (TrieMap k1 a) -> Maybe (TrieMap k2 a) -> TrieMap (Either k1 k2) a +Nothing ^ Nothing = Empty +Just m1 ^ Nothing = K1 m1 +Nothing ^ Just m2 = K2 m2 +Just m1 ^ Just m2 = Union (sizeM m1 +# sizeM m2) m1 m2 + +union :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a +union m1 m2 = Union (getSize# m1 +# getSize# m2) m1 m2 + singletonL :: (TrieKey k1, TrieKey k2, Sized a) => k1 -> a -> TrieMap (Either k1 k2) a -singletonL k a = Union (getSize# a) (singletonM k a) emptyM +singletonL k a = K1 (singletonM k a) singletonR :: (TrieKey k1, TrieKey k2, Sized a) => k2 -> a -> TrieMap (Either k1 k2) a -singletonR k a = Union (getSize# a) emptyM (singletonM k a) +singletonR k a = K2 (singletonM k a) -instance (TrieKey k1, TrieKey k2) => TrieKey (Either k1 k2) where - data TrieMap (Either k1 k2) a = Empty | Union Int# (TrieMap k1 a) (TrieMap k2 a) - data Hole (Either k1 k2) a = - LHole (Hole k1 a) (TrieMap k2 a) - | RHole (TrieMap k1 a) (Hole k2 a) +data UView k1 k2 a = UView (Maybe (TrieMap k1 a)) (Maybe (TrieMap k2 a)) +data HView k1 k2 a = Hole1 (Hole k1 a) (Maybe (TrieMap k2 a)) + | Hole2 (Maybe (TrieMap k1 a)) (Hole k2 a) +uView :: TrieMap (Either k1 k2) a -> UView k1 k2 a +uView Empty = UView Nothing Nothing +uView (K1 m1) = UView (Just m1) Nothing +uView (K2 m2) = UView Nothing (Just m2) +uView (Union _ m1 m2) = UView (Just m1) (Just m2) + +hView :: Hole (Either k1 k2) a -> HView k1 k2 a +hView (HoleX0 hole1) = Hole1 hole1 Nothing +hView (HoleX2 hole1 m2) = Hole1 hole1 (Just m2) +hView (Hole0X hole2) = Hole2 Nothing hole2 +hView (Hole1X m1 hole2) = Hole2 (Just m1) hole2 + +hole1 :: Hole k1 a -> Maybe (TrieMap k2 a) -> Hole (Either k1 k2) a +hole1 hole1 Nothing = HoleX0 hole1 +hole1 hole1 (Just m2) = HoleX2 hole1 m2 + +hole2 :: Maybe (TrieMap k1 a) -> Hole k2 a -> Hole (Either k1 k2) a +hole2 Nothing hole2 = Hole0X hole2 +hole2 (Just m1) hole2 = Hole1X m1 hole2 + +#define UVIEW uView -> UView + +instance (TrieKey k1, TrieKey k2) => TrieKey (Either k1 k2) where + {-# SPECIALIZE instance TrieKey (Either () ()) #-} + {-# SPECIALIZE instance TrieKey k => TrieKey (Either () k) #-} + {-# SPECIALIZE instance TrieKey k => TrieKey (Either k ()) #-} + Left k1 =? Left k2 = k1 =? k2 + Right k1 =? Right k2 = k1 =? k2 + _ =? _ = False + + Left k1 `cmp` Left k2 = k1 `cmp` k2 + Left{} `cmp` Right{} = LT + Right k1 `cmp` Right k2 = k1 `cmp` k2 + Right{} `cmp` Left{} = GT + + data TrieMap (Either k1 k2) a = + Empty + | K1 (TrieMap k1 a) + | K2 (TrieMap k2 a) + | Union Int# (TrieMap k1 a) (TrieMap k2 a) + data Hole (Either k1 k2) a = + HoleX0 (Hole k1 a) + | HoleX2 (Hole k1 a) (TrieMap k2 a) + | Hole0X (Hole k2 a) + | Hole1X (TrieMap k1 a) (Hole k2 a) emptyM = Empty singletonM = either singletonL singletonR - nullM Empty = True - nullM _ = False + getSimpleM (UVIEW m1 m2) = mSimple m1 `mplus` mSimple m2 where + mSimple :: TrieKey k => Maybe (TrieMap k a) -> Simple a + mSimple = maybe mzero getSimpleM sizeM Empty = 0# + sizeM (K1 m1) = sizeM m1 + sizeM (K2 m2) = sizeM m2 sizeM (Union s _ _) = s - lookupM k (Union _ m1 m2) = either (`lookupM` m1) (`lookupM` m2) k - lookupM _ _ = Nothing + lookupM (Left k) (UVIEW m1 _) = m1 >>= lookupM k + lookupM (Right k) (UVIEW _ m2) = m2 >>= lookupM k - traverseWithKeyM f (Union _ m1 m2) = (&) <$> traverseWithKeyM (f . Left) m1 <*> traverseWithKeyM (f . Right) m2 - traverseWithKeyM _ _ = pure Empty + 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 - foldrWithKeyM f (Union _ m1 m2) = foldrWithKeyM (f . Left) m1 . foldrWithKeyM (f . Right) m2 - foldrWithKeyM _ _ = id + foldrM f (UVIEW m1 m2) = fold (foldrM f) m1 . fold (foldrM f) m2 + where fold :: (a -> b -> b) -> Maybe a -> b -> b + fold = flip . foldr - foldlWithKeyM f (Union _ m1 m2) = foldlWithKeyM (f . Right) m2 . foldlWithKeyM (f . Left) m1 - foldlWithKeyM _ _ = id + foldlM f (UVIEW m1 m2) = fold (foldlM f) m2 . fold (foldlM f) m1 + where fold :: (a -> b -> b) -> Maybe a -> b -> b + fold = flip . foldr - mapWithKeyM f (Union _ m1 m2) = mapWithKeyM (f . Left) m1 & mapWithKeyM (f . Right) m2 - mapWithKeyM _ _ = 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 (Union _ m1 m2) = mapMaybeM (f . Left) m1 & mapMaybeM (f . Right) m2 - mapMaybeM _ _ = Empty + mapMaybeM f (UVIEW m1 m2) = (m1 >>= mapMaybeM' f) ^ (m2 >>= mapMaybeM' f) - mapEitherM f (Union _ m1 m2) - | (# m1L, m1R #) <- mapEitherM (f . Left) m1, - (# m2L, m2R #) <- mapEitherM (f . Right) m2 - = (# m1L & m2L, m1R & m2R #) - mapEitherM _ _ = (# Empty, Empty #) + mapEitherM f (UVIEW m1 m2) = (# m1L ^ m2L, m1R ^ m2R #) where + !(# m1L, m1R #) = mapEitherM'' f m1 + !(# m2L, m2R #) = mapEitherM'' f m2 - unionM f (Union _ m11 m12) (Union _ m21 m22) - = unionM (f . Left) m11 m21 & unionM (f . Right) m12 m22 - unionM _ Empty m2 = m2 - unionM _ m1 Empty = m1 + 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 _ Empty _ = Empty - isectM _ _ Empty = Empty - isectM f (Union _ m11 m12) (Union _ m21 m22) - = isectM (f . Left) m11 m21 & isectM (f . Right) m12 m22 + isectM f (UVIEW m11 m12) (UVIEW m21 m22) = + isectMaybe (isectM' f) m11 m21 ^ isectMaybe (isectM' f) m12 m22 - diffM _ Empty _ = Empty - diffM _ m1 Empty = m1 - diffM f (Union _ m11 m12) (Union _ m21 m22) - = diffM (f . Left) m11 m21 & diffM (f . Right) 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 _ Empty _ = True - isSubmapM (<=) (Union _ m11 m12) (Union _ m21 m22) = isSubmapM (<=) m11 m21 && isSubmapM (<=) m12 m22 - isSubmapM _ Union{} Empty = False + isSubmapM (<=) (UVIEW m11 m12) (UVIEW m21 m22) = + subMaybe (isSubmapM (<=)) m11 m21 && subMaybe (isSubmapM (<=)) m12 m22 - fromListM f = onPair (&) (fromListM (f . Left)) (fromListM (f . Right)) . partEithers + fromListM f = onPair (&) (fromListM f) (fromListM f) . partEithers - fromAscListM f = onPair (&) (fromAscListM (f . Left)) (fromAscListM (f . Right)) . partEithers + fromAscListM f = onPair (&) (fromAscListM f) (fromAscListM f) . partEithers fromDistAscListM = onPair (&) fromDistAscListM fromDistAscListM . partEithers - singleHoleM (Left k) = LHole (singleHoleM k) emptyM - singleHoleM (Right k) = RHole emptyM (singleHoleM k) - - keyM (LHole holeL _) = Left (keyM holeL) - keyM (RHole _ holeR) = Right (keyM holeR) - - beforeM a (LHole holeL _) = let mL = beforeM a holeL in - if nullM mL then Empty else Union (getSize# mL) mL emptyM - beforeM a (RHole mL holeR) = mL & beforeM a holeR + singleHoleM = either (HoleX0 . singleHoleM) (Hole0X . singleHoleM) + + beforeM a hole = case hView hole of + Hole1 h1 __ -> beforeM' a h1 ^ Nothing + Hole2 m1 h2 -> m1 ^ beforeM' a h2 - afterM a (LHole holeL mR) = afterM a holeL & mR - afterM a (RHole _ holeR) = let mR = afterM a holeR in - if nullM mR then Empty else Union (getSize# mR) emptyM mR + afterM a hole = case hView hole of + Hole1 h1 m2 -> afterM' a h1 ^ m2 + Hole2 __ h2 -> Nothing ^ afterM' a h2 - searchM k Empty = (# Nothing, singleHoleM k #) - searchM (Left k) (Union _ mL mR) = onUnboxed (`LHole` mR) (searchM k) mL - searchM (Right k) (Union _ mL mR) = onUnboxed (RHole mL) (searchM k) mR + searchM (Left k) (UVIEW m1 m2) = onSnd (`hole1` m2) (searchM' k) m1 + searchM (Right k) (UVIEW m1 m2) = onSnd (hole2 m1) (searchM' k) m2 - indexM i# (Union _ mL mR) - | i# <# sL#, (# i'#, v, holeL #) <- indexM i# mL - = (# i'#, v, LHole holeL mR #) - | (# i'#, v, holeR #) <- indexM (i# -# sL#) mR - = (# i'#, v, RHole mL holeR #) - where !sL# = getSize# mL - indexM _ _ = (# error err, error err, error err #) where - err = "Error: empty trie" + indexM i# (K1 m1) = onThird HoleX0 (indexM i#) m1 + indexM i# (K2 m2) = onThird Hole0X (indexM i#) m2 + indexM i# (Union _ m1 m2) + | i# <# s1# = onThird (`HoleX2` m2) (indexM i#) m1 + | otherwise = onThird (Hole1X m1) (indexM (i# -# s1#)) m2 + where !s1# = sizeM m1 + indexM _ _ = indexFail () - extractHoleM (Union _ mL mR) = (do - (v, holeL) <- extractHoleM mL - return (v, LHole holeL mR)) `mplus` (do - (v, holeR) <- extractHoleM mR - return (v, RHole mL holeR)) - extractHoleM _ = mzero + extractHoleM (UVIEW m1 m2) = (do + (v, h1) <- extractHoleM' m1 + return (v, hole1 h1 m2)) `mplus` (do + (v, h2) <- extractHoleM' m2 + return (v, hole2 m1 h2)) - assignM v (LHole holeL mR) = assignM v holeL & mR - assignM v (RHole mL holeR) = mL & assignM v holeR - - clearM (LHole holeL mR) = clearM holeL & mR - clearM (RHole mL holeR) = mL & clearM holeR + assignM v hole = case hView hole of + Hole1 h1 m2 -> assignM' v h1 ^ m2 + Hole2 m1 h2 -> m1 ^ assignM' v h2 + + unifyM (Left k1) a1 (Left k2) a2 = either (Left . HoleX0) (Right . K1) (unifyM k1 a1 k2 a2) + unifyM (Left k1) a1 (Right k2) a2 = Right $ singletonM k1 a1 `union` singletonM k2 a2 + unifyM (Right k2) a2 (Left k1) a1 = Right $ singletonM k1 a1 `union` singletonM k2 a2 + unifyM (Right k1) a1 (Right k2) a2 = either (Left . Hole0X) (Right . K2) (unifyM k1 a1 k2 a2) onPair :: (c -> d -> e) -> (a -> c) -> (b -> d) -> (a, b) -> e onPair f g h (a, b) = f (g a) (h b) diff --git a/Data/TrieMap/UnitMap.hs b/Data/TrieMap/UnitMap.hs --- a/Data/TrieMap/UnitMap.hs +++ b/Data/TrieMap/UnitMap.hs @@ -15,41 +15,42 @@ import Prelude hiding (foldr, foldl) instance TrieKey () where + _ =? _ = True + _ `cmp` _ = EQ + newtype TrieMap () a = Unit {getUnit :: Maybe a} data Hole () a = Hole emptyM = Unit Nothing singletonM _ = Unit . Just - nullM = isNothing . getUnit + getSimpleM (Unit m) = maybe Null Singleton m sizeM (Unit (Just a)) = getSize# a sizeM _ = 0# lookupM _ (Unit m) = m - traverseWithKeyM f (Unit m) = Unit <$> traverse (f ()) m - foldrWithKeyM f (Unit m) z = foldr (f ()) z m - foldlWithKeyM f (Unit m) z = foldl (f ()) z m - mapWithKeyM f (Unit m) = Unit (f () <$> m) - mapMaybeM f (Unit m) = Unit (m >>= f ()) - mapEitherM f (Unit (Just a)) = both Unit Unit (f ()) a - mapEitherM _ _ = (# emptyM, emptyM #) - 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) + traverseM f (Unit m) = Unit <$> traverse f m + foldrM f (Unit m) z = foldr f z m + foldlM f (Unit m) z = foldl f z 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 fromListM _ [] = Unit Nothing - fromListM f ((_, v):xs) = Unit $ Just (foldl (\ v' -> f () v' . snd) v xs) + fromListM f ((_, v):xs) = Unit $ Just (foldl (\ v' -> f v' . snd) v xs) singleHoleM _ = Hole - keyM _ = () beforeM a _ = Unit a afterM a _ = Unit a searchM _ (Unit m) = (# m, Hole #) indexM i (Unit (Just v)) = (# i, v, Hole #) - indexM _ _ = (# error err, error err, error err #) where - err = "Error: empty trie" + indexM _ _ = indexFail () + unifyM _ _ _ _ = Left Hole + extractHoleM (Unit (Just v)) = return (v, Hole) extractHoleM _ = mzero - assignM v _ = Unit (Just v) - clearM _ = emptyM + assignM v _ = Unit v diff --git a/Data/TrieMap/Utils.hs b/Data/TrieMap/Utils.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/Utils.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE Rank2Types, BangPatterns, MagicHash #-} +module Data.TrieMap.Utils (toVectorN, toVectorF) where + +import Data.Vector.Generic +import Data.Vector.Generic.Mutable +import qualified Data.Foldable +import GHC.Exts + +{-# 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 + !mv <- unsafeNew (size xs) + fold (\ x m i# -> unsafeWrite mv (I# i#) x >> m (i# +# 1#)) (\ _ -> return mv) xs 0# + +{-# INLINE toVectorF #-} +toVectorF :: (Vector v b, Data.Foldable.Foldable f) => (a -> b) -> (f a -> Int) -> f a -> v b +toVectorF g = toVectorN (\ f -> Data.Foldable.foldr (f . g)) diff --git a/Tests.hs b/Tests.hs --- a/Tests.hs +++ b/Tests.hs @@ -2,15 +2,19 @@ -- module Tests where import Control.Monad +import Debug.Trace +import Data.TrieMap.Class +import Data.TrieMap.TrieKey +import Data.TrieMap.Sized import qualified Data.TrieMap as T import qualified Data.Map as M import Test.QuickCheck import Prelude hiding (null, lookup) -type Key = [Int] -type Val = [Int] +type Key = Integer +type Val = [Integer] -main = quickCheckWith stdArgs{maxSize = 800, maxSuccess = 800} (verify M.empty T.empty) +main = quickCheckWith stdArgs{maxSize = 300, maxSuccess = 100} (verify M.empty T.empty) instance Arbitrary Op where arbitrary = oneof [ @@ -30,7 +34,7 @@ shrink (Op (Insert k v)) = [Op (Insert k' v') | k' <- shrink k, v' <- shrink v] shrink (Op (Lookup k)) = map (Op . Lookup) (shrink k) shrink (Op (Delete k)) = map (Op . Delete) (shrink k) - shrink (Op (Union ops)) = map (Op . Union) (shrink ops) + shrink (Op (Union ops)) = ops ++ map (Op . Union) (shrink ops) shrink _ = [] recurse :: Gen [Op] @@ -51,6 +55,7 @@ show (Op (Union ops)) = "Union " ++ show ops show (Op (DeleteAt i)) = "DeleteAt " ++ show i show (Op (ElemAt i)) = "ElemAt " ++ show i + show (Op (Isect ops)) = "Isect " ++ show ops data Operation r where Insert :: Key -> Val -> Operation () @@ -68,15 +73,15 @@ ElemAt :: Int -> Operation (Maybe (Key, Val)) mapFunc :: Key -> Val -> Val -mapFunc = (++) +mapFunc = (:) mapMaybeFunc :: Key -> Val -> Maybe Val -mapMaybeFunc (k:ks) xs - | even k = Just (ks ++ xs) +mapMaybeFunc k xs + | even k = Just (k:xs) mapMaybeFunc _ _ = Nothing isectFunc :: Key -> Val -> Val -> Val -isectFunc ks xs ys = ks ++ xs ++ ys +isectFunc ks xs ys = ks:xs ++ ys generateMap :: M.Map Key Val -> [Op] -> M.Map Key Val generateMap = foldl (\ mm (Op op) -> snd (operateMap mm op)) diff --git a/TrieMap.cabal b/TrieMap.cabal --- a/TrieMap.cabal +++ b/TrieMap.cabal @@ -1,39 +1,53 @@ name: TrieMap -version: 1.5.0 +version: 2.0.0 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 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. license: BSD3 license-file: LICENSE author: Louis Wasserman maintainer: wasserman.louis@gmail.com -build-Depends: base < 5.0.0.0, containers, template-haskell, bytestring, array, th-expand-syns, ghc-prim +build-Depends: base < 5.0.0.0, containers, template-haskell, bytestring, th-expand-syns, ghc-prim, vector, primitive build-type: Simple -ghc-options: -Wall -fno-warn-name-shadowing -fno-warn-orphans +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 extra-source-files: Tests.hs exposed-modules: Data.TrieMap, Data.TrieSet, + Data.TrieMap.Class, Data.TrieMap.Representation, - Data.TrieMap.Representation.TH, Data.TrieMap.Modifiers other-modules: - Data.TrieMap.Class, - Data.TrieMap.Class.Instances, - Data.TrieMap.Key, - Data.TrieMap.TrieKey, + Data.TrieMap.Utils, + Data.TrieMap.Sized, Data.TrieMap.Applicative, - Data.TrieMap.ProdMap, - Data.TrieMap.RadixTrie, - Data.TrieMap.UnionMap, - Data.TrieMap.UnitMap, - Data.TrieMap.Rep, - Data.TrieMap.Rep.Instances, - Data.TrieMap.Rep.TH, + Data.TrieMap.Representation.Class, + Data.TrieMap.Representation.TH, + Data.TrieMap.Representation.TH.Utils, + Data.TrieMap.Representation.TH.Representation, + Data.TrieMap.Representation.TH.Factorized, + Data.TrieMap.Representation.TH.ReprMonad, + Data.TrieMap.Representation.Instances, + Data.TrieMap.Representation.Instances.Basic, + Data.TrieMap.Representation.Instances.Prim, + Data.TrieMap.Representation.Instances.Foreign, + Data.TrieMap.Representation.Instances.Vectors, + Data.TrieMap.Representation.Instances.ByteString Data.TrieMap.IntMap, Data.TrieMap.OrdMap, + Data.TrieMap.UnitMap, + Data.TrieMap.ProdMap, + Data.TrieMap.UnionMap, Data.TrieMap.ReverseMap, - Data.TrieMap.Sized, - Data.TrieMap.Applicative + Data.TrieMap.Key, + Data.TrieMap.RadixTrie, + Data.TrieMap.RadixTrie.Slice, + Data.TrieMap.RadixTrie.Edge, + Data.TrieMap.Class.Instances +