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
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
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
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
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
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
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
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
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)
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]]
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 #-}
uniqueUnion :: Ord k => Map k a -> Map k a -> (Map k a, Map k a)
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
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)
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
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