-- | This module defines a way to prove that a key exists in a map, so that the
-- key can be used to index into the map without using a 'Maybe', or manually
-- handling the \"impossible\" case with 'error' or other partial functions.
--
-- To do this, @'HashMap' s k v@ has a type parameter @s@ that identifies its
-- set of keys, so that if another map has the same type parameter, you know
-- that map has the same set of keys. There is @'Key' s k@, a type of keys that
-- have been validated to belong to the set identified by @s@, and for which the
-- operation of indexing into a @'HashMap' s k v@ (only for the same @s@) can
-- proceed without failure (see '!'). The type @s@ itself has no internal
-- structure, rather it is merely a skolem type variable (rank-2 polymorphism
-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
--
-- Like "Data.HashMap.Lazy", functions in this module are strict in the keys but
-- lazy in the values. The "Data.HashMap.Strict.Refined" module reuses the same
-- 'HashMap' type but provides functions that operate strictly on the values.
--
-- = Warning
-- This module together with "Data.HashMap.Lazy" rely on 'Eq' and 'Hashable'
-- instances being lawful: that '==' is an equivalence relation, and that
-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence
-- relation; at least for the subset of values that are actually encountered at
-- runtime. If this assumption is violated, this module may not be able to
-- uphold its invariants and may throw errors. In particular beware of NaN in
-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@.
module Data.HashMap.Refined
  (
  -- * Map type
    Common.HashMap
  , Common.Key
  -- * Existentials and common proofs
  , Common.SomeHashMap(..)
  , Common.withHashMap
  , Common.SomeHashMapWith(..)
  , Common.withHashMapWith
  , Common.Some2HashMapWith(..)
  , Common.with2HashMapWith
  , SupersetProof(..)
  , EmptyProof(..)
  -- * Construction
  , Common.empty
  , singleton
  , SingletonProof(..)
  , fromSet
  , Common.fromHashMap
  , fromTraversableWithKey
  , FromTraversableProof(..)
  -- * Insertion
  , insert
  , InsertProof(..)
  , reinsert
  , insertLookupWithKey
  -- * Deletion/Update
  , Common.delete
  , adjust
  , adjustWithKey
  , update
  , updateLookupWithKey
  -- * Query
  , Common.lookup
  , (Common.!)
  , Common.member
  , Common.null
  , Common.isSubmapOfBy
  , SubsetProof(..)
  , Common.disjoint
  , DisjointProof(..)
  -- * Combine
  , zipWithKey
  , bind
  , unionWithKey
  , UnionProof(..)
  , Common.difference
  , DifferenceProof(..)
  , differenceWithKey
  , PartialDifferenceProof(..)
  , intersectionWithKey
  , IntersectionProof(..)
  -- * Traversal
  , mapWithKey
  , traverseWithKey
  , mapAccumLWithKey
  , mapAccumRWithKey
  , mapKeysWith
  , MapProof(..)
  , backpermuteKeys
  -- * Folds
  , Common.foldMapWithKey
  , Common.foldrWithKey
  , Common.foldlWithKey
  , Common.foldrWithKey'
  , Common.foldlWithKey'
  -- * Conversion
  , Common.toMap
  , Common.keysSet
  , Common.toList
  -- * Filter
  , Common.restrictKeys
  , Common.withoutKeys
  , Common.filterWithKey
  , Common.partitionWithKey
  , PartitionProof(..)
  , mapMaybeWithKey
  , mapEitherWithKey
  -- * Casts
  , Common.castKey
  , Common.cast
  , castFlavor
  ) where

import           Data.Coerce
import           Data.Container.Refined.Hashable
import           Data.Container.Refined.Proofs
import           Data.Container.Refined.Unsafe
import           Data.Functor
import           Data.HashMap.Common.Refined
  ( HashMap(..), Key, unsafeCastKey, unsafeKey, SomeHashMapWith(..)
  , Some2HashMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey
  , bind
  )
import qualified Data.HashMap.Common.Refined as Common
import qualified Data.HashMap.Lazy as HashMap
import           Data.Traversable
import           Data.Traversable.WithIndex
import           Data.Type.Coercion
import           Prelude hiding (lookup, null)
import           Refined
import           Refined.Unsafe


-- | Create a map with a single key-value pair, and return a proof that the key
-- is in the resulting map.
singleton
  :: forall k a. Hashable k
  => k -> a -> SomeHashMapWith (SingletonProof 'Hashed k) k a
singleton :: forall k a.
Hashable k =>
k -> a -> SomeHashMapWith (SingletonProof 'Hashed k) k a
singleton k
k a
v = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton k
k a
v)
  forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) a r.
Refined (InSet f r) a -> SingletonProof f a r
SingletonProof (forall k s. k -> Key s k
unsafeKey k
k)

-- | Create a map from an arbitrary traversable of key-value pairs.
fromTraversableWithKey
  :: forall t k a. (Traversable t, Hashable k)
  => (k -> a -> a -> a)
  -> t (k, a)
  -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
fromTraversableWithKey :: forall (t :: * -> *) k a.
(Traversable t, Hashable k) =>
(k -> a -> a -> a)
-> t (k, a)
-> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
fromTraversableWithKey k -> a -> a -> a
f t (k, a)
xs
  = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap HashMap k a
m) forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) (t :: * -> *) a r.
t (Refined (InSet f r) a) -> FromTraversableProof f t a r
FromTraversableProof t (Key Any k)
proof
  where
    (HashMap k a
m, t (Key Any k)
proof) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL
      (\HashMap k a
s (k
k, a
v)
        -> let !s' :: HashMap k a
s' = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith (k -> a -> a -> a
f k
k) k
k a
v HashMap k a
s in (HashMap k a
s', forall k s. k -> Key s k
unsafeKey k
k))
      forall k v. HashMap k v
