-- | @Map@ utilities.
module Unison.Util.Map
  ( alignWithKey,
    bimap,
    bitraverse,
    bitraversed,
    deleteLookup,
    deleteLookupJust,
    elemsSet,
    foldKeysCommutative,
    foldValuesCommutative,
    foldM,
    foldMapM,
    for_,
    fromSetA,
    insertLookup,
    invert,
    lookupJust,
    mergeMap,
    unionWithM,
    remap,
    thenInsertPair,
    traverseKeys,
    traverseKeysWith,
    search,
    searchr,
    swap,
    upsert,
    upsertF,
    upsertLookup,
    valuesVector,
    asList_,
  )
where

import Control.Lens hiding (bimap)
import Control.Monad qualified as Monad
import Data.Bifunctor qualified as B
import Data.Bitraversable qualified as B
import Data.Foldable (foldlM)
import Data.Map.Internal qualified as Map (Map (Bin, Tip))
import Data.Map.Merge.Strict qualified as Map
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Data.Set.Internal qualified as Set (Set (..))
import Data.These (These (..))
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Unison.Prelude hiding (bimap, foldM, for_)

-- | A common case of @Map.merge@. Like @alignWith@, but includes the key.
alignWithKey :: (Ord k) => (k -> These a b -> c) -> Map k a -> Map k b -> Map k c
alignWithKey :: forall k a b c.
Ord k =>
(k -> These a b -> c) -> Map k a -> Map k b -> Map k c
alignWithKey k -> These a b -> c
f =
  SimpleWhenMissing k a c
-> SimpleWhenMissing k b c
-> SimpleWhenMatched k a b c
-> Map k a
-> Map k b
-> Map k c
forall k a c b.
Ord k =>
SimpleWhenMissing k a c
-> SimpleWhenMissing k b c
-> SimpleWhenMatched k a b c
-> Map k a
-> Map k b
-> Map k c
Map.merge
    ((k -> a -> c) -> SimpleWhenMissing k a c
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> y) -> WhenMissing f k x y
Map.mapMissing \k
k a
x -> k -> These a b -> c
f k
k (a -> These a b
forall a b. a -> These a b
This a
x))
    ((k -> b -> c) -> SimpleWhenMissing k b c
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> y) -> WhenMissing f k x y
Map.mapMissing \k
k b
y -> k -> These a b -> c
f k
k (b -> These a b
forall a b. b -> These a b
That b
y))
    ((k -> a -> b -> c) -> SimpleWhenMatched k a b c
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> z) -> WhenMatched f k x y z
Map.zipWithMatched \k
k a
x b
y -> k -> These a b -> c
f k
k (a -> b -> These a b
forall a b. a -> b -> These a b
These a
x b
y))

bimap :: (Ord a') => (a -> a') -> (b -> b') -> Map a b -> Map a' b'
bimap :: forall a' a b b'.
Ord a' =>
(a -> a') -> (b -> b') -> Map a b -> Map a' b'
bimap a -> a'
fa b -> b'
fb = [(a', b')] -> Map a' b'
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(a', b')] -> Map a' b')
-> (Map a b -> [(a', b')]) -> Map a b -> Map a' b'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> (a', b')) -> [(a, b)] -> [(a', b')]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> a') -> (b -> b') -> (a, b) -> (a', b')
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
B.bimap a -> a'
fa b -> b'
fb) ([(a, b)] -> [(a', b')])
-> (Map a b -> [(a, b)]) -> Map a b -> [(a', b')]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
Map.toList

