-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

-- | Extra utils for "Data.Map".
module Util.Maps where
import           Prelude hiding (min, max)
import           Control.Arrow (first)
import qualified Data.Either as Either
import qualified Data.List as List
import qualified Data.Map as Map
import           Data.Map (Map)
import qualified Data.Maybe as Maybe
import qualified Data.Monoid as Monoid
import qualified Data.Set as Set

import qualified Util.Lists as Lists

import           Data.Function (on)


getM :: (Ord k, Monoid a) => Map k a -> k -> a
getM :: forall k a. (Ord k, Monoid a) => Map k a -> k -> a
getM Map k a
m k
k = forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault forall a. Monoid a => a
mempty k
k Map k a
m

filterKey :: (k -> Bool) -> Map k a -> Map k a
filterKey :: forall k a. (k -> Bool) -> Map k a -> Map k a
filterKey k -> Bool
f = forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\k
k a
_ -> k -> Bool
f k
k)

deleteKeys :: Ord k => [k] -> Map k a -> Map k a
deleteKeys :: forall k a. Ord k => [k] -> Map k a -> Map k a
deleteKeys [k]
keys Map k a
m = forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys Map k a
m (forall a. Ord a => [a] -> Set a
Set.fromList [k]
keys)

insertList :: Ord k => [(k, v)] -> Map k v -> Map k v
insertList :: forall k v. Ord k => [(k, v)] -> Map k v -> Map k v
insertList [(k, v)]
kvs Map k v
m = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\Map k v
m (k
k, v
v) -> forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
m) Map k v
m [(k, v)]
kvs

-- | Like 'Data.Map.split', except include a matched key in the above map.
split2 :: Ord k => k -> Map k a -> (Map k a, Map k a)
split2 :: forall k a. Ord k => k -> Map k a -> (Map k a, Map k a)
split2 k
k Map k a
fm = (Map k a
pre, Map k a
post')
    where
    (Map k a
pre, Maybe a
at, Map k a
post) = forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
Map.splitLookup k
k Map k a
fm
    post' :: Map k a
post' = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Map k a
post (\a
v -> forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k a
v Map k a
post) Maybe a
at

-- | Split the map into the maps below, within, and above the given range.
-- @low@ to @high@ is half-open, as usual.
split3 :: Ord k => k -> k -> Map k a -> (Map k a, Map k a, Map k a)
split3 :: forall k a.
Ord k =>
k -> k -> Map k a -> (Map k a, Map k a, Map k a)
split3 k
low k
high Map k a
fm = (Map k a
below, Map k a
within, Map k a
way_above)
    where
    (Map k a
below, Map k a
above) = forall k a. Ord k => k -> Map k a -> (Map k a, Map k a)
split2 k
low Map k a
fm
    (Map k a
within, Map k a
way_above) = forall k a. Ord k => k -> Map k a -> (Map k a, Map k a)
split2 k
high Map k a
above

-- | Return the subset of the map that is between a half-open low and high key.
within :: Ord k => k -> k -> Map k a -> Map k a
within :: forall k a. Ord k => k -> k -> Map k a -> Map k a
within k
low k
high Map k a
fm = let (Map k a
_, Map k a
m, Map k a
_) = forall k a.
Ord k =>
k -> k -> Map k a -> (Map k a, Map k a, Map k a)
split3 k
low k
high Map k a
fm in Map k a
m

-- | Find the closest key.  If two are equidistant, favor the one below.
lookupClosest :: (Ord k, Num k) => k -> Map k v -> Maybe (k, v)
lookupClosest :: forall k v. (Ord k, Num k) => k -> Map k v -> Maybe (k, v)
lookupClosest k
key Map k v
m = case (forall k v. Ord k => k -> Map k v -> Maybe (k, v)
Map.lookupLT k
key Map k v
m, forall k v. Ord k => k -> Map k v -> Maybe (k, v)
Map.lookupGE k
key Map k v
m) of
    (Just (k
k0, v
v0), Just (k
k1, v
v1))
        | k
key forall a. Num a => a -> a -> a
- k
k0 forall a. Ord a => a -> a -> Bool
<= k
k1 forall a. Num a => a -> a -> a
- k
key -> forall a. a -> Maybe a
Just (k
k0, v
v0)
        | Bool
otherwise -> forall a. a -> Maybe a
Just (k
k1, v
v1)
    (Maybe (k, v)
Nothing, Just (k, v)
kv) -> forall a. a -> Maybe a
Just (k, v)
kv
    (Just (k, v)
kv, Maybe (k, v)
Nothing) -> forall a. a -> Maybe a
Just (k, v)
kv
    (Maybe (k, v)
Nothing, Maybe (k, v)
Nothing) -> forall a. Maybe a
Nothing

invert :: Ord a => Map k a -> Map a k
invert :: forall a k. Ord a => Map k a -> Map a k
invert = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\(k
x, a
y) -> (a
y, k
x)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.assocs

-- | TODO Would it be more efficient to do 'fromListWith (++)'?
multimap :: Ord k => [(k, a)] -> Map k [a]
multimap :: forall k a. Ord k => [(k, a)] -> Map k [a]
multimap = forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\[(k, a)]
gs -> (forall a b. (a, b) -> a
fst (forall a. [a] -> a
head [(k, a)]
gs), forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(k, a)]
gs))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Bool) -> [a] -> [[a]]
List.groupBy (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> k) -> [a] -> [a]
Lists.sortOn forall a b. (a, b) -> a
fst

