diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,3 @@ +# Changelog for multi-containers + +## Unreleased changes diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Author name here (c) 2019 + +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. + + * Neither the name of Author name here nor the names of other + contributors may 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/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,1 @@ +# multi-containers 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/multi-containers.cabal b/multi-containers.cabal new file mode 100644 --- /dev/null +++ b/multi-containers.cabal @@ -0,0 +1,73 @@ +-- This file has been generated from package.yaml by hpack version 0.28.2. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: ee8ee976df6098bafd1b84099c480eac4bab9fed793ae074100bd0ab676cdc17 + +name: multi-containers +version: 0.1.0.0 +synopsis: A few variants of multimaps. +description: Please see the README on GitHub at +category: Data Structures +homepage: https://github.com/zliu41/multi-containers#readme +bug-reports: https://github.com/zliu41/multi-containers/issues +author: Ziyang Liu +maintainer: free@cofree.io +copyright: 2019 Ziyang Liu +license: BSD3 +license-file: LICENSE +build-type: Simple +cabal-version: >= 1.10 +extra-source-files: + ChangeLog.md + README.md + +source-repository head + type: git + location: https://github.com/zliu41/multi-containers + +library + exposed-modules: + Data.Multimap + Data.Multimap.Set + Data.Multimap.Table + other-modules: + Paths_multi_containers + hs-source-dirs: + src + build-depends: + base >=4.7 && <5 + , containers >=0.5.10.2 && <0.7 + default-language: Haskell2010 + +executable test-gen + main-is: Main.hs + other-modules: + Paths_multi_containers + hs-source-dirs: + test-gen + build-depends: + base >=4.7 && <5 + , containers >=0.5.10.2 && <0.7 + , directory >=1.3.0.2 && <1.4 + , extra >=1.6.9 && <1.7 + , filepath >=1.4.1.2 && <1.5 + default-language: Haskell2010 + +test-suite hspec + type: exitcode-stdio-1.0 + main-is: Main.hs + other-modules: + Data.Multimap.SetSpec + Data.Multimap.TableSpec + Data.MultimapSpec + Paths_multi_containers + hs-source-dirs: + test/hspec + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + base >=4.7 && <5 + , containers >=0.5.10.2 && <0.7 + , hspec >=2.4.8 && <2.8 + , multi-containers + default-language: Haskell2010 diff --git a/src/Data/Multimap.hs b/src/Data/Multimap.hs new file mode 100644 --- /dev/null +++ b/src/Data/Multimap.hs @@ -0,0 +1,745 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Multimap +-- Maintainer : Ziyang Liu +-- +-- Multimaps, where values behave like (non empty) lists. +-- +-- Multimaps whose values behave like sets are in "Data.Multimap.Set". +-- Multimaps whose values behave like maps (i.e., two-dimensional +-- tables) are in "Data.Multimap.Table". +-- +-- The implementation is backed by a @'Map' k ('NonEmpty' a)@. The +-- differences between @'Multimap' k a@ and @'Map' k ('NonEmpty' a)@ include: +-- +-- * 'lookup' (or '!') returns a possibly empty list. Unlike regular maps, +-- the '!' operator is total for multimaps. +-- +-- * Functions like 'map', 'adjust', 'traverse', etc., take functions on +-- individual values (e.g., @a -> b@) as opposed to, e.g., +-- @'NonEmpty' a -> 'NonEmpty' b@. +-- +-- * 'union' and 'unions' concatenate the values when there are duplicate +-- keys, rather than being left- or right-biased. +-- +-- * The 'difference' function computes list differences for values of +-- keys that exist in both maps. +-- +-- * The 'size' function returns the total number of values for all keys in +-- the multimap, not the number of distinct keys. The latter can be obtained +-- by first getting the 'keysSet' or first converting the multimap to +-- a regular map via 'toMap'. +-- +-- In the following Big-O notations, unless otherwise noted, /n/ denotes +-- the size of the multimap, /k/ denotes the number of distinct keys, and +-- /m/ denotes the maximum number of values associated with a single key. +module Data.Multimap ( + -- * Multimap type + Multimap + + -- * Construction + , empty + , singleton + , fromMap + , fromMap' + + -- ** From Unordered Lists + , fromList + + -- * Insertion + , insert + + -- * Deletion\/Update + , delete + , deleteWithValue + , deleteOne + , adjust + , adjustWithKey + , update + , update' + , updateWithKey + , updateWithKey' + , alter + , alterWithKey + + -- * Query + -- ** Lookup + , lookup + , (!) + , member + , notMember + + -- ** Size + , null + , notNull + , size + + -- * Combine + -- ** Union + , union + , unions + + -- ** Difference + , difference + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , traverseMaybeWithKey + + -- ** Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict Folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + + -- ** Lists + , toList + + -- ** Ordered lists + , toAscList + , toDescList + , toAscListBF + , toDescListBF + + -- ** Maps + , toMap + + -- * Filter + , filter + , filterWithKey + , filterKey + , filterM + , filterWithKeyM + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + ) where + +import Control.Arrow ((&&&)) +import Control.Monad (join) +import qualified Control.Monad as List (filterM) +import Data.Data (Data) +import qualified Data.Either as Either +import qualified Data.Foldable as Foldable +import Data.Functor.Classes +import qualified Data.List as List +import Data.List.NonEmpty (NonEmpty(..), (<|), nonEmpty) +import qualified Data.List.NonEmpty as Nel +import Data.Map.Lazy (Map) +import qualified Data.Map.Lazy as Map +import qualified Data.Maybe as Maybe +import Data.Semigroup (Semigroup, (<>)) +import Data.Set (Set) + +import Prelude hiding (filter, foldl, foldr, lookup, map, null) + +infixl 9 ! + +type Size = Int + +newtype Multimap k a = Multimap (Map k (NonEmpty a), Size) + deriving (Eq, Ord, Data) + +instance Eq k => Eq1 (Multimap k) where + liftEq = liftEq2 (==) + +instance Eq2 Multimap where + liftEq2 eqk eqv m n = + Map.size (toMap m) == Map.size (toMap n) + && liftEq (liftEq2 eqk eqv) (toList m) (toList n) + +instance Ord k => Ord1 (Multimap k) where + liftCompare = liftCompare2 compare + +instance Ord2 Multimap where + liftCompare2 cmpk cmpv m n = + liftCompare (liftCompare2 cmpk cmpv) (toList m) (toList n) + +instance (Show k, Show a) => Show (Multimap k a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +instance Show k => Show1 (Multimap k) where + liftShowsPrec = liftShowsPrec2 showsPrec showList + +instance Show2 Multimap where + liftShowsPrec2 spk slk spv slv d m = + showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) + where + sp = liftShowsPrec2 spk slk spv slv + sl = liftShowList2 spk slk spv slv + +instance (Ord k, Read k, Read e) => Read (Multimap k e) where + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + pure (fromList xs,t) + +instance (Ord k, Read k) => Read1 (Multimap k) where + liftReadsPrec rp rl = readsData $ + readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList + where + rp' = liftReadsPrec rp rl + rl' = liftReadList rp rl + +instance Functor (Multimap k) where + fmap = map + +instance Foldable.Foldable (Multimap k) where + foldMap = foldMapWithKey . const + {-# INLINE foldMap #-} + +instance Traversable (Multimap k) where + traverse = traverseWithKey . const + {-# INLINE traverse #-} + +instance (Ord k) => Semigroup (Multimap k a) where + (<>) = union + +instance (Ord k) => Monoid (Multimap k a) where + mempty = empty + mappend = (<>) + +------------------------------------------------------------------------------ + +-- | /O(1)/. The empty multimap. +-- +-- > size empty === 0 +empty :: Multimap k a +empty = Multimap (Map.empty, 0) + +-- | /O(1)/. A multimap with a single element. +-- +-- > singleton 1 'a' === fromList [(1, 'a')] +-- > size (singleton 1 'a') === 1 +singleton :: k -> a -> Multimap k a +singleton k a = Multimap (Map.singleton k (pure a), 1) + +-- | /O(n*log n)/ where /n/ is the length of the input list. +-- Build a multimap from a list of key\/value pairs. +-- +-- > fromList ([] :: [(Int, Char)]) === empty +fromList :: Ord k => [(k, a)] -> Multimap k a +fromList = Foldable.foldr (uncurry insert) empty + +-- | /O(1)/. +fromMap :: Map k (NonEmpty a) -> Multimap k a +fromMap m = Multimap (m, sum (fmap length m)) + +-- | /O(k)/. A key is retained only if it is associated with a +-- non-empty list of values. +-- +-- > fromMap' (Map.fromList [(1, "ab"), (2, ""), (3, "c")]) === fromList [(1, 'a'), (1, 'b'), (3, 'c')] +fromMap' :: Map k [a] -> Multimap k a +fromMap' m = Multimap (Map.mapMaybe nonEmpty m, sum (fmap length m)) + +------------------------------------------------------------------------------ + +-- | /O(log k)/. If the key exists in the multimap, the new value will be +-- prepended to the list of values for the key. +-- +-- > insert 1 'a' empty === singleton 1 'a' +-- > insert 1 'a' (fromList [(2, 'b'), (2, 'c')]) === fromList [(1, 'a'), (2, 'b'), (2, 'c')] +-- > insert 1 'a' (fromList [(1, 'b'), (2, 'c')]) === fromList [(1, 'a'), (1, 'b'), (2, 'c')] +insert :: Ord k => k -> a -> Multimap k a -> Multimap k a +insert k a (Multimap (m, _)) = fromMap (Map.alter f k m) + where + f (Just as) = Just (a <| as) + f Nothing = Just (pure a) + +-- | /O(log k)/. Delete a key and all its values from the map. +-- +-- > delete 1 (fromList [(1, 'a'), (1, 'b'), (2, 'c')]) === singleton 2 'c' +delete :: Ord k => k -> Multimap k a -> Multimap k a +delete = update' (const []) + +-- | /O(m*log k)/. Remove the first +-- occurrence of the value associated with the key, if exists. +-- +-- > deleteWithValue 1 'c' (fromList [(1, 'a'), (1, 'b'), (2, 'c')]) === fromList [(1, 'a'), (1, 'b'), (2, 'c')] +-- > deleteWithValue 1 'c' (fromList [(1, 'a'), (1, 'b'), (2, 'c'), (1, 'c')]) === fromList [(1, 'a'), (1, 'b'), (2, 'c')] +-- > deleteWithValue 1 'c' (fromList [(2, 'c'), (1, 'c')]) === singleton 2 'c' +deleteWithValue :: (Ord k, Eq a) => k -> a -> Multimap k a -> Multimap k a +deleteWithValue k a = update' (List.delete a . Nel.toList) k + +-- | /O(log k)/. Remove the first +-- value associated with the key. If the key is associated with a single value, +-- the key will be removed from the multimap. +-- +-- > deleteOne 1 (fromList [(1, 'a'), (1, 'b'), (2, 'c')]) === fromList [(1, 'b'), (2, 'c')] +-- > deleteOne 1 (fromList [(2, 'c'), (1, 'c')]) === singleton 2 'c' +deleteOne :: Ord k => k -> Multimap k a -> Multimap k a +deleteOne = update' Nel.tail + +-- | /O(m*log k)/, assuming the function @a -> a@ takes /O(1)/. +-- Update values at a specific key, if exists. +-- +-- > adjust ("new " ++) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"new a"),(1,"new b"),(2,"c")] +adjust :: Ord k => (a -> a) -> k -> Multimap k a -> Multimap k a +adjust = adjustWithKey . const + +-- | /O(m*log k)/, assuming the function @k -> a -> a@ takes /O(1)/. +-- Update values at a specific key, if exists. +-- +-- > adjustWithKey (\k x -> show k ++ ":new " ++ x) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) +-- > === fromList [(1,"1:new a"),(1,"1:new b"),(2,"c")] +adjustWithKey :: Ord k => (k -> a -> a) -> k -> Multimap k a -> Multimap k a +adjustWithKey f k (Multimap (m, sz)) = Multimap (m', sz) + where + m' = Map.adjustWithKey (fmap . f) k m + +-- | /O(m*log k)/, assuming the function @a -> 'Maybe' a@ takes /O(1)/. +-- The expression (@'update' f k map@) updates the values at key @k@, if +-- exists. If @f@ returns 'Nothing' for a value, the value is deleted. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing in do +-- > update f 1 (fromList [(1,"a"),(1, "b"),(2,"c")]) === fromList [(1,"new a"),(2, "c")] +-- > update f 1 (fromList [(1,"b"),(1, "b"),(2,"c")]) === singleton 2 "c" +update :: Ord k => (a -> Maybe a) -> k -> Multimap k a -> Multimap k a +update = updateWithKey . const + +-- | /O(log k)/, assuming the function @'NonEmpty' a -> [a]@ takes /O(1)/. +-- The expression (@'update' f k map@) updates the values at key @k@, if +-- exists. If @f@ returns 'Nothing', the key is deleted. +-- +-- > update' NonEmpty.tail 1 (fromList [(1, "a"), (1, "b"), (2, "c")]) === fromList [(1, "b"), (2, "c")] +-- > update' NonEmpty.tail 1 (fromList [(1, "a"), (2, "b")]) === singleton 2 "b" +update' :: Ord k => (NonEmpty a -> [a]) -> k -> Multimap k a -> Multimap k a +update' = updateWithKey' . const + +-- | /O(m*log k)/, assuming the function @k -> a -> 'Maybe' a@ takes /O(1)/. +-- The expression (@'updateWithKey' f k map@) updates the values at key @k@, if +-- exists. If @f@ returns 'Nothing' for a value, the value is deleted. +-- +-- > let f k x = if x == "a" then Just (show k ++ ":new a") else Nothing in do +-- > updateWithKey f 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"1:new a"),(2,"c")] +-- > updateWithKey f 1 (fromList [(1,"b"),(1,"b"),(2,"c")]) === singleton 2 "c" +updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Multimap k a -> Multimap k a +updateWithKey f = alterWithKey (Maybe.mapMaybe . f) + +-- | /O(log k)/, assuming the function @k -> 'NonEmpty' a -> [a]@ takes /O(1)/. +-- The expression (@'update' f k map@) updates the values at key @k@, if +-- exists. If @f@ returns 'Nothing', the key is deleted. +-- +-- > let f k xs = if NonEmpty.length xs == 1 then (show k : NonEmpty.toList xs) else [] in do +-- > updateWithKey' f 1 (fromList [(1, "a"), (1, "b"), (2, "c")]) === singleton 2 "c" +-- > updateWithKey' f 1 (fromList [(1, "a"), (2, "b"), (2, "c")]) === fromList [(1, "1"), (1, "a"), (2, "b"), (2, "c")] +updateWithKey' :: Ord k => (k -> NonEmpty a -> [a]) -> k -> Multimap k a -> Multimap k a +updateWithKey' f = alterWithKey g + where + g _ [] = [] + g k (a:as) = f k (a :| as) + +-- | /O(log k)/, assuming the function @[a] -> [a]@ takes /O(1)/. +-- The expression (@'alter' f k map@) alters the values at k, if exists. +-- +-- > let (f, g) = (const [], ('c':)) in do +-- > alter f 1 (fromList [(1, 'a'), (2, 'b')]) === singleton 2 'b' +-- > alter f 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b')] +-- > alter g 1 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'c'), (1, 'a'), (2, 'b')] +-- > alter g 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b'), (3, 'c')] +alter :: Ord k => ([a] -> [a]) -> k -> Multimap k a -> Multimap k a +alter = alterWithKey . const + +-- | /O(log k)/, assuming the function @k -> [a] -> [a]@ takes /O(1)/. +-- The expression (@'alterWithKey' f k map@) alters the values at k, if exists. +-- +-- > let (f, g) = (const (const []), (:) . show) in do +-- > alterWithKey f 1 (fromList [(1, "a"), (2, "b")]) === singleton 2 "b" +-- > alterWithKey f 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b")] +-- > alterWithKey g 1 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "1"), (1, "a"), (2, "b")] +-- > alterWithKey g 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b"), (3, "3")] +alterWithKey :: Ord k => (k -> [a] -> [a]) -> k -> Multimap k a -> Multimap k a +alterWithKey f k mm@(Multimap (m, _)) = case nonEmpty (f k (mm ! k)) of + Just as' -> fromMap (Map.insert k as' m) + Nothing -> fromMap (Map.delete k m) + +------------------------------------------------------------------------------ + +-- | /O(log k)/. Lookup the values at a key in the map. It returns an empty +-- list if the key is not in the map. +lookup :: Ord k => k -> Multimap k a -> [a] +lookup k (Multimap (m, _)) = maybe [] Nel.toList (Map.lookup k m) + +-- | /O(log k)/. Lookup the values at a key in the map. It returns an empty +-- list if the key is not in the map. +-- +-- > fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 3 === "ac" +-- > fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 2 === [] +(!) :: Ord k => Multimap k a -> k -> [a] +(!) = flip lookup + +-- | /O(log k)/. Is the key a member of the map? +-- +-- A key is a member of the map if and only if there is at least one value +-- associated with it. +-- +-- > member 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === True +-- > member 1 (deleteOne 1 (fromList [(2, 'c'), (1, 'c')])) === False +member :: Ord k => k -> Multimap k a -> Bool +member k (Multimap (m, _)) = Map.member k m + +-- | /O(log k)/. Is the key not a member of the map? +-- +-- A key is a member of the map if and only if there is at least one value +-- associated with it. +-- +-- > notMember 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === False +-- > notMember 1 (deleteOne 1 (fromList [(2, 'c'), (1, 'c')])) === True +notMember :: Ord k => k -> Multimap k a -> Bool +notMember k = not . member k + +-- | /O(1)/. Is the multimap empty? +-- +-- > Data.Multimap.null empty === True +-- > Data.Multimap.null (singleton 1 'a') === False +null :: Multimap k a -> Bool +null (Multimap (m, _)) = Map.null m + +-- | /O(1)/. Is the multimap non-empty? +-- +-- > notNull empty === False +-- > notNull (singleton 1 'a') === True +notNull :: Multimap k a -> Bool +notNull = not . null + +-- | The total number of values for all keys. +-- +-- @size@ is evaluated lazily. Forcing the size for the first time takes up to +-- /O(n)/ and subsequent forces take /O(1)/. +-- +-- > size empty === 0 +-- > size (singleton 1 'a') === 1 +-- > size (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === 3 +size :: Multimap k a -> Int +size (Multimap (_, sz)) = sz + +------------------------------------------------------------------------------ + +-- | Union two multimaps, concatenating values for duplicate keys. +-- +-- > union (fromList [(1,'a'),(2,'b'),(2,'c')]) (fromList [(1,'d'),(2,'b')]) +-- > === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c'),(2,'b')] +union :: Ord k => Multimap k a -> Multimap k a -> Multimap k a +union (Multimap (m1, _)) (Multimap (m2, _)) = + fromMap (Map.unionWith (<>) m1 m2) + +-- | Union a number of multimaps, concatenating values for duplicate keys. +-- +-- > unions [fromList [(1,'a'),(2,'b'),(2,'c')], fromList [(1,'d'),(2,'b')]] +-- > === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c'),(2,'b')] +unions :: (Foldable f, Ord k) => f (Multimap k a) -> Multimap k a +unions = Foldable.foldr union empty + +-- | Difference of two multimaps. +-- +-- If a key exists in the first multimap but not the second, it remains +-- unchanged in the result. If a key exists in both multimaps, a list +-- difference is performed on their values, i.e., the first occurrence +-- of each value in the second multimap is removed from the +-- first multimap. +-- +-- > difference (fromList [(1,'a'),(2,'b'),(2,'c'),(2,'b')]) (fromList [(1,'d'),(2,'b'),(2,'a')]) +-- > === fromList [(1,'a'), (2,'c'), (2,'b')] +difference :: (Ord k, Eq a) => Multimap k a -> Multimap k a -> Multimap k a +difference (Multimap (m1, _)) (Multimap (m2, _)) = fromMap $ + Map.differenceWith (\xs ys -> nonEmpty (Nel.toList xs List.\\ Nel.toList ys)) m1 m2 + +------------------------------------------------------------------------------ + +-- | /O(n)/, assuming the function @a -> b@ takes /O(1)/. +-- Map a function over all values in the map. +-- +-- > Data.Multimap.map (++ "x") (fromList [(1,"a"),(1,"a"),(2,"b")]) === fromList [(1,"ax"),(1,"ax"),(2,"bx")] +map :: (a -> b) -> Multimap k a -> Multimap k b +map = mapWithKey . const + +-- | /O(n)/, assuming the function @k -> a -> b@ takes /O(1)/. +-- Map a function over all key\/value pairs in the map. +-- +-- > mapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1,"a"),(1,"a"),(2,"b")]) === fromList [(1,"1:a"),(1,"1:a"),(2,"2:b")] +mapWithKey :: (k -> a -> b) -> Multimap k a -> Multimap k b +mapWithKey f (Multimap (m, sz)) = Multimap (Map.mapWithKey (fmap . f) m, sz) + +-- | Traverse key\/value pairs and collect the results. +-- +-- > let f k a = if odd k then Just (succ a) else Nothing in do +-- > traverseWithKey f (fromList [(1, 'a'), (1, 'b'), (3, 'b'), (3, 'c')]) === Just (fromList [(1, 'b'), (1, 'c'), (3, 'c'), (3, 'd')]) +-- > traverseWithKey f (fromList [(1, 'a'), (1, 'b'), (2, 'b')]) === Nothing +traverseWithKey :: Applicative t => (k -> a -> t b) -> Multimap k a -> t (Multimap k b) +traverseWithKey f (Multimap (m, _)) = + fromMap <$> Map.traverseWithKey (traverse . f) m + +-- | Traverse key\/value pairs and collect the 'Just' results. +traverseMaybeWithKey :: Applicative t => (k -> a -> t (Maybe b)) -> Multimap k a -> t (Multimap k b) +traverseMaybeWithKey f (Multimap (m, _)) = + fromMap <$> Map.traverseMaybeWithKey f' m + where + f' k = fmap (nonEmpty . Maybe.catMaybes) . traverse (f k) . Nel.toList + +------------------------------------------------------------------------------ + +-- | /O(n)/. Fold the values in the map using the given right-associative +-- binary operator. +-- +-- > Data.Multimap.foldr ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldr :: (a -> b -> b) -> b -> Multimap k a -> b +foldr = foldrWithKey . const + +-- | /O(n)/. Fold the values in the map using the given left-associative +-- binary operator. +-- +-- > Data.Multimap.foldl (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldl :: (a -> b -> a) -> a -> Multimap k b -> a +foldl = foldlWithKey . (const .) + +-- | /O(n)/. Fold the key\/value paris in the map using the given +-- right-associative binary operator. +-- +-- > foldrWithKey (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldrWithKey :: (k -> a -> b -> b) -> b -> Multimap k a -> b +foldrWithKey f b (Multimap (m, _)) = Map.foldrWithKey f' b m + where + f' = flip . Foldable.foldr . f + +-- | /O(n)/. Fold the key\/value paris in the map using the given +-- left-associative binary operator. +-- +-- > foldlWithKey (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldlWithKey :: (a -> k -> b -> a) -> a -> Multimap k b -> a +foldlWithKey f a (Multimap (m, _)) = Map.foldlWithKey f' a m + where + f' = flip (Foldable.foldl . flip f) + +-- | /O(n)/. A strict version of 'foldr'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > Data.Multimap.foldr' ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldr' :: (a -> b -> b) -> b -> Multimap k a -> b +foldr' = foldrWithKey' . const + +-- | /O(n)/. A strict version of 'foldl'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > Data.Multimap.foldl' (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldl' :: (a -> b -> a) -> a -> Multimap k b -> a +foldl' = foldlWithKey' . (const .) + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > foldrWithKey' (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldrWithKey' :: (k -> a -> b -> b) -> b -> Multimap k a -> b +foldrWithKey' f b (Multimap (m, _)) = Map.foldrWithKey' f' b m + where + f' = flip . Foldable.foldr . f + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > foldlWithKey' (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldlWithKey' :: (a -> k -> b -> a) -> a -> Multimap k b -> a +foldlWithKey' f a (Multimap (m, _)) = Map.foldlWithKey' f' a m + where + f' = flip (Foldable.foldl' . flip f) + +-- | /O(n)/. Fold the key\/value pairs in the map using the given monoid. +-- +-- > foldMapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1, "a"), (1, "a"), (2, "b")]) === "1:a1:a2:b" +foldMapWithKey :: Monoid m => (k -> a -> m) -> Multimap k a -> m +foldMapWithKey f (Multimap (m, _)) = Map.foldMapWithKey f' m + where + f' = Foldable.foldMap . f + +------------------------------------------------------------------------------ + +-- | /O(n)/. Return all elements of the multimap in ascending order of +-- their keys. +-- +-- > elems (fromList [(2, 'a'), (1, 'b'), (3, 'c'), (1, 'b')]) === "bbac" +-- > elems (empty :: Multimap Int Char) === [] +elems :: Multimap k a -> [a] +elems (Multimap (m, _)) = Map.elems m >>= Nel.toList + +-- | /O(k)/. Return all keys of the multimap in ascending order. +-- +-- > keys (fromList [(2, 'a'), (1, 'b'), (3, 'c'), (1, 'b')]) === [1,2,3] +-- > keys (empty :: Multimap Int Char) === [] +keys :: Multimap k a -> [k] +keys (Multimap (m, _)) = Map.keys m + +-- | /O(k)/. The set of all keys of the multimap. +-- +-- > keysSet (fromList [(2, 'a'), (1, 'b'), (3, 'c'), (1, 'b')]) === Set.fromList [1,2,3] +-- > keysSet (empty :: Multimap Int Char) === Set.empty +keysSet :: Multimap k a -> Set k +keysSet (Multimap (m, _)) = Map.keysSet m + +-- | An alias for 'toAscList'. +-- +-- > assocs (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'b'),(1,'a'),(2,'a'),(3,'c')] +assocs :: Multimap k a -> [(k, a)] +assocs = toAscList + +-- | Convert the multimap into a list of key/value pairs. +-- +-- > toList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'b'),(1,'a'),(2,'a'),(3,'c')] +toList :: Multimap k a -> [(k, a)] +toList = toAscList + +-- | Convert the multimap into a list of key/value pairs in ascending +-- order of keys. +-- +-- > toAscList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'b'),(1,'a'),(2,'a'),(3,'c')] +toAscList :: Multimap k a -> [(k, a)] +toAscList (Multimap (m, _)) = + Map.toAscList m >>= uncurry (\k -> fmap (k,) . Nel.toList) + +-- | Convert the multimap into a list of key/value pairs in descending +-- order of keys. +-- +-- > toDescList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(3,'c'),(2,'a'),(1,'b'),(1,'a')] +toDescList :: Multimap k a -> [(k, a)] +toDescList (Multimap (m, _)) = + Map.toDescList m >>= uncurry (\k -> fmap (k,) . Nel.toList) + +-- | Convert the multimap into a list of key/value pairs, in a +-- breadth-first manner, in ascending order of keys. +-- +-- > toAscListBF (fromList [("Foo",1),("Foo",2),("Foo",3),("Bar",4),("Bar",5),("Baz",6)]) +-- > === [("Bar",4),("Baz",6),("Foo",1),("Bar",5),("Foo",2),("Foo",3)] +toAscListBF :: Multimap k a -> [(k, a)] +toAscListBF (Multimap (m, _)) = + join + . List.transpose + . fmap (uncurry (\k -> fmap (k,) . Nel.toList)) + $ Map.toAscList m + +-- | Convert the multimap into a list of key/value pairs, in a +-- breadth-first manner, in descending order of keys. +-- +-- > toDescListBF (fromList [("Foo",1),("Foo",2),("Foo",3),("Bar",4),("Bar",5),("Baz",6)]) +-- > === [("Foo",1),("Baz",6),("Bar",4),("Foo",2),("Bar",5),("Foo",3)] +toDescListBF :: Multimap k a -> [(k, a)] +toDescListBF (Multimap (m, _)) = + join + . List.transpose + . fmap (uncurry (\k -> fmap (k,) . Nel.toList)) + $ Map.toDescList m + +-- | /O(1)/. Convert the multimap into a regular map. +toMap :: Multimap k a -> Map k (NonEmpty a) +toMap (Multimap (m, _)) = m + +------------------------------------------------------------------------------ + +-- | /O(n)/, assuming the predicate function takes /O(1)/. +-- Retain all values that satisfy the predicate. +-- +-- > Data.Multimap.filter (> 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === singleton 1 'b' +-- > Data.Multimap.filter (< 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === empty +filter :: (a -> Bool) -> Multimap k a -> Multimap k a +filter = filterWithKey . const + +-- | /O(k)/, assuming the predicate function takes /O(1)/. +-- Retain all keys that satisfy the predicate. +-- +-- > filterKey even (fromList [(1,'a'),(1,'b'),(2,'a')]) === singleton 2 'a' +filterKey :: (k -> Bool) -> Multimap k a -> Multimap k a +filterKey p (Multimap (m, _)) = fromMap m' + where + m' = Map.filterWithKey (const . p) m + +-- | /O(n)/, assuming the predicate function takes /O(1)/. +-- Retain all key\/value pairs that satisfy the predicate. +-- +-- > filterWithKey (\k a -> even k && a > 'a') (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'b')]) === singleton 2 'b' +filterWithKey :: (k -> a -> Bool) -> Multimap k a -> Multimap k a +filterWithKey p (Multimap (m, _)) = fromMap m' + where + m' = Map.mapMaybeWithKey (\k -> nonEmpty . Nel.filter (p k)) m + +-- | Generalized 'filter'. +-- +-- > let f a | a > 'b' = Just True +-- > | a < 'b' = Just False +-- > | a == 'b' = Nothing +-- > in do +-- > filterM f (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'c')]) === Nothing +-- > filterM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Just (fromList [(1,'c'),(2,'c')]) +filterM :: (Ord k, Applicative t) => (a -> t Bool) -> Multimap k a -> t (Multimap k a) +filterM = filterWithKeyM . const + +-- | Generalized 'filterWithKey'. +-- +-- > let f k a | even k && a > 'b' = Just True +-- > | odd k && a < 'b' = Just False +-- > | otherwise = Nothing +-- > in do +-- > filterWithKeyM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Nothing +-- > filterWithKeyM f (fromList [(1,'a'),(1,'a'),(2,'c'),(2,'c')]) === Just (fromList [(2,'c'),(2,'c')]) +filterWithKeyM :: (Ord k, Applicative t) => (k -> a -> t Bool) -> Multimap k a -> t (Multimap k a) +filterWithKeyM f = fmap fromList . List.filterM (uncurry f) . toList + +-- | /O(n)/, assuming the function @a -> 'Maybe' b@ takes /O(1)/. +-- Map values and collect the 'Just' results. +-- +-- > mapMaybe (\a -> if a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) +-- > === fromList [(1,"new a"),(2,"new a")] +mapMaybe :: (a -> Maybe b) -> Multimap k a -> Multimap k b +mapMaybe = mapMaybeWithKey . const + +-- | /O(n)/, assuming the function @k -> a -> 'Maybe' b@ takes /O(1)/. +-- Map key\/value pairs and collect the 'Just' results. +-- +-- > mapMaybeWithKey (\k a -> if k > 1 && a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) +-- > === singleton 2 "new a" +mapMaybeWithKey :: (k -> a -> Maybe b) -> Multimap k a -> Multimap k b +mapMaybeWithKey f (Multimap (m, _)) = fromMap $ + Map.mapMaybeWithKey (\k -> nonEmpty . Maybe.mapMaybe (f k) . Nel.toList) m + +-- | /O(n)/, assuming the function @a -> 'Either' b c@ takes /O(1)/. +-- Map values and separate the 'Left' and 'Right' results. +-- +-- > mapEither (\a -> if a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) +-- > === (fromList [(1,'a'),(2,'a')],fromList [(1,'c'),(2,'c')]) +mapEither :: (a -> Either b c) -> Multimap k a -> (Multimap k b, Multimap k c) +mapEither = mapEitherWithKey . const + +-- | /O(n)/, assuming the function @k -> a -> 'Either' b c@ takes /O(1)/. +-- Map key\/value pairs and separate the 'Left' and 'Right' results. +-- +-- > mapEitherWithKey (\k a -> if even k && a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) +-- > === (fromList [(2,'a')],fromList [(1,'a'),(1,'c'),(2,'c')]) +mapEitherWithKey :: (k -> a -> Either b c) -> Multimap k a -> (Multimap k b, Multimap k c) +mapEitherWithKey f (Multimap (m, _)) = + (fromMap' . Map.mapWithKey (const fst) &&& fromMap' . Map.mapWithKey (const snd)) + $ Map.mapWithKey g m + where + g k as = Either.partitionEithers $ fmap (f k) (Nel.toList as) diff --git a/src/Data/Multimap/Set.hs b/src/Data/Multimap/Set.hs new file mode 100644 --- /dev/null +++ b/src/Data/Multimap/Set.hs @@ -0,0 +1,705 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Multimap.Set +-- Maintainer : Ziyang Liu +-- +-- Multimaps, where values behave like (non empty) sets. +-- +-- Multimaps whose values behave like lists are in "Data.Multimap". +-- Multimaps whose values behave like maps (i.e., two-dimensional +-- tables) are in "Data.Multimap.Table". +-- +-- The implementation is backed by a @'Map' k ('Set' a)@. The +-- differences between @'Multimap' k a@ and @'Map' k ('Set' a)@ include: +-- +-- * A key is only present in a 'SetMultimap' if it is associated with +-- at least one values, i.e., a key is never associated with an empty +-- set of values. +-- +-- * 'lookup' (or '!') returns a possibly empty set. Unlike regular maps, +-- the '!' operator is total for multimaps. +-- +-- * Functions like 'map', 'adjust', 'update', etc., take functions on +-- individual values (e.g., @a -> b@) as opposed to, e.g., +-- @'Set' a -> 'Set' b@. +-- +-- * 'union' and 'unions' union the values when there are duplicate +-- keys, rather than being left- or right-biased. +-- +-- * The 'difference' function computes set differences for values of +-- keys that exist in both maps. +-- +-- * The 'size' function returns the total number of values for all keys in +-- the multimap, not the number of distinct keys. The latter can be obtained +-- by first getting the 'keysSet' or first converting the multimap to +-- a regular map via 'toMap'. +-- +-- In the following Big-O notations, unless otherwise noted, /n/ denotes +-- the size of the multimap, /k/ denotes the number of distinct keys, and +-- /m/ denotes the maximum number of values associated with a single key. +module Data.Multimap.Set ( + -- * Multimap type + SetMultimap + + -- * Construction + , empty + , singleton + , fromMap + + -- ** From Unordered Lists + , fromList + + -- * Insertion + , insert + + -- * Deletion\/Update + , delete + , deleteWithValue + , deleteMax + , deleteMin + , adjust + , adjustWithKey + , update + , updateWithKey + , alter + , alterWithKey + + -- * Query + -- ** Lookup + , lookup + , (!) + , member + , notMember + + -- ** Size + , null + , notNull + , size + + -- * Combine + -- ** Union + , union + , unions + + -- ** Difference + , difference + + -- * Traversal + -- ** Map + , map + , mapWithKey + + -- ** Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict Folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + + -- ** Lists + , toList + + -- ** Ordered lists + , toAscList + , toDescList + + -- ** Maps + , toMap + + -- * Filter + , filter + , filterWithKey + , filterKey + , filterM + , filterWithKeyM + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + ) where + +import Prelude hiding (filter, foldl, foldr, lookup, map, null) + +import Control.Arrow ((&&&)) +import qualified Control.Monad as List (filterM) +import Data.Data (Data) +import qualified Data.Foldable as Foldable +import Data.Functor.Classes +import Data.Map.Lazy (Map) +import qualified Data.Map.Lazy as Map +import qualified Data.Maybe as Maybe +import Data.Semigroup (Semigroup, (<>)) +import Data.Set (Set) +import qualified Data.Set as Set + +infixl 9 ! + +type Size = Int + +newtype SetMultimap k a = SetMultimap (Map k (Set a), Size) + deriving (Eq, Ord, Data) + +instance Eq k => Eq1 (SetMultimap k) where + liftEq = liftEq2 (==) + +instance Eq2 SetMultimap where + liftEq2 eqk eqv m n = + Map.size (toMap m) == Map.size (toMap n) + && liftEq (liftEq2 eqk eqv) (toList m) (toList n) + +instance Ord k => Ord1 (SetMultimap k) where + liftCompare = liftCompare2 compare + +instance Ord2 SetMultimap where + liftCompare2 cmpk cmpv m n = + liftCompare (liftCompare2 cmpk cmpv) (toList m) (toList n) + +instance (Show k, Show a) => Show (SetMultimap k a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +instance Show k => Show1 (SetMultimap k) where + liftShowsPrec = liftShowsPrec2 showsPrec showList + +instance Show2 SetMultimap where + liftShowsPrec2 spk slk spv slv d m = + showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) + where + sp = liftShowsPrec2 spk slk spv slv + sl = liftShowList2 spk slk spv slv + +instance (Ord k, Ord a, Read k, Read a) => Read (SetMultimap k a) where + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + pure (fromList xs,t) + +instance Foldable.Foldable (SetMultimap k) where + foldMap = foldMapWithKey . const + {-# INLINE foldMap #-} + +instance (Ord k, Ord a) => Semigroup (SetMultimap k a) where + (<>) = union + +instance (Ord k, Ord a) => Monoid (SetMultimap k a) where + mempty = empty + mappend = (<>) + +------------------------------------------------------------------------------ + +-- | /O(1)/. The empty multimap. +-- +-- > size empty === 0 +empty :: SetMultimap k a +empty = SetMultimap (Map.empty, 0) + +-- | /O(1)/. A multimap with a single element. +-- +-- > singleton 1 'a' === fromList [(1, 'a')] +-- > size (singleton 1 'a') === 1 +singleton :: k -> a -> SetMultimap k a +singleton k a = SetMultimap (Map.singleton k (Set.singleton a), 1) + +-- | /O(n*log n)/ where /n/ is the length of the input list. +-- Build a multimap from a list of key\/value pairs. +-- +-- > fromList ([] :: [(Int, Char)]) === empty +-- > fromList [(1, 'b'), (2, 'a'), (1, 'b')] === fromList [(1, 'b'), (2, 'a')] +fromList :: (Ord k, Ord a) => [(k, a)] -> SetMultimap k a +fromList = Foldable.foldr (uncurry insert) empty + +-- | /O(k)/. A key is retained only if it is associated with a +-- non-empty set of values. +fromMap :: Map k (Set a) -> SetMultimap k a +fromMap m = SetMultimap (m', sum (fmap Set.size m')) + where + m' = Map.filter (not . Set.null) m + +------------------------------------------------------------------------------ + +-- | /O(log m * log k)/. If the key exists in the multimap, the new value will +-- be inserted into the set of values for the key. It is a no-op if the value +-- already exists in the set. +-- +-- > insert 1 'a' empty === singleton 1 'a' +-- > insert 1 'a' (fromList [(1, 'b'), (2, 'a')]) === fromList [(1, 'a'), (1, 'b'), (2, 'a')] +-- > insert 1 'a' (fromList [(1, 'a'), (2, 'c')]) === fromList [(1, 'a'), (2, 'c')] +insert :: (Ord k, Ord a) => k -> a -> SetMultimap k a -> SetMultimap k a +insert k a (SetMultimap (m, _)) = fromMap' k (Map.alter f k m) + where + f (Just as) = Just (Set.insert a as) + f Nothing = Just (Set.singleton a) + +-- | /O(log k)/. Delete a key and all its values from the map. +-- +-- > delete 1 (fromList [(1,'a'),(1,'b'),(2,'c')]) === singleton 2 'c' +delete :: Ord k => k -> SetMultimap k a -> SetMultimap k a +delete = alter (const Set.empty) + +-- | /O(log m * log k)/. Remove the first +-- occurrence of the value associated with the key, if exists. +-- +-- > deleteWithValue 1 'c' (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(1,'b'),(2,'c')] +-- > deleteWithValue 1 'c' (fromList [(2,'c'),(1,'c')]) === singleton 2 'c' +deleteWithValue :: (Ord k, Ord a) => k -> a -> SetMultimap k a -> SetMultimap k a +deleteWithValue k a = alter (Set.delete a) k + +-- | /O(log m * log k)/. Remove the maximal value +-- associated with the key, if exists. +-- +-- > deleteMax 3 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(1,'b'),(2,'c')] +-- > deleteMax 1 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(2,'c')] +deleteMax :: (Ord k, Ord a) => k -> SetMultimap k a -> SetMultimap k a +deleteMax = alter Set.deleteMax + +-- | /O(log m * log k)/. Remove the minimal value +-- associated with the key, if exists. +-- +-- > deleteMin 3 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(1,'b'),(2,'c')] +-- > deleteMin 1 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'b'),(2,'c')] +deleteMin :: (Ord k, Ord a) => k -> SetMultimap k a -> SetMultimap k a +deleteMin = alter Set.deleteMin + +-- | /O(m * log m * log k)/, assuming the function @a -> a@ takes /O(1)/. +-- Update values at a specific key, if exists. +-- +-- Since values are sets, the result may be smaller than the original multimap. +-- +-- > adjust ("new " ++) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"new a"),(1,"new b"),(2,"c")] +-- > adjust (const "z") 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"z"),(2,"c")] +adjust :: (Ord k, Ord a) => (a -> a) -> k -> SetMultimap k a -> SetMultimap k a +adjust = adjustWithKey. const + +-- | /O(m * log m * log k)/, assuming the function @k -> a -> a@ takes /O(1)/. +-- Update values at a specific key, if exists. +-- +-- Since values are sets, the result may be smaller than the original multimap. +-- +-- > adjustWithKey (\k x -> show k ++ ":new " ++ x) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) +-- > === fromList [(1,"1:new a"),(1,"1:new b"),(2,"c")] +adjustWithKey :: (Ord k, Ord a) => (k -> a -> a) -> k -> SetMultimap k a -> SetMultimap k a +adjustWithKey f = alterWithKey (Set.map . f) + +-- | /O(m * log m * log k)/, assuming the function @a -> 'Maybe' a@ +-- takes /O(1)/. The expression (@'update' f k map@) updates the +-- values at key @k@, if exists. If @f@ returns 'Nothing' for a value, the +-- value is deleted. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing in do +-- > update f 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"new a"),(2, "c")] +-- > update f 1 (fromList [(1,"b"),(1,"c"),(2,"c")]) === singleton 2 "c" +update :: (Ord k, Ord a) => (a -> Maybe a) -> k -> SetMultimap k a -> SetMultimap k a +update = updateWithKey . const + +-- | /O(m * log m * log k)/, assuming the function @k -> a -> 'Maybe' a@ +-- takes /O(1)/. The expression (@'updateWithKey' f k map@) updates the +-- values at key @k@, if exists. If @f@ returns 'Nothing' for a value, the +-- value is deleted. +-- +-- > let f k x = if x == "a" then Just (show k ++ ":new a") else Nothing in do +-- > updateWithKey f 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"1:new a"),(2,"c")] +-- > updateWithKey f 1 (fromList [(1,"b"),(1,"c"),(2,"c")]) === singleton 2 "c" +updateWithKey :: (Ord k, Ord a) => (k -> a -> Maybe a) -> k -> SetMultimap k a -> SetMultimap k a +updateWithKey f = alterWithKey g + where + g k = catMaybes . Set.map (f k) + +-- | /O(log k)/, assuming the function @'Set' a -> 'Set' a@ takes /O(1)/. +-- The expression (@'alter' f k map@) alters the values at k, if exists. +-- +-- > let (f, g) = (const Set.empty, Set.insert 'c') in do +-- > alter f 1 (fromList [(1, 'a'), (2, 'b')]) === singleton 2 'b' +-- > alter f 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b')] +-- > alter g 1 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'c'), (1, 'a'), (2, 'b')] +-- > alter g 1 (fromList [(1, 'c'), (2, 'b')]) === fromList [(1, 'c'), (2, 'b')] +-- > alter g 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b'), (3, 'c')] +alter :: Ord k => (Set a -> Set a) -> k -> SetMultimap k a -> SetMultimap k a +alter = alterWithKey . const + +-- | /O(log k)/, assuming the function @k -> 'Set' a -> 'Set' a@ takes /O(1)/. +-- The expression (@'alterWithKey' f k map@) alters the values at k, if exists. +-- +-- > let (f, g) = (const (const Set.empty), Set.insert . show) in do +-- > alterWithKey f 1 (fromList [(1, "a"), (2, "b")]) === singleton 2 "b" +-- > alterWithKey f 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b")] +-- > alterWithKey g 1 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "1"), (1, "a"), (2, "b")] +-- > alterWithKey g 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b"), (3, "3")] +alterWithKey :: Ord k => (k -> Set a -> Set a) -> k -> SetMultimap k a -> SetMultimap k a +alterWithKey f k mm@(SetMultimap (m, _)) + | Set.null as = fromMap (Map.delete k m) + | otherwise = fromMap (Map.insert k as m) + where + as = f k (mm ! k) + +------------------------------------------------------------------------------ + +-- | /O(log k)/. Lookup the values at a key in the map. It returns an empty +-- set if the key is not in the map. +lookup :: Ord k => k -> SetMultimap k a -> Set a +lookup k (SetMultimap (m, _)) = Maybe.fromMaybe Set.empty (Map.lookup k m) + +-- | /O(log k)/. Lookup the values at a key in the map. It returns an empty +-- set if the key is not in the map. +-- +-- > fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 3 === Set.fromList "ac" +-- > fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 2 === Set.empty +(!) :: Ord k => SetMultimap k a -> k -> Set a +(!) = flip lookup + +-- | /O(log k)/. Is the key a member of the map? +-- +-- A key is a member of the map if and only if there is at least one value +-- associated with it. +-- +-- > member 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === True +-- > member 1 (deleteMax 1 (fromList [(2, 'c'), (1, 'c')])) === False +member :: Ord k => k -> SetMultimap k a -> Bool +member k (SetMultimap (m, _)) = Map.member k m + +-- | /O(log k)/. Is the key not a member of the map? +-- +-- A key is a member of the map if and only if there is at least one value +-- associated with it. +-- +-- > notMember 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === False +-- > notMember 1 (deleteMin 1 (fromList [(2, 'c'), (1, 'c')])) === True +notMember :: Ord k => k -> SetMultimap k a -> Bool +notMember k = not . member k + +-- | /O(1)/. Is the multimap empty? +-- +-- > Data.Multimap.Set.null empty === True +-- > Data.Multimap.Set.null (singleton 1 'a') === False +null :: SetMultimap k a -> Bool +null (SetMultimap (m, _)) = Map.null m + +-- | /O(1)/. Is the multimap non-empty? +-- +-- > notNull empty === False +-- > notNull (singleton 1 'a') === True +notNull :: SetMultimap k a -> Bool +notNull = not . null + +-- | The total number of values for all keys. +-- +-- @size@ is evaluated lazily. Forcing the size for the first time takes up to +-- /O(k)/ and subsequent forces take /O(1)/. +-- +-- > size empty === 0 +-- > size (singleton 1 'a') === 1 +-- > size (fromList [(1, 'a'), (2, 'b'), (2, 'c'), (2, 'b')]) === 3 +size :: SetMultimap k a -> Int +size (SetMultimap (_, sz)) = sz + +------------------------------------------------------------------------------ + +-- | Union two multimaps, unioning values for duplicate keys. +-- +-- > union (fromList [(1,'a'),(2,'b'),(2,'c')]) (fromList [(1,'d'),(2,'b')]) +-- > === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c')] +union :: (Ord k, Ord a) => SetMultimap k a -> SetMultimap k a -> SetMultimap k a +union (SetMultimap (m1, _)) (SetMultimap (m2, _)) = + fromMap (Map.unionWith Set.union m1 m2) + +-- | Union a number of multimaps, unioning values for duplicate keys. +-- +-- > unions [fromList [(1,'a'),(2,'b'),(2,'c')], fromList [(1,'d'),(2,'b')]] +-- > === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c')] +unions :: (Foldable f, Ord k, Ord a) => f (SetMultimap k a) -> SetMultimap k a +unions = Foldable.foldr union empty + +-- | Difference of two multimaps. +-- +-- If a key exists in the first multimap but not the second, it remains +-- unchanged in the result. If a key exists in both multimaps, a set +-- difference is performed on their values. +-- +-- > difference (fromList [(1,'a'),(2,'b'),(2,'c')]) (fromList [(1,'d'),(2,'b'),(2,'a')]) +-- > === fromList [(1,'a'),(2,'c')] +difference :: (Ord k, Ord a) => SetMultimap k a -> SetMultimap k a -> SetMultimap k a +difference (SetMultimap (m1, _)) (SetMultimap (m2, _)) = fromMap $ + Map.differenceWith (\xs ys -> Just (xs Set.\\ ys)) m1 m2 + +------------------------------------------------------------------------------ + +-- | /O(n * log m)/, assuming the function @a -> b@ takes /O(1)/. +-- Map a function over all values in the map. +-- +-- Since values are sets, the result may be smaller than the original multimap. +-- +-- > Data.Multimap.Set.map (++ "x") (fromList [(1,"a"),(2,"b")]) === fromList [(1,"ax"),(2,"bx")] +-- > Data.Multimap.Set.map (const "c") (fromList [(1,"a"),(1,"b"),(2,"b")]) === fromList [(1,"c"),(2,"c")] +map :: Ord b => (a -> b) -> SetMultimap k a -> SetMultimap k b +map = mapWithKey . const + +-- | /O(n * log m)/, assuming the function @k -> a -> b@ takes /O(1)/. +-- Map a function over all values in the map. +-- +-- Since values are sets, the result may be smaller than the original multimap. +-- +-- > mapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1,"a"),(2,"b")]) === fromList [(1,"1:a"),(2,"2:b")] +mapWithKey :: Ord b => (k -> a -> b) -> SetMultimap k a -> SetMultimap k b +mapWithKey f (SetMultimap (m, _)) = fromMap $ Map.mapWithKey (Set.map . f) m + +-- | /O(n)/. Fold the values in the map using the given right-associative +-- binary operator. +-- +-- > Data.Multimap.Set.foldr ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldr :: (a -> b -> b) -> b -> SetMultimap k a -> b +foldr = foldrWithKey . const + +-- | /O(n)/. Fold the values in the map using the given left-associative +-- binary operator. +-- +-- > Data.Multimap.Set.foldl (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldl :: (a -> b -> a) -> a -> SetMultimap k b -> a +foldl = foldlWithKey . (const .) + +-- | /O(n)/. Fold the key\/value paris in the map using the given +-- right-associative binary operator. +-- +-- > foldrWithKey (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldrWithKey :: (k -> a -> b -> b) -> b -> SetMultimap k a -> b +foldrWithKey f b (SetMultimap (m, _)) = Map.foldrWithKey f' b m + where + f' = flip . Set.foldr . f + +-- | /O(n)/. Fold the key\/value paris in the map using the given +-- left-associative binary operator. +-- +-- > foldlWithKey (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldlWithKey :: (a -> k -> b -> a) -> a -> SetMultimap k b -> a +foldlWithKey f a (SetMultimap (m, _)) = Map.foldlWithKey f' a m + where + f' = flip (Set.foldl . flip f) + +-- | /O(n)/. A strict version of 'foldr'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > Data.Multimap.Set.foldr' ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldr' :: (a -> b -> b) -> b -> SetMultimap k a -> b +foldr' = foldrWithKey' . const + +-- | /O(n)/. A strict version of 'foldl'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > Data.Multimap.Set.foldl' (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 +foldl' :: (a -> b -> a) -> a -> SetMultimap k b -> a +foldl' = foldlWithKey' . (const .) + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > foldrWithKey' (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldrWithKey' :: (k -> a -> b -> b) -> b -> SetMultimap k a -> b +foldrWithKey' f b (SetMultimap (m, _)) = Map.foldrWithKey f' b m + where + f' = flip . Set.foldr' . f + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > foldlWithKey' (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 +foldlWithKey' :: (a -> k -> b -> a) -> a -> SetMultimap k b -> a +foldlWithKey' f a (SetMultimap (m, _)) = Map.foldlWithKey f' a m + where + f' = flip (Set.foldl' . flip f) + +-- | /O(n)/. Fold the key\/value pairs in the map using the given monoid. +-- +-- > foldMapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1, "a"), (1, "c"), (2, "b")]) === "1:a1:c2:b" +foldMapWithKey :: Monoid m => (k -> a -> m) -> SetMultimap k a -> m +foldMapWithKey f (SetMultimap (m, _)) = Map.foldMapWithKey f' m + where + f' = Foldable.foldMap . f + +------------------------------------------------------------------------------ + +-- | /O(n)/. Return all elements of the multimap in ascending order of +-- their keys. Elements of each key appear in ascending order. +-- +-- > elems (fromList [(2,'a'),(1,'b'),(3,'d'),(3,'c'),(1,'b')]) === "bacd" +-- > elems (empty :: SetMultimap Int Char) === [] +elems :: SetMultimap k a -> [a] +elems (SetMultimap (m, _)) = Map.elems m >>= Set.toList + +-- | /O(k)/. Return all keys of the multimap in ascending order. +-- +-- > keys (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'b')]) === [1,2,3] +-- > keys (empty :: SetMultimap Int Char) === [] +keys :: SetMultimap k a -> [k] +keys (SetMultimap (m, _)) = Map.keys m + +-- | /O(k)/. The set of all keys of the multimap. +-- +-- > keysSet (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'b')]) === Set.fromList [1,2,3] +-- > keysSet (empty :: SetMultimap Int Char) === Set.empty +keysSet :: SetMultimap k a -> Set k +keysSet (SetMultimap (m, _)) = Map.keysSet m + +-- | An alias for 'toAscList'. +assocs :: SetMultimap k a -> [(k, a)] +assocs = toAscList + +-- | Convert the multimap into a list of key/value pairs. +-- +-- > toList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'a'),(1,'b'),(2,'a'),(3,'c')] +toList :: SetMultimap k a -> [(k, a)] +toList = toAscList + +-- | Convert the multimap into a list of key/value pairs in ascending +-- order of keys. Elements of each key appear in ascending order. +-- +-- > toAscList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'a'),(1,'b'),(2,'a'),(3,'c')] +toAscList :: SetMultimap k a -> [(k, a)] +toAscList (SetMultimap (m, _)) = + Map.toAscList m >>= uncurry (\k -> fmap (k,) . Set.toAscList) + +-- | Convert the multimap into a list of key/value pairs in descending +-- order of keys. Elements of each key appear in descending order. +-- +-- > toDescList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(3,'c'),(2,'a'),(1,'b'),(1,'a')] +toDescList :: SetMultimap k a -> [(k, a)] +toDescList (SetMultimap (m, _)) = + Map.toDescList m >>= uncurry (\k -> fmap (k,) . Set.toDescList) + +-- | /O(1)/. Convert the multimap into a regular map. +toMap :: SetMultimap k a -> Map k (Set a) +toMap (SetMultimap (m, _)) = m + +------------------------------------------------------------------------------ + +-- | /O(n)/, assuming the predicate function takes /O(1)/. +-- Retain all values that satisfy the predicate. A key is removed if +-- none of its values satisfies the predicate. +-- +-- > Data.Multimap.Set.filter (> 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === singleton 1 'b' +-- > Data.Multimap.Set.filter (< 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === empty +filter :: (a -> Bool) -> SetMultimap k a -> SetMultimap k a +filter = filterWithKey . const + +-- | /O(k)/, assuming the predicate function takes /O(1)/. +-- Retain all keys that satisfy the predicate. +filterKey :: (k -> Bool) -> SetMultimap k a -> SetMultimap k a +filterKey p (SetMultimap (m, _)) = fromMap m' + where + m' = Map.filterWithKey (const . p) m + +-- | /O(n)/, assuming the predicate function takes /O(1)/. +-- Retain all key\/value pairs that satisfy the predicate. A key is removed if +-- none of its values satisfies the predicate. +-- +-- > filterWithKey (\k a -> even k && a > 'a') (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'b')]) === singleton 2 'b' +filterWithKey :: (k -> a -> Bool) -> SetMultimap k a -> SetMultimap k a +filterWithKey p (SetMultimap (m, _)) = fromMap m' + where + m' = Map.mapWithKey (Set.filter . p) m + +-- | Generalized 'filter'. +-- +-- > let f a | a > 'b' = Just True +-- > | a < 'b' = Just False +-- > | a == 'b' = Nothing +-- > in do +-- > filterM f (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'c')]) === Nothing +-- > filterM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Just (fromList [(1,'c'),(2,'c')]) +filterM + :: (Ord k, Ord a, Applicative t) + => (a -> t Bool) -> SetMultimap k a -> t (SetMultimap k a) +filterM = filterWithKeyM . const + +-- | Generalized 'filterWithKey'. +-- | Generalized 'filterWithKey'. +-- +-- > let f k a | even k && a > 'b' = Just True +-- > | odd k && a < 'b' = Just False +-- > | otherwise = Nothing +-- > in do +-- > filterWithKeyM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Nothing +-- > filterWithKeyM f (fromList [(1,'a'),(3,'a'),(2,'c'),(4,'c')]) === Just (fromList [(2,'c'),(4,'c')]) +filterWithKeyM + :: (Ord k, Ord a, Applicative t) + => (k -> a -> t Bool) -> SetMultimap k a -> t (SetMultimap k a) +filterWithKeyM f = fmap fromList . List.filterM (uncurry f) . toList + +-- | /O(n * log m)/, assuming the function @a -> 'Maybe' b@ takes /O(1)/. +-- Map values and collect the 'Just' results. +-- +-- > mapMaybe (\a -> if a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) +-- > === fromList [(1,"new a"),(2,"new a")] +mapMaybe :: Ord b => (a -> Maybe b) -> SetMultimap k a -> SetMultimap k b +mapMaybe = mapMaybeWithKey . const + +-- | /O(n * log m)/, assuming the function @k -> a -> 'Maybe' b@ takes /O(1)/. +-- Map key\/value pairs and collect the 'Just' results. +-- +-- > mapMaybeWithKey (\k a -> if k > 1 && a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) +-- > === singleton 2 "new a" +mapMaybeWithKey :: Ord b => (k -> a -> Maybe b) -> SetMultimap k a -> SetMultimap k b +mapMaybeWithKey f (SetMultimap (m, _)) = fromMap $ + Map.mapWithKey (\k -> catMaybes . Set.map (f k)) m + +-- | /O(n * log m)/, assuming the function @a -> 'Either' b c@ takes /O(1)/. +-- Map values and separate the 'Left' and 'Right' results. +-- +-- > mapEither (\a -> if a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) +-- > === (fromList [(1,'a'),(2,'a')],fromList [(1,'c'),(2,'c')]) +mapEither + :: (Ord b, Ord c) + => (a -> Either b c) -> SetMultimap k a -> (SetMultimap k b, SetMultimap k c) +mapEither = mapEitherWithKey . const + +-- | /O(n * log m)/, assuming the function @k -> a -> 'Either' b c@ takes /O(1)/. +-- Map key\/value pairs and separate the 'Left' and 'Right' results. +-- +-- > mapEitherWithKey (\k a -> if even k && a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) +-- > === (fromList [(2,'a')],fromList [(1,'a'),(1,'c'),(2,'c')]) +mapEitherWithKey + :: (Ord b, Ord c) + => (k -> a -> Either b c) -> SetMultimap k a -> (SetMultimap k b, SetMultimap k c) +mapEitherWithKey f (SetMultimap (m, _)) = + (fromMap . Map.mapWithKey (const fst) &&& fromMap . Map.mapWithKey (const snd)) + $ Map.mapWithKey g m + where + g k = partitionEithers . Set.map (f k) + +------------------------------------------------------------------------------ +-- * Non exported functions +------------------------------------------------------------------------------ + +catMaybes :: Ord a => Set (Maybe a) -> Set a +catMaybes = Set.foldl' (\s -> maybe s (`Set.insert` s)) Set.empty + +partitionEithers :: (Ord a, Ord b) => Set (Either a b) -> (Set a, Set b) +partitionEithers = Set.foldr' (either left right) (Set.empty, Set.empty) + where + left a (l,r) = (Set.insert a l, r) + right b (l,r) = (l, Set.insert b r) + +fromMap' :: Ord k => k -> Map k (Set a) -> SetMultimap k a +fromMap' k m = SetMultimap (m', sum (fmap Set.size m')) + where + m' = case Map.lookup k m of + Just as | Set.null as -> Map.delete k m + _ -> m diff --git a/src/Data/Multimap/Table.hs b/src/Data/Multimap/Table.hs new file mode 100644 --- /dev/null +++ b/src/Data/Multimap/Table.hs @@ -0,0 +1,829 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE TypeFamilies #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Multimap.Table +-- Maintainer : Ziyang Liu +-- +-- The @'Table' r c a@ type represents a finite two-dimensional table +-- that associates a pair of keys (a row key of type @r@ and +-- a column key of type @c@) with a value of type @a@. +-- +-- The implementation is backed by two maps: a @'Map' r ('Map' c) a@, and +-- a @'Map' c ('Map' r) a@, called "row map" and "column map", respectively. +-- +-- It is worth noting that all functions that traverse a table, such as +-- 'foldl', 'foldr', 'foldMap' and 'traverse', are row-oriented, i.e., +-- they traverse the table row by row. To traverse a table column +-- by column, 'transpose' the table first. +-- +-- In the following Big-O notations, unless otherwise noted, /n/ denotes +-- the size of the table (i.e., the total number of values for all +-- row and column keys), /r/ denotes the number of row keys that has at +-- least one value, /c/ denotes the number of column keys that has at +-- least one value, and /k = max r c/. +module Data.Multimap.Table ( + Table + + -- * Construction + , empty + , singleton + , fromRowMap + , fromColumnMap + , transpose + + -- ** From Unordered Lists + , fromList + + -- * Deletion\/Update + , insert + , delete + , deleteRow + , deleteColumn + , adjust + , adjustWithKeys + , update + , updateWithKeys + , alter + , alterWithKeys + + -- * Query + -- ** Lookup + , lookup + , (!?) + , (!) + , hasCell + , hasRow + , hasColumn + + -- ** Size + , null + , notNull + , size + + -- * Combine + -- ** Union + , union + , unionWith + , unionWithKeys + , unions + , unionsWith + , unionsWithKeys + + -- ** Difference + , difference + + -- * Traversal + -- ** Map + , map + , mapWithKeys + , traverseWithKeys + , traverseMaybeWithKeys + + -- ** Folds + , foldr + , foldl + , foldrWithKeys + , foldlWithKeys + , foldMapWithKeys + + -- ** Strict Folds + , foldr' + , foldl' + , foldrWithKeys' + , foldlWithKeys' + + -- * Conversion + , row + , column + , rowMap + , columnMap + , rowKeys + , columnKeys + , rowKeysSet + , columnKeysSet + + -- ** Lists + , toList + + -- ** Ordered lists + , toRowAscList + , toColumnAscList + , toRowDescList + , toColumnDescList + + -- * Filter + , filter + , filterRow + , filterColumn + , filterWithKeys + + , mapMaybe + , mapMaybeWithKeys + , mapEither + , mapEitherWithKeys + ) where + +import Control.Arrow ((&&&)) +import Data.Data (Data) +import qualified Data.Foldable as Foldable +import Data.Map (Map) +import qualified Data.Map as Map +import qualified Data.Maybe as Maybe +import Data.Semigroup (Semigroup, (<>)) +import Data.Set (Set) + +import Prelude hiding (filter, foldl, foldr, lookup, map, null) + +infixl 9 !,!? + +type Size = Int + +newtype Table r c a = Table (Map r (Map c a), Map c (Map r a), Size) + deriving (Eq, Ord, Data) + +instance (Show r, Show c, Show a) => Show (Table r c a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +instance (Ord r, Ord c, Read r, Read c, Read a) => Read (Table r c a) where + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + pure (fromList xs,t) + +instance Functor (Table r c) where + fmap = map + +instance Foldable.Foldable (Table r c) where + foldMap = foldMapWithKeys . const . const + {-# INLINE foldMap #-} + +instance (Ord r, Ord c) => Traversable (Table r c) where + traverse = traverseWithKeys . const . const + {-# INLINE traverse #-} + +instance (Ord r, Ord c) => Semigroup (Table r c a) where + (<>) = union + +instance (Ord r, Ord c) => Monoid (Table r c a) where + mempty = empty + mappend = (<>) + +------------------------------------------------------------------------------ + +-- | /O(1)/. The empty table. +-- +-- > size empty === 0 +empty :: Table r c a +empty = Table (Map.empty, Map.empty, 0) + +-- | /O(1)/. A table with a single element. +-- +-- > singleton 1 'a' "a" === fromList [(1,'a',"a")] +-- > size (singleton 1 'a' "a") === 1 +singleton :: r -> c -> a -> Table r c a +singleton r c a = Table (Map.singleton r (Map.singleton c a), Map.singleton c (Map.singleton r a), 1) + +-- | Build a table from a list of key\/value pairs. +-- +-- > fromList ([] :: [(Int, Char, String)]) === empty +fromList :: (Ord r, Ord c) => [(r, c, a)] -> Table r c a +fromList = Foldable.foldr (uncurry3 insert) empty + +-- | Build a table from a row map. +-- +-- > fromRowMap (Map.fromList [(1, Map.fromList [('a',"b"),('b',"c")]), (2, Map.fromList [('a',"d")])]) +-- > === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] +fromRowMap :: (Ord r, Ord c) => Map r (Map c a) -> Table r c a +fromRowMap m = Table (m', transpose' m', size' m') + where m' = nonEmpty m + +-- | Build a table from a column map. +-- +-- > fromColumnMap (Map.fromList [(1, Map.fromList [('a',"b"),('b',"c")]), (2, Map.fromList [('a',"d")])]) +-- > === fromList [('a',1,"b"),('a',2,"d"),('b',1,"c")] +fromColumnMap :: (Ord r, Ord c) => Map c (Map r a) -> Table r c a +fromColumnMap m = Table (transpose' m', m', size' m') + where m' = nonEmpty m + +-- | Flip the row and column keys. +-- +-- > transpose (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [('a',1,"b"),('a',2,"d"),('b',1,"c")] +transpose :: Table r c a -> Table c r a +transpose (Table (rm, cm, sz)) = Table (cm, rm, sz) + +------------------------------------------------------------------------------ + +-- | /O(log k)/. Associate with value with the row key and the column key. +-- If the table already contains a value for those keys, the value is replaced. +-- +-- > insert 1 'a' "a" empty === singleton 1 'a' "a" +-- > insert 1 'a' "a" (fromList [(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"d")] +-- > insert 1 'a' "a" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"d")] +insert :: (Ord r, Ord c) => r -> c -> a -> Table r c a -> Table r c a +insert r c a (Table (rm, cm, _)) = fromMaps' r c rm' cm' + where + rm' = Map.alter f r rm + cm' = Map.alter g c cm + f = Just . maybe (Map.singleton c a) (Map.insert c a) + g = Just . maybe (Map.singleton r a) (Map.insert r a) + +-- | /O(log k)/. Remove the value associated with the given keys. +-- +-- > delete 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'b',"c"),(2,'a',"d")] +-- > delete 1 'a' (fromList [(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'b',"c"),(2,'a',"d")] +delete :: (Ord r, Ord c) => r -> c -> Table r c a -> Table r c a +delete r c (Table (rm, cm, _)) = fromMaps' r c rm' cm' + where + rm' = Map.adjust (Map.delete c) r rm + cm' = Map.adjust (Map.delete r) c cm + +-- | Remove an entire row. +-- +-- > deleteRow 1 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 2 'a' "d" +-- > deleteRow 3 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] +deleteRow :: Ord r => r -> Table r c a -> Table r c a +deleteRow r (Table (rm, cm, _)) = Table (rm', cm', size' rm') + where + rm' = Map.delete r rm + cm' = nonEmpty $ Map.map (Map.delete r) cm + +-- | Remove an entire column. +-- +-- > deleteColumn 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "c" +-- > deleteColumn 'z' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] +deleteColumn :: Ord c => c -> Table r c a -> Table r c a +deleteColumn c (Table (rm, cm, _)) = Table (rm', cm', size' cm') + where + rm' = nonEmpty $ Map.map (Map.delete c) rm + cm' = Map.delete c cm + +-- | /O(log k)/, assuming the function @a -> a@ takes /O(1)/. +-- Update the value at a specific row key and column key, if exists. +-- +-- > adjust ("new " ++) 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"new b"),(1,'b',"c"),(2,'a',"d")] +adjust :: (Ord r, Ord c) => (a -> a) -> r -> c -> Table r c a -> Table r c a +adjust = adjustWithKeys . const . const + +-- | /O(log k)/, assuming the function @r -> c -> a -> a@ takes /O(1)/. +-- Update the value at a specific row key and column key, if exists. +-- +-- > adjustWithKeys (\r c x -> show r ++ ":" ++ show c ++ ":new " ++ x) 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) +-- > === fromList [(1,'a',"1:'a':new b"),(1,'b',"c"),(2,'a',"d")] +adjustWithKeys + :: (Ord r, Ord c) + => (r -> c -> a -> a) -> r -> c -> Table r c a -> Table r c a +adjustWithKeys f = updateWithKeys (\r c a -> Just (f r c a)) + +-- | /O(log k)/, assuming the function @a -> 'Maybe' a@ takes /O(1)/. +-- The expression (@'update' f r c table@) updates the value at the given +-- row and column keys, if exists. If @f@ returns 'Nothing', the value +-- associated with those keys, if exists is deleted. +-- +-- > let f x = if x == "b" then Just "new b" else Nothing in do +-- > update f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"new b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > update f 1 'a' (fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +update :: (Ord r, Ord c) => (a -> Maybe a) -> r -> c -> Table r c a -> Table r c a +update = updateWithKeys . const . const + +-- | /O(log k)/, assuming the function @r -> c -> a -> 'Maybe' a@ takes /O(1)/. +-- The expression (@'updateWithKeys' f r c table@) updates the value at the given +-- row and column keys, if exists. If @f@ returns 'Nothing', the value +-- associated with those keys, if exists is deleted. +-- +-- > let f r c x = if x == "b" then Just (show r ++ ":" ++ show c ++ ":new b") else Nothing in do +-- > updateWithKeys f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"1:'a':new b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > updateWithKeys f 1 'a' (fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +updateWithKeys + :: (Ord r, Ord c) + => (r -> c -> a -> Maybe a) -> r -> c -> Table r c a -> Table r c a +updateWithKeys f = alterWithKeys (\r c -> (>>= f r c)) + +-- | /O(log k)/, assuming the function @'Maybe' a -> 'Maybe' a@ takes /O(1)/. +-- The expression (@'alter' f r c table@) alters the value at the given +-- row and column keys, if exists. It can be used to insert, delete +-- or update a value. +-- +-- > let (f,g,h) = (const Nothing, const (Just "hello"), fmap ('z':)) in do +-- > alter f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alter f 4 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alter f 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alter g 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"hello"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alter g 4 'e' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c"),(4,'e',"hello")] +-- > alter h 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"zb"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alter h 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +alter :: (Ord r, Ord c) => (Maybe a -> Maybe a) -> r -> c -> Table r c a -> Table r c a +alter = alterWithKeys . const . const + +-- | /O(log k)/, assuming the function @r -> c -> 'Maybe' a -> 'Maybe' a@ takes /O(1)/. +-- The expression (@'alterWithKeys' f r c table@) alters the value at the given +-- row and column keys, if exists. It can be used to insert, delete +-- or update a value. +-- +-- > let (f,g) = (\_ _ _ -> Nothing, \r c -> fmap ((show r ++ ":" ++ show c ++ ":") ++)) in do +-- > alterWithKeys f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alterWithKeys f 4 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alterWithKeys f 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alterWithKeys g 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"1:'a':b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +-- > alterWithKeys g 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] +alterWithKeys + :: (Ord r, Ord c) + => (r -> c -> Maybe a -> Maybe a) -> r -> c -> Table r c a -> Table r c a +alterWithKeys f r c tbl@(Table (rm, cm, _)) + | Just a <- f r c (lookup r c tbl) = + let rm' = Map.alter (Just . maybe (Map.singleton c a) (Map.insert c a)) r rm + cm' = Map.alter (Just . maybe (Map.singleton r a) (Map.insert r a)) c cm + in fromMaps' r c rm' cm' + | otherwise = delete r c tbl + +------------------------------------------------------------------------------ + +-- | /O(log k)/. Lookup the values at a row key and column key in the map. +lookup :: (Ord r, Ord c) => r -> c -> Table r c a -> Maybe a +lookup r c (Table (rm, _, _)) = Map.lookup r rm >>= Map.lookup c + +-- | /O(log k)/. Lookup the values at a row key and column key in the map. +-- +-- > fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")] !? (1,'a') === Just "b" +-- > fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")] !? (1,'c') === Nothing +(!?) :: (Ord r, Ord c) => Table r c a -> (r, c) -> Maybe a +(!?) = flip (uncurry lookup) + +-- | /O(log k)/. Lookup the values at a row key and column key in the map. +-- Calls 'error' if the value does not exist. +-- +-- > fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")] ! (1,'a') === "b" +(!) :: (Ord r, Ord c) => Table r c a -> (r, c) -> a +(!) tbl keys = + Maybe.fromMaybe (error "Table.!: cell does not exist") (tbl !? keys) + +-- | /O(log k)/. Is there a value associated with the given row and +-- column keys? +-- +-- > hasCell (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (1,'a') === True +-- > hasCell (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (1,'c') === False +hasCell :: (Ord r, Ord c) => Table r c a -> (r, c) -> Bool +hasCell (Table (rm, _, _)) (r, c) = + maybe False (Map.member c) (Map.lookup r rm) + +-- | /O(log r)/. Is there a row with the given row key that has at least +-- one value? +-- +-- > hasRow (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 1 === True +-- > hasRow (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 3 === False +hasRow :: Ord r => Table r c a -> r -> Bool +hasRow (Table (rm, _, _)) r = Map.member r rm + +-- | /O(log c)/. Is there a column with the given column key that has at least +-- one value? +-- +-- > hasColumn (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 'a' === True +-- > hasColumn (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 'c' === False +hasColumn :: Ord c => Table r c a -> c -> Bool +hasColumn (Table (_, cm, _)) c = Map.member c cm + +-- | /O(1)/. Is the table empty? +-- +-- > Data.Multimap.Table.null empty === True +-- > Data.Multimap.Table.null (singleton 1 'a' "a") === False +null :: Table r c a -> Bool +null (Table (rm, _, _)) = Map.null rm + +-- | /O(1)/. Is the table non-empty? +-- +-- > notNull empty === False +-- > notNull (singleton 1 'a' "a") === True +notNull :: Table r c a -> Bool +notNull = not . null + +-- | The total number of values for all row and column keys. +-- +-- @size@ is evaluated lazily. Forcing the size for the first time takes up to +-- /O(n)/ and subsequent forces take /O(1)/. +-- +-- > size empty === 0 +-- > size (singleton 1 'a' "a") === 1 +-- > size (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) === 3 +size :: Table r c a -> Int +size (Table (_, _, sz)) = sz + +------------------------------------------------------------------------------ + +-- | Union two tables, preferring values from the first table +-- upon duplicate keys. +-- +-- > union (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]) +-- > === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] +union :: (Ord r, Ord c) => Table r c a -> Table r c a -> Table r c a +union = unionWith const + +-- | Union a number of tables, preferring values from the leftmost table +-- upon duplicate keys. +-- +-- > unions [fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")], fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]] +-- > === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] +unions :: (Foldable f, Ord r, Ord c) => f (Table r c a) -> Table r c a +unions = Foldable.foldr union empty + +-- | Union two tables with a combining function for duplicate keys. +-- +-- > unionWith (++) (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]) +-- > === fromList [(1,'a',"bc"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] +unionWith :: (Ord r, Ord c) => (a -> a -> a) -> Table r c a -> Table r c a -> Table r c a +unionWith = unionWithKeys . const . const + +-- | Union two tables with a combining function for duplicate keys. +-- +-- > let f r c a a' = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ a' in do +-- > unionWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]) +-- > === fromList [(1,'a',"1:'a':b|c"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] +unionWithKeys + :: (Ord r, Ord c) + => (r -> c -> a -> a -> a) + -> Table r c a -> Table r c a -> Table r c a +unionWithKeys f (Table (rm1, cm1, _)) (Table (rm2, cm2, _)) = fromMaps rm cm + where + rm = Map.unionWithKey (Map.unionWithKey . f) rm1 rm2 + cm = Map.unionWithKey (Map.unionWithKey . flip f) cm1 cm2 + +-- | Union a number of tables with a combining function for duplicate keys. +-- +-- > unionsWith (++) [fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")], fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]] +-- > === fromList [(1,'a',"bc"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] +unionsWith :: (Foldable f, Ord r, Ord c) => (a -> a -> a) -> f (Table r c a) -> Table r c a +unionsWith f = Foldable.foldr (unionWith f) empty + +-- | Union a number of tables with a combining function for duplicate keys. +-- +-- > let f r c a a' = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ a' in do +-- > unionsWithKeys f [fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")], fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]] +-- > === fromList [(1,'a',"1:'a':b|c"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] +unionsWithKeys + :: (Foldable f, Ord r, Ord c) + => (r -> c -> a -> a -> a) + -> f (Table r c a) -> Table r c a +unionsWithKeys f = Foldable.foldr (unionWithKeys f) empty + +-- | Difference of two tables. Return values in the first table whose +-- row and column keys do not have an associated value in the second table. +-- +-- > difference (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(1,'b',"d"),(2,'b',"b")]) +-- > === singleton 2 'a' "b" +difference :: (Ord r, Ord c) => Table r c a -> Table r c a -> Table r c a +difference (Table (rm1, cm1, _)) (Table (rm2, cm2, _)) = fromMaps rm cm + where + rm = Map.differenceWith ((Just .) . Map.difference) rm1 rm2 + cm = Map.differenceWith ((Just .) . Map.difference) cm1 cm2 + +------------------------------------------------------------------------------ + +-- | /O(n)/, assuming the function @a -> b@ takes /O(1)/. +-- Map a function over all values in the table. +-- +-- > Data.Multimap.Table.map (++ "x") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) === fromList [(1,'a',"bx"),(1,'b',"cx"),(2,'a',"bx")] +map :: (a -> b) -> Table r c a -> Table r c b +map = mapWithKeys . const . const + +-- | /O(n)/, assuming the function @r -> c -> a -> b@ takes /O(1)/. +-- Map a function over all values in the table. +-- +-- > mapWithKeys (\r c x -> show r ++ ":" ++ show c ++ ":" ++ x) (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) +-- > === fromList [(1,'a',"1:'a':b"),(1,'b',"1:'b':c"),(2,'a',"2:'a':b")] +mapWithKeys :: (r -> c -> a -> b) -> Table r c a -> Table r c b +mapWithKeys f (Table (rm, cm, sz)) = Table (rm', cm', sz) + where + rm' = Map.mapWithKey (Map.mapWithKey . f) rm + cm' = Map.mapWithKey (Map.mapWithKey . flip f) cm + +-- | Traverse the (row key, column key, value) triples and collect the results. +-- +-- > let f r c a = if odd r && c > 'a' then Just (a ++ "x") else Nothing in do +-- > traverseWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) === Nothing +-- > traverseWithKeys f (fromList [(1,'b',"b"),(1,'c',"c"),(3,'d',"b")]) === Just (fromList [(1,'b',"bx"),(1,'c',"cx"),(3,'d',"bx")]) +traverseWithKeys + :: (Applicative t, Ord r, Ord c) + => (r -> c -> a -> t b) + -> Table r c a + -> t (Table r c b) +traverseWithKeys f (Table (rm, _, _)) = fromMaps <$> rm' <*> cm' + where + rm' = Map.traverseWithKey (Map.traverseWithKey . f) rm + cm' = transpose' <$> rm' + +-- | Traverse the (row key, column key, value) triples and collect the 'Just' results. +traverseMaybeWithKeys + :: (Applicative t, Ord r, Ord c) + => (r -> c -> a -> t (Maybe b)) + -> Table r c a + -> t (Table r c b) +traverseMaybeWithKeys f (Table (rm, _, _)) = fromMaps <$> rm' <*> cm' + where + rm' = Map.traverseWithKey (Map.traverseMaybeWithKey . f) rm + cm' = transpose' <$> rm' + +------------------------------------------------------------------------------ + +-- | /O(n)/. Fold the values in the table row by row using the given +-- right-associative binary operator. +-- +-- > Data.Multimap.Table.foldr (:) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "bcd" +foldr :: (a -> b -> b) -> b -> Table r c a -> b +foldr = foldrWithKeys . const . const + +-- | /O(n)/. Fold the values in the table row by row using the given +-- left-associative binary operator. +-- +-- > Data.Multimap.Table.foldl (flip (:)) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "dcb" +foldl :: (a -> b -> a) -> a -> Table r c b -> a +foldl f = foldlWithKeys (\a _ _ -> f a) + +-- | /O(n)/. Fold the (row key, column key value) triplets in the table +-- row by row using the given right-associative binary operator. +-- +-- > let f r c a b = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ b in do +-- > foldrWithKeys f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "1:'a':b|1:'b':c|2:'a':d|" +foldrWithKeys :: (r -> c -> a -> b -> b) -> b -> Table r c a -> b +foldrWithKeys f b (Table (rm, _, _)) = Map.foldrWithKey f' b rm + where + f' = flip . Map.foldrWithKey . f + +-- | /O(n)/. Fold the (row key, column key, value) triplets in the table +-- row by row using the given left-associative binary operator. +-- +-- > let f a r c b = show r ++ ":" ++ show c ++ ":" ++ b ++ "|" ++ a in do +-- > foldlWithKeys f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "2:'a':d|1:'b':c|1:'a':b|" +foldlWithKeys :: (a -> r -> c -> b -> a) -> a -> Table r c b -> a +foldlWithKeys f a (Table (rm, _, _)) = Map.foldlWithKey f' a rm + where + f' = flip (Map.foldlWithKey . flip f) + +-- | /O(n)/. A strict version of 'foldr'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > Data.Multimap.Table.foldr' (:) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "bcd" +foldr' :: (a -> b -> b) -> b -> Table r c a -> b +foldr' = foldrWithKeys' . const . const + +-- | /O(n)/. A strict version of 'foldl'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > Data.Multimap.Table.foldl' (flip (:)) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "dcb" +foldl' :: (a -> b -> a) -> a -> Table r c b -> a +foldl' f = foldlWithKeys' (\a _ _ -> f a) + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > let f r c a b = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ b in do +-- > foldrWithKeys' f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "1:'a':b|1:'b':c|2:'a':d|" +foldrWithKeys' :: (r -> c -> a -> b -> b) -> b -> Table r c a -> b +foldrWithKeys' f b (Table (rm, _, _)) = Map.foldrWithKey' f' b rm + where + f' = flip . Map.foldrWithKey' . f + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the +-- operator is evaluated before using the result in the next application. +-- This function is strict in the starting value. +-- +-- > let f a r c b = show r ++ ":" ++ show c ++ ":" ++ b ++ "|" ++ a in do +-- > foldlWithKeys' f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "2:'a':d|1:'b':c|1:'a':b|" +foldlWithKeys' :: (a -> r -> c -> b -> a) -> a -> Table r c b -> a +foldlWithKeys' f a (Table (rm, _, _)) = Map.foldlWithKey' f' a rm + where + f' = flip (Map.foldlWithKey' . flip f) + +-- | /O(n)/. Fold the (row key, column key, value) triplets in the map +-- row by row using the given monoid. +-- +-- > let f r c a = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" in do +-- > foldMapWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "1:'a':b|1:'b':c|2:'a':d|" +foldMapWithKeys :: Monoid m => (r -> c -> a -> m) -> Table r c a -> m +foldMapWithKeys f (Table (rm, _, _)) = Map.foldMapWithKey f' rm + where + f' = Map.foldMapWithKey . f + +------------------------------------------------------------------------------ + +-- | /O(r)/. Return a mapping from column keys to values for the given +-- row key. +-- +-- > row 1 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.fromList [('a',"b"),('b',"c")] +-- > row 3 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.empty +row :: Ord r => r -> Table r c a -> Map c a +row r (Table (rm, _, _)) = Map.findWithDefault Map.empty r rm + +-- | /O(c)/. Return a mapping from row keys to values for the given +-- column key. +-- +-- > column 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.fromList [(1,"b"),(2,"d")] +-- > column 'c' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.empty +column :: Ord c => c -> Table r c a -> Map r a +column c (Table (_, cm, _)) = Map.findWithDefault Map.empty c cm + +-- | Return a mapping from row keys to maps from column keys to values. +-- +-- > rowMap (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) +-- > === Map.fromList [(1, Map.fromList [('a',"b"),('b',"c")]),(2, Map.fromList [('a',"d")])] +rowMap :: Table r c a -> Map r (Map c a) +rowMap (Table (rm, _, _)) = rm + +-- | Return a mapping from column keys to maps from row keys to values. +-- +-- > columnMap (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) +-- > === Map.fromList [('a', Map.fromList [(1,"b"),(2,"d")]),('b', Map.fromList [(1,"c")])] +columnMap :: Table r c a -> Map c (Map r a) +columnMap (Table (_, cm, _)) = cm + +-- | Return, in ascending order, the list of all row keys of that have +-- at least one value in the table. +-- +-- > rowKeys (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [1,2] +rowKeys :: Table r c a -> [r] +rowKeys (Table (rm, _, _)) = Map.keys rm + +-- | Return, in ascending order, the list of all column keys of that have +-- at least one value in the table. +-- +-- > columnKeys (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === ['a','b'] +columnKeys :: Table r c a -> [c] +columnKeys (Table (_, cm, _)) = Map.keys cm + +-- | Return the set of all row keys of that have at least one value +-- in the table. +-- +-- > rowKeysSet (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Set.fromList [1,2] +rowKeysSet :: Table r c a -> Set r +rowKeysSet (Table (rm, _, _)) = Map.keysSet rm + +-- | Return the set of all column keys of that have at least one value +-- in the table. +-- +-- > columnKeysSet (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Set.fromList ['a','b'] +columnKeysSet :: Table r c a -> Set c +columnKeysSet (Table (_, cm, _)) = Map.keysSet cm + +-- | Convert the table into a list of (row key, column key, value) triples. +-- +-- > toList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] +toList :: Table r c a -> [(r, c, a)] +toList (Table (rm, _, _)) = Map.toList (Map.toList <$> rm) >>= distr + +-- | Convert the table into a list of (row key, column key, value) triples +-- in ascending order of row keys, and ascending order of column keys +-- with a row. +-- +-- > toRowAscList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] +toRowAscList :: Table r c a -> [(r, c, a)] +toRowAscList (Table (rm, _, _)) = Map.toAscList (Map.toAscList <$> rm) >>= distr + +-- | Convert the table into a list of (column key, row key, value) triples +-- in ascending order of column keys, and ascending order of row keys +-- with a column. +-- +-- > toColumnAscList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [('a',1,"b"),('a',2,"d"),('b',1,"c")] +toColumnAscList :: Table r c a -> [(c, r, a)] +toColumnAscList (Table (_, cm, _)) = Map.toAscList (Map.toAscList <$> cm) >>= distr + +-- | Convert the table into a list of (row key, column key, value) triples +-- in descending order of row keys, and descending order of column keys +-- with a row. +-- +-- > toRowDescList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [(2,'a',"d"),(1,'b',"c"),(1,'a',"b")] +toRowDescList :: Table r c a -> [(r, c, a)] +toRowDescList (Table (rm, _, _)) = Map.toDescList (Map.toDescList <$> rm) >>= distr + +-- | Convert the table into a list of (column key, row key, value) triples +-- in descending order of column keys, and descending order of row keys +-- with a column. +-- +-- > toColumnDescList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [('b',1,"c"),('a',2,"d"),('a',1,"b")] +toColumnDescList :: Table r c a -> [(c, r, a)] +toColumnDescList (Table (_, cm, _)) = Map.toDescList (Map.toDescList <$> cm) >>= distr + +------------------------------------------------------------------------------ + +-- | /O(n)/, assuming the predicate function takes /O(1)/. +-- Retain all values that satisfy the predicate. +-- +-- > Data.Multimap.Table.filter (> "c") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 2 'a' "d" +-- > Data.Multimap.Table.filter (> "d") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === empty +filter :: (a -> Bool) -> Table r c a -> Table r c a +filter = filterWithKeys . const . const + +-- | /O(r)/, assuming the predicate function takes /O(1)/. +-- Retain all rows that satisfy the predicate. +-- +-- > filterRow even (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 2 'a' "d" +filterRow :: (r -> Bool) -> Table r c a -> Table r c a +filterRow p (Table (rm, cm, _)) = Table (rm', nonEmpty cm', size' rm') + where + rm' = Map.filterWithKey (const . p) rm + cm' = Map.map (Map.filterWithKey (const . p)) cm + +-- | /O(c)/, assuming the predicate function takes /O(1)/. +-- Retain all columns that satisfy the predicate. +-- +-- > filterColumn (> 'a') (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "c" +filterColumn :: (c -> Bool) -> Table r c a -> Table r c a +filterColumn p (Table (rm, cm, _)) = Table (nonEmpty rm', cm', size' cm') + where + rm' = Map.map (Map.filterWithKey (const . p)) rm + cm' = Map.filterWithKey (const . p) cm + +-- | /O(c)/, assuming the predicate function takes /O(1)/. +-- Retain all (row key, column key, value) triples that satisfy the predicate. +-- +-- > filterWithKeys (\r c a -> odd r && c > 'a' && a > "b") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "c" +filterWithKeys + :: (r -> c -> a -> Bool) + -> Table r c a + -> Table r c a +filterWithKeys p (Table (rm, cm, _)) = fromMaps rm' cm' + where + rm' = Map.mapWithKey (Map.filterWithKey . p) rm + cm' = Map.mapWithKey (Map.filterWithKey . flip p) cm + +-- | /O(n)/, assuming the function @a -> 'Maybe' b@ takes /O(1)/. +-- Map values and collect the 'Just' results. +-- +-- > mapMaybe (\a -> if a == "a" then Just "new a" else Nothing) (fromList [(1,'a',"a"),(1,'b',"c"),(2,'b',"a")]) +-- > === fromList [(1,'a',"new a"),(2,'b',"new a")] +mapMaybe :: (a -> Maybe b) -> Table r c a -> Table r c b +mapMaybe = mapMaybeWithKeys . const . const + +-- | /O(n)/, assuming the function @r -> c -> a -> 'Maybe' b@ takes /O(1)/. +-- Map (row key, column key, value) triples and collect the 'Just' results. +-- +-- > let f r c a = if r == 1 && a == "c" then Just "new c" else Nothing in do +-- > mapMaybeWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "new c" +mapMaybeWithKeys :: (r -> c -> a -> Maybe b) -> Table r c a -> Table r c b +mapMaybeWithKeys f (Table (rm, cm, _)) = fromMaps rm' cm' + where + rm' = Map.mapWithKey (Map.mapMaybeWithKey . f) rm + cm' = Map.mapWithKey (Map.mapMaybeWithKey . flip f) cm + +-- | /O(n)/, assuming the function @a -> 'Either' a1 a2@ takes /O(1)/. +-- Map values and separate the 'Left' and 'Right' results. +-- +-- > mapEither (\a -> if a == "a" then Left a else Right a) (fromList [(1,'a',"a"),(1,'b',"c"),(2,'b',"a")]) +-- > === (fromList [(1,'a',"a"),(2,'b',"a")],fromList [(1,'b',"c")]) +mapEither :: (a -> Either a1 a2) -> Table r c a -> (Table r c a1, Table r c a2) +mapEither = mapEitherWithKeys . const . const + +-- | /O(n)/, assuming the function @r -> c -> a -> 'Either' a1 a2@ takes /O(1)/. +-- Map (row key, column key, value) triples and separate the 'Left' and 'Right' results. +-- +-- > mapEitherWithKeys (\r c a -> if r == 1 && c == 'a' then Left a else Right a) (fromList [(1,'a',"a"),(1,'b',"c"),(2,'b',"a")]) +-- > === (fromList [(1,'a',"a")],fromList [(1,'b',"c"),(2,'b',"a")]) +mapEitherWithKeys :: (r -> c -> a -> Either a1 a2) -> Table r c a -> (Table r c a1, Table r c a2) +mapEitherWithKeys f (Table (rm, cm, _)) = (fromMaps rm1 cm1, fromMaps rm2 cm2) + where + (rm1, rm2) = (fmap fst &&& fmap snd) $ Map.mapWithKey (Map.mapEitherWithKey . f) rm + (cm1, cm2) = (fmap fst &&& fmap snd) $ Map.mapWithKey (Map.mapEitherWithKey . flip f) cm + +------------------------------------------------------------------------------ +-- * Non exported functions +------------------------------------------------------------------------------ + +assoc :: (a, (b, c)) -> (a, b, c) +assoc (a, (b, c)) = (a, b, c) + +distr :: (a, [(b, c)]) -> [(a, b, c)] +distr = fmap assoc . uncurry (zip . repeat) + +-- | Build a table from a row map and a column map. +fromMaps :: Map r (Map c a) -> Map c (Map r a) -> Table r c a +fromMaps rm cm = Table (rm', cm', size' rm') + where + rm' = nonEmpty rm + cm' = nonEmpty cm + +fromMaps' :: (Ord r, Ord c) => r -> c -> Map r (Map c a) -> Map c (Map r a) -> Table r c a +fromMaps' r c rm cm = Table (rm', cm', size' rm') + where + rm' = nonEmpty' r rm + cm' = nonEmpty' c cm + +nonEmpty :: Map k1 (Map k2 a) -> Map k1 (Map k2 a) +nonEmpty = Map.filter (not . Map.null) + +nonEmpty' :: Ord k1 => k1 -> Map k1 (Map k2 a) -> Map k1 (Map k2 a) +nonEmpty' k1 m = case Map.lookup k1 m of + Just m' | Map.null m' -> Map.delete k1 m + _ -> m + +transpose' :: (Ord r, Ord c) => Map r (Map c a) -> Map c (Map r a) +transpose' = Map.foldrWithKey' f Map.empty + where + f r = Map.unionWith Map.union . Map.map (Map.singleton r) + +size' :: Map k1 (Map k2 a) -> Int +size' = sum . fmap Map.size + +uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d +uncurry3 f ~(a, b, c) = f a b c diff --git a/test-gen/Main.hs b/test-gen/Main.hs new file mode 100644 --- /dev/null +++ b/test-gen/Main.hs @@ -0,0 +1,49 @@ +module Main (main) where + +import Data.List.Extra (replace, stripPrefix, trim) +import Data.Maybe (mapMaybe) +import System.Directory +import System.FilePath + +import Prelude hiding (mod) + +main :: IO () +main = do + genTestsFor "Data.Multimap" + genTestsFor "Data.Multimap.Set" + genTestsFor "Data.Multimap.Table" + +genTestsFor :: String -> IO () +genTestsFor mod = do + let inputFile = "src" replace "." [pathSeparator] mod <.> "hs" + outputFile = "test/hspec" (replace "." [pathSeparator] mod ++ "Spec.hs") + src <- readFile inputFile + createDirectoryIfMissing True (takeDirectory outputFile) + let lns = fmap trim (lines src) + tests = mapMaybe (stripPrefix "-- > ") lns + writeFile outputFile . unlines $ header mod ++ fmap (indent 6) tests + +header :: String -> [String] +header mod = + [ "-- Generated code, do not modify by hand. Generate by running \"stack build && stack exec test-gen\"." + , "" + , "{-# OPTIONS_GHC -w #-}" + , "module " ++ mod ++ "Spec where" + , "" + , "import Test.Hspec" + , "import qualified Data.List.NonEmpty as NonEmpty" + , "import qualified Data.Map as Map" + , "import qualified Data.Set as Set" + ] ++ ["import " ++ mod] ++ + [ "" + , "(===) :: (HasCallStack, Show a, Eq a) => a -> a -> Expectation" + , "(===) = shouldBe" + , "" + , "spec :: Spec" + , "spec = do" + , " describe \"Testing " ++ mod ++ "\" $ do" + , " it \"\" $ do" + ] + +indent :: Int -> String -> String +indent n = (replicate n ' ' ++) diff --git a/test/hspec/Data/Multimap/SetSpec.hs b/test/hspec/Data/Multimap/SetSpec.hs new file mode 100644 --- /dev/null +++ b/test/hspec/Data/Multimap/SetSpec.hs @@ -0,0 +1,117 @@ +-- Generated code, do not modify by hand. Generate by running "stack build && stack exec test-gen". + +{-# OPTIONS_GHC -w #-} +module Data.Multimap.SetSpec where + +import Test.Hspec +import qualified Data.List.NonEmpty as NonEmpty +import qualified Data.Map as Map +import qualified Data.Set as Set +import Data.Multimap.Set + +(===) :: (HasCallStack, Show a, Eq a) => a -> a -> Expectation +(===) = shouldBe + +spec :: Spec +spec = do + describe "Testing Data.Multimap.Set" $ do + it "" $ do + size empty === 0 + singleton 1 'a' === fromList [(1, 'a')] + size (singleton 1 'a') === 1 + fromList ([] :: [(Int, Char)]) === empty + fromList [(1, 'b'), (2, 'a'), (1, 'b')] === fromList [(1, 'b'), (2, 'a')] + insert 1 'a' empty === singleton 1 'a' + insert 1 'a' (fromList [(1, 'b'), (2, 'a')]) === fromList [(1, 'a'), (1, 'b'), (2, 'a')] + insert 1 'a' (fromList [(1, 'a'), (2, 'c')]) === fromList [(1, 'a'), (2, 'c')] + delete 1 (fromList [(1,'a'),(1,'b'),(2,'c')]) === singleton 2 'c' + deleteWithValue 1 'c' (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(1,'b'),(2,'c')] + deleteWithValue 1 'c' (fromList [(2,'c'),(1,'c')]) === singleton 2 'c' + deleteMax 3 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(1,'b'),(2,'c')] + deleteMax 1 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(2,'c')] + deleteMin 3 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'a'),(1,'b'),(2,'c')] + deleteMin 1 (fromList [(1,'a'),(1,'b'),(2,'c')]) === fromList [(1,'b'),(2,'c')] + adjust ("new " ++) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"new a"),(1,"new b"),(2,"c")] + adjust (const "z") 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"z"),(2,"c")] + adjustWithKey (\k x -> show k ++ ":new " ++ x) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) + === fromList [(1,"1:new a"),(1,"1:new b"),(2,"c")] + let f x = if x == "a" then Just "new a" else Nothing in do + update f 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"new a"),(2, "c")] + update f 1 (fromList [(1,"b"),(1,"c"),(2,"c")]) === singleton 2 "c" + let f k x = if x == "a" then Just (show k ++ ":new a") else Nothing in do + updateWithKey f 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"1:new a"),(2,"c")] + updateWithKey f 1 (fromList [(1,"b"),(1,"c"),(2,"c")]) === singleton 2 "c" + let (f, g) = (const Set.empty, Set.insert 'c') in do + alter f 1 (fromList [(1, 'a'), (2, 'b')]) === singleton 2 'b' + alter f 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b')] + alter g 1 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'c'), (1, 'a'), (2, 'b')] + alter g 1 (fromList [(1, 'c'), (2, 'b')]) === fromList [(1, 'c'), (2, 'b')] + alter g 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b'), (3, 'c')] + let (f, g) = (const (const Set.empty), Set.insert . show) in do + alterWithKey f 1 (fromList [(1, "a"), (2, "b")]) === singleton 2 "b" + alterWithKey f 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b")] + alterWithKey g 1 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "1"), (1, "a"), (2, "b")] + alterWithKey g 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b"), (3, "3")] + fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 3 === Set.fromList "ac" + fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 2 === Set.empty + member 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === True + member 1 (deleteMax 1 (fromList [(2, 'c'), (1, 'c')])) === False + notMember 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === False + notMember 1 (deleteMin 1 (fromList [(2, 'c'), (1, 'c')])) === True + Data.Multimap.Set.null empty === True + Data.Multimap.Set.null (singleton 1 'a') === False + notNull empty === False + notNull (singleton 1 'a') === True + size empty === 0 + size (singleton 1 'a') === 1 + size (fromList [(1, 'a'), (2, 'b'), (2, 'c'), (2, 'b')]) === 3 + union (fromList [(1,'a'),(2,'b'),(2,'c')]) (fromList [(1,'d'),(2,'b')]) + === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c')] + unions [fromList [(1,'a'),(2,'b'),(2,'c')], fromList [(1,'d'),(2,'b')]] + === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c')] + difference (fromList [(1,'a'),(2,'b'),(2,'c')]) (fromList [(1,'d'),(2,'b'),(2,'a')]) + === fromList [(1,'a'),(2,'c')] + Data.Multimap.Set.map (++ "x") (fromList [(1,"a"),(2,"b")]) === fromList [(1,"ax"),(2,"bx")] + Data.Multimap.Set.map (const "c") (fromList [(1,"a"),(1,"b"),(2,"b")]) === fromList [(1,"c"),(2,"c")] + mapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1,"a"),(2,"b")]) === fromList [(1,"1:a"),(2,"2:b")] + Data.Multimap.Set.foldr ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + Data.Multimap.Set.foldl (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + foldrWithKey (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + foldlWithKey (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + Data.Multimap.Set.foldr' ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + Data.Multimap.Set.foldl' (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + foldrWithKey' (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + foldlWithKey' (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + foldMapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1, "a"), (1, "c"), (2, "b")]) === "1:a1:c2:b" + elems (fromList [(2,'a'),(1,'b'),(3,'d'),(3,'c'),(1,'b')]) === "bacd" + elems (empty :: SetMultimap Int Char) === [] + keys (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'b')]) === [1,2,3] + keys (empty :: SetMultimap Int Char) === [] + keysSet (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'b')]) === Set.fromList [1,2,3] + keysSet (empty :: SetMultimap Int Char) === Set.empty + toList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'a'),(1,'b'),(2,'a'),(3,'c')] + toAscList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'a'),(1,'b'),(2,'a'),(3,'c')] + toDescList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(3,'c'),(2,'a'),(1,'b'),(1,'a')] + Data.Multimap.Set.filter (> 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === singleton 1 'b' + Data.Multimap.Set.filter (< 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === empty + filterWithKey (\k a -> even k && a > 'a') (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'b')]) === singleton 2 'b' + let f a | a > 'b' = Just True + | a < 'b' = Just False + | a == 'b' = Nothing + in do + filterM f (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'c')]) === Nothing + filterM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Just (fromList [(1,'c'),(2,'c')]) + let f k a | even k && a > 'b' = Just True + | odd k && a < 'b' = Just False + | otherwise = Nothing + in do + filterWithKeyM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Nothing + filterWithKeyM f (fromList [(1,'a'),(3,'a'),(2,'c'),(4,'c')]) === Just (fromList [(2,'c'),(4,'c')]) + mapMaybe (\a -> if a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) + === fromList [(1,"new a"),(2,"new a")] + mapMaybeWithKey (\k a -> if k > 1 && a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) + === singleton 2 "new a" + mapEither (\a -> if a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) + === (fromList [(1,'a'),(2,'a')],fromList [(1,'c'),(2,'c')]) + mapEitherWithKey (\k a -> if even k && a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) + === (fromList [(2,'a')],fromList [(1,'a'),(1,'c'),(2,'c')]) diff --git a/test/hspec/Data/Multimap/TableSpec.hs b/test/hspec/Data/Multimap/TableSpec.hs new file mode 100644 --- /dev/null +++ b/test/hspec/Data/Multimap/TableSpec.hs @@ -0,0 +1,141 @@ +-- Generated code, do not modify by hand. Generate by running "stack build && stack exec test-gen". + +{-# OPTIONS_GHC -w #-} +module Data.Multimap.TableSpec where + +import Test.Hspec +import qualified Data.List.NonEmpty as NonEmpty +import qualified Data.Map as Map +import qualified Data.Set as Set +import Data.Multimap.Table + +(===) :: (HasCallStack, Show a, Eq a) => a -> a -> Expectation +(===) = shouldBe + +spec :: Spec +spec = do + describe "Testing Data.Multimap.Table" $ do + it "" $ do + size empty === 0 + singleton 1 'a' "a" === fromList [(1,'a',"a")] + size (singleton 1 'a' "a") === 1 + fromList ([] :: [(Int, Char, String)]) === empty + fromRowMap (Map.fromList [(1, Map.fromList [('a',"b"),('b',"c")]), (2, Map.fromList [('a',"d")])]) + === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] + fromColumnMap (Map.fromList [(1, Map.fromList [('a',"b"),('b',"c")]), (2, Map.fromList [('a',"d")])]) + === fromList [('a',1,"b"),('a',2,"d"),('b',1,"c")] + transpose (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [('a',1,"b"),('a',2,"d"),('b',1,"c")] + insert 1 'a' "a" empty === singleton 1 'a' "a" + insert 1 'a' "a" (fromList [(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"d")] + insert 1 'a' "a" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"d")] + delete 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'b',"c"),(2,'a',"d")] + delete 1 'a' (fromList [(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'b',"c"),(2,'a',"d")] + deleteRow 1 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 2 'a' "d" + deleteRow 3 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] + deleteColumn 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "c" + deleteColumn 'z' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] + adjust ("new " ++) 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === fromList [(1,'a',"new b"),(1,'b',"c"),(2,'a',"d")] + adjustWithKeys (\r c x -> show r ++ ":" ++ show c ++ ":new " ++ x) 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) + === fromList [(1,'a',"1:'a':new b"),(1,'b',"c"),(2,'a',"d")] + let f x = if x == "b" then Just "new b" else Nothing in do + update f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"new b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + update f 1 'a' (fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + let f r c x = if x == "b" then Just (show r ++ ":" ++ show c ++ ":new b") else Nothing in do + updateWithKeys f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"1:'a':new b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + updateWithKeys f 1 'a' (fromList [(1,'a',"a"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + let (f,g,h) = (const Nothing, const (Just "hello"), fmap ('z':)) in do + alter f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alter f 4 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alter f 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alter g 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"hello"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alter g 4 'e' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c"),(4,'e',"hello")] + alter h 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"zb"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alter h 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + let (f,g) = (\_ _ _ -> Nothing, \r c -> fmap ((show r ++ ":" ++ show c ++ ":") ++)) in do + alterWithKeys f 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alterWithKeys f 4 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alterWithKeys f 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alterWithKeys g 1 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"1:'a':b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + alterWithKeys g 2 'b' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")]) === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(3,'a',"c")] + fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")] !? (1,'a') === Just "b" + fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")] !? (1,'c') === Nothing + fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")] ! (1,'a') === "b" + hasCell (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (1,'a') === True + hasCell (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (1,'c') === False + hasRow (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 1 === True + hasRow (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 3 === False + hasColumn (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 'a' === True + hasColumn (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) 'c' === False + Data.Multimap.Table.null empty === True + Data.Multimap.Table.null (singleton 1 'a' "a") === False + notNull empty === False + notNull (singleton 1 'a' "a") === True + size empty === 0 + size (singleton 1 'a' "a") === 1 + size (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) === 3 + union (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]) + === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] + unions [fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")], fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]] + === fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] + unionWith (++) (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]) + === fromList [(1,'a',"bc"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] + let f r c a a' = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ a' in do + unionWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]) + === fromList [(1,'a',"1:'a':b|c"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] + unionsWith (++) [fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")], fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]] + === fromList [(1,'a',"bc"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] + let f r c a a' = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ a' in do + unionsWithKeys f [fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")], fromList [(1,'a',"c"),(2,'b',"d"),(3,'c',"e")]] + === fromList [(1,'a',"1:'a':b|c"),(1,'b',"c"),(2,'a',"b"),(2,'b',"d"),(3,'c',"e")] + difference (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) (fromList [(1,'a',"c"),(1,'b',"d"),(2,'b',"b")]) + === singleton 2 'a' "b" + Data.Multimap.Table.map (++ "x") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) === fromList [(1,'a',"bx"),(1,'b',"cx"),(2,'a',"bx")] + mapWithKeys (\r c x -> show r ++ ":" ++ show c ++ ":" ++ x) (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) + === fromList [(1,'a',"1:'a':b"),(1,'b',"1:'b':c"),(2,'a',"2:'a':b")] + let f r c a = if odd r && c > 'a' then Just (a ++ "x") else Nothing in do + traverseWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"b")]) === Nothing + traverseWithKeys f (fromList [(1,'b',"b"),(1,'c',"c"),(3,'d',"b")]) === Just (fromList [(1,'b',"bx"),(1,'c',"cx"),(3,'d',"bx")]) + Data.Multimap.Table.foldr (:) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "bcd" + Data.Multimap.Table.foldl (flip (:)) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "dcb" + let f r c a b = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ b in do + foldrWithKeys f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "1:'a':b|1:'b':c|2:'a':d|" + let f a r c b = show r ++ ":" ++ show c ++ ":" ++ b ++ "|" ++ a in do + foldlWithKeys f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "2:'a':d|1:'b':c|1:'a':b|" + Data.Multimap.Table.foldr' (:) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "bcd" + Data.Multimap.Table.foldl' (flip (:)) "" (fromList [(1,'a','b'),(1,'b','c'),(2,'a','d')]) === "dcb" + let f r c a b = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" ++ b in do + foldrWithKeys' f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "1:'a':b|1:'b':c|2:'a':d|" + let f a r c b = show r ++ ":" ++ show c ++ ":" ++ b ++ "|" ++ a in do + foldlWithKeys' f "" (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "2:'a':d|1:'b':c|1:'a':b|" + let f r c a = show r ++ ":" ++ show c ++ ":" ++ a ++ "|" in do + foldMapWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === "1:'a':b|1:'b':c|2:'a':d|" + row 1 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.fromList [('a',"b"),('b',"c")] + row 3 (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.empty + column 'a' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.fromList [(1,"b"),(2,"d")] + column 'c' (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Map.empty + rowMap (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) + === Map.fromList [(1, Map.fromList [('a',"b"),('b',"c")]),(2, Map.fromList [('a',"d")])] + columnMap (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) + === Map.fromList [('a', Map.fromList [(1,"b"),(2,"d")]),('b', Map.fromList [(1,"c")])] + rowKeys (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [1,2] + columnKeys (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === ['a','b'] + rowKeysSet (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Set.fromList [1,2] + columnKeysSet (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === Set.fromList ['a','b'] + toList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] + toRowAscList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")] + toColumnAscList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [('a',1,"b"),('a',2,"d"),('b',1,"c")] + toRowDescList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [(2,'a',"d"),(1,'b',"c"),(1,'a',"b")] + toColumnDescList (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === [('b',1,"c"),('a',2,"d"),('a',1,"b")] + Data.Multimap.Table.filter (> "c") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 2 'a' "d" + Data.Multimap.Table.filter (> "d") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === empty + filterRow even (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 2 'a' "d" + filterColumn (> 'a') (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "c" + filterWithKeys (\r c a -> odd r && c > 'a' && a > "b") (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "c" + mapMaybe (\a -> if a == "a" then Just "new a" else Nothing) (fromList [(1,'a',"a"),(1,'b',"c"),(2,'b',"a")]) + === fromList [(1,'a',"new a"),(2,'b',"new a")] + let f r c a = if r == 1 && a == "c" then Just "new c" else Nothing in do + mapMaybeWithKeys f (fromList [(1,'a',"b"),(1,'b',"c"),(2,'a',"d")]) === singleton 1 'b' "new c" + mapEither (\a -> if a == "a" then Left a else Right a) (fromList [(1,'a',"a"),(1,'b',"c"),(2,'b',"a")]) + === (fromList [(1,'a',"a"),(2,'b',"a")],fromList [(1,'b',"c")]) + mapEitherWithKeys (\r c a -> if r == 1 && c == 'a' then Left a else Right a) (fromList [(1,'a',"a"),(1,'b',"c"),(2,'b',"a")]) + === (fromList [(1,'a',"a")],fromList [(1,'b',"c"),(2,'b',"a")]) diff --git a/test/hspec/Data/MultimapSpec.hs b/test/hspec/Data/MultimapSpec.hs new file mode 100644 --- /dev/null +++ b/test/hspec/Data/MultimapSpec.hs @@ -0,0 +1,127 @@ +-- Generated code, do not modify by hand. Generate by running "stack build && stack exec test-gen". + +{-# OPTIONS_GHC -w #-} +module Data.MultimapSpec where + +import Test.Hspec +import qualified Data.List.NonEmpty as NonEmpty +import qualified Data.Map as Map +import qualified Data.Set as Set +import Data.Multimap + +(===) :: (HasCallStack, Show a, Eq a) => a -> a -> Expectation +(===) = shouldBe + +spec :: Spec +spec = do + describe "Testing Data.Multimap" $ do + it "" $ do + size empty === 0 + singleton 1 'a' === fromList [(1, 'a')] + size (singleton 1 'a') === 1 + fromList ([] :: [(Int, Char)]) === empty + fromMap' (Map.fromList [(1, "ab"), (2, ""), (3, "c")]) === fromList [(1, 'a'), (1, 'b'), (3, 'c')] + insert 1 'a' empty === singleton 1 'a' + insert 1 'a' (fromList [(2, 'b'), (2, 'c')]) === fromList [(1, 'a'), (2, 'b'), (2, 'c')] + insert 1 'a' (fromList [(1, 'b'), (2, 'c')]) === fromList [(1, 'a'), (1, 'b'), (2, 'c')] + delete 1 (fromList [(1, 'a'), (1, 'b'), (2, 'c')]) === singleton 2 'c' + deleteWithValue 1 'c' (fromList [(1, 'a'), (1, 'b'), (2, 'c')]) === fromList [(1, 'a'), (1, 'b'), (2, 'c')] + deleteWithValue 1 'c' (fromList [(1, 'a'), (1, 'b'), (2, 'c'), (1, 'c')]) === fromList [(1, 'a'), (1, 'b'), (2, 'c')] + deleteWithValue 1 'c' (fromList [(2, 'c'), (1, 'c')]) === singleton 2 'c' + deleteOne 1 (fromList [(1, 'a'), (1, 'b'), (2, 'c')]) === fromList [(1, 'b'), (2, 'c')] + deleteOne 1 (fromList [(2, 'c'), (1, 'c')]) === singleton 2 'c' + adjust ("new " ++) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"new a"),(1,"new b"),(2,"c")] + adjustWithKey (\k x -> show k ++ ":new " ++ x) 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) + === fromList [(1,"1:new a"),(1,"1:new b"),(2,"c")] + let f x = if x == "a" then Just "new a" else Nothing in do + update f 1 (fromList [(1,"a"),(1, "b"),(2,"c")]) === fromList [(1,"new a"),(2, "c")] + update f 1 (fromList [(1,"b"),(1, "b"),(2,"c")]) === singleton 2 "c" + update' NonEmpty.tail 1 (fromList [(1, "a"), (1, "b"), (2, "c")]) === fromList [(1, "b"), (2, "c")] + update' NonEmpty.tail 1 (fromList [(1, "a"), (2, "b")]) === singleton 2 "b" + let f k x = if x == "a" then Just (show k ++ ":new a") else Nothing in do + updateWithKey f 1 (fromList [(1,"a"),(1,"b"),(2,"c")]) === fromList [(1,"1:new a"),(2,"c")] + updateWithKey f 1 (fromList [(1,"b"),(1,"b"),(2,"c")]) === singleton 2 "c" + let f k xs = if NonEmpty.length xs == 1 then (show k : NonEmpty.toList xs) else [] in do + updateWithKey' f 1 (fromList [(1, "a"), (1, "b"), (2, "c")]) === singleton 2 "c" + updateWithKey' f 1 (fromList [(1, "a"), (2, "b"), (2, "c")]) === fromList [(1, "1"), (1, "a"), (2, "b"), (2, "c")] + let (f, g) = (const [], ('c':)) in do + alter f 1 (fromList [(1, 'a'), (2, 'b')]) === singleton 2 'b' + alter f 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b')] + alter g 1 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'c'), (1, 'a'), (2, 'b')] + alter g 3 (fromList [(1, 'a'), (2, 'b')]) === fromList [(1, 'a'), (2, 'b'), (3, 'c')] + let (f, g) = (const (const []), (:) . show) in do + alterWithKey f 1 (fromList [(1, "a"), (2, "b")]) === singleton 2 "b" + alterWithKey f 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b")] + alterWithKey g 1 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "1"), (1, "a"), (2, "b")] + alterWithKey g 3 (fromList [(1, "a"), (2, "b")]) === fromList [(1, "a"), (2, "b"), (3, "3")] + fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 3 === "ac" + fromList [(3, 'a'), (5, 'b'), (3, 'c')] ! 2 === [] + member 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === True + member 1 (deleteOne 1 (fromList [(2, 'c'), (1, 'c')])) === False + notMember 1 (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === False + notMember 1 (deleteOne 1 (fromList [(2, 'c'), (1, 'c')])) === True + Data.Multimap.null empty === True + Data.Multimap.null (singleton 1 'a') === False + notNull empty === False + notNull (singleton 1 'a') === True + size empty === 0 + size (singleton 1 'a') === 1 + size (fromList [(1, 'a'), (2, 'b'), (2, 'c')]) === 3 + union (fromList [(1,'a'),(2,'b'),(2,'c')]) (fromList [(1,'d'),(2,'b')]) + === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c'),(2,'b')] + unions [fromList [(1,'a'),(2,'b'),(2,'c')], fromList [(1,'d'),(2,'b')]] + === fromList [(1,'a'),(1,'d'),(2,'b'),(2,'c'),(2,'b')] + difference (fromList [(1,'a'),(2,'b'),(2,'c'),(2,'b')]) (fromList [(1,'d'),(2,'b'),(2,'a')]) + === fromList [(1,'a'), (2,'c'), (2,'b')] + Data.Multimap.map (++ "x") (fromList [(1,"a"),(1,"a"),(2,"b")]) === fromList [(1,"ax"),(1,"ax"),(2,"bx")] + mapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1,"a"),(1,"a"),(2,"b")]) === fromList [(1,"1:a"),(1,"1:a"),(2,"2:b")] + let f k a = if odd k then Just (succ a) else Nothing in do + traverseWithKey f (fromList [(1, 'a'), (1, 'b'), (3, 'b'), (3, 'c')]) === Just (fromList [(1, 'b'), (1, 'c'), (3, 'c'), (3, 'd')]) + traverseWithKey f (fromList [(1, 'a'), (1, 'b'), (2, 'b')]) === Nothing + Data.Multimap.foldr ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + Data.Multimap.foldl (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + foldrWithKey (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + foldlWithKey (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + Data.Multimap.foldr' ((+) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + Data.Multimap.foldl' (\len -> (+ len) . length) 0 (fromList [(1, "hello"), (1, "world"), (2, "!")]) === 11 + foldrWithKey' (\k a len -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + foldlWithKey' (\len k a -> length (show k) + length a + len) 0 (fromList [(1, "hello"), (1, "world"), (20, "!")]) === 15 + foldMapWithKey (\k x -> show k ++ ":" ++ x) (fromList [(1, "a"), (1, "a"), (2, "b")]) === "1:a1:a2:b" + elems (fromList [(2, 'a'), (1, 'b'), (3, 'c'), (1, 'b')]) === "bbac" + elems (empty :: Multimap Int Char) === [] + keys (fromList [(2, 'a'), (1, 'b'), (3, 'c'), (1, 'b')]) === [1,2,3] + keys (empty :: Multimap Int Char) === [] + keysSet (fromList [(2, 'a'), (1, 'b'), (3, 'c'), (1, 'b')]) === Set.fromList [1,2,3] + keysSet (empty :: Multimap Int Char) === Set.empty + assocs (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'b'),(1,'a'),(2,'a'),(3,'c')] + toList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'b'),(1,'a'),(2,'a'),(3,'c')] + toAscList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(1,'b'),(1,'a'),(2,'a'),(3,'c')] + toDescList (fromList [(2,'a'),(1,'b'),(3,'c'),(1,'a')]) === [(3,'c'),(2,'a'),(1,'b'),(1,'a')] + toAscListBF (fromList [("Foo",1),("Foo",2),("Foo",3),("Bar",4),("Bar",5),("Baz",6)]) + === [("Bar",4),("Baz",6),("Foo",1),("Bar",5),("Foo",2),("Foo",3)] + toDescListBF (fromList [("Foo",1),("Foo",2),("Foo",3),("Bar",4),("Bar",5),("Baz",6)]) + === [("Foo",1),("Baz",6),("Bar",4),("Foo",2),("Bar",5),("Foo",3)] + Data.Multimap.filter (> 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === singleton 1 'b' + Data.Multimap.filter (< 'a') (fromList [(1,'a'),(1,'b'),(2,'a')]) === empty + filterKey even (fromList [(1,'a'),(1,'b'),(2,'a')]) === singleton 2 'a' + filterWithKey (\k a -> even k && a > 'a') (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'b')]) === singleton 2 'b' + let f a | a > 'b' = Just True + | a < 'b' = Just False + | a == 'b' = Nothing + in do + filterM f (fromList [(1,'a'),(1,'b'),(2,'a'),(2,'c')]) === Nothing + filterM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Just (fromList [(1,'c'),(2,'c')]) + let f k a | even k && a > 'b' = Just True + | odd k && a < 'b' = Just False + | otherwise = Nothing + in do + filterWithKeyM f (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) === Nothing + filterWithKeyM f (fromList [(1,'a'),(1,'a'),(2,'c'),(2,'c')]) === Just (fromList [(2,'c'),(2,'c')]) + mapMaybe (\a -> if a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) + === fromList [(1,"new a"),(2,"new a")] + mapMaybeWithKey (\k a -> if k > 1 && a == "a" then Just "new a" else Nothing) (fromList [(1,"a"),(1,"b"),(2,"a"),(2,"c")]) + === singleton 2 "new a" + mapEither (\a -> if a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) + === (fromList [(1,'a'),(2,'a')],fromList [(1,'c'),(2,'c')]) + mapEitherWithKey (\k a -> if even k && a < 'b' then Left a else Right a) (fromList [(1,'a'),(1,'c'),(2,'a'),(2,'c')]) + === (fromList [(2,'a')],fromList [(1,'a'),(1,'c'),(2,'c')]) diff --git a/test/hspec/Main.hs b/test/hspec/Main.hs new file mode 100644 --- /dev/null +++ b/test/hspec/Main.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}