git.iko.soy/typeable/refined-containers
on main help sign in
14 files changed+1755−105
since splitsnapshots
unifiedsplit
.github/workflows/ci.yaml+7−6 refined-containers.cabal+3−2 src/Data/HashMap/Common/Refined.hs+108−16 src/Data/HashMap/Refined.hs+226−9 src/Data/HashMap/Strict/Refined.hs+227−10 src/Data/HashSet/Refined.hs+1−1 src/Data/IntMap/Common/Refined.hs+134−15 src/Data/IntMap/Refined.hs+225−7 src/Data/IntMap/Strict/Refined.hs+229−9 src/Data/IntSet/Refined.hs+1−1 src/Data/Map/Common/Refined.hs+141−15 src/Data/Map/Refined.hs+226−7 src/Data/Map/Strict/Refined.hs+226−6 src/Data/Set/Refined.hs+1−1
.github/workflows/ci.yaml+7−6
@@ -27,18 +27,19 @@ jobs:
2727 - "9.6.6"
2828 - "9.8.2"
2929 - "9.10.1"
30+ - "9.12.1"
3031 fail-fast: false
3132 name: Build on Linux, GHC ${{ matrix.ghc_version }}
3233 runs-on: ${{ matrix.platform }}
3334 steps:
34− - uses: actions/checkout@v3
35− - uses: haskell/actions/setup@v2
35+ - uses: actions/checkout@v4
36+ - uses: haskell-actions/setup@v2
3637 id: setup-haskell
3738 with:
3839 ghc-version: ${{ matrix.ghc_version }}
3940 cabal-version: "3.10.1.0"
4041 - name: Cache
41− uses: actions/cache@v3
42+ uses: actions/cache@v4
4243 with:
4344 path: ${{ steps.setup-haskell.outputs.cabal-store }}
4445 key: ${{ runner.os }}-stack-${{ matrix.ghc_version }}
@@ -57,15 +58,15 @@ jobs:
5758 name: Build source tarball
5859 if: github.event_name == 'release'
5960 steps:
60− - uses: actions/checkout@v3
61− - uses: haskell/actions/setup@v2
61+ - uses: actions/checkout@v4
62+ - uses: haskell-actions/setup@v2
6263 with:
6364 cabal-version: "3.10.1.0"
6465 - name: cabal sdist
6566 run: |
6667 cabal sdist --output-dir "$GITHUB_WORKSPACE/sdist"
6768 echo "dist=$(find "$GITHUB_WORKSPACE/sdist" -maxdepth 1 -type f -name '*.tar.gz')" >> "$GITHUB_ENV"
68− - uses: actions/upload-artifact@v2
69+ - uses: actions/upload-artifact@v4
6970 with:
7071 name: Source tarball
7172 path: ${{ env.dist }}
refined-containers.cabal+3−2
@@ -23,10 +23,11 @@ license: MIT
2323license-file: LICENSE
2424author: mniip@typeable.io
2525maintainer: mniip@typeable.io
26−version: 0.1.0.2
26+version: 0.1.2.0
2727build-type: Simple
2828
2929tested-with:
30+ , GHC == 9.12.1
3031 , GHC == 9.10.1
3132 , GHC == 9.8.2
3233 , GHC == 9.6.6
@@ -42,7 +43,7 @@ source-repository head
4243
4344library
4445 build-depends:
45− , base >= 4.12 && < 4.21
46+ , base >= 4.12 && < 4.22
4647 , adjunctions >= 4.4 && < 4.5
4748 , constraints >= 0.11 && < 0.15
4849 , containers >= 0.5.7 && < 0.8
src/Data/HashMap/Common/Refined.hs+108−16
@@ -4,6 +4,7 @@ module Data.HashMap.Common.Refined where
44
55import Control.Monad.Reader
66import Control.DeepSeq
7+import Data.Bifoldable
78import Data.Coerce
89import Data.Constraint (Dict(..))
910import Data.Container.Refined.Hashable
@@ -11,6 +12,7 @@ import Data.Container.Refined.Proofs
1112import Data.Container.Refined.Unsafe
1213import Data.Distributive
1314import Data.Foldable.WithIndex
15+import Data.Functor.Classes
1416import Data.Functor.Rep
1517import Data.Functor.WithIndex
1618import qualified Data.Hashable as Hashable
@@ -22,6 +24,7 @@ import Data.Reflection
2224import Data.Traversable.WithIndex
2325import Data.Type.Coercion
2426import Data.Type.Equality ((:~:)(..))
27+import Prelude hiding (zipWith)
2528import Refined
2629import Refined.Unsafe
2730import Unsafe.Coerce
@@ -44,7 +47,10 @@ import Data.Monoid (All(..))
4447-- always derive @'KnownHashSet' s k@ by pattern matching on the 'Dict' returned
4548-- by 'keysSet'.
4649newtype HashMap s k a = HashMap (HashMap.HashMap k a)
47− deriving newtype (Eq, Ord, Show, Functor, Foldable, Hashable.Hashable, NFData)
50+ deriving newtype
51+ ( Eq, Eq1, Eq2, Ord, Ord1, Ord2, Show, Show1, Show2
52+ , Functor, Foldable, Bifoldable, Hashable.Hashable, NFData
53+ )
4854 deriving stock (Traversable)
4955type role HashMap nominal nominal representational
5056
@@ -90,7 +96,7 @@ data SomeHashMap k a where
9096--
9197-- @
9298-- 'withHashMap' ('fromHashMap' ...
93−-- $ \(m :: 'HashMap' s k a) -> doSomethingWith \@s
99+-- $ \\(m :: 'HashMap' s k a) -> doSomethingWith \@s
94100-- @
95101withHashMap
96102 :: forall k a r. SomeHashMap k a -> (forall s. HashMap s k a -> r) -> r
@@ -100,6 +106,26 @@ withHashMap (SomeHashMap m) k = k m
100106fromHashMap :: forall k a. HashMap.HashMap k a -> SomeHashMap k a
101107fromHashMap m = SomeHashMap (HashMap m)
102108
109+-- | Given a set of keys @s@ known ahead of time, verify whether a regular
110+-- 'Data.HashMap.Lazy.HashMap' has exactly that set of keys.
111+verifyHashMap
112+ :: forall s k a. (Hashable k, KnownHashSet s k)
113+ => HashMap.HashMap k a -> Maybe (HashMap s k a)
114+verifyHashMap m
115+#if MIN_VERSION_unordered_containers(0, 2, 12)
116+ | HashMap.isSubmapOfBy (\_ _ -> True) m (HashSet.toMap keys)
117+ , HashMap.isSubmapOfBy (\_ _ -> True) (HashSet.toMap keys) m
118+#else
119+ | All True <- flip HashMap.foldMapWithKey m \k _
120+ -> All $ k `HashSet.member` keys
121+ , All True <- flip foldMap keys \k
122+ -> All $ k `HashMap.member` m
123+#endif
124+ = Just (HashMap m)
125+ | otherwise = Nothing
126+ where
127+ keys = reflect $ Proxy @s
128+
103129-- | An existential wrapper for a 'HashMap' with an as-yet-unknown set of keys,
104130-- together with a proof of some fact @p@ about the set. Pattern matching on it
105131-- gives you a way to refer to the set (the parameter @s@). Functions that
@@ -214,11 +240,25 @@ disjoint (HashMap m1) (HashMap m2)
214240
215241-- | Given two maps proven to have the same keys, for each key apply the
216242-- function to the associated values, to obtain a new map with the same keys.
217−zipWithKey
243+zipWith
218244 :: forall s k a b c. Hashable k
219− => (Key s k -> a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c
220−zipWithKey f (HashMap m1) (HashMap m2) = HashMap
221− $ HashMap.intersectionWithKey (f . unsafeKey) m1 m2
245+ => (a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c
246+zipWith f (HashMap m1) (HashMap m2) = HashMap $ HashMap.intersectionWith f m1 m2
247+
248+-- | Return the union of two maps. For keys that exist in both maps, the value
249+-- is taken from the first map.
250+--
251+-- @
252+-- 'union' = unionWith 'const'
253+-- @
254+union
255+ :: forall s t k a. Hashable k
256+ => HashMap s k a
257+ -> HashMap t k a
258+ -> SomeHashMapWith (UnionProof 'Hashed s t) k a
259+union (HashMap m1) (HashMap m2) = SomeHashMapWith
260+ (HashMap $ HashMap.union m1 m2)
261+ $ UnionProof unsafeSubset unsafeSubsetWith2
222262
223263-- | Remove the keys that appear in the second map from the first map.
224264difference
@@ -230,6 +270,20 @@ difference (HashMap m1) (HashMap m2)
230270 = SomeHashMapWith (HashMap $ HashMap.difference m1 m2)
231271 $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
232272
273+-- | Return the intersection of two maps, taking values from the first map.
274+--
275+-- @
276+-- 'intersection' = intersectionWith 'const'
277+-- @
278+intersection
279+ :: forall s t k a b. Hashable k
280+ => HashMap s k a
281+ -> HashMap t k b
282+ -> SomeHashMapWith (IntersectionProof 'Hashed s t) k a
283+intersection (HashMap m1) (HashMap m2) = SomeHashMapWith
284+ (HashMap $ HashMap.intersection m1 m2)
285+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
286+
233287-- | Apply a function to all values in a map, together with their corresponding
234288-- keys, that are proven to be in the map. The set of keys remains the same.
235289mapWithKey
@@ -288,6 +342,25 @@ keysSet (HashMap m) = reify (HashMap.keysSet m)
288342toList :: forall s k a. HashMap s k a -> [(Key s k, a)]
289343toList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.toList @k @a
290344
345+-- | Retain only the values that satisfy the predicate, returning a potentially
346+-- smaller map.
347+filter
348+ :: forall s k a. (a -> Bool)
349+ -> HashMap s k a
350+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
351+filter p (HashMap m) = SomeHashMapWith (HashMap $ HashMap.filter p m)
352+ $ SupersetProof unsafeSubset
353+
354+-- | Retain only the keys that satisfy the predicate, returning a potentially
355+-- smaller map.
356+filterKeys
357+ :: forall s k a. (Key s k -> Bool)
358+ -> HashMap s k a
359+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
360+filterKeys p (HashMap m) = SomeHashMapWith
361+ (HashMap $ HashMap.filterWithKey (\k _ -> p (unsafeKey k)) m)
362+ $ SupersetProof unsafeSubset
363+
291364-- | Retain only the key-value pairs that satisfy the predicate, returning a
292365-- potentially smaller map.
293366filterWithKey
@@ -315,6 +388,25 @@ withoutKeys (HashMap m) = SomeHashMapWith
315388 (HashMap $ HashMap.difference m $ HashSet.toMap $ reflect $ Proxy @t)
316389 $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
317390
391+-- | Partition a map into two disjoint submaps: those whose key-value pairs
392+-- satisfy the predicate, and those whose don't.
393+partition
394+ :: forall s k a. Hashable k -- TODO: this is only used in the proof
395+ => (a -> Bool)
396+ -> HashMap s k a
397+ -> Some2HashMapWith (PartitionProof 'Hashed s k) k a a
398+partition p (HashMap m) = Some2HashMapWith
399+ (HashMap $ HashMap.filter p m)
400+ (HashMap $ HashMap.filter (not . p) m)
401+ $ PartitionProof
402+ do \k -> case HashMap.lookup (unrefine k) m of
403+ Nothing -> error
404+ "partition: bug: Data.HashMap.Refined has been subverted"
405+ Just x -> if p x
406+ then Left $ unsafeKey $ unrefine k
407+ else Right $ unsafeKey $ unrefine k
408+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
409+
318410-- | Partition a map into two disjoint submaps: those whose key-value pairs
319411-- satisfy the predicate, and those whose don't.
320412partitionWithKey
@@ -349,15 +441,15 @@ castKey
349441castKey = castRefined
350442
351443-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps
352−-- can be interconverted too. For example, 'zipWithKey' can be implemented via
353−-- 'Data.HashMap.Refined.intersectionWithKey' by proving that the set of keys
354−-- remains unchanged:
444+-- can be interconverted too. For example, 'Data.HashMap.Refined.zipWithKey' can
445+-- be implemented via 'Data.HashMap.Refined.intersectionWithKey' by proving that
446+-- the set of keys remains unchanged:
355447--
356448-- @
357−-- 'zipWithKey'
449+-- 'Data.HashMap.Refined.zipWithKey'
358450-- :: forall s k a b c. 'Hashable' k
359451-- => ('Key' s k -> a -> b -> c) -> 'HashMap' s k a -> 'HashMap' s k b -> 'HashMap' s k c
360−-- 'zipWithKey' f m1 m2
452+-- 'Data.HashMap.Refined.zipWithKey' f m1 m2
361453-- | v'SomeHashMapWith' @r m proof <- 'Data.HashMap.Refined.intersectionWithKey' (f . 'andLeft') m1 m2
362454-- , v'IntersectionProof' p1 p2 <- proof
363455-- , ( v'Coercion' :: t'Coercion' ('HashMap' r k c) ('HashMap' s k c))
@@ -382,10 +474,10 @@ instance TraversableWithIndex (Key s k) (HashMap s k) where
382474 itraverse = traverseWithKey
383475
384476-- | Similar to the instance for functions -- zip corresponding keys. To use
385−-- '<*>'/'Control.Applicative.liftA2' without 'KnownSet' see 'zipWithKey'.
477+-- '<*>'/'Control.Applicative.liftA2' without 'KnownHashSet' see 'zipWith'.
386478instance (Hashable k, KnownHashSet s k) => Applicative (HashMap s k) where
387479 pure x = fromSet \_ -> x
388− (<*>) = zipWithKey (const id)
480+ (<*>) = zipWith id
389481
390482-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the
391483-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
@@ -394,8 +486,8 @@ bind
394486 => HashMap s k a -> (a -> HashMap s k b) -> HashMap s k b
395487bind m f = mapWithKey (\k x -> f x ! k) m
396488
397−-- | Similar to the instance for functions. To use '>>=' without 'KnownSet' see
398−-- 'bind'.
489+-- | Similar to the instance for functions. To use '>>=' without 'KnownHashSet'
490+-- see 'bind'.
399491instance (Hashable k, KnownHashSet s k) => Monad (HashMap s k) where
400492 (>>=) = bind
401493
@@ -408,7 +500,7 @@ instance (Hashable k, KnownHashSet s k)
408500
409501-- | Append the values at the corresponding keys
410502instance (Hashable k, Semigroup a) => Semigroup (HashMap s k a) where
411− (<>) = zipWithKey (const (<>))
503+ (<>) = zipWith (<>)
412504
413505instance (Hashable k, KnownHashSet s k, Monoid a)
414506 => Monoid (HashMap s k a) where
src/Data/HashMap/Refined.hs+226−9
@@ -43,18 +43,26 @@ module Data.HashMap.Refined
4343 , SingletonProof(..)
4444 , fromSet
4545 , Common.fromHashMap
46+ , Common.verifyHashMap
47+ , fromTraversable
48+ , fromTraversableWith
4649 , fromTraversableWithKey
4750 , FromTraversableProof(..)
4851 -- * Insertion
4952 , insert
53+ , insertWith
54+ , insertWithKey
5055 , InsertProof(..)
5156 , reinsert
5257 , insertLookupWithKey
5358 -- * Deletion/Update
5459 , Common.delete
60+ , adjust'
5561 , adjust
5662 , adjustWithKey
63+ , update'
5764 , update
65+ , updateWithKey
5866 , updateLookupWithKey
5967 -- * Query
6068 , Common.lookup
@@ -66,21 +74,29 @@ module Data.HashMap.Refined
6674 , Common.disjoint
6775 , DisjointProof(..)
6876 -- * Combine
77+ , zipWith
6978 , zipWithKey
7079 , bind
80+ , Common.union
81+ , unionWith
7182 , unionWithKey
7283 , UnionProof(..)
7384 , Common.difference
7485 , DifferenceProof(..)
86+ , differenceWith
7587 , differenceWithKey
7688 , PartialDifferenceProof(..)
89+ , Common.intersection
90+ , intersectionWith
7791 , intersectionWithKey
7892 , IntersectionProof(..)
7993 -- * Traversal
94+ , map
8095 , mapWithKey
8196 , traverseWithKey
8297 , mapAccumLWithKey
8398 , mapAccumRWithKey
99+ , mapKeys
84100 , mapKeysWith
85101 , MapProof(..)
86102 , backpermuteKeys
@@ -97,10 +113,15 @@ module Data.HashMap.Refined
97113 -- * Filter
98114 , Common.restrictKeys
99115 , Common.withoutKeys
116+ , Common.filter
117+ , Common.filterKeys
100118 , Common.filterWithKey
119+ , Common.partition
101120 , Common.partitionWithKey
102121 , PartitionProof(..)
122+ , mapMaybe
103123 , mapMaybeWithKey
124+ , mapEither
104125 , mapEitherWithKey
105126 -- * Casts
106127 , Common.castKey
@@ -115,7 +136,7 @@ import Data.Container.Refined.Unsafe
115136import Data.Functor
116137import Data.HashMap.Common.Refined
117138 ( HashMap(..), Key, unsafeCastKey, unsafeKey, SomeHashMapWith(..)
118− , Some2HashMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey
139+ , Some2HashMapWith(..), fromSet, (!), zipWith, mapWithKey, traverseWithKey
119140 , bind
120141 )
121142import qualified Data.HashMap.Common.Refined as Common
@@ -123,7 +144,7 @@ import qualified Data.HashMap.Lazy as HashMap
123144import Data.Traversable
124145import Data.Traversable.WithIndex
125146import Data.Type.Coercion
126−import Prelude hiding (lookup, null)
147+import Prelude hiding (lookup, map, null, zipWith)
127148import Refined
128149import Refined.Unsafe
129150
@@ -136,7 +157,42 @@ singleton
136157singleton k v = SomeHashMapWith (HashMap $ HashMap.singleton k v)
137158 $ SingletonProof (unsafeKey k)
138159
139−-- | Create a map from an arbitrary traversable of key-value pairs.
160+-- | Create a map from an arbitrary traversable of key-value pairs. If a key is
161+-- repeated, the retained value is the last one in traversal order. If you're
162+-- looking for @fromList@, this is the function you want.
163+fromTraversable
164+ :: forall t k a. (Traversable t, Hashable k)
165+ => t (k, a) -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
166+fromTraversable xs = SomeHashMapWith (HashMap m) $ FromTraversableProof proof
167+ where
168+ (m, proof) = mapAccumL
169+ (\s (k, v) -> let !s' = HashMap.insert k v s in (s', unsafeKey k))
170+ HashMap.empty
171+ xs
172+
173+-- | Create a map from an arbitrary traversable of key-value pairs, with a
174+-- function for combining values for repeated keys. The function is called as if
175+-- by 'foldl1', but flipped:
176+--
177+-- @
178+-- 'fromTraversableWith' f [(k, x1), (k, x2), (k, x3)]
179+-- = 'singleton' k (f x3 (f x2 x1))
180+-- @
181+fromTraversableWith
182+ :: forall t k a. (Traversable t, Hashable k)
183+ => (a -> a -> a)
184+ -> t (k, a)
185+ -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
186+fromTraversableWith f xs
187+ = SomeHashMapWith (HashMap m) $ FromTraversableProof proof
188+ where
189+ (m, proof) = mapAccumL
190+ (\s (k, v) -> let !s' = HashMap.insertWith f k v s in (s', unsafeKey k))
191+ HashMap.empty
192+ xs
193+
194+-- | Create a map from an arbitrary traversable of key-value pairs. Like
195+-- 'fromTraversableWith', but the combining function has access to the key.
140196fromTraversableWithKey
141197 :: forall t k a. (Traversable t, Hashable k)
142198 => (k -> a -> a -> a)
@@ -160,6 +216,35 @@ insert
160216insert k v (HashMap m) = SomeHashMapWith (HashMap $ HashMap.insert k v m)
161217 $ InsertProof (unsafeKey k) unsafeSubset
162218
219+-- | Insert a key-value pair into the map to obtain a potentially larger map,
220+-- guaranteed to contain the given key. If the key was already present, the
221+-- supplied function is used to combine the new value with the old (in that
222+-- order).
223+insertWith
224+ :: forall s k a. Hashable k
225+ => (a -> a -> a)
226+ -> k
227+ -> a
228+ -> HashMap s k a
229+ -> SomeHashMapWith (InsertProof 'Hashed k s) k a
230+insertWith f k v (HashMap m) = SomeHashMapWith
231+ (HashMap $ HashMap.insertWith f k v m)
232+ $ InsertProof (unsafeKey k) unsafeSubset
233+
234+-- | Insert a key-value pair into the map to obtain a potentially larger map,
235+-- guaranteed to contain the given key. Like 'insertWith', but the combining
236+-- function has access to the key, which is guaranteed to be in the old map.
237+insertWithKey
238+ :: forall s k a. Hashable k
239+ => (Key s k -> a -> a -> a)
240+ -> k
241+ -> a
242+ -> HashMap s k a
243+ -> SomeHashMapWith (InsertProof 'Hashed k s) k a
244+insertWithKey f k v (HashMap m) = SomeHashMapWith
245+ (HashMap $ HashMap.insertWith (f $ unsafeKey k) k v m)
246+ $ InsertProof (unsafeKey k) unsafeSubset
247+
163248-- | Overwrite a key-value pair that is known to already be in the map. The set
164249-- of keys remains the same.
165250reinsert
@@ -183,12 +268,22 @@ insertLookupWithKey f k v (HashMap m) =
183268 $ InsertProof (unsafeKey k) unsafeSubset
184269 )
185270
271+-- | If the given key is in the map, update the value at that key using the
272+-- given function. In any case, the set of keys remains the same.
273+adjust'
274+ :: forall s k a. Hashable k => (a -> a) -> k -> HashMap s k a -> HashMap s k a
275+adjust' = coerce $ HashMap.adjust @k @a
276+
186277-- | Update the value at a specific key known the be in the map using the given
187278-- function. The set of keys remains the same.
279+--
280+-- @
281+-- 'reinsert' k v = 'adjust (const v) k'
282+-- @
188283adjust
189284 :: forall s k a. Hashable k
190285 => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a
191−adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.adjust @k @a
286+adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ adjust' @s @k @a
192287
193288-- | If the given key is in the map, update the associated value using the given
194289-- function with a proof that the key was in the map; otherwise return the map
@@ -198,6 +293,17 @@ adjustWithKey
198293 => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a
199294adjustWithKey f k (HashMap m) = HashMap $ HashMap.adjust (f $ unsafeKey k) k m
200295
296+-- | If a key is present in the map, update its value or delete it using the
297+-- given function, returning a potentially smaller map.
298+update'
299+ :: forall s k a. Hashable k
300+ => (a -> Maybe a)
301+ -> k
302+ -> HashMap s k a
303+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
304+update' f k (HashMap m) = SomeHashMapWith (HashMap $ HashMap.update f k m)
305+ $ SupersetProof unsafeSubset
306+
201307-- | Update or delete a key known to be in the map using the given function,
202308-- returning a potentially smaller map.
203309update
@@ -206,9 +312,20 @@ update
206312 -> Key s k
207313 -> HashMap s k a
208314 -> SomeHashMapWith (SupersetProof 'Hashed s) k a
209−update f k (HashMap m)
210− = SomeHashMapWith (HashMap $ HashMap.update f (unrefine k) m)
211− $ SupersetProof unsafeSubset
315+update = gcoerceWith (unsafeCastKey @s @k) $ coerce $ update' @s @k @a
316+
317+-- | If a key is present in the map, update its value or delete it using the
318+-- given function with a proof that the key was in the map, returning a
319+-- potentially smaller map.
320+updateWithKey
321+ :: forall s k a. Hashable k
322+ => (Key s k -> a -> Maybe a)
323+ -> k
324+ -> HashMap s k a
325+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
326+updateWithKey f k (HashMap m) = SomeHashMapWith
327+ (HashMap $ HashMap.update (f $ unsafeKey k) k m)
328+ $ SupersetProof unsafeSubset
212329
213330-- | If the given key is in the map, update or delete it using the given
214331-- function with a proof that the key was in the map; otherwise the map is
@@ -226,6 +343,27 @@ updateLookupWithKey f k (HashMap m) =
226343 $ SupersetProof unsafeSubset
227344 )
228345
346+-- | Given two maps proven to have the same keys, for each key apply the
347+-- function to the associated values and the key, to obtain a new map with the
348+-- same keys.
349+zipWithKey
350+ :: forall s k a b c. Hashable k
351+ => (Key s k -> a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c
352+zipWithKey f (HashMap m1) (HashMap m2) = HashMap
353+ $ HashMap.intersectionWithKey (f . unsafeKey) m1 m2
354+
355+-- | Return the union of two maps, with a given combining function for keys that
356+-- exist in both maps simultaneously.
357+unionWith
358+ :: forall s t k a. Hashable k
359+ => (a -> a -> a)
360+ -> HashMap s k a
361+ -> HashMap t k a
362+ -> SomeHashMapWith (UnionProof 'Hashed s t) k a
363+unionWith f (HashMap m1) (HashMap m2) = SomeHashMapWith
364+ (HashMap $ HashMap.unionWith f m1 m2)
365+ $ UnionProof unsafeSubset unsafeSubsetWith2
366+
229367-- | Return the union of two maps, with a given combining function for keys that
230368-- exist in both maps simultaneously.
231369--
@@ -241,8 +379,20 @@ unionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
241379 (HashMap $ HashMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
242380 $ UnionProof unsafeSubset unsafeSubsetWith2
243381
244−-- | For keys that appear in both maps, the given function decides whether the
245−-- key is removed from the first map.
382+-- | Return the first map, but for keys that appear in both maps, the given
383+-- function decides whether the key is removed.
384+differenceWith
385+ :: forall s t k a b. Hashable k
386+ => (a -> b -> Maybe a)
387+ -> HashMap s k a
388+ -> HashMap t k b
389+ -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a
390+differenceWith f (HashMap m1) (HashMap m2) = SomeHashMapWith
391+ (HashMap $ HashMap.differenceWith f m1 m2)
392+ $ PartialDifferenceProof unsafeSubset unsafeSubset
393+
394+-- | Return the first map, but for keys that appear in both maps, the given
395+-- function decides whether the key is removed.
246396--
247397-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
248398-- respectively.
@@ -259,6 +409,17 @@ differenceWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
259409 (HashMap.mapWithKey (,) m2))
260410 $ PartialDifferenceProof unsafeSubset unsafeSubset
261411
412+-- | Return the intersection of two maps with the given combining function.
413+intersectionWith
414+ :: forall s t k a b c. Hashable k
415+ => (a -> b -> c)
416+ -> HashMap s k a
417+ -> HashMap t k b
418+ -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c
419+intersectionWith f (HashMap m1) (HashMap m2) = SomeHashMapWith
420+ (HashMap $ HashMap.intersectionWith f m1 m2)
421+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
422+
262423-- | Return the intersection of two maps with the given combining function.
263424--
264425-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
@@ -273,6 +434,10 @@ intersectionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
273434 (HashMap $ HashMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
274435 $ IntersectionProof unsafeSubset unsafeSubsetWith2
275436
437+-- | Apply a function to all values in a map. The set of keys remains the same.
438+map :: forall s k a b. (a -> b) -> HashMap s k a -> HashMap s k b
439+map = coerce $ HashMap.map @a @b @k
440+
276441-- | Thread an accumularing argument through the map in ascending order of
277442-- hashes.
278443mapAccumLWithKey
@@ -291,6 +456,28 @@ mapAccumRWithKey
291456 -> (a, HashMap s k c)
292457mapAccumRWithKey f = imapAccumR (flip f)
293458
459+-- | @'mapKeys' f m@ applies @f@ to each key of @m@ and collects the results
460+-- into a new map. For keys that were mapped to the same new key, the value is
461+-- picked in an unspecified way.
462+mapKeys
463+ :: forall s k1 k2 a. Hashable k2
464+ => (Key s k1 -> k2)
465+ -> HashMap s k1 a
466+ -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a
467+mapKeys g (HashMap m) = SomeHashMapWith
468+ (HashMap $ HashMap.fromList
469+ $ HashMap.foldrWithKey (\k x xs -> (g $ unsafeKey k, x) : xs) [] m)
470+ $ MapProof (unsafeKey . g) \k2 ->
471+ case HashMap.lookup (unrefine k2) backMap of
472+ Nothing -> error "mapKeys: bug: Data.HashMap.Refined has been subverted"
473+ Just k1 -> k1
474+ where
475+ ~backMap = HashMap.fromList
476+ [ (k2, unsafeKey k1)
477+ | k1 <- HashMap.keys m
478+ , let !k2 = g $ unsafeKey k1
479+ ]
480+
294481-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
295482-- results into a new map. For keys that were mapped to the same new key, @c@
296483-- acts as the combining function for corresponding values.
@@ -315,6 +502,15 @@ mapKeysWith f g (HashMap m) = SomeHashMapWith
315502 , let !k2 = g $ unsafeKey k1
316503 ]
317504
505+-- | Apply a function to all values in a map and collect only the 'Just'
506+-- results, returning a potentially smaller map.
507+mapMaybe
508+ :: forall s k a b. (a -> Maybe b)
509+ -> HashMap s k a
510+ -> SomeHashMapWith (SupersetProof 'Hashed s) k b
511+mapMaybe f (HashMap m) = SomeHashMapWith (HashMap $ HashMap.mapMaybe f m)
512+ $ SupersetProof unsafeSubset
513+
318514-- | Apply a function to all values in a map, together with their corresponding
319515-- keys, and collect only the 'Just' results, returning a potentially smaller
320516-- map.
@@ -326,6 +522,27 @@ mapMaybeWithKey f (HashMap m)
326522 = SomeHashMapWith (HashMap $ HashMap.mapMaybeWithKey (f . unsafeKey) m)
327523 $ SupersetProof unsafeSubset
328524
525+-- | Apply a function to all values in a map and collect the 'Left' and 'Right'
526+-- results into separate (disjoint) maps.
527+mapEither
528+ :: forall s k a b c. Hashable k -- TODO: this is only used in the proof
529+ => (a -> Either b c)
530+ -> HashMap s k a
531+ -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c
532+mapEither p (HashMap m)
533+ | m' <- HashMap.map p m
534+ = Some2HashMapWith
535+ (HashMap $ HashMap.mapMaybe (either Just (const Nothing)) m')
536+ (HashMap $ HashMap.mapMaybe (either (const Nothing) Just) m')
537+ $ PartitionProof
538+ do \k -> case HashMap.lookup (unrefine k) m of
539+ Nothing -> error
540+ "mapEither: bug: Data.HashMap.Refined has been subverted"
541+ Just x -> case p x of
542+ Left _ -> Left $ unsafeKey $ unrefine k
543+ Right _ -> Right $ unsafeKey $ unrefine k
544+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
545+
329546-- | Apply a function to all values in a map, together with their corresponding
330547-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
331548-- maps.
src/Data/HashMap/Strict/Refined.hs+227−10
@@ -43,18 +43,26 @@ module Data.HashMap.Strict.Refined
4343 , SingletonProof(..)
4444 , fromSet
4545 , Common.fromHashMap
46+ , Common.verifyHashMap
47+ , fromTraversable
48+ , fromTraversableWith
4649 , fromTraversableWithKey
4750 , FromTraversableProof(..)
4851 -- * Insertion
4952 , insert
53+ , insertWith
54+ , insertWithKey
5055 , InsertProof(..)
5156 , reinsert
5257 , insertLookupWithKey
5358 -- * Deletion/Update
5459 , Common.delete
60+ , adjust'
5561 , adjust
5662 , adjustWithKey
63+ , update'
5764 , update
65+ , updateWithKey
5866 , updateLookupWithKey
5967 -- * Query
6068 , Common.lookup
@@ -66,21 +74,29 @@ module Data.HashMap.Strict.Refined
6674 , Common.disjoint
6775 , DisjointProof(..)
6876 -- * Combine
77+ , zipWith
6978 , zipWithKey
7079 , bind
80+ , Common.union
81+ , unionWith
7182 , unionWithKey
7283 , UnionProof(..)
7384 , Common.difference
7485 , DifferenceProof(..)
86+ , differenceWith
7587 , differenceWithKey
7688 , PartialDifferenceProof(..)
89+ , Common.intersection
90+ , intersectionWith
7791 , intersectionWithKey
7892 , IntersectionProof(..)
7993 -- * Traversal
94+ , map
8095 , mapWithKey
8196 , traverseWithKey
8297 , mapAccumLWithKey
8398 , mapAccumRWithKey
99+ , mapKeys
84100 , mapKeysWith
85101 , MapProof(..)
86102 , backpermuteKeys
@@ -97,10 +113,15 @@ module Data.HashMap.Strict.Refined
97113 -- * Filter
98114 , Common.restrictKeys
99115 , Common.withoutKeys
116+ , Common.filter
117+ , Common.filterKeys
100118 , Common.filterWithKey
119+ , Common.partition
101120 , Common.partitionWithKey
102121 , PartitionProof(..)
122+ , mapMaybe
103123 , mapMaybeWithKey
124+ , mapEither
104125 , mapEitherWithKey
105126 -- * Casts
106127 , Common.castKey
@@ -125,7 +146,7 @@ import Data.Reflection
125146import Data.Traversable
126147import Data.Traversable.WithIndex
127148import Data.Type.Coercion
128−import Prelude hiding (lookup, null)
149+import Prelude hiding (lookup, map, null, zipWith)
129150import Refined
130151import Refined.Unsafe
131152
@@ -144,14 +165,49 @@ fromSet :: forall s k a. KnownHashSet s k => (Key s k -> a) -> HashMap s k a
144165fromSet f = HashMap $ HashMap.mapWithKey (\k _ -> f $ unsafeKey k)
145166 $ HashSet.toMap (reflect $ Proxy @s)
146167
147−-- | Create a map from an arbitrary traversable of key-value pairs.
168+-- | Create a map from an arbitrary traversable of key-value pairs. If a key is
169+-- repeated, the retained value is the last one in traversal order. If you're
170+-- looking for @fromList@, this is the function you want.
171+fromTraversable
172+ :: forall t k a. (Traversable t, Hashable k)
173+ => t (k, a) -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
174+fromTraversable xs = SomeHashMapWith (HashMap m) $ FromTraversableProof proof
175+ where
176+ (m, proof) = mapAccumL
177+ (\s (k, v) -> let !s' = HashMap.insert k v s in (s', unsafeKey k))
178+ HashMap.empty
179+ xs
180+
181+-- | Create a map from an arbitrary traversable of key-value pairs, with a
182+-- function for combining values for repeated keys. The function is called as if
183+-- by 'foldl1', but flipped:
184+--
185+-- @
186+-- 'fromTraversableWith' f [(k, x1), (k, x2), (k, x3)]
187+-- = 'singleton' k (f x3 (f x2 x1))
188+-- @
189+fromTraversableWith
190+ :: forall t k a. (Traversable t, Hashable k)
191+ => (a -> a -> a)
192+ -> t (k, a)
193+ -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
194+fromTraversableWith f xs
195+ = SomeHashMapWith (HashMap m) $ FromTraversableProof proof
196+ where
197+ (m, proof) = mapAccumL
198+ (\s (k, v) -> let !s' = HashMap.insertWith f k v s in (s', unsafeKey k))
199+ HashMap.empty
200+ xs
201+
202+-- | Create a map from an arbitrary traversable of key-value pairs. Like
203+-- 'fromTraversableWith', but the combining function has access to the key.
148204fromTraversableWithKey
149205 :: forall t k a. (Traversable t, Hashable k)
150206 => (k -> a -> a -> a)
151207 -> t (k, a)
152208 -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
153−fromTraversableWithKey f xs = SomeHashMapWith (HashMap m)
154− $ FromTraversableProof proof
209+fromTraversableWithKey f xs
210+ = SomeHashMapWith (HashMap m) $ FromTraversableProof proof
155211 where
156212 (m, proof) = mapAccumL
157213 (\s (k, v)
@@ -168,8 +224,41 @@ insert
168224insert k v (HashMap m) = SomeHashMapWith (HashMap $ HashMap.insert k v m)
169225 $ InsertProof (unsafeKey k) unsafeSubset
170226
227+-- | Insert a key-value pair into the map to obtain a potentially larger map,
228+-- guaranteed to contain the given key. If the key was already present, the
229+-- supplied function is used to combine the new value with the old (in that
230+-- order).
231+insertWith
232+ :: forall s k a. Hashable k
233+ => (a -> a -> a)
234+ -> k
235+ -> a
236+ -> HashMap s k a
237+ -> SomeHashMapWith (InsertProof 'Hashed k s) k a
238+insertWith f k v (HashMap m) = SomeHashMapWith
239+ (HashMap $ HashMap.insertWith f k v m)
240+ $ InsertProof (unsafeKey k) unsafeSubset
241+
242+-- | Insert a key-value pair into the map to obtain a potentially larger map,
243+-- guaranteed to contain the given key. Like 'insertWith', but the combining
244+-- function has access to the key, which is guaranteed to be in the old map.
245+insertWithKey
246+ :: forall s k a. Hashable k
247+ => (Key s k -> a -> a -> a)
248+ -> k
249+ -> a
250+ -> HashMap s k a
251+ -> SomeHashMapWith (InsertProof 'Hashed k s) k a
252+insertWithKey f k v (HashMap m) = SomeHashMapWith
253+ (HashMap $ HashMap.insertWith (f $ unsafeKey k) k v m)
254+ $ InsertProof (unsafeKey k) unsafeSubset
255+
171256-- | Overwrite a key-value pair that is known to already be in the map. The set
172257-- of keys remains the same.
258+--
259+-- @
260+-- 'reinsert' k v = 'adjust (const v) k'
261+-- @
173262reinsert
174263 :: forall s k a. Hashable k
175264 => Key s k -> a -> HashMap s k a -> HashMap s k a
@@ -191,12 +280,18 @@ insertLookupWithKey f k v (HashMap m) =
191280 $ InsertProof (unsafeKey k) unsafeSubset
192281 )
193282
283+-- | If the given key is in the map, update the value at that key using the
284+-- given function. In any case, the set of keys remains the same.
285+adjust'
286+ :: forall s k a. Hashable k => (a -> a) -> k -> HashMap s k a -> HashMap s k a
287+adjust' = coerce $ HashMap.adjust @k @a
288+
194289-- | Update the value at a specific key known the be in the map using the given
195290-- function. The set of keys remains the same.
196291adjust
197292 :: forall s k a. Hashable k
198293 => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a
199−adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.adjust @k @a
294+adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ adjust' @s @k @a
200295
201296-- | If the given key is in the map, update the associated value using the given
202297-- function with a proof that the key was in the map; otherwise return the map
@@ -206,6 +301,17 @@ adjustWithKey
206301 => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a
207302adjustWithKey f k (HashMap m) = HashMap $ HashMap.adjust (f $ unsafeKey k) k m
208303
304+-- | If a key is present in the map, update its value or delete it using the
305+-- given function, returning a potentially smaller map.
306+update'
307+ :: forall s k a. Hashable k
308+ => (a -> Maybe a)
309+ -> k
310+ -> HashMap s k a
311+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
312+update' f k (HashMap m) = SomeHashMapWith (HashMap $ HashMap.update f k m)
313+ $ SupersetProof unsafeSubset
314+
209315-- | Update or delete a key known to be in the map using the given function,
210316-- returning a potentially smaller map.
211317update
@@ -214,9 +320,20 @@ update
214320 -> Key s k
215321 -> HashMap s k a
216322 -> SomeHashMapWith (SupersetProof 'Hashed s) k a
217−update f k (HashMap m)
218− = SomeHashMapWith (HashMap $ HashMap.update f (unrefine k) m)
219− $ SupersetProof unsafeSubset
323+update = gcoerceWith (unsafeCastKey @s @k) $ coerce $ update' @s @k @a
324+
325+-- | If a key is present in the map, update its value or delete it using the
326+-- given function with a proof that the key was in the map, returning a
327+-- potentially smaller map.
328+updateWithKey
329+ :: forall s k a. Hashable k
330+ => (Key s k -> a -> Maybe a)
331+ -> k
332+ -> HashMap s k a
333+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
334+updateWithKey f k (HashMap m) = SomeHashMapWith
335+ (HashMap $ HashMap.update (f $ unsafeKey k) k m)
336+ $ SupersetProof unsafeSubset
220337
221338-- | If the given key is in the map, update or delete it using the given
222339-- function with a proof that the key was in the map; otherwise the map is
@@ -236,12 +353,32 @@ updateLookupWithKey f k (HashMap m) =
236353
237354-- | Given two maps proven to have the same keys, for each key apply the
238355-- function to the associated values, to obtain a new map with the same keys.
356+zipWith
357+ :: forall s k a b c. Hashable k
358+ => (a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c
359+zipWith f (HashMap m1) (HashMap m2) = HashMap $ HashMap.intersectionWith f m1 m2
360+
361+-- | Given two maps proven to have the same keys, for each key apply the
362+-- function to the associated values and the key, to obtain a new map with the
363+-- same keys.
239364zipWithKey
240365 :: forall s k a b c. Hashable k
241366 => (Key s k -> a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c
242367zipWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
243368 $ HashMap.intersectionWithKey @k @a @b @c
244369
370+-- | Return the union of two maps, with a given combining function for keys that
371+-- exist in both maps simultaneously.
372+unionWith
373+ :: forall s t k a. Hashable k
374+ => (a -> a -> a)
375+ -> HashMap s k a
376+ -> HashMap t k a
377+ -> SomeHashMapWith (UnionProof 'Hashed s t) k a
378+unionWith f (HashMap m1) (HashMap m2) = SomeHashMapWith
379+ (HashMap $ HashMap.unionWith f m1 m2)
380+ $ UnionProof unsafeSubset unsafeSubsetWith2
381+
245382-- | Return the union of two maps, with a given combining function for keys that
246383-- exist in both maps simultaneously.
247384--
@@ -257,8 +394,20 @@ unionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
257394 (HashMap $ HashMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
258395 $ UnionProof unsafeSubset unsafeSubsetWith2
259396
260−-- | For keys that appear in both maps, the given function decides whether the
261−-- key is removed from the first map.
397+-- | Return the first map, but for keys that appear in both maps, the given
398+-- function decides whether the key is removed.
399+differenceWith
400+ :: forall s t k a b. Hashable k
401+ => (a -> b -> Maybe a)
402+ -> HashMap s k a
403+ -> HashMap t k b
404+ -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a
405+differenceWith f (HashMap m1) (HashMap m2) = SomeHashMapWith
406+ (HashMap $ HashMap.differenceWith f m1 m2)
407+ $ PartialDifferenceProof unsafeSubset unsafeSubset
408+
409+-- | Return the first map, but for keys that appear in both maps, the given
410+-- function decides whether the key is removed.
262411--
263412-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
264413-- respectively.
@@ -275,6 +424,17 @@ differenceWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
275424 (HashMap.mapWithKey (,) m2))
276425 $ PartialDifferenceProof unsafeSubset unsafeSubset
277426
427+-- | Return the intersection of two maps with the given combining function.
428+intersectionWith
429+ :: forall s t k a b c. Hashable k
430+ => (a -> b -> c)
431+ -> HashMap s k a
432+ -> HashMap t k b
433+ -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c
434+intersectionWith f (HashMap m1) (HashMap m2) = SomeHashMapWith
435+ (HashMap $ HashMap.intersectionWith f m1 m2)
436+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
437+
278438-- | Return the intersection of two maps with the given combining function.
279439--
280440-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
@@ -289,6 +449,10 @@ intersectionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
289449 (HashMap $ HashMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
290450 $ IntersectionProof unsafeSubset unsafeSubsetWith2
291451
452+-- | Apply a function to all values in a map. The set of keys remains the same.
453+map :: forall s k a b. (a -> b) -> HashMap s k a -> HashMap s k b
454+map = coerce $ HashMap.map @a @b @k
455+
292456-- | Apply a function to all values in a map, together with their corresponding
293457-- keys, that are proven to be in the map. The set of keys remains the same.
294458mapWithKey
@@ -323,6 +487,29 @@ mapAccumRWithKey
323487 -> (a, HashMap s k c)
324488mapAccumRWithKey f = imapAccumR (flip f)
325489
490+-- | @'mapKeys' f m@ applies @f@ to each key of @m@ and collects the results
491+-- into a new map. For keys that were mapped to the same new key, the value is
492+-- picked in an unspecified way.
493+mapKeys
494+ :: forall s k1 k2 a. Hashable k2
495+ => (Key s k1 -> k2)
496+ -> HashMap s k1 a
497+ -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a
498+mapKeys g (HashMap m) = SomeHashMapWith
499+ (HashMap $ HashMap.fromList
500+ $ HashMap.foldrWithKey (\k x xs -> (g $ unsafeKey k, x) : xs) [] m)
501+ $ MapProof (unsafeKey . g) \k2 ->
502+ case HashMap.lookup (unrefine k2) backMap of
503+ Nothing -> error
504+ "mapKeys: bug: Data.HashMap.Strict.Refined has been subverted"
505+ Just k1 -> k1
506+ where
507+ ~backMap = HashMap.fromList
508+ [ (k2, unsafeKey k1)
509+ | k1 <- HashMap.keys m
510+ , let !k2 = g $ unsafeKey k1
511+ ]
512+
326513-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
327514-- results into a new map. For keys that were mapped to the same new key, @c@
328515-- acts as the combining function for corresponding values.
@@ -347,6 +534,15 @@ mapKeysWith f g (HashMap m) = SomeHashMapWith
347534 , let !k2 = g $ unsafeKey k1
348535 ]
349536
537+-- | Apply a function to all values in a map and collect only the 'Just'
538+-- results, returning a potentially smaller map.
539+mapMaybe
540+ :: forall s k a b. (a -> Maybe b)
541+ -> HashMap s k a
542+ -> SomeHashMapWith (SupersetProof 'Hashed s) k b
543+mapMaybe f (HashMap m) = SomeHashMapWith (HashMap $ HashMap.mapMaybe f m)
544+ $ SupersetProof unsafeSubset
545+
350546-- | Apply a function to all values in a map, together with their corresponding
351547-- keys, and collect only the 'Just' results, returning a potentially smaller
352548-- map.
@@ -358,6 +554,27 @@ mapMaybeWithKey f (HashMap m)
358554 = SomeHashMapWith (HashMap $ HashMap.mapMaybeWithKey (f . unsafeKey) m)
359555 $ SupersetProof unsafeSubset
360556
557+-- | Apply a function to all values in a map and collect the 'Left' and 'Right'
558+-- results into separate (disjoint) maps.
559+mapEither
560+ :: forall s k a b c. Hashable k -- TODO: this is only used in the proof
561+ => (a -> Either b c)
562+ -> HashMap s k a
563+ -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c
564+mapEither p (HashMap m)
565+ | m' <- HashMap.map p m
566+ = Some2HashMapWith
567+ (HashMap $ HashMap.mapMaybe (either Just (const Nothing)) m')
568+ (HashMap $ HashMap.mapMaybe (either (const Nothing) Just) m')
569+ $ PartitionProof
570+ do \k -> case HashMap.lookup (unrefine k) m of
571+ Nothing -> error
572+ "mapEither: bug: Data.HashMap.Refined has been subverted"
573+ Just x -> case p x of
574+ Left _ -> Left $ unsafeKey $ unrefine k
575+ Right _ -> Right $ unsafeKey $ unrefine k
576+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
577+
361578-- | Apply a function to all values in a map, together with their corresponding
362579-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
363580-- maps.
src/Data/HashSet/Refined.hs+1−1
@@ -165,7 +165,7 @@ data SomeHashSet a where
165165-- you a way to refer to the set (the parameter @s@), e.g.:
166166--
167167-- @
168−-- 'withHashSet' ('fromHashSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s
168+-- 'withHashSet' ('fromHashSet' ...) $ \\(_ :: 'Proxy' s) -> doSomethingWith \@s
169169-- @
170170withHashSet
171171 :: forall a r. SomeHashSet a
src/Data/IntMap/Common/Refined.hs+134−15
@@ -14,11 +14,13 @@ import Data.Functor.Rep
1414import Data.Functor.WithIndex
1515import qualified Data.Hashable as Hashable
1616import qualified Data.IntMap as IntMap
17+import qualified Data.IntSet as IntSet
1718import Data.Proxy
1819import Data.Reflection
1920import Data.Traversable.WithIndex
2021import Data.Type.Coercion
2122import Data.Type.Equality ((:~:)(..))
23+import Prelude hiding (zipWith)
2224import Refined
2325import Refined.Unsafe
2426import Unsafe.Coerce
@@ -36,6 +38,9 @@ import qualified Data.IntMap.Strict as IntMapStrict
3638import qualified Data.List as List
3739#endif
3840
41+#if MIN_VERSION_containers(0, 5, 9) || MIN_VERSION_hashable(1, 4, 0)
42+import Data.Functor.Classes
43+#endif
3944
4045-- | A wrapper around a regular 'Data.IntMap.IntMap' with a type parameter @s@
4146-- identifying the set of keys present in the map.
@@ -51,6 +56,9 @@ newtype IntMap s a = IntMap (IntMap.IntMap a)
5156 deriving newtype (Eq, Ord, Show, Functor, Foldable, NFData)
5257#if MIN_VERSION_hashable(1, 3, 4)
5358 deriving newtype (Hashable.Hashable)
59+#endif
60+#if MIN_VERSION_containers(0, 5, 9) || MIN_VERSION_hashable(1, 4, 0)
61+ deriving newtype (Eq1, Ord1, Show1)
5462#endif
5563 deriving stock (Traversable)
5664type role IntMap nominal representational
@@ -96,7 +104,8 @@ data SomeIntMap a where
96104-- parameter @s@), e.g.:
97105--
98106-- @
99−-- 'withIntMap' ('fromIntMap' ...) $ \(m :: 'IntMap' s a) -> doSomethingWith \@s
107+-- 'withIntMap' ('fromIntMap' ...)
108+-- $ \\(m :: 'IntMap' s a) -> doSomethingWith \@s
100109-- @
101110withIntMap :: forall a r. SomeIntMap a -> (forall s. IntMap s a -> r) -> r
102111withIntMap (SomeIntMap m) k = k m
@@ -105,6 +114,14 @@ withIntMap (SomeIntMap m) k = k m
105114fromIntMap :: forall a. IntMap.IntMap a -> SomeIntMap a
106115fromIntMap m = SomeIntMap (IntMap m)
107116
117+-- | Given a set of keys @s@ known ahead of time, verify whether a regular
118+-- 'Data.IntMap.IntMap' has exactly that set of keys.
119+verifyIntMap
120+ :: forall s a. KnownIntSet s => IntMap.IntMap a -> Maybe (IntMap s a)
121+verifyIntMap m
122+ | IntMap.keys m == IntSet.toList (reflect $ Proxy @s) = Just (IntMap m)
123+ | otherwise = Nothing
124+
108125-- | An existential wrapper for an 'IntMap' with an as-yet-unknown set of keys,
109126-- together with a proof of some fact @p@ about the set. Pattern matching on it
110127-- gives you a way to refer to the set (the parameter @s@). Functions that
@@ -234,23 +251,36 @@ disjoint (IntMap m1) (IntMap m2)
234251
235252-- | Given two maps proven to have the same keys, for each key apply the
236253-- function to the associated values, to obtain a new map with the same keys.
237−zipWithKey
238− :: forall s a b c. (Key s -> a -> b -> c)
254+zipWith
255+ :: forall s a b c. (a -> b -> c)
239256 -> IntMap s a
240257 -> IntMap s b
241258 -> IntMap s c
242−zipWithKey f (IntMap m1) (IntMap m2) = IntMap
243− $ IntMap.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
259+zipWith f (IntMap m1) (IntMap m2) = IntMap
260+ $ IntMap.mergeWithKey (\_ x y -> Just $ f x y)
244261 (\m -> if IntMap.null m
245262 then IntMap.empty
246− else error "zipWithKey: bug: Data.IntMap.Refined has been subverted")
263+ else error "zipWith: bug: Data.IntMap.Refined has been subverted")
247264 (\m -> if IntMap.null m
248265 then IntMap.empty
249− else error "zipWithKey: bug: Data.IntMap.Refined has been subverted")
266+ else error "zipWith: bug: Data.IntMap.Refined has been subverted")
250267 -- ^ Work around https://github.com/haskell/containers/issues/979
251268 m1
252269 m2
253270
271+-- | Return the union of two maps. For keys that exist in both maps, the value
272+-- is taken from the first map.
273+--
274+-- @
275+-- 'union' = unionWith 'const'
276+-- @
277+union
278+ :: forall s t a. IntMap s a
279+ -> IntMap t a
280+ -> SomeIntMapWith (UnionProof 'Int s t) a
281+union (IntMap m1) (IntMap m2) = SomeIntMapWith (IntMap $ IntMap.union m1 m2)
282+ $ UnionProof unsafeSubset unsafeSubsetWith2
283+
254284-- | Remove the keys that appear in the second map from the first map.
255285difference
256286 :: forall s t a b. IntMap s a
@@ -260,6 +290,19 @@ difference (IntMap m1) (IntMap m2) = SomeIntMapWith
260290 (IntMap $ IntMap.difference m1 m2)
261291 $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
262292
293+-- | Return the intersection of two maps, taking values from the first map.
294+--
295+-- @
296+-- 'intersection' = intersectionWith 'const'
297+-- @
298+intersection
299+ :: forall s t a b. IntMap s a
300+ -> IntMap t b
301+ -> SomeIntMapWith (IntersectionProof 'Int s t) a
302+intersection (IntMap m1) (IntMap m2) = SomeIntMapWith
303+ (IntMap $ IntMap.intersection m1 m2)
304+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
305+
263306-- | Apply a function to all values in a map, together with their corresponding
264307-- keys, that are proven to be in the map. The set of keys remains the same.
265308mapWithKey :: forall s a b. (Key s -> a -> b) -> IntMap s a -> IntMap s b
@@ -316,6 +359,29 @@ toList = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.toAscList @a
316359toDescList :: forall s a. IntMap s a -> [(Key s, a)]
317360toDescList = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.toDescList @a
318361
362+-- | Retain only the key-value pairs that satisfy the predicate, returning a
363+-- potentially smaller map.
364+filter
365+ :: forall s a. (a -> Bool)
366+ -> IntMap s a
367+ -> SomeIntMapWith (SupersetProof 'Int s) a
368+filter p (IntMap m) = SomeIntMapWith (IntMap $ IntMap.filter p m)
369+ $ SupersetProof unsafeSubset
370+
371+-- | Retain only the key-value pairs that satisfy the predicate, returning a
372+-- potentially smaller map.
373+filterKeys
374+ :: forall s a. (Key s -> Bool)
375+ -> IntMap s a
376+ -> SomeIntMapWith (SupersetProof 'Int s) a
377+filterKeys p (IntMap m) = SomeIntMapWith
378+#if MIN_VERSION_containers(0, 8, 0)
379+ (IntMap $ IntMap.filterKeys (p . unsafeKey) m)
380+#else
381+ (IntMap $ IntMap.filterWithKey (\k _ -> p (unsafeKey k)) m)
382+#endif
383+ $ SupersetProof unsafeSubset
384+
319385-- | Retain only the key-value pairs that satisfy the predicate, returning a
320386-- potentially smaller map.
321387filterWithKey
@@ -351,6 +417,22 @@ withoutKeys (IntMap m) = SomeIntMapWith
351417#endif
352418 $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
353419
420+-- | Partition a map into two disjoint submaps: those whose values satisfy the
421+-- predicate, and those whose don't.
422+partition
423+ :: forall s a. (a -> Bool)
424+ -> IntMap s a
425+ -> Some2IntMapWith (PartitionProof 'Int s Int) a a
426+partition p (IntMap m) = case IntMap.partition p m of
427+ (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof
428+ do \k -> case IntMap.lookup (unrefine k) m of
429+ Nothing -> error
430+ "partition: bug: Data.IntMap.Refined has been subverted"
431+ Just x -> if p x
432+ then Left $ unsafeKey $ unrefine k
433+ else Right $ unsafeKey $ unrefine k
434+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
435+
354436-- | Partition a map into two disjoint submaps: those whose key-value pairs
355437-- satisfy the predicate, and those whose don't.
356438partitionWithKey
@@ -368,6 +450,43 @@ partitionWithKey p (IntMap m)
368450 else Right $ unsafeKey $ unrefine k
369451 unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
370452
453+-- | Take the a submap of keys up to a point where the predicate stops holding.
454+--
455+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
456+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
457+-- chosen in an unspecified way.
458+takeWhileAntitone
459+ :: forall s a. (Key s -> Bool)
460+ -> IntMap s a
461+ -> SomeIntMapWith (SupersetProof 'Int s) a
462+takeWhileAntitone p (IntMap m) = SomeIntMapWith
463+#if MIN_VERSION_containers(0, 6, 7)
464+ (IntMap $ IntMap.takeWhileAntitone (p . unsafeKey) m)
465+#else
466+ (IntMap $ IntMap.fromDistinctAscList
467+ $ List.takeWhile (p . unsafeKey . fst) $ IntMap.toAscList m)
468+#endif
469+ $ SupersetProof unsafeSubset
470+
471+-- | Take the a submap of keys starting from a point where the predicate stops
472+-- holding.
473+--
474+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
475+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
476+-- chosen in an unspecified way.
477+dropWhileAntitone
478+ :: forall s a. (Key s -> Bool)
479+ -> IntMap s a
480+ -> SomeIntMapWith (SupersetProof 'Int s) a
481+dropWhileAntitone p (IntMap m) = SomeIntMapWith
482+#if MIN_VERSION_containers(0, 6, 7)
483+ (IntMap $ IntMap.dropWhileAntitone (p . unsafeKey) m)
484+#else
485+ (IntMap $ IntMap.fromDistinctAscList
486+ $ List.dropWhile (p . unsafeKey . fst) $ IntMap.toAscList m)
487+#endif
488+ $ SupersetProof unsafeSubset
489+
371490-- | Divide a map into two disjoint submaps at a point where the predicate on
372491-- the keys stops holding.
373492--
@@ -442,17 +561,17 @@ castKey
442561castKey = castRefined
443562
444563-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps
445−-- can be interconverted too. For example, 'zipWithKey' can be implemented via
446−-- 'Data.IntMap.Refined.intersectionWithKey' by proving that the set of keys
447−-- remains unchanged:
564+-- can be interconverted too. For example, 'Data.IntMap.Refined.zipWithKey' can
565+-- be implemented via 'Data.IntMap.Refined.intersectionWithKey' by proving that
566+-- the set of keys remains unchanged:
448567--
449568-- @
450−-- 'zipWithKey'
569+-- 'Data.IntMap.Refined.zipWithKey'
451570-- :: forall s a b c. ('Key' s -> a -> b -> c)
452571-- -> 'IntMap' s a
453572-- -> 'IntMap' s b
454573-- -> 'IntMap' s c
455−-- 'zipWithKey' f m1 m2
574+-- 'Data.IntMap.Refined.zipWithKey' f m1 m2
456575-- | v'SomeIntMapWith' @r m proof <- 'Data.IntMap.Refined.intersectionWithKey' (f . 'andLeft') m1 m2
457576-- , v'IntersectionProof' p1 p2 <- proof
458577-- , ( v'Coercion' :: t'Coercion' ('IntMap' r c) ('IntMap' s c))
@@ -479,10 +598,10 @@ instance TraversableWithIndex (Key s) (IntMap s) where
479598 itraverse = traverseWithKey
480599
481600-- | Similar to the instance for functions -- zip corresponding keys. To use
482−-- '<*>'/'Control.Applicative.liftA2' without 'KnownIntSet' see 'zipWithKey'.
601+-- '<*>'/'Control.Applicative.liftA2' without 'KnownIntSet' see 'zipWith'.
483602instance KnownIntSet s => Applicative (IntMap s) where
484603 pure x = fromSet \_ -> x
485− (<*>) = zipWithKey (const id)
604+ (<*>) = zipWith id
486605
487606-- | @'bind' m f@ is a map that for each key @k :: 'Key' s@, contains the
488607-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
@@ -502,7 +621,7 @@ instance KnownIntSet s => MonadReader (Key s) (IntMap s) where
502621
503622-- | Append the values at the corresponding keys
504623instance Semigroup a => Semigroup (IntMap s a) where
505− (<>) = zipWithKey (const (<>))
624+ (<>) = zipWith (<>)
506625
507626instance (KnownIntSet s, Monoid a) => Monoid (IntMap s a) where
508627 mempty = fromSet \_ -> mempty
src/Data/IntMap/Refined.hs+225−7
@@ -34,18 +34,26 @@ module Data.IntMap.Refined
3434 , SingletonProof(..)
3535 , fromSet
3636 , Common.fromIntMap
37+ , Common.verifyIntMap
38+ , fromTraversable
39+ , fromTraversableWith
3740 , fromTraversableWithKey
3841 , FromTraversableProof(..)
3942 -- * Insertion
4043 , insert
44+ , insertWith
45+ , insertWithKey
4146 , InsertProof(..)
4247 , reinsert
4348 , insertLookupWithKey
4449 -- * Deletion/Update
4550 , Common.delete
51+ , adjust'
4652 , adjust
4753 , adjustWithKey
54+ , update'
4855 , update
56+ , updateWithKey
4957 , updateLookupWithKey
5058 -- * Query
5159 , Common.lookup
@@ -61,21 +69,29 @@ module Data.IntMap.Refined
6169 , Common.disjoint
6270 , DisjointProof(..)
6371 -- * Combine
72+ , zipWith
6473 , zipWithKey
6574 , bind
75+ , Common.union
76+ , unionWith
6677 , unionWithKey
6778 , UnionProof(..)
6879 , Common.difference
6980 , DifferenceProof(..)
81+ , differenceWith
7082 , differenceWithKey
7183 , PartialDifferenceProof(..)
84+ , Common.intersection
85+ , intersectionWith
7286 , intersectionWithKey
7387 , IntersectionProof(..)
7488 -- * Traversal
89+ , map
7590 , mapWithKey
7691 , traverseWithKey
7792 , mapAccumLWithKey
7893 , mapAccumRWithKey
94+ , mapKeys
7995 , mapKeysWith
8096 , MapProof(..)
8197 , backpermuteKeys
@@ -93,12 +109,19 @@ module Data.IntMap.Refined
93109 -- * Filter
94110 , Common.restrictKeys
95111 , Common.withoutKeys
112+ , Common.filter
113+ , Common.filterKeys
96114 , Common.filterWithKey
115+ , Common.partition
97116 , Common.partitionWithKey
98117 , PartitionProof(..)
118+ , Common.takeWhileAntitone
119+ , Common.dropWhileAntitone
99120 , Common.spanAntitone
100121 , PartialPartitionProof(..)
122+ , mapMaybe
101123 , mapMaybeWithKey
124+ , mapEither
102125 , mapEitherWithKey
103126 , Common.splitLookup
104127 , SplitProof(..)
@@ -122,13 +145,13 @@ import Data.Functor
122145import qualified Data.IntMap as IntMap
123146import Data.IntMap.Common.Refined
124147 ( IntMap(..), Key, unsafeCastKey, unsafeKey, SomeIntMapWith(..)
125− , Some2IntMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey
148+ , Some2IntMapWith(..), fromSet, (!), zipWith, mapWithKey, traverseWithKey
126149 , bind
127150 )
128151import qualified Data.IntMap.Common.Refined as Common
129152import Data.Traversable
130153import Data.Type.Coercion
131−import Prelude hiding (lookup, null)
154+import Prelude hiding (lookup, map, null, zipWith)
132155import Refined
133156import Refined.Unsafe
134157
@@ -139,7 +162,42 @@ singleton :: forall a. Int -> a -> SomeIntMapWith (SingletonProof 'Int Int) a
139162singleton k v = SomeIntMapWith (IntMap $ IntMap.singleton k v)
140163 $ SingletonProof (unsafeKey k)
141164
142−-- | Create a map from an arbitrary traversable of key-value pairs.
165+-- | Create a map from an arbitrary traversable of key-value pairs. If a key is
166+-- repeated, the retained value is the last one in traversal order. If you're
167+-- looking for @fromList@, this is the function you want.
168+fromTraversable
169+ :: forall t a. Traversable t
170+ => t (Int, a) -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
171+fromTraversable xs = SomeIntMapWith (IntMap m) $ FromTraversableProof proof
172+ where
173+ (m, proof) = mapAccumL
174+ (\s (k, v) -> let !s' = IntMap.insert k v s in (s', unsafeKey k))
175+ IntMap.empty
176+ xs
177+
178+-- | Create a map from an arbitrary traversable of key-value pairs, with a
179+-- function for combining values for repeated keys. The function is called as if
180+-- by 'foldl1', but flipped:
181+--
182+-- @
183+-- 'fromTraversableWith' f [(k, x1), (k, x2), (k, x3)]
184+-- = 'singleton' k (f x3 (f x2 x1))
185+-- @
186+fromTraversableWith
187+ :: forall t a. Traversable t
188+ => (a -> a -> a)
189+ -> t (Int, a)
190+ -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
191+fromTraversableWith f xs
192+ = SomeIntMapWith (IntMap m) $ FromTraversableProof proof
193+ where
194+ (m, proof) = mapAccumL
195+ (\s (k, v) -> let !s' = IntMap.insertWith f k v s in (s', unsafeKey k))
196+ IntMap.empty
197+ xs
198+
199+-- | Create a map from an arbitrary traversable of key-value pairs. Like
200+-- 'fromTraversableWith', but the combining function has access to the key.
143201fromTraversableWithKey
144202 :: forall t a. Traversable t
145203 => (Int -> a -> a -> a)
@@ -164,6 +222,33 @@ insert
164222insert k v (IntMap m) = SomeIntMapWith (IntMap $ IntMap.insert k v m)
165223 $ InsertProof (unsafeKey k) unsafeSubset
166224
225+-- | Insert a key-value pair into the map to obtain a potentially larger map,
226+-- guaranteed to contain the given key. If the key was already present, the
227+-- supplied function is used to combine the new value with the old (in that
228+-- order).
229+insertWith
230+ :: forall s a. (a -> a -> a)
231+ -> Int
232+ -> a
233+ -> IntMap s a
234+ -> SomeIntMapWith (InsertProof 'Int Int s) a
235+insertWith f k v (IntMap m) = SomeIntMapWith
236+ (IntMap $ IntMap.insertWith f k v m)
237+ $ InsertProof (unsafeKey k) unsafeSubset
238+
239+-- | Insert a key-value pair into the map to obtain a potentially larger map,
240+-- guaranteed to contain the given key. Like 'insertWith', but the combining
241+-- function has access to the key, which is guaranteed to be in the old map.
242+insertWithKey
243+ :: forall s a. (Key s -> a -> a -> a)
244+ -> Int
245+ -> a
246+ -> IntMap s a
247+ -> SomeIntMapWith (InsertProof 'Int Int s) a
248+insertWithKey f k v (IntMap m) = SomeIntMapWith
249+ (IntMap $ IntMap.insertWithKey (f . unsafeKey) k v m)
250+ $ InsertProof (unsafeKey k) unsafeSubset
251+
167252-- | Overwrite a key-value pair that is known to already be in the map. The set
168253-- of keys remains the same.
169254reinsert
@@ -184,10 +269,19 @@ insertLookupWithKey f k v (IntMap m)
184269 (v', !m') -> ((unsafeKey k,) <$> v',)
185270 $ SomeIntMapWith (IntMap m') $ InsertProof (unsafeKey k) unsafeSubset
186271
272+-- | If the given key is in the map, update the value at that key using the
273+-- given function. In any case, the set of keys remains the same.
274+adjust' :: forall s a. (a -> a) -> Int -> IntMap s a -> IntMap s a
275+adjust' = coerce $ IntMap.adjust @a
276+
187277-- | Update the value at a specific key known the be in the map using the given
188278-- function. The set of keys remains the same.
279+--
280+-- @
281+-- 'reinsert' k v = 'adjust (const v) k'
282+-- @
189283adjust :: forall s a. (a -> a) -> Key s -> IntMap s a -> IntMap s a
190−adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.adjust @a
284+adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ adjust' @s @a
191285
192286-- | If the given key is in the map, update the associated value using the given
193287-- function with a proof that the key was in the map; otherwise return the map
@@ -197,6 +291,16 @@ adjustWithKey
197291adjustWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
198292 $ IntMap.adjustWithKey @a
199293
294+-- | If a key is present in the map, update its value or delete it using the
295+-- given function, returning a potentially smaller map.
296+update'
297+ :: forall s a. (a -> Maybe a)
298+ -> Int
299+ -> IntMap s a
300+ -> SomeIntMapWith (SupersetProof 'Int s) a
301+update' f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f k m)
302+ $ SupersetProof unsafeSubset
303+
200304-- | Update or delete a key known to be in the map using the given function,
201305-- returning a potentially smaller map.
202306update
@@ -204,7 +308,18 @@ update
204308 -> Key s
205309 -> IntMap s a
206310 -> SomeIntMapWith (SupersetProof 'Int s) a
207−update f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f (unrefine k) m)
311+update = gcoerceWith (unsafeCastKey @s) $ coerce $ update' @s @a
312+
313+-- | If a key is present in the map, update its value or delete it using the
314+-- given function with a proof that the key was in the map, returning a
315+-- potentially smaller map.
316+updateWithKey
317+ :: forall s a. (Key s -> a -> Maybe a)
318+ -> Int
319+ -> IntMap s a
320+ -> SomeIntMapWith (SupersetProof 'Int s) a
321+updateWithKey f k (IntMap m) = SomeIntMapWith
322+ (IntMap $ IntMap.updateWithKey (f . unsafeKey) k m)
208323 $ SupersetProof unsafeSubset
209324
210325-- | If the given key is in the map, update or delete it using the given
@@ -221,6 +336,37 @@ updateLookupWithKey f k (IntMap m)
221336 (v', !m') -> ((unsafeKey k,) <$> v',)
222337 $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset
223338
339+-- | Given two maps proven to have the same keys, for each key apply the
340+-- function to the associated values and the key, to obtain a new map with the
341+-- same keys.
342+zipWithKey
343+ :: forall s a b c. (Key s -> a -> b -> c)
344+ -> IntMap s a
345+ -> IntMap s b
346+ -> IntMap s c
347+zipWithKey f (IntMap m1) (IntMap m2) = IntMap
348+ $ IntMap.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
349+ (\m -> if IntMap.null m
350+ then IntMap.empty
351+ else error "zipWithKey: bug: Data.IntMap.Refined has been subverted")
352+ (\m -> if IntMap.null m
353+ then IntMap.empty
354+ else error "zipWithKey: bug: Data.IntMap.Refined has been subverted")
355+ -- ^ Work around https://github.com/haskell/containers/issues/979
356+ m1
357+ m2
358+
359+-- | Return the union of two maps, with a given combining function for keys that
360+-- exist in both maps simultaneously.
361+unionWith
362+ :: forall s t a. (a -> a -> a)
363+ -> IntMap s a
364+ -> IntMap t a
365+ -> SomeIntMapWith (UnionProof 'Int s t) a
366+unionWith f (IntMap m1) (IntMap m2) = SomeIntMapWith
367+ (IntMap $ IntMap.unionWith f m1 m2)
368+ $ UnionProof unsafeSubset unsafeSubsetWith2
369+
224370-- | Return the union of two maps, with a given combining function for keys that
225371-- exist in both maps simultaneously.
226372--
@@ -235,8 +381,19 @@ unionWithKey f (IntMap m1) (IntMap m2)
235381 = SomeIntMapWith (IntMap $ IntMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
236382 $ UnionProof unsafeSubset unsafeSubsetWith2
237383
238−-- | For keys that appear in both maps, the given function decides whether the
239−-- key is removed from the first map.
384+-- | Return the first map, but for keys that appear in both maps, the given
385+-- function decides whether the key is removed.
386+differenceWith
387+ :: forall s t a b. (a -> b -> Maybe a)
388+ -> IntMap s a
389+ -> IntMap t b
390+ -> SomeIntMapWith (PartialDifferenceProof 'Int s t) a
391+differenceWith f (IntMap m1) (IntMap m2) = SomeIntMapWith
392+ (IntMap $ IntMap.differenceWith f m1 m2)
393+ $ PartialDifferenceProof unsafeSubset unsafeSubset
394+
395+-- | Return the first map, but for keys that appear in both maps, the given
396+-- function decides whether the key is removed.
240397--
241398-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
242399-- respectively.
@@ -252,6 +409,16 @@ differenceWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
252409 (IntMap $ IntMap.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
253410 $ PartialDifferenceProof unsafeSubset unsafeSubset
254411
412+-- | Return the intersection of two maps with the given combining function.
413+intersectionWith
414+ :: forall s t a b c. (a -> b -> c)
415+ -> IntMap s a
416+ -> IntMap t b
417+ -> SomeIntMapWith (IntersectionProof 'Int s t) c
418+intersectionWith f (IntMap m1) (IntMap m2) = SomeIntMapWith
419+ (IntMap $ IntMap.intersectionWith f m1 m2)
420+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
421+
255422-- | Return the intersection of two maps with the given combining function.
256423--
257424-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
@@ -268,6 +435,10 @@ intersectionWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
268435 (IntMap $ IntMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
269436 $ IntersectionProof unsafeSubset unsafeSubsetWith2
270437
438+-- | Apply a function to all values in a map. The set of keys remains the same.
439+map :: forall s a b. (a -> b) -> IntMap s a -> IntMap s b
440+map = coerce $ IntMap.map @a @b
441+
271442-- | Thread an accumularing argument through the map in ascending order of keys.
272443mapAccumLWithKey
273444 :: forall s a b c. (a -> Key s -> b -> (a, c))
@@ -287,6 +458,27 @@ mapAccumRWithKey
287458mapAccumRWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
288459 $ IntMap.mapAccumRWithKey @a @b @c
289460
461+-- | @'mapKeys' f m@ applies @f@ to each key of @m@ and collects the results
462+-- into a new map. For keys that were mapped to the same new key, the value
463+-- corresponding to the greatest of the original keys is retained.
464+mapKeys
465+ :: forall s a. (Key s -> Int)
466+ -> IntMap s a
467+ -> SomeIntMapWith (MapProof 'Int s Int Int) a
468+mapKeys g (IntMap m)
469+ = SomeIntMapWith (IntMap $ IntMap.mapKeys (g . unsafeKey) m)
470+ $ MapProof (unsafeKey . g) \k2 ->
471+ case IntMap.lookup (unrefine k2) backMap of
472+ Nothing -> error
473+ "mapKeys: bug: Data.IntMap.Refined has been subverted"
474+ Just k1 -> k1
475+ where
476+ ~backMap = IntMap.fromList
477+ [ (k2, unsafeKey k1)
478+ | k1 <- IntMap.keys m
479+ , let !k2 = g $ unsafeKey k1
480+ ]
481+
290482-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
291483-- results into a new map. For keys that were mapped to the same new key, @c@
292484-- acts as the combining function for corresponding values.
@@ -309,6 +501,15 @@ mapKeysWith f g (IntMap m)
309501 , let !k2 = g $ unsafeKey k1
310502 ]
311503
504+-- | Apply a function to all values in a map and collect only the 'Just'
505+-- results, returning a potentially smaller map.
506+mapMaybe
507+ :: forall s a b. (a -> Maybe b)
508+ -> IntMap s a
509+ -> SomeIntMapWith (SupersetProof 'Int s) b
510+mapMaybe f (IntMap m) = SomeIntMapWith (IntMap $ IntMap.mapMaybe f m)
511+ $ SupersetProof unsafeSubset
512+
312513-- | Apply a function to all values in a map, together with their corresponding
313514-- keys, and collect only the 'Just' results, returning a potentially smaller
314515-- map.
@@ -320,6 +521,23 @@ mapMaybeWithKey f (IntMap m)
320521 = SomeIntMapWith (IntMap $ IntMap.mapMaybeWithKey (f . unsafeKey) m)
321522 $ SupersetProof unsafeSubset
322523
524+-- | Apply a function to all values in a map and collect the 'Left' and 'Right'
525+-- results into separate (disjoint) maps.
526+mapEither
527+ :: forall s a b c. (a -> Either b c)
528+ -> IntMap s a
529+ -> Some2IntMapWith (PartitionProof 'Int s Int) b c
530+mapEither p (IntMap m)
531+ = case IntMap.mapEither p m of
532+ (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof
533+ do \k -> case IntMap.lookup (unrefine k) m of
534+ Nothing -> error
535+ "mapEither: bug: Data.IntMap.Refined has been subverted"
536+ Just x -> case p x of
537+ Left _ -> Left $ unsafeKey $ unrefine k
538+ Right _ -> Right $ unsafeKey $ unrefine k
539+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
540+
323541-- | Apply a function to all values in a map, together with their corresponding
324542-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
325543-- maps.
src/Data/IntMap/Strict/Refined.hs+229−9
@@ -34,18 +34,26 @@ module Data.IntMap.Strict.Refined
3434 , SingletonProof(..)
3535 , fromSet
3636 , Common.fromIntMap
37+ , Common.verifyIntMap
38+ , fromTraversable
39+ , fromTraversableWith
3740 , fromTraversableWithKey
3841 , FromTraversableProof(..)
3942 -- * Insertion
4043 , insert
44+ , insertWith
45+ , insertWithKey
4146 , InsertProof(..)
4247 , reinsert
4348 , insertLookupWithKey
4449 -- * Deletion/Update
4550 , Common.delete
51+ , adjust'
4652 , adjust
4753 , adjustWithKey
54+ , update'
4855 , update
56+ , updateWithKey
4957 , updateLookupWithKey
5058 -- * Query
5159 , Common.lookup
@@ -61,21 +69,29 @@ module Data.IntMap.Strict.Refined
6169 , Common.disjoint
6270 , DisjointProof(..)
6371 -- * Combine
72+ , zipWith
6473 , zipWithKey
6574 , bind
75+ , Common.union
76+ , unionWith
6677 , unionWithKey
6778 , UnionProof(..)
6879 , Common.difference
6980 , DifferenceProof(..)
81+ , differenceWith
7082 , differenceWithKey
7183 , PartialDifferenceProof(..)
84+ , Common.intersection
85+ , intersectionWith
7286 , intersectionWithKey
7387 , IntersectionProof(..)
7488 -- * Traversal
89+ , map
7590 , mapWithKey
7691 , traverseWithKey
7792 , mapAccumLWithKey
7893 , mapAccumRWithKey
94+ , mapKeys
7995 , mapKeysWith
8096 , MapProof(..)
8197 , backpermuteKeys
@@ -93,12 +109,19 @@ module Data.IntMap.Strict.Refined
93109 -- * Filter
94110 , Common.restrictKeys
95111 , Common.withoutKeys
112+ , Common.filter
113+ , Common.filterKeys
96114 , Common.filterWithKey
115+ , Common.partition
97116 , Common.partitionWithKey
98117 , PartitionProof(..)
118+ , Common.takeWhileAntitone
119+ , Common.dropWhileAntitone
99120 , Common.spanAntitone
100121 , PartialPartitionProof(..)
122+ , mapMaybe
101123 , mapMaybeWithKey
124+ , mapEither
102125 , mapEitherWithKey
103126 , Common.splitLookup
104127 , SplitProof(..)
@@ -129,7 +152,7 @@ import Data.Proxy
129152import Data.Reflection
130153import Data.Traversable
131154import Data.Type.Coercion
132−import Prelude hiding (lookup, null)
155+import Prelude hiding (lookup, map, null, zipWith)
133156import Refined
134157import Refined.Unsafe
135158
@@ -145,17 +168,52 @@ singleton k v = SomeIntMapWith (IntMap $ IntMap.singleton k v)
145168fromSet :: forall s a. KnownIntSet s => (Key s -> a) -> IntMap s a
146169fromSet f = IntMap $ IntMap.fromSet (f . unsafeKey) (reflect $ Proxy @s)
147170
148−-- | Create a map from an arbitrary traversable of key-value pairs.
171+-- | Create a map from an arbitrary traversable of key-value pairs. If a key is
172+-- repeated, the retained value is the last one in traversal order. If you're
173+-- looking for @fromList@, this is the function you want.
174+fromTraversable
175+ :: forall t a. Traversable t
176+ => t (Int, a) -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
177+fromTraversable xs = SomeIntMapWith (IntMap m) $ FromTraversableProof proof
178+ where
179+ (m, proof) = mapAccumL
180+ (\s (k, v) -> let !s' = IntMap.insert k v s in (s', unsafeKey k))
181+ IntMap.empty
182+ xs
183+
184+-- | Create a map from an arbitrary traversable of key-value pairs, with a
185+-- function for combining values for repeated keys. The function is called as if
186+-- by 'foldl1', but flipped:
187+--
188+-- @
189+-- 'fromTraversableWith' f [(k, x1), (k, x2), (k, x3)]
190+-- = 'singleton' k (f x3 (f x2 x1))
191+-- @
192+fromTraversableWith
193+ :: forall t a. Traversable t
194+ => (a -> a -> a)
195+ -> t (Int, a)
196+ -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
197+fromTraversableWith f xs
198+ = SomeIntMapWith (IntMap m) $ FromTraversableProof proof
199+ where
200+ (m, proof) = mapAccumL
201+ (\s (k, v) -> let !s' = IntMap.insertWith f k v s in (s', unsafeKey k))
202+ IntMap.empty
203+ xs
204+
205+-- | Create a map from an arbitrary traversable of key-value pairs. Like
206+-- 'fromTraversableWith', but the combining function has access to the key.
149207fromTraversableWithKey
150208 :: forall t a. Traversable t
151209 => (Int -> a -> a -> a)
152210 -> t (Int, a)
153211 -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
154−fromTraversableWithKey f xs = SomeIntMapWith (IntMap m)
155− $ FromTraversableProof proof
212+fromTraversableWithKey f xs
213+ = SomeIntMapWith (IntMap m) $ FromTraversableProof proof
156214 where
157215 (m, proof) = mapAccumL
158− (\s (k, v) -> (IntMap.insertWithKey f k v s, unsafeKey k))
216+ (\s (k, v) -> let !s' = IntMap.insertWithKey f k v s in (s', unsafeKey k))
159217 IntMap.empty
160218 xs
161219
@@ -170,8 +228,39 @@ insert
170228insert k v (IntMap m) = SomeIntMapWith (IntMap $ IntMap.insert k v m)
171229 $ InsertProof (unsafeKey k) unsafeSubset
172230
231+-- | Insert a key-value pair into the map to obtain a potentially larger map,
232+-- guaranteed to contain the given key. If the key was already present, the
233+-- supplied function is used to combine the new value with the old (in that
234+-- order).
235+insertWith
236+ :: forall s a. (a -> a -> a)
237+ -> Int
238+ -> a
239+ -> IntMap s a
240+ -> SomeIntMapWith (InsertProof 'Int Int s) a
241+insertWith f k v (IntMap m) = SomeIntMapWith
242+ (IntMap $ IntMap.insertWith f k v m)
243+ $ InsertProof (unsafeKey k) unsafeSubset
244+
245+-- | Insert a key-value pair into the map to obtain a potentially larger map,
246+-- guaranteed to contain the given key. Like 'insertWith', but the combining
247+-- function has access to the key, which is guaranteed to be in the old map.
248+insertWithKey
249+ :: forall s a. (Key s -> a -> a -> a)
250+ -> Int
251+ -> a
252+ -> IntMap s a
253+ -> SomeIntMapWith (InsertProof 'Int Int s) a
254+insertWithKey f k v (IntMap m) = SomeIntMapWith
255+ (IntMap $ IntMap.insertWithKey (f . unsafeKey) k v m)
256+ $ InsertProof (unsafeKey k) unsafeSubset
257+
173258-- | Overwrite a key-value pair that is known to already be in the map. The set
174259-- of keys remains the same.
260+--
261+-- @
262+-- 'reinsert' k v = 'adjust (const v) k'
263+-- @
175264reinsert
176265 :: forall s a. Key s -> a -> IntMap s a -> IntMap s a
177266reinsert = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.insert @a
@@ -190,10 +279,15 @@ insertLookupWithKey f k v (IntMap m)
190279 (v', !m') -> ((unsafeKey k,) <$> v',)
191280 $ SomeIntMapWith (IntMap m') $ InsertProof (unsafeKey k) unsafeSubset
192281
282+-- | If the given key is in the map, update the value at that key using the
283+-- given function. In any case, the set of keys remains the same.
284+adjust' :: forall s a. (a -> a) -> Int -> IntMap s a -> IntMap s a
285+adjust' = coerce $ IntMap.adjust @a
286+
193287-- | Update the value at a specific key known the be in the map using the given
194288-- function. The set of keys remains the same.
195289adjust :: forall s a. (a -> a) -> Key s -> IntMap s a -> IntMap s a
196−adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.adjust @a
290+adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ adjust' @s @a
197291
198292-- | If the given key is in the map, update the associated value using the given
199293-- function with a proof that the key was in the map; otherwise return the map
@@ -203,6 +297,16 @@ adjustWithKey
203297adjustWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
204298 $ IntMap.adjustWithKey @a
205299
300+-- | If a key is present in the map, update its value or delete it using the
301+-- given function, returning a potentially smaller map.
302+update'
303+ :: forall s a. (a -> Maybe a)
304+ -> Int
305+ -> IntMap s a
306+ -> SomeIntMapWith (SupersetProof 'Int s) a
307+update' f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f k m)
308+ $ SupersetProof unsafeSubset
309+
206310-- | Update or delete a key known to be in the map using the given function,
207311-- returning a potentially smaller map.
208312update
@@ -210,7 +314,18 @@ update
210314 -> Key s
211315 -> IntMap s a
212316 -> SomeIntMapWith (SupersetProof 'Int s) a
213−update f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f (unrefine k) m)
317+update = gcoerceWith (unsafeCastKey @s) $ coerce $ update' @s @a
318+
319+-- | If a key is present in the map, update its value or delete it using the
320+-- given function with a proof that the key was in the map, returning a
321+-- potentially smaller map.
322+updateWithKey
323+ :: forall s a. (Key s -> a -> Maybe a)
324+ -> Int
325+ -> IntMap s a
326+ -> SomeIntMapWith (SupersetProof 'Int s) a
327+updateWithKey f k (IntMap m) = SomeIntMapWith
328+ (IntMap $ IntMap.updateWithKey (f . unsafeKey) k m)
214329 $ SupersetProof unsafeSubset
215330
216331-- | If the given key is in the map, update or delete it using the given
@@ -229,6 +344,28 @@ updateLookupWithKey f k (IntMap m)
229344
230345-- | Given two maps proven to have the same keys, for each key apply the
231346-- function to the associated values, to obtain a new map with the same keys.
347+zipWith
348+ :: forall s a b c. (a -> b -> c)
349+ -> IntMap s a
350+ -> IntMap s b
351+ -> IntMap s c
352+zipWith f (IntMap m1) (IntMap m2) = IntMap
353+ $ IntMap.mergeWithKey (\_ x y -> Just $ f x y)
354+ (\m -> if IntMap.null m
355+ then IntMap.empty
356+ else
357+ error "zipWith: bug: Data.IntMap.Strict.Refined has been subverted")
358+ (\m -> if IntMap.null m
359+ then IntMap.empty
360+ else
361+ error "zipWith: bug: Data.IntMap.Strict.Refined has been subverted")
362+ -- ^ Work around https://github.com/haskell/containers/issues/979
363+ m1
364+ m2
365+
366+-- | Given two maps proven to have the same keys, for each key apply the
367+-- function to the associated values and the key, to obtain a new map with the
368+-- same keys.
232369zipWithKey
233370 :: forall s a b c. (Key s -> a -> b -> c)
234371 -> IntMap s a
@@ -248,6 +385,17 @@ zipWithKey f (IntMap m1) (IntMap m2) = IntMap
248385 m1
249386 m2
250387
388+-- | Return the union of two maps, with a given combining function for keys that
389+-- exist in both maps simultaneously.
390+unionWith
391+ :: forall s t a. (a -> a -> a)
392+ -> IntMap s a
393+ -> IntMap t a
394+ -> SomeIntMapWith (UnionProof 'Int s t) a
395+unionWith f (IntMap m1) (IntMap m2) = SomeIntMapWith
396+ (IntMap $ IntMap.unionWith f m1 m2)
397+ $ UnionProof unsafeSubset unsafeSubsetWith2
398+
251399-- | Return the union of two maps, with a given combining function for keys that
252400-- exist in both maps simultaneously.
253401--
@@ -262,8 +410,19 @@ unionWithKey f (IntMap m1) (IntMap m2)
262410 = SomeIntMapWith (IntMap $ IntMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
263411 $ UnionProof unsafeSubset unsafeSubsetWith2
264412
265−-- | For keys that appear in both maps, the given function decides whether the
266−-- key is removed from the first map.
413+-- | Return the first map, but for keys that appear in both maps, the given
414+-- function decides whether the key is removed.
415+differenceWith
416+ :: forall s t a b. (a -> b -> Maybe a)
417+ -> IntMap s a
418+ -> IntMap t b
419+ -> SomeIntMapWith (PartialDifferenceProof 'Int s t) a
420+differenceWith f (IntMap m1) (IntMap m2) = SomeIntMapWith
421+ (IntMap $ IntMap.differenceWith f m1 m2)
422+ $ PartialDifferenceProof unsafeSubset unsafeSubset
423+
424+-- | Return the first map, but for keys that appear in both maps, the given
425+-- function decides whether the key is removed.
267426--
268427-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
269428-- respectively.
@@ -279,6 +438,16 @@ differenceWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
279438 (IntMap $ IntMap.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
280439 $ PartialDifferenceProof unsafeSubset unsafeSubset
281440
441+-- | Return the intersection of two maps with the given combining function.
442+intersectionWith
443+ :: forall s t a b c. (a -> b -> c)
444+ -> IntMap s a
445+ -> IntMap t b
446+ -> SomeIntMapWith (IntersectionProof 'Int s t) c
447+intersectionWith f (IntMap m1) (IntMap m2) = SomeIntMapWith
448+ (IntMap $ IntMap.intersectionWith f m1 m2)
449+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
450+
282451-- | Return the intersection of two maps with the given combining function.
283452--
284453-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
@@ -295,6 +464,10 @@ intersectionWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
295464 (IntMap $ IntMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
296465 $ IntersectionProof unsafeSubset unsafeSubsetWith2
297466
467+-- | Apply a function to all values in a map. The set of keys remains the same.
468+map :: forall s a b. (a -> b) -> IntMap s a -> IntMap s b
469+map = coerce $ IntMap.map @a @b
470+
298471-- | Apply a function to all values in a map, together with their corresponding
299472-- keys, that are proven to be in the map. The set of keys remains the same.
300473mapWithKey :: forall s a b. (Key s -> a -> b) -> IntMap s a -> IntMap s b
@@ -328,6 +501,27 @@ mapAccumRWithKey
328501mapAccumRWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
329502 $ IntMap.mapAccumRWithKey @a @b @c
330503
504+-- | @'mapKeys' f m@ applies @f@ to each key of @m@ and collects the results
505+-- into a new map. For keys that were mapped to the same new key, the value
506+-- corresponding to the greatest of the original keys is retained.
507+mapKeys
508+ :: forall s a. (Key s -> Int)
509+ -> IntMap s a
510+ -> SomeIntMapWith (MapProof 'Int s Int Int) a
511+mapKeys g (IntMap m)
512+ = SomeIntMapWith (IntMap $ IntMap.mapKeys (g . unsafeKey) m)
513+ $ MapProof (unsafeKey . g) \k2 ->
514+ case IntMap.lookup (unrefine k2) backMap of
515+ Nothing -> error
516+ "mapKeys: bug: Data.IntMap.Strict.Refined has been subverted"
517+ Just k1 -> k1
518+ where
519+ ~backMap = IntMap.fromList
520+ [ (k2, unsafeKey k1)
521+ | k1 <- IntMap.keys m
522+ , let !k2 = g $ unsafeKey k1
523+ ]
524+
331525-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
332526-- results into a new map. For keys that were mapped to the same new key, @c@
333527-- acts as the combining function for corresponding values.
@@ -350,6 +544,15 @@ mapKeysWith f g (IntMap m)
350544 , let !k2 = g $ unsafeKey k1
351545 ]
352546
547+-- | Apply a function to all values in a map and collect only the 'Just'
548+-- results, returning a potentially smaller map.
549+mapMaybe
550+ :: forall s a b. (a -> Maybe b)
551+ -> IntMap s a
552+ -> SomeIntMapWith (SupersetProof 'Int s) b
553+mapMaybe f (IntMap m) = SomeIntMapWith (IntMap $ IntMap.mapMaybe f m)
554+ $ SupersetProof unsafeSubset
555+
353556-- | Apply a function to all values in a map, together with their corresponding
354557-- keys, and collect only the 'Just' results, returning a potentially smaller
355558-- map.
@@ -361,6 +564,23 @@ mapMaybeWithKey f (IntMap m)
361564 = SomeIntMapWith (IntMap $ IntMap.mapMaybeWithKey (f . unsafeKey) m)
362565 $ SupersetProof unsafeSubset
363566
567+-- | Apply a function to all values in a map and collect the 'Left' and 'Right'
568+-- results into separate (disjoint) maps.
569+mapEither
570+ :: forall s a b c. (a -> Either b c)
571+ -> IntMap s a
572+ -> Some2IntMapWith (PartitionProof 'Int s Int) b c
573+mapEither p (IntMap m)
574+ = case IntMap.mapEither p m of
575+ (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof
576+ do \k -> case IntMap.lookup (unrefine k) m of
577+ Nothing -> error
578+ "mapEither: bug: Data.IntMap.Refined has been subverted"
579+ Just x -> case p x of
580+ Left _ -> Left $ unsafeKey $ unrefine k
581+ Right _ -> Right $ unsafeKey $ unrefine k
582+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
583+
364584-- | Apply a function to all values in a map, together with their corresponding
365585-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
366586-- maps.
src/Data/IntSet/Refined.hs+1−1
@@ -169,7 +169,7 @@ data SomeIntSet where
169169-- you a way to refer to the set (the parameter @s@), e.g.:
170170--
171171-- @
172−-- 'withIntSet' ('fromIntSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s
172+-- 'withIntSet' ('fromIntSet' ...) $ \\(_ :: 'Proxy' s) -> doSomethingWith \@s
173173-- @
174174withIntSet
175175 :: forall r. SomeIntSet -> (forall s. KnownIntSet s => Proxy s -> r) -> r
src/Data/Map/Common/Refined.hs+141−15
@@ -16,13 +16,19 @@ import qualified Data.Hashable as Hashable
1616import qualified Data.Map as Map
1717import Data.Proxy
1818import Data.Reflection
19+import qualified Data.Set as Set
1920import Data.Traversable.WithIndex
2021import Data.Type.Coercion
2122import Data.Type.Equality ((:~:)(..))
23+import Prelude hiding (zipWith)
2224import Refined
2325import Refined.Unsafe
2426import Unsafe.Coerce
2527
28+#if MIN_VERSION_containers(0, 6, 3)
29+import Data.Bifoldable
30+#endif
31+
2632#if MIN_VERSION_containers(0, 6, 2)
2733#elif MIN_VERSION_containers(0, 5, 8)
2834import Data.Functor.Const (Const(..))
@@ -33,6 +39,10 @@ import qualified Data.List as List
3339import qualified Data.Map.Strict as MapStrict
3440#endif
3541
42+#if MIN_VERSION_containers(0, 5, 9) || MIN_VERSION_hashable(1, 4, 0)
43+import Data.Functor.Classes
44+#endif
45+
3646
3747-- | A wrapper around a regular 'Data.Map.Map' with a type parameter @s@
3848-- identifying the set of keys present in the map.
@@ -49,6 +59,12 @@ newtype Map s k a = Map (Map.Map k a)
4959 deriving newtype (Eq, Ord, Show, Functor, Foldable, NFData)
5060#if MIN_VERSION_hashable(1, 3, 4)
5161 deriving newtype (Hashable.Hashable)
62+#endif
63+#if MIN_VERSION_containers(0, 6, 3)
64+ deriving newtype (Bifoldable)
65+#endif
66+#if MIN_VERSION_containers(0, 5, 9) || MIN_VERSION_hashable(1, 4, 0)
67+ deriving newtype (Eq1, Eq2, Ord1, Ord2, Show1, Show2)
5268#endif
5369 deriving stock (Traversable)
5470type role Map nominal nominal representational
@@ -94,7 +110,7 @@ data SomeMap k a where
94110-- parameter @s@), e.g.:
95111--
96112-- @
97−-- 'withMap' ('fromMap' ...) $ \(m :: 'Map' s k a) -> doSomethingWith \@s
113+-- 'withMap' ('fromMap' ...) $ \\(m :: 'Map' s k a) -> doSomethingWith \@s
98114-- @
99115withMap :: forall k a r. SomeMap k a -> (forall s. Map s k a -> r) -> r
100116withMap (SomeMap m) k = k m
@@ -103,6 +119,14 @@ withMap (SomeMap m) k = k m
103119fromMap :: forall k a. Map.Map k a -> SomeMap k a
104120fromMap m = SomeMap (Map m)
105121
122+-- | Given a set of keys @s@ known ahead of time, verify whether a regular
123+-- 'Data.Map.Map' has exactly that set of keys.
124+verifyMap
125+ :: forall s k a. (Eq k, KnownSet s k) => Map.Map k a -> Maybe (Map s k a)
126+verifyMap m
127+ | Map.keys m == Set.toList (reflect $ Proxy @s) = Just (Map m)
128+ | otherwise = Nothing
129+
106130-- | An existential wrapper for a 'Map' with an as-yet-unknown set of keys,
107131-- together with a proof of some fact @p@ about the set. Pattern matching on it
108132-- gives you a way to refer to the set (the parameter @s@). Functions that
@@ -236,21 +260,33 @@ disjoint (Map m1) (Map m2)
236260
237261-- | Given two maps proven to have the same keys, for each key apply the
238262-- function to the associated values, to obtain a new map with the same keys.
239−zipWithKey
263+zipWith
240264 :: forall s k a b c. Ord k
241− => (Key s k -> a -> b -> c) -> Map s k a -> Map s k b -> Map s k c
242−zipWithKey f (Map m1) (Map m2) = Map
243− $ Map.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
265+ => (a -> b -> c) -> Map s k a -> Map s k b -> Map s k c
266+zipWith f (Map m1) (Map m2) = Map
267+ $ Map.mergeWithKey (\_ x y -> Just $ f x y)
244268 (\m -> if Map.null m
245269 then Map.empty
246− else error "zipWithKey: bug: Data.Map.Refined has been subverted")
270+ else error "zipWith: bug: Data.Map.Refined has been subverted")
247271 (\m -> if Map.null m
248272 then Map.empty
249− else error "zipWithKey: bug: Data.Map.Refined has been subverted")
273+ else error "zipWith: bug: Data.Map.Refined has been subverted")
250274 -- ^ Work around https://github.com/haskell/containers/issues/979
251275 m1
252276 m2
253277
278+-- | Return the union of two maps. For keys that exist in both maps, the value
279+-- is taken from the first map.
280+--
281+-- @
282+-- 'union' = unionWith 'const'
283+-- @
284+union
285+ :: forall s t k a. Ord k
286+ => Map s k a -> Map t k a -> SomeMapWith (UnionProof 'Regular s t) k a
287+union (Map m1) (Map m2) = SomeMapWith (Map $ Map.union m1 m2)
288+ $ UnionProof unsafeSubset unsafeSubsetWith2
289+
254290-- | Remove the keys that appear in the second map from the first map.
255291difference
256292 :: forall s t k a b. Ord k
@@ -258,6 +294,20 @@ difference
258294difference (Map m1) (Map m2) = SomeMapWith (Map $ Map.difference m1 m2)
259295 $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
260296
297+-- | Return the intersection of two maps, taking values from the first map.
298+--
299+-- @
300+-- 'intersection' = intersectionWith 'const'
301+-- @
302+intersection
303+ :: forall s t k a b. Ord k
304+ => Map s k a
305+ -> Map t k b
306+ -> SomeMapWith (IntersectionProof 'Regular s t) k a
307+intersection (Map m1) (Map m2)
308+ = SomeMapWith (Map $ Map.intersection m1 m2)
309+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
310+
261311-- | Apply a function to all values in a map, together with their corresponding
262312-- keys, that are proven to be in the map. The set of keys remains the same.
263313mapWithKey :: forall s k a b. (Key s k -> a -> b) -> Map s k a -> Map s k b
@@ -314,6 +364,29 @@ toList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.toAscList @k @a
314364toDescList :: forall s k a. Map s k a -> [(Key s k, a)]
315365toDescList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.toDescList @k @a
316366
367+-- | Retain only the values that satisfy the predicate, returning a potentially
368+-- smaller map.
369+filter
370+ :: forall s k a. (a -> Bool)
371+ -> Map s k a
372+ -> SomeMapWith (SupersetProof 'Regular s) k a
373+filter p (Map m) = SomeMapWith (Map $ Map.filter p m)
374+ $ SupersetProof unsafeSubset
375+
376+-- | Retain only the keys that satisfy the predicate, returning a potentially
377+-- smaller map.
378+filterKeys
379+ :: forall s k a. (Key s k -> Bool)
380+ -> Map s k a
381+ -> SomeMapWith (SupersetProof 'Regular s) k a
382+filterKeys p (Map m) = SomeMapWith
383+#if MIN_VERSION_containers(0, 8, 0)
384+ (Map $ Map.filterKeys (p . unsafeKey) m)
385+#else
386+ (Map $ Map.filterWithKey (\k _ -> p (unsafeKey k)) m)
387+#endif
388+ $ SupersetProof unsafeSubset
389+
317390-- | Retain only the key-value pairs that satisfy the predicate, returning a
318391-- potentially smaller map.
319392filterWithKey
@@ -348,6 +421,22 @@ withoutKeys (Map m) = SomeMapWith
348421#endif
349422 $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
350423
424+-- | Partition a map into two disjoint submaps: those whose values satisfy the
425+-- predicate, and those whose don't.
426+partition
427+ :: forall s k a. Ord k -- TODO: this is only used in the proof
428+ => (a -> Bool)
429+ -> Map s k a
430+ -> Some2MapWith (PartitionProof 'Regular s k) k a a
431+partition p (Map m) = case Map.partition p m of
432+ (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof
433+ do \k -> case Map.lookup (unrefine k) m of
434+ Nothing -> error "partition: bug: Data.Map.Refined has been subverted"
435+ Just x -> if p x
436+ then Left $ unsafeKey $ unrefine k
437+ else Right $ unsafeKey $ unrefine k
438+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
439+
351440-- | Partition a map into two disjoint submaps: those whose key-value pairs
352441-- satisfy the predicate, and those whose don't.
353442partitionWithKey
@@ -365,6 +454,43 @@ partitionWithKey p (Map m) = case Map.partitionWithKey (p . unsafeKey) m of
365454 else Right $ unsafeKey $ unrefine k
366455 unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
367456
457+-- | Take the a submap of keys up to a point where the predicate stops holding.
458+--
459+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
460+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
461+-- chosen in an unspecified way.
462+takeWhileAntitone
463+ :: forall s k a. (Key s k -> Bool)
464+ -> Map s k a
465+ -> SomeMapWith (SupersetProof 'Regular s) k a
466+takeWhileAntitone p (Map m) = SomeMapWith
467+#if MIN_VERSION_containers(0, 5, 8)
468+ (Map $ Map.takeWhileAntitone (p . unsafeKey) m)
469+#else
470+ (Map $ Map.fromDistinctAscList
471+ $ List.takeWhile (p . unsafeKey . fst) $ Map.toAscList m)
472+#endif
473+ $ SupersetProof unsafeSubset
474+
475+-- | Take the a submap of keys starting from a point where the predicate stops
476+-- holding.
477+--
478+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
479+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
480+-- chosen in an unspecified way.
481+dropWhileAntitone
482+ :: forall s k a. (Key s k -> Bool)
483+ -> Map s k a
484+ -> SomeMapWith (SupersetProof 'Regular s) k a
485+dropWhileAntitone p (Map m) = SomeMapWith
486+#if MIN_VERSION_containers(0, 5, 8)
487+ (Map $ Map.dropWhileAntitone (p . unsafeKey) m)
488+#else
489+ (Map $ Map.fromDistinctAscList
490+ $ List.dropWhile (p . unsafeKey . fst) $ Map.toAscList m)
491+#endif
492+ $ SupersetProof unsafeSubset
493+
368494-- | Divide a map into two disjoint submaps at a point where the predicate on
369495-- the keys stops holding.
370496--
@@ -437,15 +563,15 @@ castKey
437563castKey = castRefined
438564
439565-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps
440−-- can be interconverted too. For example, 'zipWithKey' can be implemented via
441−-- 'Data.Map.Refined.intersectionWithKey' by proving that the set of keys
442−-- remains unchanged:
566+-- can be interconverted too. For example, 'Data.Map.Refined.zipWithKey' can be
567+-- implemented via 'Data.Map.Refined.intersectionWithKey' by proving that the
568+-- set of keys remains unchanged:
443569--
444570-- @
445−-- 'zipWithKey'
571+-- 'Data.Map.Refined.zipWithKey'
446572-- :: forall s k a b c. 'Ord' k
447573-- => ('Key' s k -> a -> b -> c) -> 'Map' s k a -> 'Map' s k b -> 'Map' s k c
448−-- 'zipWithKey' f m1 m2
574+-- 'Data.Map.Refined.zipWithKey' f m1 m2
449575-- | v'SomeMapWith' @r m proof <- 'Data.Map.Refined.intersectionWithKey' (f . 'andLeft') m1 m2
450576-- , v'IntersectionProof' p1 p2 <- proof
451577-- , ( v'Coercion' :: t'Coercion' ('Map' r k c) ('Map' s k c))
@@ -470,10 +596,10 @@ instance TraversableWithIndex (Key s k) (Map s k) where
470596 itraverse = traverseWithKey
471597
472598-- | Similar to the instance for functions -- zip corresponding keys. To use
473−-- '<*>'/'Control.Applicative.liftA2' without 'KnownSet' see 'zipWithKey'.
599+-- '<*>'/'Control.Applicative.liftA2' without 'KnownSet' see 'zipWith'.
474600instance (Ord k, KnownSet s k) => Applicative (Map s k) where
475601 pure x = fromSet \_ -> x
476− (<*>) = zipWithKey (const id)
602+ (<*>) = zipWith id
477603
478604-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the
479605-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
@@ -493,7 +619,7 @@ instance (Ord k, KnownSet s k) => MonadReader (Key s k) (Map s k) where
493619
494620-- | Append the values at the corresponding keys
495621instance (Ord k, Semigroup a) => Semigroup (Map s k a) where
496− (<>) = zipWithKey (const (<>))
622+ (<>) = zipWith (<>)
497623
498624instance (Ord k, KnownSet s k, Monoid a) => Monoid (Map s k a) where
499625 mempty = fromSet \_ -> mempty
src/Data/Map/Refined.hs+226−7
@@ -43,18 +43,26 @@ module Data.Map.Refined
4343 , SingletonProof(..)
4444 , fromSet
4545 , Common.fromMap
46+ , Common.verifyMap
47+ , fromTraversable
48+ , fromTraversableWith
4649 , fromTraversableWithKey
4750 , FromTraversableProof(..)
4851 -- * Insertion
4952 , insert
53+ , insertWith
54+ , insertWithKey
5055 , InsertProof(..)
5156 , reinsert
5257 , insertLookupWithKey
5358 -- * Deletion/Update
5459 , Common.delete
60+ , adjust'
5561 , adjust
5662 , adjustWithKey
63+ , update'
5764 , update
65+ , updateWithKey
5866 , updateLookupWithKey
5967 -- * Query
6068 , Common.lookup
@@ -70,21 +78,29 @@ module Data.Map.Refined
7078 , Common.disjoint
7179 , DisjointProof(..)
7280 -- * Combine
81+ , zipWith
7382 , zipWithKey
7483 , bind
84+ , Common.union
85+ , unionWith
7586 , unionWithKey
7687 , UnionProof(..)
7788 , Common.difference
7889 , DifferenceProof(..)
90+ , differenceWith
7991 , differenceWithKey
8092 , PartialDifferenceProof(..)
93+ , Common.intersection
94+ , intersectionWith
8195 , intersectionWithKey
8296 , IntersectionProof(..)
8397 -- * Traversal
98+ , map
8499 , mapWithKey
85100 , traverseWithKey
86101 , mapAccumLWithKey
87102 , mapAccumRWithKey
103+ , mapKeys
88104 , mapKeysWith
89105 , MapProof(..)
90106 , backpermuteKeys
@@ -102,12 +118,19 @@ module Data.Map.Refined
102118 -- * Filter
103119 , Common.restrictKeys
104120 , Common.withoutKeys
121+ , Common.filter
122+ , Common.filterKeys
105123 , Common.filterWithKey
124+ , Common.partition
106125 , Common.partitionWithKey
107126 , PartitionProof(..)
127+ , Common.takeWhileAntitone
128+ , Common.dropWhileAntitone
108129 , Common.spanAntitone
109130 , PartialPartitionProof(..)
131+ , mapMaybe
110132 , mapMaybeWithKey
133+ , mapEither
111134 , mapEitherWithKey
112135 , Common.splitLookup
113136 , SplitProof(..)
@@ -131,12 +154,12 @@ import Data.Functor
131154import qualified Data.Map as Map
132155import Data.Map.Common.Refined
133156 ( Map(..), Key, unsafeCastKey, unsafeKey, SomeMapWith(..), Some2MapWith(..)
134− , fromSet, (!), zipWithKey, mapWithKey, traverseWithKey, bind
157+ , fromSet, (!), zipWith, mapWithKey, traverseWithKey, bind
135158 )
136159import qualified Data.Map.Common.Refined as Common
137160import Data.Traversable
138161import Data.Type.Coercion
139−import Prelude hiding (lookup, null)
162+import Prelude hiding (lookup, map, null, zipWith)
140163import Refined
141164import Refined.Unsafe
142165
@@ -147,7 +170,41 @@ singleton :: forall k a. k -> a -> SomeMapWith (SingletonProof 'Regular k) k a
147170singleton k v = SomeMapWith (Map $ Map.singleton k v)
148171 $ SingletonProof (unsafeKey k)
149172
150−-- | Create a map from an arbitrary traversable of key-value pairs.
173+-- | Create a map from an arbitrary traversable of key-value pairs. If a key is
174+-- repeated, the retained value is the last one in traversal order. If you're
175+-- looking for @fromList@, this is the function you want.
176+fromTraversable
177+ :: forall t k a. (Traversable t, Ord k)
178+ => t (k, a) -> SomeMapWith (FromTraversableProof 'Regular t k) k a
179+fromTraversable xs = SomeMapWith (Map m) $ FromTraversableProof proof
180+ where
181+ (m, proof) = mapAccumL
182+ (\s (k, v) -> let !s' = Map.insert k v s in (s', unsafeKey k))
183+ Map.empty
184+ xs
185+
186+-- | Create a map from an arbitrary traversable of key-value pairs, with a
187+-- function for combining values for repeated keys. The function is called as if
188+-- by 'foldl1', but flipped:
189+--
190+-- @
191+-- 'fromTraversableWith' f [(k, x1), (k, x2), (k, x3)]
192+-- = 'singleton' k (f x3 (f x2 x1))
193+-- @
194+fromTraversableWith
195+ :: forall t k a. (Traversable t, Ord k)
196+ => (a -> a -> a)
197+ -> t (k, a)
198+ -> SomeMapWith (FromTraversableProof 'Regular t k) k a
199+fromTraversableWith f xs = SomeMapWith (Map m) $ FromTraversableProof proof
200+ where
201+ (m, proof) = mapAccumL
202+ (\s (k, v) -> let !s' = Map.insertWith f k v s in (s', unsafeKey k))
203+ Map.empty
204+ xs
205+
206+-- | Create a map from an arbitrary traversable of key-value pairs. Like
207+-- 'fromTraversableWith', but the combining function has access to the key.
151208fromTraversableWithKey
152209 :: forall t k a. (Traversable t, Ord k)
153210 => (k -> a -> a -> a)
@@ -169,8 +226,40 @@ insert
169226insert k v (Map m) = SomeMapWith (Map $ Map.insert k v m)
170227 $ InsertProof (unsafeKey k) unsafeSubset
171228
229+-- | Insert a key-value pair into the map to obtain a potentially larger map,
230+-- guaranteed to contain the given key. If the key was already present, the
231+-- supplied function is used to combine the new value with the old (in that
232+-- order).
233+insertWith
234+ :: forall s k a. Ord k
235+ => (a -> a -> a)
236+ -> k
237+ -> a
238+ -> Map s k a
239+ -> SomeMapWith (InsertProof 'Regular k s) k a
240+insertWith f k v (Map m) = SomeMapWith (Map $ Map.insertWith f k v m)
241+ $ InsertProof (unsafeKey k) unsafeSubset
242+
243+-- | Insert a key-value pair into the map to obtain a potentially larger map,
244+-- guaranteed to contain the given key. Like 'insertWith', but the combining
245+-- function has access to the key, which is guaranteed to be in the old map.
246+insertWithKey
247+ :: forall s k a. Ord k
248+ => (Key s k -> a -> a -> a)
249+ -> k
250+ -> a
251+ -> Map s k a
252+ -> SomeMapWith (InsertProof 'Regular k s) k a
253+insertWithKey f k v (Map m) = SomeMapWith
254+ (Map $ Map.insertWithKey (f . unsafeKey) k v m)
255+ $ InsertProof (unsafeKey k) unsafeSubset
256+
172257-- | Overwrite a key-value pair that is known to already be in the map. The set
173258-- of keys remains the same.
259+--
260+-- @
261+-- 'reinsert' k v = 'adjust (const v) k'
262+-- @
174263reinsert
175264 :: forall s k a. Ord k
176265 => Key s k -> a -> Map s k a -> Map s k a
@@ -191,10 +280,15 @@ insertLookupWithKey f k v (Map m)
191280 (v', !m') -> ((unsafeKey k,) <$> v',)
192281 $ SomeMapWith (Map m') $ InsertProof (unsafeKey k) unsafeSubset
193282
283+-- | If the given key is in the map, update the value at that key using the
284+-- given function. In any case, the set of keys remains the same.
285+adjust' :: forall s k a. Ord k => (a -> a) -> k -> Map s k a -> Map s k a
286+adjust' = coerce $ Map.adjust @k @a
287+
194288-- | Update the value at a specific key known the be in the map using the given
195289-- function. The set of keys remains the same.
196290adjust :: forall s k a. Ord k => (a -> a) -> Key s k -> Map s k a -> Map s k a
197−adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.adjust @k @a
291+adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ adjust' @s @k @a
198292
199293-- | If the given key is in the map, update the associated value using the given
200294-- function with a proof that the key was in the map; otherwise return the map
@@ -204,6 +298,17 @@ adjustWithKey
204298adjustWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
205299 $ Map.adjustWithKey @k @a
206300
301+-- | If a key is present in the map, update its value or delete it using the
302+-- given function, returning a potentially smaller map.
303+update'
304+ :: forall s k a. Ord k
305+ => (a -> Maybe a)
306+ -> k
307+ -> Map s k a
308+ -> SomeMapWith (SupersetProof 'Regular s) k a
309+update' f k (Map m) = SomeMapWith (Map $ Map.update f k m)
310+ $ SupersetProof unsafeSubset
311+
207312-- | Update or delete a key known to be in the map using the given function,
208313-- returning a potentially smaller map.
209314update
@@ -212,7 +317,19 @@ update
212317 -> Key s k
213318 -> Map s k a
214319 -> SomeMapWith (SupersetProof 'Regular s) k a
215−update f k (Map m) = SomeMapWith (Map $ Map.update f (unrefine k) m)
320+update = gcoerceWith (unsafeCastKey @s @k) $ coerce $ update' @s @k @a
321+
322+-- | If a key is present in the map, update its value or delete it using the
323+-- given function with a proof that the key was in the map, returning a
324+-- potentially smaller map.
325+updateWithKey
326+ :: forall s k a. Ord k
327+ => (Key s k -> a -> Maybe a)
328+ -> k
329+ -> Map s k a
330+ -> SomeMapWith (SupersetProof 'Regular s) k a
331+updateWithKey f k (Map m) = SomeMapWith
332+ (Map $ Map.updateWithKey (f . unsafeKey) k m)
216333 $ SupersetProof unsafeSubset
217334
218335-- | If the given key is in the map, update or delete it using the given
@@ -230,6 +347,35 @@ updateLookupWithKey f k (Map m)
230347 (v', !m') -> ((unsafeKey k,) <$> v',)
231348 $ SomeMapWith (Map m') $ SupersetProof unsafeSubset
232349
350+-- | Given two maps proven to have the same keys, for each key apply the
351+-- function to the associated values and the key, to obtain a new map with the
352+-- same keys.
353+zipWithKey
354+ :: forall s k a b c. Ord k
355+ => (Key s k -> a -> b -> c) -> Map s k a -> Map s k b -> Map s k c
356+zipWithKey f (Map m1) (Map m2) = Map
357+ $ Map.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
358+ (\m -> if Map.null m
359+ then Map.empty
360+ else error "zipWithKey: bug: Data.Map.Refined has been subverted")
361+ (\m -> if Map.null m
362+ then Map.empty
363+ else error "zipWithKey: bug: Data.Map.Refined has been subverted")
364+ -- ^ Work around https://github.com/haskell/containers/issues/979
365+ m1
366+ m2
367+
368+-- | Return the union of two maps, with a given combining function for keys that
369+-- exist in both maps simultaneously.
370+unionWith
371+ :: forall s t k a. Ord k
372+ => (a -> a -> a)
373+ -> Map s k a
374+ -> Map t k a
375+ -> SomeMapWith (UnionProof 'Regular s t) k a
376+unionWith f (Map m1) (Map m2) = SomeMapWith (Map $ Map.unionWith f m1 m2)
377+ $ UnionProof unsafeSubset unsafeSubsetWith2
378+
233379-- | Return the union of two maps, with a given combining function for keys that
234380-- exist in both maps simultaneously.
235381--
@@ -245,8 +391,20 @@ unionWithKey f (Map m1) (Map m2)
245391 = SomeMapWith (Map $ Map.unionWithKey (f . reallyUnsafeRefine) m1 m2)
246392 $ UnionProof unsafeSubset unsafeSubsetWith2
247393
248−-- | For keys that appear in both maps, the given function decides whether the
249−-- key is removed from the first map.
394+-- | Return the first map, but for keys that appear in both maps, the given
395+-- function decides whether the key is removed.
396+differenceWith
397+ :: forall s t k a b. Ord k
398+ => (a -> b -> Maybe a)
399+ -> Map s k a
400+ -> Map t k b
401+ -> SomeMapWith (PartialDifferenceProof 'Regular s t) k a
402+differenceWith f (Map m1) (Map m2)
403+ = SomeMapWith (Map $ Map.differenceWith f m1 m2)
404+ $ PartialDifferenceProof unsafeSubset unsafeSubset
405+
406+-- | Return the first map, but for keys that appear in both maps, the given
407+-- function decides whether the key is removed.
250408--
251409-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
252410-- respectively.
@@ -260,6 +418,17 @@ differenceWithKey f (Map m1) (Map m2)
260418 = SomeMapWith (Map $ Map.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
261419 $ PartialDifferenceProof unsafeSubset unsafeSubset
262420
421+-- | Return the intersection of two maps with the given combining function.
422+intersectionWith
423+ :: forall s t k a b c. Ord k
424+ => (a -> b -> c)
425+ -> Map s k a
426+ -> Map t k b
427+ -> SomeMapWith (IntersectionProof 'Regular s t) k c
428+intersectionWith f (Map m1) (Map m2)
429+ = SomeMapWith (Map $ Map.intersectionWith f m1 m2)
430+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
431+
263432-- | Return the intersection of two maps with the given combining function.
264433--
265434-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
@@ -274,6 +443,10 @@ intersectionWithKey f (Map m1) (Map m2)
274443 = SomeMapWith (Map $ Map.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
275444 $ IntersectionProof unsafeSubset unsafeSubsetWith2
276445
446+-- | Apply a function to all values in a map. The set of keys remains the same.
447+map :: forall s k a b. (a -> b) -> Map s k a -> Map s k b
448+map = coerce $ Map.map @a @b @k
449+
277450-- | Thread an accumularing argument through the map in ascending order of keys.
278451mapAccumLWithKey
279452 :: forall s k a b c. (a -> Key s k -> b -> (a, c))
@@ -293,6 +466,26 @@ mapAccumRWithKey
293466mapAccumRWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
294467 $ Map.mapAccumRWithKey @a @k @b @c
295468
469+-- | @'mapKeys' f m@ applies @f@ to each key of @m@ and collects the results
470+-- into a new map. For keys that were mapped to the same new key, the value
471+-- corresponding to the greatest of the original keys is retained.
472+mapKeys
473+ :: forall s k1 k2 a. Ord k2
474+ => (Key s k1 -> k2)
475+ -> Map s k1 a
476+ -> SomeMapWith (MapProof 'Regular s k1 k2) k2 a
477+mapKeys g (Map m)
478+ = SomeMapWith (Map $ Map.mapKeys (g . unsafeKey) m)
479+ $ MapProof (unsafeKey . g) \k2 -> case Map.lookup (unrefine k2) backMap of
480+ Nothing -> error "mapKeys: bug: Data.Map.Refined has been subverted"
481+ Just k1 -> k1
482+ where
483+ ~backMap = Map.fromList
484+ [ (k2, unsafeKey k1)
485+ | k1 <- Map.keys m
486+ , let !k2 = g $ unsafeKey k1
487+ ]
488+
296489-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
297490-- results into a new map. For keys that were mapped to the same new key, @c@
298491-- acts as the combining function for corresponding values.
@@ -314,6 +507,15 @@ mapKeysWith f g (Map m)
314507 , let !k2 = g $ unsafeKey k1
315508 ]
316509
510+-- | Apply a function to all values in a map and collect only the 'Just'
511+-- results, returning a potentially smaller map.
512+mapMaybe
513+ :: forall s k a b. (a -> Maybe b)
514+ -> Map s k a
515+ -> SomeMapWith (SupersetProof 'Regular s) k b
516+mapMaybe f (Map m) = SomeMapWith (Map $ Map.mapMaybe f m)
517+ $ SupersetProof unsafeSubset
518+
317519-- | Apply a function to all values in a map, together with their corresponding
318520-- keys, and collect only the 'Just' results, returning a potentially smaller
319521-- map.
@@ -325,6 +527,23 @@ mapMaybeWithKey f (Map m)
325527 = SomeMapWith (Map $ Map.mapMaybeWithKey (f . unsafeKey) m)
326528 $ SupersetProof unsafeSubset
327529
530+-- | Apply a function to all values in a map and collect the 'Left' and 'Right'
531+-- results into separate (disjoint) maps.
532+mapEither
533+ :: forall s k a b c. Ord k -- TODO: this is only used in the proof
534+ => (a -> Either b c)
535+ -> Map s k a
536+ -> Some2MapWith (PartitionProof 'Regular s k) k b c
537+mapEither p (Map m) = case Map.mapEither p m of
538+ (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof
539+ do \k -> case Map.lookup (unrefine k) m of
540+ Nothing -> error
541+ "mapEither: bug: Data.Map.Refined has been subverted"
542+ Just x -> case p x of
543+ Left _ -> Left $ unsafeKey $ unrefine k
544+ Right _ -> Right $ unsafeKey $ unrefine k
545+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
546+
328547-- | Apply a function to all values in a map, together with their corresponding
329548-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
330549-- maps.
src/Data/Map/Strict/Refined.hs+226−6
@@ -43,18 +43,26 @@ module Data.Map.Strict.Refined
4343 , SingletonProof(..)
4444 , fromSet
4545 , Common.fromMap
46+ , Common.verifyMap
47+ , fromTraversable
48+ , fromTraversableWith
4649 , fromTraversableWithKey
4750 , FromTraversableProof(..)
4851 -- * Insertion
4952 , insert
53+ , insertWith
54+ , insertWithKey
5055 , InsertProof(..)
5156 , reinsert
5257 , insertLookupWithKey
5358 -- * Deletion/Update
5459 , Common.delete
60+ , adjust'
5561 , adjust
5662 , adjustWithKey
63+ , update'
5764 , update
65+ , updateWithKey
5866 , updateLookupWithKey
5967 -- * Query
6068 , Common.lookup
@@ -70,21 +78,29 @@ module Data.Map.Strict.Refined
7078 , Common.disjoint
7179 , DisjointProof(..)
7280 -- * Combine
81+ , zipWith
7382 , zipWithKey
7483 , bind
84+ , Common.union
85+ , unionWith
7586 , unionWithKey
7687 , UnionProof(..)
7788 , Common.difference
7889 , DifferenceProof(..)
90+ , differenceWith
7991 , differenceWithKey
8092 , PartialDifferenceProof(..)
93+ , Common.intersection
94+ , intersectionWith
8195 , intersectionWithKey
8296 , IntersectionProof(..)
8397 -- * Traversal
98+ , map
8499 , mapWithKey
85100 , traverseWithKey
86101 , mapAccumLWithKey
87102 , mapAccumRWithKey
103+ , mapKeys
88104 , mapKeysWith
89105 , MapProof(..)
90106 , backpermuteKeys
@@ -102,12 +118,19 @@ module Data.Map.Strict.Refined
102118 -- * Filter
103119 , Common.restrictKeys
104120 , Common.withoutKeys
121+ , Common.filter
122+ , Common.filterKeys
105123 , Common.filterWithKey
124+ , Common.partition
106125 , Common.partitionWithKey
107126 , PartitionProof(..)
127+ , Common.takeWhileAntitone
128+ , Common.dropWhileAntitone
108129 , Common.spanAntitone
109130 , PartialPartitionProof(..)
131+ , mapMaybe
110132 , mapMaybeWithKey
133+ , mapEither
111134 , mapEitherWithKey
112135 , Common.splitLookup
113136 , SplitProof(..)
@@ -138,7 +161,7 @@ import Data.Proxy
138161import Data.Reflection
139162import Data.Traversable
140163import Data.Type.Coercion
141−import Prelude hiding (lookup, null)
164+import Prelude hiding (lookup, map, null, zipWith)
142165import Refined
143166import Refined.Unsafe
144167
@@ -154,7 +177,41 @@ singleton k v = SomeMapWith (Map $ Map.singleton k v)
154177fromSet :: forall s k a. KnownSet s k => (Key s k -> a) -> Map s k a
155178fromSet f = Map $ Map.fromSet (f . unsafeKey) (reflect $ Proxy @s)
156179
157−-- | Create a map from an arbitrary traversable of key-value pairs.
180+-- | Create a map from an arbitrary traversable of key-value pairs. If a key is
181+-- repeated, the retained value is the last one in traversal order. If you're
182+-- looking for @fromList@, this is the function you want.
183+fromTraversable
184+ :: forall t k a. (Traversable t, Ord k)
185+ => t (k, a) -> SomeMapWith (FromTraversableProof 'Regular t k) k a
186+fromTraversable xs = SomeMapWith (Map m) $ FromTraversableProof proof
187+ where
188+ (m, proof) = mapAccumL
189+ (\s (k, v) -> let !s' = Map.insert k v s in (s', unsafeKey k))
190+ Map.empty
191+ xs
192+
193+-- | Create a map from an arbitrary traversable of key-value pairs, with a
194+-- function for combining values for repeated keys. The function is called as if
195+-- by 'foldl1', but flipped:
196+--
197+-- @
198+-- 'fromTraversableWith' f [(k, x1), (k, x2), (k, x3)]
199+-- = 'singleton' k (f x3 (f x2 x1))
200+-- @
201+fromTraversableWith
202+ :: forall t k a. (Traversable t, Ord k)
203+ => (a -> a -> a)
204+ -> t (k, a)
205+ -> SomeMapWith (FromTraversableProof 'Regular t k) k a
206+fromTraversableWith f xs = SomeMapWith (Map m) $ FromTraversableProof proof
207+ where
208+ (m, proof) = mapAccumL
209+ (\s (k, v) -> let !s' = Map.insertWith f k v s in (s', unsafeKey k))
210+ Map.empty
211+ xs
212+
213+-- | Create a map from an arbitrary traversable of key-value pairs. Like
214+-- 'fromTraversableWith', but the combining function has access to the key.
158215fromTraversableWithKey
159216 :: forall t k a. (Traversable t, Ord k)
160217 => (k -> a -> a -> a)
@@ -176,8 +233,40 @@ insert
176233insert k v (Map m) = SomeMapWith (Map $ Map.insert k v m)
177234 $ InsertProof (unsafeKey k) unsafeSubset
178235
236+-- | Insert a key-value pair into the map to obtain a potentially larger map,
237+-- guaranteed to contain the given key. If the key was already present, the
238+-- supplied function is used to combine the new value with the old (in that
239+-- order).
240+insertWith
241+ :: forall s k a. Ord k
242+ => (a -> a -> a)
243+ -> k
244+ -> a
245+ -> Map s k a
246+ -> SomeMapWith (InsertProof 'Regular k s) k a
247+insertWith f k v (Map m) = SomeMapWith (Map $ Map.insertWith f k v m)
248+ $ InsertProof (unsafeKey k) unsafeSubset
249+
250+-- | Insert a key-value pair into the map to obtain a potentially larger map,
251+-- guaranteed to contain the given key. Like 'insertWith', but the combining
252+-- function has access to the key, which is guaranteed to be in the old map.
253+insertWithKey
254+ :: forall s k a. Ord k
255+ => (Key s k -> a -> a -> a)
256+ -> k
257+ -> a
258+ -> Map s k a
259+ -> SomeMapWith (InsertProof 'Regular k s) k a
260+insertWithKey f k v (Map m) = SomeMapWith
261+ (Map $ Map.insertWithKey (f . unsafeKey) k v m)
262+ $ InsertProof (unsafeKey k) unsafeSubset
263+
179264-- | Overwrite a key-value pair that is known to already be in the map. The set
180265-- of keys remains the same.
266+--
267+-- @
268+-- 'reinsert' k v = 'adjust (const v) k'
269+-- @
181270reinsert
182271 :: forall s k a. Ord k
183272 => Key s k -> a -> Map s k a -> Map s k a
@@ -198,10 +287,15 @@ insertLookupWithKey f k v (Map m)
198287 (v', !m') -> ((unsafeKey k,) <$> v',)
199288 $ SomeMapWith (Map m') $ InsertProof (unsafeKey k) unsafeSubset
200289
290+-- | If the given key is in the map, update the value at that key using the
291+-- given function. In any case, the set of keys remains the same.
292+adjust' :: forall s k a. Ord k => (a -> a) -> k -> Map s k a -> Map s k a
293+adjust' = coerce $ Map.adjust @k @a
294+
201295-- | Update the value at a specific key known the be in the map using the given
202296-- function. The set of keys remains the same.
203297adjust :: forall s k a. Ord k => (a -> a) -> Key s k -> Map s k a -> Map s k a
204−adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.adjust @k @a
298+adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ adjust' @s @k @a
205299
206300-- | If the given key is in the map, update the associated value using the given
207301-- function with a proof that the key was in the map; otherwise return the map
@@ -211,6 +305,17 @@ adjustWithKey
211305adjustWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
212306 $ Map.adjustWithKey @k @a
213307
308+-- | If a key is present in the map, update its value or delete it using the
309+-- given function, returning a potentially smaller map.
310+update'
311+ :: forall s k a. Ord k
312+ => (a -> Maybe a)
313+ -> k
314+ -> Map s k a
315+ -> SomeMapWith (SupersetProof 'Regular s) k a
316+update' f k (Map m) = SomeMapWith (Map $ Map.update f k m)
317+ $ SupersetProof unsafeSubset
318+
214319-- | Update or delete a key known to be in the map using the given function,
215320-- returning a potentially smaller map.
216321update
@@ -219,7 +324,19 @@ update
219324 -> Key s k
220325 -> Map s k a
221326 -> SomeMapWith (SupersetProof 'Regular s) k a
222−update f k (Map m) = SomeMapWith (Map $ Map.update f (unrefine k) m)
327+update = gcoerceWith (unsafeCastKey @s @k) $ coerce $ update' @s @k @a
328+
329+-- | If a key is present in the map, update its value or delete it using the
330+-- given function with a proof that the key was in the map, returning a
331+-- potentially smaller map.
332+updateWithKey
333+ :: forall s k a. Ord k
334+ => (Key s k -> a -> Maybe a)
335+ -> k
336+ -> Map s k a
337+ -> SomeMapWith (SupersetProof 'Regular s) k a
338+updateWithKey f k (Map m) = SomeMapWith
339+ (Map $ Map.updateWithKey (f . unsafeKey) k m)
223340 $ SupersetProof unsafeSubset
224341
225342-- | If the given key is in the map, update or delete it using the given
@@ -239,6 +356,24 @@ updateLookupWithKey f k (Map m)
239356
240357-- | Given two maps proven to have the same keys, for each key apply the
241358-- function to the associated values, to obtain a new map with the same keys.
359+zipWith
360+ :: forall s k a b c. Ord k
361+ => (a -> b -> c) -> Map s k a -> Map s k b -> Map s k c
362+zipWith f (Map m1) (Map m2) = Map
363+ $ Map.mergeWithKey (\_ x y -> Just $ f x y)
364+ (\m -> if Map.null m
365+ then Map.empty
366+ else error "zipWith: bug: Data.Map.Strict.Refined has been subverted")
367+ (\m -> if Map.null m
368+ then Map.empty
369+ else error "zipWith: bug: Data.Map.Strict.Refined has been subverted")
370+ -- ^ Work around https://github.com/haskell/containers/issues/979
371+ m1
372+ m2
373+
374+-- | Given two maps proven to have the same keys, for each key apply the
375+-- function to the associated values and the key, to obtain a new map with the
376+-- same keys.
242377zipWithKey
243378 :: forall s k a b c. Ord k
244379 => (Key s k -> a -> b -> c) -> Map s k a -> Map s k b -> Map s k c
@@ -254,6 +389,17 @@ zipWithKey f (Map m1) (Map m2) = Map
254389 m1
255390 m2
256391
392+-- | Return the union of two maps, with a given combining function for keys that
393+-- exist in both maps simultaneously.
394+unionWith
395+ :: forall s t k a. Ord k
396+ => (a -> a -> a)
397+ -> Map s k a
398+ -> Map t k a
399+ -> SomeMapWith (UnionProof 'Regular s t) k a
400+unionWith f (Map m1) (Map m2) = SomeMapWith (Map $ Map.unionWith f m1 m2)
401+ $ UnionProof unsafeSubset unsafeSubsetWith2
402+
257403-- | Return the union of two maps, with a given combining function for keys that
258404-- exist in both maps simultaneously.
259405--
@@ -269,8 +415,20 @@ unionWithKey f (Map m1) (Map m2)
269415 = SomeMapWith (Map $ Map.unionWithKey (f . reallyUnsafeRefine) m1 m2)
270416 $ UnionProof unsafeSubset unsafeSubsetWith2
271417
272−-- | For keys that appear in both maps, the given function decides whether the
273−-- key is removed from the first map.
418+-- | Return the first map, but for keys that appear in both maps, the given
419+-- function decides whether the key is removed.
420+differenceWith
421+ :: forall s t k a b. Ord k
422+ => (a -> b -> Maybe a)
423+ -> Map s k a
424+ -> Map t k b
425+ -> SomeMapWith (PartialDifferenceProof 'Regular s t) k a
426+differenceWith f (Map m1) (Map m2)
427+ = SomeMapWith (Map $ Map.differenceWith f m1 m2)
428+ $ PartialDifferenceProof unsafeSubset unsafeSubset
429+
430+-- | Return the first map, but for keys that appear in both maps, the given
431+-- function decides whether the key is removed.
274432--
275433-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
276434-- respectively.
@@ -284,6 +442,17 @@ differenceWithKey f (Map m1) (Map m2)
284442 = SomeMapWith (Map $ Map.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
285443 $ PartialDifferenceProof unsafeSubset unsafeSubset
286444
445+-- | Return the intersection of two maps with the given combining function.
446+intersectionWith
447+ :: forall s t k a b c. Ord k
448+ => (a -> b -> c)
449+ -> Map s k a
450+ -> Map t k b
451+ -> SomeMapWith (IntersectionProof 'Regular s t) k c
452+intersectionWith f (Map m1) (Map m2)
453+ = SomeMapWith (Map $ Map.intersectionWith f m1 m2)
454+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
455+
287456-- | Return the intersection of two maps with the given combining function.
288457--
289458-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
@@ -298,6 +467,10 @@ intersectionWithKey f (Map m1) (Map m2)
298467 = SomeMapWith (Map $ Map.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
299468 $ IntersectionProof unsafeSubset unsafeSubsetWith2
300469
470+-- | Apply a function to all values in a map. The set of keys remains the same.
471+map :: forall s k a b. (a -> b) -> Map s k a -> Map s k b
472+map = coerce $ Map.map @a @b @k
473+
301474-- | Apply a function to all values in a map, together with their corresponding
302475-- keys, that are proven to be in the map. The set of keys remains the same.
303476mapWithKey :: forall s k a b. (Key s k -> a -> b) -> Map s k a -> Map s k b
@@ -331,6 +504,27 @@ mapAccumRWithKey
331504mapAccumRWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
332505 $ Map.mapAccumRWithKey @a @k @b @c
333506
507+-- | @'mapKeys' f m@ applies @f@ to each key of @m@ and collects the results
508+-- into a new map. For keys that were mapped to the same new key, the value
509+-- corresponding to the greatest of the original keys is retained.
510+mapKeys
511+ :: forall s k1 k2 a. Ord k2
512+ => (Key s k1 -> k2)
513+ -> Map s k1 a
514+ -> SomeMapWith (MapProof 'Regular s k1 k2) k2 a
515+mapKeys g (Map m)
516+ = SomeMapWith (Map $ Map.mapKeys (g . unsafeKey) m)
517+ $ MapProof (unsafeKey . g) \k2 -> case Map.lookup (unrefine k2) backMap of
518+ Nothing -> error
519+ "mapKeys: bug: Data.Map.Strict.Refined has been subverted"
520+ Just k1 -> k1
521+ where
522+ ~backMap = Map.fromList
523+ [ (k2, unsafeKey k1)
524+ | k1 <- Map.keys m
525+ , let !k2 = g $ unsafeKey k1
526+ ]
527+
334528-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
335529-- results into a new map. For keys that were mapped to the same new key, @c@
336530-- acts as the combining function for corresponding values.
@@ -353,6 +547,15 @@ mapKeysWith f g (Map m)
353547 , let !k2 = g $ unsafeKey k1
354548 ]
355549
550+-- | Apply a function to all values in a map and collect only the 'Just'
551+-- results, returning a potentially smaller map.
552+mapMaybe
553+ :: forall s k a b. (a -> Maybe b)
554+ -> Map s k a
555+ -> SomeMapWith (SupersetProof 'Regular s) k b
556+mapMaybe f (Map m) = SomeMapWith (Map $ Map.mapMaybe f m)
557+ $ SupersetProof unsafeSubset
558+
356559-- | Apply a function to all values in a map, together with their corresponding
357560-- keys, and collect only the 'Just' results, returning a potentially smaller
358561-- map.
@@ -364,6 +567,23 @@ mapMaybeWithKey f (Map m)
364567 = SomeMapWith (Map $ Map.mapMaybeWithKey (f . unsafeKey) m)
365568 $ SupersetProof unsafeSubset
366569
570+-- | Apply a function to all values in a map and collect the 'Left' and 'Right'
571+-- results into separate (disjoint) maps.
572+mapEither
573+ :: forall s k a b c. Ord k -- TODO: this is only used in the proof
574+ => (a -> Either b c)
575+ -> Map s k a
576+ -> Some2MapWith (PartitionProof 'Regular s k) k b c
577+mapEither p (Map m) = case Map.mapEither p m of
578+ (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof
579+ do \k -> case Map.lookup (unrefine k) m of
580+ Nothing -> error
581+ "mapEither: bug: Data.Map.Strict.Refined has been subverted"
582+ Just x -> case p x of
583+ Left _ -> Left $ unsafeKey $ unrefine k
584+ Right _ -> Right $ unsafeKey $ unrefine k
585+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
586+
367587-- | Apply a function to all values in a map, together with their corresponding
368588-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
369589-- maps.
src/Data/Set/Refined.hs+1−1
@@ -185,7 +185,7 @@ data SomeSet a where
185185-- you a way to refer to the set (the parameter @s@), e.g.:
186186--
187187-- @
188−-- 'withSet' ('fromSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s
188+-- 'withSet' ('fromSet' ...) $ \\(_ :: 'Proxy' s) -> doSomethingWith \@s
189189-- @
190190withSet
191191 :: forall a r. SomeSet a -> (forall s. KnownSet s a => Proxy s -> r) -> r