-- | Like Map.fromList, but only accept the first of duplicate keys, and also
-- return the rejected duplicates.
unique :: Ord a => [(a, b)] -> (Map a b, [(a, b)])
unique :: forall a b. Ord a => [(a, b)] -> (Map a b, [(a, b)])
unique [(a, b)]
assocs = (forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(a, b)]
pairs, forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(a, b)]]
rest)
    where
    -- List.sort is stable, so only the first keys will make it into the map.
    separate :: [(a, b)] -> ([(a, b)], [[(a, b)]])
separate = forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall {a}. [a] -> (a, [a])
pair forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Bool) -> [a] -> [[a]]
List.groupBy (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> k) -> [a] -> [a]
Lists.sortOn forall a b. (a, b) -> a
fst
    pair :: [a] -> (a, [a])
pair (a
x:[a]
xs) = (a
x, [a]
xs)
    pair [] = forall a. HasCallStack => [Char] -> a
error [Char]
"[]: List.groupBy violated its postcondition"
    ([(a, b)]
pairs, [[(a, b)]]
rest) = forall {b}. [(a, b)] -> ([(a, b)], [[(a, b)]])
separate [(a, b)]
assocs

-- | Make a map, but if any keys collide, omit that key and return it along
-- with the multiple values.
unique2 :: Ord k => [(k, v)] -> (Map k v, [(k, [v])])
unique2 :: forall k v. Ord k => [(k, v)] -> (Map k v, [(k, [v])])
unique2 = forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [Either a b] -> ([a], [b])
Either.partitionEithers forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b}. (a, [b]) -> Either (a, b) (a, [b])
separate
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Ord a => [(a, b)] -> [(a, NonNull b)]
Lists.groupFst
    where
    separate :: (a, [b]) -> Either (a, b) (a, [b])
separate (a
k, [b
v]) = forall a b. a -> Either a b
Left (a
k, b
v)
    separate (a
k, [b]
vs) = forall a b. b -> Either a b
Right (a
k, [b]
vs)

-- | Given two maps, pair up the elements in @map1@ with a samed-keyed element
-- in @map2@, if there is one.  Elements that are only in @map1@ or @map2@ will
-- not be included in the output.
zipIntersection :: Ord k => Map k v1 -> Map k v2 -> [(k, v1, v2)]
zipIntersection :: forall k v1 v2. Ord k => Map k v1 -> Map k v2 -> [(k, v1, v2)]
zipIntersection Map k v1
map1 Map k v2
map2 =
    [(k
k, v1
v1, v2
v2) | (k
k, v1
v1) <- forall k a. Map k a -> [(k, a)]
Map.assocs Map k v1
map1, Just v2
v2 <- [forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v2
map2]]
    -- I could implement with 'pairs', but it would be less efficient.