HashMap.empty
      t (k, a)
xs

-- | Insert a key-value pair into the map to obtain a potentially larger map,
-- guaranteed to contain the given key. If the key was already present, the
-- associated value is replaced with the supplied value.
insert
  :: forall s k a. Hashable k
  => k -> a -> HashMap s k a -> SomeHashMapWith (InsertProof 'Hashed k s) k a
insert :: forall s k a.
Hashable k =>
k
-> a
-> HashMap s k a
-> SomeHashMapWith (InsertProof 'Hashed k s) k a
insert k
k a
v (HashMap HashMap k a
m) = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert k
k a
v HashMap k a
m)
  forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) a s r.
Refined (InSet f r) a
-> (InSet f s :-> InSet f r) -> InsertProof f a s r
InsertProof (forall k s. k -> Key s k
unsafeKey k
k) forall p q. p :-> q
unsafeSubset

-- | Overwrite a key-value pair that is known to already be in the map. The set
-- of keys remains the same.
reinsert
  :: forall s k a. Hashable k
  => Key s k -> a -> HashMap s k a -> HashMap s k a
reinsert :: forall s k a.
Hashable k =>
Key s k -> a -> HashMap s k a -> HashMap s k a
reinsert = forall {k} (a :: k) (b :: k) r.
Coercion a b -> (Coercible a b => r) -> r
gcoerceWith (forall s k. Coercion k (Key s k)
unsafeCastKey @s @k) forall a b. (a -> b) -> a -> b
$ coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert @k @a

