git.iko.soy/typeable/refined-containers
on main help sign in
Initial commit
commit413b6321b96b4e000ab61687856806c07a41df98 authormniip <mniip@mniip.com> dateWed 20 Dec 2023, 17:10 (3 years ago)
19 files changed+5881
compare with main
unifiedsplit
.github/workflows/ci.yaml+72 LICENSE+21 refined-containers.cabal+99 src/Data/Container/Refined/Conversion.hs+79 src/Data/Container/Refined/Hashable.hs+13 src/Data/Container/Refined/Proofs.hs+250 src/Data/Container/Refined/Unsafe.hs+46 src/Data/HashMap/Common/Refined.hs+426 src/Data/HashMap/Refined.hs+359 src/Data/HashMap/Strict/Refined.hs+398 src/Data/HashSet/Refined.hs+418 src/Data/IntMap/Common/Refined.hs+523 src/Data/IntMap/Refined.hs+379 src/Data/IntMap/Strict/Refined.hs+418 src/Data/IntSet/Refined.hs+502 src/Data/Map/Common/Refined.hs+515 src/Data/Map/Refined.hs+384 src/Data/Map/Strict/Refined.hs+423 src/Data/Set/Refined.hs+556
.github/workflows/ci.yamlnew file+72
@@ -0,0 +1,72 @@
1+name: CI
2+
3+on:
4+ push:
5+ branches:
6+ - "*"
7+ pull_request:
8+ branches:
9+ - "*"
10+ release:
11+ types:
12+ - prereleased
13+
14+jobs:
15+ build:
16+ strategy:
17+ matrix:
18+ platform:
19+ - ubuntu-latest
20+ ghc_version:
21+ #- "8.6.5"
22+ #- "8.8.4"
23+ #- "8.10.7"
24+ #- "9.0.2"
25+ #- "9.2.8"
26+ #- "9.4.8"
27+ - "9.6.3"
28+ #- "9.8.1"
29+ fail-fast: false
30+ name: Build on Linux, GHC ${{ matrix.ghc_version }}
31+ runs-on: ${{ matrix.platform }}
32+ steps:
33+ - uses: actions/checkout@v3
34+ - uses: haskell/actions/setup@v2
35+ with:
36+ ghc-version: ${{ matrix.ghc_version }}
37+ cabal-version: "3.10.1.0"
38+ - name: cabal configure
39+ run: cabal configure --with-ghc=ghc-${{ matrix.ghc_version }}
40+ - name: Install dependencies
41+ run: cabal build --dependencies-only --enable-tests --haddock-all
42+ - name: cabal build
43+ run: cabal build --ghc-options=-Werror
44+ - run: cabal check
45+ - run: cabal haddock
46+
47+ dist:
48+ needs: build
49+ runs-on: ubuntu-latest
50+ name: Build source tarball
51+ if: github.event_name == 'release'
52+ steps:
53+ - uses: actions/checkout@v3
54+ - uses: haskell/actions/setup@v2
55+ with:
56+ cabal-version: "3.10.1.0"
57+ - name: cabal sdist
58+ run: |
59+ cabal sdist --output-dir "$GITHUB_WORKSPACE/sdist"
60+ echo "dist=$(find "$GITHUB_WORKSPACE/sdist" -maxdepth 1 -type f -name '*.tar.gz')" >> "$GITHUB_ENV"
61+ - uses: actions/upload-artifact@v2
62+ with:
63+ name: Source tarball
64+ path: ${{ env.dist }}
65+ - uses: svenstaro/upload-release-action@2.2.1
66+ with:
67+ repo_token: ${{ secrets.GITHUB_TOKEN }}
68+ file: ${{ env.dist }}
69+ tag: ${{ github.ref }}
70+ - name: Upload tarball to hackage
71+ run: |
72+ curl -X POST 'https://hackage.haskell.org/packages/candidates/' -H "Authorization: X-ApiKey ${{ secrets.HACKAGE_KEY }}" -F "package=@$dist" -w '%{url_effective}\n'
LICENSEnew file+21
@@ -0,0 +1,21 @@
1+MIT License
2+
3+Copyright (c) 2023 Typeable
4+
5+Permission is hereby granted, free of charge, to any person obtaining a copy
6+of this software and associated documentation files (the "Software"), to deal
7+in the Software without restriction, including without limitation the rights
8+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9+copies of the Software, and to permit persons to whom the Software is
10+furnished to do so, subject to the following conditions:
11+
12+The above copyright notice and this permission notice shall be included in all
13+copies or substantial portions of the Software.
14+
15+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21+SOFTWARE.
refined-containers.cabalnew file+99
@@ -0,0 +1,99 @@
1+cabal-version: 3.0
2+name: refined-containers
3+category: Data
4+synopsis:
5+ Type-checked proof that a key exists in a container and can be safely
6+ indexed.
7+description:
8+ This package defines ways to prove that a key exists in an associative
9+ container such as a 'Map', 'IntMap', or 'HashMap'; so that the key can be
10+ used to index into the map without a 'Maybe' or manually handling the
11+ \"impossible\" case with 'error' or other partial functions.
12+
13+ To do this, the containers are tagged with a type parameter that identifies
14+ their set of keys, so that if you have another container with the same
15+ parameter, you know it has the same keys.
16+
17+ There is also a type of keys that have been proven to exist in such
18+ containers -- a refinement type. They are also tagged with a type parameter.
19+ If the type parameter of the key matches that of the container, indexing is
20+ guaranteed to proceed without failure.
21+
22+license: MIT
23+license-file: LICENSE
24+author: mniip@typeable.io
25+maintainer: mniip@typeable.io
26+version: 0.1.0.0
27+build-type: Simple
28+
29+tested-with:
30+ , GHC == 9.8.1
31+ , GHC == 9.6.3
32+ , GHC == 9.4.8
33+ , GHC == 9.2.8
34+ , GHC == 9.0.2
35+ , GHC == 8.10.7
36+ , GHC == 8.6.5
37+
38+source-repository head
39+ type: git
40+ location: https://github.com/typeable/refined-containers/
41+
42+library
43+ build-depends:
44+ , base >= 4.12 && < 4.19
45+ , adjunctions >= 4.4 && < 4.5
46+ , constraints >= 0.11 && < 0.15
47+ , containers >= 0.5.7 && < 0.8
48+ , deepseq >= 1.4 && < 1.6
49+ , distributive >= 0.5.3 && < 0.7
50+ , hashable >= 1.2.7 && < 1.5
51+ , indexed-traversable >= 0.1 && < 0.2
52+ , mtl >= 2.2.2 && < 2.4
53+ , refined >= 0.5 && < 0.9
54+ , reflection >= 2 && < 2.2
55+ , unordered-containers >= 0.2.11 && < 0.3
56+ exposed-modules:
57+ Data.HashMap.Refined
58+ Data.HashMap.Strict.Refined
59+ Data.HashSet.Refined
60+ Data.IntMap.Refined
61+ Data.IntMap.Strict.Refined
62+ Data.IntSet.Refined
63+ Data.Map.Refined
64+ Data.Map.Strict.Refined
65+ Data.Set.Refined
66+ other-modules:
67+ Data.Container.Refined.Conversion
68+ Data.Container.Refined.Hashable
69+ Data.Container.Refined.Proofs
70+ Data.Container.Refined.Unsafe
71+ Data.HashMap.Common.Refined
72+ Data.IntMap.Common.Refined
73+ Data.Map.Common.Refined
74+ hs-source-dirs: src
75+ default-language: Haskell2010
76+ default-extensions:
77+ BangPatterns
78+ BlockArguments
79+ ConstraintKinds
80+ DataKinds
81+ DeriveTraversable
82+ DerivingStrategies
83+ FlexibleContexts
84+ FlexibleInstances
85+ GADTs
86+ GeneralizedNewtypeDeriving
87+ MagicHash
88+ MultiParamTypeClasses
89+ MultiWayIf
90+ OverloadedStrings
91+ RankNTypes
92+ PatternSynonyms
93+ RoleAnnotations
94+ ScopedTypeVariables
95+ TupleSections
96+ TypeApplications
97+ TypeFamilies
98+ TypeOperators
99+ ghc-options: -Wall -Wredundant-constraints
src/Data/Container/Refined/Conversion.hsnew file+79
@@ -0,0 +1,79 @@
1+module Data.Container.Refined.Conversion where
2+
3+import Data.Constraint (Dict(..))
4+import Data.Container.Refined.Hashable
5+import Data.Container.Refined.Proofs
6+import qualified Data.HashMap.Lazy as HashMap
7+import Data.HashMap.Common.Refined
8+import qualified Data.HashSet as HashSet
9+import qualified Data.IntMap as IntMap
10+import Data.IntMap.Common.Refined
11+import qualified Data.IntSet as IntSet
12+import qualified Data.Map as Map
13+import Data.Map.Common.Refined
14+import Data.Proxy
15+import Data.Reflection
16+import qualified Data.Set as Set
17+import Data.Type.Equality ((:~:)(..))
18+import Unsafe.Coerce
19+
20+
21+unsafeWrapSet :: forall r r' a. KnownSet r' a => Proxy r' -> Set r a
22+unsafeWrapSet _ = case unsafeCoerce Refl :: r :~: r' of Refl -> Dict
23+{-# INLINE unsafeWrapSet #-}
24+
25+unsafeWrapIntSet :: forall r r'. KnownIntSet r' => Proxy r' -> IntSet r
26+unsafeWrapIntSet _ = case unsafeCoerce Refl :: r :~: r' of Refl -> Dict
27+{-# INLINE unsafeWrapIntSet #-}
28+
29+unsafeWrapHashSet :: forall r r' a. KnownHashSet r' a => Proxy r' -> HashSet r a
30+unsafeWrapHashSet _ = case unsafeCoerce Refl :: r :~: r' of Refl -> Dict
31+{-# INLINE unsafeWrapHashSet #-}
32+
33+set2IntSet :: forall s. KnownSet s Int => IntSet s
34+set2IntSet = reify
35+ (IntSet.fromDistinctAscList $ Set.toAscList $ reflect $ Proxy @s)
36+ unsafeWrapIntSet
37+
38+map2IntMap :: forall s a. Map s Int a -> IntMap s a
39+map2IntMap (Map m) = IntMap $ IntMap.fromDistinctAscList $ Map.toAscList m
40+
41+set2HashSet :: forall s a. (Hashable a, KnownSet s a) => HashSet s a
42+set2HashSet = reify
43+ (HashSet.fromList $ Set.toList $ reflect $ Proxy @s)
44+ unsafeWrapHashSet
45+
46+map2HashMap :: forall s k a. Hashable k => Map s k a -> HashMap s k a
47+map2HashMap (Map m) = HashMap $ HashMap.fromList $ Map.toList m
48+
49+intSet2Set :: forall s. KnownIntSet s => Set s Int
50+intSet2Set = reify
51+ (Set.fromDistinctAscList $ IntSet.toAscList $ reflect $ Proxy @s)
52+ unsafeWrapSet
53+
54+intMap2Map :: forall s a. IntMap s a -> Map s Int a
55+intMap2Map (IntMap m) = Map $ Map.fromDistinctAscList $ IntMap.toAscList m
56+
57+intSet2HashSet :: forall s. KnownIntSet s => HashSet s Int
58+intSet2HashSet = reify
59+ (HashSet.fromList $ IntSet.toList $ reflect $ Proxy @s)
60+ unsafeWrapHashSet
61+
62+intMap2HashMap :: forall s a. IntMap s a -> HashMap s Int a
63+intMap2HashMap (IntMap m) = HashMap $ HashMap.fromList $ IntMap.toList m
64+
65+hashSet2Set :: forall s a. (Ord a, KnownHashSet s a) => Set s a
66+hashSet2Set = reify
67+ (Set.fromList $ HashSet.toList $ reflect $ Proxy @s)
68+ unsafeWrapSet
69+
70+hashMap2Map :: forall s k a. Ord k => HashMap s k a -> Map s k a
71+hashMap2Map (HashMap m) = Map $ Map.fromList $ HashMap.toList m
72+
73+hashSet2IntSet :: forall s. KnownHashSet s Int => IntSet s
74+hashSet2IntSet = reify
75+ (IntSet.fromList $ HashSet.toList $ reflect $ Proxy @s)
76+ unsafeWrapIntSet
77+
78+hashMap2IntMap :: forall s a. HashMap s Int a -> IntMap s a
79+hashMap2IntMap (HashMap m) = IntMap $ IntMap.fromList $ HashMap.toList m
src/Data/Container/Refined/Hashable.hsnew file+13
@@ -0,0 +1,13 @@
1+{-# LANGUAGE CPP #-}
2+module Data.Container.Refined.Hashable
3+ ( Hashable
4+ ) where
5+
6+#if MIN_VERSION_hashable(1, 4, 0)
7+import Data.Hashable (Hashable)
8+#else
9+import qualified Data.Hashable as Hashable
10+
11+
12+type Hashable a = (Eq a, Hashable.Hashable a)
13+#endif
src/Data/Container/Refined/Proofs.hsnew file+250
@@ -0,0 +1,250 @@
1+{-# LANGUAGE CPP #-}
2+#if MIN_VERSION_refined(0, 7, 0)
3+#else
4+{-# LANGUAGE UndecidableInstances #-}
5+#endif
6+module Data.Container.Refined.Proofs where
7+
8+import Data.Constraint (Dict(..))
9+import Data.Container.Refined.Hashable
10+import Data.Kind
11+import Data.Reflection
12+import qualified Data.HashSet as HashSet
13+import qualified Data.IntSet as IntSet
14+import qualified Data.Set as Set
15+import Data.Type.Coercion
16+import Data.Typeable
17+import Refined
18+
19+
20+-- | A constraint evidencing that we know the contents of the set @s@ at
21+-- runtime. Whenever you see this constraint on a function, there is an actual
22+-- @'Data.Set.Set' a@ being passed around at runtime.
23+--
24+-- Given this constraint, to obtain a regular @'Data.Set.Set' a@ you can use
25+-- 'reflect'.
26+type KnownSet s a = Reifies s (Set.Set a)
27+
28+-- | A 'Data.Set.Set' whose contents are tracked by the type parameter @s@. This
29+-- is a \"singleton\": for a given @s@ there's only one value of this type.
30+--
31+-- Since this is just a 'Dict', you can freely convert between the value ('Set')
32+-- and the constraint ('KnownSet'). This library prefers to use the constraint.
33+type Set s a = Dict (KnownSet s a)
34+
35+-- | A constraint evidencing that we know the contents of the set @s@ at
36+-- runtime. Whenever you see this constraint on a function, there is an actual
37+-- 'Data.IntSet.IntSet' being passed around at runtime.
38+--
39+-- Given this constraint, to obtain a regular 'Data.IntSet.IntSet' you can use
40+-- 'reflect'.
41+type KnownIntSet s = Reifies s IntSet.IntSet
42+
43+-- | A 'Data.IntSet.IntSet' whose contents are tracked by the type parameter
44+-- @s@. This is a \"singleton\": for a given @s@ there's only one value of this
45+-- type.
46+--
47+-- Since this is just a 'Dict', you can freely convert between the value
48+-- ('IntSet') and the constraint ('KnownIntSet'). This library prefers to use
49+-- the constraint.
50+type IntSet s = Dict (KnownIntSet s)
51+
52+-- | A constraint evidencing that we know the contents of the set @s@ at
53+-- runtime. Whenever you see this constraint on a function, there is an actual
54+-- @'Data.HashSet.HashSet' a@ being passed around at runtime.
55+--
56+-- Given this constraint, to obtain a regular @'Data.HashSet.HashSet' a@ you can
57+-- use 'reflect'.
58+type KnownHashSet s a = Reifies s (HashSet.HashSet a)
59+
60+-- | A 'Data.HashSet.HashSet' whose contents are tracked by the type parameter
61+-- @s@. This is a \"singleton\": for a given @s@ there's only one value of this
62+-- type.
63+--
64+-- Since this is just a 'Dict', you can freely convert between the value
65+-- ('HashSet') and the constraint ('KnownHashSet'). This library prefers to use
66+-- the constraint.
67+type HashSet s a = Dict (KnownHashSet s a)
68+
69+-- | Disambiguate the choice of implementation for sets and maps.
70+data Flavor
71+ = Regular -- ^ 'Data.Set.Set' and 'Data.Map.Map'
72+ | Int -- ^ 'Data.IntSet.IntSet' and 'Data.IntMap.IntMap'
73+ | Hashed -- ^ 'Data.HashSet.HashSet' and 'Data.HashMap.Lazy.HashMap'
74+
75+-- | A predicate for use with "Refined", verifying that a value is an element of
76+-- the set @s@.
77+data InSet (f :: Flavor) (s :: Type) = InSet
78+
79+-- | See 'Data.Set.Refined.revealPredicate'.
80+instance (Ord a, Typeable s, KnownSet s a)
81+ => Predicate (InSet 'Regular s) a where
82+ validate p x
83+ | x `Set.member` reflect (Proxy @s) = success
84+ | otherwise = throwRefineOtherException
85+ (typeRep p)
86+ "Value is not in the Set"
87+
88+-- | See 'Data.IntSet.Refined.revealPredicate'.
89+instance (a ~ Int, Typeable s, KnownIntSet s)
90+ => Predicate (InSet 'Int s) a where
91+ validate p x
92+ | x `IntSet.member` reflect (Proxy @s) = success
93+ | otherwise = throwRefineOtherException
94+ (typeRep p)
95+ "Value is not in the IntSet"
96+
97+-- | See 'Data.HashSet.Refined.revealPredicate'.
98+instance (Hashable a, Typeable s, KnownHashSet s a)
99+ => Predicate (InSet 'Hashed s) a where
100+ validate p x
101+ | x `HashSet.member` reflect (Proxy @s) = success
102+ | otherwise = throwRefineOtherException
103+ (typeRep p)
104+ "Value is not in the HashSet"
105+
106+-- | A proof that values satisfying @p@ can be cast into values satisfying @q@.
107+--
108+-- For example, @'InSet' s ':->' 'InSet' r@ proves that \(s \subseteq r\).
109+type p :-> q = forall x. Refined p x -> Refined q x
110+infix 1 :->
111+
112+-- | Proof that the set @r@ is empty.
113+newtype EmptyProof f r = EmptyProof
114+ (forall s. InSet f r :-> InSet f s)
115+ -- ^ \(\forall s, r \subseteq s\), which is equivalent to
116+ -- \(r \subseteq \varnothing\)
117+
118+-- | Proof that @r@ contains an element of type @a@.
119+newtype SingletonProof f a r = SingletonProof
120+ (Refined (InSet f r) a) -- ^ The element that is guaranteed to be in @r@
121+
122+-- | Proof that elements of @t a@ are contained in @r@.
123+newtype FromTraversableProof f (t :: Type -> Type) a r = FromTraversableProof
124+ (t (Refined (InSet f r) a))
125+ -- ^ The original traversable, with all elements refined with a guarantee of
126+ -- being in @r@
127+
128+-- | Proof that @r@ is @s@ with @a@ inserted.
129+data InsertProof f a s r = InsertProof
130+ (Refined (InSet f r) a)
131+ -- ^ The element that was inserted and is guaranteed to be in @r@.
132+ (InSet f s :-> InSet f r) -- ^ \(s \subseteq r \)
133+
134+-- | Proof that @s@ is a subset of the set @r@.
135+newtype SubsetProof f s r = SubsetProof
136+ (InSet f s :-> InSet f r) -- ^ \(s \subseteq r\)
137+
138+-- | Proof that @s@ is a superset of the set @r@.
139+newtype SupersetProof f s r = SupersetProof
140+ (InSet f r :-> InSet f s) -- ^ \(r \subseteq s\)
141+
142+-- | Proof that @s@ and @r@ are disjoint.
143+newtype DisjointProof f s r = DisjointProof
144+ (forall t. InSet f t :-> InSet f s
145+ -> InSet f t :-> InSet f r
146+ -> forall u. InSet f t :-> InSet f u)
147+ -- ^ \(\forall t,(t\subseteq s)\land(t\subseteq r)\implies\forall u,t\subseteq u\),
148+ -- which is equivalent to \(s \cap r \subseteq \varnothing\)
149+
150+-- | Proof that unioning @s@ and @t@ gives @r@.
151+data UnionProof f s t r = UnionProof
152+ (InSet f s || InSet f t :-> InSet f r) -- ^ \(s \cup t \subseteq r\)
153+ (forall u. InSet f s :-> InSet f u
154+ -> InSet f t :-> InSet f u
155+ -> InSet f r :-> InSet f u)
156+ -- ^ \(\forall u,(s\subseteq u)\land(t\subseteq u)\implies r\subseteq u\),
157+ -- which is equivalent to \(r \subseteq s \cup u\)
158+
159+-- | Proof that if from @s@ you subtract @t@, then you get @r@.
160+data DifferenceProof f s t r = DifferenceProof
161+ (InSet f r :-> InSet f s) -- ^ \(r \subseteq s\)
162+ (forall u. InSet f u :-> InSet f r
163+ -> InSet f u :-> InSet f t
164+ -> forall v. InSet f u :-> InSet f v)
165+ -- ^ \(\forall u,(u\subseteq r)\land(u\subseteq t)\implies\forall v,u\subseteq v\),
166+ -- which is equivalent to \(r \cap t \subseteq \varnothing\)
167+ (InSet f s :-> InSet f t || InSet f r) -- ^ \(s \subseteq t \cup r\)
168+
169+-- | Proof that @r@ is obtained by removing some of @t@'s elements from @s@.
170+data PartialDifferenceProof f s t r = PartialDifferenceProof
171+ (InSet f r :-> InSet f s) -- ^ \(r \subseteq s\)
172+ (InSet f s :-> InSet f t || InSet f r) -- ^ \(s \subseteq t \cup r\)
173+
174+-- | Proof that intersecting @s@ and @t@ gives @r@.
175+data IntersectionProof f s t r = IntersectionProof
176+ (InSet f r :-> InSet f s && InSet f t) -- ^ \(r \subseteq s \cap t\)
177+ (forall u. InSet f u :-> InSet f s
178+ -> InSet f u :-> InSet f t
179+ -> InSet f u :-> InSet f r)
180+ -- ^ \(\forall u,(u\subseteq s)\land(u\subseteq t)\implies u\subseteq r\),
181+ -- which is equivalent to \(s \cap t \subseteq r\)
182+
183+-- | Proof that the cartesian product of @s@ and @t@ is @r@.
184+newtype ProductProof f s t r = ProductProof
185+ (forall a b. Coercion
186+ (Refined (InSet f s) a, Refined (InSet f t) b)
187+ (Refined (InSet f r) (a, b)))
188+ -- ^ A pair of elements from @s@ and @t@ respectively can be converted into an
189+ -- element of @r@ and back. You can use @'coerceWith' co@ and
190+ -- @'coerceWith' ('sym' co)@.
191+
192+-- | Proof that the tagged disjoint union of @s@ and @t@ is @r@.
193+newtype CoproductProof f s t r = CoproductProof
194+ (forall a b. Coercion
195+ (Either (Refined (InSet f s) a) (Refined (InSet f t) b))
196+ (Refined (InSet f r) (Either a b)))
197+ -- ^ Coproduct of elements of @s@ and @t@ can be converted into an element of
198+ -- @r@ and back. You can use @'coerceWith' co@ and @'coerceWith' ('sym' co)@.
199+
200+-- | Proof that @s@ is the union of disjoint subsets @r@ and @q@, together with
201+-- a procedure that decides which of the two an element belongs to.
202+data PartitionProof f s a r q = PartitionProof
203+ (Refined (InSet f s) a
204+ -> Either (Refined (InSet f r) a) (Refined (InSet f q) a))
205+ -- ^ Decide whether a given element of @s@ belongs to @r@ or to @q@
206+ (InSet f r || InSet f q :-> InSet f s) -- ^ \(r \cup q \subseteq s\)
207+ (forall t. InSet f r :-> InSet f t
208+ -> InSet f q :-> InSet f t
209+ -> InSet f s :-> InSet f t)
210+ -- ^ \(\forall t,(r\subseteq t)\land(q\subseteq t)\implies s\subseteq t\),
211+ -- which is equivalent to \(s \subseteq r \cup q\)
212+ (forall t. InSet f t :-> InSet f r
213+ -> InSet f t :-> InSet f q
214+ -> forall u. InSet f t :-> InSet f u)
215+ -- ^ \(\forall t,(t\subseteq r)\land(t\subseteq q)\implies\forall u,t\subseteq u\),
216+ -- which is equivalent to \(r \cap q \subseteq \varnothing\)
217+
218+-- | Proof that @s@ is the union of disjoint subsets @r@ and @q@, but without a
219+-- deciding procedure.
220+data PartialPartitionProof f s r q = PartialPartitionProof
221+ (InSet f r || InSet f q :-> InSet f s) -- ^ \(r \cup q \subseteq s\)
222+ (forall t. InSet f r :-> InSet f t
223+ -> InSet f q :-> InSet f t
224+ -> InSet f s :-> InSet f t)
225+ -- ^ \(\forall t,(r\subseteq t)\land(q\subseteq t)\implies s\subseteq t\),
226+ -- which is equivalent to \(s \subseteq r \cup q\)
227+ (forall t. InSet f t :-> InSet f r
228+ -> InSet f t :-> InSet f q
229+ -> forall u. InSet f t :-> InSet f u)
230+ -- ^ \(\forall t,(t\subseteq r)\land(t\subseteq q)\implies\forall u,t\subseteq u\),
231+ -- which is equivalent to \(r \cap q \subseteq \varnothing\)
232+
233+-- | Proof that @s@ contains disjoint subsets @r@ and @q@, along with an
234+-- optional element between them.
235+data SplitProof f s e r q = SplitProof
236+ !(Maybe e) -- ^ The element between @r@ and @q@
237+ (InSet f r || InSet f q :-> InSet f s) -- ^ \(r \cup q \subseteq s\)
238+ (forall t. InSet f t :-> InSet f r
239+ -> InSet f t :-> InSet f q
240+ -> forall u. InSet f t :-> InSet f u)
241+ -- ^ \(\forall t,(t\subseteq r)\land(t\subseteq q)\implies\forall u,t\subseteq u\),
242+ -- which is equivalent to \(r \cap q \subseteq \varnothing\)
243+
244+-- | Proof that @r@ is the direct image of @s@ under some mapping @f :: a -> b@.
245+data MapProof f s a b r = MapProof
246+ (Refined (InSet f s) a -> Refined (InSet f r) b)
247+ -- ^ Compute the image of an element of @s@ (the image is then an element of
248+ -- @r@)
249+ (Refined (InSet f r) b -> Refined (InSet f s) a)
250+ -- ^ For an element of @r@, return an arbitrary preimage from @s@
src/Data/Container/Refined/Unsafe.hsnew file+46
@@ -0,0 +1,46 @@
1+{-# LANGUAGE ImpredicativeTypes #-}
2+module Data.Container.Refined.Unsafe where
3+
4+import Data.Container.Refined.Proofs
5+import Data.Type.Coercion
6+import Refined
7+import Refined.Unsafe
8+
9+
10+unsafeSubset :: p :-> q
11+unsafeSubset = reallyUnsafeRefine . unrefine
12+{-# INLINE unsafeSubset #-}
13+
14+unsafeSubsetWith2 :: p' :-> q' -> p'' :-> q'' -> p :-> q
15+unsafeSubsetWith2 f g = reallyUnsafeRefine . unrefine
16+ . f . reallyUnsafeRefine . unrefine
17+ . g . reallyUnsafeRefine . unrefine
18+{-# INLINE unsafeSubsetWith2 #-}
19+
20+-- Because `Refined p x` is a newtype over `x`, by parametricity a `p :-> q`
21+-- can either diverge or be `id`. This ensures that it does not diverge.
22+rnfProof :: (p :-> q) -> ()
23+rnfProof f = unrefine $ f $ reallyUnsafeRefine ()
24+{-# INLINE rnfProof #-}
25+
26+-- | This function can be used to freely convert between @Element@ and @Key@
27+-- types of various flavors ('Regular', v'Int', 'Hashed'), corresponding to the
28+-- different implementations of sets and maps.
29+castFlavor
30+ :: forall (f :: Flavor) (g :: Flavor) s a. Coercion
31+ (Refined (InSet f s) a)
32+ (Refined (InSet g s) a)
33+castFlavor = sym (reallyUnsafeUnderlyingRefined @a @(InSet f s))
34+ `trans` reallyUnsafeUnderlyingRefined @a @(InSet g s)
35+{-# INLINE castFlavor #-}
36+
37+castRefined
38+ :: forall a p q. (p :-> q)
39+ -> (q :-> p)
40+ -> Coercion (Refined p a) (Refined q a)
41+castRefined f g
42+ | () <- rnfProof f
43+ , () <- rnfProof g
44+ = sym (reallyUnsafeUnderlyingRefined @a @p)
45+ `trans` reallyUnsafeUnderlyingRefined @a @q
46+{-# INLINE castRefined #-}
src/Data/HashMap/Common/Refined.hsnew file+426
@@ -0,0 +1,426 @@
1+{-# LANGUAGE CPP #-}
2+{-# LANGUAGE UndecidableInstances #-}
3+module Data.HashMap.Common.Refined where
4+
5+import Control.Monad.Reader
6+import Control.DeepSeq
7+import Data.Coerce
8+import Data.Constraint (Dict(..))
9+import Data.Container.Refined.Hashable
10+import Data.Container.Refined.Proofs
11+import Data.Container.Refined.Unsafe
12+import Data.Distributive
13+import Data.Foldable.WithIndex
14+import Data.Functor.Rep
15+import Data.Functor.WithIndex
16+import qualified Data.Hashable as Hashable
17+import qualified Data.HashMap.Lazy as HashMap
18+import qualified Data.HashMap.Strict as HashMapStrict
19+import qualified Data.HashSet as HashSet
20+import Data.Proxy
21+import Data.Reflection
22+import Data.Traversable.WithIndex
23+import Data.Type.Coercion
24+import Data.Type.Equality ((:~:)(..))
25+import Refined
26+import Refined.Unsafe
27+import Unsafe.Coerce
28+
29+#if MIN_VERSION_unordered_containers(0, 2, 12)
30+#else
31+import Data.Monoid (All(..))
32+#endif
33+
34+
35+-- | A wrapper around a regular 'Data.HashMap.HashMap' with a type parameter @s@
36+-- identifying the set of keys present in the map.
37+--
38+-- A key of type @k@ may not be present in the map, but a @'Key' s k@ is
39+-- guaranteed to be present (if the @s@ parameters match). Thus the map is
40+-- isomorphic to a (total) function @'Key' s k -> a@, which motivates many of
41+-- the instances below.
42+--
43+-- A 'HashMap' always knows its set of keys, so given @'HashMap' s k a@ we can
44+-- always derive @'KnownHashSet' s k@ by pattern matching on the 'Dict' returned
45+-- by 'keysSet'.
46+newtype HashMap s k a = HashMap (HashMap.HashMap k a)
47+ deriving newtype (Eq, Ord, Show, Functor, Foldable, Hashable.Hashable, NFData)
48+ deriving stock (Traversable)
49+type role HashMap nominal nominal representational
50+
51+-- | Convert to a regular 'Data.HashMap.HashMap', forgetting its set of keys.
52+toMap :: forall s k a. HashMap s k a -> HashMap.HashMap k a
53+toMap (HashMap m) = m
54+
55+-- | @'Key' s k@ is a key of type @k@ that has been verified to be an element
56+-- of the set @s@, and thus verified to be present in all @'HashMap' s k@ maps.
57+--
58+-- Thus, @'Key' s k@ is a \"refinement\" type of @k@, and this library
59+-- integrates with an implementation of refimenement types in "Refined", so
60+-- the machinery from there can be used to manipulate 'Key's (however see
61+-- 'Data.Set.Refined.revealPredicate').
62+--
63+-- The underlying @k@ value can be obtained with 'unrefine'. A @k@ can be
64+-- validated into an @'Key' s k@ with 'member'.
65+type Key s = Refined (InSet 'Hashed s)
66+
67+unsafeCastKey :: forall s k. Coercion k (Key s k)
68+unsafeCastKey = reallyUnsafeUnderlyingRefined
69+
70+unsafeKey :: k -> Key s k
71+unsafeKey = coerceWith unsafeCastKey
72+
73+-- | An existential wrapper for a 'HashMap' with an as-yet-unknown set of keys.
74+-- Pattern maching on it gives you a way to refer to the set (the parameter
75+-- @s@), e.g.
76+--
77+-- @
78+-- case 'fromHashMap' ... of
79+-- 'SomeHashMap' \@s m -> doSomethingWith \@s
80+--
81+-- case 'fromHashMap' ... of
82+-- 'SomeHashMap' (m :: 'HashMap' s k a) -> doSomethingWith \@s
83+-- @
84+data SomeHashMap k a where
85+ SomeHashMap :: forall s k a. !(HashMap s k a) -> SomeHashMap k a
86+
87+-- | Apply a map with an unknown set of keys to a continuation that can accept
88+-- a map with any set of keys. This gives you a way to refer to the set (the
89+-- parameter @s@), e.g.:
90+--
91+-- @
92+-- 'withHashMap' ('fromHashMap' ...
93+-- $ \(m :: 'HashMap' s k a) -> doSomethingWith \@s
94+-- @
95+withHashMap
96+ :: forall k a r. SomeHashMap k a -> (forall s. HashMap s k a -> r) -> r
97+withHashMap (SomeHashMap m) k = k m
98+
99+-- | Construct a map from a regular 'Data.HashMap.Lazy.HashMap'.
100+fromHashMap :: forall k a. HashMap.HashMap k a -> SomeHashMap k a
101+fromHashMap m = SomeHashMap (HashMap m)
102+
103+-- | An existential wrapper for a 'HashMap' with an as-yet-unknown set of keys,
104+-- together with a proof of some fact @p@ about the set. Pattern matching on it
105+-- gives you a way to refer to the set (the parameter @s@). Functions that
106+-- change the set of keys in a map will return the map in this way, together
107+-- with a proof that somehow relates the keys set to the function's inputs.
108+data SomeHashMapWith p k a where
109+ SomeHashMapWith
110+ :: forall s k a p. !(HashMap s k a) -> !(p s) -> SomeHashMapWith p k a
111+
112+-- | Apply a map with proof for an unknown set of keys to a continuation that
113+-- can accept a map with any set of keys satisfying the proof. This gives you a
114+-- way to refer to the set (the parameter @s@).
115+withHashMapWith
116+ :: forall k a r p. SomeHashMapWith p k a
117+ -> (forall s. HashMap s k a -> p s -> r)
118+ -> r
119+withHashMapWith (SomeHashMapWith m p) k = k m p
120+
121+-- | An existential wrapper for a pair of maps with as-yet-unknown sets of keys,
122+-- together with a proof of some fact @p@ relating them.
123+data Some2HashMapWith p k a b where
124+ Some2HashMapWith
125+ :: forall s t k a b p. !(HashMap s k a)
126+ -> !(HashMap t k b)
127+ -> !(p s t)
128+ -> Some2HashMapWith p k a b
129+
130+-- | Apply a pair of maps with proof for unknown sets of keys to a continuation
131+-- that can accept any pair of maps with any sets of keys satisfying the proof.
132+-- This gives you a way to refer to the sets (the parameters @s@ and @t@).
133+with2HashMapWith
134+ :: forall k a b r p. Some2HashMapWith p k a b
135+ -> (forall s t. HashMap s k a -> HashMap t k b -> p s t -> r)
136+ -> r
137+with2HashMapWith (Some2HashMapWith m1 m2 p) k = k m1 m2 p
138+
139+-- | An empty map.
140+empty :: forall k a. SomeHashMapWith (EmptyProof 'Hashed) k a
141+empty = SomeHashMapWith (HashMap HashMap.empty) $ EmptyProof unsafeSubset
142+
143+-- | Create a map from a set of keys, and a function that for each key computes
144+-- the corresponding value.
145+fromSet :: forall s k a. KnownHashSet s k => (Key s k -> a) -> HashMap s k a
146+fromSet f = HashMap $ HashMap.mapWithKey (\k _ -> f $ unsafeKey k)
147+ $ HashSet.toMap (reflect $ Proxy @s)
148+
149+-- | Delete a key and its value from the map if present, returning a potentially
150+-- smaller map.
151+delete
152+ :: forall s k a. Hashable k
153+ => k -> HashMap s k a -> SomeHashMapWith (SupersetProof 'Hashed s) k a
154+delete k (HashMap m) = SomeHashMapWith (HashMap $ HashMap.delete k m)
155+ $ SupersetProof unsafeSubset
156+
157+-- | If the key is in the map, return the proof of this, and the associated
158+-- value; otherwise return 'Nothing'.
159+lookup :: forall s k a. Hashable k => k -> HashMap s k a -> Maybe (Key s k, a)
160+lookup k (HashMap m) = (unsafeKey k,) <$> HashMap.lookup k m
161+
162+-- | Given a key that is proven to be in the map, return the associated value.
163+--
164+-- Unlike 'Data.HashMap.!' from "Data.HashMap.Lazy", this function is total, as
165+-- it is impossible to obtain a @'Key' s k@ for a key that is not in the map
166+-- @'HashMap' s k a@.
167+(!) :: forall s k a. Hashable k => HashMap s k a -> Key s k -> a
168+(!) (HashMap m) k = case HashMap.lookup (unrefine k) m of
169+ Nothing -> error "(!): bug: Data.HashMap.Refined has been subverted"
170+ Just x -> x
171+
172+-- | If a key is in the map, return the proof that it is.
173+member :: forall s k a. Hashable k => k -> HashMap s k a -> Maybe (Key s k)
174+member k (HashMap m)
175+ | k `HashMap.member` m = Just (unsafeKey k)
176+ | otherwise = Nothing
177+
178+-- | If a map is empty, return a proof that it is.
179+null :: forall s k a. HashMap s k a -> Maybe (EmptyProof 'Hashed s)
180+null (HashMap m)
181+ | HashMap.null m = Just $ EmptyProof unsafeSubset
182+ | otherwise = Nothing
183+
184+-- | If all keys of the first map are also present in the second map, and the
185+-- given function returns 'True' for their associated values, return a proof
186+-- that the keys form a subset.
187+isSubmapOfBy
188+ :: forall s t k a b. Hashable k
189+ => (a -> b -> Bool)
190+ -> HashMap s k a
191+ -> HashMap t k b
192+ -> Maybe (SubsetProof 'Hashed s t)
193+isSubmapOfBy f (HashMap m1) (HashMap m2)
194+#if MIN_VERSION_unordered_containers(0, 2, 12)
195+ | HashMap.isSubmapOfBy f m1 m2
196+#else
197+ | All True <- flip HashMap.foldMapWithKey m1
198+ \k v1 -> case HashMap.lookup k m2 of
199+ Just v2 | f v1 v2 -> mempty
200+ _ -> All False
201+#endif
202+ = Just $ SubsetProof unsafeSubset
203+ | otherwise = Nothing
204+
205+-- | If two maps are disjoint (i.e. their intersection is empty), return a proof
206+-- of that.
207+disjoint
208+ :: forall s t k a b. Hashable k
209+ => HashMap s k a -> HashMap t k b -> Maybe (DisjointProof 'Hashed s t)
210+disjoint (HashMap m1) (HashMap m2)
211+ | HashMap.null $ HashMapStrict.intersectionWith (\_ _ -> ()) m1 m2
212+ = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g
213+ | otherwise = Nothing
214+
215+-- | Given two maps proven to have the same keys, for each key apply the
216+-- function to the associated values, to obtain a new map with the same keys.
217+zipWithKey
218+ :: 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
222+
223+-- | Remove the keys that appear in the second map from the first map.
224+difference
225+ :: forall s t k a b. Hashable k
226+ => HashMap s k a
227+ -> HashMap t k b
228+ -> SomeHashMapWith (DifferenceProof 'Hashed s t) k a
229+difference (HashMap m1) (HashMap m2)
230+ = SomeHashMapWith (HashMap $ HashMap.difference m1 m2)
231+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
232+
233+-- | Apply a function to all values in a map, together with their corresponding
234+-- keys, that are proven to be in the map. The set of keys remains the same.
235+mapWithKey
236+ :: forall s k a b. (Key s k -> a -> b) -> HashMap s k a -> HashMap s k b
237+mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
238+ $ HashMap.mapWithKey @k @a @b
239+
240+-- | Map an 'Applicative' transformation with access to each value's
241+-- corresponding key and a proof that it is in the map. The set of keys remains
242+-- unchanged.
243+traverseWithKey
244+ :: forall s f k a b. Applicative f
245+ => (Key s k -> a -> f b) -> HashMap s k a -> f (HashMap s k b)
246+traverseWithKey f (HashMap m)
247+ = HashMap <$> HashMap.traverseWithKey (f . unsafeKey) m
248+
249+-- | Map each key-value pair of a map into a monoid (with proof that the key was
250+-- in the map), and combine the results using '<>'.
251+foldMapWithKey
252+ :: forall s k a m. Monoid m => (Key s k -> a -> m) -> HashMap s k a -> m
253+foldMapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
254+ $ HashMap.foldMapWithKey @m @k @a
255+
256+-- | Right associative fold with a lazy accumulator.
257+foldrWithKey
258+ :: forall s k a b. (Key s k -> a -> b -> b) -> b -> HashMap s k a -> b
259+foldrWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
260+ $ HashMap.foldrWithKey @k @a @b
261+
262+-- | Left associative fold with a lazy accumulator.
263+foldlWithKey
264+ :: forall s k a b. (b -> Key s k -> a -> b) -> b -> HashMap s k a -> b
265+foldlWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
266+ $ HashMap.foldlWithKey @b @k @a
267+
268+-- | Right associative fold with a strict accumulator.
269+foldrWithKey'
270+ :: forall s k a b. (Key s k -> a -> b -> b) -> b -> HashMap s k a -> b
271+foldrWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce
272+ $ HashMap.foldrWithKey' @k @a @b
273+
274+-- | Left associative fold with a strict accumulator.
275+foldlWithKey'
276+ :: forall s k a b. (b -> Key s k -> a -> b) -> b -> HashMap s k a -> b
277+foldlWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce
278+ $ HashMap.foldlWithKey' @b @k @a
279+
280+-- | Return the set of keys in the map, with the contents of the set still
281+-- tracked by the @s@ parameter. See "Data.HashSet.Refined".
282+keysSet :: forall s k a. HashMap s k a -> HashSet s k
283+keysSet (HashMap m) = reify (HashMap.keysSet m)
284+ \(_ :: Proxy s') -> case unsafeCoerce Refl :: s :~: s' of
285+ Refl -> Dict
286+
287+-- | Convert to a list of key-value pairs.
288+toList :: forall s k a. HashMap s k a -> [(Key s k, a)]
289+toList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.toList @k @a
290+
291+-- | Retain only the key-value pairs that satisfy the predicate, returning a
292+-- potentially smaller map.
293+filterWithKey
294+ :: forall s k a. (Key s k -> a -> Bool)
295+ -> HashMap s k a
296+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
297+filterWithKey p (HashMap m)
298+ = SomeHashMapWith (HashMap $ HashMap.filterWithKey (p . unsafeKey) m)
299+ $ SupersetProof unsafeSubset
300+
301+-- | Restrict a map to only those keys that are elements of @t@.
302+restrictKeys
303+ :: forall s t k a. (Hashable k, KnownHashSet t k)
304+ => HashMap s k a -> SomeHashMapWith (IntersectionProof 'Hashed s t) k a
305+restrictKeys (HashMap m) = SomeHashMapWith
306+ (HashMap $ HashMap.intersectionWith const m
307+ $ HashSet.toMap $ reflect $ Proxy @t)
308+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
309+
310+-- | Remove all keys that are elements of @t@ from the map.
311+withoutKeys
312+ :: forall s t k a. (Hashable k, KnownHashSet t k)
313+ => HashMap s k a -> SomeHashMapWith (DifferenceProof 'Hashed s t) k a
314+withoutKeys (HashMap m) = SomeHashMapWith
315+ (HashMap $ HashMap.difference m $ HashSet.toMap $ reflect $ Proxy @t)
316+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
317+
318+-- | Partition a map into two disjoint submaps: those whose key-value pairs
319+-- satisfy the predicate, and those whose don't.
320+partitionWithKey
321+ :: forall s k a. Hashable k -- TODO: this is only used in the proof
322+ => (Key s k -> a -> Bool)
323+ -> HashMap s k a
324+ -> Some2HashMapWith (PartitionProof 'Hashed s k) k a a
325+partitionWithKey p (HashMap m) = Some2HashMapWith
326+ (HashMap $ HashMap.filterWithKey (p . unsafeKey) m)
327+ (HashMap $ HashMap.filterWithKey ((not .) . p . unsafeKey) m)
328+ $ PartitionProof
329+ do \k -> case HashMap.lookup (unrefine k) m of
330+ Nothing -> error
331+ "partitionWithKey: bug: Data.HashMap.Refined has been subverted"
332+ Just x -> if p k x
333+ then Left $ unsafeKey $ unrefine k
334+ else Right $ unsafeKey $ unrefine k
335+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
336+
337+-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then
338+-- @s@ and @t@ actually stand for the same set and @'Key' s@ can be safely
339+-- interconverted with @'Key' t@.
340+--
341+-- The requirement that the weakenings are natural transformations ensures that
342+-- they don't actually alter the keys. To build these you can compose ':->''s
343+-- from proofs returned by functions in this module, or "Refined" functions like
344+-- 'andLeft' or 'leftOr'.
345+castKey
346+ :: forall s t k. (forall x. Key s x -> Key t x)
347+ -> (forall x. Key t x -> Key s x)
348+ -> Coercion (Key s k) (Key t k)
349+castKey = castRefined
350+
351+-- | 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:
355+--
356+-- @
357+-- 'zipWithKey'
358+-- :: forall s k a b c. 'Hashable' k
359+-- => ('Key' s k -> a -> b -> c) -> 'HashMap' s k a -> 'HashMap' s k b -> 'HashMap' s k c
360+-- 'zipWithKey' f m1 m2
361+-- | v'SomeHashMapWith' @r m proof <- 'Data.HashMap.Refined.intersectionWithKey' (f . 'andLeft') m1 m2
362+-- , v'IntersectionProof' p1 p2 <- proof
363+-- , ( v'Coercion' :: t'Coercion' ('HashMap' r k c) ('HashMap' s k c))
364+-- <- app $ 'cast' $ 'castKey' ('andLeft' . p1) (p2 'id' 'id')
365+-- = 'coerce' m
366+-- where
367+-- app :: t'Coercion' f g -> t'Coercion' (f x) (g x)
368+-- app v'Coercion' = v'Coercion'
369+-- @
370+cast
371+ :: forall s t k. (forall x. Coercion (Key s x) (Key t x))
372+ -> Coercion (HashMap s k) (HashMap t k)
373+cast Coercion = Coercion
374+
375+instance FunctorWithIndex (Key s k) (HashMap s k) where
376+ imap = mapWithKey
377+
378+instance FoldableWithIndex (Key s k) (HashMap s k) where
379+ ifoldMap = foldMapWithKey
380+
381+instance TraversableWithIndex (Key s k) (HashMap s k) where
382+ itraverse = traverseWithKey
383+
384+-- | Similar to the instance for functions -- zip corresponding keys. To use
385+-- '<*>'/'liftA2' without 'KnownSet' see 'zipWithKey'.
386+instance (Hashable k, KnownHashSet s k) => Applicative (HashMap s k) where
387+ pure x = fromSet \_ -> x
388+ (<*>) = zipWithKey (const id)
389+
390+-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the
391+-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
392+bind
393+ :: forall s k a b. Hashable k
394+ => HashMap s k a -> (a -> HashMap s k b) -> HashMap s k b
395+bind m f = mapWithKey (\k x -> f x ! k) m
396+
397+-- | Similar to the instance for functions. To use '>>=' without 'KnownSet' see
398+-- 'bind'.
399+instance (Hashable k, KnownHashSet s k) => Monad (HashMap s k) where
400+ (>>=) = bind
401+
402+-- | Similar to the instance for functions. See also
403+-- 'Data.HashMap.Refined.backpermuteKeys'.
404+instance (Hashable k, KnownHashSet s k)
405+ => MonadReader (Key s k) (HashMap s k) where
406+ ask = fromSet id
407+ local f m = mapWithKey (\k _ -> m ! f k) m
408+
409+-- | Append the values at the corresponding keys
410+instance (Hashable k, Semigroup a) => Semigroup (HashMap s k a) where
411+ (<>) = zipWithKey (const (<>))
412+
413+instance (Hashable k, KnownHashSet s k, Monoid a)
414+ => Monoid (HashMap s k a) where
415+ mempty = fromSet \_ -> mempty
416+
417+-- | Similar to the instance for functions
418+instance (Hashable k, KnownHashSet s k) => Distributive (HashMap s k) where
419+ collect = collectRep
420+ distribute = distributeRep
421+
422+-- | Witness isomorphism with functions from @'Key' s k@
423+instance (Hashable k, KnownHashSet s k) => Representable (HashMap s k) where
424+ type Rep (HashMap s k) = Key s k
425+ index = (!)
426+ tabulate = fromSet
src/Data/HashMap/Refined.hsnew file+359
@@ -0,0 +1,359 @@
1+-- | This module defines a way to prove that a key exists in a map, so that the
2+-- key can be used to index into the map without using a 'Maybe', or manually
3+-- handling the \"impossible\" case with 'error' or other partial functions.
4+--
5+-- To do this, @'HashMap' s k v@ has a type parameter @s@ that identifies its
6+-- set of keys, so that if another map has the same type parameter, you know
7+-- that map has the same set of keys. There is @'Key' s k@, a type of keys that
8+-- have been validated to belong to the set identified by @s@, and for which the
9+-- operation of indexing into a @'HashMap' s k v@ (only for the same @s@) can
10+-- proceed without failure (see '!'). The type @s@ itself has no internal
11+-- structure, rather it is merely a skolem type variable (rank-2 polymorphism
12+-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
13+--
14+-- Like "Data.HashMap.Lazy", functions in this module are strict in the keys but
15+-- lazy in the values. The "Data.HashMap.Strict.Refined" module reuses the same
16+-- 'HashMap' type but provides functions that operate strictly on the values.
17+--
18+-- = Warning
19+-- This module together with "Data.HashMap.Lazy" rely on 'Eq' and 'Hashable'
20+-- instances being lawful: that '==' is an equivalence relation, and that
21+-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence
22+-- relation; at least for the subset of values that are actually encountered at
23+-- runtime. If this assumption is violated, this module may not be able to
24+-- uphold its invariants and may throw errors. In particular beware of NaN in
25+-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@.
26+module Data.HashMap.Refined
27+ (
28+ -- * Map type
29+ Common.HashMap
30+ , Common.Key
31+ -- * Existentials and common proofs
32+ , Common.SomeHashMap(..)
33+ , Common.withHashMap
34+ , Common.SomeHashMapWith(..)
35+ , Common.withHashMapWith
36+ , Common.Some2HashMapWith(..)
37+ , Common.with2HashMapWith
38+ , SupersetProof(..)
39+ , EmptyProof(..)
40+ -- * Construction
41+ , Common.empty
42+ , singleton
43+ , SingletonProof(..)
44+ , fromSet
45+ , Common.fromHashMap
46+ , fromTraversableWithKey
47+ , FromTraversableProof(..)
48+ -- * Insertion
49+ , insert
50+ , InsertProof(..)
51+ , reinsert
52+ , insertLookupWithKey
53+ -- * Deletion/Update
54+ , Common.delete
55+ , adjust
56+ , adjustWithKey
57+ , update
58+ , updateLookupWithKey
59+ -- * Query
60+ , Common.lookup
61+ , (Common.!)
62+ , Common.member
63+ , Common.null
64+ , Common.isSubmapOfBy
65+ , SubsetProof(..)
66+ , Common.disjoint
67+ , DisjointProof(..)
68+ -- * Combine
69+ , zipWithKey
70+ , bind
71+ , unionWithKey
72+ , UnionProof(..)
73+ , Common.difference
74+ , DifferenceProof(..)
75+ , differenceWithKey
76+ , PartialDifferenceProof(..)
77+ , intersectionWithKey
78+ , IntersectionProof(..)
79+ -- * Traversal
80+ , mapWithKey
81+ , traverseWithKey
82+ , mapAccumLWithKey
83+ , mapAccumRWithKey
84+ , mapKeysWith
85+ , MapProof(..)
86+ , backpermuteKeys
87+ -- * Folds
88+ , Common.foldMapWithKey
89+ , Common.foldrWithKey
90+ , Common.foldlWithKey
91+ , Common.foldrWithKey'
92+ , Common.foldlWithKey'
93+ -- * Conversion
94+ , Common.toMap
95+ , Common.keysSet
96+ , Common.toList
97+ -- * Filter
98+ , Common.restrictKeys
99+ , Common.withoutKeys
100+ , Common.filterWithKey
101+ , Common.partitionWithKey
102+ , PartitionProof(..)
103+ , mapMaybeWithKey
104+ , mapEitherWithKey
105+ -- * Casts
106+ , Common.castKey
107+ , Common.cast
108+ , castFlavor
109+ ) where
110+
111+import Data.Coerce
112+import Data.Container.Refined.Hashable
113+import Data.Container.Refined.Proofs
114+import Data.Container.Refined.Unsafe
115+import Data.Functor
116+import Data.HashMap.Common.Refined
117+ ( HashMap(..), Key, unsafeCastKey, unsafeKey, SomeHashMapWith(..)
118+ , Some2HashMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey
119+ , bind
120+ )
121+import qualified Data.HashMap.Common.Refined as Common
122+import qualified Data.HashMap.Lazy as HashMap
123+import Data.Traversable
124+import Data.Traversable.WithIndex
125+import Data.Type.Coercion
126+import Prelude hiding (lookup, null)
127+import Refined
128+import Refined.Unsafe
129+
130+
131+-- | Create a map with a single key-value pair, and return a proof that the key
132+-- is in the resulting map.
133+singleton
134+ :: forall k a. Hashable k
135+ => k -> a -> SomeHashMapWith (SingletonProof 'Hashed k) k a
136+singleton k v = SomeHashMapWith (HashMap $ HashMap.singleton k v)
137+ $ SingletonProof (unsafeKey k)
138+
139+-- | Create a map from an arbitrary traversable of key-value pairs.
140+fromTraversableWithKey
141+ :: forall t k a. (Traversable t, Hashable k)
142+ => (k -> a -> a -> a)
143+ -> t (k, a)
144+ -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
145+fromTraversableWithKey f xs
146+ = SomeHashMapWith (HashMap m) $ FromTraversableProof proof
147+ where
148+ (m, proof) = mapAccumL
149+ (\s (k, v)
150+ -> let !s' = HashMap.insertWith (f k) k v s in (s', unsafeKey k))
151+ HashMap.empty
152+ xs
153+
154+-- | Insert a key-value pair into the map to obtain a potentially larger map,
155+-- guaranteed to contain the given key. If the key was already present, the
156+-- associated value is replaced with the supplied value.
157+insert
158+ :: forall s k a. Hashable k
159+ => k -> a -> HashMap s k a -> SomeHashMapWith (InsertProof 'Hashed k s) k a
160+insert k v (HashMap m) = SomeHashMapWith (HashMap $ HashMap.insert k v m)
161+ $ InsertProof (unsafeKey k) unsafeSubset
162+
163+-- | Overwrite a key-value pair that is known to already be in the map. The set
164+-- of keys remains the same.
165+reinsert
166+ :: forall s k a. Hashable k
167+ => Key s k -> a -> HashMap s k a -> HashMap s k a
168+reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.insert @k @a
169+
170+-- | Insert a key-value pair into the map using a combining function, and if
171+-- the key was already present, the old value is returned along with the proof
172+-- that the key was present.
173+insertLookupWithKey
174+ :: forall s k a. Hashable k
175+ => (Key s k -> a -> a -> a)
176+ -> k
177+ -> a
178+ -> HashMap s k a
179+ -> (Maybe (Key s k, a), SomeHashMapWith (InsertProof 'Hashed k s) k a)
180+insertLookupWithKey f k v (HashMap m) =
181+ ( (unsafeKey k,) <$> HashMap.lookup k m
182+ , SomeHashMapWith (HashMap $ HashMap.insertWith (f $ unsafeKey k) k v m)
183+ $ InsertProof (unsafeKey k) unsafeSubset
184+ )
185+
186+-- | Update the value at a specific key known the be in the map using the given
187+-- function. The set of keys remains the same.
188+adjust
189+ :: forall s k a. Hashable k
190+ => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a
191+adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.adjust @k @a
192+
193+-- | If the given key is in the map, update the associated value using the given
194+-- function with a proof that the key was in the map; otherwise return the map
195+-- unchanged. In any case the set of keys remains the same.
196+adjustWithKey
197+ :: forall s k a. Hashable k
198+ => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a
199+adjustWithKey f k (HashMap m) = HashMap $ HashMap.adjust (f $ unsafeKey k) k m
200+
201+-- | Update or delete a key known to be in the map using the given function,
202+-- returning a potentially smaller map.
203+update
204+ :: forall s k a. Hashable k
205+ => (a -> Maybe a)
206+ -> Key s k
207+ -> HashMap s k a
208+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
209+update f k (HashMap m)
210+ = SomeHashMapWith (HashMap $ HashMap.update f (unrefine k) m)
211+ $ SupersetProof unsafeSubset
212+
213+-- | If the given key is in the map, update or delete it using the given
214+-- function with a proof that the key was in the map; otherwise the map is
215+-- unchanged. Alongside return the new value if it was updated, or the old value
216+-- if it was deleted, and a proof that the key was in the map.
217+updateLookupWithKey
218+ :: forall s k a. Hashable k
219+ => (Key s k -> a -> Maybe a)
220+ -> k
221+ -> HashMap s k a
222+ -> (Maybe (Key s k, a), SomeHashMapWith (SupersetProof 'Hashed s) k a)
223+updateLookupWithKey f k (HashMap m) =
224+ ( (unsafeKey k,) <$> HashMap.lookup k m
225+ , SomeHashMapWith (HashMap $ HashMap.update (f $ unsafeKey k) k m)
226+ $ SupersetProof unsafeSubset
227+ )
228+
229+-- | Return the union of two maps, with a given combining function for keys that
230+-- exist in both maps simultaneously.
231+--
232+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
233+-- respectively.
234+unionWithKey
235+ :: forall s t k a. Hashable k
236+ => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a)
237+ -> HashMap s k a
238+ -> HashMap t k a
239+ -> SomeHashMapWith (UnionProof 'Hashed s t) k a
240+unionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
241+ (HashMap $ HashMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
242+ $ UnionProof unsafeSubset unsafeSubsetWith2
243+
244+-- | For keys that appear in both maps, the given function decides whether the
245+-- key is removed from the first map.
246+--
247+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
248+-- respectively.
249+differenceWithKey
250+ :: forall s t k a b. Hashable k
251+ => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a)
252+ -> HashMap s k a
253+ -> HashMap t k b
254+ -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a
255+differenceWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
256+ (HashMap $ HashMap.differenceWith
257+ (\x (k, y) -> f (reallyUnsafeRefine k) x y)
258+ m1
259+ (HashMap.mapWithKey (,) m2))
260+ $ PartialDifferenceProof unsafeSubset unsafeSubset
261+
262+-- | Return the intersection of two maps with the given combining function.
263+--
264+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
265+-- respectively.
266+intersectionWithKey
267+ :: forall s t k a b c. Hashable k
268+ => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c)
269+ -> HashMap s k a
270+ -> HashMap t k b
271+ -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c
272+intersectionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
273+ (HashMap $ HashMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
274+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
275+
276+-- | Thread an accumularing argument through the map in ascending order of
277+-- hashes.
278+mapAccumLWithKey
279+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
280+ -> a
281+ -> HashMap s k b
282+ -> (a, HashMap s k c)
283+mapAccumLWithKey f = imapAccumL (flip f)
284+
285+-- | Thread an accumularing argument through the map in descending order of
286+-- hashes.
287+mapAccumRWithKey
288+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
289+ -> a
290+ -> HashMap s k b
291+ -> (a, HashMap s k c)
292+mapAccumRWithKey f = imapAccumR (flip f)
293+
294+-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
295+-- results into a new map. For keys that were mapped to the same new key, @c@
296+-- acts as the combining function for corresponding values.
297+mapKeysWith
298+ :: forall s k1 k2 a. Hashable k2
299+ => (a -> a -> a)
300+ -> (Key s k1 -> k2)
301+ -> HashMap s k1 a
302+ -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a
303+mapKeysWith f g (HashMap m) = SomeHashMapWith
304+ (HashMap $ HashMap.fromListWith f
305+ $ HashMap.foldrWithKey (\k x xs -> (g $ unsafeKey k, x) : xs) [] m)
306+ $ MapProof (unsafeKey . g) \k2 ->
307+ case HashMap.lookup (unrefine k2) backMap of
308+ Nothing -> error
309+ "mapKeysWith: bug: Data.HashMap.Refined has been subverted"
310+ Just k1 -> k1
311+ where
312+ ~backMap = HashMap.fromList
313+ [ (k2, unsafeKey k1)
314+ | k1 <- HashMap.keys m
315+ , let !k2 = g $ unsafeKey k1
316+ ]
317+
318+-- | Apply a function to all values in a map, together with their corresponding
319+-- keys, and collect only the 'Just' results, returning a potentially smaller
320+-- map.
321+mapMaybeWithKey
322+ :: forall s k a b. (Key s k -> a -> Maybe b)
323+ -> HashMap s k a
324+ -> SomeHashMapWith (SupersetProof 'Hashed s) k b
325+mapMaybeWithKey f (HashMap m)
326+ = SomeHashMapWith (HashMap $ HashMap.mapMaybeWithKey (f . unsafeKey) m)
327+ $ SupersetProof unsafeSubset
328+
329+-- | Apply a function to all values in a map, together with their corresponding
330+-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
331+-- maps.
332+mapEitherWithKey
333+ :: forall s k a b c. Hashable k -- TODO: this is only used in the proof
334+ => (Key s k -> a -> Either b c)
335+ -> HashMap s k a
336+ -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c
337+mapEitherWithKey p (HashMap m)
338+ | m' <- HashMap.mapWithKey (p . unsafeKey) m
339+ = Some2HashMapWith
340+ (HashMap $ HashMap.mapMaybe (either Just (const Nothing)) m')
341+ (HashMap $ HashMap.mapMaybe (either (const Nothing) Just) m')
342+ $ PartitionProof
343+ do \k -> case HashMap.lookup (unrefine k) m of
344+ Nothing -> error
345+ "mapEitherWithKey: bug: Data.HashMap.Refined has been subverted"
346+ Just x -> case p k x of
347+ Left _ -> Left $ unsafeKey $ unrefine k
348+ Right _ -> Right $ unsafeKey $ unrefine k
349+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
350+
351+-- | Apply the inverse image of the given function to the keys of the given map,
352+-- so that for all @k :: 'Key' s2 k2@,c
353+-- @'backpermuteKeys' f m '!' k = m '!' f k@.
354+--
355+-- If maps are identified with functions, this computes the composition.
356+backpermuteKeys
357+ :: forall s1 s2 k1 k2 a. (Hashable k1, KnownHashSet s2 k2)
358+ => (Key s2 k2 -> Key s1 k1) -> HashMap s1 k1 a -> HashMap s2 k2 a
359+backpermuteKeys f m = fromSet \k -> m ! f k
src/Data/HashMap/Strict/Refined.hsnew file+398
@@ -0,0 +1,398 @@
1+-- | This module defines a way to prove that a key exists in a map, so that the
2+-- key can be used to index into the map without using a 'Maybe', or manually
3+-- handling the \"impossible\" case with 'error' or other partial functions.
4+--
5+-- To do this, @'HashMap' s k v@ has a type parameter @s@ that identifies its
6+-- set of keys, so that if another map has the same type parameter, you know
7+-- that map has the same set of keys. There is @'Key' s k@, a type of keys that
8+-- have been validated to belong to the set identified by @s@, and for which the
9+-- operation of indexing into a @'HashMap' s k v@ (only for the same @s@) can
10+-- proceed without failure (see '!'). The type @s@ itself has no internal
11+-- structure, rather it is merely a skolem type variable (rank-2 polymorphism
12+-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
13+--
14+-- Like "Data.HashMap.Strict", functions in this module are strict in the keys
15+-- and values. The "Data.HashMap.Refined" module reuses the same 'HashMap' type
16+-- but provides functions that operate lazily on the values.
17+--
18+-- = Warning
19+-- This module together with "Data.HashMap.Lazy" rely on 'Eq' and 'Hashable'
20+-- instances being lawful: that '==' is an equivalence relation, and that
21+-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence
22+-- relation; at least for the subset of values that are actually encountered at
23+-- runtime. If this assumption is violated, this module may not be able to
24+-- uphold its invariants and may throw errors. In particular beware of NaN in
25+-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@.
26+module Data.HashMap.Strict.Refined
27+ (
28+ -- * Map type
29+ Common.HashMap
30+ , Common.Key
31+ -- * Existentials and common proofs
32+ , Common.SomeHashMap(..)
33+ , Common.withHashMap
34+ , Common.SomeHashMapWith(..)
35+ , Common.withHashMapWith
36+ , Common.Some2HashMapWith(..)
37+ , Common.with2HashMapWith
38+ , SupersetProof(..)
39+ , EmptyProof(..)
40+ -- * Construction
41+ , Common.empty
42+ , singleton
43+ , SingletonProof(..)
44+ , fromSet
45+ , Common.fromHashMap
46+ , fromTraversableWithKey
47+ , FromTraversableProof(..)
48+ -- * Insertion
49+ , insert
50+ , InsertProof(..)
51+ , reinsert
52+ , insertLookupWithKey
53+ -- * Deletion/Update
54+ , Common.delete
55+ , adjust
56+ , adjustWithKey
57+ , update
58+ , updateLookupWithKey
59+ -- * Query
60+ , Common.lookup
61+ , (Common.!)
62+ , Common.member
63+ , Common.null
64+ , Common.isSubmapOfBy
65+ , SubsetProof(..)
66+ , Common.disjoint
67+ , DisjointProof(..)
68+ -- * Combine
69+ , zipWithKey
70+ , bind
71+ , unionWithKey
72+ , UnionProof(..)
73+ , Common.difference
74+ , DifferenceProof(..)
75+ , differenceWithKey
76+ , PartialDifferenceProof(..)
77+ , intersectionWithKey
78+ , IntersectionProof(..)
79+ -- * Traversal
80+ , mapWithKey
81+ , traverseWithKey
82+ , mapAccumLWithKey
83+ , mapAccumRWithKey
84+ , mapKeysWith
85+ , MapProof(..)
86+ , backpermuteKeys
87+ -- * Folds
88+ , Common.foldMapWithKey
89+ , Common.foldrWithKey
90+ , Common.foldlWithKey
91+ , Common.foldrWithKey'
92+ , Common.foldlWithKey'
93+ -- * Conversion
94+ , Common.toMap
95+ , Common.keysSet
96+ , Common.toList
97+ -- * Filter
98+ , Common.restrictKeys
99+ , Common.withoutKeys
100+ , Common.filterWithKey
101+ , Common.partitionWithKey
102+ , PartitionProof(..)
103+ , mapMaybeWithKey
104+ , mapEitherWithKey
105+ -- * Casts
106+ , Common.castKey
107+ , Common.cast
108+ , castFlavor
109+ ) where
110+
111+import Data.Coerce
112+import Data.Container.Refined.Hashable
113+import Data.Container.Refined.Proofs
114+import Data.Container.Refined.Unsafe
115+import Data.Functor
116+import qualified Data.HashMap.Strict as HashMap
117+import Data.HashMap.Common.Refined
118+ ( HashMap(..), Key, unsafeCastKey, unsafeKey, SomeHashMapWith(..)
119+ , Some2HashMapWith(..), (!)
120+ )
121+import qualified Data.HashMap.Common.Refined as Common
122+import qualified Data.HashSet as HashSet
123+import Data.Proxy
124+import Data.Reflection
125+import Data.Traversable
126+import Data.Traversable.WithIndex
127+import Data.Type.Coercion
128+import Prelude hiding (lookup, null)
129+import Refined
130+import Refined.Unsafe
131+
132+
133+-- | Create a map with a single key-value pair, and return a proof that the
134+-- key is in the resulting map.
135+singleton
136+ :: forall k a. Hashable k
137+ => k -> a -> SomeHashMapWith (SingletonProof 'Hashed k) k a
138+singleton k v = SomeHashMapWith (HashMap $ HashMap.singleton k v)
139+ $ SingletonProof (unsafeKey k)
140+
141+-- | Create a map from a set of keys, and a function that for each key computes
142+-- the corresponding value.
143+fromSet :: forall s k a. KnownHashSet s k => (Key s k -> a) -> HashMap s k a
144+fromSet f = HashMap $ HashMap.mapWithKey (\k _ -> f $ unsafeKey k)
145+ $ HashSet.toMap (reflect $ Proxy @s)
146+
147+-- | Create a map from an arbitrary traversable of key-value pairs.
148+fromTraversableWithKey
149+ :: forall t k a. (Traversable t, Hashable k)
150+ => (k -> a -> a -> a)
151+ -> t (k, a)
152+ -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a
153+fromTraversableWithKey f xs = SomeHashMapWith (HashMap m)
154+ $ FromTraversableProof proof
155+ where
156+ (m, proof) = mapAccumL
157+ (\s (k, v)
158+ -> let !s' = HashMap.insertWith (f k) k v s in (s', unsafeKey k))
159+ HashMap.empty
160+ xs
161+
162+-- | Insert a key-value pair into the map to obtain a potentially larger map,
163+-- guaranteed to contain the given key. If the key was already present, the
164+-- associated value is replaced with the supplied value.
165+insert
166+ :: forall s k a. Hashable k
167+ => k -> a -> HashMap s k a -> SomeHashMapWith (InsertProof 'Hashed k s) k a
168+insert k v (HashMap m) = SomeHashMapWith (HashMap $ HashMap.insert k v m)
169+ $ InsertProof (unsafeKey k) unsafeSubset
170+
171+-- | Overwrite a key-value pair that is known to already be in the map. The set
172+-- of keys remains the same.
173+reinsert
174+ :: forall s k a. Hashable k
175+ => Key s k -> a -> HashMap s k a -> HashMap s k a
176+reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.insert @k @a
177+
178+-- | Insert a key-value pair into the map using a combining function, and if
179+-- the key was already present, the old value is returned along with the proof
180+-- that the key was present.
181+insertLookupWithKey
182+ :: forall s k a. Hashable k
183+ => (Key s k -> a -> a -> a)
184+ -> k
185+ -> a
186+ -> HashMap s k a
187+ -> (Maybe (Key s k, a), SomeHashMapWith (InsertProof 'Hashed k s) k a)
188+insertLookupWithKey f k v (HashMap m) =
189+ ( (unsafeKey k,) <$> HashMap.lookup k m
190+ , SomeHashMapWith (HashMap $ HashMap.insertWith (f $ unsafeKey k) k v m)
191+ $ InsertProof (unsafeKey k) unsafeSubset
192+ )
193+
194+-- | Update the value at a specific key known the be in the map using the given
195+-- function. The set of keys remains the same.
196+adjust
197+ :: forall s k a. Hashable k
198+ => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a
199+adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.adjust @k @a
200+
201+-- | If the given key is in the map, update the associated value using the given
202+-- function with a proof that the key was in the map; otherwise return the map
203+-- unchanged. In any case the set of keys remains the same.
204+adjustWithKey
205+ :: forall s k a. Hashable k
206+ => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a
207+adjustWithKey f k (HashMap m) = HashMap $ HashMap.adjust (f $ unsafeKey k) k m
208+
209+-- | Update or delete a key known to be in the map using the given function,
210+-- returning a potentially smaller map.
211+update
212+ :: forall s k a. Hashable k
213+ => (a -> Maybe a)
214+ -> Key s k
215+ -> HashMap s k a
216+ -> SomeHashMapWith (SupersetProof 'Hashed s) k a
217+update f k (HashMap m)
218+ = SomeHashMapWith (HashMap $ HashMap.update f (unrefine k) m)
219+ $ SupersetProof unsafeSubset
220+
221+-- | If the given key is in the map, update or delete it using the given
222+-- function with a proof that the key was in the map; otherwise the map is
223+-- unchanged. Alongside return the new value if it was updated, or the old value
224+-- if it was deleted, and a proof that the key was in the map.
225+updateLookupWithKey
226+ :: forall s k a. Hashable k
227+ => (Key s k -> a -> Maybe a)
228+ -> k
229+ -> HashMap s k a
230+ -> (Maybe (Key s k, a), SomeHashMapWith (SupersetProof 'Hashed s) k a)
231+updateLookupWithKey f k (HashMap m) =
232+ ( (unsafeKey k,) <$> HashMap.lookup k m
233+ , SomeHashMapWith (HashMap $ HashMap.update (f $ unsafeKey k) k m)
234+ $ SupersetProof unsafeSubset
235+ )
236+
237+-- | Given two maps proven to have the same keys, for each key apply the
238+-- function to the associated values, to obtain a new map with the same keys.
239+zipWithKey
240+ :: forall s k a b c. Hashable k
241+ => (Key s k -> a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c
242+zipWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
243+ $ HashMap.intersectionWithKey @k @a @b @c
244+
245+-- | Return the union of two maps, with a given combining function for keys that
246+-- exist in both maps simultaneously.
247+--
248+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
249+-- respectively.
250+unionWithKey
251+ :: forall s t k a. Hashable k
252+ => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a)
253+ -> HashMap s k a
254+ -> HashMap t k a
255+ -> SomeHashMapWith (UnionProof 'Hashed s t) k a
256+unionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
257+ (HashMap $ HashMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
258+ $ UnionProof unsafeSubset unsafeSubsetWith2
259+
260+-- | For keys that appear in both maps, the given function decides whether the
261+-- key is removed from the first map.
262+--
263+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
264+-- respectively.
265+differenceWithKey
266+ :: forall s t k a b. Hashable k
267+ => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a)
268+ -> HashMap s k a
269+ -> HashMap t k b
270+ -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a
271+differenceWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
272+ (HashMap $ HashMap.differenceWith
273+ (\x (k, y) -> f (reallyUnsafeRefine k) x y)
274+ m1
275+ (HashMap.mapWithKey (,) m2))
276+ $ PartialDifferenceProof unsafeSubset unsafeSubset
277+
278+-- | Return the intersection of two maps with the given combining function.
279+--
280+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
281+-- respectively.
282+intersectionWithKey
283+ :: forall s t k a b c. Hashable k
284+ => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c)
285+ -> HashMap s k a
286+ -> HashMap t k b
287+ -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c
288+intersectionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith
289+ (HashMap $ HashMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
290+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
291+
292+-- | Apply a function to all values in a map, together with their corresponding
293+-- keys, that are proven to be in the map. The set of keys remains the same.
294+mapWithKey
295+ :: forall s k a b. (Key s k -> a -> b) -> HashMap s k a -> HashMap s k b
296+mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
297+ $ HashMap.mapWithKey @k @a @b
298+
299+-- | Map an 'Applicative' transformation with access to each value's
300+-- corresponding key and a proof that it is in the map. The set of keys remains
301+-- unchanged.
302+traverseWithKey
303+ :: forall s f k a b. Applicative f
304+ => (Key s k -> a -> f b) -> HashMap s k a -> f (HashMap s k b)
305+traverseWithKey f (HashMap m)
306+ = HashMap <$> HashMap.traverseWithKey (f . unsafeKey) m
307+
308+-- | Thread an accumularing argument through the map in ascending order of
309+-- hashes.
310+mapAccumLWithKey
311+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
312+ -> a
313+ -> HashMap s k b
314+ -> (a, HashMap s k c)
315+mapAccumLWithKey f = imapAccumL (flip f)
316+
317+-- | Thread an accumularing argument through the map in descending order of
318+-- hashes.
319+mapAccumRWithKey
320+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
321+ -> a
322+ -> HashMap s k b
323+ -> (a, HashMap s k c)
324+mapAccumRWithKey f = imapAccumR (flip f)
325+
326+-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
327+-- results into a new map. For keys that were mapped to the same new key, @c@
328+-- acts as the combining function for corresponding values.
329+mapKeysWith
330+ :: forall s k1 k2 a. Hashable k2
331+ => (a -> a -> a)
332+ -> (Key s k1 -> k2)
333+ -> HashMap s k1 a
334+ -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a
335+mapKeysWith f g (HashMap m) = SomeHashMapWith
336+ (HashMap $ HashMap.fromListWith f
337+ $ HashMap.foldrWithKey (\k x xs -> (g $ unsafeKey k, x) : xs) [] m)
338+ $ MapProof (unsafeKey . g) \k2 ->
339+ case HashMap.lookup (unrefine k2) backMap of
340+ Nothing -> error
341+ "mapKeysWith: bug: Data.HashMap.Refined has been subverted"
342+ Just k1 -> k1
343+ where
344+ ~backMap = HashMap.fromList
345+ [ (k2, unsafeKey k1)
346+ | k1 <- HashMap.keys m
347+ , let !k2 = g $ unsafeKey k1
348+ ]
349+
350+-- | Apply a function to all values in a map, together with their corresponding
351+-- keys, and collect only the 'Just' results, returning a potentially smaller
352+-- map.
353+mapMaybeWithKey
354+ :: forall s k a b. (Key s k -> a -> Maybe b)
355+ -> HashMap s k a
356+ -> SomeHashMapWith (SupersetProof 'Hashed s) k b
357+mapMaybeWithKey f (HashMap m)
358+ = SomeHashMapWith (HashMap $ HashMap.mapMaybeWithKey (f . unsafeKey) m)
359+ $ SupersetProof unsafeSubset
360+
361+-- | Apply a function to all values in a map, together with their corresponding
362+-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
363+-- maps.
364+mapEitherWithKey
365+ :: forall s k a b c. Hashable k -- TODO: this is only used in the proof
366+ => (Key s k -> a -> Either b c)
367+ -> HashMap s k a
368+ -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c
369+mapEitherWithKey p (HashMap m)
370+ | m' <- HashMap.mapWithKey (p . unsafeKey) m
371+ = Some2HashMapWith
372+ (HashMap $ HashMap.mapMaybe (either Just (const Nothing)) m')
373+ (HashMap $ HashMap.mapMaybe (either (const Nothing) Just) m')
374+ $ PartitionProof
375+ do \k -> case HashMap.lookup (unrefine k) m of
376+ Nothing -> error
377+ "mapEitherWithKey: bug: Data.HashMap.Refined has been subverted"
378+ Just x -> case p k x of
379+ Left _ -> Left $ unsafeKey $ unrefine k
380+ Right _ -> Right $ unsafeKey $ unrefine k
381+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
382+
383+-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the
384+-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
385+bind
386+ :: forall s k a b. Hashable k
387+ => HashMap s k a -> (a -> HashMap s k b) -> HashMap s k b
388+bind m f = mapWithKey (\k x -> f x ! k) m
389+
390+-- | Apply the inverse image of the given function to the keys of the given map,
391+-- so that for all @k :: 'Key' s2 k2@,
392+-- @'backpermuteKeys' f m '!' k = m '!' f k@.
393+--
394+-- If maps are identified with functions, this computes the composition.
395+backpermuteKeys
396+ :: forall s1 s2 k1 k2 a. (Hashable k1, KnownHashSet s2 k2)
397+ => (Key s2 k2 -> Key s1 k1) -> HashMap s1 k1 a -> HashMap s2 k2 a
398+backpermuteKeys f m = fromSet \k -> m ! f k
src/Data/HashSet/Refined.hsnew file+418
@@ -0,0 +1,418 @@
1+{-# LANGUAGE CPP #-}
2+-- | This module implements a way of tracking the contents of a
3+-- 'Data.HashSet.HashSet' at the type level, and provides utilities for
4+-- manipulating such sets.
5+--
6+-- The contents of a set are associated with a type parameter, e.g. @s@, so that
7+-- whenever you see the same type parameter, you know you are working with the
8+-- same set. The type @s@ itself has no internal structure, rather it is merely
9+-- a skolem type variable (rank-2 polymorphism 'Control.Monad.ST.runST' trick)
10+-- introduced by "Data.Reflection".
11+--
12+-- = Warning
13+-- This module together with "Data.HashSet" rely on 'Eq' and 'Hashable'
14+-- instances being lawful: that '==' is an equivalence relation, and that
15+-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence
16+-- relation; at least for the subset of values that are actually encountered at
17+-- runtime. If this assumption is violated, this module may not be able to
18+-- uphold its invariants and may throw errors. In particular beware of NaN in
19+-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@.
20+module Data.HashSet.Refined
21+ (
22+ -- * Set type
23+ KnownHashSet
24+ , HashSet
25+ -- * Refinement type
26+ , InSet(..)
27+ , Flavor(Hashed)
28+ , Element
29+ , revealPredicate
30+ -- * Existentials and common proofs
31+ , SomeHashSet(..)
32+ , withHashSet
33+ , SomeHashSetWith(..)
34+ , withHashSetWith
35+ , Some2HashSetWith(..)
36+ , with2HashSetWith
37+ , (:->)
38+ , SupersetProof(..)
39+ , EmptyProof(..)
40+ -- * Construction
41+ , empty
42+ , singleton
43+ , SingletonProof(..)
44+ , fromHashSet
45+ , fromTraversable
46+ , FromTraversableProof(..)
47+ -- * Insertion
48+ , insert
49+ , InsertProof(..)
50+ -- * Deletion
51+ , delete
52+ -- * Query
53+ , member
54+ , null
55+ , isSubsetOf
56+ , SubsetProof(..)
57+ , disjoint
58+ , DisjointProof(..)
59+ -- * Combine
60+ , union
61+ , UnionProof(..)
62+ , difference
63+ , DifferenceProof(..)
64+ , intersection
65+ , IntersectionProof(..)
66+ -- * Filter
67+ , filter
68+ , partition
69+ , PartitionProof(..)
70+ -- * Map
71+ , map
72+ , MapProof(..)
73+ -- * Folds
74+ , foldMap
75+ , foldr
76+ , foldl
77+ , foldr'
78+ , foldl'
79+ -- * Conversion
80+ , toList
81+ , asSet
82+ , asIntSet
83+ -- * Casts
84+ , castElement
85+ , cast
86+ , castFlavor
87+ ) where
88+
89+import Data.Coerce
90+import Data.Constraint (Dict(..))
91+import Data.Container.Refined.Conversion
92+import Data.Container.Refined.Hashable
93+import Data.Container.Refined.Proofs
94+import Data.Container.Refined.Unsafe
95+import qualified Data.Foldable as Foldable
96+import qualified Data.HashMap.Lazy as HashMap
97+import qualified Data.HashSet as HashSet
98+import Data.Proxy
99+import Data.Reflection
100+import Data.Traversable
101+import Data.Type.Coercion
102+import Data.Type.Equality ((:~:)(..))
103+import Data.Typeable (Typeable)
104+import GHC.Exts (Proxy#, proxy#)
105+import Prelude hiding (filter, foldl, foldMap, foldr, map, null)
106+import Refined
107+import Refined.Unsafe
108+import Unsafe.Coerce
109+
110+
111+-- | To use "Refined" machinery that uses the 'Predicate' typeclass you will
112+-- need to pattern match on this 'Dict'.
113+--
114+-- The reason is that in the default /fast/ implementation of reflection, we
115+-- don't have @'Typeable' s@, which "Refined" wants for pretty-printing
116+-- exceptions. We /can/ provide @'Typeable' s@, but at the cost of using the
117+-- /slow/ implementation of reflection.
118+revealPredicate
119+ :: forall s a. (Typeable a, Hashable a, KnownHashSet s a)
120+ => Dict (Predicate (InSet 'Hashed s) a)
121+revealPredicate = reifyTypeable (reflect (Proxy @s))
122+ \(_ :: Proxy s') ->
123+ reflect (Proxy @s') `seq`
124+ -- ^ Work around https://github.com/ekmett/reflection/issues/54
125+ case unsafeCoerce Refl :: s :~: s' of
126+ Refl -> Dict
127+
128+-- | @'Element' s a@ is a value of type @a@ that has been verified to be an
129+-- element of @s@.
130+--
131+-- Thus, @'Element' s a@ is a \"refinement\" type of @a@, and this library
132+-- integrates with an implementation of refimenement types in "Refined", so
133+-- the machinery from there can be used to manipulate 'Element's (however see
134+-- 'revealPredicate').
135+--
136+-- The underlying @a@ value can be obtained with 'unrefine'. An @a@ can be
137+-- validated into an @'Element' s a@ with 'member'.
138+type Element s = Refined (InSet 'Hashed s)
139+
140+unsafeCastElement :: forall s a. Coercion a (Element s a)
141+unsafeCastElement = reallyUnsafeUnderlyingRefined
142+
143+unsafeElement :: a -> Element s a
144+unsafeElement = coerceWith unsafeCastElement
145+
146+-- | An existential wrapper for an as-yet-unknown set. Pattern maching on it
147+-- gives you a way to refer to the set, e.g.
148+--
149+-- @
150+-- case 'fromHashSet' ... of
151+-- 'SomeHashSet' \@s _ -> doSomethingWith \@s
152+--
153+-- case 'fromHashSet' ... of
154+-- 'SomeHashSet' (_ :: 'Proxy#' s) -> doSomethingWith \@s
155+-- @
156+data SomeHashSet a where
157+ SomeHashSet :: forall s a. KnownHashSet s a => Proxy# s -> SomeHashSet a
158+
159+-- | Apply an unknown set to a continuation that can accept any set. This gives
160+-- you a way to refer to the set (the parameter @s@), e.g.:
161+--
162+-- @
163+-- 'withHashSet' ('fromHashSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s
164+-- @
165+withHashSet
166+ :: forall a r. SomeHashSet a
167+ -> (forall s. KnownHashSet s a => Proxy s -> r)
168+ -> r
169+withHashSet (SomeHashSet (_ :: Proxy# s)) k = k $ Proxy @s
170+
171+-- | Construct a set from a regular 'Data.HashSet.HashSet'.
172+fromHashSet :: forall a. HashSet.HashSet a -> SomeHashSet a
173+fromHashSet s = reify s \(_ :: Proxy s) -> SomeHashSet @s proxy#
174+
175+-- | An existential wrapper for an as-yet-unknown set, together with a proof of
176+-- some fact @p@ about the set. Pattern matching on it gives you a way to refer
177+-- to the set (the parameter @s@). Most functions will return a set in this way,
178+-- together with a proof that somehow relates the set to the function's inputs.
179+data SomeHashSetWith p a where
180+ SomeHashSetWith
181+ :: forall s a p. KnownHashSet s a => !(p s) -> SomeHashSetWith p a
182+
183+-- | Apply an unknown set with proof to a continuation that can accept any set
184+-- satisfying the proof. This gives you a way to refer to the set (the parameter
185+-- @s@).
186+withHashSetWith
187+ :: forall a r p. SomeHashSetWith p a
188+ -> (forall s. KnownHashSet s a => p s -> r)
189+ -> r
190+withHashSetWith (SomeHashSetWith p) k = k p
191+
192+-- | An existential wrapper for an as-yet-unknown pair of sets, together with
193+-- a proof of some fact @p@ relating them.
194+data Some2HashSetWith p a where
195+ Some2HashSetWith
196+ :: forall s t a p. (KnownHashSet s a, KnownHashSet t a)
197+ => !(p s t) -> Some2HashSetWith p a
198+
199+-- | Apply a pair of unknown sets with proof to a continuation that can accept
200+-- any pair of sets satisfying the proof. This gives you a way to refer to the
201+-- sets (the parameters @s@ and @t@).
202+with2HashSetWith
203+ :: forall a r p. Some2HashSetWith p a
204+ -> (forall s t. (KnownHashSet s a, KnownHashSet t a) => p s t -> r)
205+ -> r
206+with2HashSetWith (Some2HashSetWith p) k = k p
207+
208+-- | An empty set.
209+empty :: forall a. SomeHashSetWith (EmptyProof 'Hashed) a
210+empty = reify HashSet.empty \(_ :: Proxy r)
211+ -> SomeHashSetWith @r $ EmptyProof unsafeSubset
212+
213+-- | Create a set with a single element.
214+singleton
215+ :: forall a. Hashable a => a -> SomeHashSetWith (SingletonProof 'Hashed a) a
216+singleton x = reify (HashSet.singleton x) \(_ :: Proxy r)
217+ -> SomeHashSetWith @r $ SingletonProof $ unsafeElement x
218+
219+-- | Create a set from the elements of an arbitrary traversable.
220+fromTraversable
221+ :: forall t a. (Traversable t, Hashable a)
222+ => t a -> SomeHashSetWith (FromTraversableProof 'Hashed t a) a
223+fromTraversable xs = reify set \(_ :: Proxy r)
224+ -> SomeHashSetWith @r $ FromTraversableProof
225+ $ unsafeCoerce @(t (Element _ a)) @(t (Element r a)) proof
226+ where
227+ (set, proof) = mapAccumL
228+ (\s x -> let !s' = HashSet.insert x s in (s', unsafeElement x))
229+ HashSet.empty
230+ xs
231+
232+-- | Insert an element in a set.
233+insert
234+ :: forall s a. (Hashable a, KnownHashSet s a)
235+ => a -> SomeHashSetWith (InsertProof 'Hashed a s) a
236+insert x = reify (HashSet.insert x $ reflect $ Proxy @s) \(_ :: Proxy r)
237+ -> SomeHashSetWith @r $ InsertProof (unsafeElement x) unsafeSubset
238+
239+-- | Delete an element from a set.
240+delete
241+ :: forall s a. (Hashable a, KnownHashSet s a)
242+ => a -> SomeHashSetWith (SupersetProof 'Hashed s) a
243+delete x = reify (HashSet.delete x $ reflect $ Proxy @s) \(_ :: Proxy r)
244+ -> SomeHashSetWith @s $ SupersetProof unsafeSubset
245+
246+-- | If an element is in the set, return the proof that it is.
247+member :: forall s a. (Hashable a, KnownHashSet s a) => a -> Maybe (Element s a)
248+member x
249+ | x `HashSet.member` reflect (Proxy @s) = Just $ unsafeElement x
250+ | otherwise = Nothing
251+
252+-- | If the set is empty, return the proof that it is.
253+null :: forall s a. KnownHashSet s a => Maybe (EmptyProof 'Hashed s)
254+null
255+ | HashSet.null $ reflect $ Proxy @s = Just $ EmptyProof unsafeSubset
256+ | otherwise = Nothing
257+
258+-- | If @s@ is a subset of @t@ (or is equal to), return a proof of that.
259+isSubsetOf
260+ :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a)
261+ => Maybe (SubsetProof 'Hashed s t)
262+isSubsetOf
263+#if MIN_VERSION_unordered_containers(0, 2, 12)
264+ | reflect (Proxy @s) `HashSet.isSubsetOf` reflect (Proxy @t)
265+#else
266+ | all (`HashSet.member` reflect (Proxy @t)) (reflect (Proxy @s))
267+#endif
268+ = Just $ SubsetProof unsafeSubset
269+ | otherwise = Nothing
270+
271+-- | If @s@ and @t@ are disjoint (i.e. their intersection is empty), return a
272+-- proof of that.
273+disjoint
274+ :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a)
275+ => Maybe (DisjointProof 'Hashed s t)
276+disjoint
277+ | HashSet.null
278+ $ HashSet.intersection (reflect $ Proxy @s) (reflect $ Proxy @t)
279+ = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g
280+ | otherwise = Nothing
281+
282+-- | The union of two sets.
283+union
284+ :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a)
285+ => SomeHashSetWith (UnionProof 'Hashed s t) a
286+union = reify (reflect (Proxy @s) `HashSet.union` reflect (Proxy @t))
287+ \(_ :: Proxy r) -> SomeHashSetWith @r
288+ $ UnionProof unsafeSubset unsafeSubsetWith2
289+
290+-- unions :: ?
291+
292+-- | HashSet with elements of @s@ that are not in @t@.
293+difference
294+ :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a)
295+ => SomeHashSetWith (DifferenceProof 'Hashed s t) a
296+difference = reify (reflect (Proxy @s) `HashSet.difference` reflect (Proxy @t))
297+ \(_ :: Proxy r) -> SomeHashSetWith @r
298+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
299+
300+-- | Intersection of two sets.
301+intersection
302+ :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a)
303+ => SomeHashSetWith (IntersectionProof 'Hashed s t) a
304+intersection
305+ = reify (reflect (Proxy @s) `HashSet.intersection` reflect (Proxy @t))
306+ \(_ :: Proxy r) -> SomeHashSetWith @r
307+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
308+
309+-- | Return a subset of elements that satisfy the given predicate.
310+filter
311+ :: forall s a. KnownHashSet s a
312+ => (Element s a -> Bool) -> SomeHashSetWith (SupersetProof 'Hashed s) a
313+filter p = reify (HashSet.filter (p . unsafeElement) $ reflect $ Proxy @s)
314+ \(_ :: Proxy r) -> SomeHashSetWith @r $ SupersetProof unsafeSubset
315+
316+-- | Partition a set into two disjoint subsets: those that satisfy the
317+-- predicate, and those that don't.
318+partition
319+ :: forall s a. KnownHashSet s a
320+ => (Element s a -> Bool) -> Some2HashSetWith (PartitionProof 'Hashed s a) a
321+partition p = reify (HashSet.filter (p . unsafeElement) $ reflect $ Proxy @s)
322+ \(_ :: Proxy r)
323+ -> reify (HashSet.filter (not . p . unsafeElement) $ reflect $ Proxy @s)
324+ \(_ :: Proxy q)
325+ -> Some2HashSetWith @s @r $ PartitionProof
326+ do \x -> if p x
327+ then Left $ unsafeElement $ unrefine x
328+ else Right $ unsafeElement $ unrefine x
329+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
330+
331+-- | Apply the given function to each element of the set and collect the
332+-- results. Note that the resulting set can be smaller.
333+map
334+ :: forall s a b. (Hashable b, KnownHashSet s a)
335+ => (Element s a -> b) -> SomeHashSetWith (MapProof 'Hashed s a b) b
336+map f = reify (HashMap.keysSet m)
337+ \(_ :: Proxy r) -> SomeHashSetWith @r
338+ $ MapProof (unsafeElement . f) \y -> case HashMap.lookup (unrefine y) m of
339+ Nothing -> error "map: bug: Data.HashSet.Refined has been subverted"
340+ Just x -> x
341+ where
342+ !m = HashMap.fromList
343+ [ (y, unsafeElement x)
344+ | x <- HashSet.toList $ reflect $ Proxy @s
345+ , let !y = f $ unsafeElement x
346+ ]
347+
348+-- | Map each element of @s@ into a monoid (with proof that it was an element),
349+-- and combine the results using 'Data.Monoid.<>'.
350+foldMap :: forall s a m. (KnownHashSet s a, Monoid m) => (Element s a -> m) -> m
351+foldMap f = Foldable.foldMap (f . unsafeElement) $ reflect $ Proxy @s
352+
353+-- | Right associative fold with a lazy accumulator.
354+foldr :: forall s a b. KnownHashSet s a => (Element s a -> b -> b) -> b -> b
355+foldr f z = HashSet.foldr (f . unsafeElement) z $ reflect $ Proxy @s
356+
357+-- | Left associative fold with a lazy accumulator.
358+foldl :: forall s a b. KnownHashSet s a => (b -> Element s a -> b) -> b -> b
359+foldl f z = Foldable.foldl ((. unsafeElement) . f) z $ reflect $ Proxy @s
360+
361+-- | Right associative fold with a strict accumulator.
362+foldr' :: forall s a b. KnownHashSet s a => (Element s a -> b -> b) -> b -> b
363+foldr' f z = Foldable.foldr' (f . unsafeElement) z $ reflect $ Proxy @s
364+
365+-- | Left associative fold with a strict accumulator.
366+foldl' :: forall s a b. KnownHashSet s a => (b -> Element s a -> b) -> b -> b
367+foldl' f z = HashSet.foldl' ((. unsafeElement) . f) z $ reflect $ Proxy @s
368+
369+-- | List of elements in the set.
370+toList :: forall s a. KnownHashSet s a => [Element s a]
371+toList = gcoerceWith (unsafeCastElement @s @a) $ coerce
372+ $ HashSet.toList $ reflect $ Proxy @s
373+
374+
375+-- | Convert an 'IntSet' into a 'Set', retaining its set of elements, which can
376+-- be converted with 'castFlavor'.
377+asSet :: forall s a. (Ord a, KnownHashSet s a) => Set s a
378+asSet = hashSet2Set
379+
380+-- | Convert an 'IntSet' into a 'HashSet', retaining its set of elements, which
381+-- can be converted with 'castFlavor'.
382+asIntSet :: forall s. KnownHashSet s Int => IntSet s
383+asIntSet = hashSet2IntSet
384+
385+-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then
386+-- @s@ and @t@ actually stand for the same set and @'Element' s@ can be safely
387+-- interconverted with @'Element' t@.
388+--
389+-- The requirement that the weakenings are natural transformations ensures that
390+-- they don't actually alter the elements. To build these you can compose
391+-- ':->''s from proofs returned by functions in this module, or "Refined"
392+-- functions like 'andLeft' or 'leftOr'.
393+castElement
394+ :: forall s t a. (forall x. Element s x -> Element t x)
395+ -> (forall x. Element t x -> Element s x)
396+ -> Coercion (Element s a) (Element t a)
397+castElement = castRefined
398+
399+-- | If elements can be interconverted (e.g. as proved by 'castElement'), then
400+-- the sets can be interconverted too. For example we can establish that the
401+-- intersection of a set with itself is interconvertible with that set:
402+--
403+-- @
404+-- castIntersection
405+-- :: t'IntersectionProof' ''Data.HashSet.Refined.Hashed' s s r
406+-- -> 'Coercion' ('HashSet' r a) ('HashSet' s a)
407+-- castIntersection ( v'IntersectionProof' p1 p2)
408+-- = 'cast' $ 'castElement' ('andLeft' . p1) (p2 'id' 'id')
409+-- @
410+cast
411+ :: forall s t a. (forall x. Coercion (Element s x) (Element t x))
412+ -> Coercion (HashSet s a) (HashSet t a)
413+cast Coercion
414+#if MIN_VERSION_base(4, 15, 0)
415+ = case unsafeEqualityProof @s @t of UnsafeRefl -> Coercion
416+#else
417+ = repr $ unsafeCoerce Refl
418+#endif
src/Data/IntMap/Common/Refined.hsnew file+523
@@ -0,0 +1,523 @@
1+{-# LANGUAGE CPP #-}
2+{-# LANGUAGE UndecidableInstances #-}
3+module Data.IntMap.Common.Refined where
4+
5+import Control.Monad.Reader
6+import Control.DeepSeq
7+import Data.Coerce
8+import Data.Constraint (Dict(..))
9+import Data.Container.Refined.Proofs
10+import Data.Container.Refined.Unsafe
11+import Data.Distributive
12+import Data.Foldable.WithIndex
13+import Data.Functor.Rep
14+import Data.Functor.WithIndex
15+import qualified Data.Hashable as Hashable
16+import qualified Data.IntMap as IntMap
17+import Data.Proxy
18+import Data.Reflection
19+import Data.Traversable.WithIndex
20+import Data.Type.Coercion
21+import Data.Type.Equality ((:~:)(..))
22+import Refined
23+import Refined.Unsafe
24+import Unsafe.Coerce
25+
26+#if MIN_VERSION_containers(0, 6, 7)
27+#elif MIN_VERSION_containers(0, 6, 2)
28+import qualified Data.List as List
29+#elif MIN_VERSION_containers(0, 5, 8)
30+import Data.Functor.Const (Const(..))
31+import qualified Data.List as List
32+import Data.Monoid (Any(..))
33+import qualified Data.IntMap.Merge.Lazy as IntMap
34+#else
35+import qualified Data.IntMap.Strict as IntMapStrict
36+import qualified Data.List as List
37+#endif
38+
39+
40+-- | A wrapper around a regular 'Data.IntMap.IntMap' with a type parameter @s@
41+-- identifying the set of keys present in the map.
42+--
43+-- An t'Int' key may not be present in the map, but a @'Key' s@ is guaranteed to
44+-- be present (if the @s@ parameters match). Thus the map is isomorphic to a
45+-- (total) function @'Key' s -> a@, which motivates many of the instances below.
46+--
47+-- An 'IntMap' always knows its set of keys, so given @'IntMap' s a@ we can
48+-- always derive @'KnownIntSet' s@ by pattern matching on the 'Dict' returned by
49+-- 'keysSet'.
50+newtype IntMap s a = IntMap (IntMap.IntMap a)
51+ deriving newtype (Eq, Ord, Show, Functor, Foldable, NFData)
52+#if MIN_VERSION_hashable(1, 3, 4)
53+ deriving newtype (Hashable.Hashable)
54+#endif
55+ deriving stock (Traversable)
56+type role IntMap nominal representational
57+
58+-- | Convert to a regular 'Data.IntMap.IntMap', forgetting its set of keys.
59+toIntMap :: forall s a. IntMap s a -> IntMap.IntMap a
60+toIntMap (IntMap m) = m
61+
62+-- | @'Key' s@ is a key of type t'Int' that has been verified to be an element
63+-- of the set @s@, and thus verified to be present in all @'IntMap' s k@ maps.
64+--
65+-- Thus, @'Key' s@ is a \"refinement\" type of t'Int', and this library
66+-- integrates with an implementation of refimenement types in "Refined", so
67+-- the machinery from there can be used to manipulate 'Key's (however see
68+-- 'Data.IntSet.Refined.revealPredicate').
69+--
70+-- The underlying t'Int' value can be obtained with 'unrefine'. An t'Int' can be
71+-- validated into an @'Key' s@ with 'member'.
72+type Key s = Refined (InSet 'Int s) Int
73+
74+unsafeCastKey :: forall s. Coercion Int (Key s)
75+unsafeCastKey = reallyUnsafeUnderlyingRefined
76+
77+unsafeKey :: Int -> Key s
78+unsafeKey = coerceWith unsafeCastKey
79+
80+-- | An existential wrapper for an 'IntMap' with an as-yet-unknown set of keys.
81+-- Pattern maching on it gives you a way to refer to the set (the parameter
82+-- @s@), e.g.
83+--
84+-- @
85+-- case 'fromIntMap' ... of
86+-- 'SomeIntMap' \@s m -> doSomethingWith \@s
87+--
88+-- case 'fromIntMap' ... of
89+-- 'SomeIntMap' (m :: 'IntMap' s a) -> doSomethingWith \@s
90+-- @
91+data SomeIntMap a where
92+ SomeIntMap :: forall s a. !(IntMap s a) -> SomeIntMap a
93+
94+-- | Apply a map with an unknown set of keys to a continuation that can accept
95+-- a map with any set of keys. This gives you a way to refer to the set (the
96+-- parameter @s@), e.g.:
97+--
98+-- @
99+-- 'withIntMap' ('fromIntMap' ...) $ \(m :: 'IntMap' s a) -> doSomethingWith \@s
100+-- @
101+withIntMap :: forall a r. SomeIntMap a -> (forall s. IntMap s a -> r) -> r
102+withIntMap (SomeIntMap m) k = k m
103+
104+-- | Construct a map from a regular 'Data.IntMap.IntMap'.
105+fromIntMap :: forall a. IntMap.IntMap a -> SomeIntMap a
106+fromIntMap m = SomeIntMap (IntMap m)
107+
108+-- | An existential wrapper for an 'IntMap' with an as-yet-unknown set of keys,
109+-- together with a proof of some fact @p@ about the set. Pattern matching on it
110+-- gives you a way to refer to the set (the parameter @s@). Functions that
111+-- change the set of keys in a map will return the map in this way, together
112+-- with a proof that somehow relates the keys set to the function's inputs.
113+data SomeIntMapWith p a where
114+ SomeIntMapWith :: forall s a p. !(IntMap s a) -> !(p s) -> SomeIntMapWith p a
115+
116+-- | Apply a map with proof for an unknown set of keys to a continuation that
117+-- can accept a map with any set of keys satisfying the proof. This gives you a
118+-- way to refer to the set (the parameter @s@).
119+withIntMapWith
120+ :: forall a r p. SomeIntMapWith p a -> (forall s. IntMap s a -> p s -> r) -> r
121+withIntMapWith (SomeIntMapWith m p) k = k m p
122+
123+-- | An existential wrapper for a pair of maps with as-yet-unknown sets of keys,
124+-- together with a proof of some fact @p@ relating them.
125+data Some2IntMapWith p a b where
126+ Some2IntMapWith
127+ :: forall s t a b p. !(IntMap s a)
128+ -> !(IntMap t b)
129+ -> !(p s t)
130+ -> Some2IntMapWith p a b
131+
132+-- | Apply a pair of maps with proof for unknown sets of keys to a continuation
133+-- that can accept any pair of maps with any sets of keys satisfying the proof.
134+-- This gives you a way to refer to the sets (the parameters @s@ and @t@).
135+with2IntMapWith
136+ :: forall a b r p. Some2IntMapWith p a b
137+ -> (forall s t. IntMap s a -> IntMap t b -> p s t -> r)
138+ -> r
139+with2IntMapWith (Some2IntMapWith m1 m2 p) k = k m1 m2 p
140+
141+-- | An empty map.
142+empty :: forall a. SomeIntMapWith (EmptyProof 'Int) a
143+empty = SomeIntMapWith (IntMap IntMap.empty) $ EmptyProof unsafeSubset
144+
145+-- | Create a map from a set of keys, and a function that for each key computes
146+-- the corresponding value.
147+fromSet :: forall s a. KnownIntSet s => (Key s -> a) -> IntMap s a
148+fromSet f = IntMap $ IntMap.fromSet (f . unsafeKey) (reflect $ Proxy @s)
149+
150+-- | Delete a key and its value from the map if present, returning a potentially
151+-- smaller map.
152+delete :: forall s a. Int -> IntMap s a -> SomeIntMapWith (SupersetProof 'Int s) a
153+delete k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.delete k m)
154+ $ SupersetProof unsafeSubset
155+
156+-- | If the key is in the map, return the proof of this, and the associated
157+-- value; otherwise return 'Nothing'.
158+lookup :: forall s a. Int -> IntMap s a -> Maybe (Key s, a)
159+lookup k (IntMap m) = (unsafeKey k,) <$> IntMap.lookup k m
160+
161+-- | Given a key that is proven to be in the map, return the associated value.
162+--
163+-- Unlike 'Data.IntMap.!' from "Data.IntMap", this function is total, as it is
164+-- impossible to obtain a @'Key' s@ for a key that is not in the map
165+-- @'IntMap' s a@.
166+(!) :: forall s a. IntMap s a -> Key s -> a
167+(!) (IntMap m) k = case IntMap.lookup (unrefine k) m of
168+ Nothing -> error "(!): bug: Data.IntMap.Refined has been subverted"
169+ Just x -> x
170+
171+-- | If a key is in the map, return the proof that it is.
172+member :: forall s a. Int -> IntMap s a -> Maybe (Key s)
173+member k (IntMap m)
174+ | k `IntMap.member` m = Just (unsafeKey k)
175+ | otherwise = Nothing
176+
177+-- | Find the largest key smaller than the given one, and return the
178+-- associated key-value pair.
179+lookupLT :: forall s a. Int -> IntMap s a -> Maybe (Key s, a)
180+lookupLT = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupLT @a
181+
182+-- | Find the smallest key greater than the given one, and return the
183+-- associated key-value pair.
184+lookupGT :: forall s a. Int -> IntMap s a -> Maybe (Key s, a)
185+lookupGT = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupGT @a
186+
187+-- | Find the largest key smaller or equal to the given one, and return the
188+-- associated key-value pair.
189+lookupLE :: forall s a. Int -> IntMap s a -> Maybe (Key s, a)
190+lookupLE = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupLE @a
191+
192+-- | Find the smallest key greater or equal to the given one, and return the
193+-- associated key-value pair.
194+lookupGE :: forall s a. Int -> IntMap s a -> Maybe (Key s, a)
195+lookupGE = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupGE @a
196+
197+-- | If a map is empty, return a proof that it is.
198+null :: forall s a. IntMap s a -> Maybe (EmptyProof 'Int s)
199+null (IntMap m)
200+ | IntMap.null m = Just $ EmptyProof unsafeSubset
201+ | otherwise = Nothing
202+
203+-- | If all keys of the first map are also present in the second map, and the
204+-- given function returns 'True' for their associated values, return a proof
205+-- that the keys form a subset.
206+isSubmapOfBy
207+ :: forall s t a b. (a -> b -> Bool)
208+ -> IntMap s a
209+ -> IntMap t b
210+ -> Maybe (SubsetProof 'Int s t)
211+isSubmapOfBy f (IntMap m1) (IntMap m2)
212+ | IntMap.isSubmapOfBy f m1 m2 = Just $ SubsetProof unsafeSubset
213+ | otherwise = Nothing
214+
215+-- | If two maps are disjoint (i.e. their intersection is empty), return a proof
216+-- of that.
217+disjoint
218+ :: forall s t a b. IntMap s a -> IntMap t b -> Maybe (DisjointProof 'Int s t)
219+disjoint (IntMap m1) (IntMap m2)
220+#if MIN_VERSION_containers(0, 6, 2)
221+ | IntMap.disjoint m1 m2
222+#elif MIN_VERSION_containers(0, 5, 8)
223+ | Const (Any False) <- IntMap.mergeA
224+ (IntMap.traverseMissing \_ _ -> Const mempty)
225+ (IntMap.traverseMissing \_ _ -> Const mempty)
226+ (IntMap.zipWithAMatched \_ _ _ -> Const $ Any True)
227+ m1
228+ m2
229+#else
230+ | IntMap.null $ IntMapStrict.intersectionWith (\_ _ -> ()) m1 m2
231+#endif
232+ = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g
233+ | otherwise = Nothing
234+
235+-- | Given two maps proven to have the same keys, for each key apply the
236+-- 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)
239+ -> IntMap s a
240+ -> IntMap s b
241+ -> IntMap s c
242+zipWithKey f (IntMap m1) (IntMap m2) = IntMap
243+ $ IntMap.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
244+ (error "zipWithKey: bug: Data.IntMap.Refined has been subverted")
245+ (error "zipWithKey: bug: Data.IntMap.Refined has been subverted")
246+ m1
247+ m2
248+
249+-- | Remove the keys that appear in the second map from the first map.
250+difference
251+ :: forall s t a b. IntMap s a
252+ -> IntMap t b
253+ -> SomeIntMapWith (DifferenceProof 'Int s t) a
254+difference (IntMap m1) (IntMap m2) = SomeIntMapWith
255+ (IntMap $ IntMap.difference m1 m2)
256+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
257+
258+-- | Apply a function to all values in a map, together with their corresponding
259+-- keys, that are proven to be in the map. The set of keys remains the same.
260+mapWithKey :: forall s a b. (Key s -> a -> b) -> IntMap s a -> IntMap s b
261+mapWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
262+ $ IntMap.mapWithKey @a @b
263+
264+-- | Map an 'Applicative' transformation in ascending order of keys, with access
265+-- to each value's corresponding key and a proof that it is in the map. The set
266+-- of keys remains unchanged.
267+traverseWithKey
268+ :: forall s f a b. Applicative f
269+ => (Key s -> a -> f b) -> IntMap s a -> f (IntMap s b)
270+traverseWithKey f (IntMap m) = IntMap <$> IntMap.traverseWithKey (f . unsafeKey) m
271+
272+-- | Map each key-value pair of a map into a monoid (with proof that the key was
273+-- in the map), and combine the results using '<>'.
274+foldMapWithKey
275+ :: forall s a m. Monoid m => (Key s -> a -> m) -> IntMap s a -> m
276+foldMapWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
277+ $ IntMap.foldMapWithKey @m @a
278+
279+-- | Right associative fold with a lazy accumulator.
280+foldrWithKey :: forall s a b. (Key s -> a -> b -> b) -> b -> IntMap s a -> b
281+foldrWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
282+ $ IntMap.foldrWithKey @a @b
283+
284+-- | Left associative fold with a lazy accumulator.
285+foldlWithKey :: forall s a b. (b -> Key s -> a -> b) -> b -> IntMap s a -> b
286+foldlWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
287+ $ IntMap.foldlWithKey @b @a
288+
289+-- | Right associative fold with a strict accumulator.
290+foldrWithKey' :: forall s a b. (Key s -> a -> b -> b) -> b -> IntMap s a -> b
291+foldrWithKey' = gcoerceWith (unsafeCastKey @s) $ coerce
292+ $ IntMap.foldrWithKey' @a @b
293+
294+-- | Left associative fold with a strict accumulator.
295+foldlWithKey' :: forall s a b. (b -> Key s -> a -> b) -> b -> IntMap s a -> b
296+foldlWithKey' = gcoerceWith (unsafeCastKey @s) $ coerce
297+ $ IntMap.foldlWithKey' @b @a
298+
299+-- | Return the set of keys in the map, with the contents of the set still
300+-- tracked by the @s@ parameter. See "Data.IntSet.Refined".
301+keysSet :: forall s a. IntMap s a -> IntSet s
302+keysSet (IntMap m) = reify (IntMap.keysSet m)
303+ \(_ :: Proxy s') -> case unsafeCoerce Refl :: s :~: s' of
304+ Refl -> Dict
305+
306+-- | Convert to a list of key-value pairs in ascending order of keys.
307+toList :: forall s a. IntMap s a -> [(Key s, a)]
308+toList = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.toAscList @a
309+
310+-- | Convert to a list of key-value pairs in descending order of keys.
311+toDescList :: forall s a. IntMap s a -> [(Key s, a)]
312+toDescList = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.toDescList @a
313+
314+-- | Retain only the key-value pairs that satisfy the predicate, returning a
315+-- potentially smaller map.
316+filterWithKey
317+ :: forall s a. (Key s -> a -> Bool)
318+ -> IntMap s a
319+ -> SomeIntMapWith (SupersetProof 'Int s) a
320+filterWithKey p (IntMap m)
321+ = SomeIntMapWith (IntMap $ IntMap.filterWithKey (p . unsafeKey) m)
322+ $ SupersetProof unsafeSubset
323+
324+-- | Restrict a map to only those keys that are elements of @t@.
325+restrictKeys
326+ :: forall s t a. KnownIntSet t
327+ => IntMap s a -> SomeIntMapWith (IntersectionProof 'Int s t) a
328+restrictKeys (IntMap m) = SomeIntMapWith
329+#if MIN_VERSION_containers(0, 5, 8)
330+ (IntMap $ IntMap.restrictKeys m $ reflect $ Proxy @t)
331+#else
332+ (IntMap $ IntMap.intersectionWith const m
333+ $ IntMap.fromSet id $ reflect $ Proxy @t)
334+#endif
335+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
336+
337+-- | Remove all keys that are elements of @t@ from the map.
338+withoutKeys
339+ :: forall s t a. KnownIntSet t
340+ => IntMap s a -> SomeIntMapWith (DifferenceProof 'Int s t) a
341+withoutKeys (IntMap m) = SomeIntMapWith
342+#if MIN_VERSION_containers(0, 5, 8)
343+ (IntMap $ IntMap.withoutKeys m $ reflect $ Proxy @t)
344+#else
345+ (IntMap $ IntMap.difference m $ IntMap.fromSet id $ reflect $ Proxy @t)
346+#endif
347+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
348+
349+-- | Partition a map into two disjoint submaps: those whose key-value pairs
350+-- satisfy the predicate, and those whose don't.
351+partitionWithKey
352+ :: forall s a. (Key s -> a -> Bool)
353+ -> IntMap s a
354+ -> Some2IntMapWith (PartitionProof 'Int s Int) a a
355+partitionWithKey p (IntMap m)
356+ = case IntMap.partitionWithKey (p . unsafeKey) m of
357+ (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof
358+ do \k -> case IntMap.lookup (unrefine k) m of
359+ Nothing -> error
360+ "partitionWithKey: bug: Data.IntMap.Refined has been subverted"
361+ Just x -> if p k x
362+ then Left $ unsafeKey $ unrefine k
363+ else Right $ unsafeKey $ unrefine k
364+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
365+
366+-- | Divide a map into two disjoint submaps at a point where the predicate on
367+-- the keys stops holding.
368+--
369+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
370+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
371+-- chosen in an unspecified way.
372+spanAntitone
373+ :: forall s a. (Key s -> Bool)
374+ -> IntMap s a
375+ -> Some2IntMapWith (PartialPartitionProof 'Int s) a a
376+spanAntitone p (IntMap m) =
377+#if MIN_VERSION_containers(0, 6, 7)
378+ case IntMap.spanAntitone (p . unsafeKey) m of
379+ (m1, m2)
380+#else
381+ case List.span (p . unsafeKey . fst) $ IntMap.toAscList m of
382+ (xs1, xs2)
383+ | let m1 = IntMap.fromDistinctAscList xs1
384+ , let m2 = IntMap.fromDistinctAscList xs2
385+#endif
386+ -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartialPartitionProof
387+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
388+
389+-- | Return two disjoint submaps: those whose keys are less than the given key,
390+-- and those whose keys are greater than the given key. If the key was in the
391+-- map, also return the associated value and the proof that it was in the map.
392+splitLookup
393+ :: forall s a. Int
394+ -> IntMap s a
395+ -> Some2IntMapWith (SplitProof 'Int s (Key s, a)) a a
396+splitLookup k (IntMap m) = case IntMap.splitLookup k m of
397+ (m1, v, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ SplitProof
398+ ((unsafeKey k,) <$> v) unsafeSubset \f g -> unsafeSubsetWith2 f g
399+
400+-- | Retrieves the key-value pair corresponding to the smallest key of the map,
401+-- and the map with that pair removed; or a proof that the map was empty.
402+minViewWithKey
403+ :: forall s a. IntMap s a
404+ -> Either
405+ (EmptyProof 'Int s)
406+ ((Key s, a), SomeIntMapWith (SupersetProof 'Int s) a)
407+minViewWithKey (IntMap m) = case IntMap.minViewWithKey m of
408+ Nothing -> Left $ EmptyProof unsafeSubset
409+ Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s) $ coerce kv,)
410+ $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset
411+
412+-- | Retrieves the key-value pair corresponding to the greatest key of the map,
413+-- and the map with that pair removed; or a proof that the map was empty.
414+maxViewWithKey
415+ :: forall s a. IntMap s a
416+ -> Either
417+ (EmptyProof 'Int s)
418+ ((Key s, a), SomeIntMapWith (SupersetProof 'Int s) a)
419+maxViewWithKey (IntMap m) = case IntMap.maxViewWithKey m of
420+ Nothing -> Left $ EmptyProof unsafeSubset
421+ Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s) $ coerce kv,)
422+ $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset
423+
424+-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then
425+-- @s@ and @t@ actually stand for the same set and @'Key' s@ can be safely
426+-- interconverted with @'Key' t@.
427+--
428+-- The requirement that the weakenings are natural transformations ensures that
429+-- they don't actually alter the keys. To build these you can compose ':->''s
430+-- from proofs returned by functions in this module, or "Refined" functions like
431+-- 'andLeft' or 'leftOr'.
432+castKey
433+ :: forall s t k.
434+ (forall x. Refined (InSet 'Int s) x -> Refined (InSet 'Int t) x)
435+ -> (forall x. Refined (InSet 'Int t) x -> Refined (InSet 'Int s) x)
436+ -> Coercion (Refined (InSet 'Int s) k) (Refined (InSet 'Int t) k)
437+castKey = castRefined
438+
439+-- | 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.IntMap.Refined.intersectionWithKey' by proving that the set of keys
442+-- remains unchanged:
443+--
444+-- @
445+-- 'zipWithKey'
446+-- :: forall s a b c. ('Key' s -> a -> b -> c)
447+-- -> 'IntMap' s a
448+-- -> 'IntMap' s b
449+-- -> 'IntMap' s c
450+-- 'zipWithKey' f m1 m2
451+-- | v'SomeIntMapWith' @r m proof <- 'Data.IntMap.Refined.intersectionWithKey' (f . 'andLeft') m1 m2
452+-- , v'IntersectionProof' p1 p2 <- proof
453+-- , ( v'Coercion' :: t'Coercion' ('IntMap' r c) ('IntMap' s c))
454+-- <- app $ 'cast' $ 'castKey' ('andLeft' . p1) (p2 'id' 'id')
455+-- = 'coerce' m
456+-- where
457+-- app :: t'Coercion' f g -> t'Coercion' (f x) (g x)
458+-- app v'Coercion' = v'Coercion'
459+-- @
460+cast
461+ :: forall s t k. (forall x. Coercion
462+ (Refined (InSet 'Int s) x)
463+ (Refined (InSet 'Int t) x))
464+ -> Coercion (IntMap s k) (IntMap t k)
465+cast Coercion = Coercion
466+
467+instance FunctorWithIndex (Key s) (IntMap s) where
468+ imap = mapWithKey
469+
470+instance FoldableWithIndex (Key s) (IntMap s) where
471+ ifoldMap = foldMapWithKey
472+
473+instance TraversableWithIndex (Key s) (IntMap s) where
474+ itraverse = traverseWithKey
475+
476+-- | Similar to the instance for functions -- zip corresponding keys. To use
477+-- '<*>'/'liftA2' without 'KnownIntSet' see 'zipWithKey'.
478+instance KnownIntSet s => Applicative (IntMap s) where
479+ pure x = fromSet \_ -> x
480+ (<*>) = zipWithKey (const id)
481+
482+-- | @'bind' m f@ is a map that for each key @k :: 'Key' s@, contains the
483+-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
484+bind :: forall s a b. IntMap s a -> (a -> IntMap s b) -> IntMap s b
485+bind m f = mapWithKey (\k x -> f x ! k) m
486+
487+-- | Similar to the instance for functions. To use '>>=' without 'KnownIntSet'
488+-- see 'bind'.
489+instance KnownIntSet s => Monad (IntMap s) where
490+ (>>=) = bind
491+
492+-- | Similar to the instance for functions. See also
493+-- 'Data.IntMap.Refined.backpermuteKeys'.
494+instance KnownIntSet s => MonadReader (Key s) (IntMap s) where
495+ ask = fromSet id
496+ local f m = mapWithKey (\k _ -> m ! f k) m
497+
498+-- | Append the values at the corresponding keys
499+instance Semigroup a => Semigroup (IntMap s a) where
500+ (<>) = zipWithKey (const (<>))
501+
502+instance (KnownIntSet s, Monoid a) => Monoid (IntMap s a) where
503+ mempty = fromSet \_ -> mempty
504+
505+-- | Similar to the instance for functions
506+instance KnownIntSet s => Distributive (IntMap s) where
507+ collect = collectRep
508+ distribute = distributeRep
509+
510+-- | Witness isomorphism with functions from @'Key' s@
511+instance KnownIntSet s => Representable (IntMap s) where
512+ type Rep (IntMap s) = Key s
513+ index = (!)
514+ tabulate = fromSet
515+
516+#if MIN_VERSION_hashable(1, 3, 4)
517+#else
518+instance Hashable.Hashable a => Hashable.Hashable (IntMap s a) where
519+ hashWithSalt s (IntMap m) = IntMap.foldlWithKey'
520+ (\s' k v -> Hashable.hashWithSalt (Hashable.hashWithSalt s' k) v)
521+ (Hashable.hashWithSalt s (IntMap.size m))
522+ m
523+#endif
src/Data/IntMap/Refined.hsnew file+379
@@ -0,0 +1,379 @@
1+-- | This module defines a way to prove that a key exists in a map, so that the
2+-- key can be used to index into the map without using a 'Maybe', or manually
3+-- handling the \"impossible\" case with 'error' or other partial functions.
4+--
5+-- To do this, @'IntMap' s v@ has a type parameter @s@ that identifies its set
6+-- ofvkeys, so that if another map has the same type parameter, you know that
7+-- map has the same set of keys. There is @'Key' s@, a type of keys that have
8+-- been validated to belong to the set identified by @s@, and for which the
9+-- operation of indexing into a @'IntMap' s v@ (only for the same @s@) can
10+-- proceed without failure (see '!'). The type @s@ itself has no internal
11+-- structure, rather it is merely a skolem type variable (rank-2 polymorphism
12+-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
13+--
14+-- Like "Data.IntMap", functions in this module are strict in the keys but lazy
15+-- in the values. The "Data.IntMap.Strict.Refined" module reuses the same
16+-- 'IntMap' type but provides functions that operate strictly on the values.
17+module Data.IntMap.Refined
18+ (
19+ -- * Map type
20+ Common.IntMap
21+ , Common.Key
22+ -- * Existentials and common proofs
23+ , Common.SomeIntMap(..)
24+ , Common.withIntMap
25+ , Common.SomeIntMapWith(..)
26+ , Common.withIntMapWith
27+ , Common.Some2IntMapWith(..)
28+ , Common.with2IntMapWith
29+ , SupersetProof(..)
30+ , EmptyProof(..)
31+ -- * Construction
32+ , Common.empty
33+ , singleton
34+ , SingletonProof(..)
35+ , fromSet
36+ , Common.fromIntMap
37+ , fromTraversableWithKey
38+ , FromTraversableProof(..)
39+ -- * Insertion
40+ , insert
41+ , InsertProof(..)
42+ , reinsert
43+ , insertLookupWithKey
44+ -- * Deletion/Update
45+ , Common.delete
46+ , adjust
47+ , adjustWithKey
48+ , update
49+ , updateLookupWithKey
50+ -- * Query
51+ , Common.lookup
52+ , (Common.!)
53+ , Common.member
54+ , Common.lookupLT
55+ , Common.lookupGT
56+ , Common.lookupLE
57+ , Common.lookupGE
58+ , Common.null
59+ , Common.isSubmapOfBy
60+ , SubsetProof(..)
61+ , Common.disjoint
62+ , DisjointProof(..)
63+ -- * Combine
64+ , zipWithKey
65+ , bind
66+ , unionWithKey
67+ , UnionProof(..)
68+ , Common.difference
69+ , DifferenceProof(..)
70+ , differenceWithKey
71+ , PartialDifferenceProof(..)
72+ , intersectionWithKey
73+ , IntersectionProof(..)
74+ -- * Traversal
75+ , mapWithKey
76+ , traverseWithKey
77+ , mapAccumLWithKey
78+ , mapAccumRWithKey
79+ , mapKeysWith
80+ , MapProof(..)
81+ , backpermuteKeys
82+ -- * Folds
83+ , Common.foldMapWithKey
84+ , Common.foldrWithKey
85+ , Common.foldlWithKey
86+ , Common.foldrWithKey'
87+ , Common.foldlWithKey'
88+ -- * Conversion
89+ , Common.toIntMap
90+ , Common.keysSet
91+ , Common.toList
92+ , Common.toDescList
93+ -- * Filter
94+ , Common.restrictKeys
95+ , Common.withoutKeys
96+ , Common.filterWithKey
97+ , Common.partitionWithKey
98+ , PartitionProof(..)
99+ , Common.spanAntitone
100+ , PartialPartitionProof(..)
101+ , mapMaybeWithKey
102+ , mapEitherWithKey
103+ , Common.splitLookup
104+ , SplitProof(..)
105+ -- * Min/Max
106+ , updateMinWithKey
107+ , updateMaxWithKey
108+ , adjustMinWithKey
109+ , adjustMaxWithKey
110+ , Common.minViewWithKey
111+ , Common.maxViewWithKey
112+ -- * Casts
113+ , Common.castKey
114+ , Common.cast
115+ , castFlavor
116+ ) where
117+
118+import Data.Coerce
119+import Data.Container.Refined.Proofs
120+import Data.Container.Refined.Unsafe
121+import Data.Functor
122+import qualified Data.IntMap as IntMap
123+import Data.IntMap.Common.Refined
124+ ( IntMap(..), Key, unsafeCastKey, unsafeKey, SomeIntMapWith(..)
125+ , Some2IntMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey
126+ , bind
127+ )
128+import qualified Data.IntMap.Common.Refined as Common
129+import Data.Traversable
130+import Data.Type.Coercion
131+import Prelude hiding (lookup, null)
132+import Refined
133+import Refined.Unsafe
134+
135+
136+-- | Create a map with a single key-value pair, and return a proof that the
137+-- key is in the resulting map.
138+singleton :: forall a. Int -> a -> SomeIntMapWith (SingletonProof 'Int Int) a
139+singleton k v = SomeIntMapWith (IntMap $ IntMap.singleton k v)
140+ $ SingletonProof (unsafeKey k)
141+
142+-- | Create a map from an arbitrary traversable of key-value pairs.
143+fromTraversableWithKey
144+ :: forall t a. Traversable t
145+ => (Int -> a -> a -> a)
146+ -> t (Int, a)
147+ -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
148+fromTraversableWithKey f xs
149+ = SomeIntMapWith (IntMap m) $ FromTraversableProof proof
150+ where
151+ (m, proof) = mapAccumL
152+ (\s (k, v) -> let !s' = IntMap.insertWithKey f k v s in (s', unsafeKey k))
153+ IntMap.empty
154+ xs
155+
156+-- | Insert a key-value pair into the map to obtain a potentially larger map,
157+-- guaranteed to contain the given key. If the key was already present, the
158+-- associated value is replaced with the supplied value.
159+insert
160+ :: forall s a. Int
161+ -> a
162+ -> IntMap s a
163+ -> SomeIntMapWith (InsertProof 'Int Int s) a
164+insert k v (IntMap m) = SomeIntMapWith (IntMap $ IntMap.insert k v m)
165+ $ InsertProof (unsafeKey k) unsafeSubset
166+
167+-- | Overwrite a key-value pair that is known to already be in the map. The set
168+-- of keys remains the same.
169+reinsert
170+ :: forall s a. Key s -> a -> IntMap s a -> IntMap s a
171+reinsert = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.insert @a
172+
173+-- | Insert a key-value pair into the map using a combining function, and if
174+-- the key was already present, the old value is returned along with the proof
175+-- that the key was present.
176+insertLookupWithKey
177+ :: forall s a. (Key s -> a -> a -> a)
178+ -> Int
179+ -> a
180+ -> IntMap s a
181+ -> (Maybe (Key s, a), SomeIntMapWith (InsertProof 'Int Int s) a)
182+insertLookupWithKey f k v (IntMap m)
183+ = case IntMap.insertLookupWithKey (f . unsafeKey) k v m of
184+ (v', !m') -> ((unsafeKey k,) <$> v',)
185+ $ SomeIntMapWith (IntMap m') $ InsertProof (unsafeKey k) unsafeSubset
186+
187+-- | Update the value at a specific key known the be in the map using the given
188+-- function. The set of keys remains the same.
189+adjust :: forall s a. (a -> a) -> Key s -> IntMap s a -> IntMap s a
190+adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.adjust @a
191+
192+-- | If the given key is in the map, update the associated value using the given
193+-- function with a proof that the key was in the map; otherwise return the map
194+-- unchanged. In any case the set of keys remains the same.
195+adjustWithKey
196+ :: forall s a. (Key s -> a -> a) -> Int -> IntMap s a -> IntMap s a
197+adjustWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
198+ $ IntMap.adjustWithKey @a
199+
200+-- | Update or delete a key known to be in the map using the given function,
201+-- returning a potentially smaller map.
202+update
203+ :: forall s a. (a -> Maybe a)
204+ -> Key s
205+ -> IntMap s a
206+ -> SomeIntMapWith (SupersetProof 'Int s) a
207+update f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f (unrefine k) m)
208+ $ SupersetProof unsafeSubset
209+
210+-- | If the given key is in the map, update or delete it using the given
211+-- function with a proof that the key was in the map; otherwise the map is
212+-- unchanged. Alongside return the new value if it was updated, or the old value
213+-- if it was deleted, and a proof that the key was in the map.
214+updateLookupWithKey
215+ :: forall s a. (Key s -> a -> Maybe a)
216+ -> Int
217+ -> IntMap s a
218+ -> (Maybe (Key s, a), SomeIntMapWith (SupersetProof 'Int s) a)
219+updateLookupWithKey f k (IntMap m)
220+ = case IntMap.updateLookupWithKey (f . unsafeKey) k m of
221+ (v', !m') -> ((unsafeKey k,) <$> v',)
222+ $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset
223+
224+-- | Return the union of two maps, with a given combining function for keys that
225+-- exist in both maps simultaneously.
226+--
227+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
228+-- respectively.
229+unionWithKey
230+ :: forall s t a. (Refined (InSet 'Int s && InSet 'Int t) Int -> a -> a -> a)
231+ -> IntMap s a
232+ -> IntMap t a
233+ -> SomeIntMapWith (UnionProof 'Int s t) a
234+unionWithKey f (IntMap m1) (IntMap m2)
235+ = SomeIntMapWith (IntMap $ IntMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
236+ $ UnionProof unsafeSubset unsafeSubsetWith2
237+
238+-- | For keys that appear in both maps, the given function decides whether the
239+-- key is removed from the first map.
240+--
241+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
242+-- respectively.
243+differenceWithKey
244+ :: forall s t a b. (Refined (InSet 'Int s && InSet 'Int t) Int
245+ -> a
246+ -> b
247+ -> Maybe a)
248+ -> IntMap s a
249+ -> IntMap t b
250+ -> SomeIntMapWith (PartialDifferenceProof 'Int s t) a
251+differenceWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
252+ (IntMap $ IntMap.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
253+ $ PartialDifferenceProof unsafeSubset unsafeSubset
254+
255+-- | Return the intersection of two maps with the given combining function.
256+--
257+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
258+-- respectively.
259+intersectionWithKey
260+ :: forall s t a b c. (Refined (InSet 'Int s && InSet 'Int t) Int
261+ -> a
262+ -> b
263+ -> c)
264+ -> IntMap s a
265+ -> IntMap t b
266+ -> SomeIntMapWith (IntersectionProof 'Int s t) c
267+intersectionWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
268+ (IntMap $ IntMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
269+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
270+
271+-- | Thread an accumularing argument through the map in ascending order of keys.
272+mapAccumLWithKey
273+ :: forall s a b c. (a -> Key s -> b -> (a, c))
274+ -> a
275+ -> IntMap s b
276+ -> (a, IntMap s c)
277+mapAccumLWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
278+ $ IntMap.mapAccumWithKey @a @b @c
279+
280+-- | Thread an accumularing argument through the map in descending order of
281+-- keys.
282+mapAccumRWithKey
283+ :: forall s a b c. (a -> Key s -> b -> (a, c))
284+ -> a
285+ -> IntMap s b
286+ -> (a, IntMap s c)
287+mapAccumRWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
288+ $ IntMap.mapAccumRWithKey @a @b @c
289+
290+-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
291+-- results into a new map. For keys that were mapped to the same new key, @c@
292+-- acts as the combining function for corresponding values.
293+mapKeysWith
294+ :: forall s a. (a -> a -> a)
295+ -> (Key s -> Int)
296+ -> IntMap s a
297+ -> SomeIntMapWith (MapProof 'Int s Int Int) a
298+mapKeysWith f g (IntMap m)
299+ = SomeIntMapWith (IntMap $ IntMap.mapKeysWith f (g . unsafeKey) m)
300+ $ MapProof (unsafeKey . g) \k2 ->
301+ case IntMap.lookup (unrefine k2) backMap of
302+ Nothing -> error
303+ "mapKeysWith: bug: Data.IntMap.Refined has been subverted"
304+ Just k1 -> k1
305+ where
306+ ~backMap = IntMap.fromList
307+ [ (k2, unsafeKey k1)
308+ | k1 <- IntMap.keys m
309+ , let !k2 = g $ unsafeKey k1
310+ ]
311+
312+-- | Apply a function to all values in a map, together with their corresponding
313+-- keys, and collect only the 'Just' results, returning a potentially smaller
314+-- map.
315+mapMaybeWithKey
316+ :: forall s a b. (Key s -> a -> Maybe b)
317+ -> IntMap s a
318+ -> SomeIntMapWith (SupersetProof 'Int s) b
319+mapMaybeWithKey f (IntMap m)
320+ = SomeIntMapWith (IntMap $ IntMap.mapMaybeWithKey (f . unsafeKey) m)
321+ $ SupersetProof unsafeSubset
322+
323+-- | Apply a function to all values in a map, together with their corresponding
324+-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
325+-- maps.
326+mapEitherWithKey
327+ :: forall s a b c. (Key s -> a -> Either b c)
328+ -> IntMap s a
329+ -> Some2IntMapWith (PartitionProof 'Int s Int) b c
330+mapEitherWithKey p (IntMap m)
331+ = case IntMap.mapEitherWithKey (p . unsafeKey) m of
332+ (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof
333+ do \k -> case IntMap.lookup (unrefine k) m of
334+ Nothing -> error
335+ "mapEitherWithKey: bug: Data.IntMap.Refined has been subverted"
336+ Just x -> case p k x of
337+ Left _ -> Left $ unsafeKey $ unrefine k
338+ Right _ -> Right $ unsafeKey $ unrefine k
339+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
340+
341+-- | Update or delete the value at the smallest key, returning a potentially
342+-- smaller map.
343+updateMinWithKey
344+ :: forall s a. (Key s -> a -> Maybe a)
345+ -> IntMap s a
346+ -> SomeIntMapWith (SupersetProof 'Int s) a
347+updateMinWithKey f (IntMap m)
348+ = SomeIntMapWith (IntMap $ IntMap.updateMinWithKey (f . unsafeKey) m)
349+ $ SupersetProof unsafeSubset
350+
351+-- | Update or delete the value at the largest key, returning a potentially
352+-- smaller map.
353+updateMaxWithKey
354+ :: forall s a. (Key s -> a -> Maybe a)
355+ -> IntMap s a
356+ -> SomeIntMapWith (SupersetProof 'Int s) a
357+updateMaxWithKey f (IntMap m)
358+ = SomeIntMapWith (IntMap $ IntMap.updateMaxWithKey (f . unsafeKey) m)
359+ $ SupersetProof unsafeSubset
360+
361+-- | Adjust the value at the smallest key. The set of keys remains unchanged.
362+adjustMinWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a
363+adjustMinWithKey f (IntMap m)
364+ = IntMap $ IntMap.updateMinWithKey ((Just .) . f . unsafeKey) m
365+
366+-- | Adjust the value at the greatest key. The set of keys remains unchanged.
367+adjustMaxWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a
368+adjustMaxWithKey f (IntMap m)
369+ = IntMap $ IntMap.updateMaxWithKey ((Just .) . f . unsafeKey) m
370+
371+-- | Apply the inverse image of the given function to the keys of the given map,
372+-- so that for all @k :: 'Key' s2@,
373+-- @'backpermuteKeys' f m '!' k = m '!' f k@.
374+--
375+-- If maps are identified with functions, this computes the composition.
376+backpermuteKeys
377+ :: forall s1 s2 a. KnownIntSet s2
378+ => (Key s2 -> Key s1) -> IntMap s1 a -> IntMap s2 a
379+backpermuteKeys f m = fromSet \k -> m ! f k
src/Data/IntMap/Strict/Refined.hsnew file+418
@@ -0,0 +1,418 @@
1+-- | This module defines a way to prove that a key exists in a map, so that the
2+-- key can be used to index into the map without using a 'Maybe', or manually
3+-- handling the \"impossible\" case with 'error' or other partial functions.
4+--
5+-- To do this, @'IntMap' s v@ has a type parameter @s@ that identifies its set
6+-- of keys, so that if another map has the same type parameter, you know that
7+-- map has the same set of keys. There is @'Key' s@, a type of keys that have
8+-- been validated to belong to the set identified by @s@, and for which the
9+-- operation of indexing into a @'IntMap' s v@ (only for the same @s@) can
10+-- proceed without failure (see '!'). The type @s@ itself has no internal
11+-- structure, rather it is merely a skolem type variable (rank-2 polymorphism
12+-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
13+--
14+-- Like "Data.IntMap.Strict", functions in this module are strict in the keys
15+-- and values. The "Data.IntMap.Refined" module reuses the same 'IntMap' type
16+-- but provides functions that operate lazily on the values.
17+module Data.IntMap.Strict.Refined
18+ (
19+ -- * Map type
20+ Common.IntMap
21+ , Common.Key
22+ -- * Existentials and common proofs
23+ , Common.SomeIntMap(..)
24+ , Common.withIntMap
25+ , Common.SomeIntMapWith(..)
26+ , Common.withIntMapWith
27+ , Common.Some2IntMapWith(..)
28+ , Common.with2IntMapWith
29+ , SupersetProof(..)
30+ , EmptyProof(..)
31+ -- * Construction
32+ , Common.empty
33+ , singleton
34+ , SingletonProof(..)
35+ , fromSet
36+ , Common.fromIntMap
37+ , fromTraversableWithKey
38+ , FromTraversableProof(..)
39+ -- * Insertion
40+ , insert
41+ , InsertProof(..)
42+ , reinsert
43+ , insertLookupWithKey
44+ -- * Deletion/Update
45+ , Common.delete
46+ , adjust
47+ , adjustWithKey
48+ , update
49+ , updateLookupWithKey
50+ -- * Query
51+ , Common.lookup
52+ , (Common.!)
53+ , Common.member
54+ , Common.lookupLT
55+ , Common.lookupGT
56+ , Common.lookupLE
57+ , Common.lookupGE
58+ , Common.null
59+ , Common.isSubmapOfBy
60+ , SubsetProof(..)
61+ , Common.disjoint
62+ , DisjointProof(..)
63+ -- * Combine
64+ , zipWithKey
65+ , bind
66+ , unionWithKey
67+ , UnionProof(..)
68+ , Common.difference
69+ , DifferenceProof(..)
70+ , differenceWithKey
71+ , PartialDifferenceProof(..)
72+ , intersectionWithKey
73+ , IntersectionProof(..)
74+ -- * Traversal
75+ , mapWithKey
76+ , traverseWithKey
77+ , mapAccumLWithKey
78+ , mapAccumRWithKey
79+ , mapKeysWith
80+ , MapProof(..)
81+ , backpermuteKeys
82+ -- * Folds
83+ , Common.foldMapWithKey
84+ , Common.foldrWithKey
85+ , Common.foldlWithKey
86+ , Common.foldrWithKey'
87+ , Common.foldlWithKey'
88+ -- * Conversion
89+ , Common.toIntMap
90+ , Common.keysSet
91+ , Common.toList
92+ , Common.toDescList
93+ -- * Filter
94+ , Common.restrictKeys
95+ , Common.withoutKeys
96+ , Common.filterWithKey
97+ , Common.partitionWithKey
98+ , PartitionProof(..)
99+ , Common.spanAntitone
100+ , PartialPartitionProof(..)
101+ , mapMaybeWithKey
102+ , mapEitherWithKey
103+ , Common.splitLookup
104+ , SplitProof(..)
105+ -- * Min/Max
106+ , updateMinWithKey
107+ , updateMaxWithKey
108+ , adjustMinWithKey
109+ , adjustMaxWithKey
110+ , Common.minViewWithKey
111+ , Common.maxViewWithKey
112+ -- * Casts
113+ , Common.castKey
114+ , Common.cast
115+ , castFlavor
116+ ) where
117+
118+import Data.Coerce
119+import Data.Container.Refined.Proofs
120+import Data.Container.Refined.Unsafe
121+import Data.Functor
122+import qualified Data.IntMap.Strict as IntMap
123+import Data.IntMap.Common.Refined
124+ ( IntMap(..), Key, unsafeCastKey, unsafeKey, SomeIntMapWith(..)
125+ , Some2IntMapWith(..), (!)
126+ )
127+import qualified Data.IntMap.Common.Refined as Common
128+import Data.Proxy
129+import Data.Reflection
130+import Data.Traversable
131+import Data.Type.Coercion
132+import Prelude hiding (lookup, null)
133+import Refined
134+import Refined.Unsafe
135+
136+
137+-- | Create a map with a single key-value pair, and return a proof that the
138+-- key is in the resulting map.
139+singleton :: forall a. Int -> a -> SomeIntMapWith (SingletonProof 'Int Int) a
140+singleton k v = SomeIntMapWith (IntMap $ IntMap.singleton k v)
141+ $ SingletonProof (unsafeKey k)
142+
143+-- | Create a map from a set of keys, and a function that for each key computes
144+-- the corresponding value.
145+fromSet :: forall s a. KnownIntSet s => (Key s -> a) -> IntMap s a
146+fromSet f = IntMap $ IntMap.fromSet (f . unsafeKey) (reflect $ Proxy @s)
147+
148+-- | Create a map from an arbitrary traversable of key-value pairs.
149+fromTraversableWithKey
150+ :: forall t a. Traversable t
151+ => (Int -> a -> a -> a)
152+ -> t (Int, a)
153+ -> SomeIntMapWith (FromTraversableProof 'Int t Int) a
154+fromTraversableWithKey f xs = SomeIntMapWith (IntMap m)
155+ $ FromTraversableProof proof
156+ where
157+ (m, proof) = mapAccumL
158+ (\s (k, v) -> (IntMap.insertWithKey f k v s, unsafeKey k))
159+ IntMap.empty
160+ xs
161+
162+-- | Insert a key-value pair into the map to obtain a potentially larger map,
163+-- guaranteed to contain the given key. If the key was already present, the
164+-- associated value is replaced with the supplied value.
165+insert
166+ :: forall s a. Int
167+ -> a
168+ -> IntMap s a
169+ -> SomeIntMapWith (InsertProof 'Int Int s) a
170+insert k v (IntMap m) = SomeIntMapWith (IntMap $ IntMap.insert k v m)
171+ $ InsertProof (unsafeKey k) unsafeSubset
172+
173+-- | Overwrite a key-value pair that is known to already be in the map. The set
174+-- of keys remains the same.
175+reinsert
176+ :: forall s a. Key s -> a -> IntMap s a -> IntMap s a
177+reinsert = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.insert @a
178+
179+-- | Insert a key-value pair into the map using a combining function, and if
180+-- the key was already present, the old value is returned along with the proof
181+-- that the key was present.
182+insertLookupWithKey
183+ :: forall s a. (Key s -> a -> a -> a)
184+ -> Int
185+ -> a
186+ -> IntMap s a
187+ -> (Maybe (Key s, a), SomeIntMapWith (InsertProof 'Int Int s) a)
188+insertLookupWithKey f k v (IntMap m)
189+ = case IntMap.insertLookupWithKey (f . unsafeKey) k v m of
190+ (v', !m') -> ((unsafeKey k,) <$> v',)
191+ $ SomeIntMapWith (IntMap m') $ InsertProof (unsafeKey k) unsafeSubset
192+
193+-- | Update the value at a specific key known the be in the map using the given
194+-- function. The set of keys remains the same.
195+adjust :: forall s a. (a -> a) -> Key s -> IntMap s a -> IntMap s a
196+adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.adjust @a
197+
198+-- | If the given key is in the map, update the associated value using the given
199+-- function with a proof that the key was in the map; otherwise return the map
200+-- unchanged. In any case the set of keys remains the same.
201+adjustWithKey
202+ :: forall s a. (Key s -> a -> a) -> Int -> IntMap s a -> IntMap s a
203+adjustWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
204+ $ IntMap.adjustWithKey @a
205+
206+-- | Update or delete a key known to be in the map using the given function,
207+-- returning a potentially smaller map.
208+update
209+ :: forall s a. (a -> Maybe a)
210+ -> Key s
211+ -> IntMap s a
212+ -> SomeIntMapWith (SupersetProof 'Int s) a
213+update f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f (unrefine k) m)
214+ $ SupersetProof unsafeSubset
215+
216+-- | If the given key is in the map, update or delete it using the given
217+-- function with a proof that the key was in the map; otherwise the map is
218+-- unchanged. Alongside return the new value if it was updated, or the old value
219+-- if it was deleted, and a proof that the key was in the map.
220+updateLookupWithKey
221+ :: forall s a. (Key s -> a -> Maybe a)
222+ -> Int
223+ -> IntMap s a
224+ -> (Maybe (Key s, a), SomeIntMapWith (SupersetProof 'Int s) a)
225+updateLookupWithKey f k (IntMap m)
226+ = case IntMap.updateLookupWithKey (f . unsafeKey) k m of
227+ (v', !m') -> ((unsafeKey k,) <$> v',)
228+ $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset
229+
230+-- | Given two maps proven to have the same keys, for each key apply the
231+-- function to the associated values, to obtain a new map with the same keys.
232+zipWithKey
233+ :: forall s a b c. (Key s -> a -> b -> c)
234+ -> IntMap s a
235+ -> IntMap s b
236+ -> IntMap s c
237+zipWithKey f (IntMap m1) (IntMap m2) = IntMap
238+ $ IntMap.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
239+ (error "zipWithKey: bug: Data.IntMap.Strict.Refined has been subverted")
240+ (error "zipWithKey: bug: Data.IntMap.Strict.Refined has been subverted")
241+ m1
242+ m2
243+
244+-- | Return the union of two maps, with a given combining function for keys that
245+-- exist in both maps simultaneously.
246+--
247+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
248+-- respectively.
249+unionWithKey
250+ :: forall s t a. (Refined (InSet 'Int s && InSet 'Int t) Int -> a -> a -> a)
251+ -> IntMap s a
252+ -> IntMap t a
253+ -> SomeIntMapWith (UnionProof 'Int s t) a
254+unionWithKey f (IntMap m1) (IntMap m2)
255+ = SomeIntMapWith (IntMap $ IntMap.unionWithKey (f . reallyUnsafeRefine) m1 m2)
256+ $ UnionProof unsafeSubset unsafeSubsetWith2
257+
258+-- | For keys that appear in both maps, the given function decides whether the
259+-- key is removed from the first map.
260+--
261+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
262+-- respectively.
263+differenceWithKey
264+ :: forall s t a b. (Refined (InSet 'Int s && InSet 'Int t) Int
265+ -> a
266+ -> b
267+ -> Maybe a)
268+ -> IntMap s a
269+ -> IntMap t b
270+ -> SomeIntMapWith (PartialDifferenceProof 'Int s t) a
271+differenceWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
272+ (IntMap $ IntMap.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
273+ $ PartialDifferenceProof unsafeSubset unsafeSubset
274+
275+-- | Return the intersection of two maps with the given combining function.
276+--
277+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@
278+-- respectively.
279+intersectionWithKey
280+ :: forall s t a b c. (Refined (InSet 'Int s && InSet 'Int t) Int
281+ -> a
282+ -> b
283+ -> c)
284+ -> IntMap s a
285+ -> IntMap t b
286+ -> SomeIntMapWith (IntersectionProof 'Int s t) c
287+intersectionWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith
288+ (IntMap $ IntMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
289+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
290+
291+-- | Apply a function to all values in a map, together with their corresponding
292+-- keys, that are proven to be in the map. The set of keys remains the same.
293+mapWithKey :: forall s a b. (Key s -> a -> b) -> IntMap s a -> IntMap s b
294+mapWithKey = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.mapWithKey @a @b
295+
296+-- | Map an 'Applicative' transformation in ascending order of keys, with access
297+-- to each value's corresponding key and a proof that it is in the map. The set
298+-- of keys remains unchanged.
299+traverseWithKey
300+ :: forall s f a b. Applicative f
301+ => (Key s -> a -> f b) -> IntMap s a -> f (IntMap s b)
302+traverseWithKey f (IntMap m)
303+ = IntMap <$> IntMap.traverseWithKey (f . unsafeKey) m
304+
305+-- | Thread an accumularing argument through the map in ascending order of keys.
306+mapAccumLWithKey
307+ :: forall s a b c. (a -> Key s -> b -> (a, c))
308+ -> a
309+ -> IntMap s b
310+ -> (a, IntMap s c)
311+mapAccumLWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
312+ $ IntMap.mapAccumWithKey @a @b @c
313+
314+-- | Thread an accumularing argument through the map in descending order of
315+-- keys.
316+mapAccumRWithKey
317+ :: forall s a b c. (a -> Key s -> b -> (a, c))
318+ -> a
319+ -> IntMap s b
320+ -> (a, IntMap s c)
321+mapAccumRWithKey = gcoerceWith (unsafeCastKey @s) $ coerce
322+ $ IntMap.mapAccumRWithKey @a @b @c
323+
324+-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
325+-- results into a new map. For keys that were mapped to the same new key, @c@
326+-- acts as the combining function for corresponding values.
327+mapKeysWith
328+ :: forall s a. (a -> a -> a)
329+ -> (Key s -> Int)
330+ -> IntMap s a
331+ -> SomeIntMapWith (MapProof 'Int s Int Int) a
332+mapKeysWith f g (IntMap m)
333+ = SomeIntMapWith (IntMap $ IntMap.mapKeysWith f (g . unsafeKey) m)
334+ $ MapProof (unsafeKey . g) \k2 ->
335+ case IntMap.lookup (unrefine k2) backMap of
336+ Nothing -> error
337+ "mapKeysWith: bug: Data.IntMap.Strict.Refined has been subverted"
338+ Just k1 -> k1
339+ where
340+ ~backMap = IntMap.fromList
341+ [ (k2, unsafeKey k1)
342+ | k1 <- IntMap.keys m
343+ , let !k2 = g $ unsafeKey k1
344+ ]
345+
346+-- | Apply a function to all values in a map, together with their corresponding
347+-- keys, and collect only the 'Just' results, returning a potentially smaller
348+-- map.
349+mapMaybeWithKey
350+ :: forall s a b. (Key s -> a -> Maybe b)
351+ -> IntMap s a
352+ -> SomeIntMapWith (SupersetProof 'Int s) b
353+mapMaybeWithKey f (IntMap m)
354+ = SomeIntMapWith (IntMap $ IntMap.mapMaybeWithKey (f . unsafeKey) m)
355+ $ SupersetProof unsafeSubset
356+
357+-- | Apply a function to all values in a map, together with their corresponding
358+-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
359+-- maps.
360+mapEitherWithKey
361+ :: forall s a b c. (Key s -> a -> Either b c)
362+ -> IntMap s a
363+ -> Some2IntMapWith (PartitionProof 'Int s Int) b c
364+mapEitherWithKey p (IntMap m)
365+ = case IntMap.mapEitherWithKey (p . unsafeKey) m of
366+ (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof
367+ do \k -> case IntMap.lookup (unrefine k) m of
368+ Nothing -> error
369+ "mapEitherWithKey: bug: Data.IntMap.Strict.Refined has been subverted"
370+ Just x -> case p k x of
371+ Left _ -> Left $ unsafeKey $ unrefine k
372+ Right _ -> Right $ unsafeKey $ unrefine k
373+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
374+
375+-- | Update or delete the value at the smallest key, returning a potentially
376+-- smaller map.
377+updateMinWithKey
378+ :: forall s a. (Key s -> a -> Maybe a)
379+ -> IntMap s a
380+ -> SomeIntMapWith (SupersetProof 'Int s) a
381+updateMinWithKey f (IntMap m)
382+ = SomeIntMapWith (IntMap $ IntMap.updateMinWithKey (f . unsafeKey) m)
383+ $ SupersetProof unsafeSubset
384+
385+-- | Update or delete the value at the largest key, returning a potentially
386+-- smaller map.
387+updateMaxWithKey
388+ :: forall s a. (Key s -> a -> Maybe a)
389+ -> IntMap s a
390+ -> SomeIntMapWith (SupersetProof 'Int s) a
391+updateMaxWithKey f (IntMap m)
392+ = SomeIntMapWith (IntMap $ IntMap.updateMaxWithKey (f . unsafeKey) m)
393+ $ SupersetProof unsafeSubset
394+
395+-- | Adjust the value at the smallest key. The set of keys remains unchanged.
396+adjustMinWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a
397+adjustMinWithKey f (IntMap m)
398+ = IntMap $ IntMap.updateMinWithKey ((Just .) . f . unsafeKey) m
399+
400+-- | Adjust the value at the greatest key. The set of keys remains unchanged.
401+adjustMaxWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a
402+adjustMaxWithKey f (IntMap m)
403+ = IntMap $ IntMap.updateMaxWithKey ((Just .) . f . unsafeKey) m
404+
405+-- | @'bind' m f@ is a map that for each key @k :: 'Key' s@, contains the
406+-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
407+bind :: forall s a b. IntMap s a -> (a -> IntMap s b) -> IntMap s b
408+bind m f = mapWithKey (\k x -> f x ! k) m
409+
410+-- | Apply the inverse image of the given function to the keys of the given map,
411+-- so that for all @k :: 'Key' s2@,
412+-- @'backpermuteKeys' f m '!' k = m '!' f k@.
413+--
414+-- If maps are identified with functions, this computes the composition.
415+backpermuteKeys
416+ :: forall s1 s2 a. KnownIntSet s2
417+ => (Key s2 -> Key s1) -> IntMap s1 a -> IntMap s2 a
418+backpermuteKeys f m = fromSet \k -> m ! f k
src/Data/IntSet/Refined.hsnew file+502
@@ -0,0 +1,502 @@
1+{-# LANGUAGE CPP #-}
2+-- | This module implements a way of tracking the contents of an
3+-- 'Data.IntSet.IntSet' at the type level, and provides utilities for
4+-- manipulating such sets.
5+--
6+-- The contents of a set are associated with a type parameter, e.g. @s@, so that
7+-- whenever you see the same type parameter, you know you are working with the
8+-- same set. The type @s@ itself has no internal structure, rather it is merely
9+-- a skolem type variable (rank-2 polymorphism 'Control.Monad.ST.runST' trick)
10+-- introduced by "Data.Reflection".
11+module Data.IntSet.Refined
12+ (
13+ -- * Set type
14+ KnownIntSet
15+ , IntSet
16+ -- * Refinement type
17+ , InSet(..)
18+ , Flavor(Int)
19+ , Element
20+ , revealPredicate
21+ -- * Existentials and common proofs
22+ , SomeIntSet(..)
23+ , withIntSet
24+ , SomeIntSetWith(..)
25+ , withIntSetWith
26+ , Some2IntSetWith(..)
27+ , with2IntSetWith
28+ , (:->)
29+ , SupersetProof(..)
30+ , EmptyProof(..)
31+ -- * Construction
32+ , empty
33+ , singleton
34+ , SingletonProof(..)
35+ , fromIntSet
36+ , fromTraversable
37+ , FromTraversableProof(..)
38+ -- * Insertion
39+ , insert
40+ , InsertProof(..)
41+ -- * Deletion
42+ , delete
43+ -- * Query
44+ , member
45+ , lookupLT
46+ , lookupGT
47+ , lookupLE
48+ , lookupGE
49+ , null
50+ , isSubsetOf
51+ , SubsetProof(..)
52+ , disjoint
53+ , DisjointProof(..)
54+ -- * Combine
55+ , union
56+ , UnionProof(..)
57+ , difference
58+ , DifferenceProof(..)
59+ , intersection
60+ , IntersectionProof(..)
61+ -- * Filter
62+ , filter
63+ , partition
64+ , PartitionProof(..)
65+ , spanAntitone
66+ , PartialPartitionProof(..)
67+ , splitMember
68+ , SplitProof(..)
69+ -- * Map
70+ , map
71+ , MapProof(..)
72+ -- * Folds
73+ , foldMap
74+ , foldr
75+ , foldl
76+ , foldr'
77+ , foldl'
78+ -- * Min/Max
79+ , minView
80+ , maxView
81+ -- * Conversion
82+ , toList
83+ , toDescList
84+ , asSet
85+ , asHashSet
86+ -- * Casts
87+ , castElement
88+ , cast
89+ , castFlavor
90+ ) where
91+
92+import Data.Coerce
93+import Data.Constraint (Dict(..))
94+import Data.Container.Refined.Conversion
95+import Data.Container.Refined.Proofs
96+import Data.Container.Refined.Unsafe
97+import qualified Data.Foldable as Foldable
98+import qualified Data.IntMap as IntMap
99+import qualified Data.IntSet as IntSet
100+import Data.Proxy
101+import Data.Reflection
102+import Data.Traversable
103+import Data.Type.Coercion
104+import Data.Type.Equality ((:~:)(..))
105+import GHC.Exts (Proxy#, proxy#)
106+import Prelude hiding (filter, foldl, foldMap, foldr, map, null)
107+import Refined
108+import Refined.Unsafe
109+import Unsafe.Coerce
110+
111+#if MIN_VERSION_containers(0, 6, 7)
112+#else
113+import qualified Data.List as List
114+#endif
115+
116+
117+-- | To use "Refined" machinery that uses the 'Predicate' typeclass you will
118+-- need to pattern match on this 'Dict'.
119+--
120+-- The reason is that in the default /fast/ implementation of reflection, we
121+-- don't have @'Data.Typeable.Typeable' s@, which "Refined" wants for
122+-- pretty-printing exceptions. We /can/ provide @'Data.TypeableTypeable' s@, but
123+-- at the cost of using the /slow/ implementation of reflection.
124+revealPredicate
125+ :: forall s. KnownIntSet s => Dict (Predicate (InSet 'Int s) Int)
126+revealPredicate = reifyTypeable (reflect (Proxy @s))
127+ \(_ :: Proxy s') ->
128+ reflect (Proxy @s') `seq`
129+ -- ^ Work around https://github.com/ekmett/reflection/issues/54
130+ case unsafeCoerce Refl :: s :~: s' of
131+ Refl -> Dict
132+
133+-- | @'Element' s@ is an t'Int' that has been verified to be an element of @s@.
134+--
135+-- Thus, @'Element' s@ is a \"refinement\" type of t'Int', and this library
136+-- integrates with an implementation of refimenement types in "Refined", so
137+-- the machinery from there can be used to manipulate 'Element's (however see
138+-- 'revealPredicate').
139+--
140+-- The underlying t'Int' can be obtained with 'unrefine'. An t'Int' can be
141+-- validated into an @'Element' s@ with 'member'.
142+type Element s = Refined (InSet 'Int s) Int
143+
144+unsafeCastElement :: forall s. Coercion Int (Element s)
145+unsafeCastElement = reallyUnsafeUnderlyingRefined
146+
147+unsafeElement :: Int -> Element s
148+unsafeElement = coerceWith unsafeCastElement
149+
150+-- | An existential wrapper for an as-yet-unknown set. Pattern maching on it
151+-- gives you a way to refer to the set, e.g.
152+--
153+-- @
154+-- case 'fromIntSet' ... of
155+-- 'SomeIntSet' \@s _ -> doSomethingWith \@s
156+--
157+-- case 'fromIntSet' ... of
158+-- 'SomeIntSet' (_ :: 'Proxy#' s) -> doSomethingWith \@s
159+-- @
160+data SomeIntSet where
161+ SomeIntSet :: forall s. KnownIntSet s => Proxy# s -> SomeIntSet
162+
163+-- | Apply an unknown set to a continuation that can accept any set. This gives
164+-- you a way to refer to the set (the parameter @s@), e.g.:
165+--
166+-- @
167+-- 'withIntSet' ('fromIntSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s
168+-- @
169+withIntSet
170+ :: forall r. SomeIntSet -> (forall s. KnownIntSet s => Proxy s -> r) -> r
171+withIntSet (SomeIntSet (_ :: Proxy# s)) k = k $ Proxy @s
172+
173+-- | Construct a set from a regular 'Data.IntSet.IntSet'.
174+fromIntSet :: IntSet.IntSet -> SomeIntSet
175+fromIntSet s = reify s \(_ :: Proxy s) -> SomeIntSet @s proxy#
176+
177+-- | An existential wrapper for an as-yet-unknown set, together with a proof of
178+-- some fact @p@ about the set. Pattern matching on it gives you a way to refer
179+-- to the set (the parameter @s@). Most functions will return a set in this way,
180+-- together with a proof that somehow relates the set to the function's inputs.
181+data SomeIntSetWith p where
182+ SomeIntSetWith :: forall s p. KnownIntSet s => !(p s) -> SomeIntSetWith p
183+
184+-- | Apply an unknown set with proof to a continuation that can accept any set
185+-- satisfying the proof. This gives you a way to refer to the set (the parameter
186+-- @s@).
187+withIntSetWith
188+ :: forall r p. SomeIntSetWith p -> (forall s. KnownIntSet s => p s -> r) -> r
189+withIntSetWith (SomeIntSetWith p) k = k p
190+
191+-- | An existential wrapper for an as-yet-unknown pair of sets, together with
192+-- a proof of some fact @p@ relating them.
193+data Some2IntSetWith p where
194+ Some2IntSetWith
195+ :: forall s t p. (KnownIntSet s, KnownIntSet t)
196+ => !(p s t) -> Some2IntSetWith p
197+
198+-- | Apply a pair of unknown sets with proof to a continuation that can accept
199+-- any pair of sets satisfying the proof. This gives you a way to refer to the
200+-- sets (the parameters @s@ and @t@).
201+with2IntSetWith
202+ :: forall r p. Some2IntSetWith p
203+ -> (forall s t. (KnownIntSet s, KnownIntSet t) => p s t -> r)
204+ -> r
205+with2IntSetWith (Some2IntSetWith p) k = k p
206+
207+-- | An empty set.
208+empty :: SomeIntSetWith (EmptyProof 'Int)
209+empty = reify IntSet.empty \(_ :: Proxy r)
210+ -> SomeIntSetWith @r $ EmptyProof unsafeSubset
211+
212+
213+-- | Create a set with a single element.
214+singleton :: Int -> SomeIntSetWith (SingletonProof 'Int Int)
215+singleton x = reify (IntSet.singleton x) \(_ :: Proxy r)
216+ -> SomeIntSetWith @r $ SingletonProof $ unsafeElement x
217+
218+-- | Create a set from the elements of an arbitrary traversable.
219+fromTraversable
220+ :: forall t. Traversable t
221+ => t Int -> SomeIntSetWith (FromTraversableProof 'Int t Int)
222+fromTraversable xs = reify set \(_ :: Proxy r)
223+ -> SomeIntSetWith @r $ FromTraversableProof
224+ $ unsafeCoerce @(t (Element _)) @(t (Element r)) proof
225+ where
226+ (set, proof) = mapAccumL
227+ (\s x -> let !s' = IntSet.insert x s in (s', unsafeElement x))
228+ IntSet.empty
229+ xs
230+
231+-- | Insert an element in a set.
232+insert :: forall s. KnownIntSet s
233+ => Int -> SomeIntSetWith (InsertProof 'Int Int s)
234+insert x = reify (IntSet.insert x $ reflect $ Proxy @s) \(_ :: Proxy r)
235+ -> SomeIntSetWith @r $ InsertProof (unsafeElement x) unsafeSubset
236+
237+-- | Delete an element from a set.
238+delete :: forall s. KnownIntSet s
239+ => Int -> SomeIntSetWith (SupersetProof 'Int s)
240+delete x = reify (IntSet.delete x $ reflect $ Proxy @s) \(_ :: Proxy r)
241+ -> SomeIntSetWith @s $ SupersetProof unsafeSubset
242+
243+-- | If an element is in the set, return the proof that it is.
244+member :: forall s. KnownIntSet s => Int -> Maybe (Element s)
245+member x
246+ | x `IntSet.member` reflect (Proxy @s) = Just $ unsafeElement x
247+ | otherwise = Nothing
248+
249+-- | Find the largest element smaller than the given one.
250+lookupLT :: forall s. KnownIntSet s => Int -> Maybe (Element s)
251+lookupLT x = gcoerceWith (unsafeCastElement @s) $ coerce
252+ $ IntSet.lookupLT x (reflect $ Proxy @s)
253+
254+-- | Find the smallest element greater than the given one.
255+lookupGT :: forall s. KnownIntSet s => Int -> Maybe (Element s)
256+lookupGT x = gcoerceWith (unsafeCastElement @s) $ coerce
257+ $ IntSet.lookupGT x (reflect $ Proxy @s)
258+
259+-- | Find the largest element smaller or equal to the given one.
260+lookupLE :: forall s. KnownIntSet s => Int -> Maybe (Element s)
261+lookupLE x = gcoerceWith (unsafeCastElement @s) $ coerce
262+ $ IntSet.lookupLE x (reflect $ Proxy @s)
263+
264+-- | Find the smallest element greater or equal to the given one.
265+lookupGE :: forall s. KnownIntSet s => Int -> Maybe (Element s)
266+lookupGE x = gcoerceWith (unsafeCastElement @s) $ coerce
267+ $ IntSet.lookupGE x (reflect $ Proxy @s)
268+
269+-- | If the set is empty, return the proof that it is.
270+null :: forall s. KnownIntSet s => Maybe (EmptyProof 'Int s)
271+null
272+ | IntSet.null $ reflect $ Proxy @s = Just $ EmptyProof unsafeSubset
273+ | otherwise = Nothing
274+
275+-- | If @s@ is a subset of @t@ (or is equal to), return a proof of that.
276+isSubsetOf
277+ :: forall s t. (KnownIntSet s, KnownIntSet t) => Maybe (SubsetProof 'Int s t)
278+isSubsetOf
279+ | reflect (Proxy @s) `IntSet.isSubsetOf` reflect (Proxy @t)
280+ = Just $ SubsetProof unsafeSubset
281+ | otherwise = Nothing
282+
283+-- | If @s@ and @t@ are disjoint (i.e. their intersection is empty), return a
284+-- proof of that.
285+disjoint
286+ :: forall s t. (KnownIntSet s, KnownIntSet t)
287+ => Maybe (DisjointProof 'Int s t)
288+disjoint
289+#if MIN_VERSION_containers(0, 5, 11)
290+ | IntSet.disjoint (reflect $ Proxy @s) (reflect $ Proxy @t)
291+#else
292+ | IntSet.null $ IntSet.intersection (reflect $ Proxy @s) (reflect $ Proxy @t)
293+#endif
294+ = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g
295+ | otherwise = Nothing
296+
297+-- | The union of two sets.
298+union
299+ :: forall s t. (KnownIntSet s, KnownIntSet t)
300+ => SomeIntSetWith (UnionProof 'Int s t)
301+union = reify (reflect (Proxy @s) `IntSet.union` reflect (Proxy @t))
302+ \(_ :: Proxy r) -> SomeIntSetWith @r
303+ $ UnionProof unsafeSubset unsafeSubsetWith2
304+
305+-- unions :: ?
306+
307+-- | Set with elements of @s@ that are not in @t@.
308+difference
309+ :: forall s t. (KnownIntSet s, KnownIntSet t)
310+ => SomeIntSetWith (DifferenceProof 'Int s t)
311+difference = reify (reflect (Proxy @s) `IntSet.difference` reflect (Proxy @t))
312+ \(_ :: Proxy r) -> SomeIntSetWith @r
313+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
314+
315+-- | Intersection of two sets.
316+intersection
317+ :: forall s t. (KnownIntSet s, KnownIntSet t)
318+ => SomeIntSetWith (IntersectionProof 'Int s t)
319+intersection
320+ = reify (reflect (Proxy @s) `IntSet.intersection` reflect (Proxy @t))
321+ \(_ :: Proxy r) -> SomeIntSetWith @r
322+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
323+
324+-- | Return a subset of elements that satisfy the given predicate.
325+filter
326+ :: forall s. KnownIntSet s
327+ => (Element s -> Bool) -> SomeIntSetWith (SupersetProof 'Int s)
328+filter p = reify (IntSet.filter (p . unsafeElement) $ reflect $ Proxy @s)
329+ \(_ :: Proxy r) -> SomeIntSetWith @r $ SupersetProof unsafeSubset
330+
331+-- | Partition a set into two disjoint subsets: those that satisfy the
332+-- predicate, and those that don't.
333+partition
334+ :: forall s. KnownIntSet s
335+ => (Element s -> Bool) -> Some2IntSetWith (PartitionProof 'Int s Int)
336+partition p = case IntSet.partition (p . unsafeElement) $ reflect $ Proxy @s of
337+ (r, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q)
338+ -> Some2IntSetWith @s @r $ PartitionProof
339+ do \x -> if p x
340+ then Left $ unsafeElement $ unrefine x
341+ else Right $ unsafeElement $ unrefine x
342+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
343+
344+-- | Divide a set into two disjoint subsets at a point where the predicate stops
345+-- holding.
346+--
347+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
348+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
349+-- chosen in an unspecified way.
350+spanAntitone
351+ :: forall s. KnownIntSet s
352+ => (Element s -> Bool) -> Some2IntSetWith (PartialPartitionProof 'Int s)
353+spanAntitone p =
354+#if MIN_VERSION_containers(0, 6, 7)
355+ case IntSet.spanAntitone (p . unsafeElement) $ reflect $ Proxy @s of
356+ (r, q)
357+#else
358+ case List.span (p . unsafeElement)
359+ $ IntSet.toAscList $ reflect $ Proxy @s of
360+ (rs, qs)
361+ | let r = IntSet.fromDistinctAscList rs
362+ , let q = IntSet.fromDistinctAscList qs
363+#endif
364+ -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q)
365+ -> Some2IntSetWith @r @q $ PartialPartitionProof
366+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
367+
368+-- | Return two disjoint subsets: those less than the given element, and those
369+-- greater than the given element; along with the proof that the given element
370+-- was in the set, if it was.
371+splitMember
372+ :: forall s. KnownIntSet s
373+ => Int -> Some2IntSetWith (SplitProof 'Int s (Element s))
374+splitMember x = case IntSet.splitMember x $ reflect $ Proxy @s of
375+ (r, m, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q)
376+ -> Some2IntSetWith @r @q $ SplitProof
377+ (if m then Just (unsafeElement x) else Nothing)
378+ unsafeSubset \f g -> unsafeSubsetWith2 f g
379+
380+-- | Apply the given function to each element of the set and collect the
381+-- results. Note that the resulting set can be smaller.
382+map
383+ :: forall s. KnownIntSet s
384+ => (Element s -> Int) -> SomeIntSetWith (MapProof 'Int s Int Int)
385+map f = reify (IntMap.keysSet m) \(_ :: Proxy r) -> SomeIntSetWith @r
386+ $ MapProof (unsafeElement . f) \y -> case IntMap.lookup (unrefine y) m of
387+ Nothing -> error "map: bug: Data.IntSet.Refined has been subverted"
388+ Just x -> x
389+ where
390+ !m = IntMap.fromList
391+ [ (y, unsafeElement x)
392+ | x <- IntSet.toList $ reflect $ Proxy @s
393+ , let !y = f $ unsafeElement x
394+ ]
395+
396+
397+-- | Map each element of @s@ into a monoid (with proof that it was an element),
398+-- and combine the results using 'Data.Monoid.<>'.
399+foldMap :: forall s m. (KnownIntSet s, Monoid m) => (Element s -> m) -> m
400+foldMap f = go $ reflect $ Proxy @s
401+ where
402+ go s = case IntSet.splitRoot s of
403+ [s'] -> Foldable.foldMap (f . unsafeElement) $ IntSet.toAscList s'
404+ xs -> Foldable.foldMap go xs
405+
406+-- | Right associative fold with a lazy accumulator.
407+foldr :: forall s a. KnownIntSet s => (Element s -> a -> a) -> a -> a
408+foldr f z = IntSet.foldr (f . unsafeElement) z $ reflect $ Proxy @s
409+
410+-- | Left associative fold with a lazy accumulator.
411+foldl :: forall s a. KnownIntSet s => (a -> Element s -> a) -> a -> a
412+foldl f z = IntSet.foldl ((. unsafeElement) . f) z $ reflect $ Proxy @s
413+
414+-- | Right associative fold with a strict accumulator.
415+foldr' :: forall s a. KnownIntSet s => (Element s -> a -> a) -> a -> a
416+foldr' f z = IntSet.foldr' (f . unsafeElement) z $ reflect $ Proxy @s
417+
418+-- | Left associative fold with a strict accumulator.
419+foldl' :: forall s a. KnownIntSet s => (a -> Element s -> a) -> a -> a
420+foldl' f z = IntSet.foldl' ((. unsafeElement) . f) z $ reflect $ Proxy @s
421+
422+-- | Retrieves the smallest element of the set, and the set with that element
423+-- removed; or a proof that the set was empty.
424+minView
425+ :: forall s. KnownIntSet s
426+ => Either
427+ (EmptyProof 'Int s)
428+ (Element s, SomeIntSetWith (SupersetProof 'Int s))
429+minView = case IntSet.minView $ reflect $ Proxy @s of
430+ Nothing -> Left $ EmptyProof unsafeSubset
431+ Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r)
432+ -> SomeIntSetWith @r $ SupersetProof unsafeSubset
433+
434+-- | Retrieves the greatest element of the set, and the set with that element
435+-- removed; or a proof that the set was empty.
436+maxView
437+ :: forall s. KnownIntSet s
438+ => Either
439+ (EmptyProof 'Int s)
440+ (Element s, SomeIntSetWith (SupersetProof 'Int s))
441+maxView = case IntSet.maxView $ reflect $ Proxy @s of
442+ Nothing -> Left $ EmptyProof unsafeSubset
443+ Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r)
444+ -> SomeIntSetWith @r $ SupersetProof unsafeSubset
445+
446+-- | List of elements in the set in ascending order.
447+toList :: forall s. KnownIntSet s => [Element s]
448+toList = gcoerceWith (unsafeCastElement @s) $ coerce
449+ $ IntSet.toAscList $ reflect $ Proxy @s
450+
451+-- | List of elements in the set in descending order.
452+toDescList :: forall s. KnownIntSet s => [Element s]
453+toDescList = gcoerceWith (unsafeCastElement @s) $ coerce
454+ $ IntSet.toDescList $ reflect $ Proxy @s
455+
456+-- | Convert an 'IntSet' into a 'Set', retaining its set of elements, which can
457+-- be converted with 'castFlavor'.
458+asSet :: forall s. KnownIntSet s => Set s Int
459+asSet = intSet2Set
460+
461+-- | Convert an 'IntSet' into a 'HashSet', retaining its set of elements, which
462+-- can be converted with 'castFlavor'.
463+asHashSet :: forall s. KnownIntSet s => HashSet s Int
464+asHashSet = intSet2HashSet
465+
466+-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then
467+-- @s@ and @t@ actually stand for the same set and @'Element' s@ can be safely
468+-- interconverted with @'Element' t@.
469+--
470+-- The requirement that the weakenings are natural transformations ensures that
471+-- they don't actually alter the elements. To build these you can compose
472+-- ':->''s from proofs returned by functions in this module, or "Refined"
473+-- functions like 'andLeft' or 'leftOr'.
474+castElement
475+ :: forall s t a.
476+ (forall x. Refined (InSet 'Int s) x -> Refined (InSet 'Int t) x)
477+ -> (forall x. Refined (InSet 'Int t) x -> Refined (InSet 'Int s) x)
478+ -> Coercion (Refined (InSet 'Int s) a) (Refined (InSet 'Int t) a)
479+castElement = castRefined
480+
481+-- | If elements can be interconverted (e.g. as proved by 'castElement'), then
482+-- the sets can be interconverted too. For example we can establish that the
483+-- intersection of a set with itself is interconvertible with that set:
484+--
485+-- @
486+-- castIntersection
487+-- :: t'IntersectionProof' ''Data.IntSet.Refined.Int' s s r
488+-- -> 'Coercion' ('IntSet' r) ('IntSet' s)
489+-- castIntersection ( v'IntersectionProof' p1 p2)
490+-- = 'cast' $ 'castElement' ('andLeft' . p1) (p2 'id' 'id')
491+-- @
492+cast
493+ :: forall s t. (forall x. Coercion
494+ (Refined (InSet 'Int s) x)
495+ (Refined (InSet 'Int t) x))
496+ -> Coercion (IntSet s) (IntSet t)
497+cast Coercion
498+#if MIN_VERSION_base(4, 15, 0)
499+ = case unsafeEqualityProof @s @t of UnsafeRefl -> Coercion
500+#else
501+ = repr $ unsafeCoerce Refl
502+#endif
src/Data/Map/Common/Refined.hsnew file+515
@@ -0,0 +1,515 @@
1+{-# LANGUAGE CPP #-}
2+{-# LANGUAGE UndecidableInstances #-}
3+module Data.Map.Common.Refined where
4+
5+import Control.Monad.Reader
6+import Control.DeepSeq
7+import Data.Coerce
8+import Data.Constraint (Dict(..))
9+import Data.Container.Refined.Proofs
10+import Data.Container.Refined.Unsafe
11+import Data.Distributive
12+import Data.Foldable.WithIndex
13+import Data.Functor.Rep
14+import Data.Functor.WithIndex
15+import qualified Data.Hashable as Hashable
16+import qualified Data.Map as Map
17+import Data.Proxy
18+import Data.Reflection
19+import Data.Traversable.WithIndex
20+import Data.Type.Coercion
21+import Data.Type.Equality ((:~:)(..))
22+import Refined
23+import Refined.Unsafe
24+import Unsafe.Coerce
25+
26+#if MIN_VERSION_containers(0, 6, 2)
27+#elif MIN_VERSION_containers(0, 5, 8)
28+import Data.Functor.Const (Const(..))
29+import Data.Monoid (Any(..))
30+import qualified Data.Map.Merge.Lazy as Map
31+#else
32+import qualified Data.List as List
33+import qualified Data.Map.Strict as MapStrict
34+#endif
35+
36+
37+-- | A wrapper around a regular 'Data.Map.Map' with a type parameter @s@
38+-- identifying the set of keys present in the map.
39+--
40+-- A key of type @k@ may not be present in the map, but a @'Key' s k@ is
41+-- guaranteed to be present (if the @s@ parameters match). Thus the map is
42+-- isomorphic to a (total) function @'Key' s k -> a@, which motivates many of
43+-- the instances below.
44+--
45+-- A 'Map' always knows its set of keys, so given @'Map' s k a@ we can always
46+-- derive @'KnownSet' s k@ by pattern matching on the 'Dict' returned by
47+-- 'keysSet'.
48+newtype Map s k a = Map (Map.Map k a)
49+ deriving newtype (Eq, Ord, Show, Functor, Foldable, NFData)
50+#if MIN_VERSION_hashable(1, 3, 4)
51+ deriving newtype (Hashable.Hashable)
52+#endif
53+ deriving stock (Traversable)
54+type role Map nominal nominal representational
55+
56+-- | Convert to a regular 'Data.Map.Map', forgetting its set of keys.
57+toMap :: forall s k a. Map s k a -> Map.Map k a
58+toMap (Map m) = m
59+
60+-- | @'Key' s k@ is a key of type @k@ that has been verified to be an element
61+-- of the set @s@, and thus verified to be present in all @'Map' s k@ maps.
62+--
63+-- Thus, @'Key' s k@ is a \"refinement\" type of @k@, and this library
64+-- integrates with an implementation of refimenement types in "Refined", so
65+-- the machinery from there can be used to manipulate 'Key's (however see
66+-- 'Data.Set.Refined.revealPredicate').
67+--
68+-- The underlying @k@ value can be obtained with 'unrefine'. A @k@ can be
69+-- validated into an @'Key' s k@ with 'member'.
70+type Key s = Refined (InSet 'Regular s)
71+
72+unsafeCastKey :: forall s k. Coercion k (Key s k)
73+unsafeCastKey = reallyUnsafeUnderlyingRefined
74+
75+unsafeKey :: k -> Key s k
76+unsafeKey = coerceWith unsafeCastKey
77+
78+-- | An existential wrapper for a 'Map' with an as-yet-unknown set of keys.
79+-- Pattern maching on it gives you a way to refer to the set (the parameter
80+-- @s@), e.g.
81+--
82+-- @
83+-- case 'fromMap' ... of
84+-- 'SomeMap' \@s m -> doSomethingWith \@s
85+--
86+-- case 'fromMap' ... of
87+-- 'SomeMap' (m :: 'Map' s k a) -> doSomethingWith \@s
88+-- @
89+data SomeMap k a where
90+ SomeMap :: forall s k a. !(Map s k a) -> SomeMap k a
91+
92+-- | Apply a map with an unknown set of keys to a continuation that can accept
93+-- a map with any set of keys. This gives you a way to refer to the set (the
94+-- parameter @s@), e.g.:
95+--
96+-- @
97+-- 'withMap' ('fromMap' ...) $ \(m :: 'Map' s k a) -> doSomethingWith \@s
98+-- @
99+withMap :: forall k a r. SomeMap k a -> (forall s. Map s k a -> r) -> r
100+withMap (SomeMap m) k = k m
101+
102+-- | Construct a map from a regular 'Data.Map.Map'.
103+fromMap :: forall k a. Map.Map k a -> SomeMap k a
104+fromMap m = SomeMap (Map m)
105+
106+-- | An existential wrapper for a 'Map' with an as-yet-unknown set of keys,
107+-- together with a proof of some fact @p@ about the set. Pattern matching on it
108+-- gives you a way to refer to the set (the parameter @s@). Functions that
109+-- change the set of keys in a map will return the map in this way, together
110+-- with a proof that somehow relates the keys set to the function's inputs.
111+data SomeMapWith p k a where
112+ SomeMapWith :: forall s k a p. !(Map s k a) -> !(p s) -> SomeMapWith p k a
113+
114+-- | Apply a map with proof for an unknown set of keys to a continuation that
115+-- can accept a map with any set of keys satisfying the proof. This gives you a
116+-- way to refer to the set (the parameter @s@).
117+withMapWith
118+ :: forall k a r p. SomeMapWith p k a -> (forall s. Map s k a -> p s -> r) -> r
119+withMapWith (SomeMapWith m p) k = k m p
120+
121+-- | An existential wrapper for a pair of maps with as-yet-unknown sets of keys,
122+-- together with a proof of some fact @p@ relating them.
123+data Some2MapWith p k a b where
124+ Some2MapWith
125+ :: forall s t k a b p. !(Map s k a)
126+ -> !(Map t k b)
127+ -> !(p s t)
128+ -> Some2MapWith p k a b
129+
130+-- | Apply a pair of maps with proof for unknown sets of keys to a continuation
131+-- that can accept any pair of maps with any sets of keys satisfying the proof.
132+-- This gives you a way to refer to the sets (the parameters @s@ and @t@).
133+with2MapWith
134+ :: forall k a b r p. Some2MapWith p k a b
135+ -> (forall s t. Map s k a -> Map t k b -> p s t -> r)
136+ -> r
137+with2MapWith (Some2MapWith m1 m2 p) k = k m1 m2 p
138+
139+-- | An empty map.
140+empty :: forall k a. SomeMapWith (EmptyProof 'Regular) k a
141+empty = SomeMapWith (Map Map.empty) $ EmptyProof unsafeSubset
142+
143+-- | Create a map from a set of keys, and a function that for each key computes
144+-- the corresponding value.
145+fromSet :: forall s k a. KnownSet s k => (Key s k -> a) -> Map s k a
146+fromSet f = Map $ Map.fromSet (f . unsafeKey) (reflect $ Proxy @s)
147+
148+-- | Delete a key and its value from the map if present, returning a potentially
149+-- smaller map.
150+delete
151+ :: forall s k a. Ord k
152+ => k -> Map s k a -> SomeMapWith (SupersetProof 'Regular s) k a
153+delete k (Map m) = SomeMapWith (Map $ Map.delete k m)
154+ $ SupersetProof unsafeSubset
155+
156+-- | If the key is in the map, return the proof of this, and the associated
157+-- value; otherwise return 'Nothing'.
158+lookup :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a)
159+lookup k (Map m) = (unsafeKey k,) <$> Map.lookup k m
160+
161+-- | Given a key that is proven to be in the map, return the associated value.
162+--
163+-- Unlike 'Data.Map.!' from "Data.Map", this function is total, as it is
164+-- impossible to obtain a @'Key' s k@ for a key that is not in the map
165+-- @'Map' s k a@.
166+(!) :: forall s k a. Ord k => Map s k a -> Key s k -> a
167+(!) (Map m) k = case Map.lookup (unrefine k) m of
168+ Nothing -> error "(!): bug: Data.Map.Refined has been subverted"
169+ Just x -> x
170+
171+-- | If a key is in the map, return the proof that it is.
172+member :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k)
173+member k (Map m)
174+ | k `Map.member` m = Just (unsafeKey k)
175+ | otherwise = Nothing
176+
177+-- | Find the largest key smaller than the given one, and return the
178+-- associated key-value pair.
179+lookupLT :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a)
180+lookupLT = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupLT @k @a
181+
182+-- | Find the smallest key greater than the given one, and return the
183+-- associated key-value pair.
184+lookupGT :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a)
185+lookupGT = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupGT @k @a
186+
187+-- | Find the largest key smaller or equal to the given one, and return the
188+-- associated key-value pair.
189+lookupLE :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a)
190+lookupLE = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupLE @k @a
191+
192+-- | Find the smallest key greater or equal to the given one, and return the
193+-- associated key-value pair.
194+lookupGE :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a)
195+lookupGE = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupGE @k @a
196+
197+-- | If a map is empty, return a proof that it is.
198+null :: forall s k a. Map s k a -> Maybe (EmptyProof 'Regular s)
199+null (Map m)
200+ | Map.null m = Just $ EmptyProof unsafeSubset
201+ | otherwise = Nothing
202+
203+-- | If all keys of the first map are also present in the second map, and the
204+-- given function returns 'True' for their associated values, return a proof
205+-- that the keys form a subset.
206+isSubmapOfBy
207+ :: forall s t k a b. Ord k
208+ => (a -> b -> Bool)
209+ -> Map s k a
210+ -> Map t k b
211+ -> Maybe (SubsetProof 'Regular s t)
212+isSubmapOfBy f (Map m1) (Map m2)
213+ | Map.isSubmapOfBy f m1 m2 = Just $ SubsetProof unsafeSubset
214+ | otherwise = Nothing
215+
216+-- | If two maps are disjoint (i.e. their intersection is empty), return a proof
217+-- of that.
218+disjoint
219+ :: forall s t k a b. Ord k
220+ => Map s k a -> Map t k b -> Maybe (DisjointProof 'Regular s t)
221+disjoint (Map m1) (Map m2)
222+#if MIN_VERSION_containers(0, 6, 2)
223+ | Map.disjoint m1 m2
224+#elif MIN_VERSION_containers(0, 5, 8)
225+ | Const (Any False) <- Map.mergeA
226+ (Map.traverseMissing \_ _ -> Const mempty)
227+ (Map.traverseMissing \_ _ -> Const mempty)
228+ (Map.zipWithAMatched \_ _ _ -> Const $ Any True)
229+ m1
230+ m2
231+#else
232+ | Map.null $ MapStrict.intersectionWith (\_ _ -> ()) m1 m2
233+#endif
234+ = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g
235+ | otherwise = Nothing
236+
237+-- | Given two maps proven to have the same keys, for each key apply the
238+-- function to the associated values, to obtain a new map with the same keys.
239+zipWithKey
240+ :: 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)
244+ (error "zipWithKey: bug: Data.Map.Refined has been subverted")
245+ (error "zipWithKey: bug: Data.Map.Refined has been subverted")
246+ m1
247+ m2
248+
249+-- | Remove the keys that appear in the second map from the first map.
250+difference
251+ :: forall s t k a b. Ord k
252+ => Map s k a -> Map t k b -> SomeMapWith (DifferenceProof 'Regular s t) k a
253+difference (Map m1) (Map m2) = SomeMapWith (Map $ Map.difference m1 m2)
254+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
255+
256+-- | Apply a function to all values in a map, together with their corresponding
257+-- keys, that are proven to be in the map. The set of keys remains the same.
258+mapWithKey :: forall s k a b. (Key s k -> a -> b) -> Map s k a -> Map s k b
259+mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
260+ $ Map.mapWithKey @k @a @b
261+
262+-- | Map an 'Applicative' transformation in ascending order of keys, with access
263+-- to each value's corresponding key and a proof that it is in the map. The set
264+-- of keys remains unchanged.
265+traverseWithKey
266+ :: forall s f k a b. Applicative f
267+ => (Key s k -> a -> f b) -> Map s k a -> f (Map s k b)
268+traverseWithKey f (Map m) = Map <$> Map.traverseWithKey (f . unsafeKey) m
269+
270+-- | Map each key-value pair of a map into a monoid (with proof that the key was
271+-- in the map), and combine the results using '<>'.
272+foldMapWithKey
273+ :: forall s k a m. Monoid m => (Key s k -> a -> m) -> Map s k a -> m
274+foldMapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
275+ $ Map.foldMapWithKey @m @k @a
276+
277+-- | Right associative fold with a lazy accumulator.
278+foldrWithKey :: forall s k a b. (Key s k -> a -> b -> b) -> b -> Map s k a -> b
279+foldrWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
280+ $ Map.foldrWithKey @k @a @b
281+
282+-- | Left associative fold with a lazy accumulator.
283+foldlWithKey :: forall s k a b. (b -> Key s k -> a -> b) -> b -> Map s k a -> b
284+foldlWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
285+ $ Map.foldlWithKey @b @k @a
286+
287+-- | Right associative fold with a strict accumulator.
288+foldrWithKey' :: forall s k a b. (Key s k -> a -> b -> b) -> b -> Map s k a -> b
289+foldrWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce
290+ $ Map.foldrWithKey' @k @a @b
291+
292+-- | Left associative fold with a strict accumulator.
293+foldlWithKey' :: forall s k a b. (b -> Key s k -> a -> b) -> b -> Map s k a -> b
294+foldlWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce
295+ $ Map.foldlWithKey' @b @k @a
296+
297+-- | Return the set of keys in the map, with the contents of the set still
298+-- tracked by the @s@ parameter. See "Data.Set.Refined".
299+keysSet :: forall s k a. Map s k a -> Set s k
300+keysSet (Map m) = reify (Map.keysSet m)
301+ \(_ :: Proxy s') -> case unsafeCoerce Refl :: s :~: s' of
302+ Refl -> Dict
303+
304+-- | Convert to a list of key-value pairs in ascending order of keys.
305+toList :: forall s k a. Map s k a -> [(Key s k, a)]
306+toList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.toAscList @k @a
307+
308+-- | Convert to a list of key-value pairs in descending order of keys.
309+toDescList :: forall s k a. Map s k a -> [(Key s k, a)]
310+toDescList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.toDescList @k @a
311+
312+-- | Retain only the key-value pairs that satisfy the predicate, returning a
313+-- potentially smaller map.
314+filterWithKey
315+ :: forall s k a. (Key s k -> a -> Bool)
316+ -> Map s k a
317+ -> SomeMapWith (SupersetProof 'Regular s) k a
318+filterWithKey p (Map m)
319+ = SomeMapWith (Map $ Map.filterWithKey (p . unsafeKey) m)
320+ $ SupersetProof unsafeSubset
321+
322+-- | Restrict a map to only those keys that are elements of @t@.
323+restrictKeys
324+ :: forall s t k a. (Ord k, KnownSet t k)
325+ => Map s k a -> SomeMapWith (IntersectionProof 'Regular s t) k a
326+restrictKeys (Map m) = SomeMapWith
327+#if MIN_VERSION_containers(0, 5, 8)
328+ (Map $ Map.restrictKeys m $ reflect $ Proxy @t)
329+#else
330+ (Map $ Map.intersectionWith const m $ Map.fromSet id $ reflect $ Proxy @t)
331+#endif
332+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
333+
334+-- | Remove all keys that are elements of @t@ from the map.
335+withoutKeys
336+ :: forall s t k a. (Ord k, KnownSet t k)
337+ => Map s k a -> SomeMapWith (DifferenceProof 'Regular s t) k a
338+withoutKeys (Map m) = SomeMapWith
339+#if MIN_VERSION_containers(0, 5, 8)
340+ (Map $ Map.withoutKeys m $ reflect $ Proxy @t)
341+#else
342+ (Map $ Map.difference m $ Map.fromSet id $ reflect $ Proxy @t)
343+#endif
344+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
345+
346+-- | Partition a map into two disjoint submaps: those whose key-value pairs
347+-- satisfy the predicate, and those whose don't.
348+partitionWithKey
349+ :: forall s k a. Ord k -- TODO: this is only used in the proof
350+ => (Key s k -> a -> Bool)
351+ -> Map s k a
352+ -> Some2MapWith (PartitionProof 'Regular s k) k a a
353+partitionWithKey p (Map m) = case Map.partitionWithKey (p . unsafeKey) m of
354+ (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof
355+ do \k -> case Map.lookup (unrefine k) m of
356+ Nothing
357+ -> error "partitionWithKey: bug: Data.Map.Refined has been subverted"
358+ Just x -> if p k x
359+ then Left $ unsafeKey $ unrefine k
360+ else Right $ unsafeKey $ unrefine k
361+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
362+
363+-- | Divide a map into two disjoint submaps at a point where the predicate on
364+-- the keys stops holding.
365+--
366+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
367+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
368+-- chosen in an unspecified way.
369+spanAntitone
370+ :: forall s k a. (Key s k -> Bool)
371+ -> Map s k a
372+ -> Some2MapWith (PartialPartitionProof 'Regular s) k a a
373+spanAntitone p (Map m) =
374+#if MIN_VERSION_containers(0, 5, 8)
375+ case Map.spanAntitone (p . unsafeKey) m of
376+ (m1, m2)
377+#else
378+ case List.span (p . unsafeKey . fst) $ Map.toAscList m of
379+ (xs1, xs2)
380+ | let m1 = Map.fromDistinctAscList xs1
381+ , let m2 = Map.fromDistinctAscList xs2
382+#endif
383+ -> Some2MapWith (Map m1) (Map m2) $ PartialPartitionProof
384+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
385+
386+-- | Return two disjoint submaps: those whose keys are less than the given key,
387+-- and those whose keys are greater than the given key. If the key was in the
388+-- map, also return the associated value and the proof that it was in the map.
389+splitLookup
390+ :: forall s k a. Ord k
391+ => k -> Map s k a -> Some2MapWith (SplitProof 'Regular s (Key s k, a)) k a a
392+splitLookup k (Map m) = case Map.splitLookup k m of
393+ (!m1, v, !m2) -> Some2MapWith (Map m1) (Map m2) $ SplitProof
394+ ((unsafeKey k,) <$> v) unsafeSubset \f g -> unsafeSubsetWith2 f g
395+
396+-- | Retrieves the key-value pair corresponding to the smallest key of the map,
397+-- and the map with that pair removed; or a proof that the map was empty.
398+minViewWithKey
399+ :: forall s k a. Map s k a
400+ -> Either
401+ (EmptyProof 'Regular s)
402+ ((Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a)
403+minViewWithKey (Map m) = case Map.minViewWithKey m of
404+ Nothing -> Left $ EmptyProof unsafeSubset
405+ Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s @k) $ coerce kv,)
406+ $ SomeMapWith (Map m') $ SupersetProof unsafeSubset
407+
408+-- | Retrieves the key-value pair corresponding to the greatest key of the map,
409+-- and the map with that pair removed; or a proof that the map was empty.
410+maxViewWithKey
411+ :: forall s k a. Map s k a
412+ -> Either
413+ (EmptyProof 'Regular s)
414+ ((Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a)
415+maxViewWithKey (Map m) = case Map.maxViewWithKey m of
416+ Nothing -> Left $ EmptyProof unsafeSubset
417+ Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s @k) $ coerce kv,)
418+ $ SomeMapWith (Map m') $ SupersetProof unsafeSubset
419+
420+-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then
421+-- @s@ and @t@ actually stand for the same set and @'Key' s@ can be safely
422+-- interconverted with @'Key' t@.
423+--
424+-- The requirement that the weakenings are natural transformations ensures that
425+-- they don't actually alter the keys. To build these you can compose ':->''s
426+-- from proofs returned by functions in this module, or "Refined" functions like
427+-- 'andLeft' or 'leftOr'.
428+castKey
429+ :: forall s t k. (forall x. Key s x -> Key t x)
430+ -> (forall x. Key t x -> Key s x)
431+ -> Coercion (Key s k) (Key t k)
432+castKey = castRefined
433+
434+-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps
435+-- can be interconverted too. For example, 'zipWithKey' can be implemented via
436+-- 'Data.Map.Refined.intersectionWithKey' by proving that the set of keys
437+-- remains unchanged:
438+--
439+-- @
440+-- 'zipWithKey'
441+-- :: forall s k a b c. 'Ord' k
442+-- => ('Key' s k -> a -> b -> c) -> 'Map' s k a -> 'Map' s k b -> 'Map' s k c
443+-- 'zipWithKey' f m1 m2
444+-- | v'SomeMapWith' @r m proof <- 'Data.Map.Refined.intersectionWithKey' (f . 'andLeft') m1 m2
445+-- , v'IntersectionProof' p1 p2 <- proof
446+-- , ( v'Coercion' :: t'Coercion' ('Map' r k c) ('Map' s k c))
447+-- <- app $ 'cast' $ 'castKey' ('andLeft' . p1) (p2 'id' 'id')
448+-- = 'coerce' m
449+-- where
450+-- app :: t'Coercion' f g -> t'Coercion' (f x) (g x)
451+-- app v'Coercion' = v'Coercion'
452+-- @
453+cast
454+ :: forall s t k. (forall x. Coercion (Key s x) (Key t x))
455+ -> Coercion (Map s k) (Map t k)
456+cast Coercion = Coercion
457+
458+instance FunctorWithIndex (Key s k) (Map s k) where
459+ imap = mapWithKey
460+
461+instance FoldableWithIndex (Key s k) (Map s k) where
462+ ifoldMap = foldMapWithKey
463+
464+instance TraversableWithIndex (Key s k) (Map s k) where
465+ itraverse = traverseWithKey
466+
467+-- | Similar to the instance for functions -- zip corresponding keys. To use
468+-- '<*>'/'liftA2' without 'KnownSet' see 'zipWithKey'.
469+instance (Ord k, KnownSet s k) => Applicative (Map s k) where
470+ pure x = fromSet \_ -> x
471+ (<*>) = zipWithKey (const id)
472+
473+-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the
474+-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
475+bind :: forall s k a b. Ord k => Map s k a -> (a -> Map s k b) -> Map s k b
476+bind m f = mapWithKey (\k x -> f x ! k) m
477+
478+-- | Similar to the instance for functions. To use '>>=' without 'KnownSet' see
479+-- 'bind'.
480+instance (Ord k, KnownSet s k) => Monad (Map s k) where
481+ (>>=) = bind
482+
483+-- | Similar to the instance for functions. See also
484+-- 'Data.Map.Refined.backpermuteKeys'.
485+instance (Ord k, KnownSet s k) => MonadReader (Key s k) (Map s k) where
486+ ask = fromSet id
487+ local f m = mapWithKey (\k _ -> m ! f k) m
488+
489+-- | Append the values at the corresponding keys
490+instance (Ord k, Semigroup a) => Semigroup (Map s k a) where
491+ (<>) = zipWithKey (const (<>))
492+
493+instance (Ord k, KnownSet s k, Monoid a) => Monoid (Map s k a) where
494+ mempty = fromSet \_ -> mempty
495+
496+-- | Similar to the instance for functions
497+instance (Ord k, KnownSet s k) => Distributive (Map s k) where
498+ collect = collectRep
499+ distribute = distributeRep
500+
501+-- | Witness isomorphism with functions from @'Key' s k@
502+instance (Ord k, KnownSet s k) => Representable (Map s k) where
503+ type Rep (Map s k) = Key s k
504+ index = (!)
505+ tabulate = fromSet
506+
507+#if MIN_VERSION_hashable(1, 3, 4)
508+#else
509+instance (Hashable.Hashable a, Hashable.Hashable k)
510+ => Hashable.Hashable (Map s k a) where
511+ hashWithSalt s (Map m) = Map.foldlWithKey'
512+ (\s' k v -> Hashable.hashWithSalt (Hashable.hashWithSalt s' k) v)
513+ (Hashable.hashWithSalt s (Map.size m))
514+ m
515+#endif
src/Data/Map/Refined.hsnew file+384
@@ -0,0 +1,384 @@
1+-- | This module defines a way to prove that a key exists in a map, so that the
2+-- key can be used to index into the map without using a 'Maybe', or manually
3+-- handling the \"impossible\" case with 'error' or other partial functions.
4+--
5+-- To do this, @'Map' s k v@ has a type parameter @s@ that identifies its set of
6+-- keys, so that if another map has the same type parameter, you know that map
7+-- has the same set of keys. There is @'Key' s k@, a type of keys that have been
8+-- validated to belong to the set identified by @s@, and for which the operation
9+-- of indexing into a @'Map' s k v@ (only for the same @s@) can proceed without
10+-- failure (see '!'). The type @s@ itself has no internal structure, rather it
11+-- is merely a skolem type variable (rank-2 polymorphism
12+-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
13+--
14+-- Like "Data.Map", functions in this module are strict in the keys but lazy in
15+-- the values. The "Data.Map.Strict.Refined" module reuses the same 'Map' type
16+-- but provides functions that operate strictly on the values.
17+--
18+-- = Warning
19+-- This module together with "Data.Map" rely on 'Eq' and 'Ord' instances being
20+-- lawful: that '==' is an equivalence relation, and that the 'Ord' operations
21+-- define a total order on the quotient defined by this equivalence relation; at
22+-- least for the subset of keys that are actually encountered at runtime. If
23+-- this assumption is violated, this module may not be able to uphold its
24+-- invariants and may throw errors. In particular beware of NaN in 'Float' and
25+-- 'Double'.
26+module Data.Map.Refined
27+ (
28+ -- * Map type
29+ Common.Map
30+ , Common.Key
31+ -- * Existentials and common proofs
32+ , Common.SomeMap(..)
33+ , Common.withMap
34+ , Common.SomeMapWith(..)
35+ , Common.withMapWith
36+ , Common.Some2MapWith(..)
37+ , Common.with2MapWith
38+ , SupersetProof(..)
39+ , EmptyProof(..)
40+ -- * Construction
41+ , Common.empty
42+ , singleton
43+ , SingletonProof(..)
44+ , fromSet
45+ , Common.fromMap
46+ , fromTraversableWithKey
47+ , FromTraversableProof(..)
48+ -- * Insertion
49+ , insert
50+ , InsertProof(..)
51+ , reinsert
52+ , insertLookupWithKey
53+ -- * Deletion/Update
54+ , Common.delete
55+ , adjust
56+ , adjustWithKey
57+ , update
58+ , updateLookupWithKey
59+ -- * Query
60+ , Common.lookup
61+ , (Common.!)
62+ , Common.member
63+ , Common.lookupLT
64+ , Common.lookupGT
65+ , Common.lookupLE
66+ , Common.lookupGE
67+ , Common.null
68+ , Common.isSubmapOfBy
69+ , SubsetProof(..)
70+ , Common.disjoint
71+ , DisjointProof(..)
72+ -- * Combine
73+ , zipWithKey
74+ , bind
75+ , unionWithKey
76+ , UnionProof(..)
77+ , Common.difference
78+ , DifferenceProof(..)
79+ , differenceWithKey
80+ , PartialDifferenceProof(..)
81+ , intersectionWithKey
82+ , IntersectionProof(..)
83+ -- * Traversal
84+ , mapWithKey
85+ , traverseWithKey
86+ , mapAccumLWithKey
87+ , mapAccumRWithKey
88+ , mapKeysWith
89+ , MapProof(..)
90+ , backpermuteKeys
91+ -- * Folds
92+ , Common.foldMapWithKey
93+ , Common.foldrWithKey
94+ , Common.foldlWithKey
95+ , Common.foldrWithKey'
96+ , Common.foldlWithKey'
97+ -- * Conversion
98+ , Common.toMap
99+ , Common.keysSet
100+ , Common.toList
101+ , Common.toDescList
102+ -- * Filter
103+ , Common.restrictKeys
104+ , Common.withoutKeys
105+ , Common.filterWithKey
106+ , Common.partitionWithKey
107+ , PartitionProof(..)
108+ , Common.spanAntitone
109+ , PartialPartitionProof(..)
110+ , mapMaybeWithKey
111+ , mapEitherWithKey
112+ , Common.splitLookup
113+ , SplitProof(..)
114+ -- * Min/Max
115+ , updateMinWithKey
116+ , updateMaxWithKey
117+ , adjustMinWithKey
118+ , adjustMaxWithKey
119+ , Common.minViewWithKey
120+ , Common.maxViewWithKey
121+ -- * Casts
122+ , Common.castKey
123+ , Common.cast
124+ , castFlavor
125+ ) where
126+
127+import Data.Coerce
128+import Data.Container.Refined.Proofs
129+import Data.Container.Refined.Unsafe
130+import Data.Functor
131+import qualified Data.Map as Map
132+import Data.Map.Common.Refined
133+ ( Map(..), Key, unsafeCastKey, unsafeKey, SomeMapWith(..), Some2MapWith(..)
134+ , fromSet, (!), zipWithKey, mapWithKey, traverseWithKey, bind
135+ )
136+import qualified Data.Map.Common.Refined as Common
137+import Data.Traversable
138+import Data.Type.Coercion
139+import Prelude hiding (lookup, null)
140+import Refined
141+import Refined.Unsafe
142+
143+
144+-- | Create a map with a single key-value pair, and return a proof that the
145+-- key is in the resulting map.
146+singleton :: forall k a. k -> a -> SomeMapWith (SingletonProof 'Regular k) k a
147+singleton k v = SomeMapWith (Map $ Map.singleton k v)
148+ $ SingletonProof (unsafeKey k)
149+
150+-- | Create a map from an arbitrary traversable of key-value pairs.
151+fromTraversableWithKey
152+ :: forall t k a. (Traversable t, Ord k)
153+ => (k -> a -> a -> a)
154+ -> t (k, a)
155+ -> SomeMapWith (FromTraversableProof 'Regular t k) k a
156+fromTraversableWithKey f xs = SomeMapWith (Map m) $ FromTraversableProof proof
157+ where
158+ (m, proof) = mapAccumL
159+ (\s (k, v) -> let !s' = Map.insertWithKey f k v s in (s', unsafeKey k))
160+ Map.empty
161+ xs
162+
163+-- | Insert a key-value pair into the map to obtain a potentially larger map,
164+-- guaranteed to contain the given key. If the key was already present, the
165+-- associated value is replaced with the supplied value.
166+insert
167+ :: forall s k a. Ord k
168+ => k -> a -> Map s k a -> SomeMapWith (InsertProof 'Regular k s) k a
169+insert k v (Map m) = SomeMapWith (Map $ Map.insert k v m)
170+ $ InsertProof (unsafeKey k) unsafeSubset
171+
172+-- | Overwrite a key-value pair that is known to already be in the map. The set
173+-- of keys remains the same.
174+reinsert
175+ :: forall s k a. Ord k
176+ => Key s k -> a -> Map s k a -> Map s k a
177+reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.insert @k @a
178+
179+-- | Insert a key-value pair into the map using a combining function, and if
180+-- the key was already present, the old value is returned along with the proof
181+-- that the key was present.
182+insertLookupWithKey
183+ :: forall s k a. Ord k
184+ => (Key s k -> a -> a -> a)
185+ -> k
186+ -> a
187+ -> Map s k a
188+ -> (Maybe (Key s k, a), SomeMapWith (InsertProof 'Regular k s) k a)
189+insertLookupWithKey f k v (Map m)
190+ = case Map.insertLookupWithKey (f . unsafeKey) k v m of
191+ (v', !m') -> ((unsafeKey k,) <$> v',)
192+ $ SomeMapWith (Map m') $ InsertProof (unsafeKey k) unsafeSubset
193+
194+-- | Update the value at a specific key known the be in the map using the given
195+-- function. The set of keys remains the same.
196+adjust :: 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
198+
199+-- | If the given key is in the map, update the associated value using the given
200+-- function with a proof that the key was in the map; otherwise return the map
201+-- unchanged. In any case the set of keys remains the same.
202+adjustWithKey
203+ :: forall s k a. Ord k => (Key s k -> a -> a) -> k -> Map s k a -> Map s k a
204+adjustWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
205+ $ Map.adjustWithKey @k @a
206+
207+-- | Update or delete a key known to be in the map using the given function,
208+-- returning a potentially smaller map.
209+update
210+ :: forall s k a. Ord k
211+ => (a -> Maybe a)
212+ -> Key s k
213+ -> Map s k a
214+ -> SomeMapWith (SupersetProof 'Regular s) k a
215+update f k (Map m) = SomeMapWith (Map $ Map.update f (unrefine k) m)
216+ $ SupersetProof unsafeSubset
217+
218+-- | If the given key is in the map, update or delete it using the given
219+-- function with a proof that the key was in the map; otherwise the map is
220+-- unchanged. Alongside return the new value if it was updated, or the old value
221+-- if it was deleted, and a proof that the key was in the map.
222+updateLookupWithKey
223+ :: forall s k a. Ord k
224+ => (Key s k -> a -> Maybe a)
225+ -> k
226+ -> Map s k a
227+ -> (Maybe (Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a)
228+updateLookupWithKey f k (Map m)
229+ = case Map.updateLookupWithKey (f . unsafeKey) k m of
230+ (v', !m') -> ((unsafeKey k,) <$> v',)
231+ $ SomeMapWith (Map m') $ SupersetProof unsafeSubset
232+
233+-- | Return the union of two maps, with a given combining function for keys that
234+-- exist in both maps simultaneously.
235+--
236+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
237+-- respectively.
238+unionWithKey
239+ :: forall s t k a. Ord k
240+ => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> a -> a)
241+ -> Map s k a
242+ -> Map t k a
243+ -> SomeMapWith (UnionProof 'Regular s t) k a
244+unionWithKey f (Map m1) (Map m2)
245+ = SomeMapWith (Map $ Map.unionWithKey (f . reallyUnsafeRefine) m1 m2)
246+ $ UnionProof unsafeSubset unsafeSubsetWith2
247+
248+-- | For keys that appear in both maps, the given function decides whether the
249+-- key is removed from the first map.
250+--
251+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
252+-- respectively.
253+differenceWithKey
254+ :: forall s t k a b. Ord k
255+ => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> Maybe a)
256+ -> Map s k a
257+ -> Map t k b
258+ -> SomeMapWith (PartialDifferenceProof 'Regular s t) k a
259+differenceWithKey f (Map m1) (Map m2)
260+ = SomeMapWith (Map $ Map.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
261+ $ PartialDifferenceProof unsafeSubset unsafeSubset
262+
263+-- | Return the intersection of two maps with the given combining function.
264+--
265+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
266+-- respectively.
267+intersectionWithKey
268+ :: forall s t k a b c. Ord k
269+ => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> c)
270+ -> Map s k a
271+ -> Map t k b
272+ -> SomeMapWith (IntersectionProof 'Regular s t) k c
273+intersectionWithKey f (Map m1) (Map m2)
274+ = SomeMapWith (Map $ Map.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
275+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
276+
277+-- | Thread an accumularing argument through the map in ascending order of keys.
278+mapAccumLWithKey
279+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
280+ -> a
281+ -> Map s k b
282+ -> (a, Map s k c)
283+mapAccumLWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
284+ $ Map.mapAccumWithKey @a @k @b @c
285+
286+-- | Thread an accumularing argument through the map in descending order of
287+-- keys.
288+mapAccumRWithKey
289+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
290+ -> a
291+ -> Map s k b
292+ -> (a, Map s k c)
293+mapAccumRWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
294+ $ Map.mapAccumRWithKey @a @k @b @c
295+
296+-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
297+-- results into a new map. For keys that were mapped to the same new key, @c@
298+-- acts as the combining function for corresponding values.
299+mapKeysWith
300+ :: forall s k1 k2 a. Ord k2
301+ => (a -> a -> a)
302+ -> (Key s k1 -> k2)
303+ -> Map s k1 a
304+ -> SomeMapWith (MapProof 'Regular s k1 k2) k2 a
305+mapKeysWith f g (Map m)
306+ = SomeMapWith (Map $ Map.mapKeysWith f (g . unsafeKey) m)
307+ $ MapProof (unsafeKey . g) \k2 -> case Map.lookup (unrefine k2) backMap of
308+ Nothing -> error "mapKeysWith: bug: Data.Map.Refined has been subverted"
309+ Just k1 -> k1
310+ where
311+ ~backMap = Map.fromList
312+ [ (k2, unsafeKey k1)
313+ | k1 <- Map.keys m
314+ , let !k2 = g $ unsafeKey k1
315+ ]
316+
317+-- | Apply a function to all values in a map, together with their corresponding
318+-- keys, and collect only the 'Just' results, returning a potentially smaller
319+-- map.
320+mapMaybeWithKey
321+ :: forall s k a b. (Key s k -> a -> Maybe b)
322+ -> Map s k a
323+ -> SomeMapWith (SupersetProof 'Regular s) k b
324+mapMaybeWithKey f (Map m)
325+ = SomeMapWith (Map $ Map.mapMaybeWithKey (f . unsafeKey) m)
326+ $ SupersetProof unsafeSubset
327+
328+-- | Apply a function to all values in a map, together with their corresponding
329+-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
330+-- maps.
331+mapEitherWithKey
332+ :: forall s k a b c. Ord k -- TODO: this is only used in the proof
333+ => (Key s k -> a -> Either b c)
334+ -> Map s k a
335+ -> Some2MapWith (PartitionProof 'Regular s k) k b c
336+mapEitherWithKey p (Map m) = case Map.mapEitherWithKey (p . unsafeKey) m of
337+ (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof
338+ do \k -> case Map.lookup (unrefine k) m of
339+ Nothing
340+ -> error "mapEitherWithKey: bug: Data.Map.Refined has been subverted"
341+ Just x -> case p k x of
342+ Left _ -> Left $ unsafeKey $ unrefine k
343+ Right _ -> Right $ unsafeKey $ unrefine k
344+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
345+
346+-- | Update or delete the value at the smallest key, returning a potentially
347+-- smaller map.
348+updateMinWithKey
349+ :: forall s k a. (Key s k -> a -> Maybe a)
350+ -> Map s k a
351+ -> SomeMapWith (SupersetProof 'Regular s) k a
352+updateMinWithKey f (Map m)
353+ = SomeMapWith (Map $ Map.updateMinWithKey (f . unsafeKey) m)
354+ $ SupersetProof unsafeSubset
355+
356+-- | Update or delete the value at the largest key, returning a potentially
357+-- smaller map.
358+updateMaxWithKey
359+ :: forall s k a. (Key s k -> a -> Maybe a)
360+ -> Map s k a
361+ -> SomeMapWith (SupersetProof 'Regular s) k a
362+updateMaxWithKey f (Map m)
363+ = SomeMapWith (Map $ Map.updateMaxWithKey (f . unsafeKey) m)
364+ $ SupersetProof unsafeSubset
365+
366+-- | Adjust the value at the smallest key. The set of keys remains unchanged.
367+adjustMinWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a
368+adjustMinWithKey f (Map m)
369+ = Map $ Map.updateMinWithKey ((Just .) . f . unsafeKey) m
370+
371+-- | Adjust the value at the greatest key. The set of keys remains unchanged.
372+adjustMaxWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a
373+adjustMaxWithKey f (Map m)
374+ = Map $ Map.updateMaxWithKey ((Just .) . f . unsafeKey) m
375+
376+-- | Apply the inverse image of the given function to the keys of the given map,
377+-- so that for all @k :: 'Key' s2 k2@,
378+-- @'backpermuteKeys' f m '!' k = m '!' f k@.
379+--
380+-- If maps are identified with functions, this computes the composition.
381+backpermuteKeys
382+ :: forall s1 s2 k1 k2 a. (Ord k1, KnownSet s2 k2)
383+ => (Key s2 k2 -> Key s1 k1) -> Map s1 k1 a -> Map s2 k2 a
384+backpermuteKeys f m = fromSet \k -> m ! f k
src/Data/Map/Strict/Refined.hsnew file+423
@@ -0,0 +1,423 @@
1+-- | This module defines a way to prove that a key exists in a map, so that the
2+-- key can be used to index into the map without using a 'Maybe', or manually
3+-- handling the \"impossible\" case with 'error' or other partial functions.
4+--
5+-- To do this, @'Map' s k v@ has a type parameter @s@ that identifies its set of
6+-- keys, so that if another map has the same type parameter, you know that map
7+-- has the same set of keys. There is @'Key' s k@, a type of keys that have been
8+-- validated to belong to the set identified by @s@, and for which the operation
9+-- of indexing into a @'Map' s k v@ (only for the same @s@) can proceed without
10+-- failure (see '!'). The type @s@ itself has no internal structure, rather it
11+-- is merely a skolem type variable (rank-2 polymorphism
12+-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection".
13+--
14+-- Like "Data.Map.Strict", functions in this module are strict in the keys and
15+-- values. The "Data.Map.Refined" module reuses the same 'Map' type but provides
16+-- functions that operate lazily on the values.
17+--
18+-- = Warning
19+-- This module together with "Data.Map.Strict" rely on 'Eq' and 'Ord' instances
20+-- being lawful: that '==' is an equivalence relation, and that the 'Ord'
21+-- operations define a total order on the quotient defined by this equivalence
22+-- relation; at least for the subset of keys that are actually encountered at
23+-- runtime. If this assumption is violated, this module may not be able to
24+-- uphold its invariants and may throw errors. In particular beware of NaN in
25+-- 'Float' and 'Double'.
26+module Data.Map.Strict.Refined
27+ (
28+ -- * Map type
29+ Common.Map
30+ , Common.Key
31+ -- * Existentials and common proofs
32+ , Common.SomeMap(..)
33+ , Common.withMap
34+ , Common.SomeMapWith(..)
35+ , Common.withMapWith
36+ , Common.Some2MapWith(..)
37+ , Common.with2MapWith
38+ , SupersetProof(..)
39+ , EmptyProof(..)
40+ -- * Construction
41+ , Common.empty
42+ , singleton
43+ , SingletonProof(..)
44+ , fromSet
45+ , Common.fromMap
46+ , fromTraversableWithKey
47+ , FromTraversableProof(..)
48+ -- * Insertion
49+ , insert
50+ , InsertProof(..)
51+ , reinsert
52+ , insertLookupWithKey
53+ -- * Deletion/Update
54+ , Common.delete
55+ , adjust
56+ , adjustWithKey
57+ , update
58+ , updateLookupWithKey
59+ -- * Query
60+ , Common.lookup
61+ , (Common.!)
62+ , Common.member
63+ , Common.lookupLT
64+ , Common.lookupGT
65+ , Common.lookupLE
66+ , Common.lookupGE
67+ , Common.null
68+ , Common.isSubmapOfBy
69+ , SubsetProof(..)
70+ , Common.disjoint
71+ , DisjointProof(..)
72+ -- * Combine
73+ , zipWithKey
74+ , bind
75+ , unionWithKey
76+ , UnionProof(..)
77+ , Common.difference
78+ , DifferenceProof(..)
79+ , differenceWithKey
80+ , PartialDifferenceProof(..)
81+ , intersectionWithKey
82+ , IntersectionProof(..)
83+ -- * Traversal
84+ , mapWithKey
85+ , traverseWithKey
86+ , mapAccumLWithKey
87+ , mapAccumRWithKey
88+ , mapKeysWith
89+ , MapProof(..)
90+ , backpermuteKeys
91+ -- * Folds
92+ , Common.foldMapWithKey
93+ , Common.foldrWithKey
94+ , Common.foldlWithKey
95+ , Common.foldrWithKey'
96+ , Common.foldlWithKey'
97+ -- * Conversion
98+ , Common.toMap
99+ , Common.keysSet
100+ , Common.toList
101+ , Common.toDescList
102+ -- * Filter
103+ , Common.restrictKeys
104+ , Common.withoutKeys
105+ , Common.filterWithKey
106+ , Common.partitionWithKey
107+ , PartitionProof(..)
108+ , Common.spanAntitone
109+ , PartialPartitionProof(..)
110+ , mapMaybeWithKey
111+ , mapEitherWithKey
112+ , Common.splitLookup
113+ , SplitProof(..)
114+ -- * Min/Max
115+ , updateMinWithKey
116+ , updateMaxWithKey
117+ , adjustMinWithKey
118+ , adjustMaxWithKey
119+ , Common.minViewWithKey
120+ , Common.maxViewWithKey
121+ -- * Casts
122+ , Common.castKey
123+ , Common.cast
124+ , castFlavor
125+ ) where
126+
127+import Data.Coerce
128+import Data.Container.Refined.Proofs
129+import Data.Container.Refined.Unsafe
130+import Data.Functor
131+import qualified Data.Map.Strict as Map
132+import Data.Map.Common.Refined
133+ ( Map(..), Key, unsafeCastKey, unsafeKey, SomeMapWith(..), Some2MapWith(..)
134+ , (!)
135+ )
136+import qualified Data.Map.Common.Refined as Common
137+import Data.Proxy
138+import Data.Reflection
139+import Data.Traversable
140+import Data.Type.Coercion
141+import Prelude hiding (lookup, null)
142+import Refined
143+import Refined.Unsafe
144+
145+
146+-- | Create a map with a single key-value pair, and return a proof that the
147+-- key is in the resulting map.
148+singleton :: forall k a. k -> a -> SomeMapWith (SingletonProof 'Regular k) k a
149+singleton k v = SomeMapWith (Map $ Map.singleton k v)
150+ $ SingletonProof (unsafeKey k)
151+
152+-- | Create a map from a set of keys, and a function that for each key computes
153+-- the corresponding value.
154+fromSet :: forall s k a. KnownSet s k => (Key s k -> a) -> Map s k a
155+fromSet f = Map $ Map.fromSet (f . unsafeKey) (reflect $ Proxy @s)
156+
157+-- | Create a map from an arbitrary traversable of key-value pairs.
158+fromTraversableWithKey
159+ :: forall t k a. (Traversable t, Ord k)
160+ => (k -> a -> a -> a)
161+ -> t (k, a)
162+ -> SomeMapWith (FromTraversableProof 'Regular t k) k a
163+fromTraversableWithKey f xs = SomeMapWith (Map m) $ FromTraversableProof proof
164+ where
165+ (m, proof) = mapAccumL
166+ (\s (k, v) -> let !s' = Map.insertWithKey f k v s in (s', unsafeKey k))
167+ Map.empty
168+ xs
169+
170+-- | Insert a key-value pair into the map to obtain a potentially larger map,
171+-- guaranteed to contain the given key. If the key was already present, the
172+-- associated value is replaced with the supplied value.
173+insert
174+ :: forall s k a. Ord k
175+ => k -> a -> Map s k a -> SomeMapWith (InsertProof 'Regular k s) k a
176+insert k v (Map m) = SomeMapWith (Map $ Map.insert k v m)
177+ $ InsertProof (unsafeKey k) unsafeSubset
178+
179+-- | Overwrite a key-value pair that is known to already be in the map. The set
180+-- of keys remains the same.
181+reinsert
182+ :: forall s k a. Ord k
183+ => Key s k -> a -> Map s k a -> Map s k a
184+reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.insert @k @a
185+
186+-- | Insert a key-value pair into the map using a combining function, and if
187+-- the key was already present, the old value is returned along with the proof
188+-- that the key was present.
189+insertLookupWithKey
190+ :: forall s k a. Ord k
191+ => (Key s k -> a -> a -> a)
192+ -> k
193+ -> a
194+ -> Map s k a
195+ -> (Maybe (Key s k, a), SomeMapWith (InsertProof 'Regular k s) k a)
196+insertLookupWithKey f k v (Map m)
197+ = case Map.insertLookupWithKey (f . unsafeKey) k v m of
198+ (v', !m') -> ((unsafeKey k,) <$> v',)
199+ $ SomeMapWith (Map m') $ InsertProof (unsafeKey k) unsafeSubset
200+
201+-- | Update the value at a specific key known the be in the map using the given
202+-- function. The set of keys remains the same.
203+adjust :: 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
205+
206+-- | If the given key is in the map, update the associated value using the given
207+-- function with a proof that the key was in the map; otherwise return the map
208+-- unchanged. In any case the set of keys remains the same.
209+adjustWithKey
210+ :: forall s k a. Ord k => (Key s k -> a -> a) -> k -> Map s k a -> Map s k a
211+adjustWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
212+ $ Map.adjustWithKey @k @a
213+
214+-- | Update or delete a key known to be in the map using the given function,
215+-- returning a potentially smaller map.
216+update
217+ :: forall s k a. Ord k
218+ => (a -> Maybe a)
219+ -> Key s k
220+ -> Map s k a
221+ -> SomeMapWith (SupersetProof 'Regular s) k a
222+update f k (Map m) = SomeMapWith (Map $ Map.update f (unrefine k) m)
223+ $ SupersetProof unsafeSubset
224+
225+-- | If the given key is in the map, update or delete it using the given
226+-- function with a proof that the key was in the map; otherwise the map is
227+-- unchanged. Alongside return the new value if it was updated, or the old value
228+-- if it was deleted, and a proof that the key was in the map.
229+updateLookupWithKey
230+ :: forall s k a. Ord k
231+ => (Key s k -> a -> Maybe a)
232+ -> k
233+ -> Map s k a
234+ -> (Maybe (Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a)
235+updateLookupWithKey f k (Map m)
236+ = case Map.updateLookupWithKey (f . unsafeKey) k m of
237+ (v', !m') -> ((unsafeKey k,) <$> v',)
238+ $ SomeMapWith (Map m') $ SupersetProof unsafeSubset
239+
240+-- | Given two maps proven to have the same keys, for each key apply the
241+-- function to the associated values, to obtain a new map with the same keys.
242+zipWithKey
243+ :: forall s k a b c. Ord k
244+ => (Key s k -> a -> b -> c) -> Map s k a -> Map s k b -> Map s k c
245+zipWithKey f (Map m1) (Map m2) = Map
246+ $ Map.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y)
247+ (error "zipWithKey: bug: Data.Map.Strict.Refined has been subverted")
248+ (error "zipWithKey: bug: Data.Map.Strict.Refined has been subverted")
249+ m1
250+ m2
251+
252+-- | Return the union of two maps, with a given combining function for keys that
253+-- exist in both maps simultaneously.
254+--
255+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
256+-- respectively.
257+unionWithKey
258+ :: forall s t k a. Ord k
259+ => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> a -> a)
260+ -> Map s k a
261+ -> Map t k a
262+ -> SomeMapWith (UnionProof 'Regular s t) k a
263+unionWithKey f (Map m1) (Map m2)
264+ = SomeMapWith (Map $ Map.unionWithKey (f . reallyUnsafeRefine) m1 m2)
265+ $ UnionProof unsafeSubset unsafeSubsetWith2
266+
267+-- | For keys that appear in both maps, the given function decides whether the
268+-- key is removed from the first map.
269+--
270+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
271+-- respectively.
272+differenceWithKey
273+ :: forall s t k a b. Ord k
274+ => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> Maybe a)
275+ -> Map s k a
276+ -> Map t k b
277+ -> SomeMapWith (PartialDifferenceProof 'Regular s t) k a
278+differenceWithKey f (Map m1) (Map m2)
279+ = SomeMapWith (Map $ Map.differenceWithKey (f . reallyUnsafeRefine) m1 m2)
280+ $ PartialDifferenceProof unsafeSubset unsafeSubset
281+
282+-- | Return the intersection of two maps with the given combining function.
283+--
284+-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@
285+-- respectively.
286+intersectionWithKey
287+ :: forall s t k a b c. Ord k
288+ => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> c)
289+ -> Map s k a
290+ -> Map t k b
291+ -> SomeMapWith (IntersectionProof 'Regular s t) k c
292+intersectionWithKey f (Map m1) (Map m2)
293+ = SomeMapWith (Map $ Map.intersectionWithKey (f . reallyUnsafeRefine) m1 m2)
294+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
295+
296+-- | Apply a function to all values in a map, together with their corresponding
297+-- keys, that are proven to be in the map. The set of keys remains the same.
298+mapWithKey :: forall s k a b. (Key s k -> a -> b) -> Map s k a -> Map s k b
299+mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
300+ $ Map.mapWithKey @k @a @b
301+
302+-- | Map an 'Applicative' transformation in ascending order of keys, with access
303+-- to each value's corresponding key and a proof that it is in the map. The set
304+-- of keys remains unchanged.
305+traverseWithKey
306+ :: forall s f k a b. Applicative f
307+ => (Key s k -> a -> f b) -> Map s k a -> f (Map s k b)
308+traverseWithKey f (Map m) = Map <$> Map.traverseWithKey (f . unsafeKey) m
309+
310+-- | Thread an accumularing argument through the map in ascending order of keys.
311+mapAccumLWithKey
312+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
313+ -> a
314+ -> Map s k b
315+ -> (a, Map s k c)
316+mapAccumLWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
317+ $ Map.mapAccumWithKey @a @k @b @c
318+
319+-- | Thread an accumularing argument through the map in descending order of
320+-- keys.
321+mapAccumRWithKey
322+ :: forall s k a b c. (a -> Key s k -> b -> (a, c))
323+ -> a
324+ -> Map s k b
325+ -> (a, Map s k c)
326+mapAccumRWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce
327+ $ Map.mapAccumRWithKey @a @k @b @c
328+
329+-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the
330+-- results into a new map. For keys that were mapped to the same new key, @c@
331+-- acts as the combining function for corresponding values.
332+mapKeysWith
333+ :: forall s k1 k2 a. Ord k2
334+ => (a -> a -> a)
335+ -> (Key s k1 -> k2)
336+ -> Map s k1 a
337+ -> SomeMapWith (MapProof 'Regular s k1 k2) k2 a
338+mapKeysWith f g (Map m)
339+ = SomeMapWith (Map $ Map.mapKeysWith f (g . unsafeKey) m)
340+ $ MapProof (unsafeKey . g) \k2 -> case Map.lookup (unrefine k2) backMap of
341+ Nothing -> error
342+ "mapKeysWith: bug: Data.Map.Strict.Refined has been subverted"
343+ Just k1 -> k1
344+ where
345+ ~backMap = Map.fromList
346+ [ (k2, unsafeKey k1)
347+ | k1 <- Map.keys m
348+ , let !k2 = g $ unsafeKey k1
349+ ]
350+
351+-- | Apply a function to all values in a map, together with their corresponding
352+-- keys, and collect only the 'Just' results, returning a potentially smaller
353+-- map.
354+mapMaybeWithKey
355+ :: forall s k a b. (Key s k -> a -> Maybe b)
356+ -> Map s k a
357+ -> SomeMapWith (SupersetProof 'Regular s) k b
358+mapMaybeWithKey f (Map m)
359+ = SomeMapWith (Map $ Map.mapMaybeWithKey (f . unsafeKey) m)
360+ $ SupersetProof unsafeSubset
361+
362+-- | Apply a function to all values in a map, together with their corresponding
363+-- keys, and collect the 'Left' and 'Right' results into separate (disjoint)
364+-- maps.
365+mapEitherWithKey
366+ :: forall s k a b c. Ord k -- TODO: this is only used in the proof
367+ => (Key s k -> a -> Either b c)
368+ -> Map s k a
369+ -> Some2MapWith (PartitionProof 'Regular s k) k b c
370+mapEitherWithKey p (Map m) = case Map.mapEitherWithKey (p . unsafeKey) m of
371+ (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof
372+ do \k -> case Map.lookup (unrefine k) m of
373+ Nothing -> error
374+ "mapEitherWithKey: bug: Data.Map.Strict.Refined has been subverted"
375+ Just x -> case p k x of
376+ Left _ -> Left $ unsafeKey $ unrefine k
377+ Right _ -> Right $ unsafeKey $ unrefine k
378+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
379+
380+-- | Update or delete the value at the smallest key, returning a potentially
381+-- smaller map.
382+updateMinWithKey
383+ :: forall s k a. (Key s k -> a -> Maybe a)
384+ -> Map s k a
385+ -> SomeMapWith (SupersetProof 'Regular s) k a
386+updateMinWithKey f (Map m)
387+ = SomeMapWith (Map $ Map.updateMinWithKey (f . unsafeKey) m)
388+ $ SupersetProof unsafeSubset
389+
390+-- | Update or delete the value at the largest key, returning a potentially
391+-- smaller map.
392+updateMaxWithKey
393+ :: forall s k a. (Key s k -> a -> Maybe a)
394+ -> Map s k a
395+ -> SomeMapWith (SupersetProof 'Regular s) k a
396+updateMaxWithKey f (Map m)
397+ = SomeMapWith (Map $ Map.updateMaxWithKey (f . unsafeKey) m)
398+ $ SupersetProof unsafeSubset
399+
400+-- | Adjust the value at the smallest key. The set of keys remains unchanged.
401+adjustMinWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a
402+adjustMinWithKey f (Map m)
403+ = Map $ Map.updateMinWithKey ((Just .) . f . unsafeKey) m
404+
405+-- | Adjust the value at the greatest key. The set of keys remains unchanged.
406+adjustMaxWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a
407+adjustMaxWithKey f (Map m)
408+ = Map $ Map.updateMaxWithKey ((Just .) . f . unsafeKey) m
409+
410+-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the
411+-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions.
412+bind :: forall s k a b. Ord k => Map s k a -> (a -> Map s k b) -> Map s k b
413+bind m f = mapWithKey (\k x -> f x ! k) m
414+
415+-- | Apply the inverse image of the given function to the keys of the given map,
416+-- so that for all @k :: 'Key' s2 k2@,
417+-- @'backpermuteKeys' f m '!' k = m '!' f k@.
418+--
419+-- If maps are identified with functions, this computes the composition.
420+backpermuteKeys
421+ :: forall s1 s2 k1 k2 a. (Ord k1, KnownSet s2 k2)
422+ => (Key s2 k2 -> Key s1 k1) -> Map s1 k1 a -> Map s2 k2 a
423+backpermuteKeys f m = fromSet \k -> m ! f k
src/Data/Set/Refined.hsnew file+556
@@ -0,0 +1,556 @@
1+{-# LANGUAGE CPP #-}
2+-- | This module implements a way of tracking the contents of a 'Data.Set.Set'
3+-- at the type level, and provides utilities for manipulating such sets.
4+--
5+-- The contents of a set are associated with a type parameter, e.g. @s@, so that
6+-- whenever you see the same type parameter, you know you are working with the
7+-- same set. The type @s@ itself has no internal structure, rather it is merely
8+-- a skolem type variable (rank-2 polymorphism 'Control.Monad.ST.runST' trick)
9+-- introduced by "Data.Reflection".
10+--
11+-- = Warning
12+-- This module together with "Data.Set" rely on 'Eq' and 'Ord' instances being
13+-- lawful: that '==' is an equivalence relation, and that the 'Ord' operations
14+-- define a total order on the quotient defined by this equivalence relation; at
15+-- least for the subset of values that are actually encountered at runtime. If
16+-- this assumption is violated, this module may not be able to uphold its
17+-- invariants and may throw errors. In particular beware of NaN in 'Float' and
18+-- 'Double'.
19+module Data.Set.Refined
20+ (
21+ -- * Set type
22+ KnownSet
23+ , Set
24+ -- * Refinement type
25+ , InSet(..)
26+ , Flavor(Regular)
27+ , Element
28+ , revealPredicate
29+ -- * Existentials and common proofs
30+ , SomeSet(..)
31+ , withSet
32+ , SomeSetWith(..)
33+ , withSetWith
34+ , Some2SetWith(..)
35+ , with2SetWith
36+ , (:->)
37+ , SupersetProof(..)
38+ , EmptyProof(..)
39+ -- * Construction
40+ , empty
41+ , singleton
42+ , SingletonProof(..)
43+ , fromSet
44+ , fromTraversable
45+ , FromTraversableProof(..)
46+ -- * Insertion
47+ , insert
48+ , InsertProof(..)
49+ -- * Deletion
50+ , delete
51+ -- * Query
52+ , member
53+ , lookupLT
54+ , lookupGT
55+ , lookupLE
56+ , lookupGE
57+ , null
58+ , isSubsetOf
59+ , SubsetProof(..)
60+ , disjoint
61+ , DisjointProof(..)
62+ -- * Combine
63+ , union
64+ , UnionProof(..)
65+ , difference
66+ , DifferenceProof(..)
67+ , intersection
68+ , IntersectionProof(..)
69+ , cartesianProduct
70+ , ProductProof(..)
71+ , disjointUnion
72+ , CoproductProof(..)
73+ -- * Filter
74+ , filter
75+ , partition
76+ , PartitionProof(..)
77+ , spanAntitone
78+ , PartialPartitionProof(..)
79+ , splitMember
80+ , SplitProof(..)
81+ -- * Map
82+ , map
83+ , MapProof(..)
84+ -- * Folds
85+ , foldMap
86+ , foldr
87+ , foldl
88+ , foldr'
89+ , foldl'
90+ -- * Min/Max
91+ , minView
92+ , maxView
93+ -- * Conversion
94+ , toList
95+ , toDescList
96+ , asIntSet
97+ , asHashSet
98+ -- * Casts
99+ , castElement
100+ , cast
101+ , castFlavor
102+ ) where
103+
104+import Data.Coerce
105+import Data.Constraint (Dict(..))
106+import Data.Container.Refined.Conversion
107+import Data.Container.Refined.Hashable
108+import Data.Container.Refined.Proofs
109+import Data.Container.Refined.Unsafe
110+import qualified Data.Foldable as Foldable
111+import qualified Data.Map as Map
112+import Data.Proxy
113+import Data.Reflection
114+import qualified Data.Set as Set
115+import Data.Traversable
116+import Data.Type.Coercion
117+import Data.Type.Equality ((:~:)(..))
118+import Data.Typeable (Typeable)
119+import GHC.Exts (Proxy#, proxy#)
120+import Prelude hiding (filter, foldl, foldMap, foldr, map, null)
121+import Refined
122+import Refined.Unsafe
123+import Unsafe.Coerce
124+
125+#if MIN_VERSION_containers(0, 5, 8)
126+#else
127+import qualified Data.List as List
128+#endif
129+
130+
131+-- | To use "Refined" machinery that uses the 'Predicate' typeclass you will
132+-- need to pattern match on this 'Dict'.
133+--
134+-- The reason is that in the default /fast/ implementation of reflection, we
135+-- don't have @'Typeable' s@, which "Refined" wants for pretty-printing
136+-- exceptions. We /can/ provide @'Typeable' s@, but at the cost of using the
137+-- /slow/ implementation of reflection.
138+revealPredicate
139+ :: forall s a. (Typeable a, Ord a, KnownSet s a)
140+ => Dict (Predicate (InSet 'Regular s) a)
141+revealPredicate = reifyTypeable (reflect (Proxy @s))
142+ \(_ :: Proxy s') ->
143+ reflect (Proxy @s') `seq`
144+ -- ^ Work around https://github.com/ekmett/reflection/issues/54
145+ case unsafeCoerce Refl :: s :~: s' of
146+ Refl -> Dict
147+
148+-- | @'Element' s a@ is a value of type @a@ that has been verified to be an
149+-- element of @s@.
150+--
151+-- Thus, @'Element' s a@ is a \"refinement\" type of @a@, and this library
152+-- integrates with an implementation of refimenement types in "Refined", so
153+-- the machinery from there can be used to manipulate 'Element's (however see
154+-- 'revealPredicate').
155+--
156+-- The underlying @a@ value can be obtained with 'unrefine'. An @a@ can be
157+-- validated into an @'Element' s a@ with 'member'.
158+type Element s = Refined (InSet 'Regular s)
159+
160+unsafeCastElement :: forall s a. Coercion a (Element s a)
161+unsafeCastElement = reallyUnsafeUnderlyingRefined
162+
163+unsafeElement :: a -> Element s a
164+unsafeElement = coerceWith unsafeCastElement
165+
166+-- | An existential wrapper for an as-yet-unknown set. Pattern maching on it
167+-- gives you a way to refer to the set, e.g.
168+--
169+-- @
170+-- case 'fromSet' ... of
171+-- 'SomeSet' \@s _ -> doSomethingWith \@s
172+--
173+-- case 'fromSet' ... of
174+-- 'SomeSet' (_ :: 'Proxy#' s) -> doSomethingWith \@s
175+-- @
176+data SomeSet a where
177+ SomeSet :: forall s a. KnownSet s a => Proxy# s -> SomeSet a
178+
179+-- | Apply an unknown set to a continuation that can accept any set. This gives
180+-- you a way to refer to the set (the parameter @s@), e.g.:
181+--
182+-- @
183+-- 'withSet' ('fromSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s
184+-- @
185+withSet
186+ :: forall a r. SomeSet a -> (forall s. KnownSet s a => Proxy s -> r) -> r
187+withSet (SomeSet (_ :: Proxy# s)) k = k $ Proxy @s
188+
189+-- | Construct a set from a regular 'Data.Set.Set'.
190+fromSet :: forall a. Set.Set a -> SomeSet a
191+fromSet s = reify s \(_ :: Proxy s) -> SomeSet @s proxy#
192+
193+-- | An existential wrapper for an as-yet-unknown set, together with a proof of
194+-- some fact @p@ about the set. Pattern matching on it gives you a way to refer
195+-- to the set (the parameter @s@). Most functions will return a set in this way,
196+-- together with a proof that somehow relates the set to the function's inputs.
197+data SomeSetWith p a where
198+ SomeSetWith :: forall s a p. KnownSet s a => !(p s) -> SomeSetWith p a
199+
200+-- | Apply an unknown set with proof to a continuation that can accept any set
201+-- satisfying the proof. This gives you a way to refer to the set (the parameter
202+-- @s@).
203+withSetWith
204+ :: forall a r p. SomeSetWith p a -> (forall s. KnownSet s a => p s -> r) -> r
205+withSetWith (SomeSetWith p) k = k p
206+
207+-- | An existential wrapper for an as-yet-unknown pair of sets, together with
208+-- a proof of some fact @p@ relating them.
209+data Some2SetWith p a where
210+ Some2SetWith
211+ :: forall s t a p. (KnownSet s a, KnownSet t a)
212+ => !(p s t) -> Some2SetWith p a
213+
214+-- | Apply a pair of unknown sets with proof to a continuation that can accept
215+-- any pair of sets satisfying the proof. This gives you a way to refer to the
216+-- sets (the parameters @s@ and @t@).
217+with2SetWith
218+ :: forall a r p. Some2SetWith p a
219+ -> (forall s t. (KnownSet s a, KnownSet t a) => p s t -> r)
220+ -> r
221+with2SetWith (Some2SetWith p) k = k p
222+
223+-- | An empty set.
224+empty :: forall a. SomeSetWith (EmptyProof 'Regular) a
225+empty = reify Set.empty \(_ :: Proxy r)
226+ -> SomeSetWith @r $ EmptyProof unsafeSubset
227+
228+-- | Create a set with a single element.
229+singleton :: forall a. a -> SomeSetWith (SingletonProof 'Regular a) a
230+singleton x = reify (Set.singleton x) \(_ :: Proxy r)
231+ -> SomeSetWith @r $ SingletonProof $ unsafeElement x
232+
233+-- | Create a set from the elements of an arbitrary traversable.
234+fromTraversable
235+ :: forall t a. (Traversable t, Ord a)
236+ => t a -> SomeSetWith (FromTraversableProof 'Regular t a) a
237+fromTraversable xs = reify set \(_ :: Proxy r)
238+ -> SomeSetWith @r $ FromTraversableProof
239+ $ unsafeCoerce @(t (Element _ a)) @(t (Element r a)) proof
240+ where
241+ (set, proof) = mapAccumL
242+ (\s x -> let !s' = Set.insert x s in (s', unsafeElement x))
243+ Set.empty
244+ xs
245+
246+-- | Insert an element in a set.
247+insert
248+ :: forall s a. (Ord a, KnownSet s a)
249+ => a -> SomeSetWith (InsertProof 'Regular a s) a
250+insert x = reify (Set.insert x $ reflect $ Proxy @s) \(_ :: Proxy r)
251+ -> SomeSetWith @r $ InsertProof (unsafeElement x) unsafeSubset
252+
253+-- | Delete an element from a set.
254+delete
255+ :: forall s a. (Ord a, KnownSet s a)
256+ => a -> SomeSetWith (SupersetProof 'Regular s) a
257+delete x = reify (Set.delete x $ reflect $ Proxy @s) \(_ :: Proxy r)
258+ -> SomeSetWith @s $ SupersetProof unsafeSubset
259+
260+-- | If an element is in the set, return the proof that it is.
261+member :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a)
262+member x
263+ | x `Set.member` reflect (Proxy @s) = Just $ unsafeElement x
264+ | otherwise = Nothing
265+
266+-- | Find the largest element smaller than the given one.
267+lookupLT :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a)
268+lookupLT x = gcoerceWith (unsafeCastElement @s @a) $ coerce
269+ $ Set.lookupLT x (reflect $ Proxy @s)
270+
271+-- | Find the smallest element greater than the given one.
272+lookupGT :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a)
273+lookupGT x = gcoerceWith (unsafeCastElement @s @a) $ coerce
274+ $ Set.lookupGT x (reflect $ Proxy @s)
275+
276+-- | Find the largest element smaller or equal to the given one.
277+lookupLE :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a)
278+lookupLE x = gcoerceWith (unsafeCastElement @s @a) $ coerce
279+ $ Set.lookupLE x (reflect $ Proxy @s)
280+
281+-- | Find the smallest element greater or equal to the given one.
282+lookupGE :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a)
283+lookupGE x = gcoerceWith (unsafeCastElement @s @a) $ coerce
284+ $ Set.lookupGE x (reflect $ Proxy @s)
285+
286+-- | If the set is empty, return the proof that it is.
287+null :: forall s a. KnownSet s a => Maybe (EmptyProof 'Regular s)
288+null
289+ | Set.null $ reflect $ Proxy @s = Just $ EmptyProof unsafeSubset
290+ | otherwise = Nothing
291+
292+-- | If @s@ is a subset of @t@ (or is equal to), return a proof of that.
293+isSubsetOf
294+ :: forall s t a. (Ord a, KnownSet s a, KnownSet t a)
295+ => Maybe (SubsetProof 'Regular s t)
296+isSubsetOf
297+ | reflect (Proxy @s) `Set.isSubsetOf` reflect (Proxy @t)
298+ = Just $ SubsetProof unsafeSubset
299+ | otherwise = Nothing
300+
301+-- | If @s@ and @t@ are disjoint (i.e. their intersection is empty), return a
302+-- proof of that.
303+disjoint
304+ :: forall s t a. (Ord a, KnownSet s a, KnownSet t a)
305+ => Maybe (DisjointProof 'Regular s t)
306+disjoint
307+#if MIN_VERSION_containers(0, 5, 11)
308+ | Set.disjoint (reflect $ Proxy @s) (reflect $ Proxy @t)
309+#else
310+ | Set.null $ Set.intersection (reflect $ Proxy @s) (reflect $ Proxy @t)
311+#endif
312+ = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g
313+ | otherwise = Nothing
314+
315+-- | The union of two sets.
316+union
317+ :: forall s t a. (Ord a, KnownSet s a, KnownSet t a)
318+ => SomeSetWith (UnionProof 'Regular s t) a
319+union = reify (reflect (Proxy @s) `Set.union` reflect (Proxy @t))
320+ \(_ :: Proxy r) -> SomeSetWith @r $ UnionProof unsafeSubset unsafeSubsetWith2
321+
322+-- unions :: ?
323+
324+-- | Set with elements of @s@ that are not in @t@.
325+difference
326+ :: forall s t a. (Ord a, KnownSet s a, KnownSet t a)
327+ => SomeSetWith (DifferenceProof 'Regular s t) a
328+difference = reify (reflect (Proxy @s) `Set.difference` reflect (Proxy @t))
329+ \(_ :: Proxy r) -> SomeSetWith @r
330+ $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset
331+
332+-- | Intersection of two sets.
333+intersection
334+ :: forall s t a. (Ord a, KnownSet s a, KnownSet t a)
335+ => SomeSetWith (IntersectionProof 'Regular s t) a
336+intersection = reify (reflect (Proxy @s) `Set.intersection` reflect (Proxy @t))
337+ \(_ :: Proxy r) -> SomeSetWith @r
338+ $ IntersectionProof unsafeSubset unsafeSubsetWith2
339+
340+-- | Cartesian product of two sets. The elements are all pairs @(x, y)@ for each
341+-- @x@ from @s@ and for each @y@ from @t@.
342+cartesianProduct
343+ :: forall s t a b. (KnownSet s a, KnownSet t b)
344+ => SomeSetWith (ProductProof 'Regular s t) (a, b)
345+cartesianProduct = reify
346+#if MIN_VERSION_containers(0, 5, 11)
347+ (reflect (Proxy @s) `Set.cartesianProduct` reflect (Proxy @t))
348+#else
349+ (Set.fromDistinctAscList $ (,) <$> Set.toAscList (reflect $ Proxy @s)
350+ <*> Set.toAscList (reflect $ Proxy @t))
351+#endif
352+ \(_ :: Proxy r) -> SomeSetWith @r $ ProductProof let
353+ proof :: forall x y. Coercion
354+ (Refined (InSet 'Regular s) x, Refined (InSet 'Regular t) y)
355+ (Refined (InSet 'Regular r) (x, y))
356+ !proof
357+ | Coercion <- reallyUnsafeUnderlyingRefined @x @(InSet 'Regular s)
358+ , Coercion <- reallyUnsafeUnderlyingRefined @y @(InSet 'Regular t)
359+ = Coercion `trans`
360+ reallyUnsafeUnderlyingRefined @(x, y) @(InSet 'Regular r)
361+ in proof
362+
363+-- | Disjoint union of two sets. Includes @'Left' x@ for each @x@ from @s@, and
364+-- @'Right' y@ for each @y@ from @t@.
365+disjointUnion
366+ :: forall s t a b. (KnownSet s a, KnownSet t b)
367+ => SomeSetWith (CoproductProof 'Regular s t) (Either a b)
368+disjointUnion = reify
369+#if MIN_VERSION_containers(0, 5, 11)
370+ (reflect (Proxy @s) `Set.disjointUnion` reflect (Proxy @t))
371+#else
372+ (Set.fromDistinctAscList $ (Left <$> Set.toAscList (reflect $ Proxy @s))
373+ ++ (Right <$> Set.toAscList (reflect $ Proxy @t)))
374+#endif
375+ \(_ :: Proxy r) -> SomeSetWith @r $ CoproductProof let
376+ proof :: forall x y. Coercion
377+ (Either (Refined (InSet 'Regular s) x) (Refined (InSet 'Regular t) y))
378+ (Refined (InSet 'Regular r) (Either x y))
379+ !proof
380+ | Coercion <- reallyUnsafeUnderlyingRefined @x @(InSet 'Regular s)
381+ , Coercion <- reallyUnsafeUnderlyingRefined @y @(InSet 'Regular t)
382+ = Coercion `trans`
383+ reallyUnsafeUnderlyingRefined @(Either x y) @(InSet 'Regular r)
384+ in proof
385+
386+-- | Return a subset of elements that satisfy the given predicate.
387+filter
388+ :: forall s a. KnownSet s a
389+ => (Element s a -> Bool) -> SomeSetWith (SupersetProof 'Regular s) a
390+filter p = reify (Set.filter (p . unsafeElement) $ reflect $ Proxy @s)
391+ \(_ :: Proxy r) -> SomeSetWith @r $ SupersetProof unsafeSubset
392+
393+-- | Partition a set into two disjoint subsets: those that satisfy the
394+-- predicate, and those that don't.
395+partition
396+ :: forall s a. KnownSet s a
397+ => (Element s a -> Bool) -> Some2SetWith (PartitionProof 'Regular s a) a
398+partition p = case Set.partition (p . unsafeElement) $ reflect $ Proxy @s of
399+ (r, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q)
400+ -> Some2SetWith @s @r $ PartitionProof
401+ do \x -> if p x
402+ then Left $ unsafeElement $ unrefine x
403+ else Right $ unsafeElement $ unrefine x
404+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
405+
406+-- | Divide a set into two disjoint subsets at a point where the predicate stops
407+-- holding.
408+--
409+-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then
410+-- this point is uniquely defined. If @p@ is not antitone, a splitting point is
411+-- chosen in an unspecified way.
412+spanAntitone
413+ :: forall s a. KnownSet s a
414+ => (Element s a -> Bool) -> Some2SetWith (PartialPartitionProof 'Regular s) a
415+spanAntitone p =
416+#if MIN_VERSION_containers(0, 5, 8)
417+ case Set.spanAntitone (p . unsafeElement) $ reflect $ Proxy @s of
418+ (r, q)
419+#else
420+ case List.span (p . unsafeElement)
421+ $ Set.toAscList $ reflect $ Proxy @s of
422+ (rs, qs)
423+ | let r = Set.fromDistinctAscList rs
424+ , let q = Set.fromDistinctAscList qs
425+#endif
426+ -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q)
427+ -> Some2SetWith @r @q $ PartialPartitionProof
428+ unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g
429+
430+-- | Return two disjoint subsets: those less than the given element, and those
431+-- greater than the given element; along with the proof that the given element
432+-- was in the set, if it was.
433+splitMember
434+ :: forall s a. (Ord a, KnownSet s a)
435+ => a -> Some2SetWith (SplitProof 'Regular s (Element s a)) a
436+splitMember x = case Set.splitMember x $ reflect $ Proxy @s of
437+ (r, m, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q)
438+ -> Some2SetWith @r @q $ SplitProof
439+ (if m then Just (unsafeElement x) else Nothing)
440+ unsafeSubset \f g -> unsafeSubsetWith2 f g
441+
442+-- | Apply the given function to each element of the set and collect the
443+-- results. Note that the resulting set can be smaller.
444+map
445+ :: forall s a b. (Ord b, KnownSet s a)
446+ => (Element s a -> b) -> SomeSetWith (MapProof 'Regular s a b) b
447+map f = reify (Map.keysSet m) \(_ :: Proxy r) -> SomeSetWith @r
448+ $ MapProof (unsafeElement . f) \y -> case Map.lookup (unrefine y) m of
449+ Nothing -> error "map: bug: Data.Set.Refined has been subverted"
450+ Just x -> x
451+ where
452+ !m = Map.fromList
453+ [ (y, unsafeElement x)
454+ | x <- Set.toList $ reflect $ Proxy @s
455+ , let !y = f $ unsafeElement x
456+ ]
457+
458+-- | Map each element of @s@ into a monoid (with proof that it was an element),
459+-- and combine the results using 'Data.Monoid.<>'.
460+foldMap :: forall s a m. (KnownSet s a, Monoid m) => (Element s a -> m) -> m
461+foldMap f = Foldable.foldMap (f . unsafeElement) $ reflect $ Proxy @s
462+
463+-- | Right associative fold with a lazy accumulator.
464+foldr :: forall s a b. KnownSet s a => (Element s a -> b -> b) -> b -> b
465+foldr f z = Set.foldr (f . unsafeElement) z $ reflect $ Proxy @s
466+
467+-- | Left associative fold with a lazy accumulator.
468+foldl :: forall s a b. KnownSet s a => (b -> Element s a -> b) -> b -> b
469+foldl f z = Set.foldl ((. unsafeElement) . f) z $ reflect $ Proxy @s
470+
471+-- | Right associative fold with a strict accumulator.
472+foldr' :: forall s a b. KnownSet s a => (Element s a -> b -> b) -> b -> b
473+foldr' f z = Set.foldr' (f . unsafeElement) z $ reflect $ Proxy @s
474+
475+-- | Left associative fold with a strict accumulator.
476+foldl' :: forall s a b. KnownSet s a => (b -> Element s a -> b) -> b -> b
477+foldl' f z = Set.foldl' ((. unsafeElement) . f) z $ reflect $ Proxy @s
478+
479+-- | Retrieves the smallest element of the set, and the set with that element
480+-- removed; or a proof that the set was empty.
481+minView
482+ :: forall s a. KnownSet s a
483+ => Either
484+ (EmptyProof 'Regular s)
485+ (Element s a, SomeSetWith (SupersetProof 'Regular s) a)
486+minView = case Set.minView $ reflect $ Proxy @s of
487+ Nothing -> Left $ EmptyProof unsafeSubset
488+ Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r)
489+ -> SomeSetWith @r $ SupersetProof unsafeSubset
490+
491+-- | Retrieves the greatest element of the set, and the set with that element
492+-- removed; or a proof that the set was empty.
493+maxView
494+ :: forall s a. KnownSet s a
495+ => Either
496+ (EmptyProof 'Regular s)
497+ (Element s a, SomeSetWith (SupersetProof 'Regular s) a)
498+maxView = case Set.maxView $ reflect $ Proxy @s of
499+ Nothing -> Left $ EmptyProof unsafeSubset
500+ Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r)
501+ -> SomeSetWith @r $ SupersetProof unsafeSubset
502+
503+-- | List of elements in the set in ascending order.
504+toList :: forall s a. KnownSet s a => [Element s a]
505+toList = gcoerceWith (unsafeCastElement @s @a) $ coerce
506+ $ Set.toAscList $ reflect $ Proxy @s
507+
508+-- | List of elements in the set in descending order.
509+toDescList :: forall s a. KnownSet s a => [Element s a]
510+toDescList = gcoerceWith (unsafeCastElement @s @a) $ coerce
511+ $ Set.toDescList $ reflect $ Proxy @s
512+
513+-- | Convert a 'Set' into an 'IntSet', retaining its set of elements, which can
514+-- be converted with 'castFlavor'.
515+asIntSet :: forall s. KnownSet s Int => IntSet s
516+asIntSet = set2IntSet
517+
518+-- | Convert a 'Set' into a 'HashSet', retaining its set of elements, which can
519+-- be converted with 'castFlavor'.
520+asHashSet :: forall s a. (Hashable a, KnownSet s a) => HashSet s a
521+asHashSet = set2HashSet
522+
523+-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then
524+-- @s@ and @t@ actually stand for the same set and @'Element' s@ can be safely
525+-- interconverted with @'Element' t@.
526+--
527+-- The requirement that the weakenings are natural transformations ensures that
528+-- they don't actually alter the elements. To build these you can compose
529+-- ':->''s from proofs returned by functions in this module, or "Refined"
530+-- functions like 'andLeft' or 'leftOr'.
531+castElement
532+ :: forall s t a. (forall x. Element s x -> Element t x)
533+ -> (forall x. Element t x -> Element s x)
534+ -> Coercion (Element s a) (Element t a)
535+castElement = castRefined
536+
537+-- | If elements can be interconverted (e.g. as proved by 'castElement'), then
538+-- the sets can be interconverted too. For example we can establish that the
539+-- intersection of a set with itself is interconvertible with that set:
540+--
541+-- @
542+-- castIntersection
543+-- :: t'IntersectionProof' ''Regular' s s r
544+-- -> 'Coercion' ('Set' r a) ('Set' s a)
545+-- castIntersection ( v'IntersectionProof' p1 p2)
546+-- = 'cast' $ 'castElement' ('andLeft' . p1) (p2 'id' 'id')
547+-- @
548+cast
549+ :: forall s t a. (forall x. Coercion (Element s x) (Element t x))
550+ -> Coercion (Set s a) (Set t a)
551+cast Coercion
552+#if MIN_VERSION_base(4, 15, 0)
553+ = case unsafeEqualityProof @s @t of UnsafeRefl -> Coercion
554+#else
555+ = repr $ unsafeCoerce Refl
556+#endif