-- | Pair up elements from each map with equal keys.
pairs :: Ord k => Map k v1 -> Map k v2 -> [(k, Lists.Paired v1 v2)]
pairs :: forall k v1 v2.
Ord k =>
Map k v1 -> Map k v2 -> [(k, Paired v1 v2)]
pairs Map k v1
map1 Map k v2
map2 = forall k a b. Ord k => [(k, a)] -> [(k, b)] -> [(k, Paired a b)]
Lists.pairSorted (forall k a. Map k a -> [(k, a)]
Map.toAscList Map k v1
map1) (forall k a. Map k a -> [(k, a)]
Map.toAscList Map k v2
map2)
{-# INLINE pairs #-}

paired :: Ord k => Map k v1 -> Map k v2 -> Map k (Lists.Paired v1 v2)
paired :: forall k v1 v2.
Ord k =>
Map k v1 -> Map k v2 -> Map k (Paired v1 v2)
paired Map k v1
map1 Map k v2
map2 = forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList (forall k v1 v2.
Ord k =>
Map k v1 -> Map k v2 -> [(k, Paired v1 v2)]
pairs Map k v1
map1 Map k v2
map2)
{-# INLINE paired #-}

-- | Like 'Map.union', but also return a map of rejected duplicate keys from
-- the map on the right.
uniqueUnion :: Ord k => Map k a -> Map k a -> (Map k a, Map k a)
    -- ^ (union, rejected)
uniqueUnion :: forall k a. Ord k => Map k a -> Map k a -> (Map k a, Map k a)
uniqueUnion Map k a
fm1 Map k a
fm2 = (forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map k a
fm1 Map k a
fm2, Map k a
rejected)
    where rejected :: Map k a
rejected = forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.intersection Map k a
fm2 Map k a
fm1

-- | Like 'Map.unions', but return a map of the rejected keys.  Like
-- Map.unions, the first key in the list wins.  If there are multiple
-- conflicting keys, only the first one will show up in the reject map.
uniqueUnions :: Ord k => [Map k a] -> (Map k a, Map k a)
uniqueUnions :: forall k a. Ord k => [Map k a] -> (Map k a, Map k a)
uniqueUnions = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (forall k a. Map k a
Map.empty, forall k a. Map k a
Map.empty) forall a b. (a -> b) -> a -> b
$
    \(Map k a
collect, Map k a
rejected) Map k a
fm ->
        let (Map k a
collect2, Map k a
rejected2) = forall k a. Ord k => Map k a -> Map k a -> (Map k a, Map k a)
uniqueUnion Map k a
collect Map k a
fm
        in (Map k a
collect2, forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map k a
rejected Map k a
rejected2)

-- | The Data.Map Monoid instance is just a plain union, and doesn't mappend
-- the values.
mappend :: (Ord k, Monoid.Monoid a) => Map k a -> Map k a -> Map k a
mappend :: forall k a. (Ord k, Monoid a) => Map k a -> Map k a -> Map k a
mappend = forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. Monoid a => a -> a -> a
Monoid.mappend

-- | Like @Map.unionsWith Monoid.mappend@, but collects the values together,
-- in case mconcat is more efficient than successive mappends.
mconcat :: (Ord k, Monoid.Monoid a) => [Map k a] -> Map k a
mconcat :: forall k a. (Ord k, Monoid a) => [Map k a] -> Map k a
mconcat [Map k a]
ms = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(k
k, forall a. Monoid a => [a] -> a
Monoid.mconcat (k -> [a]
get k
k)) | k
k <- [k]
ks]
    where
    ks :: [k]
ks = forall a. Set a -> [a]
Set.toList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall k a. Map k a -> Set k
Map.keysSet [Map k a]
ms
    get :: k -> [a]
get k
k = forall a b. (a -> Maybe b) -> [a] -> [b]
Maybe.mapMaybe (forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k) [Map k a]
ms