diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,16 @@ +Copyright (c) 2009, Louis Wasserman +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the +documentation and/or other materials provided with the distribution. + * The name of Louis Wasserman may not be used to endorse or promote products derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/TrieMap.cabal b/TrieMap.cabal new file mode 100644 --- /dev/null +++ b/TrieMap.cabal @@ -0,0 +1,33 @@ +name: TrieMap +version: 0.0.1.0 +license: BSD3 +license-file: LICENSE +maintainer: wasserman.louis@gmail.com +category: Data Structures +synopsis: An implementation of generalized tries with sophisticated map type inference. +description: Generalized trie implementation that automatically infers map types. Keys must implement the class 'TrieMap.Algebraic.Algebraic', which + declares that they are isomorphic to an /algebraic type/, + defined recursively as follows: + . + * () and 'Int' are algebraic types. + . + * If @'Ord' a@, then @'Ordered' a@ is an algebraic type. + . + * If @a,b@ are algebraic types, then so are @(a, b)@ and @Either a b@. + . + * If @a@ is algebraic, so is @[a]@. + . + This package exports almost the entire collection of methods available in Data.Map, and several new methods as well. In addition, each method will automatically infer the correct map type. + +build-type: Simple +build-depends: + base >= 4 && <= 5, containers +exposed-modules: + TrieMap + TrieMap.Algebraic +other-modules: + TrieMap.TrieAlgebraic + TrieMap.Applicative + TrieMap.Reflection + TrieMap.RadixTrie + TrieMap.MapTypes diff --git a/TrieMap.hs b/TrieMap.hs new file mode 100644 --- /dev/null +++ b/TrieMap.hs @@ -0,0 +1,923 @@ +{-# LANGUAGE FlexibleContexts, TypeFamilies #-} + +-- | We will use the following terminology: +-- +-- An /algebraic/ type is a type isomorphic to an algebraic type, as defined in the package description. This isomorphism is +-- declared via the type class 'Algebraic', where @'Alg' k@ is algebraic. It is assumed for purposes of ordering that +-- this isomorphism is order- and equality-preserving. We also require that if @k@ is algebraic, @'Alg' k ~ k@. +-- +-- These methods will automatically infer the correct type of a 'TrieMap' on any given argument. For example, +-- +-- @'fromList' [((\"alphabet\", 'Just' (0.2 :: 'Double'), 'True'), \"wxyz\")]@ +-- +-- returns a variable of type +-- +-- @'TrieMap' ('String', 'Double', 'Bool') ('RadixTrie' 'Int' 'Data.IntMap.IntMap' \``ProdMap`\` 'UnionMap' 'Maybe' ('Data.Map.Map' 'Double') \``ProdMap`\` 'UnionMap' 'Maybe' 'Maybe') 'String'@ +-- +-- The inference was done entirely automatically. Note also: +-- +-- * @'Alg' 'Char' ~ 'Int'@: the 'Algebraic' instance for 'Char' maps characters to their ASCII representations, so an 'IntMap' can be used. +-- +-- * @'Alg' ('Maybe' a) ~ 'Either' () ('Alg' a)@; a 'TrieMap' on a 'Maybe' key type simply gets a space for one extra (possible) value. +-- +-- * @'Alg' 'Double' ~ 'Ordered' 'Double'@; the 'Algebraic' instance for 'Double' tells "TrieMap" to just use a regular 'Data.Map.Map' +-- and the default ordering for 'Double's. +-- +-- * @'Alg' 'Bool' ~ 'Either' () ()@, so a 'TrieMap' on a 'Bool' takes the form of -- essentially -- a pair of 'Maybe's. +-- +-- * @'Alg' (a, b, c) ~ ('Alg' a, ('Alg' b, 'Alg' c))@, so tuple types get handled by a sequence of map products. +-- +-- (If you plan to use these maps in type arguments, it is strongly suggested that you either reproduce the context +-- @('Algebraic' k, 'TrieKey' ('Alg' k) m) => TrieMap k m a@, or you create a type alias!) +-- + +-- The following is a general attempt to describe the runtime of operations supported by 'TrieMap's. +-- +-- * Lookup operations take /O(log n)/ for 'Ordered' keys, /O(max(log n, W))/ for 'Int' keys, /O(l)/ times lookup cost for @k@ +-- for keys of type @[k]@, and otherwise will take @O(1)@ over the total cost of their components. +-- +-- * Insertion operations take roughly the same asymptotic time as lookup operations. +-- +-- * Traversal operations take /O(n)/ for all map types, with obviously greater overhead for use of specialized +-- 'Applicative' functors. +-- +-- * Set operations (union, intersection, difference) take /O(m + n)/ in all cases. + +module TrieMap ( + -- * Map type + TrieMap, + TrieKey, + Algebraic (..), + -- * Map instances + ProdMap, UnionMap, RadixTrie, + -- * Operators + (!), + (\\), + -- * Query + null, + size, + member, + notMember, + lookup, + find, + findWithDefault, + -- * Construction + empty, + singleton, + -- * Insertion + insert, + insertWith, + insertWithKey, + insertLookupWithKey, + -- * Delete/Update + delete, + update, + updateWithKey, + updateLookupWithKey, + alter, + alterLookup, + -- * Combine + -- ** Union/Symmetric Difference + union, + unionWith, + unionWithKey, + unions, + unionsWith, + unionsWithKey, + unionMaybeWith, + unionMaybeWithKey, + symDifference, + -- ** Intersection + intersection, + intersectionWith, + intersectionWithKey, + intersectionMaybeWith, + intersectionMaybeWithKey, + -- ** Difference + difference, + differenceWith, + differenceWithKey, + -- * Traversal + -- ** Map + map, + mapWithKey, + mapApp, + mapAppWithKey, + mapMaybe, + mapMaybeWithKey, + mapEither, + mapEitherWithKey, + mapKeys, + mapKeysWith, + mapKeysMonotonic, + -- ** Fold + fold, + foldWithKey, + -- * Conversion + elems, + keys, + assocs, + -- ** Lists + fromList, + fromListWith, + fromListWithKey, + -- ** Ordered lists + fromAscList, + fromAscListWith, + fromAscListWithKey, + fromDistinctAscList, + -- * Filter + filter, + filterWithKey, + partition, + partitionWithKey, + split, + splitLookup, + -- * Submap + isSubmapOf, + isSubmapOfBy, + -- * Min/Max + findMin, + getMin, + findMax, + getMax, + deleteMin, + deleteMax, + deleteFindMin, + deleteFindMax, + updateMin, + updateMax, + updateMinWithKey, + updateMaxWithKey, + minView, + maxView, + minViewWithKey, + maxViewWithKey) where + +import Control.Monad +import Data.Monoid +import Data.Traversable +import TrieMap.Applicative +import TrieMap.Algebraic +import TrieMap.TrieAlgebraic +import TrieMap.RadixTrie +import TrieMap.Reflection +import Control.Applicative hiding (Alternative(..)) +import Data.Maybe hiding (mapMaybe) +import Data.Map (Map) +import Data.IntMap (IntMap) +import Data.Foldable hiding (fold, find) +import GHC.Exts +-- import TrieMap.FixPoint +-- import TrieMap.FixPoint.Algebraic +-- import TrieMap.Reflection +import Prelude hiding (lookup, foldr, null, filter, foldl, map) +import qualified Prelude as Prelude + +-- | A 'TrieMap' is a size-tracking wrapper around a generalized trie map. +data TrieMap k m a = TrieMap {sizeMap :: Int, trieMap :: m a} + +instance (Eq k, Eq a, Algebraic k, TrieKey (Alg k) m) => Eq (TrieMap k m a) where + (==) = (==) `on` assocs + +instance (Ord k, Ord a, Algebraic k, TrieKey (Alg k) m) => Ord (TrieMap k m a) where + compare = compare `on` assocs + +instance (Show k, Show a, Algebraic k, TrieKey (Alg k) m) => Show (TrieMap k m a) where + show m = "fromList " ++ show (assocs m) + +instance (Algebraic k, Algebraic a, TrieKey (Alg k) m) => Algebraic (TrieMap k m a) where + type Alg (TrieMap k m a) = (Int, [(Alg k, Alg a)]) + toAlg (TrieMap n m) = (n, build (\ c n -> foldWithKeyAlg (\ k a -> c (k, toAlg a)) n m)) + fromAlg (n, xs) = TrieMap n $ fromDistAscListAlg [(k, fromAlg a) | (k, a) <- xs] + +instance Functor m => Functor (TrieMap k m) where + fmap f (TrieMap n m) = TrieMap n (fmap f m) + +instance Foldable m => Foldable (TrieMap k m) where + foldr f z = foldr f z . trieMap + foldl f z = foldl f z . trieMap + foldMap f = foldMap f . trieMap + +instance Traversable m => Traversable (TrieMap k m) where + traverse f (TrieMap n m) = TrieMap n <$> traverse f m + +instance (Algebraic k, TrieKey (Alg k) m) => Monoid (TrieMap k m a) where + mempty = empty + mappend = union + +mkTrieMap :: (Algebraic k, TrieKey (Alg k) m) => m a -> TrieMap k m a +mkTrieMap m = TrieMap (sizeAlg m) m + +-- | Lookup the value of a key in the map. +-- +-- The function will return the corresponding value as @('Just' value)@, +-- or 'Nothing' if the key isn't in the map. +lookup :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> Maybe a +lookup k = lookupAlg (toAlg k) . trieMap + +-- | Is the key a member of the map? See also 'notMember'. +-- +-- > member 5 (fromList [(5,'a'), (3,'b')]) == True +-- > member 1 (fromList [(5,'a'), (3,'b')]) == False +member :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> Bool +member = isJust .: lookup + +-- | Is the key not a member of the map? See also 'member'. +-- +-- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False +-- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True +notMember :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> Bool +notMember = not .: member + +-- | Find the value at a key. +-- Calls 'error' when the element can not be found. + +find :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> a +find = findWithDefault $ error "TrieMap.find: element not in the map" + +-- | 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. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' +findWithDefault :: (Algebraic k, TrieKey (Alg k) m) => a -> k -> TrieMap k m a -> a +findWithDefault v = fromMaybe v .: lookup + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +singleton :: (Algebraic k, TrieKey (Alg k) m) => k -> a -> TrieMap k m a +singleton k v = TrieMap 1 (insertAlg (toAlg k) v emptyAlg) + +-- | Find the value at a key. +-- Calls 'error' when the element can not be found. +-- +-- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map +-- > fromList [(5,'a'), (3,'b')] ! 5 == 'a' +(!) :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> k -> a +m ! k = fromMaybe (error "element not in the map") (lookup k m) + +empty :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a +empty = TrieMap 0 emptyAlg + +-- | Check if the specified map is empty. +null :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Bool +null = nullAlg . trieMap + +-- | Returns the size of the specified map. +size :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Int +size = sizeMap + +-- | Build a map from a list of key\/value pairs. See also 'fromAscList'. +-- If the list contains more than one value for the same key, the last value +-- for the key is retained. +-- +-- > fromList [] == empty +-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")] +-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")] +fromList :: (Algebraic k, TrieKey (Alg k) m) => [(k, a)] -> TrieMap k m a +fromList = fromListWith const + +-- | Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] +-- > fromListWith (++) [] == empty +fromListWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> a -> a) -> [(k, a)] -> TrieMap k m a +fromListWith = fromListWithKey . const + +-- | Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] +-- > fromListWithKey f [] == empty +fromListWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> a) -> [(k, a)] -> TrieMap k m a +fromListWithKey f xs = mkTrieMap $ fromListAlg (f . fromAlg) [(toAlg k, v) | (k, v) <- xs] + +-- | /O(n)/. Build a map from an ascending list in linear time. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +fromAscList :: (Algebraic k, TrieKey (Alg k) m) => [(k, a)] -> TrieMap k m a +fromAscList = fromAscListWith const + +-- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] +fromAscListWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> a -> a) -> [(k, a)] -> TrieMap k m a +fromAscListWith = fromAscListWithKey . const + +-- | /O(n)/. Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +fromAscListWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> a) -> [(k, a)] -> TrieMap k m a +fromAscListWithKey f xs = mkTrieMap $ fromAscListAlg (f . fromAlg) [(toAlg k, v) | (k, v) <- xs] + +-- | /O(n)/. Build a map from an ascending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +fromDistinctAscList :: (Algebraic k, TrieKey (Alg k) m) => [(k, a)] -> TrieMap k m a +fromDistinctAscList xs = TrieMap (length xs) $ fromDistAscListAlg [(toAlg k, v) | (k, v) <- xs] + +-- | Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value. 'insert' is equivalent to +-- @'insertWith' 'const'@. +-- +-- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')] +-- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')] +-- > insert 5 'x' empty == singleton 5 'x' +insert :: (Algebraic k, TrieKey (Alg k) m) => k -> a -> TrieMap k m a -> TrieMap k m a +insert = insertWith const + +-- | Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key, f new_value old_value)@. +-- +-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] +-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" +insertWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> a -> a) -> k -> a -> TrieMap k m a -> TrieMap k m a +insertWith = insertWithKey . const + +-- | Insert with a function, combining key, new value and old value. +-- @'insertWithKey' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key,f key new_value old_value)@. +-- Note that the key passed to f is the same key passed to 'insertWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] +-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" +insertWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> a) -> k -> a -> TrieMap k m a -> TrieMap k m a +insertWithKey f k = snd .: insertLookupWithKey f k + +-- | Combines insert operation with old value retrieval. +-- The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- and the second element equal to (@'insertWithKey' f k x map@). +insertLookupWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> a) -> k -> a -> TrieMap k m a -> (Maybe a, TrieMap k m a) +insertLookupWithKey f k v (TrieMap n m) = case alterLookupAlg (\ v' -> (v', Just $ maybe v (f k v) v')) (toAlg k) m of + (old, m') -> (old, TrieMap (if isJust old then n else n + 1) m') + +-- | The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +update :: (Algebraic k, TrieKey (Alg k) m) => (a -> Maybe a) -> k -> TrieMap k m a -> TrieMap k m a +update = updateWithKey . const + +-- | The expression (@'updateWithKey' f k map@) updates the +-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', +-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound +-- to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +updateWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Maybe a) -> k -> TrieMap k m a -> TrieMap k m a +updateWithKey f = snd .: updateLookupWithKey f + +-- | Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") +updateLookupWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Maybe a) -> k -> TrieMap k m a -> (Maybe a, TrieMap k m a) +updateLookupWithKey f k (TrieMap n m) = + case alterLookupAlg (\ v -> let v' = v >>= f k in ((isNothing v' && isJust v, maybe v Just v'), v')) (toAlg k) m of + ((del, res), m') -> (res, TrieMap (if del then n - 1 else n) m') + +-- | Delete a key and its value from the map. When the key is not +-- a member of the map, the original map is returned. +-- +-- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > delete 5 empty == empty +-- +-- 'delete' is equivalent to @'alter' ('const' 'Nothing')@. +delete :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> TrieMap k m a +delete = alter (const Nothing) + +-- | The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. +-- 'alter' can be used to insert, delete, or update a value in a 'Map'. +-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] +alter :: (Algebraic k, TrieKey (Alg k) m) => (Maybe a -> Maybe a) -> k -> TrieMap k m a -> TrieMap k m a +alter f k = snd . alterLookup f k + +-- | The expression (@'alterLookup' f k map@) alters the value @x@ at @k@, or absence thereof, and returns the old value. +-- 'alterLookup' can be used to insert, delete, or update a value in a 'Map'. +-- +-- In short : @alterLookup f k m = (lookup k m, alter f k m)@. +alterLookup :: (Algebraic k, TrieKey (Alg k) m) => (Maybe a -> Maybe a) -> k -> TrieMap k m a -> (Maybe a, TrieMap k m a) +alterLookup f k (TrieMap n m) = case alterLookupAlg g (toAlg k) m of + ((old, delta), m') -> (old, TrieMap (n + delta) m') + where g v = let fv = f v in ((v, just1 fv - just1 v), fv) + just1 = maybe 0 (const 1) + +-- | /O(n)/. Map a function over all values in the map. +-- +-- > let f key x = (show key) ++ ":" ++ x +-- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")] +mapWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> b) -> TrieMap k m a -> TrieMap k m b +mapWithKey f = unId . mapAppWithKey (Id .: f) + +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] +map :: (Algebraic k, TrieKey (Alg k) m) => (a -> b) -> TrieMap k m a -> TrieMap k m b +map = fmap + +-- | Essentially equivalent to 'traverse' with a function that takes both the key and the value as arguments. +mapAppWithKey :: (Algebraic k, TrieKey (Alg k) m, Applicative f) => + (k -> a -> f b) -> TrieMap k m a -> f (TrieMap k m b) +mapAppWithKey f (TrieMap n m) = TrieMap n <$> mapAppAlg (f . fromAlg) m + +-- | Equivalent to 'traverse'. +mapApp :: (Algebraic k, TrieKey (Alg k) m, Applicative f) => (a -> f b) -> TrieMap k m a -> f (TrieMap k m b) +mapApp = traverse + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +-- +-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing +-- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" +mapMaybeWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Maybe b) -> TrieMap k m a -> TrieMap k m b +mapMaybeWithKey f = mkTrieMap . mapMaybeAlg (f . fromAlg) . trieMap + +-- | /O(n)/. Map values and collect the 'Just' results. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a" +mapMaybe :: (Algebraic k, TrieKey (Alg k) m) => (a -> Maybe b) -> TrieMap k m a -> TrieMap k m b +mapMaybe = mapMaybeWithKey . const + +-- | /O(n)/. Map values and separate the 'Left' and 'Right' results. +-- +-- > let f a = if a < "c" then Left a else Right a +-- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")]) +-- > +-- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +mapEither :: (Algebraic k, TrieKey (Alg k) m) => (a -> Either b c) -> TrieMap k m a -> (TrieMap k m b, TrieMap k m c) +mapEither = mapEitherWithKey . const + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- +-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) +-- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")]) +-- > +-- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) +mapEitherWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Either b c) -> TrieMap k m a -> (TrieMap k m b, TrieMap k m c) +mapEitherWithKey f (TrieMap _ m) = (mkTrieMap mL, mkTrieMap mR) + where (mL, mR) = mapEitherAlg (f . fromAlg) m + +-- | +-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the value at the smallest of +-- these keys is retained. +-- +-- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")] +-- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c" +-- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c" +mapKeys :: (Algebraic k1, Algebraic k2, TrieKey (Alg k1) m1, TrieKey (Alg k2) m2) => + (k1 -> k2) -> TrieMap k1 m1 a -> TrieMap k2 m2 a +mapKeys = mapKeysWith const + +-- | +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" +mapKeysWith :: (Algebraic k1, Algebraic k2, TrieKey (Alg k1) m1, TrieKey (Alg k2) m2) => + (a -> a -> a) -> (k1 -> k2) -> TrieMap k1 m1 a -> TrieMap k2 m2 a +mapKeysWith f g m = fromListWith f [(g k, v) | (k, v) <- assocs m] + +-- | /O(n)/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has better performance than 'mapKeys'. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] +-- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True +-- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False +mapKeysMonotonic :: (Algebraic k1, Algebraic k2, TrieKey (Alg k1) m1, TrieKey (Alg k2) m2) => + (k1 -> k2) -> TrieMap k1 m1 a -> TrieMap k2 m2 a +mapKeysMonotonic f (TrieMap n m) = TrieMap n $ fromDistAscListAlg [(toAlg (f (fromAlg k)), v) | (k, v) <- assocsAlg m] + +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +filterWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Bool) -> TrieMap k m a -> TrieMap k m a +filterWithKey p = mapMaybeWithKey (\ k v -> if p k v then Just v else Nothing) + +-- | /O(n)/. Filter all values that satisfy the predicate. +-- +-- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty +-- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty +filter :: (Algebraic k, TrieKey (Alg k) m) => (a -> Bool) -> TrieMap k m a -> TrieMap k m a +filter = filterWithKey . const + +-- | /O(n)/. Partition the map according to a predicate. The first +-- map contains all elements that satisfy the predicate, the second all +-- elements that fail the predicate. +-- +-- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") +-- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) +partition :: (Algebraic k, TrieKey (Alg k) m) => (a -> Bool) -> TrieMap k m a -> (TrieMap k m a, TrieMap k m a) +partition = partitionWithKey . const + +-- | /O(n)/. Partition the map according to a predicate. The first +-- map contains all elements that satisfy the predicate, the second all +-- elements that fail the predicate. +-- +-- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b") +-- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) +partitionWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Bool) -> TrieMap k m a -> (TrieMap k m a, TrieMap k m a) +partitionWithKey p = mapEitherWithKey (\ k v -> (if p k v then Left else Right) v) + +{-# INLINE assocs #-} +-- | /O(n)/. Return all key\/value pairs in the map in ascending key order. +-- +-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > assocs empty == [] +assocs :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> [(k, a)] +assocs m = build (\ c n -> foldWithKey (curry c) n m) + +-- | /O(n)/. Return all keys of the map in ascending order. +-- +-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] +-- > keys empty == [] +keys :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> [k] +keys m = Prelude.map fst (assocs m) + +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- +-- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] +-- > elems empty == [] +elems :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> [a] +elems = toList + +-- | /O(n)/. Fold the values in the map, such that +-- @'fold' f z == 'Prelude.foldr' f z . 'elems'@. +-- For example, +-- +-- > elems map = fold (:) [] map +-- +-- > let f a len = len + (length a) +-- > fold f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +fold :: TrieKey k m => (a -> b -> b) -> b -> TrieMap k m a -> b +fold = foldr + +-- | /O(n)/. Fold the keys and values in the map, such that +-- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'assocs'@. +-- For example, +-- +-- > keys map = foldWithKey (\k x ks -> k:ks) [] map +-- +-- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" +-- > foldWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)" +foldWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> b -> b) -> b -> TrieMap k m a -> b +foldWithKey f z = foldWithKeyAlg (f . fromAlg) z . trieMap + +-- | /O(n+m)/. Union with a combining function that may discard some elements. +unionMaybeWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> Maybe a) -> TrieMap k m a -> TrieMap k m a -> TrieMap k m a +unionMaybeWithKey f = mkTrieMap .: unionMaybeAlg (f . fromAlg) `on` trieMap + +-- | /O(n+m)/. +-- Union with a combining function. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] +unionWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> a) -> TrieMap k m a -> TrieMap k m a -> TrieMap k m a +unionWithKey f = unionMaybeWithKey (\ k x y -> Just (f k x y)) + +-- | /O(n+m)/. Union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] +unionWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> a -> a) -> TrieMap k m a -> TrieMap k m a -> TrieMap k m a +unionWith = unionWithKey . const + +-- | /O(n+m)/. Union with a combining function that may discard some elements. +unionMaybeWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> a -> Maybe a) -> TrieMap k m a -> TrieMap k m a -> TrieMap k m a +unionMaybeWith = unionMaybeWithKey . const + +-- | /O(n+m)/. +-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@. +-- It prefers @t1@ when duplicate keys are encountered, +-- i.e. (@'union' == 'unionWith' 'const'@). +-- +-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] +union :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m a -> TrieMap k m a +union = unionWith const + +unions :: (Algebraic k, TrieKey (Alg k) m) => [TrieMap k m a] -> TrieMap k m a +unions = unionsWith const + +unionsWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> a -> a) -> [TrieMap k m a] -> TrieMap k m a +unionsWith = unionsWithKey . const + +unionsWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> a -> a) -> [TrieMap k m a] -> TrieMap k m a +unionsWithKey f = mkTrieMap . foldl' (unionMaybeAlg (\ k x y -> Just (f (fromAlg k) x y))) emptyAlg + . Prelude.map trieMap + +-- | O(n+m). Symmetric difference. Equivalent to @'unionMaybeWith' (\ _ _ -> Nothing)@. +symDifference :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m a -> TrieMap k m a +symDifference = unionMaybeWith (\ _ _ -> Nothing) + +-- | /O(n+m)/. Intersection of two maps with a combining function that may discard some elements. +intersectionMaybeWithKey :: (Algebraic k, TrieKey (Alg k) m) => + (k -> a -> b -> Maybe c) -> TrieMap k m a -> TrieMap k m b -> TrieMap k m c +intersectionMaybeWithKey f (TrieMap _ m1) (TrieMap _ m2) = mkTrieMap $ intersectAlg (f . fromAlg) m1 m2 + +-- | /O(n+m)/. Intersection with a combining function. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" +intersectionWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> b -> c) -> TrieMap k m a -> TrieMap k m b -> TrieMap k m c +intersectionWithKey f = intersectionMaybeWithKey (\ k x y -> Just (f k x y)) + +-- | /O(n+m)/. Intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" +intersectionWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> b -> c) -> TrieMap k m a -> TrieMap k m b -> TrieMap k m c +intersectionWith f = intersectionMaybeWith (Just .: f) + +-- | /O(n+m)/. Intersection of two maps with a combining function that may discard some elements. +intersectionMaybeWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> b -> Maybe c) -> TrieMap k m a -> TrieMap k m b -> TrieMap k m c +intersectionMaybeWith = intersectionMaybeWithKey . const + +-- | /O(n+m)/. Intersection of two maps. +-- Return data in the first map for the keys existing in both maps. +-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). +-- +-- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" +intersection :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m b -> TrieMap k m a +intersection = intersectionWith const + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) +-- > == singleton 3 "3:b|B" +differenceWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> b -> Maybe a) -> TrieMap k m a -> TrieMap k m b -> TrieMap k m a +differenceWithKey f (TrieMap _ m1) (TrieMap _ m2) = mkTrieMap $ differenceAlg (f . fromAlg) m1 m2 + +-- | /O(n+m)/. Difference with a combining function. +-- When two equal keys are +-- encountered, the combining function is applied to the values of these keys. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) +-- > == singleton 3 "b:B" +differenceWith :: (Algebraic k, TrieKey (Alg k) m) => (a -> b -> Maybe a) -> TrieMap k m a -> TrieMap k m b -> TrieMap k m a +differenceWith f (TrieMap _ m1) (TrieMap _ m2) = mkTrieMap $ differenceAlg (const f) m1 m2 + +-- | /O(n+m)/. Difference of two maps. +-- Return elements of the first map not existing in the second map. +-- +-- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" +difference :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m b -> TrieMap k m a +difference = differenceWith (\ _ _ -> Nothing) + +-- | Same as 'difference'. +(\\) :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m b -> TrieMap k m a +(\\) = difference + +-- | The minimal key of the map. Calls 'error' if the map is empty. +-- +-- > findMin (fromList [(5,"a"), (3,"b")]) == (3,"b") +-- > findMin empty Error: empty map has no minimal element +findMin :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> (k, a) +findMin = fromMaybe (error "empty map has no minimal element") . getMin + +-- | The minimal key of the map, if any. Returns 'Nothing' if the map is empty. +-- +-- > getMin (fromList [(5,"a"), (3,"b")]) == Just (3,"b") +-- > getMin empty == Nothing +getMin :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Maybe (k, a) +getMin = fst <.> minViewWithKey + +-- | The maximal key of the map. Calls 'error' is the map is empty. +-- +-- > findMax (fromList [(5,"a"), (3,"b")]) == (5,"a") +-- > findMax empty Error: empty map has no maximal element +findMax :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> (k, a) +findMax = fromMaybe (error "empty map has no maximal element") . getMax + +-- | The maximal key of the map, if any. Returns 'Nothing' if the map is empty. +-- +-- > getMax (fromList [(5,"a"), (3,"b")]) == Just (5,"a") +-- > getMax empty == Nothing +getMax :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Maybe (k, a) +getMax = fst <.> maxViewWithKey + +-- | Delete the minimal key. Returns an empty map if the map is empty. +-- +-- > deleteMin (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(5,"a"), (7,"c")] +-- > deleteMin empty == empty +deleteMin :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m a +deleteMin m0@(TrieMap n m) = maybe m0 (TrieMap (n-1) . snd) $ getMinAlg m + +-- | Delete the maximal key. Returns an empty map if the map is empty. +-- +-- > deleteMax (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(3,"b"), (5,"a")] +-- > deleteMax empty == empty +deleteMax :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> TrieMap k m a +deleteMax m0@(TrieMap n m) = maybe m0 (TrieMap (n-1) . snd) $ getMaxAlg m + +-- | Delete and find the minimal element. +-- +-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) +-- > deleteFindMin Error: can not return the minimal element of an empty map +deleteFindMin :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> ((k, a), TrieMap k m a) +deleteFindMin = fromMaybe (error "cannot return the minimal element of an empty map") . minViewWithKey + +checkNothing :: Maybe a -> (Bool, Maybe a) +checkNothing x = (isNothing x, x) + +-- | Delete and find the maximal element. +-- +-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) +-- > deleteFindMax empty Error: can not return the maximal element of an empty map +deleteFindMax :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> ((k, a), TrieMap k m a) +deleteFindMax = fromMaybe (error "cannot return the maximal element of an empty map") . maxViewWithKey + +-- | Update the value at the minimal key. +-- +-- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")] +-- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +updateMin :: (Algebraic k, TrieKey (Alg k) m) => (a -> Maybe a) -> TrieMap k m a -> TrieMap k m a +updateMin f (TrieMap n m) = TrieMap (if del then n-1 else n) m' + where (del, m') = updateMinAlg (const (checkNothing . f)) m + +-- | Update the value at the maximal key. +-- +-- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")] +-- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +updateMax :: (Algebraic k, TrieKey (Alg k) m) => (a -> Maybe a) -> TrieMap k m a -> TrieMap k m a +updateMax f (TrieMap n m) = TrieMap (if del then n-1 else n) m' + where (del, m') = updateMaxAlg (const (checkNothing . f)) m + +-- | Update the value at the minimal key. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +updateMinWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Maybe a) -> TrieMap k m a -> TrieMap k m a +updateMinWithKey f (TrieMap n m) = TrieMap (if del then n-1 else n) m' + where (del, m') = updateMinAlg (checkNothing .: f . fromAlg) m + +-- | Update the value at the maximal key. +-- +-- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")] +-- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +updateMaxWithKey :: (Algebraic k, TrieKey (Alg k) m) => (k -> a -> Maybe a) -> TrieMap k m a -> TrieMap k m a +updateMaxWithKey f (TrieMap n m) = TrieMap (if del then n-1 else n) m' + where (del, m') = updateMaxAlg (checkNothing .: f . fromAlg) m + +-- | Retrieves the value associated with the minimal key of the +-- map, and the map stripped of that element, or 'Nothing' if passed an +-- empty map. +-- +-- > minView (fromList [(5,"a"), (3,"b")]) == Just ("b", singleton 5 "a") +-- > minView empty == Nothing +minView :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Maybe (a, TrieMap k m a) +minView (TrieMap n m) = do + (~(_, v), m') <- getMinAlg m + return (v, TrieMap (n-1) m') + +-- | Retrieves the value associated with the maximal key of the +-- map, and the map stripped of that element, or 'Nothing' if passed an +-- +-- > maxView (fromList [(5,"a"), (3,"b")]) == Just ("a", singleton 3 "b") +-- > maxView empty == Nothing +maxView :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Maybe (a, TrieMap k m a) +maxView (TrieMap n m) = do + (~(_, v), m') <- getMaxAlg m + return (v, TrieMap (n-1) m') + +-- | Retrieves the minimal (key,value) pair of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +-- +-- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a") +-- > minViewWithKey empty == Nothing +minViewWithKey :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Maybe ((k, a), TrieMap k m a) +minViewWithKey (TrieMap n m) = do + (~(k, v), m') <- getMinAlg m + return ((fromAlg k, v), TrieMap (n-1) m') + +-- | Retrieves the maximal (key,value) pair of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +-- +-- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b") +-- > maxViewWithKey empty == Nothing +maxViewWithKey :: (Algebraic k, TrieKey (Alg k) m) => TrieMap k m a -> Maybe ((k, a), TrieMap k m a) +maxViewWithKey (TrieMap n m) = do + ~(~(k, v), m') <- getMaxAlg m + return ((fromAlg k, v), TrieMap (n-1) m') + +-- | /O(n+m)/. +-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +-- +isSubmapOf :: (Algebraic k, TrieKey (Alg k) m, Eq a) => TrieMap k m a -> TrieMap k m a -> Bool +isSubmapOf = isSubmapOfBy (==) + +{- | /O(n+m)/. + The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if + all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when + applied to their respective values. For example, the following + expressions are all 'True': + + > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)]) + + But the following are all 'False': + + > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (<) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)]) + +-} +isSubmapOfBy :: (Algebraic k, TrieKey (Alg k) m) => (a -> b -> Bool) -> TrieMap k m a -> TrieMap k m b -> Bool +isSubmapOfBy (<=) (TrieMap n1 m1) (TrieMap n2 m2) = (Prelude.<=) n1 n2 && isSubmapAlg (<=) m1 m2 + +-- | The expression (@'split' k map@) is a pair @(map1,map2)@ where +-- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. +-- Any key equal to @k@ is found in neither @map1@ nor @map2@. +-- +-- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")]) +-- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a") +-- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") +-- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty) +-- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty) +split :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> (TrieMap k m a, TrieMap k m a) +split k m = case splitLookup k m of + (mL, _, mR) -> (mL, mR) + +-- | The expression (@'splitLookup' k map@) splits a map just +-- like 'split' but also returns @'lookup' k map@. +-- +-- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")]) +-- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a") +-- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a") +-- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty) +-- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty) +splitLookup :: (Algebraic k, TrieKey (Alg k) m) => k -> TrieMap k m a -> (TrieMap k m a, Maybe a, TrieMap k m a) +splitLookup k (TrieMap n m) = case splitLookupAlg (\ v -> (Nothing, Just v, Nothing)) (toAlg k) m of + (mL, v, mR) -> (mkTrieMap mL, v, mkTrieMap mR) +-- TODO: Somehow, avoid the mkTrieMap call. Is this possible? I don't think so, without a sophisticated range-mconcat operation +-- with monoids or some crazy shit like that. diff --git a/TrieMap/Algebraic.hs b/TrieMap/Algebraic.hs new file mode 100644 --- /dev/null +++ b/TrieMap/Algebraic.hs @@ -0,0 +1,122 @@ +{-# LANGUAGE UndecidableInstances, TypeFamilies, TypeSynonymInstances #-} + +module TrieMap.Algebraic (Algebraic(..), Ordered(..)) where + +import Data.Char +import Data.Maybe +import Data.IntSet (IntSet) +import Data.Set(Set) +import qualified Data.IntSet as ISet +import qualified Data.Set as Set +import Data.IntMap (IntMap) +import Data.Map (Map) +import qualified Data.IntMap as IMap +import qualified Data.Map as Map +import qualified Data.Foldable as Fold +import GHC.Exts (build) + +import TrieMap.TrieAlgebraic + +-- | 'Algebraic' refers to a type with an algebraic representation, armed with methods to convert in each direction. +-- 'toAlg' and 'fromAlg' should preserve equality and ordering. +class Algebraic k where + -- | @'Alg' k@ is a fully decomposed representation of k into algebraic pieces. + type Alg k + toAlg :: k -> Alg k + fromAlg :: Alg k -> k + +instance (Algebraic k1, Algebraic k2) => Algebraic (k1, k2) where + type Alg (k1, k2) = (Alg k1, Alg k2) + toAlg (k1, k2) = (toAlg k1, toAlg k2) + fromAlg (k1, k2) = (fromAlg k1, fromAlg k2) + +instance (Algebraic a, Algebraic b, Algebraic c) => Algebraic (a, b, c) where + type Alg (a, b, c) = (Alg a, (Alg b, Alg c)) + toAlg (a, b, c) = toAlg (a, (b, c)) + fromAlg x = case fromAlg x of + (a, (b, c)) -> (a, b, c) + +instance (Algebraic a, Algebraic b, Algebraic c, Algebraic d) => Algebraic (a, b, c, d) where + type Alg (a, b, c, d) = (Alg a, (Alg b, (Alg c, Alg d))) + toAlg (a, b, c, d) = toAlg (a, (b, (c, d))) + fromAlg x = case fromAlg x of + (a, (b, (c, d))) -> (a, b, c, d) + +instance (Algebraic a, Algebraic b, Algebraic c, Algebraic d, Algebraic e) => Algebraic (a, b, c, d, e) where + type Alg (a, b, c, d, e) = (Alg a, (Alg b, (Alg c, (Alg d, Alg e)))) + toAlg (a, b, c, d, e) = toAlg (a, (b, (c, (d, e)))) + fromAlg x = case fromAlg x of + (a, (b, (c, (d, e)))) -> (a, b, c, d, e) + +instance (Algebraic k1, Algebraic k2) => Algebraic (Either k1 k2) where + type Alg (Either k1 k2) = Either (Alg k1) (Alg k2) + toAlg = either (Left . toAlg) (Right . toAlg) + fromAlg = either (Left . fromAlg) (Right . fromAlg) + +instance Algebraic k => Algebraic [k] where + type Alg [k] = [Alg k] + toAlg = map toAlg + fromAlg = map fromAlg + +instance Algebraic () where + type Alg () = () + toAlg = id + fromAlg = id + +instance Algebraic a => Algebraic (Maybe a) where + type Alg (Maybe a) = Either () (Alg a) + toAlg Nothing = Left () + toAlg (Just a) = Right (toAlg a) + fromAlg (Left _) = Nothing + fromAlg (Right a) = Just (fromAlg a) + +instance Algebraic Bool where + type Alg Bool = Alg (Maybe ()) + toAlg b = toAlg $ if b then Just () else Nothing + fromAlg = maybe False (const True) . fromAlg' + where fromAlg' = fromAlg :: Alg (Maybe ()) -> Maybe () + +instance Algebraic Int where + type Alg Int = Int + toAlg = id + fromAlg = id + +instance Algebraic Char where + type Alg Char = Int + toAlg = ord + fromAlg = chr + +instance Algebraic Float where + type Alg Float = Ordered Float + toAlg = Ord + fromAlg = unOrd + +instance Algebraic Double where + type Alg Double = Ordered Double + toAlg = Ord + fromAlg = unOrd + +instance Algebraic Rational where + type Alg Rational = Ordered Rational + toAlg = Ord + fromAlg = unOrd + +instance (Algebraic k, Algebraic v) => Algebraic (Map k v) where + type Alg (Map k v) = [(Alg k, Alg v)] + toAlg m = build (\ c n -> Map.foldWithKey (\ k v -> c (toAlg k, toAlg v)) n m) + fromAlg xs = Map.fromDistinctAscList [(fromAlg k, fromAlg v) | (k, v) <- xs] + +instance Algebraic v => Algebraic (IntMap v) where + type Alg (IntMap v) = [(Int, Alg v)] + toAlg m = build (\ c n -> IMap.foldWithKey (\ k v -> c (k, toAlg v)) n m) + fromAlg xs = IMap.fromDistinctAscList [(k, fromAlg v) | (k, v) <- xs] + +instance Algebraic a => Algebraic (Set a) where + type Alg (Set a) = [Alg a] + toAlg s = build (\ c n -> Fold.foldr (c . toAlg) n s) + fromAlg = Set.fromDistinctAscList . map fromAlg + +instance Algebraic IntSet where + type Alg IntSet = [Int] + toAlg = ISet.toList + fromAlg = ISet.fromDistinctAscList diff --git a/TrieMap/Applicative.hs b/TrieMap/Applicative.hs new file mode 100644 --- /dev/null +++ b/TrieMap/Applicative.hs @@ -0,0 +1,32 @@ +module TrieMap.Applicative(Id(..), (.:), (<.>), on, build) where + +import Control.Monad +import Control.Applicative +import Data.Traversable (sequenceA) +import GHC.Exts (build) + +newtype Id a = Id {unId :: a} + +instance Functor Id where + fmap f (Id x) = Id (f x) + +instance Applicative Id where + pure = return + (<*>) = ap + +instance Monad Id where + return = Id + m >>= k = k (unId m) + +(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d +(.:) = (.) . (.) + +(<.>) :: Functor f => (b -> c) -> (a -> f b) -> (a -> f c) +(<.>) = (.) . (<$>) + +on :: (b -> b -> c) -> (a -> b) -> a -> a -> c +(f `on` g) x y = f (g x) (g y) + +infixr 9 <.> +infixr 9 .: +infixr 8 `on` diff --git a/TrieMap/MapTypes.hs b/TrieMap/MapTypes.hs new file mode 100644 --- /dev/null +++ b/TrieMap/MapTypes.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE TypeOperators, FlexibleContexts, StandaloneDeriving #-} + +module TrieMap.MapTypes where + +import Data.Foldable +import Data.Traversable +import Control.Applicative +import Prelude hiding (foldl, foldr) + +-- | 'ProdMap' is used to hold a map on the product of two key types. +newtype ProdMap m1 m2 v = PMap {unPMap :: m1 (m2 v)} deriving (Eq, Ord) + +-- | 'UnionMap' is used to hold a map on the sum of two key types. +data UnionMap m1 m2 v = m1 v :+: m2 v deriving (Eq, Ord) + +data Edge k m v = Edge [k] (Maybe v) (m (Edge k m v)) +type MEdge k m v = Maybe (Edge k m v) + +-- | 'RadixTrie' is used to hold a map on a list of keys. +newtype RadixTrie k m v = Radix {unRad :: MEdge k m v} + +infixr 5 `ProdMap` +infixr 5 :+: + +instance (Functor m1, Functor m2) => Functor (ProdMap m1 m2) where + fmap f (PMap m) = PMap (fmap (fmap f) m) + +instance (Foldable m1, Foldable m2) => Foldable (ProdMap m1 m2) where + foldr f z (PMap m) = foldr (flip (foldr f)) z m + foldl f z (PMap m) = foldl (foldl f) z m + +instance (Traversable m1, Traversable m2) => Traversable (ProdMap m1 m2) where + traverse f (PMap m) = PMap <$> traverse (traverse f) m + +instance (Functor m1, Functor m2) => Functor (UnionMap m1 m2) where + fmap f (m1 :+: m2) = fmap f m1 :+: fmap f m2 + +instance (Foldable m1, Foldable m2) => Foldable (UnionMap m1 m2) where + foldr f z (m1 :+: m2) = foldr f (foldr f z m2) m1 + foldl f z (m1 :+: m2) = foldl f (foldl f z m1) m2 + +instance (Traversable m1, Traversable m2) => Traversable (UnionMap m1 m2) where + traverse f (m1 :+: m2) = liftA2 (:+:) (traverse f m1) (traverse f m2) + +instance Functor m => Functor (Edge k m) where + fmap f (Edge ks v ts) = Edge ks (fmap f v) (fmap (fmap f) ts) + +instance Functor m => Functor (RadixTrie k m) where + fmap f (Radix e) = Radix (fmap (fmap f) e) + +instance Foldable m => Foldable (Edge k m) where + foldr f z (Edge _ v ts) = foldr (flip (foldr f)) (foldr f z v) ts + foldl f z (Edge _ v ts) = foldl f (foldl (foldl f) z ts) v + +instance Foldable m => Foldable (RadixTrie k m) where + foldr f z (Radix e) = foldr (flip (foldr f)) z e + foldl f z (Radix e) = foldl (foldl f) z e + +instance Traversable m => Traversable (Edge k m) where + traverse f (Edge ks v ts) = + liftA2 (Edge ks) (traverse f v) (traverse (traverse f) ts) + +instance Traversable m => Traversable (RadixTrie k m) where + traverse f (Radix e) = Radix <$> traverse (traverse f) e diff --git a/TrieMap/RadixTrie.hs b/TrieMap/RadixTrie.hs new file mode 100644 --- /dev/null +++ b/TrieMap/RadixTrie.hs @@ -0,0 +1,284 @@ +{-# LANGUAGE MultiParamTypeClasses, UndecidableInstances, FlexibleContexts, StandaloneDeriving, PatternGuards #-} + +module TrieMap.RadixTrie (RadixTrie) where + +import Control.Applicative hiding (Alternative(..)) +import Control.Monad +import Data.Foldable +import Data.Traversable +import Data.Monoid +import Data.Maybe +import Data.Ord +import Data.Sequence (Seq, (|>)) +import qualified Data.Sequence as Seq + +import TrieMap.MapTypes +import TrieMap.TrieAlgebraic +import TrieMap.Applicative + +import Prelude hiding (null, foldr, all) + +instance (Eq k, Eq v, TrieKey k m) => Eq (Edge k m v) where + Edge ks1 v1 ts1 == Edge ks2 v2 ts2 = ks1 == ks2 && v1 == v2 && assocsAlg ts1 == assocsAlg ts2 + +instance (Ord k, Ord v, TrieKey k m) => Ord (Edge k m v) where + Edge ks1 v1 ts1 `compare` Edge ks2 v2 ts2 = + compare ks1 ks2 `mappend` compare v1 v2 `mappend` comparing assocsAlg ts1 ts2 + +deriving instance (Eq k, Eq v, TrieKey k m) => Eq (RadixTrie k m v) +deriving instance (Ord k, Ord v, TrieKey k m) => Ord (RadixTrie k m v) +deriving instance (Show k, Show v, Functor m, Show (m String)) => Show (RadixTrie k m v) + +instance (Show k, Show v, Functor m, Show (m String)) => Show (Edge k m v) where + show (Edge k v ts) = "Edge " ++ show k ++ " " ++ show v ++ " " ++ show (fmap show ts) + +instance (Ord k, TrieKey k m) => TrieKey [k] (RadixTrie k m) where + emptyAlg = Radix Nothing + nullAlg = isNothing . unRad + getSingleAlg (Radix e) = e >>= getSingleEdge + guardNullAlg (Radix e) = do e <- guardNullEdge =<< e + return (Radix (Just e)) + lookupAlg ks = unRad >=> lookupEdge (==) ks +-- sizeAlg (Radix e) = maybe 0 sizeEdge e + alterLookupAlg f k = fmap Radix . + maybe (fmap (maybeSingleEdge k) $ f Nothing) + (alterLookupEdge (==) f k) . unRad + foldWithKeyAlg f z = foldr (flip (foldWithKeyEdge f)) z . unRad + mapMaybeAlg f (Radix e) = Radix (e >>= mapMaybeEdge f) + mapEitherAlg f (Radix Nothing) = (emptyAlg, emptyAlg) + mapEitherAlg f (Radix (Just e)) = (Radix e1, Radix e2) + where (e1, e2) = mapEitherEdge f e +-- mapMaybeAlg f (Radix e) = (Radix . join) <$> traverse (mapAppMaybeEdge f) e + mapAppAlg f = fmap Radix . traverse (mapAppEdge f) . unRad + unionMaybeAlg f (Radix e1) (Radix e2) = Radix (unionMaybe (unionMaybeEdge f) e1 e2) + intersectAlg f (Radix e1) (Radix e2) = Radix (intersectMaybe (intersectEdge f) e1 e2) + differenceAlg f (Radix e1) (Radix e2) = Radix (differenceMaybe (differenceEdge f) e1 e2) + + getMinAlg (Radix e) = fmap (fmap Radix . getMinEdge) e + getMaxAlg (Radix e) = fmap (fmap Radix . getMaxEdge) e +-- updateMinAlg f (Radix e) = Radix $ e >>= updateMinEdge f +-- updateMaxAlg f (Radix e) = Radix $ e >>= updateMaxEdge f + + fromListAlg f xs = Radix (edgeFromList f xs) + fromAscListAlg f xs = Radix (edgeFromAscList f xs) + fromDistAscListAlg = fromAscListAlg (\ _ v _ -> v) + + isSubmapAlg (<=) (Radix e1) (Radix e2) = isSubmapAlg (isSubmapEdge (<=)) e1 e2 + + valid (Radix e) = maybe True validEdge e + + splitLookupAlg _ _ (Radix Nothing) = (Radix Nothing, Nothing, Radix Nothing) + splitLookupAlg f k (Radix (Just e)) = case splitEdge f k e of + (eL, ans, eR) -> (Radix eL, ans, Radix eR) + +-- sizeEdge :: Edge k m v -> Int +-- sizeEdge (Edge n _ _ _) = n + +-- edge :: TrieKey k m => [k] -> Maybe v -> m (Edge k m v) -> Edge k m v +-- edge ks v ts = Edge (maybe id (const (+1)) v $ foldl' (\ n e -> n + sizeEdge e) 0 ts) ks v ts + +lookupEdge :: TrieKey k m => (k -> k -> Bool) -> [k] -> Edge k m v -> Maybe v +lookupEdge (==) ks (Edge ls v ts) = procEdge ks ls where + procEdge (k:ks) (l:ls) + | k == l = procEdge ks ls + procEdge (k:ks) [] = lookupAlg k ts >>= lookupEdge (==) ks + procEdge [] [] = v + procEdge _ _ = Nothing + +edgeFromList :: (Eq k, TrieKey k m) => ([k] -> v -> v -> v) -> [([k], v)] -> MEdge k m v +edgeFromList f xs = guardNullEdge $ Edge [] v0 $ mapMaybeAlg (\ k -> edgeFromList (f . (k:))) $ fromListAlg (const (flip (++))) ys + where part ([], v) (v0, ys) = (Just $ maybe v (flip (f []) v) v0, ys) + part (k:ks, v) (v0, ys) = (v0, (k, [(ks, v)]):ys) + (v0, ys) = foldr part (Nothing, []) xs + +edgeFromAscList :: (Eq k, TrieKey k m) => ([k] -> v -> v -> v) -> [([k], v)] -> MEdge k m v +edgeFromAscList _ [] = Nothing +edgeFromAscList f xs = Just $ case groupHead f xs of + (Nothing, [(k, ~(Edge ks v ts))]) + -> Edge (k:ks) v ts + (ans, xs') -> Edge [] ans (fromDistAscListAlg xs') + +groupHead :: (Eq k, TrieKey k m) => ([k] -> v -> v -> v) -> [([k], v)] -> (Maybe v, [(k, Edge k m v)]) +groupHead f (([], v):xs) = case groupHead f xs of + (v', ans) -> (Just $ maybe v (f [] v) v', ans) +groupHead f ((k:ks, v):xs) = (Nothing, groupHead' k (Seq.singleton (ks, v)) xs) where + groupHead' k0 xs ((k:ks, v):ys) + | k == k0 = groupHead' k0 (xs |> (ks, v)) ys + | otherwise = (k0, fromJust $ edgeFromAscList (f . (k0:)) (toList xs)):groupHead' k (Seq.singleton (ks, v)) ys + groupHead' k0 xs [] = [(k0, fromJust $ edgeFromAscList (f . (k0:)) (toList xs))] + groupHead' _ _ _ = error "Violation of ascending invariant!" +groupHead _ [] = (Nothing, []) + {-guardNullEdge $ Edge [] v0 $ mapMaybeAlg (\ k -> edgeFromAscList (f . (k:))) $ fromAscListAlg (const (flip (++))) ys + where part ([], v) (v0, ys) = (Just $ maybe v (flip (f []) v) v0, ys) + part (k:ks, v) (v0, ys) = (v0, (k, [(ks, v)]):ys) + (v0, ys) = foldr part (Nothing, []) xs-} + +maybeSingleEdge :: TrieKey k m => [k] -> Maybe v -> MEdge k m v +maybeSingleEdge ks = fmap (\ v -> Edge ks (Just v) emptyAlg) + +getSingleEdge :: (TrieKey k m) => Edge k m v -> Maybe ([k], v) +getSingleEdge (Edge ks (Just v) ts) + | nullAlg ts = Just (ks, v) +getSingleEdge (Edge ks Nothing ts) = do + (x, e') <- getSingleAlg ts + (xs, v) <- getSingleEdge e' + return (ks ++ x:xs, v) +getSingleEdge _ = Nothing + +{-# INLINE guardNullEdge #-} +guardNullEdge :: TrieKey k m => Edge k m v -> MEdge k m v +guardNullEdge (Edge ks Nothing ts) + | nullAlg ts = Nothing + | Just (x, Edge xs v ts') <- getSingleAlg ts + = Just (Edge (ks ++ x:xs) v ts') +guardNullEdge e = Just e + +alterLookupEdge :: (TrieKey k m) => (k -> k -> Bool) -> + (Maybe v -> (a, Maybe v)) -> [k] -> Edge k m v -> (a, MEdge k m v) +alterLookupEdge (==) f ks0 e@(Edge ls0 v ts) = procEdge 0 ks0 ls0 where + procEdge i _ _ | i `seq` False = undefined + procEdge i (k:ks) (l:ls) + | k == l = procEdge (i+1) ks ls + | otherwise = fmap (Just . g) $ f Nothing + where g Nothing = e + g (Just v') = Edge (take i ks0) Nothing $ + fromListAlg' [(k, Edge ks (Just v') emptyAlg), (l, Edge ls v ts)] + procEdge i (k:ks) [] = proc (alterLookupAlg g k ts) where + g Nothing = maybeSingleEdge ks <$> f Nothing + g (Just e') = alterLookupEdge (==) f ks e' + proc = fmap (guardNullEdge . Edge ls0 v) + procEdge i [] (l:ls) = fmap (Just . g) $ f Nothing + where g Nothing = e + g (Just v') = Edge ks0 (Just v') $ insertAlg l (Edge ls v ts) emptyAlg + procEdge i [] [] = (ans, guardNullEdge (Edge ks0 fv ts)) + where (ans, fv) = f v + +foldWithKeyEdge :: TrieKey k m => ([k] -> v -> x -> x) -> x -> Edge k m v -> x +foldWithKeyEdge f z (Edge ks v ts) = + foldr (f ks) (foldWithKeyAlg (\ x -> flip (foldWithKeyEdge (\ xs -> f (ks ++ x:xs)))) z ts) v + +mapMaybeEdge :: (TrieKey k m) => ([k] -> v -> Maybe w) -> Edge k m v -> MEdge k m w +mapMaybeEdge f (Edge ks v ts) = guardNullEdge $ + Edge ks (join $ traverse (f ks) v) (mapMaybeAlg (\ x -> mapMaybeEdge (\ xs -> f (ks ++ x:xs))) ts) + +mapEitherEdge :: TrieKey k m => ([k] -> a -> Either b c) -> Edge k m a -> (MEdge k m b, MEdge k m c) +mapEitherEdge f (Edge ks v ts) = + (guardNullEdge $ Edge ks vL tsL, guardNullEdge $ Edge ks vR tsR) + where (vL, vR) = case fmap (f ks) v of + Nothing -> (Nothing, Nothing) + Just (Left v) -> (Just v, Nothing) + Just (Right v) -> (Nothing, Just v) + ts' = mapWithKeyAlg (\ x -> mapEitherEdge (\ xs -> f (ks ++ x:xs))) ts + tsL = mapMaybeAlg (const fst) ts' + tsR = mapMaybeAlg (const snd) ts' + +mapAppEdge :: (Applicative f, TrieKey k m) => ([k] -> v -> f w) -> Edge k m v -> f (Edge k m w) +mapAppEdge f (Edge ks v ts) = liftA2 (Edge ks) (traverse (f ks) v) (mapAppAlg (\ x -> mapAppEdge (\ xs -> f (ks ++ x:xs))) ts) + +unionMaybeEdge :: (Eq k, TrieKey k m) => ([k] -> v -> v -> Maybe v) -> Edge k m v -> Edge k m v -> MEdge k m v +unionMaybeEdge f (Edge ks0 vK tsK) (Edge ls0 vL tsL) = procEdge 0 ks0 ls0 where + procEdge i _ _ | i `seq` False = undefined + procEdge i (k:ks) (l:ls) + | k == l = procEdge (i+1) ks ls + | otherwise = Just $ Edge (take i ks0) Nothing $ fromListAlg' [(k, Edge ks vK tsK), (l, Edge ls vL tsL)] + procEdge _ [] (l:ls) = guardNullEdge $ Edge ks0 vK $ alterAlg g l tsK + where g Nothing = Just (Edge ls vL tsL) + g (Just e') = unionMaybeEdge (\ ls' -> f (ks0 ++ l:ls')) e' (Edge ls vL tsL) + procEdge _ (k:ks) [] = guardNullEdge $ Edge ls0 vL $ alterAlg g k tsL + where g Nothing = Just $ Edge ks vK tsK + g (Just e') = unionMaybeEdge (\ ks' -> f (ls0 ++ k:ks')) (Edge ks vK tsK) e' + procEdge _ [] [] = guardNullEdge $ Edge ks0 (unionMaybe (f ks0) vK vL) $ + unionMaybeAlg (\ x -> unionMaybeEdge (\ xs -> f (ks0 ++ x:xs))) tsK tsL + +intersectEdge :: (Eq k, TrieKey k m) => ([k] -> a -> b -> Maybe c) -> Edge k m a -> Edge k m b -> MEdge k m c +intersectEdge f (Edge ks0 vK tsK) (Edge ls0 vL tsL) = procEdge ks0 ls0 where + procEdge (k:ks) (l:ls) + | k == l = procEdge ks ls + | otherwise = Nothing + procEdge (k:ks) [] = do + e' <- lookupAlg k tsL + Edge xs vX tsX <- intersectEdge (\ ks' -> f (ls0 ++ k:ks')) (Edge ks vK tsK) e' + return (Edge (ls0 ++ k:xs) vX tsX) + procEdge [] (l:ls) = do + e' <- lookupAlg l tsK + Edge xs vX tsX <- intersectEdge (\ ls' -> f (ks0 ++ l:ls')) e' (Edge ls vL tsL) + return (Edge (ks0 ++ l:xs) vX tsX) + procEdge [] [] = guardNullEdge $ Edge ks0 (intersectMaybe (f ks0) vK vL) $ + intersectAlg (\ x -> intersectEdge (\ xs -> f (ks0 ++ x:xs))) tsK tsL + +differenceEdge :: (Eq k, TrieKey k m) => ([k] -> v -> w -> Maybe v) -> Edge k m v -> Edge k m w -> MEdge k m v +differenceEdge f e@(Edge ks0 vK tsK) (Edge ls0 vL tsL) = procEdge ks0 ls0 where + procEdge (k:ks) (l:ls) + | k == l = procEdge ks ls + procEdge (k:ks) [] + | Just e' <- lookupAlg k tsL + = do Edge xs vX tsX <- differenceEdge (\ ks' -> f (ls0 ++ k:ks')) (Edge ks vK tsK) e' + return (Edge (ls0 ++ k:xs) vX tsX) + procEdge [] (l:ls) = guardNullEdge $ Edge ks0 vK $ alterAlg g l tsK + where g Nothing = Nothing + g (Just e') = differenceEdge (\ ls' -> f (ks0 ++ l:ls')) e' (Edge ls vL tsL) + procEdge [] [] = guardNullEdge $ Edge ks0 (differenceMaybe (f ks0) vK vL) $ + differenceAlg (\ x -> differenceEdge (\ xs -> f (ks0 ++ x:xs))) tsK tsL + procEdge _ _ = Just e + +getMinEdge :: TrieKey k m => Edge k m v -> (([k], v), MEdge k m v) +getMinEdge (Edge ks (Just v) ts) = ((ks, v), guardNullEdge $ Edge ks Nothing ts) +getMinEdge (Edge ks _ ts) + | Just ((l, e), ts') <- getMinAlg ts, ((ls, v), e') <- getMinEdge e + = ((ks ++ l:ls, v), guardNullEdge $ Edge ks Nothing $ maybe ts' (\ e' -> snd $ updateMinAlg (\ _ _ -> (False, Just e')) ts) e') +getMinEdge _ = error "Uncompacted edge" + +getMaxEdge :: TrieKey k m => Edge k m v -> (([k], v), MEdge k m v) +getMaxEdge (Edge ks v0 ts) + | Just ((l, e), ts') <- getMaxAlg ts, ((ls, v), e') <- getMaxEdge e + = ((ks ++ l:ls, v), guardNullEdge $ Edge ks v0 $ maybe ts' (\ e' -> snd $ updateMaxAlg (\ _ _ -> (False, Just e')) ts) e') +getMaxEdge (Edge ks (Just v) ts) = ((ks, v), guardNullEdge $ Edge ks Nothing ts) +getMaxEdge _ = error "Uncompacted edge" + +updateMinEdge :: TrieKey k m => ([k] -> v -> (Bool, Maybe v)) -> Edge k m v -> (Bool, MEdge k m v) +updateMinEdge f (Edge ks (Just v) ts) = fmap (\ v' -> guardNullEdge $ Edge ks v' ts) (f ks v) +updateMinEdge f (Edge ks Nothing ts) + = fmap (guardNullEdge . Edge ks Nothing) $ updateMinAlg (\ l -> updateMinEdge (\ ls -> f (ks ++ l:ls))) ts + +updateMaxEdge :: TrieKey k m => ([k] -> v -> (Bool, Maybe v)) -> Edge k m v -> (Bool, MEdge k m v) +updateMaxEdge f (Edge ks (Just v) ts) + | nullAlg ts = fmap (\ v' -> guardNullEdge $ Edge ks v' ts) (f ks v) +updateMaxEdge f (Edge ks v ts) = + fmap (guardNullEdge . Edge ks v) $ updateMaxAlg (\ l -> updateMaxEdge (\ ls -> f (ks ++ l:ls))) ts + +isSubmapEdge :: TrieKey k m => (a -> b -> Bool) -> Edge k m a -> Edge k m b -> Bool +isSubmapEdge (<=) (Edge ks vK tsK) (Edge ls vL tsL) = procEdge ks ls where + procEdge (k:ks) (l:ls) + | k == l = procEdge ks ls + procEdge (k:ks) [] + | Just e <- lookupAlg k tsL + = isSubmapEdge (<=) (Edge ks vK tsK) e + procEdge [] [] + | Nothing <- vK = isSubmapAlg (isSubmapEdge (<=)) tsK tsL + | Just x <- vK, Just y <- vL, x <= y + = isSubmapAlg (isSubmapEdge (<=)) tsK tsL + procEdge _ _ = False +validEdge :: TrieKey k m => Edge k m v -> Bool +validEdge (Edge _ Nothing m) + | nullAlg m = False + | Just{} <- getSingleAlg m + = False +validEdge (Edge _ _ m) + = valid m && all validEdge m + + +splitEdge :: (Ord k, TrieKey k m) => (a -> (Maybe a, Maybe b, Maybe a)) -> [k] -> Edge k m a -> (MEdge k m a, Maybe b, MEdge k m a) +splitEdge f ks0 (Edge ls0 v ts) = procEdge ks0 ls0 where + procEdge (k:ks) (l:ls) = case compare k l of + LT -> (Nothing, Nothing, Just (Edge ls0 v ts)) + EQ -> procEdge ks ls + GT -> (Just (Edge ks0 v ts), Nothing, Nothing) + procEdge (k:ks) [] = case splitLookupAlg (splitEdge f ks) k ts of + (tsL, ans, tsR) -> (guardNullEdge $ Edge ls0 Nothing tsL, ans, guardNullEdge $ Edge ls0 v tsR) + procEdge [] (l:ls) = (Nothing, Nothing, Just $ Edge ls0 v ts) + procEdge [] [] + | Just v <- v, (vL, ans, vR) <- f v + = (fmap (\ v' -> Edge ls0 (Just v') emptyAlg) vL, ans, + guardNullEdge $ Edge ls0 vR ts) + | otherwise = (Nothing, Nothing, Just (Edge ls0 v ts)) + diff --git a/TrieMap/Reflection.hs b/TrieMap/Reflection.hs new file mode 100644 --- /dev/null +++ b/TrieMap/Reflection.hs @@ -0,0 +1,25 @@ +{-# LANGUAGE TypeFamilies, FlexibleContexts, UndecidableInstances #-} + +module TrieMap.Reflection where + +-- import TrieMap.Fixpoint +import TrieMap.TrieAlgebraic +import TrieMap.Algebraic +import TrieMap.Applicative +import TrieMap.RadixTrie() +import qualified TrieMap.TrieAlgebraic as TA + +instance Algebraic (m1 (m2 v)) => Algebraic (ProdMap m1 m2 v) where + type Alg (ProdMap m1 m2 v) = Alg (m1 (m2 v)) + toAlg (PMap m) = toAlg m + fromAlg = PMap . fromAlg + +instance (Algebraic (m1 v), Algebraic (m2 v)) => Algebraic (UnionMap m1 m2 v) where + type Alg (UnionMap m1 m2 v) = (Alg (m1 v), Alg (m2 v)) + toAlg (m1 :+: m2) = (toAlg m1, toAlg m2) + fromAlg (m1, m2) = fromAlg m1 :+: fromAlg m2 + +instance (Ord k, Algebraic k, Algebraic v, TrieKey k m) => Algebraic (RadixTrie k m v) where + type Alg (RadixTrie k m v) = Alg [([k], v)] + toAlg m = toAlg (build (\ c n -> foldWithKeyAlg (curry c) n m)) + fromAlg = fromDistAscListAlg . fromAlg diff --git a/TrieMap/TrieAlgebraic.hs b/TrieMap/TrieAlgebraic.hs new file mode 100644 --- /dev/null +++ b/TrieMap/TrieAlgebraic.hs @@ -0,0 +1,364 @@ +{-# LANGUAGE TypeOperators, MultiParamTypeClasses, FunctionalDependencies, UndecidableInstances, PatternGuards #-} + +module TrieMap.TrieAlgebraic (TrieKey (..), ProdMap (..), UnionMap(..), RadixTrie(..), Edge (..), Ordered (..), unionMaybe, intersectMaybe, differenceMaybe, mapWithKeyAlg, assocsAlg, insertAlg, alterAlg, fromListAlg') where + +import Data.Traversable +import Data.Foldable +import Data.Either +import Data.Sequence (Seq) +import Data.Maybe +import Data.Monoid +import Data.IntMap (IntMap) +import Data.Map (Map) +import qualified Data.Sequence as Seq +import qualified Data.IntMap as IMap +import qualified Data.Map as Map + +import Control.Monad +import Control.Applicative hiding (Alternative(..)) + +import GHC.Exts (build) + +import TrieMap.Applicative +-- import TrieMap.Algebraic (Ordered (..)) +import TrieMap.MapTypes +import Prelude hiding (foldr, foldl, all, any) + +newtype Ordered k = Ord {unOrd :: k} deriving (Eq, Ord) + +instance Show k => Show (Ordered k) where + show = show . unOrd + showsPrec x = showsPrec x . unOrd + +instance Functor Ordered where + fmap f (Ord x) = Ord (f x) + +-- | TrieKey defines a bijection between map types and algebraic key types. +class (Eq a, Foldable m, Traversable m) => TrieKey a m | a -> m, m -> a where + emptyAlg :: m v + nullAlg :: m v -> Bool + sizeAlg :: m v -> Int + getSingleAlg :: m v -> Maybe (a, v) + guardNullAlg :: m v -> Maybe (m v) + {-# SPECIALIZE alterAlg :: (Maybe v -> Id (b, Maybe v)) -> a -> m v -> Id (b, m v) #-} + alterLookupAlg :: (Maybe v -> (b, Maybe v)) -> a -> m v -> (b, m v) + lookupAlg :: a -> m v -> Maybe v + foldWithKeyAlg :: (a -> v -> x -> x) -> x -> m v -> x + mapAppAlg :: Applicative f => (a -> v -> f w) -> m v -> f (m w) + mapMaybeAlg :: (a -> v -> Maybe w) -> m v -> m w + mapEitherAlg :: (a -> v -> Either x y) -> m v -> (m x, m y) + unionMaybeAlg :: (a -> v -> v -> Maybe v) -> m v -> m v -> m v + intersectAlg :: (a -> v -> w -> Maybe x) -> m v -> m w -> m x + differenceAlg :: (a -> v -> w -> Maybe v) -> m v -> m w -> m v + fromDistAscListAlg :: [(a, v)] -> m v + fromAscListAlg :: (a -> v -> v -> v) -> [(a, v)] -> m v + fromListAlg :: (a -> v -> v -> v) -> [(a, v)] -> m v + getMinAlg :: m v -> Maybe ((a, v), m v) + getMaxAlg :: m v -> Maybe ((a, v), m v) + updateMinAlg :: (a -> v -> (Bool, Maybe v)) -> m v -> (Bool, m v) + updateMaxAlg :: (a -> v -> (Bool, Maybe v)) -> m v -> (Bool, m v) + valid :: m v -> Bool + isSubmapAlg :: (v -> w -> Bool) -> m v -> m w -> Bool + splitLookupAlg :: (v -> (Maybe v, Maybe x, Maybe v)) -> a -> m v -> (m v, Maybe x, m v) + + lookupAlg k = fst . alterLookupAlg (\ v -> (v, v)) k + guardNullAlg m + | nullAlg m = Nothing + | otherwise = Just m + fromListAlg f = foldr (\ (k, v) -> alterAlg (Just . maybe v (f k v)) k) emptyAlg + fromAscListAlg _ [] = emptyAlg + fromAscListAlg f ((k, v):xs) = fromDistAscListAlg (distinct k v xs) where + distinct k v ((k', v'):xs) + | k == k' = distinct k (f k v v') xs + | otherwise = (k, v):distinct k' v' xs + distinct k v [] = [(k, v)] + fromDistAscListAlg = fromListAlg' + sizeAlg = foldl' (\ n _ -> n + 1) 0 + + updateMinAlg f m = maybe (False, m) (\ ((k, v), m') -> maybe m' (\ v' -> insertAlg k v' m) <$> f k v) (getMinAlg m) + updateMaxAlg f m = maybe (False, m) (\ ((k, v), m') -> maybe m' (\ v' -> insertAlg k v' m) <$> f k v) (getMaxAlg m) + valid = (`seq` True) + +fromListAlg' :: TrieKey k m => [(k, v)] -> m v +fromListAlg' = fromListAlg (const const) + +singletonAlg :: TrieKey k m => k -> v -> m v +singletonAlg k v = insertAlg k v emptyAlg + +mapWithKeyAlg :: TrieKey k m => (k -> v -> w) -> m v -> m w +mapWithKeyAlg f m = unId (mapAppAlg (\ k v -> Id (f k v)) m) + +-- mapMaybeWithKeyAlg :: TrieKey k m => (k -> v -> Maybe w) -> m v -> m w +-- mapMaybeWithKeyAlg f m = unId (mapAppMaybeAlg (\ k v -> Id (f k v)) m) + +insertAlg :: TrieKey k m => k -> v -> m v -> m v +insertAlg k v = alterAlg (const (Just v)) k + +alterAlg :: TrieKey k m => (Maybe v -> Maybe v) -> k -> m v -> m v +alterAlg f k = snd . alterLookupAlg (\ x -> ((), f x)) k + +-- alterLookupAlg :: TrieKey k m => (Maybe a -> (b, Maybe a)) -> k -> m a -> (b, m a) +-- alterLookupAlg f = unId .: alterAppAlg (Id . f) + +foldrAlg :: TrieKey k m => (a -> b -> b) -> b -> m a -> b +foldrAlg = foldWithKeyAlg . const + +unionMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a +unionMaybe f (Just x) (Just y) = f x y +unionMaybe _ Nothing y = y +unionMaybe _ x Nothing = x + +intersectMaybe :: (a -> b -> Maybe c) -> Maybe a -> Maybe b -> Maybe c +intersectMaybe f (Just x) (Just y) = f x y +intersectMaybe _ _ _ = Nothing + +differenceMaybe :: (a -> b -> Maybe a) -> Maybe a -> Maybe b -> Maybe a +differenceMaybe _ Nothing _ = Nothing +differenceMaybe _ x Nothing = x +differenceMaybe f (Just x) (Just y) = f x y + +filterLeft :: a -> Either b c -> Maybe b +filterLeft _ (Left x) = Just x +filterLeft _ _ = Nothing + +filterRight :: a -> Either b c -> Maybe c +filterRight _ (Right x) = Just x +filterRight _ _ = Nothing + +{-# INLINE assocsAlg #-} +assocsAlg :: TrieKey k m => m a -> [(k, a)] +assocsAlg m = build (\ c n -> foldWithKeyAlg (\ k v xs -> (k,v) `c` xs) n m) + +instance (Eq a1, Eq a2, TrieKey a1 m1, TrieKey a2 m2) => TrieKey (a1, a2) (m1 `ProdMap` m2) where + emptyAlg = PMap emptyAlg + nullAlg (PMap m) = nullAlg m + sizeAlg (PMap m) = foldl' (\ n m -> n + sizeAlg m) 0 m + getSingleAlg (PMap m) = do (k1, m') <- getSingleAlg m + (k2, v) <- getSingleAlg m' + return ((k1, k2), v) + alterLookupAlg f (k1, k2) (PMap m) = PMap <$> alterLookupAlg g k1 m + where g = fmap guardNullAlg . alterLookupAlg f k2 . fromMaybe emptyAlg + lookupAlg (k1, k2) (PMap m) = lookupAlg k1 m >>= lookupAlg k2 + foldWithKeyAlg f z (PMap m) = foldWithKeyAlg (\ k1 -> flip (foldWithKeyAlg (\ k2 -> f (k1, k2)))) z m + mapAppAlg f (PMap m) = + PMap <$> mapAppAlg (\ k1 -> mapAppAlg (\ k2 -> f (k1, k2))) m + mapMaybeAlg f (PMap m) = + PMap $ mapMaybeAlg (\ k1 -> guardNullAlg . mapMaybeAlg (\ k2 -> f (k1, k2))) m + mapEitherAlg f (PMap m) = (PMap (fmap fst m'), PMap (fmap snd m')) + where m' = mapWithKeyAlg (\ k1 -> mapEitherAlg (\ k2 -> f (k1, k2))) m + unionMaybeAlg f (PMap m1) (PMap m2) = + PMap (unionMaybeAlg (\ k1 -> guardNullAlg .: unionMaybeAlg (\ k2 -> f (k1, k2))) m1 m2) + intersectAlg f (PMap m1) (PMap m2) = + PMap (intersectAlg (\ k1 -> guardNullAlg .: intersectAlg (\ k2 -> f (k1, k2))) m1 m2) + differenceAlg f (PMap m1) (PMap m2) = + PMap (differenceAlg (\ k1 -> guardNullAlg .: differenceAlg (\ k2 -> f (k1, k2))) m1 m2) + fromListAlg f xs = PMap $ mapWithKeyAlg (\ k1 -> fromListAlg (\ k2 -> f (k1, k2))) $ + fromListAlg (const (++)) [(k1, [(k2, v)]) | ((k1, k2), v) <- xs] + fromDistAscListAlg xs = PMap $ fromDistAscListAlg [(k1, fromDistAscListAlg ys) | (k1, ys) <- breakFst xs] + fromAscListAlg f xs = PMap $ fromDistAscListAlg [(k1, fromAscListAlg (\ k2 -> f (k1, k2)) ys) | (k1, ys) <- breakFst xs] + getMinAlg (PMap m) = do + ((k1, m'), m1') <- getMinAlg m + ((k2, v), m2') <- getMinAlg m' + return (((k1, k2), v), PMap (maybe m1' (\ m2' -> insertAlg k1 m2' m) (guardNullAlg m2'))) + getMaxAlg (PMap m) = do + ((k1, m'), m1') <- getMaxAlg m + ((k2, v), m2') <- getMaxAlg m' + return (((k1, k2), v), PMap (maybe m1' (\ m2' -> insertAlg k1 m2' m) (guardNullAlg m2'))) + updateMinAlg f (PMap m) = + PMap <$> updateMinAlg (\ k1 -> guardNullAlg <.> updateMinAlg (\ k2 -> f (k1, k2))) m + updateMaxAlg f (PMap m) = + PMap <$> updateMaxAlg (\ k1 -> guardNullAlg <.> updateMaxAlg (\ k2 -> f (k1, k2))) m + isSubmapAlg (<=) (PMap m1) (PMap m2) = + isSubmapAlg (isSubmapAlg (<=)) m1 m2 + + splitLookupAlg f (k1, k2) (PMap m) = case splitLookupAlg g k1 m of + (mL, ans, mR) -> (PMap mL, ans, PMap mR) + where g m' = case splitLookupAlg f k2 m' of + (mL, ans, mR) -> (guardNullAlg mL, ans, guardNullAlg mR) + + valid (PMap m) = valid m && all valid m && not (any nullAlg m) + +breakFst :: (Eq k1, Eq k2) => [((k1, k2), v)] -> [(k1, [(k2, v)])] +breakFst [] = [] +breakFst (((k1, k2), x):xs) = breakFst' k1 (Seq.singleton (k2, x)) xs where + breakFst' k xs (((k', k2), x):xss) + | k == k' = breakFst' k ((Seq.|>) xs (k2, x)) xss + | otherwise = (k, toList xs):breakFst' k' (Seq.singleton (k2, x)) xss + breakFst' k xs [] = [(k, toList xs)] + +instance (TrieKey a1 m1, TrieKey a2 m2) => TrieKey (Either a1 a2) (m1 `UnionMap` m2) where + emptyAlg = emptyAlg :+: emptyAlg + nullAlg (m1 :+: m2) = nullAlg m1 && nullAlg m2 + sizeAlg (m1 :+: m2) = sizeAlg m1 + sizeAlg m2 + getSingleAlg (m1 :+: m2) = case (getSingleAlg m1, getSingleAlg m2) of + (Just (k, v), Nothing) -> Just (Left k, v) + (Nothing, Just (k, v)) -> Just (Right k, v) + _ -> Nothing + alterLookupAlg f (Left k) (m1 :+: m2) = + fmap (:+: m2) $ alterLookupAlg f k m1 + alterLookupAlg f (Right k) (m1 :+: m2) = + fmap (m1 :+:) $ alterLookupAlg f k m2 + lookupAlg k (m1 :+: m2) = either (`lookupAlg` m1) (`lookupAlg` m2) k + foldWithKeyAlg f z (m1 :+: m2) = foldWithKeyAlg (f . Left) (foldWithKeyAlg (f . Right) z m2) m1 + mapAppAlg f (m1 :+: m2) = + liftA2 (:+:) (mapAppAlg (f . Left) m1) (mapAppAlg (f . Right) m2) + mapMaybeAlg f (m1 :+: m2) = mapMaybeAlg (f . Left) m1 :+: mapMaybeAlg (f . Right) m2 + mapEitherAlg f (m1 :+: m2) = (m1L :+: m2L, m1R :+: m2R) + where (m1L, m1R) = mapEitherAlg (f . Left) m1 + (m2L, m2R) = mapEitherAlg (f . Right) m2 + unionMaybeAlg f (m11 :+: m12) (m21 :+: m22) + = unionMaybeAlg (f . Left) m11 m21 :+: unionMaybeAlg (f . Right) m12 m22 + intersectAlg f (m11 :+: m12) (m21 :+: m22) + = intersectAlg (f . Left) m11 m21 :+: intersectAlg (f . Right) m12 m22 + differenceAlg f (m11 :+: m12) (m21 :+: m22) + = differenceAlg (f . Left) m11 m21 :+: differenceAlg (f . Right) m12 m22 + fromListAlg f xs = fromListAlg (f . Left) ys :+: fromListAlg (f . Right) zs + where (ys, zs) = partitionEithers (map pullEither xs) + fromAscListAlg f xs = fromAscListAlg (f . Left) ys :+: fromAscListAlg (f . Right) zs + where (ys, zs) = partitionEithers (map pullEither xs) + fromDistAscListAlg xs = fromDistAscListAlg ys :+: fromDistAscListAlg zs + where (ys, zs) = partitionEithers (map pullEither xs) + getMinAlg (m1 :+: m2) + | Just ((k, v), m1') <- getMinAlg m1 + = Just ((Left k, v), m1' :+: m2) + | Just ((k, v), m2') <- getMinAlg m2 + = Just ((Right k, v), m1 :+: m2') + getMinAlg _ = Nothing + getMaxAlg (m1 :+: m2) = getFirst $ First + (do ((k, v), m2') <- getMaxAlg m2 + return ((Right k, v), m1 :+: m2')) `mappend` First + (do ((k, v), m1') <- getMaxAlg m1 + return ((Left k, v), m1' :+: m2)) + updateMinAlg f (m1 :+: m2) + | nullAlg m1 = fmap (m1 :+:) (updateMinAlg (f . Right) m2) + | otherwise = fmap (:+: m2) (updateMinAlg (f . Left) m1) + updateMaxAlg f (m1 :+: m2) + | nullAlg m2 = fmap (:+: m2) (updateMaxAlg (f . Left) m1) + | otherwise = fmap (m1 :+:) (updateMaxAlg (f . Right) m2) + isSubmapAlg (<=) (m11 :+: m12) (m21 :+: m22) = + isSubmapAlg (<=) m11 m21 && isSubmapAlg (<=) m12 m22 + valid (m1 :+: m2) = valid m1 && valid m2 + splitLookupAlg f (Left k) (m1 :+: m2) = case splitLookupAlg f k m1 of + (m1L, ans, m1R) -> (m1L :+: emptyAlg, ans, m1R :+: m2) + splitLookupAlg f (Right k) (m1 :+: m2) = case splitLookupAlg f k m2 of + (m2L, ans, m2R) -> (m1 :+: m2L, ans, emptyAlg :+: m2R) + +pullEither :: (Either k1 k2, v) -> Either (k1, v) (k2, v) +pullEither (Left k, v) = Left (k, v) +pullEither (Right k, v) = Right (k, v) + +instance TrieKey Int IntMap where + emptyAlg = IMap.empty + nullAlg = IMap.null + sizeAlg = IMap.size + getSingleAlg m + | IMap.size m == 1, [(k, v)] <- IMap.toList m + = Just (k, v) + getSingleAlg _ = Nothing + lookupAlg = IMap.lookup + alterLookupAlg f k m = fmap (\ v' -> IMap.alter (const v') k m) (f x) + where x = IMap.lookup k m + foldWithKeyAlg = IMap.foldWithKey + mapAppAlg = sequenceA .: IMap.mapWithKey + mapMaybeAlg = IMap.mapMaybeWithKey + mapEitherAlg = IMap.mapEitherWithKey + unionMaybeAlg f m1 m2 = IMap.mapMaybe (either Just id) (IMap.unionWithKey g (fmap Left m1) (fmap Left m2)) where + g k (Left v1) (Left v2) = Right (f k v1 v2) + g k (Right v) _ = Right v + g k _ (Right v) = Right v + intersectAlg f m1 m2 = IMap.mapMaybe (either (const Nothing) Just) $ IMap.intersectionWithKey g (fmap Left m1) m2 where + g k (Left x) = maybe (Left x) Right . f k x + g _ (Right x) = const (Right x) + differenceAlg = IMap.differenceWithKey + fromListAlg = IMap.fromListWithKey + fromAscListAlg = IMap.fromAscListWithKey + fromDistAscListAlg = IMap.fromDistinctAscList + getMinAlg = IMap.minViewWithKey + getMaxAlg = IMap.maxViewWithKey + updateMinAlg f m = case IMap.minViewWithKey m of + Just ((k, v), m') -> let (ans, v') = f k v in (ans, maybe m' (\ v' -> IMap.updateMin (const v') m) v') + _ -> (False, m) + updateMaxAlg f m = case IMap.maxViewWithKey m of + Just ((k, v), m') -> let (ans, v') = f k v in (ans, maybe m' (\ v' -> IMap.updateMax (const v') m) v') + _ -> (False, m) + isSubmapAlg = IMap.isSubmapOfBy + splitLookupAlg f k m = case IMap.splitLookup k m of + (mL, Nothing, mR) -> (mL, Nothing, mR) + (mL, Just v, mR) -> case f v of + (vL, ans, vR) -> (maybe mL (flip (IMap.insert k) mL) vL, ans, maybe mR (flip (IMap.insert k) mR) vR) + +instance Ord k => TrieKey (Ordered k) (Map k) where + emptyAlg = Map.empty + nullAlg = Map.null + sizeAlg = Map.size + getSingleAlg m + | Map.size m == 1, (k, v) <- Map.findMin m + = Just (Ord k, v) + lookupAlg = Map.lookup . unOrd + alterLookupAlg f (Ord k) m = fmap (\ v -> Map.alter (const v) k m) (f x) + where x = Map.lookup k m + foldWithKeyAlg f = Map.foldWithKey (f . Ord) + mapAppAlg f = sequenceA . Map.mapWithKey (f . Ord) + mapMaybeAlg f = Map.mapMaybeWithKey (f . Ord) + mapEitherAlg f = Map.mapEitherWithKey (f . Ord) + unionMaybeAlg f m1 m2 = Map.mapMaybe (either Just id) (Map.unionWithKey g (fmap Left m1) (fmap Left m2)) where + g k (Left v1) (Left v2) = Right (f (Ord k) v1 v2) + g k (Right v) _ = Right v + g k _ (Right v) = Right v + intersectAlg f = Map.mapMaybe id .: Map.intersectionWithKey (f . Ord) + differenceAlg f = Map.differenceWithKey (f . Ord) + fromListAlg f xs = Map.fromListWithKey (f . Ord) [(k, v) | (Ord k, v) <- xs] + fromAscListAlg f xs = Map.fromAscListWithKey (f . Ord) [(k, v) | (Ord k, v) <- xs] + fromDistAscListAlg xs = Map.fromDistinctAscList [(k, v) | (Ord k, v) <- xs] + getMinAlg m = do ~(~(k, v), m') <- Map.minViewWithKey m + return ((Ord k, v), m') + getMaxAlg m = do ~(~(k, v), m') <- Map.maxViewWithKey m + return ((Ord k, v), m') + updateMinAlg f m + | Map.null m = (False, m) + | otherwise = case Map.findMin m of + (k, v) -> let (ans, v') = f (Ord k) v in (ans, Map.updateMin (const v') m) + updateMaxAlg f m + | Map.null m = (False, m) + | otherwise = case Map.findMin m of + (k, v) -> let (ans, v') = f (Ord k) v in (ans, Map.updateMax (const v') m) + isSubmapAlg = Map.isSubmapOfBy + splitLookupAlg f (Ord k) m = case Map.splitLookup k m of + (mL, Nothing, mR) -> (mL, Nothing, mR) + (mL, Just v, mR) -> case f v of + (vL, ans, vR) -> (maybe mL (flip (Map.insert k) mL) vL, ans, maybe mR (flip (Map.insert k) mR) vR) + +instance TrieKey () Maybe where + emptyAlg = Nothing + nullAlg = isNothing + sizeAlg = maybe 0 (const 1) + getSingleAlg = fmap ((,) ()) + lookupAlg _ = id + alterLookupAlg f _ = f + foldWithKeyAlg f = foldr (f ()) + mapAppAlg f = traverse (f ()) + mapMaybeAlg f = (>>= f ()) + mapEitherAlg _ Nothing = (Nothing, Nothing) + mapEitherAlg f (Just v) = case f () v of + Left v -> (Just v, Nothing) + Right v -> (Nothing, Just v) + unionMaybeAlg f = unionMaybe (f ()) + intersectAlg f = intersectMaybe (f ()) + differenceAlg f = differenceMaybe (f ()) + fromListAlg _ [] = Nothing + fromListAlg f ((_, v):xs) = Just (foldr (f () . snd) v xs) + fromAscListAlg = fromListAlg + getMinAlg = fmap g where + g v = (((), v), Nothing) + getMaxAlg = fmap g where + g v = (((), v), Nothing) + updateMinAlg f = maybe (False, Nothing) (f ()) + updateMaxAlg f = maybe (False, Nothing) (f ()) + isSubmapAlg _ Nothing _ = True + isSubmapAlg _ _ Nothing = False + isSubmapAlg (<=) (Just x) (Just y) = x <= y + splitLookupAlg f _ (Just v) = f v + splitLookupAlg _ _ _ = (Nothing, Nothing, Nothing) + +first :: (a -> c) -> (a, b) -> (c, b) +first f (x, y) = (f x, y)