bitraverse :: (Applicative f, Ord a') => (a -> f a') -> (b -> f b') -> Map a b -> f (Map a' b')
bitraverse :: forall (f :: * -> *) a' a b b'.
(Applicative f, Ord a') =>
(a -> f a') -> (b -> f b') -> Map a b -> f (Map a' b')
bitraverse a -> f a'
fa b -> f b'
fb = ([(a', b')] -> Map a' b') -> f [(a', b')] -> f (Map a' b')
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(a', b')] -> Map a' b'
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (f [(a', b')] -> f (Map a' b'))
-> (Map a b -> f [(a', b')]) -> Map a b -> f (Map a' b')
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> f (a', b')) -> [(a, b)] -> f [(a', b')]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((a -> f a') -> (b -> f b') -> (a, b) -> f (a', b')
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> (a, b) -> f (c, d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
B.bitraverse a -> f a'
fa b -> f b'
fb) ([(a, b)] -> f [(a', b')])
-> (Map a b -> [(a, b)]) -> Map a b -> f [(a', b')]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
Map.toList

bitraversed :: (Ord a', Ord k') => Traversal k k' a a' -> Traversal v v' a a' -> Traversal (Map k v) (Map k' v') a a'
bitraversed :: forall a' k' k a v v'.
(Ord a', Ord k') =>
Traversal k k' a a'
-> Traversal v v' a a' -> Traversal (Map k v) (Map k' v') a a'
bitraversed Traversal k k' a a'
keyT Traversal v v' a a'
valT a -> f a'
f Map k v
m =
  (k -> f k') -> (v -> f v') -> Map k v -> f (Map k' v')
forall (f :: * -> *) a' a b b'.
(Applicative f, Ord a') =>
(a -> f a') -> (b -> f b') -> Map a b -> f (Map a' b')
bitraverse ((a -> f a') -> k -> f k'
Traversal k k' a a'
keyT a -> f a'
f) ((a -> f a') -> v -> f v'
Traversal v v' a a'
valT a -> f a'
f) Map k v
m

-- | Traverse a map as a list of key-value pairs.
-- Note: This can have unexpected results if the result contains duplicate keys.
asList_ :: (Ord k') => Traversal (Map k v) (Map k' v') [(k, v)] [(k', v')]
asList_ :: forall k' k v v'.
Ord k' =>
Traversal (Map k v) (Map k' v') [(k, v)] [(k', v')]
asList_ [(k, v)] -> f [(k', v')]
f Map k v
s =
  Map k v
s
    Map k v -> (Map k v -> [(k, v)]) -> [(k, v)]
forall a b. a -> (a -> b) -> b
& Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
Map.toList
    [(k, v)] -> ([(k, v)] -> f [(k', v')]) -> f [(k', v')]
forall a b. a -> (a -> b) -> b
& [(k, v)] -> f [(k', v')]
f
    f [(k', v')] -> ([(k', v')] -> Map k' v') -> f (Map k' v')
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> [(k', v')] -> Map k' v'
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList

-- | Search a map, given a monotone ordering function on keys. Summarizes the key/value pairs of the (possibly empty)
-- contiguous block that compares equal.
search :: (Monoid m) => (k -> v -> m) -> (k -> Ordering) -> Map k v -> m
search :: forall m k v.
Monoid m =>
(k -> v -> m) -> (k -> Ordering) -> Map k v -> m
search k -> v -> m
f k -> Ordering
keyOrdering =
  Map k v -> m
go
  where
    go :: Map k v -> m
go = \case
      Map.Bin Size
_ k
k v
v Map k v
l Map k v
r ->
        case k -> Ordering
keyOrdering k
k of
          Ordering
EQ -> Map k v -> m
goL Map k v
l m -> m -> m
forall a. Semigroup a => a -> a -> a
<> k -> v -> m
f k
k v
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> Map k v -> m
goR Map k v
r
          Ordering
LT -> Map k v -> m
go Map k v
r
          Ordering
GT -> Map k v -> m
go Map k v
l
      Map k v
Map.Tip -> m
forall a. Monoid a => a
mempty

    goL :: Map k v -> m
goL = \case
      Map.Bin Size
_ k
k v
v Map k v
l Map k v
r ->
        case k -> Ordering
keyOrdering k
k of
          Ordering
EQ -> Map k v -> m
goL Map k v
l m -> m -> m
forall a. Semigroup a => a -> a -> a
<> k -> v -> m
f k
k v
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (k -> v -> m -> m) -> m -> Map k v -> m
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\k
k v
v m
acc -> k -> v -> m
f k
k v
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> m
acc) m
forall a. Monoid a => a
mempty Map k v
r
          Ordering
LT -> Map k v -> m
goL Map k v
r
          Ordering
GT -> [Char] -> m
forall a. HasCallStack => [Char] -> a
error [Char]
"predicate not monotone with respect to ordering"
      Map k v
Map.Tip -> m
forall a. Monoid a => a
mempty

    goR :: Map k v -> m
goR = \case
      Map.Bin Size
_ k
k v
v Map k v
l Map k v
r ->
        case k -> Ordering
keyOrdering k
k of
          Ordering
EQ -> (k -> v -> m -> m) -> m -> Map k v -> m
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\k
k v
v m
acc -> k -> v -> m
f k
k v
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> m
acc) m
forall a. Monoid a => a
mempty Map k v
l m -> m -> m
forall a. Semigroup a => a -> a -> a
<> k -> v -> m
f k
k v
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> Map k v -> m
goR Map k v
r
          Ordering
GT -> Map k v -> m
goR Map k v
l
          Ordering
LT -> [Char] -> m
forall a. HasCallStack => [Char] -> a
error [Char]
"predicate not monotone with respect to ordering"
      Map k v
Map.Tip -> m
forall a. Monoid a => a
mempty

searchr :: (k -> v -> acc -> acc) -> acc -> (k -> Ordering) -> Map k v -> acc
searchr :: forall k v acc.
(k -> v -> acc -> acc) -> acc -> (k -> Ordering) -> Map k v -> acc
searchr k -> v -> acc -> acc
f acc
z k -> Ordering
keyOrdering =
  acc -> Map k v -> acc
go acc
z
  where
    go :: acc -> Map k v -> acc
go acc
acc = \case
      Map.Bin Size
_ k
k v
v Map k v
l Map k v
r ->
        case k -> Ordering
keyOrdering k
k of
          Ordering
EQ -> acc -> Map k v -> acc
goL (k -> v -> acc -> acc
f k
k v
v (acc -> Map k v -> acc
goR acc
acc Map k v
r)) Map k v
l -- goL l <> f k v <> goR r
          Ordering
LT -> acc -> Map k v -> acc
go acc
acc Map k v
r
          Ordering
GT -> acc -> Map k v -> acc
go acc
acc Map k v
l
      Map k v
Map.Tip -> acc
acc

    goL :: acc -> Map k v -> acc
goL acc
acc = \case
      Map.Bin Size
_ k
k v
v Map k v
l Map k v
r ->
        case k -> Ordering
keyOrdering k
k of
          Ordering
EQ -> acc -> Map k v -> acc
goL (k -> v -> acc -> acc
f k
k v
v ((k -> v -> acc -> acc) -> acc -> Map k v -> acc
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey k -> v -> acc -> acc
f acc
acc Map k v
r)) Map k v
l
          Ordering
LT -> acc -> Map k v -> acc
goL acc
acc Map k v
r
          Ordering
GT -> [Char] -> acc
forall a. HasCallStack => [Char] -> a
error [Char]
"predicate not monotone with respect to ordering"
      Map k v
Map.Tip -> acc
acc

    goR :: acc -> Map k v -> acc
goR acc
acc = \case
      Map.Bin Size
_ k
k v
v Map k v
l Map k v
r ->
        case k -> Ordering
keyOrdering k
k of
          Ordering
EQ -> (k -> v -> acc -> acc) -> acc -> Map k v -> acc
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey k -> v -> acc -> acc
f (k -> v -> acc -> acc
f k
k v
v (acc -> Map k v -> acc
goR acc
acc Map k v
r)) Map k v
l
          Ordering
GT -> acc -> Map k v -> acc
goR acc
acc Map k v
l
          Ordering
LT -> [Char] -> acc
forall a. HasCallStack => [Char] -> a
error [Char]
"predicate not monotone with respect to ordering"
      Map k v
Map.Tip -> acc
acc

-- | 'swap' throws away data if the input contains duplicate values
swap :: (Ord b) => Map a b -> Map b a
swap :: forall b a. Ord b => Map a b -> Map b a
swap =
  (Map b a -> a -> b -> Map b a) -> Map b a -> Map a b -> Map b a
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
Map.foldlWithKey' (\Map b a
z a
a b
b -> b -> a -> Map b a -> Map b a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert b
b a
a Map b a
z) Map b a
forall a. Monoid a => a
mempty

-- | Like 'Map.insert', but returns the old value as well.
insertLookup :: (Ord k) => k -> v -> Map k v -> (Maybe v, Map k v)
insertLookup :: forall k v. Ord k => k -> v -> Map k v -> (Maybe v, Map k v)
insertLookup k
k v
v =
  (Maybe v -> v) -> k -> Map k v -> (Maybe v, Map k v)
forall k v.
Ord k =>
(Maybe v -> v) -> k -> Map k v -> (Maybe v, Map k v)
upsertLookup (v -> Maybe v -> v
forall a b. a -> b -> a
const v
v) k
k

-- | Invert a map's keys and values. This probably only makes sense with injective maps, but otherwise, later key/value
-- pairs (ordered by the original map's keys) overwrite earlier ones.
invert :: (Ord v) => Map k v -> Map v k
invert :: forall b a. Ord b => Map a b -> Map b a
invert =
  (Map v k -> k -> v -> Map v k) -> Map v k -> Map k v -> Map v k
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
Map.foldlWithKey' (\Map v k
m k
k v
v -> v -> k -> Map v k -> Map v k
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert v
v k
k Map v k
m) Map v k
forall k a. Map k a
Map.empty

-- | Like 'Map.lookup', but asserts the key is in the map.
lookupJust :: (Ord k, Show k) => k -> Map k v -> v
lookupJust :: forall k v. (Ord k, Show k) => k -> Map k v -> v
lookupJust k
k =
  v -> k -> Map k v -> v
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault ([Char] -> v
forall a. HasCallStack => [Char] -> a
error ([Char] -> [Char] -> [Char]
reportBug [Char]
"E147567" ([Char]
"Missing key: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ k -> [Char]
forall a. Show a => a -> [Char]
show k
k))) k
k

-- | Upsert an element into a map.
upsert :: (Ord k) => (Maybe v -> v) -> k -> Map k v -> Map k v
upsert :: forall k v. Ord k => (Maybe v -> v) -> k -> Map k v -> Map k v
upsert Maybe v -> v
f =
  (Maybe v -> Maybe v) -> k -> Map k v -> Map k v
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (v -> Maybe v
forall a. a -> Maybe a
Just (v -> Maybe v) -> (Maybe v -> v) -> Maybe v -> Maybe v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe v -> v
f)

-- | Upsert an element into a map.
upsertF :: (Functor f, Ord k) => (Maybe v -> f v) -> k -> Map k v -> f (Map k v)
upsertF :: forall (f :: * -> *) k v.
(Functor f, Ord k) =>
(Maybe v -> f v) -> k -> Map k v -> f (Map k v)
upsertF Maybe v -> f v
f =
  (Maybe v -> f (Maybe v)) -> k -> Map k v -> f (Map k v)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF ((v -> Maybe v) -> f v -> f (Maybe v)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap v -> Maybe v
forall a. a -> Maybe a
Just (f v -> f (Maybe v)) -> (Maybe v -> f v) -> Maybe v -> f (Maybe v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe v -> f v
f)

-- | Like 'upsert', but returns the old value as well.
upsertLookup :: (Ord k) => (Maybe v -> v) -> k -> Map k v -> (Maybe v, Map k v)
upsertLookup :: forall k v.
Ord k =>
(Maybe v -> v) -> k -> Map k v -> (Maybe v, Map k v)
upsertLookup Maybe v -> v
f =
  (Maybe v -> (Maybe v, v)) -> k -> Map k v -> (Maybe v, Map k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
(Maybe v -> f v) -> k -> Map k v -> f (Map k v)
upsertF (\Maybe v
v -> (Maybe v
v, Maybe v -> v
f Maybe v
v))

valuesVector :: Map k v -> Vector v
valuesVector :: forall k v. Map k v -> Vector v
valuesVector =
  [v] -> Vector v
forall a. [a] -> Vector a
Vector.fromList ([v] -> Vector v) -> (Map k v -> [v]) -> Map k v -> Vector v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> [v]
forall k a. Map k a -> [a]
Map.elems

-- | Like 'Map.delete', but returns the value as well.
deleteLookup :: (Ord k) => k -> Map k v -> (Maybe v, Map k v)
deleteLookup :: forall k v. Ord k => k -> Map k v -> (Maybe v, Map k v)
deleteLookup =
  (Maybe v -> (Maybe v, Maybe v))
-> k -> Map k v -> (Maybe v, Map k v)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF (,Maybe v
forall a. Maybe a
Nothing)

-- | Like 'deleteLookup', but asserts the value is in the map prior to deletion.
deleteLookupJust :: (HasCallStack, Ord k) => k -> Map k v -> (v, Map k v)
deleteLookupJust :: forall k v. (HasCallStack, Ord k) => k -> Map k v -> (v, Map k v)
deleteLookupJust =
  (Maybe v -> (v, Maybe v)) -> k -> Map k v -> (v, Map k v)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF ((v, Maybe v) -> (v -> (v, Maybe v)) -> Maybe v -> (v, Maybe v)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> (v, Maybe v)
forall a. HasCallStack => [Char] -> a
error ([Char] -> [Char] -> [Char]
reportBug [Char]
"E525283" [Char]
"deleteLookupJust: element not found")) (,Maybe v
forall a. Maybe a
Nothing))

-- | Like 'Map.elems', but return the values as a set.
elemsSet :: (Ord v) => Map k v -> Set v
elemsSet :: forall v k. Ord v => Map k v -> Set v
elemsSet =
  [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList ([v] -> Set v) -> (Map k v -> [v]) -> Map k v -> Set v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> [v]
forall k a. Map k a -> [a]
Map.elems

-- | Fold the keys of a map strictly with a "commutative" combining function that doesn't receive the elements in any
-- particular order.
foldKeysCommutative :: (k -> acc -> acc) -> acc -> Map k v -> acc
foldKeysCommutative :: forall k acc v. (k -> acc -> acc) -> acc -> Map k v -> acc
foldKeysCommutative k -> acc -> acc
f =
  let go :: acc -> [Map k v] -> acc
go !acc
acc = \case
        Map.Bin Size
_ k
k v
_ Map k v
l Map k v
r : [Map k v]
xs -> acc -> [Map k v] -> acc
go (k -> acc -> acc
f k
k acc
acc) (Map k v
l Map k v -> [Map k v] -> [Map k v]
forall a. a -> [a] -> [a]
: Map k v
r Map k v -> [Map k v] -> [Map k v]
forall a. a -> [a] -> [a]
: [Map k v]
xs)
        Map k v
Map.Tip : [Map k v]
xs -> acc -> [Map k v] -> acc
go acc
acc [Map k v]
xs
        [] -> acc
acc
   in \acc
z Map k v
xs -> acc -> [Map k v] -> acc
go acc
z [Map k v
xs]

-- | Fold the values of a map strictly with a "commutative" combining function that doesn't receive the elements in any
-- particular order.
foldValuesCommutative :: (v -> acc -> acc) -> acc -> Map k v -> acc
foldValuesCommutative :: forall v acc k. (v -> acc -> acc) -> acc -> Map k v -> acc
foldValuesCommutative v -> acc -> acc
f =
  let go :: acc -> [Map k v] -> acc
go !acc
acc = \case
        Map.Bin Size
_ k
_ v
v Map k v
l Map k v
r : [Map k v]
xs -> acc -> [Map k v] -> acc
go (v -> acc -> acc
f v
v acc
acc) (Map k v
l Map k v -> [Map k v] -> [Map k v]
forall a. a -> [a] -> [a]
: Map k v
r Map k v -> [Map k v] -> [Map k v]
forall a. a -> [a] -> [a]
: [Map k v]
xs)
        Map k v
Map.Tip : [Map k v]
xs -> acc -> [Map k v] -> acc
go acc
acc [Map k v]
xs
        [] -> acc
acc
   in \acc
z Map k v
xs -> acc -> [Map k v] -> acc
go acc
z [Map k v
xs]

-- | Like 'Map.foldlWithKey'', but with a monadic accumulator.
foldM :: (Monad m) => (acc -> k -> v -> m acc) -> acc -> Map k v -> m acc
foldM :: forall (m :: * -> *) acc k v.
Monad m =>
(acc -> k -> v -> m acc) -> acc -> Map k v -> m acc
foldM acc -> k -> v -> m acc
f acc
acc0 =
  acc -> Map k v -> m acc
go acc
acc0
  where
    go :: acc -> Map k v -> m acc
go !acc
acc = \case
      Map k v
Map.Tip -> acc -> m acc
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure acc
acc
      Map.Bin Size
_ k
k v
v Map k v
xs Map k v
ys -> do
        acc
acc1 <- acc -> Map k v -> m acc
go acc
acc Map k v
xs
        acc
acc2 <- acc -> k -> v -> m acc
f acc
acc1 k
k v
v
        acc -> Map k v -> m acc
go acc
acc2 Map k v
ys

-- | Construct a map from a foldable container by mapping each element to monadic action that returns a key and a value.
--
-- The map is constructed from the left: if two elements map to the same key, the second will overwrite the first.
foldMapM :: (Ord k, Monad m, Foldable t) => (a -> m (k, v)) -> t a -> m (Map k v)
foldMapM :: forall k (m :: * -> *) (t :: * -> *) a v.
(Ord k, Monad m, Foldable t) =>
(a -> m (k, v)) -> t a -> m (Map k v)
foldMapM a -> m (k, v)
f =
  (Map k v -> a -> m (Map k v)) -> Map k v -> t a -> m (Map k v)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Map k v -> a -> m (Map k v)
g Map k v
forall k a. Map k a
Map.empty
  where
    g :: Map k v -> a -> m (Map k v)
g Map k v
acc a
x = do
      (k
k, v
v) <- a -> m (k, v)
f a
x
      Map k v -> m (Map k v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map k v -> m (Map k v)) -> Map k v -> m (Map k v)
forall a b. (a -> b) -> a -> b
$! 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
acc

-- | Run a monadic action for each key/value pair in a map.
for_ :: (Monad m) => Map k v -> (k -> v -> m ()) -> m ()
for_ :: forall (m :: * -> *) k v.
Monad m =>
Map k v -> (k -> v -> m ()) -> m ()
for_ Map k v
m k -> v -> m ()
f =
  Map k v -> m ()
go Map k v
m
  where
    go :: Map k v -> m ()
go = \case
      Map k v
Map.Tip -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Map.Bin Size
_ k
k v
v Map k v
xs Map k v
ys -> do
        Map k v -> m ()
go Map k v
xs
        k -> v -> m ()
f k
k v
v
        Map k v -> m ()
go Map k v
ys

-- | Like 'Map.fromSet', but in an applicative functor.
fromSetA :: (Applicative m) => (k -> m a) -> Set k -> m (Map k a)
fromSetA :: forall (m :: * -> *) k a.
Applicative m =>
(k -> m a) -> Set k -> m (Map k a)
fromSetA k -> m a
f =
  Set k -> m (Map k a)
go
  where
    go :: Set k -> m (Map k a)
go = \case
      Set k
Set.Tip -> Map k a -> m (Map k a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k a
forall k a. Map k a
Map.Tip
      Set.Bin Size
n k
k Set k
l Set k
r -> (\a
v Map k a
l' Map k a
r' -> Size -> k -> a -> Map k a -> Map k a -> Map k a
forall k a. Size -> k -> a -> Map k a -> Map k a -> Map k a
Map.Bin Size
n k
k a
v Map k a
l' Map k a
r') (a -> Map k a -> Map k a -> Map k a)
-> m a -> m (Map k a -> Map k a -> Map k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> m a
f k
k m (Map k a -> Map k a -> Map k a)
-> m (Map k a) -> m (Map k a -> Map k a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set k -> m (Map k a)
go Set k
l m (Map k a -> Map k a) -> m (Map k a) -> m (Map k a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set k -> m (Map k a)
go Set k
r

unionWithM ::
  forall m k a.
  (Monad m, Ord k) =>
  (a -> a -> m a) ->
  Map k a ->
  Map k a ->
  m (Map k a)
unionWithM :: forall (m :: * -> *) k a.
(Monad m, Ord k) =>
(a -> a -> m a) -> Map k a -> Map k a -> m (Map k a)
unionWithM a -> a -> m a
f Map k a
m1 Map k a
m2 =
  (Map k a -> (k, a) -> m (Map k a))
-> Map k a -> [(k, a)] -> m (Map k a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
Monad.foldM Map k a -> (k, a) -> m (Map k a)
go Map k a
m1 ([(k, a)] -> m (Map k a)) -> [(k, a)] -> m (Map k a)
forall a b. (a -> b) -> a -> b
$ Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k a
m2
  where
    go :: Map k a -> (k, a) -> m (Map k a)
    go :: Map k a -> (k, a) -> m (Map k a)
go Map k a
m1 (k
k, a
a2) = case k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k a
m1 of
      Just a
a1 -> do a
a <- a -> a -> m a
f a
a1 a
a2; pure $ 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
a Map k a
m1
      Maybe a
Nothing -> Map k a -> m (Map k a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map k a -> m (Map k a)) -> Map k a -> m (Map k a)
forall a b. (a -> b) -> a -> b
$ 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
a2 Map k a
m1

-- | Reconstruct a map entirely, given a function from old key/value to new key/value.
--
-- @
-- remap f = Map.fromList . map f . Map.toList
-- @
remap :: (Ord k1) => ((k0, v0) -> (k1, v1)) -> Map k0 v0 -> Map k1 v1
remap :: forall k1 k0 v0 v1.
Ord k1 =>
((k0, v0) -> (k1, v1)) -> Map k0 v0 -> Map k1 v1
remap (k0, v0) -> (k1, v1)
f =
  [(k1, v1)] -> Map k1 v1
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(k1, v1)] -> Map k1 v1)
-> (Map k0 v0 -> [(k1, v1)]) -> Map k0 v0 -> Map k1 v1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k0, v0) -> (k1, v1)) -> [(k0, v0)] -> [(k1, v1)]
forall a b. (a -> b) -> [a] -> [b]
map (k0, v0) -> (k1, v1)
f ([(k0, v0)] -> [(k1, v1)])
-> (Map k0 v0 -> [(k0, v0)]) -> Map k0 v0 -> [(k1, v1)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k0 v0 -> [(k0, v0)]
forall k a. Map k a -> [(k, a)]
Map.toList

-- | Insert a pair in postfix-style.
thenInsertPair :: (Ord k) => Map k v -> (k, v) -> Map k v
thenInsertPair :: forall k v. Ord k => Map k v -> (k, v) -> Map k v
thenInsertPair 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

traverseKeys :: (Applicative f, Ord k') => (k -> f k') -> Map k v -> f (Map k' v)
traverseKeys :: forall (f :: * -> *) k' k v.
(Applicative f, Ord k') =>
(k -> f k') -> Map k v -> f (Map k' v)
traverseKeys k -> f k'
f = (k -> f k') -> (v -> f v) -> Map k v -> f (Map k' v)
forall (f :: * -> *) a' a b b'.
(Applicative f, Ord a') =>
(a -> f a') -> (b -> f b') -> Map a b -> f (Map a' b')
bitraverse k -> f k'
f v -> f v
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

traverseKeysWith :: (Applicative f, Ord k') => (v -> v -> v) -> (k -> f k') -> Map k v -> f (Map k' v)
traverseKeysWith :: forall (f :: * -> *) k' v k.
(Applicative f, Ord k') =>
(v -> v -> v) -> (k -> f k') -> Map k v -> f (Map k' v)
traverseKeysWith v -> v -> v
combine k -> f k'
f Map k v
m =
  (v -> v -> v) -> [(k', v)] -> Map k' v
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith v -> v -> v
combine ([(k', v)] -> Map k' v) -> f [(k', v)] -> f (Map k' v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k v
m [(k, v)] -> ([(k, v)] -> f [(k', v)]) -> f [(k', v)]
forall a b. a -> (a -> b) -> b
& ((k, v) -> f (k', v)) -> [(k, v)] -> f [(k', v)]
forall (f :: * -> *) a b.
Traversable f =>
IndexedTraversal Size (f a) (f b) a b
IndexedTraversal Size [(k, v)] [(k', v)] (k, v) (k', v)
traversed (((k, v) -> f (k', v)) -> [(k, v)] -> f [(k', v)])
-> ((k -> f k') -> (k, v) -> f (k', v))
-> (k -> f k')
-> [(k, v)]
-> f [(k', v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k -> f k') -> (k, v) -> f (k', v)
forall s t a b. Field1 s t a b => Lens s t a b
Lens (k, v) (k', v) k k'
_1 ((k -> f k') -> [(k, v)] -> f [(k', v)])
-> (k -> f k') -> [(k, v)] -> f [(k', v)]
forall {k} (f :: k -> *) s (t :: k) a (b :: k).
LensLike f s t a b -> LensLike f s t a b
%%~ k -> f k'
f)

-- | @mergeMap@ is like a @foldMap@ version of @merge@: summarize the merging of two maps together as a monoidal value.
mergeMap ::
  forall a b k m.
  (Monoid m, Ord k) =>
  -- | Function to apply when a key exists in the first map, but not the second.
  (k -> a -> m) ->
  -- | Function to apply when a key exists in the second map, but not the first.
  (k -> b -> m) ->
  -- | Function to apply when a key exists in both maps.
  (k -> a -> b -> m) ->
  Map k a ->
  Map k b ->
  m
mergeMap :: forall a b k m.
(Monoid m, Ord k) =>
(k -> a -> m)
-> (k -> b -> m) -> (k -> a -> b -> m) -> Map k a -> Map k b -> m
mergeMap k -> a -> m
f k -> b -> m
g k -> a -> b -> m
h =
  forall a b. Coercible a b => a -> b
forall a b. Coercible a b => a -> b
coerce @(Map k a -> Map k b -> Const m (Map k ())) do
    WhenMissing (Const m) k a ()
-> WhenMissing (Const m) k b ()
-> WhenMatched (Const m) k a b ()
-> Map k a
-> Map k b
-> Const m (Map k ())
forall (f :: * -> *) k a c b.
(Applicative f, Ord k) =>
WhenMissing f k a c
-> WhenMissing f k b c
-> WhenMatched f k a b c
-> Map k a
-> Map k b
-> f (Map k c)
Map.mergeA
      ((k -> a -> Const m ()) -> WhenMissing (Const m) k a ()
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> f y) -> WhenMissing f k x y
Map.traverseMissing ((k -> a -> m) -> k -> a -> Const m ()
forall a b. Coercible a b => a -> b
coerce k -> a -> m
f))
      ((k -> b -> Const m ()) -> WhenMissing (Const m) k b ()
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> f y) -> WhenMissing f k x y
Map.traverseMissing ((k -> b -> m) -> k -> b -> Const m ()
forall a b. Coercible a b => a -> b
coerce k -> b -> m
g))
      ((k -> a -> b -> Const m ()) -> WhenMatched (Const m) k a b ()
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> f z) -> WhenMatched f k x y z
Map.zipWithAMatched ((k -> a -> b -> m) -> k -> a -> b -> Const m ()
forall a b. Coercible a b => a -> b
coerce k -> a -> b -> m
h))