-- | Insert a key-value pair into the map using a combining function, and if
-- the key was already present, the old value is returned along with the proof
-- that the key was present.
insertLookupWithKey
  :: forall s k a. Hashable k
  => (Key s k -> a -> a -> a)
  -> k
  -> a
  -> HashMap s k a
  -> (Maybe (Key s k, a), SomeHashMapWith (InsertProof 'Hashed k s) k a)
insertLookupWithKey :: forall s k a.
Hashable k =>
(Key s k -> a -> a -> a)
-> k
-> a
-> HashMap s k a
-> (Maybe (Key s k, a),
    SomeHashMapWith (InsertProof 'Hashed k s) k a)
insertLookupWithKey Key s k -> a -> a -> a
f k
k a
v (HashMap HashMap k a
m) =
  ( (forall k s. k -> Key s k
unsafeKey k
k,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k HashMap k a
m
  , forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith (Key s k -> a -> a -> a
f forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey k
k) k
k a
v HashMap k a
m)
    forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) a s r.
Refined (InSet f r) a
-> (InSet f s :-> InSet f r) -> InsertProof f a s r
InsertProof (forall k s. k -> Key s k
unsafeKey k
k) forall p q. p :-> q
unsafeSubset
  )

-- | Update the value at a specific key known the be in the map using the given
-- function. The set of keys remains the same.
adjust
  :: forall s k a. Hashable k
  => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a
adjust :: forall s k a.
Hashable k =>
(a -> a) -> Key s k -> HashMap s k a -> HashMap s k a
adjust = forall {k} (a :: k) (b :: k) r.
Coercion a b -> (Coercible a b => r) -> r
gcoerceWith (forall s k. Coercion k (Key s k)
unsafeCastKey @s @k) forall a b. (a -> b) -> a -> b
$ coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HashMap.adjust @k @a

-- | If the given key is in the map, update the associated value using the given
-- function with a proof that the key was in the map; otherwise return the map
-- unchanged. In any case the set of keys remains the same.
adjustWithKey
  :: forall s k a. Hashable k
  => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a
adjustWithKey :: forall s k a.
Hashable k =>
(Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a
adjustWithKey Key s k -> a -> a
f k
k (HashMap HashMap k a
m) = forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HashMap.adjust (Key s k -> a -> a
f forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey k
k) k
k HashMap k a
m

-- | Update or delete a key known to be in the map using the given function,
-- returning a potentially smaller map.
update
  :: forall s k a. Hashable k
  => (a -> Maybe a)
  -> Key s k
  -> HashMap s k a
  -> SomeHashMapWith (SupersetProof 'Hashed s) k a
update :: forall s k a.
Hashable k =>
(a -> Maybe a)
-> Key s k
-> HashMap s k a
-> SomeHashMapWith (SupersetProof 'Hashed s) k a
update a -> Maybe a
f Key s k
k (HashMap HashMap k a
m)
  = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k a.
(Eq k, Hashable k) =>
(a -> Maybe a) -> k -> HashMap k a -> HashMap k a
HashMap.update a -> Maybe a
f (forall {k} (p :: k) x. Refined p x -> x
unrefine Key s k
k) HashMap k a
m)
    forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s r.
(InSet f r :-> InSet f s) -> SupersetProof f s r
SupersetProof forall p q. p :-> q
unsafeSubset

-- | If the given key is in the map, update or delete it using the given
-- function with a proof that the key was in the map; otherwise the map is
-- unchanged. Alongside return the new value if it was updated, or the old value
-- if it was deleted, and a proof that the key was in the map.
updateLookupWithKey
  :: forall s k a. Hashable k
  => (Key s k -> a -> Maybe a)
  -> k
  -> HashMap s k a
  -> (Maybe (Key s k, a), SomeHashMapWith (SupersetProof 'Hashed s) k a)
updateLookupWithKey :: forall s k a.
Hashable k =>
(Key s k -> a -> Maybe a)
-> k
-> HashMap s k a
-> (Maybe (Key s k, a),
    SomeHashMapWith (SupersetProof 'Hashed s) k a)
updateLookupWithKey Key s k -> a -> Maybe a
f k
k (HashMap HashMap k a
m) =
  ( (forall k s. k -> Key s k
unsafeKey k
k,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k HashMap k a
m
  , forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k a.
(Eq k, Hashable k) =>
(a -> Maybe a) -> k -> HashMap k a -> HashMap k a
HashMap.update (Key s k -> a -> Maybe a
f forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey k
k) k
k HashMap k a
m)
    forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s r.
(InSet f r :-> InSet f s) -> SupersetProof f s r
SupersetProof forall p q. p :-> q
unsafeSubset
  )

-- | Return the union of two maps, with a given combining function for keys that
-- exist in both maps simultaneously.
--
-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
-- respectively.
unionWithKey
  :: forall s t k a. Hashable k
  => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a)
  -> HashMap s k a
  -> HashMap t k a
  -> SomeHashMapWith (UnionProof 'Hashed s t) k a
unionWithKey :: forall s t k a.
Hashable k =>
(Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a)
-> HashMap s k a
-> HashMap t k a
-> SomeHashMapWith (UnionProof 'Hashed s t) k a
unionWithKey Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a
f (HashMap HashMap k a
m1) (HashMap HashMap k a
m2) = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith
  (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(k -> v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
HashMap.unionWithKey (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} x (p :: k). x -> Refined p x
reallyUnsafeRefine) HashMap k a
m1 HashMap k a
m2)
  forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s t r.
((InSet f s || InSet f t) :-> InSet f r)
-> (forall u.
    (InSet f s :-> InSet f u)
    -> (InSet f t :-> InSet f u) -> InSet f r :-> InSet f u)
-> UnionProof f s t r
UnionProof forall p q. p :-> q
unsafeSubset forall p' q' p'' q'' p q. (p' :-> q') -> (p'' :-> q'') -> p :-> q
unsafeSubsetWith2

-- | For keys that appear in both maps, the given function decides whether the
-- key is removed from the first map.
--
-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
-- respectively.
differenceWithKey
  :: forall s t k a b. Hashable k
  => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a)
  -> HashMap s k a
  -> HashMap t k b
  -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a
differenceWithKey :: forall s t k a b.
Hashable k =>
(Refined (InSet 'Hashed s && InSet 'Hashed t) k
 -> a -> b -> Maybe a)
-> HashMap s k a
-> HashMap t k b
-> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a
differenceWithKey Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a
f (HashMap HashMap k a
m1) (HashMap HashMap k b
m2) = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith
  (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v w.
(Eq k, Hashable k) =>
(v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v
HashMap.differenceWith
    (\a
x (k
k, b
y) -> Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a
f (forall {k} x (p :: k). x -> Refined p x
reallyUnsafeRefine k
k) a
x b
y)
    HashMap k a
m1
    (forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapWithKey (,) HashMap k b
m2))
  forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s t r.
(InSet f r :-> InSet f s)
-> (InSet f s :-> (InSet f t || InSet f r))
-> PartialDifferenceProof f s t r
PartialDifferenceProof forall p q. p :-> q
unsafeSubset forall p q. p :-> q
unsafeSubset

-- | Return the intersection of two maps with the given combining function.
--
-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
-- respectively.
intersectionWithKey
  :: forall s t k a b c. Hashable k
  => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c)
  -> HashMap s k a
  -> HashMap t k b
  -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c
intersectionWithKey :: forall s t k a b c.
Hashable k =>
(Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c)
-> HashMap s k a
-> HashMap t k b
-> SomeHashMapWith (IntersectionProof 'Hashed s t) k c
intersectionWithKey Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c
f (HashMap HashMap k a
m1) (HashMap HashMap k b
m2) = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith
  (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v1 v2 v3.
(Eq k, Hashable k) =>
(k -> v1 -> v2 -> v3)
-> HashMap k v1 -> HashMap k v2 -> HashMap k v3
HashMap.intersectionWithKey (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} x (p :: k). x -> Refined p x
reallyUnsafeRefine) HashMap k a
m1 HashMap k b
m2)
  forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s t r.
(InSet f r :-> (InSet f s && InSet f t))
-> (forall u.
    (InSet f u :-> InSet f s)
    -> (InSet f u :-> InSet f t) -> InSet f u :-> InSet f r)
-> IntersectionProof f s t r
IntersectionProof forall p q. p :-> q
unsafeSubset forall p' q' p'' q'' p q. (p' :-> q') -> (p'' :-> q'') -> p :-> q
unsafeSubsetWith2

-- | Thread an accumularing argument through the map in ascending order of
-- hashes.
mapAccumLWithKey
  :: forall s k a b c. (a -> Key s k -> b -> (a, c))
  -> a
  -> HashMap s k b
  -> (a, HashMap s k c)
mapAccumLWithKey :: forall s k a b c.
(a -> Key s k -> b -> (a, c))
-> a -> HashMap s k b -> (a, HashMap s k c)
mapAccumLWithKey a -> Key s k -> b -> (a, c)
f = forall i (t :: * -> *) s a b.
TraversableWithIndex i t =>
(i -> s -> a -> (s, b)) -> s -> t a -> (s, t b)
imapAccumL (forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Key s k -> b -> (a, c)
f)

-- | Thread an accumularing argument through the map in descending order of
-- hashes.
mapAccumRWithKey
  :: forall s k a b c. (a -> Key s k -> b -> (a, c))
  -> a
  -> HashMap s k b
  -> (a, HashMap s k c)
mapAccumRWithKey :: forall s k a b c.
(a -> Key s k -> b -> (a, c))
-> a -> HashMap s k b -> (a, HashMap s k c)
mapAccumRWithKey a -> Key s k -> b -> (a, c)
f = forall i (t :: * -> *) s a b.
TraversableWithIndex i t =>
(i -> s -> a -> (s, b)) -> s -> t a -> (s, t b)
imapAccumR (forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Key s k -> b -> (a, c)
f)

-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
-- results into a new map. For keys that were mapped to the same new key, @c@
-- acts as the combining function for corresponding values.
mapKeysWith
  :: forall s k1 k2 a. Hashable k2
  => (a -> a -> a)
  -> (Key s k1 -> k2)
  -> HashMap s k1 a
  -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a
mapKeysWith :: forall s k1 k2 a.
Hashable k2 =>
(a -> a -> a)
-> (Key s k1 -> k2)
-> HashMap s k1 a
-> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a
mapKeysWith a -> a -> a
f Key s k1 -> k2
g (HashMap HashMap k1 a
m) = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith
  (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
HashMap.fromListWith a -> a -> a
f
    forall a b. (a -> b) -> a -> b
$ forall k v a. (k -> v -> a -> a) -> a -> HashMap k v -> a
HashMap.foldrWithKey (\k1
k a
x [(k2, a)]
xs -> (Key s k1 -> k2
g forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey k1
k, a
x) forall a. a -> [a] -> [a]
: [(k2, a)]
xs) [] HashMap k1 a
m)
  forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s a b r.
(Refined (InSet f s) a -> Refined (InSet f r) b)
-> (Refined (InSet f r) b -> Refined (InSet f s) a)
-> MapProof f s a b r
MapProof (forall k s. k -> Key s k
unsafeKey forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key s k1 -> k2
g) \Refined (InSet 'Hashed Any) k2
k2 ->
    case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup (forall {k} (p :: k) x. Refined p x -> x
unrefine Refined (InSet 'Hashed Any) k2
k2) HashMap k2 (Key s k1)
backMap of
      Maybe (Key s k1)
Nothing -> forall a. HasCallStack => [Char] -> a
error
        [Char]
"mapKeysWith: bug: Data.HashMap.Refined has been subverted"
      Just Key s k1
k1 -> Key s k1
k1
  where
    ~HashMap k2 (Key s k1)
backMap = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList
      [ (k2
k2, forall k s. k -> Key s k
unsafeKey k1
k1)
      | k1
k1 <- forall k v. HashMap k v -> [k]
HashMap.keys HashMap k1 a
m
      , let !k2 :: k2
k2 = Key s k1 -> k2
g forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey k1
k1
      ]

-- | Apply a function to all values in a map, together with their corresponding
-- keys, and collect only the 'Just' results, returning a potentially smaller
-- map.
mapMaybeWithKey
  :: forall s k a b. (Key s k -> a -> Maybe b)
  -> HashMap s k a
  -> SomeHashMapWith (SupersetProof 'Hashed s) k b
mapMaybeWithKey :: forall s k a b.
(Key s k -> a -> Maybe b)
-> HashMap s k a -> SomeHashMapWith (SupersetProof 'Hashed s) k b
mapMaybeWithKey Key s k -> a -> Maybe b
f (HashMap HashMap k a
m)
  = forall s k a (p :: * -> *).
HashMap s k a -> p s -> SomeHashMapWith p k a
SomeHashMapWith (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall k v1 v2.
(k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapMaybeWithKey (Key s k -> a -> Maybe b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k s. k -> Key s k
unsafeKey) HashMap k a
m)
    forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s r.
(InSet f r :-> InSet f s) -> SupersetProof f s r
SupersetProof forall p q. p :-> q
unsafeSubset

-- | Apply a function to all values in a map, together with their corresponding
-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
-- maps.
mapEitherWithKey
  :: forall s k a b c. Hashable k -- TODO: this is only used in the proof
  => (Key s k -> a -> Either b c)
  -> HashMap s k a
  -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c
mapEitherWithKey :: forall s k a b c.
Hashable k =>
(Key s k -> a -> Either b c)
-> HashMap s k a
-> Some2HashMapWith (PartitionProof 'Hashed s k) k b c
mapEitherWithKey Key s k -> a -> Either b c
p (HashMap HashMap k a
m)
  | HashMap k (Either b c)
m' <- forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapWithKey (Key s k -> a -> Either b c
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k s. k -> Key s k
unsafeKey) HashMap k a
m
  = forall s t k a b (p :: * -> * -> *).
HashMap s k a -> HashMap t k b -> p s t -> Some2HashMapWith p k a b
Some2HashMapWith
    (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall v1 v2 k. (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapMaybe (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. a -> Maybe a
Just (forall a b. a -> b -> a
const forall a. Maybe a
Nothing)) HashMap k (Either b c)
m')
    (forall s k a. HashMap k a -> HashMap s k a
HashMap forall a b. (a -> b) -> a -> b
$ forall v1 v2 k. (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapMaybe (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) forall a. a -> Maybe a
Just) HashMap k (Either b c)
m')
    forall a b. (a -> b) -> a -> b
$ forall (f :: Flavor) s a r q.
(Refined (InSet f s) a
 -> Either (Refined (InSet f r) a) (Refined (InSet f q) a))
-> ((InSet f r || InSet f q) :-> InSet f s)
-> (forall t.
    (InSet f r :-> InSet f t)
    -> (InSet f q :-> InSet f t) -> InSet f s :-> InSet f t)
-> (forall t.
    (InSet f t :-> InSet f r)
    -> (InSet f t :-> InSet f q) -> forall u. InSet f t :-> InSet f u)
-> PartitionProof f s a r q
PartitionProof
      do \Key s k
k -> case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup (forall {k} (p :: k) x. Refined p x -> x
unrefine Key s k
k) HashMap k a
m of
          Maybe a
Nothing -> forall a. HasCallStack => [Char] -> a
error
            [Char]
"mapEitherWithKey: bug: Data.HashMap.Refined has been subverted"
          Just a
x -> case Key s k -> a -> Either b c
p Key s k
k a
x of
            Left b
_ -> forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey forall a b. (a -> b) -> a -> b
$ forall {k} (p :: k) x. Refined p x -> x
unrefine Key s k
k
            Right c
_ -> forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ forall k s. k -> Key s k
unsafeKey forall a b. (a -> b) -> a -> b
$ forall {k} (p :: k) x. Refined p x -> x
unrefine Key s k
k
      forall p q. p :-> q
unsafeSubset forall p' q' p'' q'' p q. (p' :-> q') -> (p'' :-> q'') -> p :-> q
unsafeSubsetWith2 \InSet 'Hashed t :-> InSet 'Hashed Any
f InSet 'Hashed t :-> InSet 'Hashed Any
g -> forall p' q' p'' q'' p q. (p' :-> q') -> (p'' :-> q'') -> p :-> q
unsafeSubsetWith2 InSet 'Hashed t :-> InSet 'Hashed Any
f InSet 'Hashed t :-> InSet 'Hashed Any
g

-- | Apply the inverse image of the given function to the keys of the given map,
-- so that for all @k :: 'Key' s2 k2@,c
-- @'backpermuteKeys' f m '!' k = m '!' f k@.
--
-- If maps are identified with functions, this computes the composition.
backpermuteKeys
  :: forall s1 s2 k1 k2 a. (Hashable k1, KnownHashSet s2 k2)
  => (Key s2 k2 -> Key s1 k1) -> HashMap s1 k1 a -> HashMap s2 k2 a
backpermuteKeys :: forall s1 s2 k1 k2 a.
(Hashable k1, KnownHashSet s2 k2) =>
(Key s2 k2 -> Key s1 k1) -> HashMap s1 k1 a -> HashMap s2 k2 a
backpermuteKeys Key s2 k2 -> Key s1 k1
f HashMap s1 k1 a
m = forall s k a. KnownHashSet s k => (Key s k -> a) -> HashMap s k a
fromSet \Key s2 k2
k -> HashMap s1 k1 a
m forall s k a. Hashable k => HashMap s k a -> Key s k -> a
! Key s2 k2 -> Key s1 k1
f Key s2 k2
k