-- 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.Seq as Seq

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 = a -> k -> Map k a -> a
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault a
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 = (k -> a -> Bool) -> Map k a -> Map k a
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 = Map k a -> Set k -> Map k a
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys Map k a
m ([k] -> Set k
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 = (Map k v -> (k, v) -> Map k v) -> Map k v -> [(k, v)] -> Map k v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\Map k v
m (k
k, v
v) -> k -> v -> Map k v -> Map k 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) = k -> Map k a -> (Map k a, Maybe a, Map k a)
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' = Map k a -> (a -> Map k a) -> Maybe a -> Map k a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Map k a
post (\a
v -> k -> a -> Map k a -> Map k a
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) = k -> Map k a -> (Map k a, Map k a)
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) = k -> Map k a -> (Map k a, Map k a)
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
_) = k -> k -> Map k a -> (Map k a, Map k a, 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 (k -> Map k v -> Maybe (k, v)
forall k v. Ord k => k -> Map k v -> Maybe (k, v)
Map.lookupLT k
key Map k v
m, k -> Map k v -> Maybe (k, v)
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 k -> k -> k
forall a. Num a => a -> a -> a
- k
k0 k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
k1 k -> k -> k
forall a. Num a => a -> a -> a
- k
key -> (k, v) -> Maybe (k, v)
forall a. a -> Maybe a
Just (k
k0, v
v0)
        | Bool
otherwise -> (k, v) -> Maybe (k, v)
forall a. a -> Maybe a
Just (k
k1, v
v1)
    (Maybe (k, v)
Nothing, Just (k, v)
kv) -> (k, v) -> Maybe (k, v)
forall a. a -> Maybe a
Just (k, v)
kv
    (Just (k, v)
kv, Maybe (k, v)
Nothing) -> (k, v) -> Maybe (k, v)
forall a. a -> Maybe a
Just (k, v)
kv
    (Maybe (k, v)
Nothing, Maybe (k, v)
Nothing) -> Maybe (k, v)
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 = [(a, k)] -> Map a k
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(a, k)] -> Map a k)
-> (Map k a -> [(a, k)]) -> Map k a -> Map a k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, a) -> (a, k)) -> [(k, a)] -> [(a, k)]
forall a b. (a -> b) -> [a] -> [b]
map (\(k
x, a
y) -> (a
y, k
x)) ([(k, a)] -> [(a, k)])
-> (Map k a -> [(k, a)]) -> Map k a -> [(a, k)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k a -> [(k, a)]
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 = [(k, [a])] -> Map k [a]
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(k, [a])] -> Map k [a])
-> ([(k, a)] -> [(k, [a])]) -> [(k, a)] -> Map k [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(k, a)] -> (k, [a])) -> [[(k, a)]] -> [(k, [a])]
forall a b. (a -> b) -> [a] -> [b]
map (\[(k, a)]
gs -> ((k, a) -> k
forall a b. (a, b) -> a
fst ([(k, a)] -> (k, a)
forall a. [a] -> a
head [(k, a)]
gs), ((k, a) -> a) -> [(k, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (k, a) -> a
forall a b. (a, b) -> b
snd [(k, a)]
gs))
    ([[(k, a)]] -> [(k, [a])])
-> ([(k, a)] -> [[(k, a)]]) -> [(k, a)] -> [(k, [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, a) -> (k, a) -> Bool) -> [(k, a)] -> [[(k, a)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
List.groupBy (k -> k -> Bool
forall a. Eq a => a -> a -> Bool
(==) (k -> k -> Bool) -> ((k, a) -> k) -> (k, a) -> (k, a) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (k, a) -> k
forall a b. (a, b) -> a
fst) ([(k, a)] -> [[(k, a)]])
-> ([(k, a)] -> [(k, a)]) -> [(k, a)] -> [[(k, a)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, a) -> k) -> [(k, a)] -> [(k, a)]
forall k a. Ord k => (a -> k) -> [a] -> [a]
Seq.sort_on (k, a) -> k
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 = ([(a, b)] -> Map a b
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(a, b)]
pairs, [[(a, b)]] -> [(a, b)]
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 = [((a, b), [(a, b)])] -> ([(a, b)], [[(a, b)]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([((a, b), [(a, b)])] -> ([(a, b)], [[(a, b)]]))
-> ([(a, b)] -> [((a, b), [(a, b)])])
-> [(a, b)]
-> ([(a, b)], [[(a, b)]])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(a, b)] -> ((a, b), [(a, b)]))
-> [[(a, b)]] -> [((a, b), [(a, b)])]
forall a b. (a -> b) -> [a] -> [b]
map [(a, b)] -> ((a, b), [(a, b)])
forall {a}. [a] -> (a, [a])
pair ([[(a, b)]] -> [((a, b), [(a, b)])])
-> ([(a, b)] -> [[(a, b)]]) -> [(a, b)] -> [((a, b), [(a, b)])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> (a, b) -> Bool) -> [(a, b)] -> [[(a, b)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
List.groupBy (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (a -> a -> Bool) -> ((a, b) -> a) -> (a, b) -> (a, b) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, b) -> a
forall a b. (a, b) -> a
fst) ([(a, b)] -> [[(a, b)]])
-> ([(a, b)] -> [(a, b)]) -> [(a, b)] -> [[(a, b)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> a) -> [(a, b)] -> [(a, b)]
forall k a. Ord k => (a -> k) -> [a] -> [a]
Seq.sort_on (a, b) -> a
forall a b. (a, b) -> a
fst
    pair :: [a] -> (a, [a])
pair (a
x:[a]
xs) = (a
x, [a]
xs)
    pair [] = [Char] -> (a, [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"[]: List.groupBy violated its postcondition"
    ([(a, b)]
pairs, [[(a, b)]]
rest) = [(a, b)] -> ([(a, b)], [[(a, b)]])
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 = ([(k, v)] -> Map k v)
-> ([(k, v)], [(k, [v])]) -> (Map k v, [(k, [v])])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first [(k, v)] -> Map k v
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList (([(k, v)], [(k, [v])]) -> (Map k v, [(k, [v])]))
-> ([(k, v)] -> ([(k, v)], [(k, [v])]))
-> [(k, v)]
-> (Map k v, [(k, [v])])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Either (k, v) (k, [v])] -> ([(k, v)], [(k, [v])])
forall a b. [Either a b] -> ([a], [b])
Either.partitionEithers ([Either (k, v) (k, [v])] -> ([(k, v)], [(k, [v])]))
-> ([(k, v)] -> [Either (k, v) (k, [v])])
-> [(k, v)]
-> ([(k, v)], [(k, [v])])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, [v]) -> Either (k, v) (k, [v]))
-> [(k, [v])] -> [Either (k, v) (k, [v])]
forall a b. (a -> b) -> [a] -> [b]
map (k, [v]) -> Either (k, v) (k, [v])
forall {a} {b}. (a, [b]) -> Either (a, b) (a, [b])
separate
    ([(k, [v])] -> [Either (k, v) (k, [v])])
-> ([(k, v)] -> [(k, [v])]) -> [(k, v)] -> [Either (k, v) (k, [v])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, v)] -> [(k, [v])]
forall a b. Ord a => [(a, b)] -> [(a, NonNull b)]
Seq.group_fst
    where
    separate :: (a, [b]) -> Either (a, b) (a, [b])
separate (a
k, [b
v]) = (a, b) -> Either (a, b) (a, [b])
forall a b. a -> Either a b
Left (a
k, b
v)
    separate (a
k, [b]
vs) = (a, [b]) -> Either (a, b) (a, [b])
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) <- Map k v1 -> [(k, v1)]
forall k a. Map k a -> [(k, a)]
Map.assocs Map k v1
map1, Just v2
v2 <- [k -> Map k v2 -> Maybe 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, Seq.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 = [(k, v1)] -> [(k, v2)] -> [(k, Paired v1 v2)]
forall k a b. Ord k => [(k, a)] -> [(k, b)] -> [(k, Paired a b)]
Seq.pair_sorted (Map k v1 -> [(k, v1)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map k v1
map1) (Map k v2 -> [(k, v2)]
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 (Seq.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 = [(k, Paired v1 v2)] -> Map k (Paired v1 v2)
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList (Map k v1 -> Map k v2 -> [(k, Paired v1 v2)]
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 = (Map k a -> Map k a -> Map k a
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 = Map k a -> Map k a -> Map k a
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 = (((Map k a, Map k a) -> Map k a -> (Map k a, Map k a))
 -> (Map k a, Map k a) -> [Map k a] -> (Map k a, Map k a))
-> (Map k a, Map k a)
-> ((Map k a, Map k a) -> Map k a -> (Map k a, Map k a))
-> [Map k a]
-> (Map k a, Map k a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Map k a, Map k a) -> Map k a -> (Map k a, Map k a))
-> (Map k a, Map k a) -> [Map k a] -> (Map k a, Map k a)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (Map k a
forall k a. Map k a
Map.empty, Map k a
forall k a. Map k a
Map.empty) (((Map k a, Map k a) -> Map k a -> (Map k a, Map k a))
 -> [Map k a] -> (Map k a, Map k a))
-> ((Map k a, Map k a) -> Map k a -> (Map k a, Map k a))
-> [Map k a]
-> (Map k a, Map k a)
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) = Map k a -> Map k a -> (Map k a, Map k a)
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, Map k a -> Map k a -> Map k a
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 = (a -> a -> a) -> Map k a -> Map k a -> Map k a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith a -> a -> a
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 = [(k, a)] -> Map k a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(k
k, [a] -> a
forall a. Monoid a => [a] -> a
Monoid.mconcat (k -> [a]
get k
k)) | k
k <- [k]
ks]
    where
    ks :: [k]
ks = Set k -> [k]
forall a. Set a -> [a]
Set.toList (Set k -> [k]) -> Set k -> [k]
forall a b. (a -> b) -> a -> b
$ [Set k] -> Set k
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set k] -> Set k) -> [Set k] -> Set k
forall a b. (a -> b) -> a -> b
$ (Map k a -> Set k) -> [Map k a] -> [Set k]
forall a b. (a -> b) -> [a] -> [b]
map Map k a -> Set k
forall k a. Map k a -> Set k
Map.keysSet [Map k a]
ms
    get :: k -> [a]
get k
k = (Map k a -> Maybe a) -> [Map k a] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
Maybe.mapMaybe (k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k) [Map k a]
ms