diff --git a/Control/Monad/Ends.hs b/Control/Monad/Ends.hs new file mode 100644 --- /dev/null +++ b/Control/Monad/Ends.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Control.Monad.Ends where + +import Control.Monad + +newtype First a = First {getFirst :: Maybe a} deriving (Functor, Monad) +newtype Last a = Last {getLast :: Maybe a} deriving (Functor, Monad) + +instance MonadPlus First where + mzero = First Nothing + First Nothing `mplus` m = m + m `mplus` _ = m + +instance MonadPlus Last where + mzero = Last Nothing + m `mplus` Last Nothing = m + _ `mplus` m = m 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, RecordWildCards #-} +{-# LANGUAGE UnboxedTuples #-} module Data.TrieMap ( -- * Map type @@ -125,18 +125,21 @@ maxViewWithKey ) where +import Control.Monad.Ends + import Data.TrieMap.Class import Data.TrieMap.Class.Instances() import Data.TrieMap.TrieKey -import Data.TrieMap.Applicative import Data.TrieMap.Representation import Data.TrieMap.Representation.Instances () import Data.TrieMap.Sized +import Data.TrieMap.Utils import Control.Applicative hiding (empty) import Control.Monad +import qualified Data.Foldable as F import Data.Maybe hiding (mapMaybe) -import Data.Monoid(Monoid(..), First(..), Last(..)) +import Data.Monoid(Monoid(..)) import GHC.Exts (build) @@ -183,7 +186,7 @@ -- 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) = getValue <$> lookupM (toRep k) m +lookup k (TMap m) = option (lookupM (toRep k) m) Nothing (Just . getValue) -- | The expression @('findWithDefault' def k map)@ returns the value at key @k@ or returns default value @def@ -- when the key is not in the map. @@ -201,11 +204,11 @@ -- @'lookup' k ('alter' f k m) = f ('lookup' k m)@. {-# INLINE alter #-} alter :: TKey k => (Maybe a -> Maybe a) -> k -> TMap k a -> TMap k a -alter f k m = case search k m of - (Nothing, hole) -> case f Nothing of - Nothing -> m - Just a' -> assign a' hole - (a, hole) -> fillHole (f a) hole +alter f k (TMap m) = TMap $ searchMC (toRep k) m nomatch match where + nomatch hole = case f Nothing of + Nothing -> m + Just a' -> assignM (Assoc k a') hole + match (Assoc _ a) hole = fillHoleM (Assoc k <$> f (Just a)) hole -- | Insert a new key and value in the map. -- If the key is already present in the map, the associated value is @@ -245,8 +248,8 @@ -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" {-# INLINE insertWithKey #-} insertWithKey :: TKey k => (k -> a -> a -> a) -> k -> a -> TMap k a -> TMap k a -insertWithKey f k a m = snd (insertLookupWithKey f k a m) - +insertWithKey f k a (TMap m) = + TMap (insertWithM (\ (Assoc _ a0) -> Assoc k (f k a a0)) (toRep k) (Assoc k a) m) -- | Combines insert operation with old value retrieval. -- The expression (@'insertLookupWithKey' f k x map@) @@ -328,13 +331,13 @@ -- value to the highest. {-# INLINE foldrWithKey #-} foldrWithKey :: TKey k => (k -> a -> b -> b) -> b -> TMap k a -> b -foldrWithKey f z (TMap m) = foldrM (\ (Assoc k a) -> f k a) m z +foldrWithKey f z (TMap m) = F.foldr (\ (Assoc k a) -> f k a) z m -- | 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) = foldlM (\ z (Assoc k a) -> f z k a) m z +foldlWithKey f z (TMap m) = F.foldl (\ z (Assoc k a) -> f z k a) z m -- | Map each key\/element pair to an action, evaluate these actions from left to right, and collect the results. {-# INLINE traverseWithKey #-} @@ -346,7 +349,7 @@ -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] {-# INLINE map #-} map :: TKey k => (a -> b) -> TMap k a -> TMap k b -map f = mapWithKey (const f) +map = fmap -- | Map a function over all values in the map. -- @@ -597,10 +600,9 @@ updateMax = updateMaxWithKey . const {-# INLINE updateHelper #-} -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 - (Assoc k a, loc) <- extractHoleM m - return (Assoc k <$> f k a, loc) +updateHelper :: (TKey k, Functor m, MonadPlus m) => + (k -> a -> Maybe a) -> TMap k a -> m (Maybe (Assoc k a), Hole (Rep k) (Assoc k a)) +updateHelper f (TMap m) = fmap (\ (Assoc k a, loc) -> (Assoc k <$> f k a, loc)) (extractHoleM m) -- | Update the value at the minimal key. -- @@ -610,7 +612,7 @@ updateMinWithKey :: TKey k => (k -> a -> Maybe a) -> TMap k a -> TMap k a updateMinWithKey f m = fromMaybe m $ do (a, loc) <- getFirst $ updateHelper f m - return (TMap (afterM a loc)) + return (TMap (afterMM a loc)) -- | Update the value at the maximal key. -- @@ -620,7 +622,7 @@ updateMaxWithKey :: TKey k => (k -> a -> Maybe a) -> TMap k a -> TMap k a updateMaxWithKey f m = fromMaybe m $ do (a, loc) <- getLast $ updateHelper f m - return (TMap (afterM a loc)) + return (TMap (beforeMM a loc)) -- | Delete and find the minimal element. -- @@ -916,7 +918,7 @@ -- > keysSet empty == Data.TrieSet.empty {-# INLINE keysSet #-} keysSet :: TKey k => TMap k a -> TSet k -keysSet m = TSet (() <$ m) +keysSet (TMap m) = TSet (fmapM (\ (Assoc k _) -> Elem k) m) -- | /O(1)/. The key marking the position of the \"hole\" in the map. {-# INLINE key #-} @@ -926,12 +928,12 @@ -- | @'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 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 hole) -- | Search the map for the given key, returning the -- corresponding value (if any) and an updatable location for that key. @@ -947,9 +949,9 @@ -- @'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 - (# Just (Assoc k a), hole #) -> (Just a, TLoc k hole) - (# _, hole #) -> (Nothing, TLoc k hole) +search k (TMap m) = searchMC (toRep k) m nomatch match where + nomatch hole = (Nothing, TLoc k hole) + match (Assoc k a) hole = (Just a, 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. @@ -968,14 +970,12 @@ index i m | i < 0 || i >= size m = error "TrieMap.index: index out of range" -index i (TMap m) = case indexM (unbox i) m of +index i (TMap m) = case indexM i m of (# _, 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 - (Assoc k a, hole) <- extractHoleM m - return (a, TLoc k hole) +extract :: (TKey k, Functor m, MonadPlus m) => TMap k a -> m (a, TLocation k a) +extract (TMap m) = fmap (\ (Assoc k a, hole) -> (a, TLoc k hole)) (extractHoleM m) -- | /O(log n)/. Return the value and an updatable location for the -- least key in the map, or 'Nothing' if the map is empty. @@ -1015,14 +1015,14 @@ -- @'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 k hole) = TMap (assignM (Just $ Assoc k a) hole) +assign a (TLoc k hole) = TMap (assignM (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 (assignM Nothing hole) +clear (TLoc _ hole) = TMap (clearM hole) {-# INLINE fillHole #-} fillHole :: TKey k => Maybe a -> TLocation k a -> TMap k a diff --git a/Data/TrieMap/Applicative.hs b/Data/TrieMap/Applicative.hs deleted file mode 100644 --- a/Data/TrieMap/Applicative.hs +++ /dev/null @@ -1,68 +0,0 @@ -{-# LANGUAGE StandaloneDeriving, GeneralizedNewtypeDeriving #-} - -module Data.TrieMap.Applicative where - -import Control.Applicative -import Control.Monad - -import Data.Monoid hiding (Dual) - -instance Functor First where - fmap f (First m) = First (fmap f m) - -instance Functor Last where - fmap f (Last m) = Last (fmap f m) - -instance Monad First where - return = First . return - First m >>= k = First (m >>= getFirst . k) - -instance Monad Last where - return = Last . return - Last m >>= k = Last (m >>= getLast . k) - -instance MonadPlus First where - mzero = mempty - mplus = mappend - -instance MonadPlus Last where - mzero = mempty - mplus = mappend - -(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d -(f .: g) x y = f (g x y) - -(<.>) :: Functor f => (b -> c) -> (a -> f b) -> a -> f c -f <.> g = fmap f . g - -(<.:>) :: Functor f => (c -> d) -> (a -> b -> f c) -> a -> b -> f d -(f <.:> g) x y = f <$> g x y - -instance Applicative First where - pure = return - (<*>) = ap - -instance Alternative First where - empty = mempty - (<|>) = mplus - -instance Applicative Last where - pure = return - (<*>) = ap - -instance Alternative Last where - empty = mempty - (<|>) = mplus - -newtype DualPlus f a = DualPlus {runDualPlus :: f a} deriving (Functor, Applicative, Monad) -newtype Dual f a = Dual {runDual :: f a} deriving (Functor) - -instance Applicative f => Applicative (Dual f) where - pure = Dual . pure - Dual f <*> Dual a = Dual (a <**> f) - Dual f *> Dual g = Dual (g <* f) - Dual f <* Dual g = Dual (g *> f) - -instance MonadPlus m => MonadPlus (DualPlus m) where - mzero = DualPlus mzero - DualPlus m `mplus` DualPlus k = DualPlus (k `mplus` m) diff --git a/Data/TrieMap/Class.hs b/Data/TrieMap/Class.hs --- a/Data/TrieMap/Class.hs +++ b/Data/TrieMap/Class.hs @@ -6,28 +6,38 @@ import Data.TrieMap.Representation.Class import Data.TrieMap.Sized -import Control.Applicative -import Data.Foldable hiding (foldrM, foldlM) +import Data.Functor +import Data.Foldable import Data.Traversable -import Prelude hiding (foldr) +import Prelude hiding (foldr, foldl, foldl1, foldr1) +-- | A map from keys @k@ to values @a@, backed by a trie. newtype TMap k a = TMap {getTMap :: TrieMap (Rep k) (Assoc k a)} -newtype TSet a = TSet (TMap a ()) +-- | A set of values @a@, backed by a trie. +newtype TSet a = TSet {getTSet :: TrieMap (Rep a) (Elem 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)@. +-- create a 'Repr' instance that will satisfy @'TrieKey' ('Rep' k)@, possibly using the Template Haskell methods +-- provided by "Data.TrieMap.Representation". class (Repr k, TrieKey (Rep k)) => TKey k instance (Repr k, TrieKey (Rep k)) => TKey k instance TKey k => Functor (TMap k) where - fmap = fmapDefault + fmap f (TMap m) = TMap (fmapM (fmap f) m) instance TKey k => Foldable (TMap k) where - foldr f z (TMap m) = foldrM (\ (Assoc _ a) -> f a) m z + foldMap f (TMap m) = foldMap (foldMap f) m + foldr f z (TMap m) = foldr (flip $ foldr f) z m + foldl f z (TMap m) = foldl (foldl f) z m + foldr1 f (TMap m) = getElem (foldr1 f' m') where + f' (Elem a) (Elem b) = Elem (f a b) + m' = fmapM (\ (Assoc _ a) -> Elem a) m + foldl1 f (TMap m) = getElem (foldl1 f' m') where + f' (Elem a) (Elem b) = Elem (f a b) + m' = fmapM (\ (Assoc _ a) -> Elem a) m instance TKey k => Traversable (TMap k) where - traverse f (TMap m) = TMap <$> traverseM (\ (Assoc k a) -> Assoc k <$> f a) m + traverse f (TMap m) = TMap <$> traverseM (traverse f) 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 @@ -6,7 +6,7 @@ import Data.TrieMap.Sized () import Data.TrieMap.ReverseMap () import Data.TrieMap.RadixTrie () -import Data.TrieMap.IntMap () +import Data.TrieMap.WordMap () import Data.TrieMap.OrdMap () import Data.TrieMap.ProdMap () import Data.TrieMap.UnionMap () diff --git a/Data/TrieMap/IntMap.hs b/Data/TrieMap/IntMap.hs deleted file mode 100644 --- a/Data/TrieMap/IntMap.hs +++ /dev/null @@ -1,318 +0,0 @@ -{-# LANGUAGE UnboxedTuples, BangPatterns, TypeFamilies, PatternGuards, MagicHash, CPP #-} -{-# OPTIONS -funbox-strict-fields #-} -module Data.TrieMap.IntMap () where - -import Data.TrieMap.TrieKey -import Data.TrieMap.Sized - -import Control.Applicative -import Control.Monad hiding (join) - -import Data.Bits -import Data.Maybe hiding (mapMaybe) -import Data.Word - -import GHC.Exts - -import Prelude hiding (lookup, null, foldl, foldr) - -#include "MachDeps.h" -type Nat = Word - -type Prefix = Word -type Mask = Word -type Key = Word -type Size = Int# - -data Path a = Root - | LeftBin !Prefix !Mask !(Path a) !(TrieMap Word a) - | RightBin !Prefix !Mask !(TrieMap Word a) !(Path a) - --- | @'TrieMap' 'Word' a@ is based on "Data.IntMap". -instance TrieKey Word where - (=?) = (==) - cmp = compare - - data TrieMap Word a = Nil - | Tip !Size !Key a - | Bin !Size !Prefix !Mask !(TrieMap Word a) !(TrieMap Word a) - data Hole Word a = Hole !Key !(Path a) - emptyM = Nil - singletonM = singleton - getSimpleM Nil = Null - getSimpleM (Tip _ _ a) = Singleton a - getSimpleM _ = NonSimple - sizeM = size - lookupM = lookup - traverseM = traverse - foldrM = foldr - foldlM = foldl - fmapM = mapWithKey - mapMaybeM = mapMaybe - mapEitherM = mapEither - unionM = unionWith - isectM = intersectionWith - diffM = differenceWith - isSubmapM = isSubmapOfBy - - singleHoleM k = Hole k Root - beforeM a (Hole k path) = before (singletonMaybe k a) path where - before t Root = t - before t (LeftBin _ _ path _) = before t path - before t (RightBin p m l path) = before (bin p m l t) path - afterM a (Hole k path) = after (singletonMaybe k a) path where - 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 = 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 - = search (LeftBin p m path r) l - | otherwise - = search (RightBin p m l path) r - search path t@(Tip _ ky y) - | k == ky = (# Just y, path #) - | otherwise = (# Nothing, branchHole k ky path t #) - search path _ = (# Nothing, path #) - indexM i# t = indexT i# t Root where - 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) - | otherwise = indexT (i# -# sl#) r (RightBin p m l path) - where !sl# = size l - extractHoleM = extractHole Root where - extractHole _ Nil = mzero - extractHole path (Tip _ kx x) = return (x, Hole kx path) - extractHole path (Bin _ p m l r) = - extractHole (LeftBin p m path r) l `mplus` - extractHole (RightBin p m l path) r - 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 - - {-# INLINE unifyM #-} - unifyM = unify - -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 :: Word -> Nat -natFromInt = id - -intFromNat :: Nat -> Word -intFromNat = id - -shiftRL :: Nat -> Key -> Nat --- #if __GLASGOW_HASKELL__ -{-------------------------------------------------------------------- - GHC: use unboxing to get @shiftRL@ inlined. ---------------------------------------------------------------------} --- shiftRL (W# x) (I# i) --- = W# (shiftRL# x i) --- #else -shiftRL x i = shiftR x (fromIntegral i) --- #endif - -size :: TrieMap Word a -> Int# -size Nil = 0# -size (Tip sz _ _) = sz -size (Bin sz _ _ _ _) = sz - -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 Word a -singleton k a = Tip (getSize# a) k a - -singletonMaybe :: Sized a => Key -> Maybe a -> TrieMap Word a -singletonMaybe k = maybe Nil (singleton k) - -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 x - Bin _ p m l r -> bin p m <$> traverse f l <*> traverse f r - -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 _ _ x -> f x - Nil -> id - -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 _ _ x -> flip f x - Nil -> id - -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 x) -mapWithKey _ _ = Nil - -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 x) -mapMaybe _ _ = Nil - -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 #) - -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 (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 (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 (unionWith f t1 l2) r2 - | otherwise = bin p2 m2 l2 (unionWith f t1 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 (intersectionWith f l1 l2) (intersectionWith f r1 r2) - | otherwise = Nil - where - intersection1 | nomatch p2 p1 m1 = Nil - | zero p2 m1 = intersectionWith f l1 t2 - | otherwise = intersectionWith f r1 t2 - - intersection2 | nomatch p1 p2 m2 = Nil - | zero p1 m2 = intersectionWith f t1 l2 - | otherwise = intersectionWith f t1 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 (differenceWith f l1 l2) (differenceWith f r1 r2) - | otherwise = t1 - where - difference1 | nomatch p2 p1 m1 = t1 - | 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 = differenceWith f t1 l2 - | otherwise = differenceWith f t1 r2 - -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 - else isSubmapOfBy (<=) t1 r2) - | otherwise = (p1==p2) && isSubmapOfBy (<=) l1 l2 && isSubmapOfBy (<=) r1 r2 -isSubmapOfBy _ (Bin _ _ _ _ _) _ - = False -isSubmapOfBy (<=) (Tip _ k x) t - = maybe False (x <=) (lookup (natFromInt k) t) -isSubmapOfBy _ Nil _ - = True - -mask :: Key -> Mask -> Prefix -mask i m - = maskW (natFromInt i) (natFromInt m) - -zero :: Key -> Mask -> Bool -zero i m - = (natFromInt i) .&. (natFromInt m) == 0 - -nomatch,match :: Key -> Prefix -> Mask -> Bool -nomatch i p m - = (mask i m) /= p - -match i p m - = (mask i m) == p - -zeroN :: Nat -> Nat -> Bool -zeroN i m = (i .&. m) == 0 - -maskW :: Nat -> Nat -> Prefix -maskW i m - = intFromNat (i .&. (complement (m-1) `xor` m)) - -shorter :: Mask -> Mask -> Bool -shorter m1 m2 - = (natFromInt m1) > (natFromInt m2) - -branchMask :: Prefix -> Prefix -> Mask -branchMask p1 p2 - = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2)) - -highestBitMask :: Nat -> Nat -highestBitMask x0 - = case (x0 .|. shiftRL x0 1) of - x1 -> case (x1 .|. shiftRL x1 2) of - x2 -> case (x2 .|. shiftRL x2 4) of - x3 -> case (x3 .|. shiftRL x3 8) of - x4 -> case (x4 .|. shiftRL x4 16) of -#if WORD_SIZE_IN_BITS > 32 - x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms - x6 -> (x6 `xor` (shiftRL x6 1)) -#else - x5 -> x5 `xor` shiftRL x5 1 -#endif - -{-# 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 - where - m = branchMask p1 p2 - p = mask p1 m - -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,53 +1,61 @@ -{-# LANGUAGE TypeFamilies, UnboxedTuples #-} - +{-# LANGUAGE TypeFamilies, MagicHash, CPP, FlexibleInstances #-} +{-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.Key () where -import Control.Applicative +import Data.Functor +import Data.Foldable import Data.TrieMap.Class import Data.TrieMap.TrieKey +import Data.TrieMap.Sized import Data.TrieMap.Representation.Class import Data.TrieMap.Modifiers -import Data.TrieMap.ProdMap() -import Data.TrieMap.UnionMap() -import Data.TrieMap.IntMap() -import Data.TrieMap.OrdMap() -import Data.TrieMap.RadixTrie() +import Prelude hiding (foldr, foldl, foldr1, foldl1) +keyMap :: (TKey k, Sized a) => TrieMap (Rep k) a -> TrieMap (Key k) a +keyMap m = KeyMap (sizeM m) m + +#define KMAP(m) KeyMap{tMap = m} + +instance TKey k => Foldable (TrieMap (Key k)) where + foldMap f KMAP(m) = foldMap f m + foldr f z KMAP(m) = foldr f z m + foldl f z KMAP(m) = foldl f z m + foldr1 f KMAP(m) = foldr1 f m + foldl1 f KMAP(m) = foldl1 f m + -- | @'TrieMap' ('Key' k) a@ is a wrapper around a @TrieMap (Rep k) a@. 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) + data TrieMap (Key k) a = KeyMap {sz :: !Int, tMap :: !(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) - getSimpleM (KeyMap m) = getSimpleM m - sizeM (KeyMap m) = sizeM m - lookupM (Key k) (KeyMap m) = lookupM (toRep k) m - 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 + emptyM = KeyMap 0 emptyM + singletonM (Key k) a = KeyMap (getSize a) (singletonM (toRep k) a) + getSimpleM KMAP(m) = getSimpleM m + sizeM = sz + lookupM (Key k) KMAP(m) = lookupM (toRep k) m + traverseM f KMAP(m) = keyMap <$> traverseM f m + fmapM f KMAP(m) = keyMap (fmapM f m) + mapMaybeM f KMAP(m) = keyMap (mapMaybeM f m) + mapEitherM f KMAP(m) = both keyMap keyMap (mapEitherM f) m + unionM f KMAP(m1) KMAP(m2) = keyMap (unionM f m1 m2) + isectM f KMAP(m1) KMAP(m2) = keyMap (isectM f m1 m2) + diffM f KMAP(m1) KMAP(m2) = keyMap (diffM f m1 m2) + isSubmapM (<=) KMAP(m1) KMAP(m2) = isSubmapM (<=) m1 m2 singleHoleM (Key k) = KeyHole (singleHoleM (toRep k)) - beforeM a (KeyHole hole) = KeyMap (beforeM a hole) - afterM a (KeyHole hole) = KeyMap (afterM a hole) - 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) + beforeM (KeyHole hole) = keyMap (beforeM hole) + beforeWithM a (KeyHole hole) = keyMap (beforeWithM a hole) + afterM (KeyHole hole) = keyMap (afterM hole) + afterWithM a (KeyHole hole) = keyMap (afterWithM a hole) + searchMC (Key k) KMAP(m) = mapSearch KeyHole (searchMC (toRep k) m) + indexM i KMAP(m) = onThird KeyHole (indexM i) m + extractHoleM KMAP(m) = fmap KeyHole <$> extractHoleM m + 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) + insertWithM f (Key k) a KMAP(m) = keyMap (insertWithM f (toRep k) a m) + fromListM f xs = keyMap (fromListM f [(toRep k, a) | (Key k, a) <- xs]) + fromAscListM f xs = keyMap (fromAscListM f [(toRep k, a) | (Key k, a) <- xs]) + fromDistAscListM xs = keyMap (fromDistAscListM [(toRep k, a) | (Key k, a) <- xs]) diff --git a/Data/TrieMap/Modifiers.hs b/Data/TrieMap/Modifiers.hs --- a/Data/TrieMap/Modifiers.hs +++ b/Data/TrieMap/Modifiers.hs @@ -7,6 +7,11 @@ newtype Rev k = Rev {getRev :: k} deriving (Eq) instance Ord k => Ord (Rev k) where compare (Rev a) (Rev b) = compare b a + Rev a < Rev b = b < a + Rev a <= Rev b = b <= a + (>) = flip (<) + (>=) = flip (<=) + instance Functor Ordered where fmap f (Ord a) = Ord (f a) @@ -16,6 +21,18 @@ newtype Key k = Key {getKey :: k} +instance (Repr k, Eq (Rep k)) => Eq (Key k) where + Key a == Key b = toRep a == toRep b + +instance (Repr k, Ord (Rep k)) => Ord (Key k) where + Key a `compare` Key b = toRep a `compare` toRep b + Key a < Key b = toRep a < toRep b + Key a <= Key b = toRep a <= toRep b + (>) = flip (<) + (>=) = flip (<=) + instance Repr k => Repr (Key k) where type Rep (Key k) = Rep k + type RepList (Key k) = RepList k toRep (Key k) = toRep k + toRepList ks = toRepList [k | Key k <- ks] diff --git a/Data/TrieMap/OrdMap.hs b/Data/TrieMap/OrdMap.hs --- a/Data/TrieMap/OrdMap.hs +++ b/Data/TrieMap/OrdMap.hs @@ -1,157 +1,185 @@ -{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, PatternGuards, MagicHash, CPP, TupleSections #-} - +{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, PatternGuards, MagicHash, CPP, TupleSections, NamedFieldPuns, FlexibleInstances #-} +{-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.OrdMap () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized import Data.TrieMap.Modifiers -import Control.Applicative -import Control.Monad hiding (join, fmap) - -import Prelude hiding (lookup, foldr, foldl, fmap) +import Control.Applicative (Applicative(..), (<$>)) +import Control.Monad hiding (join) -import GHC.Exts +import Data.Foldable +import Data.Monoid -#define DELTA 5# -#define RATIO 2# +import Prelude hiding (lookup, foldr, foldl, foldr1, foldl1, map) -type OrdMap k = TrieMap (Ordered k) +#define DELTA 5 +#define RATIO 2 data Path k a = Root - | LeftBin k a !(Path k a) !(OrdMap k a) - | RightBin k a !(OrdMap k a) !(Path k a) + | LeftBin k a !(Path k a) !(SNode k a) + | RightBin k a !(SNode k a) !(Path k a) -singletonMaybe :: Sized a => k -> Maybe a -> OrdMap k a -singletonMaybe k = maybe Tip (singleton k) +data Node k a = + Tip + | Bin k a !(SNode k a) !(SNode k a) +data SNode k a = SNode{sz :: !Int, count :: !Int, node :: Node k a} +#define TIP SNode{node=Tip} +#define BIN(args) SNode{node=Bin args} + +instance Sized a => Sized (Node k a) where + getSize# m = unbox $ case m of + Tip -> 0 + Bin _ a l r -> getSize a + getSize l + getSize r + +instance Sized (SNode k a) where + getSize# SNode{sz} = unbox sz + +nCount :: Node k a -> Int +nCount Tip = 0 +nCount (Bin _ _ l r) = 1 + count l + count r + +sNode :: Sized a => Node k a -> SNode k a +sNode !n = SNode (getSize n) (nCount n) n + +tip :: SNode k a +tip = SNode 0 0 Tip + -- | @'TrieMap' ('Ordered' k) a@ is based on "Data.Map". 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) + newtype TrieMap (Ordered k) a = OrdMap (SNode k a) data Hole (Ordered k) a = Empty k !(Path k a) - | Full k !(Path k a) !(OrdMap k a) !(OrdMap k a) - emptyM = Tip - singletonM (Ord k) = singleton k - lookupM (Ord k) = lookup k - 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 [(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 = hedgeUnion f (const LT) (const GT) m1 m2 - isectM = isect - diffM _ Tip _ = Tip - diffM _ m1 Tip = m1 - diffM f m1 m2 = hedgeDiff f (const LT) (const GT) m1 m2 + | Full k !(Path k a) !(SNode k a) !(SNode k a) + emptyM = OrdMap tip + singletonM (Ord k) a = OrdMap (singleton k a) + lookupM (Ord k) (OrdMap m) = lookup k m + getSimpleM (OrdMap m) = case m of + TIP -> Null + BIN(_ a TIP TIP) + -> Singleton a + _ -> NonSimple + sizeM (OrdMap m) = sz m + traverseM f (OrdMap m) = OrdMap <$> traverse f m + fmapM f (OrdMap m) = OrdMap (map f m) + mapMaybeM f (OrdMap m) = OrdMap (mapMaybe f m) + mapEitherM f (OrdMap m) = both OrdMap OrdMap (mapEither f) m + isSubmapM (<=) (OrdMap m1) (OrdMap m2) = isSubmap (<=) m1 m2 + fromAscListM f xs = OrdMap $ fromAscList f [(k, a) | (Ord k, a) <- xs] + fromDistAscListM xs = OrdMap $ fromDistinctAscList [(k, a) | (Ord k, a) <- xs] + unionM f (OrdMap m1) (OrdMap m2) = OrdMap $ hedgeUnion f (const LT) (const GT) m1 m2 + isectM f (OrdMap m1) (OrdMap m2) = OrdMap $ isect f m1 m2 + diffM f (OrdMap m1) (OrdMap m2) = OrdMap $ hedgeDiff f (const LT) (const GT) m1 m2 singleHoleM (Ord k) = Empty k Root - beforeM a (Empty k path) = before (singletonMaybe k a) path - beforeM a (Full k path l _) = before t path - where t = case a of - Nothing -> l - Just a -> insertMax k a l - afterM a (Empty k path) = after (singletonMaybe k a) path - afterM a (Full k path _ r) = after t path - where t = case a of - Nothing -> r - Just a -> insertMin k a r - searchM (Ord k) = search k Root - indexM i# = indexT Root i# where - indexT path i# (Bin _ kx x l r) - | i# <# sl# = indexT (LeftBin kx x path r) i# l - | i# <# sx# = (# i# -# sl#, x, Full kx path l r #) - | otherwise = indexT (RightBin kx x l path) (i# -# sx#) r - where !sl# = size# l - !sx# = getSize# x +# sl# + beforeM (Empty _ path) = OrdMap $ before tip path + beforeM (Full _ path l _) = OrdMap $ before l path + beforeWithM a (Empty k path) = OrdMap $ before (singleton k a) path + beforeWithM a (Full k path l _) = OrdMap $ before (insertMax k a l) path + afterM (Empty _ path) = OrdMap $ after tip path + afterM (Full _ path _ r) = OrdMap $ after r path + afterWithM a (Empty k path) = OrdMap $ after (singleton k a) path + afterWithM a (Full k path _ r) = OrdMap $ after (insertMin k a r) path + searchMC (Ord k) (OrdMap m) = search k m + indexM i (OrdMap m) = indexT Root i m where + indexT path i BIN(kx x l r) + | i < sl = indexT (LeftBin kx x path r) i l + | i < sx = (# i - sl, x, Full kx path l r #) + | otherwise = indexT (RightBin kx x l path) (i - sx) r + where !sl = getSize l + !sx = getSize x + sl indexT _ _ _ = indexFail () - extractHoleM = extractHole Root where - extractHole path (Bin _ kx x l r) = + extractHoleM (OrdMap m) = extractHole Root m 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 (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 + clearM (Empty _ path) = OrdMap $ rebuild tip path + clearM (Full _ path l r) = OrdMap $ rebuild (merge l r) path + assignM x (Empty k path) = OrdMap $ rebuild (singleton k x) path + assignM x (Full k path l r) = OrdMap $ rebuild (join k x l r) path + + unifierM (Ord k') (Ord k) a = case compare k' k of + EQ -> Nothing + LT -> Just $ Empty k' (LeftBin k a Root tip) + GT -> Just $ Empty k' (RightBin k a tip Root) -rebuild :: Sized a => OrdMap k a -> Path k a -> OrdMap k a +rebuild :: Sized a => SNode k a -> Path k a -> SNode k a rebuild t Root = t rebuild t (LeftBin kx x path r) = rebuild (balance kx x t r) path rebuild t (RightBin kx x l path) = rebuild (balance kx x l t) path -lookup :: Ord k => k -> OrdMap k a -> Maybe a -lookup k (Bin _ k' v l r) = case compare k k' of +lookup :: Ord k => k -> SNode k a -> Lookup a +lookup k = look where + look BIN(kx x l r) = case compare k kx of LT -> lookup k l - EQ -> Just v + EQ -> some x GT -> lookup k r -lookup _ _ = Nothing + look _ = none -singleton :: Sized a => k -> a -> OrdMap k a -singleton k a = Bin (getSize# a) k a Tip Tip +singleton :: Sized a => k -> a -> SNode k a +singleton k a = bin k a tip tip -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 +traverse :: (Applicative f, Sized b) => (a -> f b) -> SNode k a -> f (SNode k b) +traverse _ TIP = pure tip +traverse f BIN(k a l r) = balance k <$> f a <*> traverse f l <*> traverse f r -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 +instance Foldable (SNode k) where + foldMap _ TIP = mempty + foldMap f BIN(_ a l r) = foldMap f l `mappend` f a `mappend` foldMap f r -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 + foldr _ z TIP = z + foldr f z BIN(_ a l r) = foldr f (a `f` foldr f z r) l + foldl _ z TIP = z + foldl f z BIN(_ a l r) = foldl f (foldl f z l `f` a) r + + foldr1 _ TIP = foldr1Empty + foldr1 f BIN(_ a l TIP) = foldr f a l + foldr1 f BIN(_ a l r) = foldr f (a `f` foldr1 f r) l + + foldl1 _ TIP = foldl1Empty + foldl1 f BIN(_ a TIP r) = foldl f a r + foldl1 f BIN(_ a l r) = foldl f (foldl1 f l `f` a) r -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 +instance Foldable (TrieMap (Ordered k)) where + foldMap f (OrdMap m) = foldMap f m + foldr f z (OrdMap m) = foldr f z m + foldl f z (OrdMap m) = foldl f z m + foldl1 f (OrdMap m) = foldl1 f m + foldr1 f (OrdMap m) = foldr1 f m -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 +map :: (Ord k, Sized b) => (a -> b) -> SNode k a -> SNode k b +map f BIN(k a l r) = join k (f a) (map f l) (map f r) +map _ _ = tip +mapMaybe :: (Ord k, Sized b) => (a -> Maybe b) -> SNode k a -> SNode k b +mapMaybe f BIN(k a l r) = joinMaybe k (f a) (mapMaybe f l) (mapMaybe f r) +mapMaybe _ _ = tip + 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) = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #) + SNode k a -> (# SNode k b, SNode k c #) +mapEither f BIN(k a l r) = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #) where !(# aL, aR #) = f a; !(# lL, lR #) = mapEither f l; !(# rL, rR #) = mapEither f r -mapEither _ _ = (# Tip, Tip #) +mapEither _ _ = (# tip, tip #) -splitLookup :: (Ord k, Sized a) => k -> OrdMap k a -> (# OrdMap k a, Maybe a, OrdMap k a #) -splitLookup k m = case m of - Tip -> (# Tip, Nothing, Tip #) - Bin _ kx x l r -> case compare k kx of - LT -> let !(# lL, ans, lR #) = splitLookup k l in (# lL, ans, join kx x lR r #) - EQ -> (# l, Just x, r #) - GT -> let !(# rL, ans, rR #) = splitLookup k r in (# join kx x l rL, ans, rR #) +splitLookup :: (Ord k, Sized a) => k -> SNode k a -> (SNode k a -> Maybe a -> SNode k a -> r) -> r +splitLookup k t cont = search k t (split Nothing) (split . Just) where + split v (Empty _ path) = cont (before tip path) v (after tip path) + split v (Full _ path l r) = cont (before l path) v (after r path) -isSubmap :: (Ord k, Sized a, Sized b) => LEq a b -> LEq (OrdMap k a) (OrdMap k b) -isSubmap _ Tip _ = True -isSubmap _ _ Tip = False -isSubmap (<=) (Bin _ kx x l r) t = case found of - Nothing -> False - Just y -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt - where !(# lt, found, gt #) = splitLookup kx t +isSubmap :: (Ord k, Sized a, Sized b) => LEq a b -> LEq (SNode k a) (SNode k b) +isSubmap _ TIP _ = True +isSubmap _ _ TIP = False +isSubmap (<=) BIN(kx x l r) t = splitLookup kx t result + where result _ Nothing _ = False + result tl (Just y) tr = x <= y && isSubmap (<=) l tl && isSubmap (<=) r tr -fromAscList :: (Eq k, Sized a) => (a -> a -> a) -> [(k, a)] -> OrdMap k a +fromAscList :: (Eq k, Sized a) => (a -> a -> a) -> [(k, a)] -> SNode k a fromAscList f xs = fromDistinctAscList (combineEq xs) where combineEq (x:xs) = combineEq' x xs combineEq [] = [] @@ -161,12 +189,12 @@ | kz == kx = combineEq' (kx, f xx zz) xs | otherwise = (kz,zz):combineEq' x xs -fromDistinctAscList :: Sized a => [(k, a)] -> OrdMap k a +fromDistinctAscList :: Sized a => [(k, a)] -> SNode k a fromDistinctAscList xs = build const (length xs) xs where -- 1) use continutations so that we use heap space instead of stack space. -- 2) special case for n==5 to build bushier trees. - build c 0 xs' = c Tip xs' + build c 0 xs' = c tip xs' build c 5 xs' = case xs' of ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx) -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx @@ -183,12 +211,12 @@ hedgeUnion :: (Ord k, Sized a) => (a -> a -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) - -> OrdMap k a -> OrdMap k a -> OrdMap k a -hedgeUnion _ _ _ t1 Tip + -> SNode k a -> SNode k a -> SNode k a +hedgeUnion _ _ _ t1 TIP = t1 -hedgeUnion _ 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) -hedgeUnion f cmplo cmphi (Bin _ kx x l r) t2 +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 @@ -199,58 +227,54 @@ Nothing -> Just x Just (_,y) -> f x y -filterGt :: (Ord k, Sized a) => (k -> Ordering) -> OrdMap k a -> OrdMap k a -filterGt _ Tip = Tip -filterGt cmp (Bin _ kx x l r) +filterGt :: (Ord k, Sized a) => (k -> Ordering) -> SNode k a -> SNode k a +filterGt _ TIP = tip +filterGt cmp BIN(kx x l r) = case cmp kx of LT -> join kx x (filterGt cmp l) r GT -> filterGt cmp r EQ -> r -filterLt :: (Ord k, Sized a) => (k -> Ordering) -> OrdMap k a -> OrdMap k a -filterLt _ Tip = Tip -filterLt cmp (Bin _ kx x l r) +filterLt :: (Ord k, Sized a) => (k -> Ordering) -> SNode k a -> SNode k a +filterLt _ TIP = tip +filterLt cmp BIN(kx x l r) = case cmp kx of LT -> filterLt cmp l GT -> join kx x l (filterLt cmp r) EQ -> l -trim :: (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k a -trim _ _ Tip = Tip -trim cmplo cmphi t@(Bin _ kx _ l r) - = case cmplo kx of - LT -> case cmphi kx of - GT -> t - _ -> trim cmplo cmphi l - _ -> trim cmplo cmphi r +trim :: (k -> Ordering) -> (k -> Ordering) -> SNode k a -> SNode k a +trim cmplo cmphi = trimmer where + trimmer TIP = tip + trimmer t@BIN(kx _ l r) = case (cmplo kx, cmphi kx) of + (LT, GT) -> t + (LT, _) -> trimmer l + _ -> trimmer r -trimLookupLo :: Ord k => k -> (k -> Ordering) -> OrdMap k a -> (Maybe (k,a), OrdMap k a) -trimLookupLo _ _ Tip = (Nothing,Tip) -trimLookupLo lo cmphi t@(Bin _ kx x l r) +trimLookupLo :: Ord k => k -> (k -> Ordering) -> SNode k a -> (Maybe (k,a), SNode k a) +trimLookupLo _ _ TIP = (Nothing,tip) +trimLookupLo lo cmphi t@BIN(kx x l r) = case compare lo kx of LT -> case cmphi kx of - GT -> ((lo,) <$> lookup lo t, t) + GT -> (option (lookup lo t) Nothing (\ a -> Just (lo, a)), t) _ -> trimLookupLo lo cmphi l GT -> trimLookupLo lo cmphi r EQ -> (Just (kx,x),trim (compare lo) cmphi r) -isect :: (Ord k, Sized a, Sized b, Sized c) => (a -> b -> Maybe c) -> OrdMap k a -> OrdMap k b -> OrdMap k c -isect f t1@Bin{} (Bin _ k2 x2 l2 r2) - = 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 +isect :: (Ord k, Sized a, Sized b, Sized c) => (a -> b -> Maybe c) -> SNode k a -> SNode k b -> SNode k c +isect f t1@BIN(_ _ _ _) BIN(k2 x2 l2 r2) = splitLookup k2 t1 result where + result tl found tr = joinMaybe k2 (found >>= \ x1' -> f x1' x2) (isect f tl l2) (isect f tr r2) +isect _ _ _ = tip hedgeDiff :: (Ord k, Sized a) => (a -> b -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) - -> OrdMap k a -> OrdMap k b -> OrdMap k a -hedgeDiff _ _ _ Tip _ - = Tip -hedgeDiff _ cmplo cmphi (Bin _ kx x l r) Tip + -> SNode k a -> SNode k b -> SNode k a +hedgeDiff _ _ _ TIP _ + = tip +hedgeDiff _ cmplo cmphi BIN(kx x l r) TIP = join kx x (filterGt cmplo l) (filterLt cmphi r) -hedgeDiff 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) -> @@ -264,128 +288,122 @@ 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 :: (Ord k, Sized a) => k -> Maybe a -> SNode k a -> SNode k a -> SNode k a joinMaybe kx = maybe merge (join kx) -join :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a -join kx x Tip r = insertMin kx x r -join kx x l Tip = insertMax kx x l -join kx x l@(Bin sL# ky y ly ry) r@(Bin sR# kz z lz rz) - | DELTA *# sL# <=# sR# = balance kz z (join kx x l lz) rz - | DELTA *# sR# <=# sL# = balance ky y ly (join kx x ry r) - | otherwise = bin kx x l r +join :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a +join kx x TIP r = insertMin kx x r +join kx x l TIP = insertMax kx x l +join kx x l@(SNode _ sL (Bin ky y ly ry)) r@(SNode _ sR (Bin kz z lz rz)) + | DELTA * sL <= sR = balance kz z (join kx x l lz) rz + | DELTA * sR <= sL = balance ky y ly (join kx x ry r) + | otherwise = bin kx x l r -- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -insertMax kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balance ky y l (insertMax kx x r) +insertMax,insertMin :: Sized a => k -> a -> SNode k a -> SNode k a +insertMax kx x = insMax where + insMax TIP = singleton kx x + insMax BIN(ky y l r) + = balance ky y l (insMax r) -insertMin kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balance ky y (insertMin kx x l) r +insertMin kx x = insMin where + insMin TIP = singleton kx x + insMin BIN(ky y l r) + = balance ky y (insMin l) r {-------------------------------------------------------------------- [merge l r]: merges two trees. --------------------------------------------------------------------} -merge :: Sized a => OrdMap k a -> OrdMap k a -> OrdMap k a -merge Tip r = r -merge l Tip = l -merge l@(Bin sL# kx x lx rx) r@(Bin sR# ky y ly ry) - | DELTA *# sL# <=# sR# = balance ky y (merge l ly) ry - | DELTA *# sR# <=# sL# = balance kx x lx (merge rx r) - | otherwise = glue l r +merge :: Sized a => SNode k a -> SNode k a -> SNode k a +merge TIP r = r +merge l TIP = l +merge l@(SNode _ sL (Bin kx x lx rx)) r@(SNode _ sR (Bin ky y ly ry)) + | DELTA * sL <= sR = balance ky y (merge l ly) ry + | DELTA * sR <= sL = balance kx x lx (merge rx r) + | otherwise = glue l r {-------------------------------------------------------------------- [glue l r]: glues two trees together. Assumes that [l] and [r] are already balanced with respect to each other. --------------------------------------------------------------------} -glue :: Sized a => OrdMap k a -> OrdMap k a -> OrdMap k a -glue Tip r = r -glue l Tip = l +glue :: Sized a => SNode k a -> SNode k a -> SNode k a +glue TIP r = r +glue l TIP = l glue 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' + | count l > count r = let !(# f, l' #) = deleteFindMax balance l in f l' r + | otherwise = let !(# f, r' #) = deleteFindMin balance r in f l r' -deleteFindMin :: Sized a => (k -> a -> (# x, Maybe a #)) -> OrdMap k a -> (# x, OrdMap k a #) +deleteFindMin :: Sized a => (k -> a -> x) -> SNode k a -> (# x, SNode k a #) deleteFindMin f t = case t of - 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 #) + BIN(k x TIP r) -> (# f k x, r #) + 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 :: Sized a => (k -> a -> x) -> SNode k a -> (# x, SNode k a #) deleteFindMax f t = case t of - 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# -size# (Bin sz _ _ _ _) = sz + BIN(k x l TIP) -> (# f k x, l #) + 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 #) -balance :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a +balance :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a balance k x l r - | sR# >=# (DELTA *# sL#) = rotateL k x l r - | sL# >=# (DELTA *# sR#) = rotateR k x l r - | otherwise = Bin sX# k x l r + | sR >= (DELTA * sL) = rotateL k x l r + | sL >= (DELTA * sR) = rotateR k x l r + | otherwise = bin k x l r where - !sL# = size# l - !sR# = size# r - !sX# = sL# +# sR# +# getSize# x + !sL = count l + !sR = count r -- rotate -rotateL :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a -rotateL k x l r@(Bin _ _ _ ly ry) - | sL# <# (RATIO *# sR#) = singleL k x l r - | otherwise = doubleL k x l r - where !sL# = size# ly - !sR# = size# ry -rotateL _ _ _ Tip = error "rotateL Tip" +rotateL :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a +rotateL k x l r@BIN(_ _ ly ry) + | sL < (RATIO * sR) = singleL k x l r + | otherwise = doubleL k x l r + where !sL = count ly + !sR = count ry +rotateL k x l TIP = insertMax k x l -rotateR :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a -rotateR k x l@(Bin _ _ _ ly ry) r - | sR# <# (RATIO *# sL#) = singleR k x l r - | otherwise = doubleR k x l r - where !sL# = size# ly - !sR# = size# ry -rotateR _ _ _ _ = error "rotateR Tip" +rotateR :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a +rotateR k x l@BIN(_ _ ly ry) r + | sR < (RATIO * sL) = singleR k x l r + | otherwise = doubleR k x l r + where !sL = count ly + !sR = count ry +rotateR k x TIP r = insertMin k x r -- basic rotations -singleL, singleR :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a -singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 -singleL k1 x1 t1 Tip = bin k1 x1 t1 Tip -singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) -singleR k1 x1 Tip t2 = bin k1 x1 Tip t2 +singleL, singleR :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a +singleL k1 x1 t1 BIN(k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 +singleL k1 x1 t1 TIP = bin k1 x1 t1 tip +singleR k1 x1 BIN(k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) +singleR k1 x1 TIP t2 = bin k1 x1 tip t2 -doubleL, doubleR :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a -doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4) +doubleL, doubleR :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a +doubleL k1 x1 t1 BIN(k2 x2 BIN(k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4) doubleL k1 x1 t1 t2 = singleL k1 x1 t1 t2 -doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4) +doubleR k1 x1 BIN(k2 x2 t1 BIN(k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4) doubleR k1 x1 t1 t2 = singleR k1 x1 t1 t2 -bin :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a +bin :: Sized a => k -> a -> SNode k a -> SNode k a -> SNode k a bin k x l r - = Bin (size# l +# size# r +# getSize# x) k x l r + = sNode (Bin k x l r) -before :: Sized a => OrdMap k a -> Path k a -> OrdMap k a +before :: Sized a => SNode k a -> Path k a -> SNode k a before t (LeftBin _ _ path _) = before t path before t (RightBin k a l path) = before (join k a l t) path before t _ = t -after :: Sized a => OrdMap k a -> Path k a -> OrdMap k a +after :: Sized a => SNode k a -> Path k a -> SNode k a after t (LeftBin k a path r) = after (join k a t r) path after t (RightBin _ _ _ path) = after t path after t _ = t -search :: Ord k => k -> Path k a -> OrdMap k a -> (# Maybe a, Hole (Ordered k) a #) -search k path Tip = (# Nothing, Empty k path #) -search k path (Bin _ kx x l r) = case compare k kx of - LT -> search k (LeftBin kx x path r) l - EQ -> (# Just x, Full k path l r #) - GT -> search k (RightBin kx x l path) r +search :: Ord k => k -> SNode k a -> SearchCont (Hole (Ordered k) a) a r +search k t f g = searcher Root t where + searcher path TIP = f (Empty k path) + searcher path BIN(kx x l r) = case compare k kx of + LT -> searcher (LeftBin kx x path r) l + EQ -> g x (Full k path l r) + GT -> searcher (RightBin kx x l path) r diff --git a/Data/TrieMap/ProdMap.hs b/Data/TrieMap/ProdMap.hs --- a/Data/TrieMap/ProdMap.hs +++ b/Data/TrieMap/ProdMap.hs @@ -1,23 +1,26 @@ -{-# LANGUAGE UnboxedTuples, TupleSections, PatternGuards, TypeFamilies #-} +{-# LANGUAGE UnboxedTuples, TupleSections, PatternGuards, TypeFamilies, FlexibleInstances #-} module Data.TrieMap.ProdMap () where import Data.TrieMap.Sized import Data.TrieMap.TrieKey -import Control.Applicative - +import Control.Monad +import Data.Functor import Data.Foldable hiding (foldlM, foldrM) -import Data.Monoid import Data.Sequence ((|>)) import qualified Data.Sequence as Seq +import Prelude hiding (foldl, foldl1, foldr, foldr1) + +instance (TrieKey k1, TrieKey k2) => Foldable (TrieMap (k1, k2)) where + foldMap f (PMap m) = foldMap (foldMap f) m + foldr f z (PMap m) = foldr (flip $ foldr f) z m + foldl f z (PMap m) = foldl (foldl f) z m + -- | @'TrieMap' (k1, k2) a@ is implemented as a @'TrieMap' k1 ('TrieMap' k2 a)@. 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) @@ -27,8 +30,6 @@ sizeM (PMap m) = sizeM m lookupM (k1, k2) (PMap m) = lookupM k1 m >>= lookupM k2 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 @@ -36,17 +37,21 @@ 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) + insertWithM f (k1, k2) a (PMap m) = PMap (insertWithM f' k1 (singletonM k2 a) m) where + f' = insertWithM f k2 a fromAscListM f xs = PMap (fromDistAscListM [(a, fromAscListM f ys) | (a, Elem ys) <- breakFst xs]) fromDistAscListM xs = PMap (fromDistAscListM [(a, fromDistAscListM ys) | (a, Elem ys) <- breakFst xs]) singleHoleM (k1, k2) = PHole (singleHoleM k1) (singleHoleM k2) - 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 + beforeM (PHole hole1 hole2) = PMap (beforeMM (gNull beforeM hole2) hole1) + beforeWithM a (PHole hole1 hole2) = PMap (beforeWithM (beforeWithM a hole2) hole1) + afterM (PHole hole1 hole2) = PMap (afterMM (gNull afterM hole2) hole1) + afterWithM a (PHole hole1 hole2) = PMap (afterWithM (afterWithM a hole2) hole1) + searchMC (k1, k2) (PMap m) f g = searchMC k1 m f' g' where + f' hole1 = f (PHole hole1 (singleHoleM k2)) + g' m' hole1 = mapSearch (PHole hole1) (searchMC k2 m') f g indexM i (PMap m) = onThird (PHole hole1) (indexM i') m' where !(# i', m', hole1 #) = indexM i m extractHoleM (PMap m) = do @@ -54,16 +59,20 @@ (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) + clearM (PHole hole1 hole2) = PMap (fillHoleM (clearM' hole2) hole1) + assignM a (PHole hole1 hole2) = PMap (assignM (assignM a hole2) hole1) + + unifierM (k1', k2') (k1, k2) a = case unifierM k1' k1 (singletonM k2 a) of + Just hole1 -> Just (PHole hole1 (singleHoleM k2')) + Nothing -> PHole (singleHoleM k1) <$> unifierM k2' k2 a -breakFst :: TrieKey k1 => [((k1, k2), a)] -> [(k1, Elem [(k2, a)])] +gNull :: TrieKey k => (x -> TrieMap k a) -> x -> Maybe (TrieMap k a) +gNull = (guardNullM .) + +breakFst :: Eq k1 => [((k1, k2), a)] -> [(k1, Elem [(k2, a)])] breakFst [] = [] breakFst (((a, b),v):xs) = breakFst' a (Seq.singleton (b, v)) xs where breakFst' a vs (((a', b'), v'):xs) - | a =? a' = breakFst' a' (vs |> (b', v')) xs + | 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 @@ -5,126 +5,129 @@ import Data.TrieMap.TrieKey import Data.TrieMap.Sized -import Control.Applicative +import Data.Functor +import Data.Foldable (Foldable(..)) 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 Data.Word -import Data.TrieMap.RadixTrie.Slice import Data.TrieMap.RadixTrie.Edge +import Data.TrieMap.RadixTrie.Label import Prelude hiding (length, and, zip, zipWith, foldr, foldl) +instance TrieKey k => Foldable (TrieMap (Vector k)) where + foldMap f (Radix m) = foldMap (foldMap f) m + foldr f z (Radix m) = foldl (foldr f) z m + foldl f z (Radix m) = foldl (foldl f) z m + -- | @'TrieMap' ('Vector' k) a@ is a traditional radix trie. 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) - 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)) + singletonM ks a = Radix (Just (singletonEdge ks a)) getSimpleM (Radix Nothing) = Null getSimpleM (Radix (Just e)) = getSimpleEdge e - sizeM (Radix m) = getSize# m - lookupM ks (Radix m) = m >>= lookupEdge ks + sizeM (Radix m) = getSize m + lookupM ks (Radix m) = liftMaybe m >>= lookupEdge ks fmapM f (Radix m) = Radix (mapEdge f <$> m) mapMaybeM f (Radix m) = Radix (m >>= mapMaybeEdge f) mapEitherM f (Radix e) = both Radix Radix (mapEitherMaybe (mapEitherEdge f)) e traverseM f (Radix m) = Radix <$> traverse (traverseEdge f) m - foldrM f (Radix m) z = foldr (foldrEdge f) z m - foldlM f (Radix m) z = foldl (foldlEdge f) z m - unionM f (Radix m1) (Radix m2) = Radix (unionMaybe (unionEdge f) m1 m2) isectM f (Radix m1) (Radix m2) = Radix (isectMaybe (isectEdge f) m1 m2) diffM f (Radix m1) (Radix m2) = Radix (diffMaybe (diffEdge f) m1 m2) isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubEdge (<=)) m1 m2 - 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 () + singleHoleM ks = Hole (singleLoc ks) + {-# INLINE searchMC #-} + searchMC ks (Radix (Just e)) = mapSearch Hole (searchEdgeC ks e) + searchMC ks _ = \ f _ -> f (singleHoleM ks) + indexM i (Radix (Just e)) = onThird Hole (indexEdge i e) root + indexM _ _ = indexFail () - assignM a (Hole loc) = Radix (fillHoleEdge a loc) + clearM (Hole loc) = Radix (clearEdge loc) + {-# INLINE assignM #-} + assignM a (Hole loc) = Radix (Just (assignEdge a loc)) - extractHoleM (Radix (Just e)) = do - (a, loc) <- extractEdgeLoc e Root - return (a, Hole loc) + extractHoleM (Radix (Just e)) = fmap Hole <$> extractEdgeLoc e root extractHoleM _ = mzero - beforeM a (Hole loc) = Radix (beforeEdge a loc) - afterM a (Hole loc) = Radix (afterEdge a loc) + beforeM (Hole loc) = Radix (beforeEdge Nothing loc) + beforeWithM a (Hole loc) = Radix (beforeEdge (Just a) loc) + afterM (Hole loc) = Radix (afterEdge Nothing loc) + afterWithM a (Hole loc) = Radix (afterEdge (Just a) loc) - unifyM ks1 a1 ks2 a2 = either (Left . Hole) (Right . Radix . Just) (unifyEdge (v2S ks1) a1 (v2S ks2) a2) - + insertWithM f ks v (Radix e) = Radix (Just (maybe (singletonEdge ks v) (insertEdge f ks v) e)) + fromListM _ [] = emptyM + fromListM f ((k, a):xs) = Radix (Just (roll (singletonEdge k a) xs)) where + roll !e [] = e + roll !e ((ks, a):xs) = roll (insertEdge (f a) ks a e) xs + type WordVec = S.Vector Word -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) +instance Foldable (TrieMap (S.Vector Word)) where + foldMap f (WRadix m) = foldMap (foldMap f) m + foldr f z (WRadix m) = foldl (foldr f) z m + foldl f z (WRadix m) = foldl (foldl f) z m -- | @'TrieMap' ('S.Vector' Word) a@ is a traditional radix trie specialized for word arrays. instance TrieKey (S.Vector Word) where - 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) - 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)) + singletonM ks a = WRadix (Just (singletonEdge ks a)) getSimpleM (WRadix Nothing) = Null getSimpleM (WRadix (Just e)) = getSimpleEdge e - sizeM (WRadix m) = getSize# m - lookupM ks (WRadix m) = m >>= lookupEdge ks + sizeM (WRadix m) = getSize m + lookupM ks (WRadix m) = liftMaybe m >>= lookupEdge ks fmapM f (WRadix m) = WRadix (mapEdge f <$> m) mapMaybeM f (WRadix m) = WRadix (m >>= mapMaybeEdge f) mapEitherM f (WRadix e) = both WRadix WRadix (mapEitherMaybe (mapEitherEdge f)) e traverseM f (WRadix m) = WRadix <$> traverse (traverseEdge f) m - foldrM f (WRadix m) z = foldr (foldrEdge f) z m - foldlM f (WRadix m) z = foldl (foldlEdge f) z m - unionM f (WRadix m1) (WRadix m2) = WRadix (unionMaybe (unionEdge f) m1 m2) isectM f (WRadix m1) (WRadix m2) = WRadix (isectMaybe (isectEdge f) m1 m2) diffM f (WRadix m1) (WRadix m2) = WRadix (diffMaybe (diffEdge f) m1 m2) - + isSubmapM (<=) (WRadix m1) (WRadix m2) = subMaybe (isSubEdge (<=)) m1 m2 - 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 #) + singleHoleM ks = WHole (singleLoc ks) + {-# INLINE searchMC #-} + searchMC ks (WRadix (Just e)) f g = searchEdgeC ks e f' g' where + f' loc = f (WHole loc) + g' a loc = g a (WHole loc) + searchMC ks _ f _ = f (singleHoleM ks) + indexM i (WRadix (Just e)) = onThird WHole (indexEdge i e) root indexM _ (WRadix Nothing) = indexFail () - assignM a (WHole loc) = WRadix (fillHoleEdge a loc) - + clearM (WHole loc) = WRadix (clearEdge loc) + {-# INLINE assignM #-} + assignM a (WHole loc) = WRadix (Just (assignEdge a loc)) + extractHoleM (WRadix (Just e)) = do - (a, loc) <- extractEdgeLoc e Root + (a, loc) <- extractEdgeLoc e root return (a, WHole loc) extractHoleM _ = mzero - beforeM a (WHole loc) = WRadix (beforeEdge a loc) - afterM a (WHole loc) = WRadix (afterEdge a loc) + beforeM (WHole loc) = WRadix (beforeEdge Nothing loc) + beforeWithM a (WHole loc) = WRadix (beforeEdge (Just a) loc) + afterM (WHole loc) = WRadix (afterEdge Nothing loc) + afterWithM a (WHole loc) = WRadix (afterEdge (Just a) loc) - unifyM ks1 a1 ks2 a2 = either (Left . WHole) (Right . WRadix . Just) (unifyEdge (v2S ks1) a1 (v2S ks2) a2) + insertWithM f ks v (WRadix e) = WRadix (Just (maybe (singletonEdge ks v) (insertEdge f ks v) e)) + {-# INLINE fromListM #-} + fromListM _ [] = emptyM + fromListM f ((k, a):xs) = WRadix (Just (roll (singletonEdge k a) xs)) where + roll !e [] = e + roll !e ((ks, a):xs) = roll (insertEdge (f a) ks a e) xs diff --git a/Data/TrieMap/RadixTrie/Edge.hs b/Data/TrieMap/RadixTrie/Edge.hs --- a/Data/TrieMap/RadixTrie/Edge.hs +++ b/Data/TrieMap/RadixTrie/Edge.hs @@ -1,269 +1,265 @@ -{-# LANGUAGE MagicHash, BangPatterns, UnboxedTuples, PatternGuards, CPP #-} +{-# LANGUAGE MagicHash, BangPatterns, UnboxedTuples, PatternGuards, CPP, ViewPatterns #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.RadixTrie.Edge where import Data.TrieMap.Sized import Data.TrieMap.TrieKey +import Data.TrieMap.WordMap () +import Data.TrieMap.RadixTrie.Label import Data.TrieMap.RadixTrie.Slice -import Data.TrieMap.IntMap () -import Data.TrieMap.Applicative () import Control.Applicative import Control.Monad + +import Data.Foldable +import Data.Monoid 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 Data.Vector.Generic (length) +import qualified Data.Vector (Vector) +import qualified Data.Vector.Storable (Vector) 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 +#define EDGE(args) (eView -> Edge args) +#define LOC(args) !(locView -> Loc args) -{-# 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 +{-# SPECIALIZE lookupEdge :: + TrieKey k => V() -> V(Edge) a -> Lookup a, + U() -> U(Edge) a -> Lookup a #-} +lookupEdge :: (Eq k, Label v k) => v k -> Edge v k a -> Lookup a lookupEdge = lookupE where - lookupE !ks (Edge _ ls v ts) = if kLen < lLen then Nothing else matchSliceV matcher matches ks ls where + lookupE !ks !EDGE(_ ls v ts) = if kLen < lLen then none else matchSlice matcher matches ks ls where !kLen = length ks - !lLen = len ls + !lLen = length ls matcher k l z - | k =? l = z - | otherwise = Nothing + | k == l = z + | otherwise = none matches _ _ - | kLen == lLen = v - | otherwise = do e' <- lookupM (ks `unsafeIndex` lLen) ts - lookupE (unsafeDrop (lLen + 1) ks) e' + | kLen == lLen = liftMaybe v + | (_, k, ks') <- splitSlice lLen ks + = lookupM k ts >>= lookupE ks' -{-# 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) +{-# INLINE searchEdgeC #-} +searchEdgeC :: (Eq k, Label v k) => v k -> Edge v k a -> (EdgeLoc v k a -> r) -> (a -> EdgeLoc v k a -> r) -> r +searchEdgeC ks0 e f g = searchE ks0 e root where + searchE !ks !e@EDGE(_ !ls !v ts) path = matcher 0 where + !kLen = length ks + !lLen = length ls + !len = min kLen lLen + {-# INLINE kk #-} + kk = ks !$ lLen + matcher !i + | i < len = let k = ks !$ i; l = ls !$ i in case unifierM k l (dropEdge (i+1) e) of + Nothing -> matcher (i+1) + Just tHole -> f (loc (dropSlice (i+1) ks) emptyM (deep path (takeSlice i ls) Nothing tHole)) + matcher _ + | kLen < lLen + = let lPre = takeSlice kLen ls; l = ls !$ kLen; e' = dropEdge (kLen + 1) e in + f (loc lPre (singletonM l e') path) + | kLen == lLen + = maybe f g v (loc ls ts path) + | otherwise = let + ks' = dropSlice (lLen + 1) ks + f' tHole = f (loc ks' emptyM (deep path ls v tHole)) + g' e' tHole = searchE ks' e' (deep path ls v tHole) + in searchMC kk ts f' g' -{-# 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 +{-# SPECIALIZE mapEdge :: + (TrieKey k, Sized b) => (a -> b) -> V(Edge) a -> V(Edge) b, + Sized b => (a -> b) -> U(Edge) a -> U(Edge) b #-} +mapEdge :: (Label v k, Sized b) => (a -> b) -> Edge v k a -> Edge v k b mapEdge f = mapE where - mapE (Edge _ ks v ts) = edge ks (f <$> v) (fmapM mapE ts) + 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 +{-# SPECIALIZE mapMaybeEdge :: + (TrieKey k, Sized b) => (a -> Maybe b) -> V(Edge) a -> V(MEdge) b, + Sized b => (a -> Maybe b) -> U(Edge) a -> U(MEdge) b #-} +mapMaybeEdge :: (Label v k, Sized b) => (a -> Maybe b) -> Edge v k a -> MEdge v k b mapMaybeEdge f = mapMaybeE where - mapMaybeE (Edge _ ks v ts) = compact (edge ks (v >>= f) (mapMaybeM mapMaybeE ts)) + mapMaybeE EDGE(_ ks v ts) = cEdge 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) => +{-# SPECIALIZE mapEitherEdge :: + (TrieKey k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> V(Edge) a -> (# V(MEdge) b, V(MEdge) c #), + (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> U(Edge) a -> (# U(MEdge) b, U(MEdge) c #) #-} +mapEitherEdge :: (Label v k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> Edge v k a -> (# MEdge v k b, MEdge v k c #) mapEitherEdge f = mapEitherE where - mapEitherE (Edge _ ks v ts) = (# compact (edge ks vL tsL), compact (edge ks vR tsR) #) + mapEitherE !EDGE(_ ks v ts) = (# cEdge ks vL tsL, cEdge ks vR tsR #) where !(# vL, vR #) = mapEitherMaybe f v !(# tsL, tsR #) = mapEitherM mapEitherE ts -{-# SPECIALIZE traverseEdge :: (Applicative f, Sized b) => - (a -> f b) -> U(Edge) a -> f (U(Edge) b) #-} -traverseEdge :: (TrieKey k, Applicative f, Sized b) => +{-# SPECIALIZE traverseEdge :: + (TrieKey k, Applicative f, Sized b) => (a -> f b) -> V(Edge) a -> f (V(Edge) b), + (Applicative f, Sized b) => (a -> f b) -> U(Edge) a -> f (U(Edge) b) #-} +traverseEdge :: (Label v k, Applicative f, Sized b) => (a -> f b) -> Edge v k a -> f (Edge v k b) traverseEdge f = traverseE where - traverseE (Edge _ ks v ts) = edge ks <$> traverse f v <*> traverseM traverseE ts + traverseE e = case eView e of + Edge _ ks Nothing ts -> edge ks Nothing <$> traverseM traverseE ts + Edge _ ks (Just v) ts -> edge ks . Just <$> f v <*> traverseM traverseE ts -{-# 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 +instance Label v k => Foldable (EView v k) where + {-# INLINE foldr #-} + {-# INLINE foldl #-} + {-# INLINE foldMap #-} + foldMap f (Edge _ _ Nothing ts) = foldMap (foldMap f) ts + foldMap f (Edge _ _ (Just v) ts) = f v `mappend` foldMap (foldMap f) ts + foldr f z (Edge _ _ v ts) = foldr f (foldr (flip $ foldr f) z ts) v + foldl f z (Edge _ _ v ts) = foldl (foldl f) (foldl f z v) ts -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) +instance Label v k => Foldable (Edge v k) where + {-# SPECIALIZE instance TrieKey k => Foldable (V(Edge)) #-} + {-# SPECIALIZE instance Foldable (U(Edge)) #-} + foldMap f e = foldMap f (eView e) + foldr f z e = foldr f z (eView e) + foldl f z e = foldl f z (eView e) -{-# 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 +{-# INLINE assignEdge #-} +assignEdge :: (Label v k, Sized a) => a -> EdgeLoc v k a -> Edge v k a +assignEdge v LOC(ks ts path) = assign (edge ks (Just v) ts) 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 assign :: + (TrieKey k, Sized a) => V(Edge) a -> V(Path) a -> V(Edge) a, + Sized a => U(Edge) a -> U(Path) a -> U(Edge) a #-} +assign :: (Label v k, Sized a) => Edge v k a -> Path v k a -> Edge v k a +assign !e path = case pView path of + Root -> e + Deep path ks v tHole + -> assign (edge ks v (assignM e tHole)) path -{-# 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) => +{-# SPECIALIZE clearEdge :: + (TrieKey k, Sized a) => V(EdgeLoc) a -> V(MEdge) a, + Sized a => U(EdgeLoc) a -> U(MEdge) a #-} +clearEdge :: (Label v k, Sized a) => EdgeLoc v k a -> MEdge v k a +clearEdge LOC(ks ts path) = rebuild (cEdge ks Nothing ts) path where + rebuild !e path = case pView path of + Root -> e + Deep path ks v tHole + -> rebuild (cEdge ks v (fillHoleM e tHole)) path + +{-# SPECIALIZE unionEdge :: + (TrieKey k, Sized a) => (a -> a -> Maybe a) -> V(Edge) a -> V(Edge) a -> V(MEdge) a, + Sized a => (a -> a -> Maybe a) -> U(Edge) a -> U(Edge) a -> U(MEdge) a #-} +unionEdge :: (Label v k, Sized a) => (a -> a -> Maybe a) -> Edge v k a -> Edge v k a -> MEdge v k a unionEdge f = unionE where - eK@(Edge _ ks0 vK tsK) `unionE` eL@(Edge _ ls0 vL tsL) = iMatchSlice matcher matches ks0 ls0 where + unionE !eK@EDGE(_ ks0 vK tsK) !eL@EDGE(_ ls0 vL tsL) = iMatchSlice matcher matches ks0 ls0 where matcher i k l z = case unifyM k eK' l eL' of - Left{} -> z - Right ts -> Just (edge (takeSlice i ks0) Nothing ts) + Nothing -> z + Just 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 + EQ -> cEdge ks0 (unionMaybe f vK vL) $ unionM unionE tsK tsL + LT -> searchMC l tsK nomatch match where + eL' = dropEdge (kLen + 1) eL; l = ls0 !$ kLen + nomatch holeKT = cEdge ks0 vK $ assignM eL' holeKT + match eK' holeKT = cEdge ks0 vK $ fillHoleM (eK' `unionE` eL') holeKT + GT -> searchMC k tsL nomatch match where + eK' = dropEdge (lLen + 1) eK; k = ks0 !$ lLen + nomatch holeLT = cEdge ls0 vL $ assignM eK' holeLT + match eL' holeLT = cEdge ls0 vL $ fillHoleM (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) => +{-# SPECIALIZE isectEdge :: + (TrieKey k, Sized c) => (a -> b -> Maybe c) -> V(Edge) a -> V(Edge) b -> V(MEdge) c, + Sized c => (a -> b -> Maybe c) -> U(Edge) a -> U(Edge) b -> U(MEdge) c #-} +isectEdge :: (Eq k, Label 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 + isectE !eK@EDGE(_ ks0 vK tsK) !eL@EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where + matcher k l z = guard (k == l) >> z matches kLen lLen = case compare kLen lLen of EQ -> compact $ edge ks0 (isectMaybe f vK vL) $ isectM isectE tsK tsL LT -> let l = ls0 !$ kLen in do - eK' <- lookupM l tsK + eK' <- toMaybe $ lookupM l tsK let eL' = dropEdge (kLen + 1) eL unDropEdge (kLen + 1) <$> eK' `isectE` eL' GT -> let k = ks0 !$ lLen in do - eL' <- lookupM k tsL + eL' <- toMaybe $ 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) => +{-# SPECIALIZE diffEdge :: + (TrieKey k, Sized a) => (a -> b -> Maybe a) -> V(Edge) a -> V(Edge) b -> V(MEdge) a, + Sized a => (a -> b -> Maybe a) -> U(Edge) a -> U(Edge) b -> U(MEdge) a #-} +diffEdge :: (Eq k, Label v k, Sized a) => (a -> b -> Maybe a) -> Edge v k a -> Edge v k b -> MEdge v k a diffEdge f = diffE where - eK@(Edge _ ks0 vK tsK) `diffE` eL@(Edge _ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where + diffE !eK@EDGE(_ ks0 vK tsK) !eL@EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where matcher k l z - | k =? l = z + | 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') + EQ -> cEdge ks0 (diffMaybe f vK vL) $ diffM diffE tsK tsL + LT -> searchMC l tsK nomatch match where + l = ls0 !$ kLen; eL' = dropEdge (kLen + 1) eL + nomatch _ = Just eK + match eK' holeKT = cEdge ks0 vK $ fillHoleM (eK' `diffE` eL') holeKT + GT -> let k = ks0 !$ lLen; eK' = dropEdge (lLen + 1) eK in + option (lookupM k tsL) (Just eK) (\ eL' -> fmap (unDropEdge (lLen + 1)) (eK' `diffE` eL')) -{-# 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) +{-# SPECIALIZE isSubEdge :: + TrieKey k => LEq a b -> LEq (V(Edge) a) (V(Edge) b), + LEq a b -> LEq (U(Edge) a) (U(Edge) b) #-} +isSubEdge :: (Eq k, Label v k) => LEq a b -> LEq (Edge v k a) (Edge v k b) isSubEdge (<=) = isSubE where - eK@(Edge _ ks0 vK tsK) `isSubE` (Edge _ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where - matcher k l z = k =? l && z + isSubE !eK@EDGE(_ ks0 vK tsK) !EDGE(_ ls0 vL tsL) = matchSlice matcher matches ks0 ls0 where + matcher k l z = k == l && z matches kLen lLen = case compare kLen lLen of LT -> False EQ -> subMaybe (<=) vK vL && isSubmapM isSubE tsK tsL - GT -> let k = ks0 !$ lLen in case lookupM k tsL of - Nothing -> False - Just eL' -> isSubE (dropEdge (lLen + 1) eK) eL' + GT -> let k = ks0 !$ lLen in option (lookupM k tsL) False (isSubE (dropEdge (lLen + 1) eK)) -{-# 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 beforeEdge :: + (TrieKey k, Sized a) => Maybe a -> V(EdgeLoc) a -> V(MEdge) a, + Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} +beforeEdge :: (Label v k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a +beforeEdge v LOC(ks ts path) = buildBefore (cEdge ks v ts) path where + buildBefore !e path = case pView path of + Root -> e + Deep path ks v tHole -> buildBefore (cEdge ks v $ beforeMM e tHole) path -{-# SPECIALIZE afterEdge :: 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 afterEdge :: + (TrieKey k, Sized a) => Maybe a -> V(EdgeLoc) a -> V(MEdge) a, + Sized a => Maybe a -> U(EdgeLoc) a -> U(MEdge) a #-} +afterEdge :: (Label v k, Sized a) => Maybe a -> EdgeLoc v k a -> MEdge v k a +afterEdge v LOC(ks ts path) = buildAfter (cEdge ks v ts) path where + buildAfter !e path = case pView path of + Root -> e + Deep path ks v tHole + -> buildAfter (cEdge ks v $ afterMM e tHole) path -{-# 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 +{-# SPECIALIZE extractEdgeLoc :: + (TrieKey k, Functor m, MonadPlus m) => V(Edge) a -> V(Path) a -> m (a, V(EdgeLoc) a), + (Functor m, MonadPlus m) => U(Edge) a -> U(Path) a -> m (a, U(EdgeLoc) a) #-} +extractEdgeLoc :: (Label v k, Functor m, MonadPlus m) => Edge v k a -> Path v k a -> m (a, EdgeLoc v k a) +extractEdgeLoc !EDGE(_ ks v ts) path = case v of Nothing -> extractTS - Just a -> return (a, Loc ks ts path) `mplus` 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) + 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 #) +{-# SPECIALIZE INLINE indexEdge :: + (TrieKey k, Sized a) => Int -> V(Edge) a -> V(Path) a -> (# Int, a, V(EdgeLoc) a #), + Sized a => Int -> U(Edge) a -> U(Path) a -> (# Int, a, U(EdgeLoc) a #) #-} +indexEdge :: (Label v k, Sized a) => Int -> Edge v k a -> Path v k a -> (# Int, a, EdgeLoc v k a #) indexEdge = indexE where - indexE i# (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 + indexE !i e path = case eView e of + Edge _ ks v@(Just a) ts + | i < sv -> (# i, a, loc ks ts path #) + | (# i', e', tHole #) <- indexM (i - sv) ts + -> indexE i' e' (deep path ks v tHole) + where !sv = getSize a + Edge _ ks Nothing ts + -> indexE i' e' (deep path ks Nothing tHole) + where !(# i', e', tHole #) = indexM i ts -{-# SPECIALIZE 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) +{-# SPECIALIZE insertEdge :: + (TrieKey k, Sized a) => (a -> a) -> V() -> a -> V(Edge) a -> V(Edge) a, + Sized a => (a -> a) -> U() -> a -> U(Edge) a -> U(Edge) a #-} +insertEdge :: (Label v k, Sized a) => (a -> a) -> v k -> a -> Edge v k a -> Edge v k a +insertEdge f ks v e = searchEdgeC ks e nomatch match where + nomatch = assignEdge v + match = assignEdge . f diff --git a/Data/TrieMap/RadixTrie/Label.hs b/Data/TrieMap/RadixTrie/Label.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/RadixTrie/Label.hs @@ -0,0 +1,152 @@ +{-# LANGUAGE MagicHash, TypeFamilies, MultiParamTypeClasses, FlexibleInstances, BangPatterns, CPP, ViewPatterns #-} +{-# OPTIONS -funbox-strict-fields #-} +module Data.TrieMap.RadixTrie.Label where + +import Data.TrieMap.TrieKey +import Data.TrieMap.Sized +import Data.TrieMap.RadixTrie.Slice +import Data.TrieMap.WordMap + +import Data.Word +import Data.Vector.Generic +import qualified Data.Vector as V +import qualified Data.Vector.Storable as S + +import Prelude hiding (length) + +#define V(ty) (ty (V.Vector) (k)) +#define U(ty) (ty (S.Vector) Word) + +class (Vector v k, TrieKey k) => Label v k where + data Edge v k :: * -> * + data Path v k :: * -> * + data EdgeLoc v k :: * -> * + edge :: Sized a => v k -> Maybe a -> Branch v k a -> Edge v k a + edge' :: Int -> v k -> Maybe a -> Branch v k a -> Edge v k a + root :: Path v k a + deep :: Path v k a -> v k -> Maybe a -> BHole v k a -> Path v k a + loc :: v k -> Branch v k a -> Path v k a -> EdgeLoc v k a + + eView :: Edge v k a -> EView v k a + pView :: Path v k a -> PView v k a + locView :: EdgeLoc v k a -> LocView v k a + +type BHole v k a = Hole k (Edge v k a) + +type Branch v k a = TrieMap k (Edge v k a) +data EView v k a = + Edge Int (v k) (Maybe a) (Branch v k a) +data LocView v k a = Loc !( v k) (Branch v k a) (Path v k a) +data PView v k a = Root + | Deep (Path v k a) (v k) (Maybe a) (BHole v k a) +type MEdge v k a = Maybe (Edge v k a) + +instance Sized (EView v k a) where + getSize# (Edge sz _ _ _) = unbox sz + +instance Label v k => Sized (Edge v k a) where + {-# SPECIALIZE instance TrieKey k => Sized (Edge V.Vector k a) #-} + getSize# e = getSize# (eView e) + +instance TrieKey k => Label V.Vector k where + data Edge V.Vector k a = + VEdge Int !(V()) (V(Branch) a) + | VEdgeX Int !(V()) a (V(Branch) a) + data Path V.Vector k a = + VRoot + | VDeep (V(Path) a) !(V()) (V(BHole) a) + | VDeepX (V(Path) a) !(V()) a (V(BHole) a) + data EdgeLoc V.Vector k a = VLoc !(V()) (V(Branch) a) (V(Path) a) + + edge !ks Nothing ts = VEdge (sizeM ts) ks ts + edge !ks (Just a) ts = VEdgeX (sizeM ts + getSize a) ks a ts + edge' s !ks Nothing ts = VEdge s ks ts + edge' s !ks (Just a) ts = VEdgeX s ks a ts + + root = VRoot + deep path !ks Nothing tHole = VDeep path ks tHole + deep path !ks (Just a) tHole = VDeepX path ks a tHole + + loc = VLoc + + eView (VEdge s ks ts) = Edge s ks Nothing ts + eView (VEdgeX s ks v ts) = Edge s ks (Just v) ts + pView VRoot = Root + pView (VDeep path ks tHole) = Deep path ks Nothing tHole + pView (VDeepX path ks v tHole) = Deep path ks (Just v) tHole + locView (VLoc ks ts path) = Loc ks ts path + +instance Label S.Vector Word where + data Edge S.Vector Word a = + SEdge !Int !(U()) !(SNode (U(Edge) a)) + | SEdgeX !Int !(U()) a !(SNode (U(Edge) a)) + data Path S.Vector Word a = + SRoot + | SDeep (U(Path) a) !(U()) !(WHole (U(Edge) a)) + | SDeepX (U(Path) a) !(U()) a !(WHole (U(Edge) a)) + data EdgeLoc S.Vector Word a = + SLoc !(U()) !(SNode (U(Edge) a)) (U(Path) a) + + edge !ks Nothing ts = SEdge (sizeM ts) ks (getWordMap ts) + edge !ks (Just v) ts = SEdgeX (getSize v + sizeM ts) ks v (getWordMap ts) + edge' sz !ks Nothing ts = SEdge sz ks (getWordMap ts) + edge' sz !ks (Just v) ts = SEdgeX sz ks v (getWordMap ts) + + root = SRoot + deep path !ks Nothing tHole = SDeep path ks (getHole tHole) + deep path !ks (Just v) tHole = SDeepX path ks v (getHole tHole) + + loc ks ts path = SLoc ks (getWordMap ts) path + + eView (SEdge s ks ts) = Edge s ks Nothing (WordMap ts) + eView (SEdgeX s ks v ts) = Edge s ks (Just v) (WordMap ts) + pView SRoot = Root + pView (SDeep path ks tHole) = Deep path ks Nothing (Hole tHole) + pView (SDeepX path ks v tHole) = Deep path ks (Just v) (Hole tHole) + locView (SLoc ks ts path) = Loc ks (WordMap ts) path + +{-# SPECIALIZE singletonEdge :: + (TrieKey k, Sized a) => V() -> a -> V(Edge) a, + Sized a => U() -> a -> U(Edge) a #-} +singletonEdge :: (Label v k, Sized a) => v k -> a -> Edge v k a +singletonEdge ks a = edge ks (Just a) emptyM + +{-# SPECIALIZE singleLoc :: + TrieKey k => V() -> V(EdgeLoc) a, + U() -> U(EdgeLoc) a #-} +singleLoc :: Label v k => v k -> EdgeLoc v k a +singleLoc ks = loc ks emptyM root + +{-# SPECIALIZE getSimpleEdge :: + TrieKey k => V(Edge) a -> Simple a, + U(Edge) a -> Simple a #-} +getSimpleEdge :: Label v k => Edge v k a -> Simple a +getSimpleEdge !(eView -> Edge _ _ v ts) + | nullM ts = maybe Null Singleton v + | otherwise = NonSimple + +{-# SPECIALIZE INLINE dropEdge :: + TrieKey k => Int -> V(Edge) a -> V(Edge) a, + Int -> U(Edge) a -> U(Edge) a #-} +{-# SPECIALIZE INLINE unDropEdge :: + TrieKey k => Int -> V(Edge) a -> V(Edge) a, + Int -> U(Edge) a -> U(Edge) a #-} +dropEdge, unDropEdge :: Label v k => Int -> Edge v k a -> Edge v k a +dropEdge !n !(eView -> Edge sz# ks v ts) = edge' sz# (dropSlice n ks) v ts +unDropEdge !n !(eView -> Edge sz# ks v ts) = edge' sz# (unDropSlice n ks) v ts + +{-# SPECIALIZE compact :: + TrieKey k => V(Edge) a -> V(MEdge) a, + U(Edge) a -> U(MEdge) a #-} +compact :: Label v k => Edge v k a -> MEdge v k a +compact !e@(eView -> Edge _ ks Nothing ts) = case getSimpleM ts of + Null -> Nothing + Singleton e' -> Just (unDropEdge (length ks + 1) e') + NonSimple -> Just e +compact e = Just e + +{-# SPECIALIZE cEdge :: + (TrieKey k, Sized a) => V() -> Maybe a -> V(Branch) a -> V(MEdge) a, + Sized a => U() -> Maybe a -> U(Branch) a -> U(MEdge) a #-} +cEdge :: (Label v k, Sized a) => v k -> Maybe a -> Branch v k a -> MEdge v k a +cEdge ks v ts = compact (edge ks v ts) diff --git a/Data/TrieMap/RadixTrie/Slice.hs b/Data/TrieMap/RadixTrie/Slice.hs --- a/Data/TrieMap/RadixTrie/Slice.hs +++ b/Data/TrieMap/RadixTrie/Slice.hs @@ -8,41 +8,28 @@ 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 :: Vector v a => Int -> v a -> (v a, a, 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 +takeSlice :: Vector v a => Int -> v a -> v a +takeSlice !n xs = assert (n <= length xs) $ unsafeTake n xs -{-# INLINE v2S #-} -v2S :: Vector v a => v a -> Slice v a -v2S xs = Slice xs 0 (length xs) +dropSlice :: Vector v a => Int -> v a -> v a +dropSlice !n xs = assert (n <= length xs) $ unsafeDrop n 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)) +-- | Do you have any idea how unsafe this method is? No, because you're STILL SANE ENOUGH TO READ THIS. +unDropSlice :: Vector v a => Int -> v a -> v a +unDropSlice !n = unsafeDrop (-n) {-# 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)) +matchSlice :: (Vector v a, Vector v b) => (a -> b -> z -> z) -> (Int -> Int -> z) -> v a -> v b -> z +matchSlice f z !xs !ys = foldr (\ (a, b) -> f a b) (z (length xs) (length ys)) (V.zip (convert xs) (convert ys)) {-# 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)) +iMatchSlice :: (Vector v a, Vector v b) => (Int -> a -> b -> z -> z) -> (Int -> Int -> z) -> v a -> v b -> z +iMatchSlice f z !xs !ys = ifoldr (\ i (a, b) -> f i a b) (z (length xs) (length ys)) (V.zip (convert xs) (convert ys)) {-# INLINE (!$) #-} -(!$) :: Vector v a => Slice v a -> Int -> a -Slice xs i n !$ j = assert (j >= 0 && j < n) $ unsafeIndex xs (i + j) +(!$) :: Vector v a => v a -> Int -> a +xs !$ j = assert (j >= 0 && j < length xs) $ unsafeIndex xs j diff --git a/Data/TrieMap/Representation/Class.hs b/Data/TrieMap/Representation/Class.hs --- a/Data/TrieMap/Representation/Class.hs +++ b/Data/TrieMap/Representation/Class.hs @@ -1,6 +1,8 @@ {-# LANGUAGE TypeFamilies #-} module Data.TrieMap.Representation.Class where +import Data.Vector + -- | 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. -- @@ -12,5 +14,17 @@ -- 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 + type Rep a + type RepList a + toRep :: a -> Rep a + toRepList :: [a] -> RepList a + +type DRepList a = Vector (Rep a) +dToRepList :: Repr a => [a] -> DRepList a +dToRepList = fromList . Prelude.map toRep + +instance Repr a => Repr [a] where + type Rep [a] = RepList a + type RepList [a] = Vector (RepList a) + toRep = toRepList + toRepList = dToRepList diff --git a/Data/TrieMap/Representation/Instances.hs b/Data/TrieMap/Representation/Instances.hs --- a/Data/TrieMap/Representation/Instances.hs +++ b/Data/TrieMap/Representation/Instances.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE TemplateHaskell, QuasiQuotes, TypeFamilies, FlexibleInstances #-} +{-# LANGUAGE TemplateHaskell, QuasiQuotes, TypeFamilies, FlexibleInstances, CPP, UndecidableInstances #-} module Data.TrieMap.Representation.Instances () where import Data.Tree @@ -21,18 +21,25 @@ import Data.TrieMap.Representation.Instances.Foreign () import Data.TrieMap.Representation.TH +#define DefList(ty) \ + type RepList (ty) = DRepList (ty); \ + toRepList = dToRepList + 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 + DefList(S.Set a) 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 + DefList(M.Map k a) instance Repr a => Repr (Seq.Seq a) where type Rep (Seq.Seq a) = V.Vector (Rep a) toRep = toVectorF toRep Seq.length + DefList(Seq.Seq a) genRepr ''Tree genRepr ''Ratio @@ -42,6 +49,7 @@ 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) + DefList(Integer) unroll :: Integer -> S.Vector Word unroll x = S.reverse (S.unfoldr split x) diff --git a/Data/TrieMap/Representation/Instances/Basic.hs b/Data/TrieMap/Representation/Instances/Basic.hs --- a/Data/TrieMap/Representation/Instances/Basic.hs +++ b/Data/TrieMap/Representation/Instances/Basic.hs @@ -1,39 +1,23 @@ {-# LANGUAGE TemplateHaskell, TypeFamilies #-} module Data.TrieMap.Representation.Instances.Basic () where -import Data.TrieMap.Representation.Class import Data.TrieMap.Representation.TH - -import Control.Monad +import Data.TrieMap.Representation.Class -import qualified Data.Vector as V +import Data.Word 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]) +$(fmap concat $ mapM (genRepr . tupleTypeName) [2..10]) genOrdRepr ''Float genOrdRepr ''Double genRepr ''Maybe genRepr ''Either -genRepr ''Bool -genRepr ''() genRepr ''Ordering + +instance Repr () where + type Rep () = () + toRep _ = () + type RepList () = Word + toRepList = fromIntegral . length diff --git a/Data/TrieMap/Representation/Instances/ByteString.hs b/Data/TrieMap/Representation/Instances/ByteString.hs --- a/Data/TrieMap/Representation/Instances/ByteString.hs +++ b/Data/TrieMap/Representation/Instances/ByteString.hs @@ -12,10 +12,16 @@ import Data.Vector.Storable +-- | @'Rep' 'ByteString' = 'Rep' ('Vector' 'Word8')@ instance Repr ByteString where - type Rep ByteString = (Vector Word, Word) + type Rep ByteString = Rep (Vector Word8) toRep (PS fp off len) = toRep (unsafeFromForeignPtr fp off len) + type RepList ByteString = DRepList ByteString + toRepList = dToRepList +-- | @'Rep' 'L.ByteString' = 'Rep' ('Vector' 'Word8')@ instance Repr L.ByteString where - type Rep L.ByteString = (Vector Word, Word) + type Rep L.ByteString = Rep (Vector Word8) toRep = toRep . B.concat . L.toChunks + type RepList L.ByteString = DRepList L.ByteString + toRepList = dToRepList diff --git a/Data/TrieMap/Representation/Instances/Prim.hs b/Data/TrieMap/Representation/Instances/Prim.hs --- a/Data/TrieMap/Representation/Instances/Prim.hs +++ b/Data/TrieMap/Representation/Instances/Prim.hs @@ -1,52 +1,73 @@ {-# LANGUAGE ScopedTypeVariables, BangPatterns, TypeFamilies, UndecidableInstances, CPP #-} -module Data.TrieMap.Representation.Instances.Prim (i2w) where +module Data.TrieMap.Representation.Instances.Prim () where #include "MachDeps.h" import Data.TrieMap.Representation.Class +import Data.TrieMap.Representation.Instances.Vectors import Data.Word import Data.Int import Data.Char import Data.Bits +import Data.Vector.Storable +import qualified Data.Vector.Unboxed as U +import Prelude hiding (map) +#define WDOC(ty) {-| @'Rep' 'ty' = 'Word'@ -} + +WDOC(Char) instance Repr Char where type Rep Char = Word + type RepList Char = Vector Word toRep = fromIntegral . ord + toRepList xs = toRep (fromList xs) #define WREPR(wTy) \ instance Repr wTy where { \ type Rep wTy = Word; \ - toRep = fromIntegral} + toRep = fromIntegral; \ + type RepList wTy = Rep (Vector wTy);\ + toRepList xs = toRep (fromList xs)} WREPR(Word) +WDOC(Word8) WREPR(Word8) +WDOC(Word16) WREPR(Word16) +WDOC(Word32) WREPR(Word32) #if WORD_SIZE_IN_BITS < 64 +-- | @'Rep' 'Word64' = ('Word', 'Word')@ 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 + type RepList Word64 = Vector Word + toRepList xs = toRep (fromList xs) #else +WDOC(Word64) 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)} + toRep = toRep . (i2w :: iTy -> wTy); \ + type RepList iTy = Rep (Vector wTy); \ + toRepList xs = toRep (fromList xs)} IREPR(Int8,Word8) IREPR(Int16,Word16) IREPR(Int32,Word32) IREPR(Int64,Word64) +-- | @'Rep' 'Int' = 'Word'@, by way of a careful translation of their domains to avoid overflow. IREPR(Int,Word) + +instance Repr Bool where + type Rep Bool = Either () () + toRep False = Left () + toRep True = Right () + type RepList Bool = (Vector Word, Word) + toRepList xs = toRep (U.fromList xs) diff --git a/Data/TrieMap/Representation/Instances/Vectors.hs b/Data/TrieMap/Representation/Instances/Vectors.hs --- a/Data/TrieMap/Representation/Instances/Vectors.hs +++ b/Data/TrieMap/Representation/Instances/Vectors.hs @@ -1,5 +1,6 @@ -{-# LANGUAGE TypeFamilies, FlexibleInstances, CPP, BangPatterns, UndecidableInstances, TemplateHaskell #-} -module Data.TrieMap.Representation.Instances.Vectors () where +{-# LANGUAGE TypeFamilies, FlexibleInstances, CPP, BangPatterns, ScopedTypeVariables, UndecidableInstances, FlexibleContexts #-} +{-# OPTIONS -funbox-strict-fields #-} +module Data.TrieMap.Representation.Instances.Vectors (i2w) where import Control.Monad.Primitive @@ -11,102 +12,103 @@ import Foreign.Ptr import Foreign.ForeignPtr -import Data.Vector.Generic (convert) +import Data.Vector.Generic (convert, stream, unstream) 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.Vector.Fusion.Stream.Monadic +import Data.Vector.Fusion.Stream.Size + +import Data.TrieMap.Utils import Data.TrieMap.Representation.Class -import Data.TrieMap.Representation.Instances.Prim -import Language.Haskell.TH.Syntax +import Prelude hiding (length) +import GHC.Exts #include "MachDeps.h" +#define DefList(ty) \ + type RepList (ty) = DRepList (ty); \ + toRepList = dToRepList + instance Repr a => Repr (V.Vector a) where type Rep (V.Vector a) = V.Vector (Rep a) toRep = V.map toRep + DefList(V.Vector a) 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]]) + DefList(S.Vector Word) --- Just a version of 'rem' for modding by powers of 2. -remPow :: Int -> Int -> Int -remPow n d = n .&. (d - 1) +{-# INLINE unsafeCastStorable #-} +unsafeCastStorable :: (Storable a, Storable b) => (Int -> Int) -> S.Vector a -> S.Vector b +unsafeCastStorable f xs = unsafeInlineST $ do + S.MVector ptr n fp <- S.unsafeThaw xs + let n' = f n + S.unsafeFreeze (S.MVector (castPtr ptr) n' (castForeignPtr fp)) -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) +wordSize :: Int +wordSize = bitSize (0 :: Word) -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 VEC_WORD_INST(vec,wTy) \ + instance Repr (vec wTy) where { \ + type Rep (vec wTy) = Rep (S.Vector wTy); \ + toRep xs = toHangingVector xs;\ + DefList(vec wTy)} +#define HANGINSTANCE(wTy) \ + instance Repr (S.Vector wTy) where { \ + type Rep (S.Vector wTy) = (S.Vector Word, Word);\ + {-# INLINE toRep #-}; \ + toRep xs = toHangingVector xs; \ + DefList(S.Vector wTy) }; \ + VEC_WORD_INST(P.Vector,wTy); \ + VEC_WORD_INST(U.Vector,wTy) -#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) +{-# INLINE toHangingVector #-} +toHangingVector :: (G.Vector v w, Bits w, Integral w, Storable w) => v w -> (S.Vector Word, Word) +toHangingVector xs = let !ys = unstream (packStream (stream xs)) in (S.unsafeInit ys, S.unsafeLast ys) +-- | @'Rep' ('S.Vector' 'Word8') = 'S.Vector' 'Word'@, by packing multiple 'Word8's into each 'Word' for space efficiency. HANGINSTANCE(Word8) +-- | @'Rep' ('S.Vector' 'Word16') = 'S.Vector' 'Word'@, by packing multiple 'Word16's into each 'Word' for space efficiency. 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 + toRep xs = unsafeCastStorable id xs + DefList (S.Vector Word32) +instance Repr (U.Vector Word32) where + type Rep (U.Vector Word32) = S.Vector Word + toRep xs = unsafeCastStorable id (convert xs) + DefList (U.Vector Word32) +instance Repr (P.Vector Word32) where + type Rep (P.Vector Word32) = S.Vector Word + toRep xs = unsafeCastStorable id (convert xs) + DefList (P.Vector Word32) #elif WORD_SIZE_IN_BITS > 32 HANGINSTANCE(Word32) #endif +#if WORD_SIZE_IN_BITS == 32 +-- | @'Rep' ('S.Vector' 'Word64') = 'S.Vector' 'Word'@, by viewing each 'Word64' as two 'Word's. +#else +-- | @'Rep' ('S.Vector' 'Word64') = 'S.Vector' 'Word'@ +#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 + toRep xs = unsafeCastStorable (ratio *) xs 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) + DefList(S.Vector Word64) -#define VEC_INT_INST(vec,iTy,wTy) \ - instance Repr (vec iTy) where { \ +#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)} + toRep xs = (toRep :: S.Vector wTy -> Rep (S.Vector wTy)) (convert (G.map (i2w :: iTy -> wTy) xs)); \ + DefList(vec iTy)} #define VEC_INT_INSTANCES(iTy,wTy) \ VEC_INT_INST(S.Vector,iTy,wTy); \ VEC_INT_INST(P.Vector,iTy,wTy); \ @@ -121,10 +123,80 @@ #define VEC_ENUM_INST(ty, vec) \ instance Repr (vec ty) where { \ type Rep (vec ty) = S.Vector Word; \ - toRep = convert . G.map (fromIntegral . fromEnum)} + {-# INLINE toRep #-}; \ + toRep xs = convert (G.map (fromIntegral . fromEnum) xs);\ + DefList(vec ty)} #define VEC_ENUM_INSTANCES(ty) \ VEC_ENUM_INST(ty,S.Vector); \ VEC_ENUM_INST(ty,P.Vector); \ VEC_ENUM_INST(ty,U.Vector) +-- | @'Rep' ('S.Vector' 'Char') = 'S.Vector' 'Word'@ VEC_ENUM_INSTANCES(Char) + +-- | 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 + +data PackState s = PackState !Word !Int s | Last !Int | End +{-# ANN type PackState ForceSpecConstr #-} + +{-# INLINE packStream #-} +packStream :: forall m w . (Bits w, Integral w, Storable w, Monad m) => Stream m w -> Stream m Word +packStream (Stream step s0 size) = Stream step' s0' size' + where !ratio = wordSize `quoPow` bitSize (0 :: w) + size' = 1 + case size of + Exact n -> Exact $ (n + ratio - 1) `quoPow` ratio + Max n -> Max $ (n + ratio - 1) `quoPow` ratio + Unknown -> Unknown + s0' = PackState 0 ratio s0 + step' End = return Done + step' (Last i) = return $ Yield (fromIntegral i) End + step' (PackState w 0 s) = return $ Yield w (PackState 0 ratio s) + step' (PackState w i s) = do + s' <- step s + case s' of + Done | i == ratio -> return $ Skip (Last 0) + | otherwise -> return $ Yield (w .<<. (i * bitSize (0 :: w))) (Last (ratio - i)) + Skip s' -> return $ Skip (PackState w i s') + Yield ww s' -> return $ Skip (PackState ((w .<<. bitSize (0 :: w)) .|. fromIntegral ww) (i-1) s') + +instance Repr (S.Vector Bool) where + type Rep (S.Vector Bool) = (S.Vector Word, Word) + toRep = boolVecToRep + DefList(S.Vector Bool) + +instance Repr (U.Vector Bool) where + type Rep (U.Vector Bool) = (S.Vector Word, Word) + {-# INLINE toRep #-} + toRep xs = boolVecToRep xs + DefList(U.Vector Bool) + +{-# INLINE boolVecToRep #-} +boolVecToRep :: G.Vector v Bool => v Bool -> (S.Vector Word, Word) +boolVecToRep xs = let !ys = unstream (packBoolStream (stream xs)) in (S.unsafeInit ys, S.unsafeLast ys) + +{-# INLINE packBoolStream #-} +packBoolStream :: Monad m => Stream m Bool -> Stream m Word +packBoolStream (Stream step s0 size) = Stream step' s0' size' + where !ratio = wordSize + size' = 1 + case size of + Exact n -> Exact $ (n + ratio - 1) `quoPow` ratio + Max n -> Max $ (n + ratio - 1) `quoPow` ratio + Unknown -> Unknown + s0' = PackState 0 ratio s0 + toW False = 0 + toW True = 1 + step' End = return Done + step' (Last i) = return $ Yield (fromIntegral i) End + step' (PackState w 0 s) = return $ Yield w (PackState 0 ratio s) + step' (PackState w i s) = do + s' <- step s + case s' of + Done | i == ratio -> return $ Skip (Last 0) + | otherwise -> return $ Yield (w .<<. i) (Last (ratio - i)) + Skip s' -> return $ Skip (PackState w i s') + Yield ww s' -> return $ Skip (PackState ((w .<<. 1) .|. toW ww) (i-1) s') 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 @@ -23,7 +23,7 @@ getDataForName :: Quasi m => Name -> m (Cxt, Type, [AlgCon]) getDataForName tycon = do TyConI dec <- qReify tycon - let theTyp = compose tycon . map tyVarBndrVar + let theTyp = compose tycon . map (mkName . nameBase . tyVarBndrVar) case dec of DataD cxt _ tyvars cons _ -> return (cxt, theTyp tyvars, map algCon cons) diff --git a/Data/TrieMap/Representation/TH/Representation.hs b/Data/TrieMap/Representation/TH/Representation.hs --- a/Data/TrieMap/Representation/TH/Representation.hs +++ b/Data/TrieMap/Representation/TH/Representation.hs @@ -112,7 +112,9 @@ [InstanceD cxt (ConT ''Repr `AppT` ty) [TySynInstD ''Rep [ty] reprType, FunD 'toRep - (map caseToClause cases)]] + (map caseToClause cases), + TySynInstD ''RepList [ty] (ConT ''V.Vector `AppT` reprType), + ValD (VarP 'toRepList) (NormalB (VarE 'dToRepList)) []]] return reprType recursiveRepr :: Quasi m => Type -> Exp -> m Representation diff --git a/Data/TrieMap/ReverseMap.hs b/Data/TrieMap/ReverseMap.hs --- a/Data/TrieMap/ReverseMap.hs +++ b/Data/TrieMap/ReverseMap.hs @@ -1,23 +1,42 @@ -{-# LANGUAGE TypeFamilies, MagicHash, UnboxedTuples #-} +{-# LANGUAGE TypeFamilies, MagicHash, UnboxedTuples, GeneralizedNewtypeDeriving, FlexibleInstances #-} module Data.TrieMap.ReverseMap () where import Control.Applicative +import Control.Monad +import Control.Monad.Ends -import Data.TrieMap.Applicative +import Data.Foldable +import qualified Data.Monoid as M + import Data.TrieMap.TrieKey import Data.TrieMap.Modifiers import Data.TrieMap.Sized -import GHC.Exts +import Prelude hiding (foldr, foldl, foldr1, foldl1) +newtype DualPlus m a = DualPlus {runDualPlus :: m a} deriving (Functor, Monad) +newtype Dual f a = Dual {runDual :: f a} deriving (Functor) + +instance Applicative f => Applicative (Dual f) where + pure a = Dual (pure a) + Dual f <*> Dual x = Dual (x <**> f) + +instance MonadPlus m => MonadPlus (DualPlus m) where + mzero = DualPlus mzero + DualPlus m `mplus` DualPlus k = DualPlus (k `mplus` m) + +instance TrieKey k => Foldable (TrieMap (Rev k)) where + foldMap f (RevMap m) = M.getDual (foldMap (M.Dual . f) m) + foldr f z (RevMap m) = foldl (flip f) z m + foldl f z (RevMap m) = foldr (flip f) z m + foldr1 f (RevMap m) = foldl1 (flip f) m + foldl1 f (RevMap m) = foldr1 (flip f) m + -- | @'TrieMap' ('Rev' k) a@ is a wrapper around a @'TrieMap' k a@ that reverses the order of the operations. instance TrieKey k => TrieKey (Rev k) where newtype TrieMap (Rev k) a = RevMap (TrieMap k a) newtype Hole (Rev k) a = RHole (Hole k a) - 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 @@ -27,9 +46,6 @@ 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) @@ -38,21 +54,26 @@ isSubmapM (<=) (RevMap m1) (RevMap m2) = isSubmapM (<=) m1 m2 singleHoleM (Rev k) = RHole (singleHoleM k) - 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) + beforeM (RHole hole) = RevMap (afterM hole) + beforeWithM a (RHole hole) = RevMap (afterWithM a hole) + afterM (RHole hole) = RevMap (beforeM hole) + afterWithM a (RHole hole) = RevMap (beforeWithM a hole) + searchMC (Rev k) (RevMap m) = mapSearch RHole (searchMC k m) + indexM i (RevMap m) = case indexM (revIndex i m) m of + (# i', a, hole #) -> (# revIndex i' a, a, RHole hole #) + where revIndex :: Sized a => Int -> a -> Int + revIndex i a = getSize a - 1 - i + + extractHoleM (RevMap m) = fmap RHole <$> runDualPlus (extractHoleM m) + firstHoleM (RevMap m) = First (fmap RHole <$> getLast (lastHoleM m)) + lastHoleM (RevMap m) = Last (fmap RHole <$> getFirst (firstHoleM m)) + assignM v (RHole m) = RevMap (assignM v m) + clearM (RHole m) = RevMap (clearM m) + insertWithM f (Rev k) a (RevMap m) = RevMap (insertWithM f k a m) fromListM f xs = RevMap (fromListM f [(k, a) | (Rev k, a) <- xs]) fromAscListM f xs = RevMap (fromAscListM (flip f) [(k, a) | (Rev k, a) <- reverse xs]) fromDistAscListM xs = RevMap (fromDistAscListM [(k, a) | (Rev k, a) <- reverse xs]) - unifyM (Rev k1) a1 (Rev k2) a2 = either (Left . RHole) (Right . RevMap) (unifyM k1 a1 k2 a2) - -revIndex :: Sized a => Int# -> a -> Int# -revIndex i# a = getSize# a -# 1# -# i# + unifierM (Rev k') (Rev k) a = RHole <$> unifierM k' k a diff --git a/Data/TrieMap/Sized.hs b/Data/TrieMap/Sized.hs --- a/Data/TrieMap/Sized.hs +++ b/Data/TrieMap/Sized.hs @@ -1,15 +1,17 @@ -{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MagicHash, DeriveFunctor, DeriveFoldable, DeriveTraversable #-} module Data.TrieMap.Sized where +import Data.Foldable +import Data.Traversable import GHC.Exts class Sized a where getSize# :: a -> Int# -data Assoc k a = Assoc {getK :: k, getValue :: a} +data Assoc k a = Assoc {getK :: k, getValue :: a} deriving (Functor, Foldable, Traversable) -newtype Elem a = Elem a +newtype Elem a = Elem {getElem :: a} deriving (Functor, Foldable, Traversable) instance Sized (Elem a) where getSize# _ = 1# @@ -21,8 +23,10 @@ getSize# (Just a) = getSize# a getSize# _ = 0# +{-# INLINE getSize #-} getSize :: Sized a => a -> Int getSize a = I# (getSize# a) +{-# INLINE unbox #-} unbox :: Int -> Int# unbox (I# i#) = i# diff --git a/Data/TrieMap/TrieKey.hs b/Data/TrieMap/TrieKey.hs --- a/Data/TrieMap/TrieKey.hs +++ b/Data/TrieMap/TrieKey.hs @@ -1,24 +1,45 @@ -{-# LANGUAGE TupleSections, TypeFamilies, UnboxedTuples, MagicHash #-} +{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash, FlexibleContexts, TupleSections, Rank2Types #-} module Data.TrieMap.TrieKey where import Data.TrieMap.Sized +import Data.TrieMap.Utils -import Control.Applicative +import Control.Applicative (Applicative) import Control.Monad +import Control.Monad.Ends -import Data.Monoid import Data.Foldable hiding (foldrM, foldlM) +import qualified Data.List as L import Prelude hiding (foldr, foldl) import GHC.Exts type LEq a b = a -> b -> Bool -type Unified k a = Either (Hole k a) (TrieMap k a) +type SearchCont h a r = (h -> r) -> (a -> h -> r) -> r +type Lookup a = Maybe a data Simple a = Null | Singleton a | NonSimple +class (Functor f, Monad f) => Option f where + none :: f a + some :: a -> f a + option :: f a -> r -> (a -> r) -> r + +instance Option Maybe where + none = Nothing + some = Just + option m a f = maybe a f m + +{-# INLINE [0] liftMaybe #-} +liftMaybe :: Option f => Maybe a -> f a +liftMaybe = maybe none some + +{-# INLINE [0] toMaybe #-} +toMaybe :: Option f => f a -> Maybe a +toMaybe x = option x Nothing Just + instance Monad Simple where return = Singleton Null >>= _ = Null @@ -31,35 +52,29 @@ simple `mplus` Null = simple _ `mplus` _ = NonSimple +{-# INLINE onSnd #-} onSnd :: (c -> d) -> (a -> (# b, c #)) -> a -> (# b, d #) onSnd g f a = case f a of (# b, c #) -> (# b, g c #) -onThird :: (d -> e) -> (a -> (# Int#, c, d #)) -> a -> (# Int#, c, e #) +{-# INLINE onThird #-} +onThird :: (d -> e) -> (a -> (# Int, c, d #)) -> a -> (# Int, c, e #) onThird g f a = case f a of (# b, c, d #) -> (# b, c, g d #) -instance TrieKey k => Foldable (TrieMap k) where - foldr f = flip $ foldrM f - foldl f = flip $ foldlM f - -- | A @TrieKey k@ instance implies that @k@ is a standardized representation for which a -- generalized trie structure can be derived. -class TrieKey k where - (=?) :: k -> k -> Bool - cmp :: k -> k -> Ordering - +class (Ord k, Foldable (TrieMap k)) => TrieKey k where data TrieMap k :: * -> * emptyM :: TrieMap k a singletonM :: Sized a => k -> a -> TrieMap k a getSimpleM :: TrieMap k a -> Simple a - sizeM :: Sized a => TrieMap k a -> Int# - lookupM :: k -> TrieMap k a -> Maybe a + sizeM# :: Sized a => TrieMap k a -> Int# + sizeM :: Sized a => TrieMap k a -> Int + lookupM :: k -> TrieMap k a -> Lookup a fmapM :: Sized b => (a -> b) -> TrieMap k a -> TrieMap k b traverseM :: (Applicative f, Sized b) => (a -> f b) -> TrieMap k a -> f (TrieMap k b) - foldrM :: (a -> b -> b) -> TrieMap k a -> b -> b - foldlM :: (b -> a -> b) -> TrieMap k a -> b -> b mapMaybeM :: Sized b => (a -> Maybe b) -> TrieMap k a -> TrieMap k b mapEitherM :: (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> TrieMap k a -> (# TrieMap k b, TrieMap k c #) unionM :: Sized a => (a -> a -> Maybe a) -> TrieMap k a -> TrieMap k a -> TrieMap k a @@ -67,42 +82,83 @@ (a -> b -> Maybe c) -> TrieMap k a -> TrieMap k b -> TrieMap k c diffM :: Sized a => (a -> b -> Maybe a) -> TrieMap k a -> TrieMap k b -> TrieMap k a isSubmapM :: (Sized a, Sized b) => LEq a b -> LEq (TrieMap k a) (TrieMap k b) + fromListM, fromAscListM :: Sized a => (a -> a -> a) -> [(k, a)] -> TrieMap k a fromDistAscListM :: Sized a => [(k, a)] -> TrieMap k a + insertWithM :: (TrieKey k, Sized a) => (a -> a) -> k -> a -> TrieMap k a -> TrieMap k a data Hole k :: * -> * singleHoleM :: k -> Hole k a - beforeM :: 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 => Maybe a -> Hole k a -> TrieMap k a + beforeM, afterM :: Sized a => Hole k a -> TrieMap k a + beforeWithM, afterWithM :: Sized a => a -> Hole k a -> TrieMap k a + searchMC :: k -> TrieMap k a -> SearchCont (Hole k a) a r + indexM :: Sized a => Int -> TrieMap k a -> (# Int, a, Hole k a #) + indexM# :: Sized a => Int# -> TrieMap k a -> (# Int#, a, Hole k a #) - fromListM f = foldr (\ (k, a) -> insertWithM f k a) emptyM + -- By combining rewrite rules and these NOINLINE pragmas, we automatically derive + -- specializations of functions for every instance of TrieKey. + extractHoleM :: (Functor m, MonadPlus m) => Sized a => TrieMap k a -> m (a, Hole k a) + {-# NOINLINE firstHoleM #-} + {-# NOINLINE lastHoleM #-} + {-# NOINLINE sizeM# #-} + {-# NOINLINE indexM# #-} + sizeM# m = unbox (inline sizeM m) + indexM# i# m = case inline indexM (I# i#) m of + (# I# i'#, a, hole #) -> (# i'#, a, hole #) + firstHoleM :: Sized a => TrieMap k a -> First (a, Hole k a) + firstHoleM m = inline extractHoleM m + lastHoleM :: Sized a => TrieMap k a -> Last (a, Hole k a) + lastHoleM m = inline extractHoleM m + + insertWithM f k a m = inline searchMC k m (assignM a) (assignM . f) + + assignM :: Sized a => a -> Hole k a -> TrieMap k a + clearM :: Sized a => Hole k a -> TrieMap k a + unifierM :: Sized a => k -> k -> a -> Maybe (Hole k a) + + fromListM f = L.foldl' (\ m (k, a) -> insertWithM (f a) k a m) emptyM fromAscListM = fromListM fromDistAscListM = fromAscListM const - - unifyM :: Sized a => k -> a -> k -> a -> Unified k a + unifierM k' k a = searchMC k' (singletonM k a) Just (\ _ _ -> Nothing) instance (TrieKey k, Sized a) => Sized (TrieMap k a) where - getSize# = sizeM + getSize# = sizeM# -singletonM' :: (TrieKey k, Sized a) => k -> Maybe a -> TrieMap k a -singletonM' k = maybe emptyM (singletonM k) +foldl1Empty :: a +foldl1Empty = error "Error: cannot call foldl1 on an empty map" +foldr1Empty :: a +foldr1Empty = error "Error: cannot call foldr1 on an empty map" + +{-# INLINE fillHoleM #-} +fillHoleM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a +fillHoleM = maybe clearM assignM + +{-# INLINE mapSearch #-} +mapSearch :: (hole -> hole') -> SearchCont hole a r -> SearchCont hole' a r +mapSearch f run nomatch match = run nomatch' match' where + nomatch' hole = nomatch (f hole) + match' a hole = match a (f hole) + +{-# INLINE unifyM #-} +unifyM :: (TrieKey k, Sized a) => k -> a -> k -> a -> Maybe (TrieMap k a) +unifyM k1 a1 k2 a2 = case unifierM k1 k2 a2 of + Nothing -> Nothing + Just hole -> Just $ inline assignM a1 hole + +insertWithM' :: (TrieKey k, Sized a) => (a -> a) -> k -> a -> Maybe (TrieMap k a) -> TrieMap k a +insertWithM' f k a = maybe (singletonM k a) (insertWithM f k a) + mapMaybeM' :: (TrieKey k, Sized b) => (a -> Maybe b) -> TrieMap k a -> Maybe (TrieMap k b) -mapMaybeM' f = guardNullM . mapMaybeM f +mapMaybeM' = guardNullM .: mapMaybeM mapEitherM' :: (TrieKey k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> TrieMap k a -> (# Maybe (TrieMap k b), Maybe (TrieMap k c) #) -mapEitherM' f = both guardNullM guardNullM (mapEitherM f) +mapEitherM' = both guardNullM guardNullM . mapEitherM mapEitherM'' :: (TrieKey k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> Maybe (TrieMap k a) -> (# Maybe (TrieMap k b), Maybe (TrieMap k c) #) -mapEitherM'' f = mapEitherMaybe (mapEitherM' f) +mapEitherM'' = mapEitherMaybe . mapEitherM' unionM' :: (TrieKey k, Sized a) => (a -> a -> Maybe a) -> TrieMap k a -> TrieMap k a -> Maybe (TrieMap k a) unionM' f m1 m2 = guardNullM (unionM f m1 m2) @@ -113,33 +169,30 @@ diffM' :: (TrieKey k, Sized a) => (a -> b -> Maybe a) -> TrieMap k a -> TrieMap k b -> Maybe (TrieMap k a) diffM' f m1 m2 = guardNullM (diffM f m1 m2) -beforeM' :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> Maybe (TrieMap k a) -beforeM' v hole = guardNullM (beforeM v hole) - -afterM' :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> Maybe (TrieMap k a) -afterM' v hole = guardNullM (afterM v hole) - -searchM' :: TrieKey k => k -> Maybe (TrieMap k a) -> (# Maybe a, Hole k a #) -searchM' k Nothing = (# Nothing, singleHoleM k #) -searchM' k (Just m) = searchM k m +{-# INLINE beforeMM #-} +beforeMM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a +beforeMM = maybe beforeM beforeWithM -extractHoleM' :: (TrieKey k, MonadPlus m, Sized a) => Maybe (TrieMap k a) -> m (a, Hole k a) -extractHoleM' Nothing = mzero -extractHoleM' (Just m) = extractHoleM m +{-# INLINE afterMM #-} +afterMM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a +afterMM = maybe afterM afterWithM -{-# INLINE assignM' #-} -assignM' :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> Maybe (TrieMap k a) -assignM' v@Just{} hole = Just (assignM v hole) -assignM' Nothing hole = guardNullM (assignM Nothing hole) +clearM' :: (TrieKey k, Sized a) => Hole k a -> Maybe (TrieMap k a) +clearM' hole = guardNullM (clearM hole) {-# 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 #) -> case f Nothing of - Nothing -> m - a -> assignM a hole - (# a, hole #) -> assignM (f a) hole +alterM f k m = searchMC k m g h where + g hole = case f Nothing of + Nothing -> m + Just a -> assignM a hole + h = fillHoleM . f . Just +{-# INLINE searchMC' #-} +searchMC' :: TrieKey k => k -> Maybe (TrieMap k a) -> (Hole k a -> r) -> (a -> Hole k a -> r) -> r +searchMC' k Nothing f _ = f (singleHoleM k) +searchMC' k (Just m) f g = searchMC k m f g + nullM :: TrieKey k => TrieMap k a -> Bool nullM m = case getSimpleM m of Null -> True @@ -159,11 +212,7 @@ (# x, y #) -> (# g1 x, g2 y #) elemsM :: TrieKey k => TrieMap k a -> [a] -elemsM m = build (\ f z -> foldrM f m z) - -insertWithM :: (TrieKey k, Sized a) => (a -> a -> a) -> k -> a -> TrieMap k a -> TrieMap k a -insertWithM f k a m = case searchM k m of - (# a', hole #) -> assignM (Just $ maybe a (f a) a') hole +elemsM m = build (\ f z -> foldr f z m) mapEitherMaybe :: (a -> (# Maybe b, Maybe c #)) -> Maybe a -> (# Maybe b, Maybe c #) mapEitherMaybe f (Just a) = f a @@ -189,6 +238,18 @@ subMaybe (<=) (Just a) (Just b) = a <= b subMaybe _ _ _ = False -indexFail :: a -> (# Int#, b, c #) +indexFail :: a -> (# Int, b, c #) indexFail _ = (# error err, error err, error err #) where err = "Error: not a valid index" + +{-# RULES + "extractHoleM/First" [0] extractHoleM = firstHoleM; + "extractHoleM/Last" [0] extractHoleM = lastHoleM; + "sizeM" [0] forall m . sizeM m = I# (sizeM# m); + "indexM" [0] forall i m . indexM i m = case indexM# (unbox i) m of { + (# i'#, a, m #) -> (# I# i'#, a, m #)}; + "getSimpleM/emptyM" getSimpleM emptyM = Null; + "getSimpleM/singletonM" forall k a . getSimpleM (singletonM k a) = Singleton a; + "toMaybe" forall f . toMaybe f = f; + "liftMaybe" forall m . liftMaybe m = m; + #-} diff --git a/Data/TrieMap/UnionMap.hs b/Data/TrieMap/UnionMap.hs --- a/Data/TrieMap/UnionMap.hs +++ b/Data/TrieMap/UnionMap.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns, MagicHash, CPP, BangPatterns #-} +{-# LANGUAGE UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns, MagicHash, CPP, BangPatterns, FlexibleInstances #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.UnionMap () where @@ -9,9 +9,9 @@ import Control.Applicative import Control.Monad -import Data.Foldable (foldr) -import Prelude hiding (foldr, (^)) -import GHC.Exts +import Data.Monoid +import Data.Foldable (Foldable(..)) +import Prelude hiding (foldr, foldr1, foldl, foldl1, (^)) (&) :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a m1 & m2 = guardNullM m1 ^ guardNullM m2 @@ -21,10 +21,10 @@ Nothing ^ Nothing = Empty Just m1 ^ Nothing = K1 m1 Nothing ^ Just m2 = K2 m2 -Just m1 ^ Just m2 = Union (sizeM m1 +# sizeM m2) m1 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 +union m1 m2 = Union (sizeM m1 + getSize m2) m1 m2 singletonL :: (TrieKey k1, TrieKey k2, Sized a) => k1 -> a -> TrieMap (Either k1 k2) a singletonL k a = K1 (singletonM k a) @@ -34,7 +34,7 @@ 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) + | Hole2 (Maybe (TrieMap k1 a)) (Hole k2 a) uView :: TrieMap (Either k1 k2) a -> UView k1 k2 a uView Empty = UView Nothing Nothing @@ -58,26 +58,38 @@ #define UVIEW uView -> UView +instance (TrieKey k1, TrieKey k2) => Foldable (UView k1 k2) where + {-# INLINE foldr #-} + {-# INLINE foldl #-} + {-# INLINE foldMap #-} + foldMap f (UView m1 m2) = foldMap (foldMap f) m1 `mappend` foldMap (foldMap f) m2 + foldr f z (UView m1 m2) = foldl (foldr f) (foldl (foldr f) z m2) m1 + foldl f z (UView m1 m2) = foldl (foldl f) (foldl (foldl f) z m1) m2 + +instance (TrieKey k1, TrieKey k2) => Foldable (TrieMap (Either k1 k2)) where + foldMap f m = foldMap f (uView m) + foldr f z m = foldr f z (uView m) + foldl f z m = foldl f z (uView m) + + foldl1 _ Empty = foldl1Empty + foldl1 f (K1 m1) = foldl1 f m1 + foldl1 f (K2 m2) = foldl1 f m2 + foldl1 f (Union _ m1 m2) = foldl f (foldl1 f m1) m2 + + foldr1 _ Empty = foldr1Empty + foldr1 f (K1 m1) = foldr1 f m1 + foldr1 f (K2 m2) = foldr1 f m2 + foldr1 f (Union _ m1 m2) = foldr f (foldr1 f m2) m1 + -- | @'TrieMap' ('Either' k1 k2) a@ is essentially a @(TrieMap k1 a, TrieMap k2 a)@, but -- specialized for the cases where one or both maps are empty. instance (TrieKey k1, TrieKey k2) => TrieKey (Either k1 k2) where - {-# SPECIALIZE instance TrieKey (Either () ()) #-} - {-# 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 - + {-# SPECIALIZE instance TrieKey (Either () ()) #-} data TrieMap (Either k1 k2) a = Empty | K1 (TrieMap k1 a) | K2 (TrieMap k2 a) - | Union Int# (TrieMap k1 a) (TrieMap k2 a) + | 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) @@ -91,27 +103,19 @@ mSimple :: TrieKey k => Maybe (TrieMap k a) -> Simple a mSimple = maybe mzero getSimpleM - sizeM Empty = 0# + sizeM Empty = 0 sizeM (K1 m1) = sizeM m1 sizeM (K2 m2) = sizeM m2 sizeM (Union s _ _) = s - lookupM (Left k) (UVIEW m1 _) = m1 >>= lookupM k - lookupM (Right k) (UVIEW _ m2) = m2 >>= lookupM k + lookupM (Left k) (UVIEW m1 _) = liftMaybe m1 >>= lookupM k + lookupM (Right k) (UVIEW _ m2) = liftMaybe m2 >>= lookupM k traverseM f (Union _ m1 m2) = union <$> traverseM f m1 <*> traverseM f m2 traverseM f (K1 m1) = K1 <$> traverseM f m1 traverseM f (K2 m2) = K2 <$> traverseM f m2 traverseM _ _ = pure Empty - 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 - - 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 - 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) @@ -138,47 +142,61 @@ isSubmapM (<=) (UVIEW m11 m12) (UVIEW m21 m22) = subMaybe (isSubmapM (<=)) m11 m21 && subMaybe (isSubmapM (<=)) m12 m22 + insertWithM f (Left k) a (UVIEW m1 m2) + = Just (insertWithM' f k a m1) ^ m2 + insertWithM f (Right k) a (UVIEW m1 m2) + = m1 ^ Just (insertWithM' f k a m2) fromListM f = onPair (&) (fromListM f) (fromListM f) . partEithers - fromAscListM f = onPair (&) (fromAscListM f) (fromAscListM f) . partEithers - fromDistAscListM = onPair (&) fromDistAscListM fromDistAscListM . partEithers singleHoleM = either (HoleX0 . singleHoleM) (Hole0X . singleHoleM) - beforeM a hole = case hView hole of - Hole1 h1 __ -> beforeM' a h1 ^ Nothing - Hole2 m1 h2 -> m1 ^ beforeM' a h2 + beforeM hole = case hView hole of + Hole1 h1 __ -> guardNullM (beforeM h1) ^ Nothing + Hole2 m1 h2 -> m1 ^ guardNullM (beforeM h2) + beforeWithM a hole = case hView hole of + Hole1 h1 __ -> K1 (beforeWithM a h1) + Hole2 m1 h2 -> m1 ^ Just (beforeWithM a h2) - afterM a hole = case hView hole of - Hole1 h1 m2 -> afterM' a h1 ^ m2 - Hole2 __ h2 -> Nothing ^ afterM' a h2 + afterM hole = case hView hole of + Hole1 h1 m2 -> guardNullM (afterM h1) ^ m2 + Hole2 __ h2 -> Nothing ^ guardNullM (afterM h2) + afterWithM a hole = case hView hole of + Hole1 h1 m2 -> Just (afterWithM a h1) ^ m2 + Hole2 __ h2 -> K2 (afterWithM a h2) - searchM (Left k) (UVIEW m1 m2) = onSnd (`hole1` m2) (searchM' k) m1 - searchM (Right k) (UVIEW m1 m2) = onSnd (hole2 m1) (searchM' k) m2 + searchMC (Left k) (UVIEW m1 m2) = mapSearch (`hole1` m2) (searchMC' k m1) + searchMC (Right k) (UVIEW m1 m2) = mapSearch (hole2 m1) (searchMC' k m2) - indexM i# (K1 m1) = onThird HoleX0 (indexM i#) m1 - indexM i# (K2 m2) = onThird Hole0X (indexM i#) m2 - indexM i# (Union _ m1 m2) - | i# <# s1# = onThird (`HoleX2` m2) (indexM i#) m1 - | otherwise = onThird (Hole1X m1) (indexM (i# -# s1#)) m2 - where !s1# = sizeM m1 + indexM i (K1 m1) = onThird HoleX0 (indexM i) m1 + indexM i (K2 m2) = onThird Hole0X (indexM i) m2 + indexM i (Union _ m1 m2) + | i < s1 = onThird (`HoleX2` m2) (indexM i) m1 + | otherwise = onThird (Hole1X m1) (indexM (i - s1)) m2 + where !s1 = sizeM m1 indexM _ _ = indexFail () - extractHoleM (UVIEW m1 m2) = (do - (v, h1) <- extractHoleM' m1 - return (v, hole1 h1 m2)) `mplus` (do - (v, h2) <- extractHoleM' m2 - return (v, hole2 m1 h2)) + extractHoleM (UVIEW !m1 !m2) = holes1 `mplus` holes2 where + holes1 = holes extractHoleM (`hole1` m2) m1 + holes2 = holes extractHoleM (hole2 m1) m2 + clearM hole = case hView hole of + Hole1 h1 m2 -> clearM' h1 ^ m2 + Hole2 m1 h2 -> m1 ^ clearM' h2 assignM v hole = case hView hole of - Hole1 h1 m2 -> assignM' v h1 ^ m2 - Hole2 m1 h2 -> m1 ^ assignM' v h2 + Hole1 h1 m2 -> Just (assignM v h1) ^ m2 + Hole2 m1 h2 -> m1 ^ Just (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) + unifierM (Left k') (Left k) a = HoleX0 <$> unifierM k' k a + unifierM (Left k') (Right k) a = Just $ HoleX2 (singleHoleM k') (singletonM k a) + unifierM (Right k') (Left k) a = Just $ Hole1X (singletonM k a) (singleHoleM k') + unifierM (Right k') (Right k) a = Hole0X <$> unifierM k' k a + +{-# INLINE holes #-} +holes :: (Functor m, Functor f, MonadPlus m) => (a -> m (f b)) -> (b -> c) -> Maybe a -> m (f c) +holes k f (Just a) = fmap f <$> k a +holes _ _ Nothing = mzero onPair :: (c -> d -> e) -> (a -> c) -> (b -> d) -> (a, b) -> e onPair f g h (a, b) = f (g a) (h b) diff --git a/Data/TrieMap/UnitMap.hs b/Data/TrieMap/UnitMap.hs --- a/Data/TrieMap/UnitMap.hs +++ b/Data/TrieMap/UnitMap.hs @@ -1,36 +1,37 @@ -{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash #-} +{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash, FlexibleInstances #-} -module Data.TrieMap.UnitMap where +module Data.TrieMap.UnitMap () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized -import Control.Applicative +import Data.Functor import Control.Monad import Data.Foldable import Data.Traversable import Data.Maybe -import Prelude hiding (foldr, foldl) +import Prelude hiding (foldr, foldl, foldr1, foldl1) +instance Foldable (TrieMap ()) where + foldMap f (Unit m) = foldMap f m + foldr f z (Unit m) = foldr f z m + foldl f z (Unit m) = foldl f z m + foldr1 f (Unit m) = foldr1 f m + foldl1 f (Unit m) = foldl1 f m + -- | @'TrieMap' () a@ is implemented as @'Maybe' a@. instance TrieKey () where - _ =? _ = True - _ `cmp` _ = EQ - - newtype TrieMap () a = Unit {getUnit :: Maybe a} + newtype TrieMap () a = Unit (Maybe a) data Hole () a = Hole emptyM = Unit Nothing - singletonM _ = Unit . Just + singletonM _ = single getSimpleM (Unit m) = maybe Null Singleton m - sizeM (Unit (Just a)) = getSize# a - sizeM _ = 0# - lookupM _ (Unit m) = m + sizeM (Unit m) = getSize m + lookupM _ (Unit m) = liftMaybe m 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 @@ -38,20 +39,30 @@ 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) + insertWithM f _ a (Unit m) = Unit (Just (maybe a f m)) + fromListM _ [] = emptyM + fromListM f ((_, v):xs) = single (foldl (\ v' -> f v' . snd) v xs) + singleHoleM _ = Hole - beforeM a _ = Unit a - afterM a _ = Unit a - searchM _ (Unit m) = (# m, Hole #) + beforeM _ = emptyM + afterM _ = emptyM + beforeWithM a _ = single a + afterWithM a _ = single a + + searchMC _ (Unit (Just v)) _ g = g v Hole + searchMC _ _ f _ = f Hole indexM i (Unit (Just v)) = (# i, v, Hole #) indexM _ _ = indexFail () - unifyM _ _ _ _ = Left Hole + unifierM _ _ _ = Nothing extractHoleM (Unit (Just v)) = return (v, Hole) extractHoleM _ = mzero - assignM v _ = Unit v + clearM _ = emptyM + assignM v _ = single v + +single :: a -> TrieMap () a +single = Unit . Just diff --git a/Data/TrieMap/Utils.hs b/Data/TrieMap/Utils.hs --- a/Data/TrieMap/Utils.hs +++ b/Data/TrieMap/Utils.hs @@ -1,9 +1,12 @@ {-# LANGUAGE Rank2Types, BangPatterns, MagicHash #-} -module Data.TrieMap.Utils (toVectorN, toVectorF) where +module Data.TrieMap.Utils where +import Data.Bits +import qualified Data.Foldable + import Data.Vector.Generic import Data.Vector.Generic.Mutable -import qualified Data.Foldable + import GHC.Exts {-# INLINE toVectorN #-} @@ -15,3 +18,33 @@ {-# 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)) + +{-# INLINE quoPow #-} +quoPow :: Int -> Int -> Int +n `quoPow` 1 = n +n `quoPow` 2 = n `shiftR` 1 +n `quoPow` 4 = n `shiftR` 2 +n `quoPow` 8 = n `shiftR` 3 +n `quoPow` 16 = n `shiftR` 4 +n `quoPow` 32 = n `shiftR` 5 +n `quoPow` 64 = n `shiftR` 6 +n `quoPow` k = n `quot` k + +{-# INLINE remPow #-} +remPow :: Int -> Int -> Int +n `remPow` k = if k .&. (k-1) == 0 then n .&. (k-1) else n `rem` k + +compl :: Word -> Word +compl (W# w#) = W# (not# w#) + +(.<<.) :: Word -> Int -> Word +W# w# .<<. I# i# = W# (uncheckedShiftL# w# i#) + +(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d +(f .: g) a b = f (g a b) + +{-# RULES + "or 0" forall w# . or# w# 0## = w#; + "0 or" forall w# . or# 0## w# = w#; + "plusAddr 0" forall a# . plusAddr# a# 0# = a#; + #-} diff --git a/Data/TrieMap/WordMap.hs b/Data/TrieMap/WordMap.hs new file mode 100644 --- /dev/null +++ b/Data/TrieMap/WordMap.hs @@ -0,0 +1,356 @@ +{-# LANGUAGE UnboxedTuples, BangPatterns, TypeFamilies, PatternGuards, MagicHash, CPP, NamedFieldPuns, FlexibleInstances #-} +{-# OPTIONS -funbox-strict-fields #-} +module Data.TrieMap.WordMap (SNode, WHole, TrieMap(WordMap), Hole(Hole), getWordMap, getHole) where + +import Data.TrieMap.TrieKey +import Data.TrieMap.Sized + +import Control.Exception (assert) +import Control.Applicative (Applicative(..), (<$>)) +import Control.Monad hiding (join) + +import Data.Bits +import Data.Foldable +import Data.Maybe hiding (mapMaybe) +import Data.Monoid +import Data.TrieMap.Utils + +import GHC.Exts + +import Prelude hiding (lookup, null, map, foldl, foldr, foldl1, foldr1) + +#include "MachDeps.h" +#define NIL SNode{node = Nil} +#define TIP(args) SNode{node = (Tip args)} +#define BIN(args) SNode{node = (Bin args)} + +type Nat = Word + +type Prefix = Word +type Mask = Word +type Key = Word +type Size = Int + +data Path a = Root + | LeftBin !Prefix !Mask (Path a) !(SNode a) + | RightBin !Prefix !Mask !(SNode a) (Path a) + +data SNode a = SNode {sz :: !Size, node :: (Node a)} +{-# ANN type SNode ForceSpecConstr #-} +data Node a = Nil | Tip !Key a | Bin !Prefix !Mask !(SNode a) !(SNode a) +{-# ANN type Node ForceSpecConstr #-} + +instance Sized (SNode a) where + getSize# SNode{sz} = unbox sz + +instance Sized a => Sized (Node a) where + getSize# t = unbox $ case t of + Nil -> 0 + Tip _ a -> getSize a + Bin _ _ l r -> getSize l + getSize r + +{-# INLINE sNode #-} +sNode :: Sized a => Node a -> SNode a +sNode !n = SNode (getSize n) n + +data WHole a = WHole !Key (Path a) + +{-# INLINE hole #-} +hole :: Key -> Path a -> Hole Word a +hole k path = Hole (WHole k path) + +#define HOLE(args) (Hole (WHole args)) + +-- | @'TrieMap' 'Word' a@ is based on "Data.IntMap". +instance TrieKey Word where + newtype TrieMap Word a = WordMap {getWordMap :: SNode a} + newtype Hole Word a = Hole {getHole :: WHole a} + emptyM = WordMap nil + singletonM k a = WordMap (singleton k a) + getSimpleM (WordMap (SNode _ n)) = case n of + Nil -> Null + Tip _ a -> Singleton a + _ -> NonSimple + sizeM (WordMap t) = getSize t + lookupM k (WordMap m) = lookup k m + traverseM f (WordMap m) = WordMap <$> traverse f m + fmapM f (WordMap m) = WordMap (map f m) + mapMaybeM f (WordMap m) = WordMap (mapMaybe f m) + mapEitherM f (WordMap m) = both WordMap WordMap (mapEither f) m + unionM f (WordMap m1) (WordMap m2) = WordMap (unionWith f m1 m2) + isectM f (WordMap m1) (WordMap m2) = WordMap (intersectionWith f m1 m2) + diffM f (WordMap m1) (WordMap m2) = WordMap (differenceWith f m1 m2) + isSubmapM (<=) (WordMap m1) (WordMap m2) = isSubmapOfBy (<=) m1 m2 + + singleHoleM k = hole k Root + beforeM HOLE(_ path) = WordMap (before nil path) + beforeWithM a HOLE(k path) = WordMap (before (singleton k a) path) + afterM HOLE(_ path) = WordMap (after nil path) + afterWithM a HOLE(k path) = WordMap (after (singleton k a) path) + + {-# INLINE searchMC #-} + searchMC !k (WordMap t) = mapSearch (hole k) (searchC k t) + indexM i (WordMap m) = indexT i m Root where + indexT !i TIP(kx x) path = (# i, x, hole kx path #) + indexT !i BIN(p m l r) path + | i < sl = indexT i l (LeftBin p m path r) + | otherwise = indexT (i - sl) r (RightBin p m l path) + where !sl = getSize l + indexT _ NIL _ = indexFail () + extractHoleM (WordMap m) = extractHole Root m where + extractHole _ (SNode _ Nil) = mzero + extractHole path TIP(kx x) = return (x, hole kx path) + extractHole path BIN(p m l r) = + extractHole (LeftBin p m path r) l `mplus` + extractHole (RightBin p m l path) r + clearM HOLE(_ path) = WordMap (assign nil path) + {-# INLINE assignM #-} + assignM v HOLE(kx path) = WordMap (assign (singleton kx v) path) + + {-# INLINE unifierM #-} + unifierM k' k a = Hole <$> unifier k' k a + +{-# INLINE searchC #-} +searchC :: Key -> SNode a -> SearchCont (Path a) a r +searchC !k t notfound found = seek Root t where + seek path t@BIN(p m l r) + | nomatch k p m = notfound (branchHole k p path t) + | zero k m + = seek (LeftBin p m path r) l + | otherwise + = seek (RightBin p m l path) r + seek path t@TIP(ky y) + | k == ky = found y path + | otherwise = notfound (branchHole k ky path t) + seek path NIL = notfound path + +before, after :: SNode a -> Path a -> SNode a +before !t Root = t +before !t (LeftBin _ _ path _) = before t path +before !t (RightBin p m l path) = before (bin p m l t) path +after !t Root = t +after !t (RightBin _ _ _ path) = after t path +after !t (LeftBin p m path r) = after (bin p m t r) path + +assign :: Sized a => SNode a -> Path a -> SNode a +assign NIL Root = nil +assign NIL (LeftBin _ _ path r) = assign' r path +assign NIL (RightBin _ _ l path) = assign' l path +assign t Root = t +assign t (LeftBin p m path r) = assign' (bin' p m t r) path +assign t (RightBin p m l path) = assign' (bin' p m l t) path + +assign' :: Sized a => SNode a -> Path a -> SNode a +assign' !t Root = t +assign' !t (LeftBin p m path r) = assign' (bin' p m t r) path +assign' !t (RightBin p m l path) = assign' (bin' p m l t) path + +branchHole :: Key -> Prefix -> Path a -> SNode 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 + +lookup :: Key -> SNode a -> Lookup a +lookup !k = look where + look BIN(_ m l r) = look (if zeroN k m then l else r) + look TIP(kx x) + | k == kx = some x + look _ = none + +singleton :: Sized a => Key -> a -> SNode a +singleton k a = sNode (Tip k a) + +singletonMaybe :: Sized a => Key -> Maybe a -> SNode a +singletonMaybe k = maybe nil (singleton k) + +traverse :: (Applicative f, Sized b) => (a -> f b) -> SNode a -> f (SNode b) +traverse f = trav where + trav NIL = pure nil + trav TIP(kx x) = singleton kx <$> f x + trav BIN(p m l r) = bin' p m <$> trav l <*> trav r + +instance Foldable SNode where + foldMap _ NIL = mempty + foldMap f TIP(_ x) = f x + foldMap f BIN(_ _ l r) = foldMap f l `mappend` foldMap f r + + foldr f z BIN(_ _ l r) = foldr f (foldr f z r) l + foldr f z TIP(_ x) = f x z + foldr _ z NIL = z + + foldl f z BIN(_ _ l r) = foldl f (foldl f z l) r + foldl f z TIP(_ x) = f z x + foldl _ z NIL = z + + foldr1 _ NIL = foldr1Empty + foldr1 _ TIP(_ x) = x + foldr1 f BIN(_ _ l r) = foldr f (foldr1 f r) l + + foldl1 _ NIL = foldl1Empty + foldl1 _ TIP(_ x) = x + foldl1 f BIN(_ _ l r) = foldl f (foldl1 f l) r + +instance Foldable (TrieMap Word) where + foldMap f (WordMap m) = foldMap f m + foldr f z (WordMap m) = foldr f z m + foldl f z (WordMap m) = foldl f z m + foldr1 f (WordMap m) = foldr1 f m + foldl1 f (WordMap m) = foldl1 f m + +map :: Sized b => (a -> b) -> SNode a -> SNode b +map f BIN(p m l r) = bin' p m (map f l) (map f r) +map f TIP(kx x) = singleton kx (f x) +map _ _ = nil + +mapMaybe :: Sized b => (a -> Maybe b) -> SNode a -> SNode b +mapMaybe f BIN(p m l r) = bin p m (mapMaybe f l) (mapMaybe f r) +mapMaybe f TIP(kx x) = singletonMaybe kx (f x) +mapMaybe _ _ = nil + +mapEither :: (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) -> + SNode a -> (# SNode b, SNode c #) +mapEither f BIN(p m l r) = both (bin p m lL) (bin p m lR) (mapEither f) r + where !(# lL, lR #) = mapEither f l +mapEither f TIP(kx x) = both (singletonMaybe kx) (singletonMaybe kx) f x +mapEither _ _ = (# nil, nil #) + +unionWith :: Sized a => (a -> a -> Maybe a) -> SNode a -> SNode a -> SNode a +unionWith f n1@(SNode _ t1) n2@(SNode _ t2) = case (t1, t2) of + (Nil, _) -> n2 + (_, Nil) -> n1 + (Tip k x, _) -> alter (maybe (Just x) (f x)) k n2 + (_, Tip k x) -> alter (maybe (Just x) (`f` x)) k n1 + (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) + | shorter m1 m2 -> union1 + | shorter m2 m1 -> union2 + | p1 == p2 -> bin p1 m1 (unionWith f l1 l2) (unionWith f r1 r2) + | otherwise -> join p1 n1 p2 n2 + where + union1 | nomatch p2 p1 m1 = join p1 n1 p2 n2 + | zero p2 m1 = bin p1 m1 (unionWith f l1 n2) r1 + | otherwise = bin p1 m1 l1 (unionWith f r1 n2) + + union2 | nomatch p1 p2 m2 = join p1 n1 p2 n2 + | zero p1 m2 = bin p2 m2 (unionWith f n1 l2) r2 + | otherwise = bin p2 m2 l2 (unionWith f n1 r2) + +{-# INLINE alter #-} +alter :: Sized a => (Maybe a -> Maybe a) -> Key -> SNode a -> SNode a +alter f k t = getWordMap $ alterM f k (WordMap t) + +intersectionWith :: Sized c => (a -> b -> Maybe c) -> SNode a -> SNode b -> SNode c +intersectionWith f n1@(SNode _ t1) n2@(SNode _ t2) = case (t1, t2) of + (Nil, _) -> nil + (_, Nil) -> nil + (Tip k x, _) -> option (lookup k n2) nil (singletonMaybe k . f x) + (_, Tip k y) -> option (lookup k n1) nil (singletonMaybe k . flip f y) + (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) + | shorter m1 m2 -> intersection1 + | shorter m2 m1 -> intersection2 + | p1 == p2 -> bin p1 m1 (intersectionWith f l1 l2) (intersectionWith f r1 r2) + | otherwise -> nil + where + intersection1 | nomatch p2 p1 m1 = nil + | zero p2 m1 = intersectionWith f l1 n2 + | otherwise = intersectionWith f r1 n2 + + intersection2 | nomatch p1 p2 m2 = nil + | zero p1 m2 = intersectionWith f n1 l2 + | otherwise = intersectionWith f n1 r2 + +differenceWith :: Sized a => (a -> b -> Maybe a) -> SNode a -> SNode b -> SNode a +differenceWith f n1@(SNode _ t1) n2@(SNode _ t2) = case (t1, t2) of + (Nil, _) -> nil + (_, Nil) -> n1 + (Tip k x, _) -> option (lookup k n2) n1 (singletonMaybe k . f x) + (_, Tip k y) -> alter (>>= flip f y) k n1 + (Bin p1 m1 l1 r1, Bin p2 m2 l2 r2) + | shorter m1 m2 -> difference1 + | shorter m2 m1 -> difference2 + | p1 == p2 -> bin p1 m1 (differenceWith f l1 l2) (differenceWith f r1 r2) + | otherwise -> n1 + where + difference1 | nomatch p2 p1 m1 = n1 + | zero p2 m1 = bin p1 m1 (differenceWith f l1 n2) r1 + | otherwise = bin p1 m1 l1 (differenceWith f r1 n2) + + difference2 | nomatch p1 p2 m2 = n1 + | zero p1 m2 = differenceWith f n1 l2 + | otherwise = differenceWith f n1 r2 + +isSubmapOfBy :: LEq a b -> LEq (SNode a) (SNode b) +isSubmapOfBy (<=) t1@BIN(p1 m1 l1 r1) BIN(p2 m2 l2 r2) + | shorter m1 m2 = False + | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy (<=) t1 l2 + else isSubmapOfBy (<=) t1 r2) + | otherwise = (p1==p2) && isSubmapOfBy (<=) l1 l2 && isSubmapOfBy (<=) r1 r2 +isSubmapOfBy _ BIN(_ _ _ _) _ = False +isSubmapOfBy (<=) TIP(k x) t2 = option (lookup k t2) False (x <=) +isSubmapOfBy _ NIL _ = True + +zero :: Key -> Mask -> Bool +zero i m + = i .&. m == 0 + +nomatch,match :: Key -> Prefix -> Mask -> Bool +nomatch i p m + = (mask i m) /= p + +match i p m + = (mask i m) == p + +zeroN :: Nat -> Nat -> Bool +zeroN i m = (i .&. m) == 0 + +mask :: Nat -> Nat -> Prefix +mask i m + = i .&. compl ((m-1) .|. m) + +shorter :: Mask -> Mask -> Bool +shorter m1 m2 + = m1 > m2 + +branchMask :: Prefix -> Prefix -> Mask +branchMask p1 p2 + = highestBitMask (p1 `xor` p2) + +highestBitMask :: Nat -> Nat +highestBitMask x0 + = case (x0 .|. shiftR x0 1) of + x1 -> case (x1 .|. shiftR x1 2) of + x2 -> case (x2 .|. shiftR x2 4) of + x3 -> case (x3 .|. shiftR x3 8) of + x4 -> case (x4 .|. shiftR x4 16) of + x5 -> case (x5 .|. shiftR x5 32) of -- for 64 bit platforms + x6 -> (x6 `xor` (shiftR x6 1)) + +{-# INLINE join #-} +join :: Prefix -> SNode a -> Prefix -> SNode a -> SNode a +join p1 t1 p2 t2 + | zero p1 m = bin' p m t1 t2 + | otherwise = bin' p m t2 t1 + where + m = branchMask p1 p2 + p = mask p1 m + +nil :: SNode a +nil = SNode 0 Nil + +bin :: Prefix -> Mask -> SNode a -> SNode a -> SNode a +bin p m l@(SNode sl tl) r@(SNode sr tr) = case (tl, tr) of + (Nil, _) -> r + (_, Nil) -> l + _ -> SNode (sl + sr) (Bin p m l r) + +bin' :: Prefix -> Mask -> SNode a -> SNode a -> SNode a +bin' p m l@SNode{sz=sl} r@SNode{sz=sr} = assert (nonempty l && nonempty r) $ SNode (sl + sr) (Bin p m l r) + where nonempty NIL = False + nonempty _ = True + +{-# INLINE unifier #-} +unifier :: Sized a => Key -> Key -> a -> Maybe (WHole a) +unifier k' k a + | k' == k = Nothing + | otherwise = Just (WHole k' $ branchHole k' k Root (singleton k a)) diff --git a/Data/TrieSet.hs b/Data/TrieSet.hs --- a/Data/TrieSet.hs +++ b/Data/TrieSet.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE UnboxedTuples #-} module Data.TrieSet ( -- * Set type TSet, @@ -29,7 +30,6 @@ map, mapMonotonic, -- * Fold - fold, foldl, foldr, -- * Min/Max @@ -42,6 +42,8 @@ minView, maxView, -- * Conversion + -- ** Map + mapSet, -- ** List elems, toList, @@ -52,19 +54,24 @@ fromDistinctAscList) where -import qualified Data.TrieMap as M import Data.TrieMap.Class +import Data.TrieMap.Class.Instances () +import Data.TrieMap.TrieKey +import Data.TrieMap.Representation.Class +import Data.TrieMap.Sized +import Data.TrieMap.Utils -import Control.Applicative hiding (empty) -import Control.Arrow +import Control.Monad.Ends import Data.Maybe -import Data.Monoid +import qualified Data.Foldable as F +import Data.Monoid (Monoid (..)) +import GHC.Exts import Prelude hiding (foldr, foldl, map, filter, null) instance TKey a => Eq (TSet a) where - s1 == s2 = s1 `isSubsetOf` s2 && size s1 == size s2 + s1 == s2 = size s1 == size s2 && s1 `isSubsetOf` s2 instance (TKey a, Ord a) => Ord (TSet a) where s1 `compare` s2 = elems s1 `compare` elems s2 @@ -76,98 +83,192 @@ mempty = empty mappend = union +-- | The empty 'TSet'. empty :: TKey a => TSet a -empty = TSet M.empty +empty = TSet emptyM +-- | Insert an element into the 'TSet'. insert :: TKey a => a -> TSet a -> TSet a -insert a (TSet s) = TSet (M.insert a () s) +insert a (TSet s) = TSet (insertWithM (const (Elem a)) (toRep a) (Elem a) s) +-- | Delete an element from the 'TSet'. delete :: TKey a => a -> TSet a -> TSet a -delete a (TSet s) = TSet (M.delete a s) +delete a (TSet s) = TSet (searchMC (toRep a) s clearM (const clearM)) +-- | /O(1)/. Create a singleton set. singleton :: TKey a => a -> TSet a -singleton a = insert a empty +singleton a = TSet (singletonM (toRep a) (Elem a)) +-- | The union of two 'TSet's, preferring the first set when +-- equal elements are encountered. union :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `union` TSet s2 = TSet (s1 `M.union` s2) +TSet s1 `union` TSet s2 = TSet (unionM (const . Just) s1 s2) +-- | The symmetric difference of two 'TSet's. symmetricDifference :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `symmetricDifference` TSet s2 = TSet (M.unionMaybeWith (\ _ _ -> Nothing) s1 s2) +TSet s1 `symmetricDifference` TSet s2 = TSet (unionM (\ _ _ -> Nothing) s1 s2) +-- | Difference of two 'TSet's. difference :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `difference` TSet s2 = TSet (s1 `M.difference` s2) +TSet s1 `difference` TSet s2 = TSet (diffM (\ _ _ -> Nothing) s1 s2) +-- | Intersection of two 'TSet's. Elements of the result come from the first set. intersection :: TKey a => TSet a -> TSet a -> TSet a -TSet s1 `intersection` TSet s2 = TSet (s1 `M.intersection` s2) +TSet s1 `intersection` TSet s2 = TSet (isectM (const . Just) s1 s2) +-- | Filter all elements that satisfy the predicate. filter :: TKey a => (a -> Bool) -> TSet a -> TSet a -filter p (TSet s) = TSet (M.filterWithKey (\ k _ -> p k) s) +filter p (TSet s) = TSet (mapMaybeM (\ (Elem a) -> if p a then Just (Elem a) else Nothing) s) +-- | Partition the set into two sets, one with all elements that satisfy +-- the predicate and one with all elements that don't satisfy the predicate. +-- See also 'split'. partition :: TKey a => (a -> Bool) -> TSet a -> (TSet a, TSet a) -partition p (TSet s) = (TSet *** TSet) (M.partitionWithKey (\ k _ -> p k) s) +partition p (TSet s) = case mapEitherM f s of + (# s1, s2 #) -> (TSet s1, TSet s2) + where f e@(Elem a) + | p a = (# Just e, Nothing #) + | otherwise = (# Nothing, Just e #) +-- | The expression (@'split' x set@) is a pair @(set1,set2)@ +-- where @set1@ comprises the elements of @set@ less than @x@ and @set2@ +-- comprises the elements of @set@ greater than @x@. split :: TKey a => a -> TSet a -> (TSet a, TSet a) split a s = case splitMember a s of (sL, _, sR) -> (sL, sR) +-- | Performs a 'split' but also returns whether the pivot +-- element was found in the original set. splitMember :: TKey a => a -> TSet a -> (TSet a, Bool, TSet a) -splitMember a (TSet s) = case M.splitLookup a s of - (sL, x, sR) -> (TSet sL, isJust x, TSet sR) +splitMember a (TSet s) = searchMC (toRep a) s nomatch match where + nomatch hole = (TSet (beforeM hole), False, TSet (afterM hole)) + match _ hole = (TSet (beforeM hole), True, TSet (afterM hole)) +-- | +-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. +-- +-- It's worth noting that the size of the result may be smaller if, +-- for some @(x,y)@, @x \/= y && f x == f y@ map :: (TKey a, TKey b) => (a -> b) -> TSet a -> TSet b -map f (TSet s) = TSet (M.mapKeys f s) +map f s = fromList [f x | x <- elems s] +-- | +-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is monotonic. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapMonotonic f s == map f s +-- > where ls = toList s mapMonotonic :: (TKey a, TKey b) => (a -> b) -> TSet a -> TSet b -mapMonotonic f (TSet s) = TSet (M.mapKeysMonotonic f s) +mapMonotonic f s = fromAscList [f x | x <- toAscList s] -fold, foldr :: TKey a => (a -> b -> b) -> b -> TSet a -> b -fold = foldr -foldr f z (TSet s) = M.foldrWithKey (const . f) z s +-- | Post-order fold. +foldr :: TKey a => (a -> b -> b) -> b -> TSet a -> b +foldr f z (TSet s) = F.foldr (flip $ F.foldr f) z s +-- | Pre-order fold. foldl :: TKey b => (a -> b -> a) -> a -> TSet b -> a -foldl f z (TSet s) = M.foldlWithKey (\ z a _ -> f z a) z s +foldl f z (TSet s) = F.foldl (F.foldl f) z s -findMin, findMax :: TKey a => TSet a -> a +-- | The minimal element of the set. +findMin :: TKey a => TSet a -> a findMin = fst . deleteFindMin + +-- | The maximal element of the set. +findMax :: TKey a => TSet a -> a findMax = fst . deleteFindMax -deleteMin, deleteMax :: TKey a => TSet a -> TSet a +-- | Delete the minimal element. +deleteMin :: TKey a => TSet a -> TSet a deleteMin s = maybe s snd (minView s) + +-- | Delete the maximal element. +deleteMax :: TKey a => TSet a -> TSet a deleteMax s = maybe s snd (maxView s) -deleteFindMin, deleteFindMax :: TKey a => TSet a -> (a, TSet a) +-- | Delete and find the minimal element. +-- +-- > 'deleteFindMin' set = ('findMin' set, 'deleteMin' set) +deleteFindMin :: TKey a => TSet a -> (a, TSet a) deleteFindMin = fromJust . minView + +-- | Delete and find the maximal element. +-- +-- > 'deleteFindMax' set = ('findMax' set, 'deleteMax' set) +deleteFindMax :: TKey a => TSet a -> (a, TSet a) deleteFindMax = fromJust . maxView -minView, maxView :: TKey a => TSet a -> Maybe (a, TSet a) -minView (TSet s) = (fst *** TSet) <$> M.minViewWithKey s -maxView (TSet s) = (fst *** TSet) <$> M.maxViewWithKey s +-- | Retrieves the minimal key of the set, and the set +-- stripped of that element, or 'Nothing' if passed an empty set. +minView :: TKey a => TSet a -> Maybe (a, TSet a) +minView (TSet s) = case getFirst (extractHoleM s) of + Nothing -> Nothing + Just (Elem a, hole) -> Just (a, TSet (afterM hole)) -elems, toList, toAscList :: TKey a => TSet a -> [a] -elems (TSet s) = M.keys s -toList = elems -toAscList = toList +-- | Retrieves the maximal key of the set, and the set +-- stripped of that element, or 'Nothing' if passed an empty set. +maxView :: TKey a => TSet a -> Maybe (a, TSet a) +maxView (TSet s) = case getLast (extractHoleM s) of + Nothing -> Nothing + Just (Elem a, hole) -> Just (a, TSet (beforeM hole)) -fromList, fromAscList, fromDistinctAscList :: TKey a => [a] -> TSet a -fromList xs = TSet (M.fromList [(x, ()) | x <- xs]) -fromAscList xs = TSet (M.fromAscList [(x, ()) | x <- xs]) -fromDistinctAscList xs = TSet (M.fromDistinctAscList [(x, ()) | x <- xs]) +{-# INLINE elems #-} +-- | See 'toAscList'. +elems :: TKey a => TSet a -> [a] +elems = toAscList +{-# INLINE toList #-} +-- | See 'toAscList'. +toList :: TKey a => TSet a -> [a] +toList = toAscList +{-# INLINE toAscList #-} +-- | Convert the set to an ascending list of elements. +toAscList :: TKey a => TSet a -> [a] +toAscList s = build (\ c n -> foldr c n s) +-- | Create a set from a list of elements. +fromList :: TKey a => [a] -> TSet a +fromList xs = TSet (fromListM const [(toRep x, Elem x) | x <- xs]) + +-- | Build a set from an ascending list in linear time. +-- /The precondition (input list is ascending) is not checked./ +fromAscList :: TKey a => [a] -> TSet a +fromAscList xs = TSet (fromAscListM const [(toRep x, Elem x) | x <- xs]) + +-- | /O(n)/. Build a set from an ascending list of distinct elements in linear time. +-- /The precondition (input list is strictly ascending) is not checked./ +fromDistinctAscList :: TKey a => [a] -> TSet a +fromDistinctAscList xs = TSet (fromDistAscListM [(toRep x, Elem x) | x <- xs]) + +-- | /O(1)/. Is this the empty set? null :: TKey a => TSet a -> Bool -null (TSet s) = M.null s +null (TSet s) = nullM s +-- | /O(1)/. The number of elements in the set. size :: TKey a => TSet a -> Int -size (TSet s) = M.size s +size (TSet s) = getSize s +-- | Is the element in the set? member :: TKey a => a -> TSet a -> Bool -member a (TSet s) = a `M.member` s +member a (TSet s) = option (lookupM (toRep a) s) False (const True) +-- | Is the element not in the set? notMember :: TKey a => a -> TSet a -> Bool -notMember a = not . member a +notMember = not .: member -isSubsetOf, isProperSubsetOf :: TKey a => TSet a -> TSet a -> Bool -TSet s1 `isSubsetOf` TSet s2 = M.isSubmapOfBy (\ _ _ -> True) s1 s2 +-- | Is this a subset? @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@. +isSubsetOf :: TKey a => TSet a -> TSet a -> Bool +TSet s1 `isSubsetOf` TSet s2 = isSubmapM (\ _ _ -> True) s1 s2 + +-- | Is this a proper subset? (ie. a subset but not equal). +isProperSubsetOf :: TKey a => TSet a -> TSet a -> Bool s1 `isProperSubsetOf` s2 = size s1 < size s2 && s1 `isSubsetOf` s2 +-- | See 'difference'. (\\) :: TKey a => TSet a -> TSet a -> TSet a (\\) = difference + +{-# INLINE [1] mapSet #-} +-- | Generate a 'TMap' by mapping on the elements of a 'TSet'. +mapSet :: TKey a => (a -> b) -> TSet a -> TMap a b +mapSet f (TSet s) = TMap (fmapM (\ (Elem a) -> Assoc a (f a)) s) diff --git a/Tests.hs b/Tests.hs --- a/Tests.hs +++ b/Tests.hs @@ -1,21 +1,51 @@ -{-# LANGUAGE TemplateHaskell, TypeFamilies, GADTs, ExistentialQuantification, CPP, ViewPatterns #-} --- module Tests where +{-# LANGUAGE TemplateHaskell, TypeFamilies, GADTs, ExistentialQuantification, CPP, UndecidableInstances #-} +module Tests (main) where + import Control.Monad -import Debug.Trace -import Data.TrieMap.Class -import Data.TrieMap.TrieKey -import Data.TrieMap.Sized +import Control.Applicative import qualified Data.TrieMap as T import qualified Data.Map as M +import Data.List (foldl') +import Data.TrieMap.Representation import Test.QuickCheck import Prelude hiding (null, lookup) +import Data.ByteString (ByteString, pack) +import qualified Data.ByteString as BS +type Val = [Int] -type Key = Integer -type Val = [Integer] +main :: IO () +main = quickCheckWith stdArgs{maxSuccess = 1000} (verify M.empty T.empty .&&. conjoin concretes) -main = quickCheckWith stdArgs{maxSize = 300, maxSuccess = 100} (verify M.empty T.empty) +data Key = A (ByteString, Int) | B Int ByteString | C [Bool] | D [Char] | E (Either String Double) deriving (Eq, Ord, Show) +data Key' = A' (ByteString, Int) | B' Int ByteString | C' [Bool] | D' [Char] | E' (Either String Double) deriving (Eq, Ord, Show) + +hash :: Key -> Int +hash (A (bs, i)) = BS.foldl' (\ i w -> i * 31 + fromIntegral w) i bs +hash (B i bs) = BS.foldl' (\ i w -> i * 61 + fromIntegral w) i bs +hash (C bs) = length bs +hash (D cs) = foldl' (\ i w -> i * 91 + fromEnum w) 0 cs +hash (E (Left cs)) = foldl' (\ i w -> i * 255 + fromEnum w) 0 cs +hash (E (Right i)) = fst (properFraction i) + +instance Arbitrary Key where + arbitrary = oneof [A <$> arbitrary, + B <$> arbitrary <*> arbitrary, + C <$> arbitrary, + D <$> arbitrary, + E <$> arbitrary] + +instance Arbitrary Key' where + arbitrary = oneof [A' <$> arbitrary, + B' <$> arbitrary <*> arbitrary, + C' <$> arbitrary, + D' <$> arbitrary, + E' <$> arbitrary] + +instance Arbitrary ByteString where + arbitrary = liftM pack arbitrary + instance Arbitrary Op where arbitrary = oneof [ liftM Op (liftM2 Insert arbitrary arbitrary), @@ -30,7 +60,9 @@ liftM Op (liftM Union recurse), liftM Op (liftM Isect recurse), liftM (Op . ElemAt) (arbitrary `suchThat` (>= 0)), - liftM (Op . DeleteAt) (arbitrary `suchThat` (>= 0))] + liftM (Op . DeleteAt) (arbitrary `suchThat` (>= 0)), + return (Op UpdateMin), + return (Op UpdateMax)] 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) @@ -56,6 +88,8 @@ show (Op (DeleteAt i)) = "DeleteAt " ++ show i show (Op (ElemAt i)) = "ElemAt " ++ show i show (Op (Isect ops)) = "Isect " ++ show ops + show (Op UpdateMax) = "UpdateMax" + show (Op UpdateMin) = "UpdateMin" data Operation r where Insert :: Key -> Val -> Operation () @@ -71,17 +105,20 @@ Isect :: [Op] -> Operation () DeleteAt :: Int -> Operation () ElemAt :: Int -> Operation (Maybe (Key, Val)) + UpdateMax :: Operation () + UpdateMin :: Operation () mapFunc :: Key -> Val -> Val -mapFunc = (:) +mapFunc ks xs = fromIntegral (hash ks):xs mapMaybeFunc :: Key -> Val -> Maybe Val -mapMaybeFunc k xs - | even k = Just (k:xs) +mapMaybeFunc ks xs + | even h = Just (fromIntegral h:xs) + where h = hash ks mapMaybeFunc _ _ = Nothing isectFunc :: Key -> Val -> Val -> Val -isectFunc ks xs ys = ks:xs ++ ys +isectFunc ks xs ys = [fromIntegral $ hash ks] ++ xs ++ ys generateMap :: M.Map Key Val -> [Op] -> M.Map Key Val generateMap = foldl (\ mm (Op op) -> snd (operateMap mm op)) @@ -105,6 +142,8 @@ operateMap m (DeleteAt i) = if M.null m then ((), m) else ((), M.deleteAt (i `mod` M.size m) m) operateMap m (ElemAt i) = if M.null m then (Nothing, m) else (Just $ M.elemAt (i `mod` M.size m) m, m) operateMap m (Isect ops) = ((), M.intersectionWithKey isectFunc m (generateMap M.empty ops)) +operateMap m (UpdateMin) = ((), M.updateMinWithKey mapMaybeFunc m) +operateMap m (UpdateMax) = ((), M.updateMaxWithKey mapMaybeFunc m) generateTMap :: T.TMap Key Val -> [Op] -> T.TMap Key Val generateTMap = foldl (\ m (Op op) -> snd (operateTMap m op)) @@ -131,6 +170,8 @@ operateTMap m (ElemAt i) | T.null m = (Nothing, m) | otherwise = (Just $ T.elemAt (i `mod` T.size m) m, m) +operateTMap m UpdateMin = ((), T.updateMinWithKey mapMaybeFunc m) +operateTMap m UpdateMax = ((), T.updateMaxWithKey mapMaybeFunc m) #define VERIFYOP(operation) verifyOp op@operation{} m tm = \ case (operateMap m op, operateTMap tm op) of \ @@ -150,9 +191,27 @@ VERIFYOP(DeleteAt) VERIFYOP(ElemAt) VERIFYOP(Isect) +VERIFYOP(UpdateMin) +VERIFYOP(UpdateMax) verify :: M.Map Key Val -> T.TMap Key Val -> [Op] -> Bool verify m tm (Op op:ops) = case verifyOp op m tm of Nothing -> False Just (m', tm') -> verify m' tm' ops verify _ _ [] = True + +concretes :: [Property] +concretes = [ + printTestCase "extending by a single 0 makes a difference" + (T.intersection (T.singleton (BS.pack [0]) "a") (T.singleton (BS.pack [0,0]) "b") == T.empty), + printTestCase "comparisons are correct" + (let input = [(BS.pack [0], "a"), (BS.pack [0,0,0,0,0], "a")] in T.assocs (T.fromList input) == input), + printTestCase "comparisons are correct" + (let input = [(BS.pack [0], "a"), (BS.pack [0,0,0,0,maxBound], "a")] in T.assocs (T.fromList input) == input), + printTestCase "genOptRepr is consistent with equality" (\ a b -> ((a :: Key') == b) == (toRep a == toRep b)), + printTestCase "deleteAt works for OrdMap" + (let input = [(1.4 :: Double, 'a'), (-4.0, 'b')] in T.assocs (T.deleteAt 0 (T.fromList input)) == [(1.4, 'a')]) + ] + +$(genRepr ''Key) +$(genOptRepr ''Key') diff --git a/TrieMap.cabal b/TrieMap.cabal --- a/TrieMap.cabal +++ b/TrieMap.cabal @@ -1,5 +1,5 @@ name: TrieMap -version: 2.0.3 +version: 3.0.0 cabal-version: >= 1.6 tested-with: GHC category: Algorithms @@ -9,6 +9,13 @@ The most recent release combines zipper-based ideas from recently proposed changes to Data.Map, as well as heavily optimized ByteString and Vector instances based on the vector package. + + Since version 2, unit tests and benchmarks have been taken much more seriously, and major optimizations + have been made. + + Compared to Data.Map and Data.Set, on e.g. @ByteString@s, TrieMaps support 6-12x faster @union@, + @intersection@, and @difference@ operations, 2x faster @lookup@, but 2x slower @toList@, and 4x slower @filter@. + Other operations are closely tied. license: BSD3 license-file: LICENSE author: Louis Wasserman @@ -32,10 +39,10 @@ Data.TrieMap.Representation, Data.TrieMap.Modifiers other-modules: + Control.Monad.Ends, Data.TrieMap.TrieKey, Data.TrieMap.Utils, Data.TrieMap.Sized, - Data.TrieMap.Applicative, Data.TrieMap.Representation.Class, Data.TrieMap.Representation.TH, Data.TrieMap.Representation.TH.Utils, @@ -48,7 +55,7 @@ Data.TrieMap.Representation.Instances.Foreign, Data.TrieMap.Representation.Instances.Vectors, Data.TrieMap.Representation.Instances.ByteString - Data.TrieMap.IntMap, + Data.TrieMap.WordMap, Data.TrieMap.OrdMap, Data.TrieMap.UnitMap, Data.TrieMap.ProdMap, @@ -58,5 +65,6 @@ Data.TrieMap.RadixTrie, Data.TrieMap.RadixTrie.Slice, Data.TrieMap.RadixTrie.Edge, + Data.TrieMap.RadixTrie.Label, Data.TrieMap.Class.Instances }