From 413b6321b96b4e000ab61687856806c07a41df98 Mon Sep 17 00:00:00 2001 From: mniip Date: Wed, 20 Dec 2023 17:10:58 +0100 Subject: [PATCH] Initial commit --- .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 +++++++++++++++++++++++ 19 files changed, 5881 insertions(+) create mode 100644 .github/workflows/ci.yaml create mode 100644 LICENSE create mode 100644 refined-containers.cabal create mode 100644 src/Data/Container/Refined/Conversion.hs create mode 100644 src/Data/Container/Refined/Hashable.hs create mode 100644 src/Data/Container/Refined/Proofs.hs create mode 100644 src/Data/Container/Refined/Unsafe.hs create mode 100644 src/Data/HashMap/Common/Refined.hs create mode 100644 src/Data/HashMap/Refined.hs create mode 100644 src/Data/HashMap/Strict/Refined.hs create mode 100644 src/Data/HashSet/Refined.hs create mode 100644 src/Data/IntMap/Common/Refined.hs create mode 100644 src/Data/IntMap/Refined.hs create mode 100644 src/Data/IntMap/Strict/Refined.hs create mode 100644 src/Data/IntSet/Refined.hs create mode 100644 src/Data/Map/Common/Refined.hs create mode 100644 src/Data/Map/Refined.hs create mode 100644 src/Data/Map/Strict/Refined.hs create mode 100644 src/Data/Set/Refined.hs diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml new file mode 100644 index 0000000..bac98a6 --- /dev/null +++ b/.github/workflows/ci.yaml @@ -0,0 +1,72 @@ +name: CI + +on: + push: + branches: + - "*" + pull_request: + branches: + - "*" + release: + types: + - prereleased + +jobs: + build: + strategy: + matrix: + platform: + - ubuntu-latest + ghc_version: + #- "8.6.5" + #- "8.8.4" + #- "8.10.7" + #- "9.0.2" + #- "9.2.8" + #- "9.4.8" + - "9.6.3" + #- "9.8.1" + fail-fast: false + name: Build on Linux, GHC ${{ matrix.ghc_version }} + runs-on: ${{ matrix.platform }} + steps: + - uses: actions/checkout@v3 + - uses: haskell/actions/setup@v2 + with: + ghc-version: ${{ matrix.ghc_version }} + cabal-version: "3.10.1.0" + - name: cabal configure + run: cabal configure --with-ghc=ghc-${{ matrix.ghc_version }} + - name: Install dependencies + run: cabal build --dependencies-only --enable-tests --haddock-all + - name: cabal build + run: cabal build --ghc-options=-Werror + - run: cabal check + - run: cabal haddock + + dist: + needs: build + runs-on: ubuntu-latest + name: Build source tarball + if: github.event_name == 'release' + steps: + - uses: actions/checkout@v3 + - uses: haskell/actions/setup@v2 + with: + cabal-version: "3.10.1.0" + - name: cabal sdist + run: | + cabal sdist --output-dir "$GITHUB_WORKSPACE/sdist" + echo "dist=$(find "$GITHUB_WORKSPACE/sdist" -maxdepth 1 -type f -name '*.tar.gz')" >> "$GITHUB_ENV" + - uses: actions/upload-artifact@v2 + with: + name: Source tarball + path: ${{ env.dist }} + - uses: svenstaro/upload-release-action@2.2.1 + with: + repo_token: ${{ secrets.GITHUB_TOKEN }} + file: ${{ env.dist }} + tag: ${{ github.ref }} + - name: Upload tarball to hackage + run: | + curl -X POST 'https://hackage.haskell.org/packages/candidates/' -H "Authorization: X-ApiKey ${{ secrets.HACKAGE_KEY }}" -F "package=@$dist" -w '%{url_effective}\n' diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..72d29ee --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Typeable + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/refined-containers.cabal b/refined-containers.cabal new file mode 100644 index 0000000..9797299 --- /dev/null +++ b/refined-containers.cabal @@ -0,0 +1,99 @@ +cabal-version: 3.0 +name: refined-containers +category: Data +synopsis: + Type-checked proof that a key exists in a container and can be safely + indexed. +description: + This package defines ways to prove that a key exists in an associative + container such as a 'Map', 'IntMap', or 'HashMap'; so that the key can be + used to index into the map without a 'Maybe' or manually handling the + \"impossible\" case with 'error' or other partial functions. + + To do this, the containers are tagged with a type parameter that identifies + their set of keys, so that if you have another container with the same + parameter, you know it has the same keys. + + There is also a type of keys that have been proven to exist in such + containers -- a refinement type. They are also tagged with a type parameter. + If the type parameter of the key matches that of the container, indexing is + guaranteed to proceed without failure. + +license: MIT +license-file: LICENSE +author: mniip@typeable.io +maintainer: mniip@typeable.io +version: 0.1.0.0 +build-type: Simple + +tested-with: + , GHC == 9.8.1 + , GHC == 9.6.3 + , GHC == 9.4.8 + , GHC == 9.2.8 + , GHC == 9.0.2 + , GHC == 8.10.7 + , GHC == 8.6.5 + +source-repository head + type: git + location: https://github.com/typeable/refined-containers/ + +library + build-depends: + , base >= 4.12 && < 4.19 + , adjunctions >= 4.4 && < 4.5 + , constraints >= 0.11 && < 0.15 + , containers >= 0.5.7 && < 0.8 + , deepseq >= 1.4 && < 1.6 + , distributive >= 0.5.3 && < 0.7 + , hashable >= 1.2.7 && < 1.5 + , indexed-traversable >= 0.1 && < 0.2 + , mtl >= 2.2.2 && < 2.4 + , refined >= 0.5 && < 0.9 + , reflection >= 2 && < 2.2 + , unordered-containers >= 0.2.11 && < 0.3 + exposed-modules: + Data.HashMap.Refined + Data.HashMap.Strict.Refined + Data.HashSet.Refined + Data.IntMap.Refined + Data.IntMap.Strict.Refined + Data.IntSet.Refined + Data.Map.Refined + Data.Map.Strict.Refined + Data.Set.Refined + other-modules: + Data.Container.Refined.Conversion + Data.Container.Refined.Hashable + Data.Container.Refined.Proofs + Data.Container.Refined.Unsafe + Data.HashMap.Common.Refined + Data.IntMap.Common.Refined + Data.Map.Common.Refined + hs-source-dirs: src + default-language: Haskell2010 + default-extensions: + BangPatterns + BlockArguments + ConstraintKinds + DataKinds + DeriveTraversable + DerivingStrategies + FlexibleContexts + FlexibleInstances + GADTs + GeneralizedNewtypeDeriving + MagicHash + MultiParamTypeClasses + MultiWayIf + OverloadedStrings + RankNTypes + PatternSynonyms + RoleAnnotations + ScopedTypeVariables + TupleSections + TypeApplications + TypeFamilies + TypeOperators + ghc-options: -Wall -Wredundant-constraints diff --git a/src/Data/Container/Refined/Conversion.hs b/src/Data/Container/Refined/Conversion.hs new file mode 100644 index 0000000..9957d36 --- /dev/null +++ b/src/Data/Container/Refined/Conversion.hs @@ -0,0 +1,79 @@ +module Data.Container.Refined.Conversion where + +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Hashable +import Data.Container.Refined.Proofs +import qualified Data.HashMap.Lazy as HashMap +import Data.HashMap.Common.Refined +import qualified Data.HashSet as HashSet +import qualified Data.IntMap as IntMap +import Data.IntMap.Common.Refined +import qualified Data.IntSet as IntSet +import qualified Data.Map as Map +import Data.Map.Common.Refined +import Data.Proxy +import Data.Reflection +import qualified Data.Set as Set +import Data.Type.Equality ((:~:)(..)) +import Unsafe.Coerce + + +unsafeWrapSet :: forall r r' a. KnownSet r' a => Proxy r' -> Set r a +unsafeWrapSet _ = case unsafeCoerce Refl :: r :~: r' of Refl -> Dict +{-# INLINE unsafeWrapSet #-} + +unsafeWrapIntSet :: forall r r'. KnownIntSet r' => Proxy r' -> IntSet r +unsafeWrapIntSet _ = case unsafeCoerce Refl :: r :~: r' of Refl -> Dict +{-# INLINE unsafeWrapIntSet #-} + +unsafeWrapHashSet :: forall r r' a. KnownHashSet r' a => Proxy r' -> HashSet r a +unsafeWrapHashSet _ = case unsafeCoerce Refl :: r :~: r' of Refl -> Dict +{-# INLINE unsafeWrapHashSet #-} + +set2IntSet :: forall s. KnownSet s Int => IntSet s +set2IntSet = reify + (IntSet.fromDistinctAscList $ Set.toAscList $ reflect $ Proxy @s) + unsafeWrapIntSet + +map2IntMap :: forall s a. Map s Int a -> IntMap s a +map2IntMap (Map m) = IntMap $ IntMap.fromDistinctAscList $ Map.toAscList m + +set2HashSet :: forall s a. (Hashable a, KnownSet s a) => HashSet s a +set2HashSet = reify + (HashSet.fromList $ Set.toList $ reflect $ Proxy @s) + unsafeWrapHashSet + +map2HashMap :: forall s k a. Hashable k => Map s k a -> HashMap s k a +map2HashMap (Map m) = HashMap $ HashMap.fromList $ Map.toList m + +intSet2Set :: forall s. KnownIntSet s => Set s Int +intSet2Set = reify + (Set.fromDistinctAscList $ IntSet.toAscList $ reflect $ Proxy @s) + unsafeWrapSet + +intMap2Map :: forall s a. IntMap s a -> Map s Int a +intMap2Map (IntMap m) = Map $ Map.fromDistinctAscList $ IntMap.toAscList m + +intSet2HashSet :: forall s. KnownIntSet s => HashSet s Int +intSet2HashSet = reify + (HashSet.fromList $ IntSet.toList $ reflect $ Proxy @s) + unsafeWrapHashSet + +intMap2HashMap :: forall s a. IntMap s a -> HashMap s Int a +intMap2HashMap (IntMap m) = HashMap $ HashMap.fromList $ IntMap.toList m + +hashSet2Set :: forall s a. (Ord a, KnownHashSet s a) => Set s a +hashSet2Set = reify + (Set.fromList $ HashSet.toList $ reflect $ Proxy @s) + unsafeWrapSet + +hashMap2Map :: forall s k a. Ord k => HashMap s k a -> Map s k a +hashMap2Map (HashMap m) = Map $ Map.fromList $ HashMap.toList m + +hashSet2IntSet :: forall s. KnownHashSet s Int => IntSet s +hashSet2IntSet = reify + (IntSet.fromList $ HashSet.toList $ reflect $ Proxy @s) + unsafeWrapIntSet + +hashMap2IntMap :: forall s a. HashMap s Int a -> IntMap s a +hashMap2IntMap (HashMap m) = IntMap $ IntMap.fromList $ HashMap.toList m diff --git a/src/Data/Container/Refined/Hashable.hs b/src/Data/Container/Refined/Hashable.hs new file mode 100644 index 0000000..1c14827 --- /dev/null +++ b/src/Data/Container/Refined/Hashable.hs @@ -0,0 +1,13 @@ +{-# LANGUAGE CPP #-} +module Data.Container.Refined.Hashable + ( Hashable + ) where + +#if MIN_VERSION_hashable(1, 4, 0) +import Data.Hashable (Hashable) +#else +import qualified Data.Hashable as Hashable + + +type Hashable a = (Eq a, Hashable.Hashable a) +#endif diff --git a/src/Data/Container/Refined/Proofs.hs b/src/Data/Container/Refined/Proofs.hs new file mode 100644 index 0000000..5594f20 --- /dev/null +++ b/src/Data/Container/Refined/Proofs.hs @@ -0,0 +1,250 @@ +{-# LANGUAGE CPP #-} +#if MIN_VERSION_refined(0, 7, 0) +#else +{-# LANGUAGE UndecidableInstances #-} +#endif +module Data.Container.Refined.Proofs where + +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Hashable +import Data.Kind +import Data.Reflection +import qualified Data.HashSet as HashSet +import qualified Data.IntSet as IntSet +import qualified Data.Set as Set +import Data.Type.Coercion +import Data.Typeable +import Refined + + +-- | A constraint evidencing that we know the contents of the set @s@ at +-- runtime. Whenever you see this constraint on a function, there is an actual +-- @'Data.Set.Set' a@ being passed around at runtime. +-- +-- Given this constraint, to obtain a regular @'Data.Set.Set' a@ you can use +-- 'reflect'. +type KnownSet s a = Reifies s (Set.Set a) + +-- | A 'Data.Set.Set' whose contents are tracked by the type parameter @s@. This +-- is a \"singleton\": for a given @s@ there's only one value of this type. +-- +-- Since this is just a 'Dict', you can freely convert between the value ('Set') +-- and the constraint ('KnownSet'). This library prefers to use the constraint. +type Set s a = Dict (KnownSet s a) + +-- | A constraint evidencing that we know the contents of the set @s@ at +-- runtime. Whenever you see this constraint on a function, there is an actual +-- 'Data.IntSet.IntSet' being passed around at runtime. +-- +-- Given this constraint, to obtain a regular 'Data.IntSet.IntSet' you can use +-- 'reflect'. +type KnownIntSet s = Reifies s IntSet.IntSet + +-- | A 'Data.IntSet.IntSet' whose contents are tracked by the type parameter +-- @s@. This is a \"singleton\": for a given @s@ there's only one value of this +-- type. +-- +-- Since this is just a 'Dict', you can freely convert between the value +-- ('IntSet') and the constraint ('KnownIntSet'). This library prefers to use +-- the constraint. +type IntSet s = Dict (KnownIntSet s) + +-- | A constraint evidencing that we know the contents of the set @s@ at +-- runtime. Whenever you see this constraint on a function, there is an actual +-- @'Data.HashSet.HashSet' a@ being passed around at runtime. +-- +-- Given this constraint, to obtain a regular @'Data.HashSet.HashSet' a@ you can +-- use 'reflect'. +type KnownHashSet s a = Reifies s (HashSet.HashSet a) + +-- | A 'Data.HashSet.HashSet' whose contents are tracked by the type parameter +-- @s@. This is a \"singleton\": for a given @s@ there's only one value of this +-- type. +-- +-- Since this is just a 'Dict', you can freely convert between the value +-- ('HashSet') and the constraint ('KnownHashSet'). This library prefers to use +-- the constraint. +type HashSet s a = Dict (KnownHashSet s a) + +-- | Disambiguate the choice of implementation for sets and maps. +data Flavor + = Regular -- ^ 'Data.Set.Set' and 'Data.Map.Map' + | Int -- ^ 'Data.IntSet.IntSet' and 'Data.IntMap.IntMap' + | Hashed -- ^ 'Data.HashSet.HashSet' and 'Data.HashMap.Lazy.HashMap' + +-- | A predicate for use with "Refined", verifying that a value is an element of +-- the set @s@. +data InSet (f :: Flavor) (s :: Type) = InSet + +-- | See 'Data.Set.Refined.revealPredicate'. +instance (Ord a, Typeable s, KnownSet s a) + => Predicate (InSet 'Regular s) a where + validate p x + | x `Set.member` reflect (Proxy @s) = success + | otherwise = throwRefineOtherException + (typeRep p) + "Value is not in the Set" + +-- | See 'Data.IntSet.Refined.revealPredicate'. +instance (a ~ Int, Typeable s, KnownIntSet s) + => Predicate (InSet 'Int s) a where + validate p x + | x `IntSet.member` reflect (Proxy @s) = success + | otherwise = throwRefineOtherException + (typeRep p) + "Value is not in the IntSet" + +-- | See 'Data.HashSet.Refined.revealPredicate'. +instance (Hashable a, Typeable s, KnownHashSet s a) + => Predicate (InSet 'Hashed s) a where + validate p x + | x `HashSet.member` reflect (Proxy @s) = success + | otherwise = throwRefineOtherException + (typeRep p) + "Value is not in the HashSet" + +-- | A proof that values satisfying @p@ can be cast into values satisfying @q@. +-- +-- For example, @'InSet' s ':->' 'InSet' r@ proves that \(s \subseteq r\). +type p :-> q = forall x. Refined p x -> Refined q x +infix 1 :-> + +-- | Proof that the set @r@ is empty. +newtype EmptyProof f r = EmptyProof + (forall s. InSet f r :-> InSet f s) + -- ^ \(\forall s, r \subseteq s\), which is equivalent to + -- \(r \subseteq \varnothing\) + +-- | Proof that @r@ contains an element of type @a@. +newtype SingletonProof f a r = SingletonProof + (Refined (InSet f r) a) -- ^ The element that is guaranteed to be in @r@ + +-- | Proof that elements of @t a@ are contained in @r@. +newtype FromTraversableProof f (t :: Type -> Type) a r = FromTraversableProof + (t (Refined (InSet f r) a)) + -- ^ The original traversable, with all elements refined with a guarantee of + -- being in @r@ + +-- | Proof that @r@ is @s@ with @a@ inserted. +data InsertProof f a s r = InsertProof + (Refined (InSet f r) a) + -- ^ The element that was inserted and is guaranteed to be in @r@. + (InSet f s :-> InSet f r) -- ^ \(s \subseteq r \) + +-- | Proof that @s@ is a subset of the set @r@. +newtype SubsetProof f s r = SubsetProof + (InSet f s :-> InSet f r) -- ^ \(s \subseteq r\) + +-- | Proof that @s@ is a superset of the set @r@. +newtype SupersetProof f s r = SupersetProof + (InSet f r :-> InSet f s) -- ^ \(r \subseteq s\) + +-- | Proof that @s@ and @r@ are disjoint. +newtype DisjointProof f s r = DisjointProof + (forall t. InSet f t :-> InSet f s + -> InSet f t :-> InSet f r + -> forall u. InSet f t :-> InSet f u) + -- ^ \(\forall t,(t\subseteq s)\land(t\subseteq r)\implies\forall u,t\subseteq u\), + -- which is equivalent to \(s \cap r \subseteq \varnothing\) + +-- | Proof that unioning @s@ and @t@ gives @r@. +data UnionProof f s t r = UnionProof + (InSet f s || InSet f t :-> InSet f r) -- ^ \(s \cup t \subseteq r\) + (forall u. InSet f s :-> InSet f u + -> InSet f t :-> InSet f u + -> InSet f r :-> InSet f u) + -- ^ \(\forall u,(s\subseteq u)\land(t\subseteq u)\implies r\subseteq u\), + -- which is equivalent to \(r \subseteq s \cup u\) + +-- | Proof that if from @s@ you subtract @t@, then you get @r@. +data DifferenceProof f s t r = DifferenceProof + (InSet f r :-> InSet f s) -- ^ \(r \subseteq s\) + (forall u. InSet f u :-> InSet f r + -> InSet f u :-> InSet f t + -> forall v. InSet f u :-> InSet f v) + -- ^ \(\forall u,(u\subseteq r)\land(u\subseteq t)\implies\forall v,u\subseteq v\), + -- which is equivalent to \(r \cap t \subseteq \varnothing\) + (InSet f s :-> InSet f t || InSet f r) -- ^ \(s \subseteq t \cup r\) + +-- | Proof that @r@ is obtained by removing some of @t@'s elements from @s@. +data PartialDifferenceProof f s t r = PartialDifferenceProof + (InSet f r :-> InSet f s) -- ^ \(r \subseteq s\) + (InSet f s :-> InSet f t || InSet f r) -- ^ \(s \subseteq t \cup r\) + +-- | Proof that intersecting @s@ and @t@ gives @r@. +data IntersectionProof f s t r = IntersectionProof + (InSet f r :-> InSet f s && InSet f t) -- ^ \(r \subseteq s \cap t\) + (forall u. InSet f u :-> InSet f s + -> InSet f u :-> InSet f t + -> InSet f u :-> InSet f r) + -- ^ \(\forall u,(u\subseteq s)\land(u\subseteq t)\implies u\subseteq r\), + -- which is equivalent to \(s \cap t \subseteq r\) + +-- | Proof that the cartesian product of @s@ and @t@ is @r@. +newtype ProductProof f s t r = ProductProof + (forall a b. Coercion + (Refined (InSet f s) a, Refined (InSet f t) b) + (Refined (InSet f r) (a, b))) + -- ^ A pair of elements from @s@ and @t@ respectively can be converted into an + -- element of @r@ and back. You can use @'coerceWith' co@ and + -- @'coerceWith' ('sym' co)@. + +-- | Proof that the tagged disjoint union of @s@ and @t@ is @r@. +newtype CoproductProof f s t r = CoproductProof + (forall a b. Coercion + (Either (Refined (InSet f s) a) (Refined (InSet f t) b)) + (Refined (InSet f r) (Either a b))) + -- ^ Coproduct of elements of @s@ and @t@ can be converted into an element of + -- @r@ and back. You can use @'coerceWith' co@ and @'coerceWith' ('sym' co)@. + +-- | Proof that @s@ is the union of disjoint subsets @r@ and @q@, together with +-- a procedure that decides which of the two an element belongs to. +data PartitionProof f s a r q = PartitionProof + (Refined (InSet f s) a + -> Either (Refined (InSet f r) a) (Refined (InSet f q) a)) + -- ^ Decide whether a given element of @s@ belongs to @r@ or to @q@ + (InSet f r || InSet f q :-> InSet f s) -- ^ \(r \cup q \subseteq s\) + (forall t. InSet f r :-> InSet f t + -> InSet f q :-> InSet f t + -> InSet f s :-> InSet f t) + -- ^ \(\forall t,(r\subseteq t)\land(q\subseteq t)\implies s\subseteq t\), + -- which is equivalent to \(s \subseteq r \cup q\) + (forall t. InSet f t :-> InSet f r + -> InSet f t :-> InSet f q + -> forall u. InSet f t :-> InSet f u) + -- ^ \(\forall t,(t\subseteq r)\land(t\subseteq q)\implies\forall u,t\subseteq u\), + -- which is equivalent to \(r \cap q \subseteq \varnothing\) + +-- | Proof that @s@ is the union of disjoint subsets @r@ and @q@, but without a +-- deciding procedure. +data PartialPartitionProof f s r q = PartialPartitionProof + (InSet f r || InSet f q :-> InSet f s) -- ^ \(r \cup q \subseteq s\) + (forall t. InSet f r :-> InSet f t + -> InSet f q :-> InSet f t + -> InSet f s :-> InSet f t) + -- ^ \(\forall t,(r\subseteq t)\land(q\subseteq t)\implies s\subseteq t\), + -- which is equivalent to \(s \subseteq r \cup q\) + (forall t. InSet f t :-> InSet f r + -> InSet f t :-> InSet f q + -> forall u. InSet f t :-> InSet f u) + -- ^ \(\forall t,(t\subseteq r)\land(t\subseteq q)\implies\forall u,t\subseteq u\), + -- which is equivalent to \(r \cap q \subseteq \varnothing\) + +-- | Proof that @s@ contains disjoint subsets @r@ and @q@, along with an +-- optional element between them. +data SplitProof f s e r q = SplitProof + !(Maybe e) -- ^ The element between @r@ and @q@ + (InSet f r || InSet f q :-> InSet f s) -- ^ \(r \cup q \subseteq s\) + (forall t. InSet f t :-> InSet f r + -> InSet f t :-> InSet f q + -> forall u. InSet f t :-> InSet f u) + -- ^ \(\forall t,(t\subseteq r)\land(t\subseteq q)\implies\forall u,t\subseteq u\), + -- which is equivalent to \(r \cap q \subseteq \varnothing\) + +-- | Proof that @r@ is the direct image of @s@ under some mapping @f :: a -> b@. +data MapProof f s a b r = MapProof + (Refined (InSet f s) a -> Refined (InSet f r) b) + -- ^ Compute the image of an element of @s@ (the image is then an element of + -- @r@) + (Refined (InSet f r) b -> Refined (InSet f s) a) + -- ^ For an element of @r@, return an arbitrary preimage from @s@ diff --git a/src/Data/Container/Refined/Unsafe.hs b/src/Data/Container/Refined/Unsafe.hs new file mode 100644 index 0000000..3f1e847 --- /dev/null +++ b/src/Data/Container/Refined/Unsafe.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE ImpredicativeTypes #-} +module Data.Container.Refined.Unsafe where + +import Data.Container.Refined.Proofs +import Data.Type.Coercion +import Refined +import Refined.Unsafe + + +unsafeSubset :: p :-> q +unsafeSubset = reallyUnsafeRefine . unrefine +{-# INLINE unsafeSubset #-} + +unsafeSubsetWith2 :: p' :-> q' -> p'' :-> q'' -> p :-> q +unsafeSubsetWith2 f g = reallyUnsafeRefine . unrefine + . f . reallyUnsafeRefine . unrefine + . g . reallyUnsafeRefine . unrefine +{-# INLINE unsafeSubsetWith2 #-} + +-- Because `Refined p x` is a newtype over `x`, by parametricity a `p :-> q` +-- can either diverge or be `id`. This ensures that it does not diverge. +rnfProof :: (p :-> q) -> () +rnfProof f = unrefine $ f $ reallyUnsafeRefine () +{-# INLINE rnfProof #-} + +-- | This function can be used to freely convert between @Element@ and @Key@ +-- types of various flavors ('Regular', v'Int', 'Hashed'), corresponding to the +-- different implementations of sets and maps. +castFlavor + :: forall (f :: Flavor) (g :: Flavor) s a. Coercion + (Refined (InSet f s) a) + (Refined (InSet g s) a) +castFlavor = sym (reallyUnsafeUnderlyingRefined @a @(InSet f s)) + `trans` reallyUnsafeUnderlyingRefined @a @(InSet g s) +{-# INLINE castFlavor #-} + +castRefined + :: forall a p q. (p :-> q) + -> (q :-> p) + -> Coercion (Refined p a) (Refined q a) +castRefined f g + | () <- rnfProof f + , () <- rnfProof g + = sym (reallyUnsafeUnderlyingRefined @a @p) + `trans` reallyUnsafeUnderlyingRefined @a @q +{-# INLINE castRefined #-} diff --git a/src/Data/HashMap/Common/Refined.hs b/src/Data/HashMap/Common/Refined.hs new file mode 100644 index 0000000..0afef4d --- /dev/null +++ b/src/Data/HashMap/Common/Refined.hs @@ -0,0 +1,426 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.HashMap.Common.Refined where + +import Control.Monad.Reader +import Control.DeepSeq +import Data.Coerce +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Hashable +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Distributive +import Data.Foldable.WithIndex +import Data.Functor.Rep +import Data.Functor.WithIndex +import qualified Data.Hashable as Hashable +import qualified Data.HashMap.Lazy as HashMap +import qualified Data.HashMap.Strict as HashMapStrict +import qualified Data.HashSet as HashSet +import Data.Proxy +import Data.Reflection +import Data.Traversable.WithIndex +import Data.Type.Coercion +import Data.Type.Equality ((:~:)(..)) +import Refined +import Refined.Unsafe +import Unsafe.Coerce + +#if MIN_VERSION_unordered_containers(0, 2, 12) +#else +import Data.Monoid (All(..)) +#endif + + +-- | A wrapper around a regular 'Data.HashMap.HashMap' with a type parameter @s@ +-- identifying the set of keys present in the map. +-- +-- A key of type @k@ may not be present in the map, but a @'Key' s k@ is +-- guaranteed to be present (if the @s@ parameters match). Thus the map is +-- isomorphic to a (total) function @'Key' s k -> a@, which motivates many of +-- the instances below. +-- +-- A 'HashMap' always knows its set of keys, so given @'HashMap' s k a@ we can +-- always derive @'KnownHashSet' s k@ by pattern matching on the 'Dict' returned +-- by 'keysSet'. +newtype HashMap s k a = HashMap (HashMap.HashMap k a) + deriving newtype (Eq, Ord, Show, Functor, Foldable, Hashable.Hashable, NFData) + deriving stock (Traversable) +type role HashMap nominal nominal representational + +-- | Convert to a regular 'Data.HashMap.HashMap', forgetting its set of keys. +toMap :: forall s k a. HashMap s k a -> HashMap.HashMap k a +toMap (HashMap m) = m + +-- | @'Key' s k@ is a key of type @k@ that has been verified to be an element +-- of the set @s@, and thus verified to be present in all @'HashMap' s k@ maps. +-- +-- Thus, @'Key' s k@ is a \"refinement\" type of @k@, and this library +-- integrates with an implementation of refimenement types in "Refined", so +-- the machinery from there can be used to manipulate 'Key's (however see +-- 'Data.Set.Refined.revealPredicate'). +-- +-- The underlying @k@ value can be obtained with 'unrefine'. A @k@ can be +-- validated into an @'Key' s k@ with 'member'. +type Key s = Refined (InSet 'Hashed s) + +unsafeCastKey :: forall s k. Coercion k (Key s k) +unsafeCastKey = reallyUnsafeUnderlyingRefined + +unsafeKey :: k -> Key s k +unsafeKey = coerceWith unsafeCastKey + +-- | An existential wrapper for a 'HashMap' with an as-yet-unknown set of keys. +-- Pattern maching on it gives you a way to refer to the set (the parameter +-- @s@), e.g. +-- +-- @ +-- case 'fromHashMap' ... of +-- 'SomeHashMap' \@s m -> doSomethingWith \@s +-- +-- case 'fromHashMap' ... of +-- 'SomeHashMap' (m :: 'HashMap' s k a) -> doSomethingWith \@s +-- @ +data SomeHashMap k a where + SomeHashMap :: forall s k a. !(HashMap s k a) -> SomeHashMap k a + +-- | Apply a map with an unknown set of keys to a continuation that can accept +-- a map with any set of keys. This gives you a way to refer to the set (the +-- parameter @s@), e.g.: +-- +-- @ +-- 'withHashMap' ('fromHashMap' ... +-- $ \(m :: 'HashMap' s k a) -> doSomethingWith \@s +-- @ +withHashMap + :: forall k a r. SomeHashMap k a -> (forall s. HashMap s k a -> r) -> r +withHashMap (SomeHashMap m) k = k m + +-- | Construct a map from a regular 'Data.HashMap.Lazy.HashMap'. +fromHashMap :: forall k a. HashMap.HashMap k a -> SomeHashMap k a +fromHashMap m = SomeHashMap (HashMap m) + +-- | An existential wrapper for a 'HashMap' with an as-yet-unknown set of keys, +-- together with a proof of some fact @p@ about the set. Pattern matching on it +-- gives you a way to refer to the set (the parameter @s@). Functions that +-- change the set of keys in a map will return the map in this way, together +-- with a proof that somehow relates the keys set to the function's inputs. +data SomeHashMapWith p k a where + SomeHashMapWith + :: forall s k a p. !(HashMap s k a) -> !(p s) -> SomeHashMapWith p k a + +-- | Apply a map with proof for an unknown set of keys to a continuation that +-- can accept a map with any set of keys satisfying the proof. This gives you a +-- way to refer to the set (the parameter @s@). +withHashMapWith + :: forall k a r p. SomeHashMapWith p k a + -> (forall s. HashMap s k a -> p s -> r) + -> r +withHashMapWith (SomeHashMapWith m p) k = k m p + +-- | An existential wrapper for a pair of maps with as-yet-unknown sets of keys, +-- together with a proof of some fact @p@ relating them. +data Some2HashMapWith p k a b where + Some2HashMapWith + :: forall s t k a b p. !(HashMap s k a) + -> !(HashMap t k b) + -> !(p s t) + -> Some2HashMapWith p k a b + +-- | Apply a pair of maps with proof for unknown sets of keys to a continuation +-- that can accept any pair of maps with any sets of keys satisfying the proof. +-- This gives you a way to refer to the sets (the parameters @s@ and @t@). +with2HashMapWith + :: forall k a b r p. Some2HashMapWith p k a b + -> (forall s t. HashMap s k a -> HashMap t k b -> p s t -> r) + -> r +with2HashMapWith (Some2HashMapWith m1 m2 p) k = k m1 m2 p + +-- | An empty map. +empty :: forall k a. SomeHashMapWith (EmptyProof 'Hashed) k a +empty = SomeHashMapWith (HashMap HashMap.empty) $ EmptyProof unsafeSubset + +-- | Create a map from a set of keys, and a function that for each key computes +-- the corresponding value. +fromSet :: forall s k a. KnownHashSet s k => (Key s k -> a) -> HashMap s k a +fromSet f = HashMap $ HashMap.mapWithKey (\k _ -> f $ unsafeKey k) + $ HashSet.toMap (reflect $ Proxy @s) + +-- | Delete a key and its value from the map if present, returning a potentially +-- smaller map. +delete + :: forall s k a. Hashable k + => k -> HashMap s k a -> SomeHashMapWith (SupersetProof 'Hashed s) k a +delete k (HashMap m) = SomeHashMapWith (HashMap $ HashMap.delete k m) + $ SupersetProof unsafeSubset + +-- | If the key is in the map, return the proof of this, and the associated +-- value; otherwise return 'Nothing'. +lookup :: forall s k a. Hashable k => k -> HashMap s k a -> Maybe (Key s k, a) +lookup k (HashMap m) = (unsafeKey k,) <$> HashMap.lookup k m + +-- | Given a key that is proven to be in the map, return the associated value. +-- +-- Unlike 'Data.HashMap.!' from "Data.HashMap.Lazy", this function is total, as +-- it is impossible to obtain a @'Key' s k@ for a key that is not in the map +-- @'HashMap' s k a@. +(!) :: forall s k a. Hashable k => HashMap s k a -> Key s k -> a +(!) (HashMap m) k = case HashMap.lookup (unrefine k) m of + Nothing -> error "(!): bug: Data.HashMap.Refined has been subverted" + Just x -> x + +-- | If a key is in the map, return the proof that it is. +member :: forall s k a. Hashable k => k -> HashMap s k a -> Maybe (Key s k) +member k (HashMap m) + | k `HashMap.member` m = Just (unsafeKey k) + | otherwise = Nothing + +-- | If a map is empty, return a proof that it is. +null :: forall s k a. HashMap s k a -> Maybe (EmptyProof 'Hashed s) +null (HashMap m) + | HashMap.null m = Just $ EmptyProof unsafeSubset + | otherwise = Nothing + +-- | If all keys of the first map are also present in the second map, and the +-- given function returns 'True' for their associated values, return a proof +-- that the keys form a subset. +isSubmapOfBy + :: forall s t k a b. Hashable k + => (a -> b -> Bool) + -> HashMap s k a + -> HashMap t k b + -> Maybe (SubsetProof 'Hashed s t) +isSubmapOfBy f (HashMap m1) (HashMap m2) +#if MIN_VERSION_unordered_containers(0, 2, 12) + | HashMap.isSubmapOfBy f m1 m2 +#else + | All True <- flip HashMap.foldMapWithKey m1 + \k v1 -> case HashMap.lookup k m2 of + Just v2 | f v1 v2 -> mempty + _ -> All False +#endif + = Just $ SubsetProof unsafeSubset + | otherwise = Nothing + +-- | If two maps are disjoint (i.e. their intersection is empty), return a proof +-- of that. +disjoint + :: forall s t k a b. Hashable k + => HashMap s k a -> HashMap t k b -> Maybe (DisjointProof 'Hashed s t) +disjoint (HashMap m1) (HashMap m2) + | HashMap.null $ HashMapStrict.intersectionWith (\_ _ -> ()) m1 m2 + = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g + | otherwise = Nothing + +-- | Given two maps proven to have the same keys, for each key apply the +-- function to the associated values, to obtain a new map with the same keys. +zipWithKey + :: forall s k a b c. Hashable k + => (Key s k -> a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c +zipWithKey f (HashMap m1) (HashMap m2) = HashMap + $ HashMap.intersectionWithKey (f . unsafeKey) m1 m2 + +-- | Remove the keys that appear in the second map from the first map. +difference + :: forall s t k a b. Hashable k + => HashMap s k a + -> HashMap t k b + -> SomeHashMapWith (DifferenceProof 'Hashed s t) k a +difference (HashMap m1) (HashMap m2) + = SomeHashMapWith (HashMap $ HashMap.difference m1 m2) + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, that are proven to be in the map. The set of keys remains the same. +mapWithKey + :: forall s k a b. (Key s k -> a -> b) -> HashMap s k a -> HashMap s k b +mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.mapWithKey @k @a @b + +-- | Map an 'Applicative' transformation with access to each value's +-- corresponding key and a proof that it is in the map. The set of keys remains +-- unchanged. +traverseWithKey + :: forall s f k a b. Applicative f + => (Key s k -> a -> f b) -> HashMap s k a -> f (HashMap s k b) +traverseWithKey f (HashMap m) + = HashMap <$> HashMap.traverseWithKey (f . unsafeKey) m + +-- | Map each key-value pair of a map into a monoid (with proof that the key was +-- in the map), and combine the results using '<>'. +foldMapWithKey + :: forall s k a m. Monoid m => (Key s k -> a -> m) -> HashMap s k a -> m +foldMapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.foldMapWithKey @m @k @a + +-- | Right associative fold with a lazy accumulator. +foldrWithKey + :: forall s k a b. (Key s k -> a -> b -> b) -> b -> HashMap s k a -> b +foldrWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.foldrWithKey @k @a @b + +-- | Left associative fold with a lazy accumulator. +foldlWithKey + :: forall s k a b. (b -> Key s k -> a -> b) -> b -> HashMap s k a -> b +foldlWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.foldlWithKey @b @k @a + +-- | Right associative fold with a strict accumulator. +foldrWithKey' + :: forall s k a b. (Key s k -> a -> b -> b) -> b -> HashMap s k a -> b +foldrWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.foldrWithKey' @k @a @b + +-- | Left associative fold with a strict accumulator. +foldlWithKey' + :: forall s k a b. (b -> Key s k -> a -> b) -> b -> HashMap s k a -> b +foldlWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.foldlWithKey' @b @k @a + +-- | Return the set of keys in the map, with the contents of the set still +-- tracked by the @s@ parameter. See "Data.HashSet.Refined". +keysSet :: forall s k a. HashMap s k a -> HashSet s k +keysSet (HashMap m) = reify (HashMap.keysSet m) + \(_ :: Proxy s') -> case unsafeCoerce Refl :: s :~: s' of + Refl -> Dict + +-- | Convert to a list of key-value pairs. +toList :: forall s k a. HashMap s k a -> [(Key s k, a)] +toList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.toList @k @a + +-- | Retain only the key-value pairs that satisfy the predicate, returning a +-- potentially smaller map. +filterWithKey + :: forall s k a. (Key s k -> a -> Bool) + -> HashMap s k a + -> SomeHashMapWith (SupersetProof 'Hashed s) k a +filterWithKey p (HashMap m) + = SomeHashMapWith (HashMap $ HashMap.filterWithKey (p . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Restrict a map to only those keys that are elements of @t@. +restrictKeys + :: forall s t k a. (Hashable k, KnownHashSet t k) + => HashMap s k a -> SomeHashMapWith (IntersectionProof 'Hashed s t) k a +restrictKeys (HashMap m) = SomeHashMapWith + (HashMap $ HashMap.intersectionWith const m + $ HashSet.toMap $ reflect $ Proxy @t) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Remove all keys that are elements of @t@ from the map. +withoutKeys + :: forall s t k a. (Hashable k, KnownHashSet t k) + => HashMap s k a -> SomeHashMapWith (DifferenceProof 'Hashed s t) k a +withoutKeys (HashMap m) = SomeHashMapWith + (HashMap $ HashMap.difference m $ HashSet.toMap $ reflect $ Proxy @t) + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Partition a map into two disjoint submaps: those whose key-value pairs +-- satisfy the predicate, and those whose don't. +partitionWithKey + :: forall s k a. Hashable k -- TODO: this is only used in the proof + => (Key s k -> a -> Bool) + -> HashMap s k a + -> Some2HashMapWith (PartitionProof 'Hashed s k) k a a +partitionWithKey p (HashMap m) = Some2HashMapWith + (HashMap $ HashMap.filterWithKey (p . unsafeKey) m) + (HashMap $ HashMap.filterWithKey ((not .) . p . unsafeKey) m) + $ PartitionProof + do \k -> case HashMap.lookup (unrefine k) m of + Nothing -> error + "partitionWithKey: bug: Data.HashMap.Refined has been subverted" + Just x -> if p k x + then Left $ unsafeKey $ unrefine k + else Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then +-- @s@ and @t@ actually stand for the same set and @'Key' s@ can be safely +-- interconverted with @'Key' t@. +-- +-- The requirement that the weakenings are natural transformations ensures that +-- they don't actually alter the keys. To build these you can compose ':->''s +-- from proofs returned by functions in this module, or "Refined" functions like +-- 'andLeft' or 'leftOr'. +castKey + :: forall s t k. (forall x. Key s x -> Key t x) + -> (forall x. Key t x -> Key s x) + -> Coercion (Key s k) (Key t k) +castKey = castRefined + +-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps +-- can be interconverted too. For example, 'zipWithKey' can be implemented via +-- 'Data.HashMap.Refined.intersectionWithKey' by proving that the set of keys +-- remains unchanged: +-- +-- @ +-- 'zipWithKey' +-- :: forall s k a b c. 'Hashable' k +-- => ('Key' s k -> a -> b -> c) -> 'HashMap' s k a -> 'HashMap' s k b -> 'HashMap' s k c +-- 'zipWithKey' f m1 m2 +-- | v'SomeHashMapWith' @r m proof <- 'Data.HashMap.Refined.intersectionWithKey' (f . 'andLeft') m1 m2 +-- , v'IntersectionProof' p1 p2 <- proof +-- , ( v'Coercion' :: t'Coercion' ('HashMap' r k c) ('HashMap' s k c)) +-- <- app $ 'cast' $ 'castKey' ('andLeft' . p1) (p2 'id' 'id') +-- = 'coerce' m +-- where +-- app :: t'Coercion' f g -> t'Coercion' (f x) (g x) +-- app v'Coercion' = v'Coercion' +-- @ +cast + :: forall s t k. (forall x. Coercion (Key s x) (Key t x)) + -> Coercion (HashMap s k) (HashMap t k) +cast Coercion = Coercion + +instance FunctorWithIndex (Key s k) (HashMap s k) where + imap = mapWithKey + +instance FoldableWithIndex (Key s k) (HashMap s k) where + ifoldMap = foldMapWithKey + +instance TraversableWithIndex (Key s k) (HashMap s k) where + itraverse = traverseWithKey + +-- | Similar to the instance for functions -- zip corresponding keys. To use +-- '<*>'/'liftA2' without 'KnownSet' see 'zipWithKey'. +instance (Hashable k, KnownHashSet s k) => Applicative (HashMap s k) where + pure x = fromSet \_ -> x + (<*>) = zipWithKey (const id) + +-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the +-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions. +bind + :: forall s k a b. Hashable k + => HashMap s k a -> (a -> HashMap s k b) -> HashMap s k b +bind m f = mapWithKey (\k x -> f x ! k) m + +-- | Similar to the instance for functions. To use '>>=' without 'KnownSet' see +-- 'bind'. +instance (Hashable k, KnownHashSet s k) => Monad (HashMap s k) where + (>>=) = bind + +-- | Similar to the instance for functions. See also +-- 'Data.HashMap.Refined.backpermuteKeys'. +instance (Hashable k, KnownHashSet s k) + => MonadReader (Key s k) (HashMap s k) where + ask = fromSet id + local f m = mapWithKey (\k _ -> m ! f k) m + +-- | Append the values at the corresponding keys +instance (Hashable k, Semigroup a) => Semigroup (HashMap s k a) where + (<>) = zipWithKey (const (<>)) + +instance (Hashable k, KnownHashSet s k, Monoid a) + => Monoid (HashMap s k a) where + mempty = fromSet \_ -> mempty + +-- | Similar to the instance for functions +instance (Hashable k, KnownHashSet s k) => Distributive (HashMap s k) where + collect = collectRep + distribute = distributeRep + +-- | Witness isomorphism with functions from @'Key' s k@ +instance (Hashable k, KnownHashSet s k) => Representable (HashMap s k) where + type Rep (HashMap s k) = Key s k + index = (!) + tabulate = fromSet diff --git a/src/Data/HashMap/Refined.hs b/src/Data/HashMap/Refined.hs new file mode 100644 index 0000000..552606f --- /dev/null +++ b/src/Data/HashMap/Refined.hs @@ -0,0 +1,359 @@ +-- | This module defines a way to prove that a key exists in a map, so that the +-- key can be used to index into the map without using a 'Maybe', or manually +-- handling the \"impossible\" case with 'error' or other partial functions. +-- +-- To do this, @'HashMap' s k v@ has a type parameter @s@ that identifies its +-- set of keys, so that if another map has the same type parameter, you know +-- that map has the same set of keys. There is @'Key' s k@, a type of keys that +-- have been validated to belong to the set identified by @s@, and for which the +-- operation of indexing into a @'HashMap' s k v@ (only for the same @s@) can +-- proceed without failure (see '!'). The type @s@ itself has no internal +-- structure, rather it is merely a skolem type variable (rank-2 polymorphism +-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection". +-- +-- Like "Data.HashMap.Lazy", functions in this module are strict in the keys but +-- lazy in the values. The "Data.HashMap.Strict.Refined" module reuses the same +-- 'HashMap' type but provides functions that operate strictly on the values. +-- +-- = Warning +-- This module together with "Data.HashMap.Lazy" rely on 'Eq' and 'Hashable' +-- instances being lawful: that '==' is an equivalence relation, and that +-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence +-- relation; at least for the subset of values that are actually encountered at +-- runtime. If this assumption is violated, this module may not be able to +-- uphold its invariants and may throw errors. In particular beware of NaN in +-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@. +module Data.HashMap.Refined + ( + -- * Map type + Common.HashMap + , Common.Key + -- * Existentials and common proofs + , Common.SomeHashMap(..) + , Common.withHashMap + , Common.SomeHashMapWith(..) + , Common.withHashMapWith + , Common.Some2HashMapWith(..) + , Common.with2HashMapWith + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , Common.empty + , singleton + , SingletonProof(..) + , fromSet + , Common.fromHashMap + , fromTraversableWithKey + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + , reinsert + , insertLookupWithKey + -- * Deletion/Update + , Common.delete + , adjust + , adjustWithKey + , update + , updateLookupWithKey + -- * Query + , Common.lookup + , (Common.!) + , Common.member + , Common.null + , Common.isSubmapOfBy + , SubsetProof(..) + , Common.disjoint + , DisjointProof(..) + -- * Combine + , zipWithKey + , bind + , unionWithKey + , UnionProof(..) + , Common.difference + , DifferenceProof(..) + , differenceWithKey + , PartialDifferenceProof(..) + , intersectionWithKey + , IntersectionProof(..) + -- * Traversal + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , MapProof(..) + , backpermuteKeys + -- * Folds + , Common.foldMapWithKey + , Common.foldrWithKey + , Common.foldlWithKey + , Common.foldrWithKey' + , Common.foldlWithKey' + -- * Conversion + , Common.toMap + , Common.keysSet + , Common.toList + -- * Filter + , Common.restrictKeys + , Common.withoutKeys + , Common.filterWithKey + , Common.partitionWithKey + , PartitionProof(..) + , mapMaybeWithKey + , mapEitherWithKey + -- * Casts + , Common.castKey + , Common.cast + , castFlavor + ) where + +import Data.Coerce +import Data.Container.Refined.Hashable +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Functor +import Data.HashMap.Common.Refined + ( HashMap(..), Key, unsafeCastKey, unsafeKey, SomeHashMapWith(..) + , Some2HashMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey + , bind + ) +import qualified Data.HashMap.Common.Refined as Common +import qualified Data.HashMap.Lazy as HashMap +import Data.Traversable +import Data.Traversable.WithIndex +import Data.Type.Coercion +import Prelude hiding (lookup, null) +import Refined +import Refined.Unsafe + + +-- | Create a map with a single key-value pair, and return a proof that the key +-- is in the resulting map. +singleton + :: forall k a. Hashable k + => k -> a -> SomeHashMapWith (SingletonProof 'Hashed k) k a +singleton k v = SomeHashMapWith (HashMap $ HashMap.singleton k v) + $ SingletonProof (unsafeKey k) + +-- | Create a map from an arbitrary traversable of key-value pairs. +fromTraversableWithKey + :: forall t k a. (Traversable t, Hashable k) + => (k -> a -> a -> a) + -> t (k, a) + -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a +fromTraversableWithKey f xs + = SomeHashMapWith (HashMap m) $ FromTraversableProof proof + where + (m, proof) = mapAccumL + (\s (k, v) + -> let !s' = HashMap.insertWith (f k) k v s in (s', unsafeKey k)) + HashMap.empty + xs + +-- | Insert a key-value pair into the map to obtain a potentially larger map, +-- guaranteed to contain the given key. If the key was already present, the +-- associated value is replaced with the supplied value. +insert + :: forall s k a. Hashable k + => k -> a -> HashMap s k a -> SomeHashMapWith (InsertProof 'Hashed k s) k a +insert k v (HashMap m) = SomeHashMapWith (HashMap $ HashMap.insert k v m) + $ InsertProof (unsafeKey k) unsafeSubset + +-- | Overwrite a key-value pair that is known to already be in the map. The set +-- of keys remains the same. +reinsert + :: forall s k a. Hashable k + => Key s k -> a -> HashMap s k a -> HashMap s k a +reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.insert @k @a + +-- | Insert a key-value pair into the map using a combining function, and if +-- the key was already present, the old value is returned along with the proof +-- that the key was present. +insertLookupWithKey + :: forall s k a. Hashable k + => (Key s k -> a -> a -> a) + -> k + -> a + -> HashMap s k a + -> (Maybe (Key s k, a), SomeHashMapWith (InsertProof 'Hashed k s) k a) +insertLookupWithKey f k v (HashMap m) = + ( (unsafeKey k,) <$> HashMap.lookup k m + , SomeHashMapWith (HashMap $ HashMap.insertWith (f $ unsafeKey k) k v m) + $ InsertProof (unsafeKey k) unsafeSubset + ) + +-- | Update the value at a specific key known the be in the map using the given +-- function. The set of keys remains the same. +adjust + :: forall s k a. Hashable k + => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a +adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.adjust @k @a + +-- | If the given key is in the map, update the associated value using the given +-- function with a proof that the key was in the map; otherwise return the map +-- unchanged. In any case the set of keys remains the same. +adjustWithKey + :: forall s k a. Hashable k + => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a +adjustWithKey f k (HashMap m) = HashMap $ HashMap.adjust (f $ unsafeKey k) k m + +-- | Update or delete a key known to be in the map using the given function, +-- returning a potentially smaller map. +update + :: forall s k a. Hashable k + => (a -> Maybe a) + -> Key s k + -> HashMap s k a + -> SomeHashMapWith (SupersetProof 'Hashed s) k a +update f k (HashMap m) + = SomeHashMapWith (HashMap $ HashMap.update f (unrefine k) m) + $ SupersetProof unsafeSubset + +-- | If the given key is in the map, update or delete it using the given +-- function with a proof that the key was in the map; otherwise the map is +-- unchanged. Alongside return the new value if it was updated, or the old value +-- if it was deleted, and a proof that the key was in the map. +updateLookupWithKey + :: forall s k a. Hashable k + => (Key s k -> a -> Maybe a) + -> k + -> HashMap s k a + -> (Maybe (Key s k, a), SomeHashMapWith (SupersetProof 'Hashed s) k a) +updateLookupWithKey f k (HashMap m) = + ( (unsafeKey k,) <$> HashMap.lookup k m + , SomeHashMapWith (HashMap $ HashMap.update (f $ unsafeKey k) k m) + $ SupersetProof unsafeSubset + ) + +-- | Return the union of two maps, with a given combining function for keys that +-- exist in both maps simultaneously. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +unionWithKey + :: forall s t k a. Hashable k + => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a) + -> HashMap s k a + -> HashMap t k a + -> SomeHashMapWith (UnionProof 'Hashed s t) k a +unionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith + (HashMap $ HashMap.unionWithKey (f . reallyUnsafeRefine) m1 m2) + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- | For keys that appear in both maps, the given function decides whether the +-- key is removed from the first map. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +differenceWithKey + :: forall s t k a b. Hashable k + => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a) + -> HashMap s k a + -> HashMap t k b + -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a +differenceWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith + (HashMap $ HashMap.differenceWith + (\x (k, y) -> f (reallyUnsafeRefine k) x y) + m1 + (HashMap.mapWithKey (,) m2)) + $ PartialDifferenceProof unsafeSubset unsafeSubset + +-- | Return the intersection of two maps with the given combining function. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +intersectionWithKey + :: forall s t k a b c. Hashable k + => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c) + -> HashMap s k a + -> HashMap t k b + -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c +intersectionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith + (HashMap $ HashMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Thread an accumularing argument through the map in ascending order of +-- hashes. +mapAccumLWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> HashMap s k b + -> (a, HashMap s k c) +mapAccumLWithKey f = imapAccumL (flip f) + +-- | Thread an accumularing argument through the map in descending order of +-- hashes. +mapAccumRWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> HashMap s k b + -> (a, HashMap s k c) +mapAccumRWithKey f = imapAccumR (flip f) + +-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the +-- results into a new map. For keys that were mapped to the same new key, @c@ +-- acts as the combining function for corresponding values. +mapKeysWith + :: forall s k1 k2 a. Hashable k2 + => (a -> a -> a) + -> (Key s k1 -> k2) + -> HashMap s k1 a + -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a +mapKeysWith f g (HashMap m) = SomeHashMapWith + (HashMap $ HashMap.fromListWith f + $ HashMap.foldrWithKey (\k x xs -> (g $ unsafeKey k, x) : xs) [] m) + $ MapProof (unsafeKey . g) \k2 -> + case HashMap.lookup (unrefine k2) backMap of + Nothing -> error + "mapKeysWith: bug: Data.HashMap.Refined has been subverted" + Just k1 -> k1 + where + ~backMap = HashMap.fromList + [ (k2, unsafeKey k1) + | k1 <- HashMap.keys m + , let !k2 = g $ unsafeKey k1 + ] + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect only the 'Just' results, returning a potentially smaller +-- map. +mapMaybeWithKey + :: forall s k a b. (Key s k -> a -> Maybe b) + -> HashMap s k a + -> SomeHashMapWith (SupersetProof 'Hashed s) k b +mapMaybeWithKey f (HashMap m) + = SomeHashMapWith (HashMap $ HashMap.mapMaybeWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect the 'Left' and 'Right' results into separate (disjoint) +-- maps. +mapEitherWithKey + :: forall s k a b c. Hashable k -- TODO: this is only used in the proof + => (Key s k -> a -> Either b c) + -> HashMap s k a + -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c +mapEitherWithKey p (HashMap m) + | m' <- HashMap.mapWithKey (p . unsafeKey) m + = Some2HashMapWith + (HashMap $ HashMap.mapMaybe (either Just (const Nothing)) m') + (HashMap $ HashMap.mapMaybe (either (const Nothing) Just) m') + $ PartitionProof + do \k -> case HashMap.lookup (unrefine k) m of + Nothing -> error + "mapEitherWithKey: bug: Data.HashMap.Refined has been subverted" + Just x -> case p k x of + Left _ -> Left $ unsafeKey $ unrefine k + Right _ -> Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Apply the inverse image of the given function to the keys of the given map, +-- so that for all @k :: 'Key' s2 k2@,c +-- @'backpermuteKeys' f m '!' k = m '!' f k@. +-- +-- If maps are identified with functions, this computes the composition. +backpermuteKeys + :: forall s1 s2 k1 k2 a. (Hashable k1, KnownHashSet s2 k2) + => (Key s2 k2 -> Key s1 k1) -> HashMap s1 k1 a -> HashMap s2 k2 a +backpermuteKeys f m = fromSet \k -> m ! f k diff --git a/src/Data/HashMap/Strict/Refined.hs b/src/Data/HashMap/Strict/Refined.hs new file mode 100644 index 0000000..24c52f6 --- /dev/null +++ b/src/Data/HashMap/Strict/Refined.hs @@ -0,0 +1,398 @@ +-- | This module defines a way to prove that a key exists in a map, so that the +-- key can be used to index into the map without using a 'Maybe', or manually +-- handling the \"impossible\" case with 'error' or other partial functions. +-- +-- To do this, @'HashMap' s k v@ has a type parameter @s@ that identifies its +-- set of keys, so that if another map has the same type parameter, you know +-- that map has the same set of keys. There is @'Key' s k@, a type of keys that +-- have been validated to belong to the set identified by @s@, and for which the +-- operation of indexing into a @'HashMap' s k v@ (only for the same @s@) can +-- proceed without failure (see '!'). The type @s@ itself has no internal +-- structure, rather it is merely a skolem type variable (rank-2 polymorphism +-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection". +-- +-- Like "Data.HashMap.Strict", functions in this module are strict in the keys +-- and values. The "Data.HashMap.Refined" module reuses the same 'HashMap' type +-- but provides functions that operate lazily on the values. +-- +-- = Warning +-- This module together with "Data.HashMap.Lazy" rely on 'Eq' and 'Hashable' +-- instances being lawful: that '==' is an equivalence relation, and that +-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence +-- relation; at least for the subset of values that are actually encountered at +-- runtime. If this assumption is violated, this module may not be able to +-- uphold its invariants and may throw errors. In particular beware of NaN in +-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@. +module Data.HashMap.Strict.Refined + ( + -- * Map type + Common.HashMap + , Common.Key + -- * Existentials and common proofs + , Common.SomeHashMap(..) + , Common.withHashMap + , Common.SomeHashMapWith(..) + , Common.withHashMapWith + , Common.Some2HashMapWith(..) + , Common.with2HashMapWith + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , Common.empty + , singleton + , SingletonProof(..) + , fromSet + , Common.fromHashMap + , fromTraversableWithKey + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + , reinsert + , insertLookupWithKey + -- * Deletion/Update + , Common.delete + , adjust + , adjustWithKey + , update + , updateLookupWithKey + -- * Query + , Common.lookup + , (Common.!) + , Common.member + , Common.null + , Common.isSubmapOfBy + , SubsetProof(..) + , Common.disjoint + , DisjointProof(..) + -- * Combine + , zipWithKey + , bind + , unionWithKey + , UnionProof(..) + , Common.difference + , DifferenceProof(..) + , differenceWithKey + , PartialDifferenceProof(..) + , intersectionWithKey + , IntersectionProof(..) + -- * Traversal + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , MapProof(..) + , backpermuteKeys + -- * Folds + , Common.foldMapWithKey + , Common.foldrWithKey + , Common.foldlWithKey + , Common.foldrWithKey' + , Common.foldlWithKey' + -- * Conversion + , Common.toMap + , Common.keysSet + , Common.toList + -- * Filter + , Common.restrictKeys + , Common.withoutKeys + , Common.filterWithKey + , Common.partitionWithKey + , PartitionProof(..) + , mapMaybeWithKey + , mapEitherWithKey + -- * Casts + , Common.castKey + , Common.cast + , castFlavor + ) where + +import Data.Coerce +import Data.Container.Refined.Hashable +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Functor +import qualified Data.HashMap.Strict as HashMap +import Data.HashMap.Common.Refined + ( HashMap(..), Key, unsafeCastKey, unsafeKey, SomeHashMapWith(..) + , Some2HashMapWith(..), (!) + ) +import qualified Data.HashMap.Common.Refined as Common +import qualified Data.HashSet as HashSet +import Data.Proxy +import Data.Reflection +import Data.Traversable +import Data.Traversable.WithIndex +import Data.Type.Coercion +import Prelude hiding (lookup, null) +import Refined +import Refined.Unsafe + + +-- | Create a map with a single key-value pair, and return a proof that the +-- key is in the resulting map. +singleton + :: forall k a. Hashable k + => k -> a -> SomeHashMapWith (SingletonProof 'Hashed k) k a +singleton k v = SomeHashMapWith (HashMap $ HashMap.singleton k v) + $ SingletonProof (unsafeKey k) + +-- | Create a map from a set of keys, and a function that for each key computes +-- the corresponding value. +fromSet :: forall s k a. KnownHashSet s k => (Key s k -> a) -> HashMap s k a +fromSet f = HashMap $ HashMap.mapWithKey (\k _ -> f $ unsafeKey k) + $ HashSet.toMap (reflect $ Proxy @s) + +-- | Create a map from an arbitrary traversable of key-value pairs. +fromTraversableWithKey + :: forall t k a. (Traversable t, Hashable k) + => (k -> a -> a -> a) + -> t (k, a) + -> SomeHashMapWith (FromTraversableProof 'Hashed t k) k a +fromTraversableWithKey f xs = SomeHashMapWith (HashMap m) + $ FromTraversableProof proof + where + (m, proof) = mapAccumL + (\s (k, v) + -> let !s' = HashMap.insertWith (f k) k v s in (s', unsafeKey k)) + HashMap.empty + xs + +-- | Insert a key-value pair into the map to obtain a potentially larger map, +-- guaranteed to contain the given key. If the key was already present, the +-- associated value is replaced with the supplied value. +insert + :: forall s k a. Hashable k + => k -> a -> HashMap s k a -> SomeHashMapWith (InsertProof 'Hashed k s) k a +insert k v (HashMap m) = SomeHashMapWith (HashMap $ HashMap.insert k v m) + $ InsertProof (unsafeKey k) unsafeSubset + +-- | Overwrite a key-value pair that is known to already be in the map. The set +-- of keys remains the same. +reinsert + :: forall s k a. Hashable k + => Key s k -> a -> HashMap s k a -> HashMap s k a +reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.insert @k @a + +-- | Insert a key-value pair into the map using a combining function, and if +-- the key was already present, the old value is returned along with the proof +-- that the key was present. +insertLookupWithKey + :: forall s k a. Hashable k + => (Key s k -> a -> a -> a) + -> k + -> a + -> HashMap s k a + -> (Maybe (Key s k, a), SomeHashMapWith (InsertProof 'Hashed k s) k a) +insertLookupWithKey f k v (HashMap m) = + ( (unsafeKey k,) <$> HashMap.lookup k m + , SomeHashMapWith (HashMap $ HashMap.insertWith (f $ unsafeKey k) k v m) + $ InsertProof (unsafeKey k) unsafeSubset + ) + +-- | Update the value at a specific key known the be in the map using the given +-- function. The set of keys remains the same. +adjust + :: forall s k a. Hashable k + => (a -> a) -> Key s k -> HashMap s k a -> HashMap s k a +adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ HashMap.adjust @k @a + +-- | If the given key is in the map, update the associated value using the given +-- function with a proof that the key was in the map; otherwise return the map +-- unchanged. In any case the set of keys remains the same. +adjustWithKey + :: forall s k a. Hashable k + => (Key s k -> a -> a) -> k -> HashMap s k a -> HashMap s k a +adjustWithKey f k (HashMap m) = HashMap $ HashMap.adjust (f $ unsafeKey k) k m + +-- | Update or delete a key known to be in the map using the given function, +-- returning a potentially smaller map. +update + :: forall s k a. Hashable k + => (a -> Maybe a) + -> Key s k + -> HashMap s k a + -> SomeHashMapWith (SupersetProof 'Hashed s) k a +update f k (HashMap m) + = SomeHashMapWith (HashMap $ HashMap.update f (unrefine k) m) + $ SupersetProof unsafeSubset + +-- | If the given key is in the map, update or delete it using the given +-- function with a proof that the key was in the map; otherwise the map is +-- unchanged. Alongside return the new value if it was updated, or the old value +-- if it was deleted, and a proof that the key was in the map. +updateLookupWithKey + :: forall s k a. Hashable k + => (Key s k -> a -> Maybe a) + -> k + -> HashMap s k a + -> (Maybe (Key s k, a), SomeHashMapWith (SupersetProof 'Hashed s) k a) +updateLookupWithKey f k (HashMap m) = + ( (unsafeKey k,) <$> HashMap.lookup k m + , SomeHashMapWith (HashMap $ HashMap.update (f $ unsafeKey k) k m) + $ SupersetProof unsafeSubset + ) + +-- | Given two maps proven to have the same keys, for each key apply the +-- function to the associated values, to obtain a new map with the same keys. +zipWithKey + :: forall s k a b c. Hashable k + => (Key s k -> a -> b -> c) -> HashMap s k a -> HashMap s k b -> HashMap s k c +zipWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.intersectionWithKey @k @a @b @c + +-- | Return the union of two maps, with a given combining function for keys that +-- exist in both maps simultaneously. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +unionWithKey + :: forall s t k a. Hashable k + => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> a -> a) + -> HashMap s k a + -> HashMap t k a + -> SomeHashMapWith (UnionProof 'Hashed s t) k a +unionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith + (HashMap $ HashMap.unionWithKey (f . reallyUnsafeRefine) m1 m2) + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- | For keys that appear in both maps, the given function decides whether the +-- key is removed from the first map. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +differenceWithKey + :: forall s t k a b. Hashable k + => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> Maybe a) + -> HashMap s k a + -> HashMap t k b + -> SomeHashMapWith (PartialDifferenceProof 'Hashed s t) k a +differenceWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith + (HashMap $ HashMap.differenceWith + (\x (k, y) -> f (reallyUnsafeRefine k) x y) + m1 + (HashMap.mapWithKey (,) m2)) + $ PartialDifferenceProof unsafeSubset unsafeSubset + +-- | Return the intersection of two maps with the given combining function. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +intersectionWithKey + :: forall s t k a b c. Hashable k + => (Refined (InSet 'Hashed s && InSet 'Hashed t) k -> a -> b -> c) + -> HashMap s k a + -> HashMap t k b + -> SomeHashMapWith (IntersectionProof 'Hashed s t) k c +intersectionWithKey f (HashMap m1) (HashMap m2) = SomeHashMapWith + (HashMap $ HashMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, that are proven to be in the map. The set of keys remains the same. +mapWithKey + :: forall s k a b. (Key s k -> a -> b) -> HashMap s k a -> HashMap s k b +mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ HashMap.mapWithKey @k @a @b + +-- | Map an 'Applicative' transformation with access to each value's +-- corresponding key and a proof that it is in the map. The set of keys remains +-- unchanged. +traverseWithKey + :: forall s f k a b. Applicative f + => (Key s k -> a -> f b) -> HashMap s k a -> f (HashMap s k b) +traverseWithKey f (HashMap m) + = HashMap <$> HashMap.traverseWithKey (f . unsafeKey) m + +-- | Thread an accumularing argument through the map in ascending order of +-- hashes. +mapAccumLWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> HashMap s k b + -> (a, HashMap s k c) +mapAccumLWithKey f = imapAccumL (flip f) + +-- | Thread an accumularing argument through the map in descending order of +-- hashes. +mapAccumRWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> HashMap s k b + -> (a, HashMap s k c) +mapAccumRWithKey f = imapAccumR (flip f) + +-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the +-- results into a new map. For keys that were mapped to the same new key, @c@ +-- acts as the combining function for corresponding values. +mapKeysWith + :: forall s k1 k2 a. Hashable k2 + => (a -> a -> a) + -> (Key s k1 -> k2) + -> HashMap s k1 a + -> SomeHashMapWith (MapProof 'Hashed s k1 k2) k2 a +mapKeysWith f g (HashMap m) = SomeHashMapWith + (HashMap $ HashMap.fromListWith f + $ HashMap.foldrWithKey (\k x xs -> (g $ unsafeKey k, x) : xs) [] m) + $ MapProof (unsafeKey . g) \k2 -> + case HashMap.lookup (unrefine k2) backMap of + Nothing -> error + "mapKeysWith: bug: Data.HashMap.Refined has been subverted" + Just k1 -> k1 + where + ~backMap = HashMap.fromList + [ (k2, unsafeKey k1) + | k1 <- HashMap.keys m + , let !k2 = g $ unsafeKey k1 + ] + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect only the 'Just' results, returning a potentially smaller +-- map. +mapMaybeWithKey + :: forall s k a b. (Key s k -> a -> Maybe b) + -> HashMap s k a + -> SomeHashMapWith (SupersetProof 'Hashed s) k b +mapMaybeWithKey f (HashMap m) + = SomeHashMapWith (HashMap $ HashMap.mapMaybeWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect the 'Left' and 'Right' results into separate (disjoint) +-- maps. +mapEitherWithKey + :: forall s k a b c. Hashable k -- TODO: this is only used in the proof + => (Key s k -> a -> Either b c) + -> HashMap s k a + -> Some2HashMapWith (PartitionProof 'Hashed s k) k b c +mapEitherWithKey p (HashMap m) + | m' <- HashMap.mapWithKey (p . unsafeKey) m + = Some2HashMapWith + (HashMap $ HashMap.mapMaybe (either Just (const Nothing)) m') + (HashMap $ HashMap.mapMaybe (either (const Nothing) Just) m') + $ PartitionProof + do \k -> case HashMap.lookup (unrefine k) m of + Nothing -> error + "mapEitherWithKey: bug: Data.HashMap.Refined has been subverted" + Just x -> case p k x of + Left _ -> Left $ unsafeKey $ unrefine k + Right _ -> Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the +-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions. +bind + :: forall s k a b. Hashable k + => HashMap s k a -> (a -> HashMap s k b) -> HashMap s k b +bind m f = mapWithKey (\k x -> f x ! k) m + +-- | Apply the inverse image of the given function to the keys of the given map, +-- so that for all @k :: 'Key' s2 k2@, +-- @'backpermuteKeys' f m '!' k = m '!' f k@. +-- +-- If maps are identified with functions, this computes the composition. +backpermuteKeys + :: forall s1 s2 k1 k2 a. (Hashable k1, KnownHashSet s2 k2) + => (Key s2 k2 -> Key s1 k1) -> HashMap s1 k1 a -> HashMap s2 k2 a +backpermuteKeys f m = fromSet \k -> m ! f k diff --git a/src/Data/HashSet/Refined.hs b/src/Data/HashSet/Refined.hs new file mode 100644 index 0000000..d6f297a --- /dev/null +++ b/src/Data/HashSet/Refined.hs @@ -0,0 +1,418 @@ +{-# LANGUAGE CPP #-} +-- | This module implements a way of tracking the contents of a +-- 'Data.HashSet.HashSet' at the type level, and provides utilities for +-- manipulating such sets. +-- +-- The contents of a set are associated with a type parameter, e.g. @s@, so that +-- whenever you see the same type parameter, you know you are working with the +-- same set. The type @s@ itself has no internal structure, rather it is merely +-- a skolem type variable (rank-2 polymorphism 'Control.Monad.ST.runST' trick) +-- introduced by "Data.Reflection". +-- +-- = Warning +-- This module together with "Data.HashSet" rely on 'Eq' and 'Hashable' +-- instances being lawful: that '==' is an equivalence relation, and that +-- 'Data.Hashable.hashWithSalt' is defined on the quotient by this equivalence +-- relation; at least for the subset of values that are actually encountered at +-- runtime. If this assumption is violated, this module may not be able to +-- uphold its invariants and may throw errors. In particular beware of NaN in +-- 'Float' and 'Double', and, if using @hashable < 1.3@, beware of @0@ and @-0@. +module Data.HashSet.Refined + ( + -- * Set type + KnownHashSet + , HashSet + -- * Refinement type + , InSet(..) + , Flavor(Hashed) + , Element + , revealPredicate + -- * Existentials and common proofs + , SomeHashSet(..) + , withHashSet + , SomeHashSetWith(..) + , withHashSetWith + , Some2HashSetWith(..) + , with2HashSetWith + , (:->) + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , empty + , singleton + , SingletonProof(..) + , fromHashSet + , fromTraversable + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + -- * Deletion + , delete + -- * Query + , member + , null + , isSubsetOf + , SubsetProof(..) + , disjoint + , DisjointProof(..) + -- * Combine + , union + , UnionProof(..) + , difference + , DifferenceProof(..) + , intersection + , IntersectionProof(..) + -- * Filter + , filter + , partition + , PartitionProof(..) + -- * Map + , map + , MapProof(..) + -- * Folds + , foldMap + , foldr + , foldl + , foldr' + , foldl' + -- * Conversion + , toList + , asSet + , asIntSet + -- * Casts + , castElement + , cast + , castFlavor + ) where + +import Data.Coerce +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Conversion +import Data.Container.Refined.Hashable +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import qualified Data.Foldable as Foldable +import qualified Data.HashMap.Lazy as HashMap +import qualified Data.HashSet as HashSet +import Data.Proxy +import Data.Reflection +import Data.Traversable +import Data.Type.Coercion +import Data.Type.Equality ((:~:)(..)) +import Data.Typeable (Typeable) +import GHC.Exts (Proxy#, proxy#) +import Prelude hiding (filter, foldl, foldMap, foldr, map, null) +import Refined +import Refined.Unsafe +import Unsafe.Coerce + + +-- | To use "Refined" machinery that uses the 'Predicate' typeclass you will +-- need to pattern match on this 'Dict'. +-- +-- The reason is that in the default /fast/ implementation of reflection, we +-- don't have @'Typeable' s@, which "Refined" wants for pretty-printing +-- exceptions. We /can/ provide @'Typeable' s@, but at the cost of using the +-- /slow/ implementation of reflection. +revealPredicate + :: forall s a. (Typeable a, Hashable a, KnownHashSet s a) + => Dict (Predicate (InSet 'Hashed s) a) +revealPredicate = reifyTypeable (reflect (Proxy @s)) + \(_ :: Proxy s') -> + reflect (Proxy @s') `seq` + -- ^ Work around https://github.com/ekmett/reflection/issues/54 + case unsafeCoerce Refl :: s :~: s' of + Refl -> Dict + +-- | @'Element' s a@ is a value of type @a@ that has been verified to be an +-- element of @s@. +-- +-- Thus, @'Element' s a@ is a \"refinement\" type of @a@, and this library +-- integrates with an implementation of refimenement types in "Refined", so +-- the machinery from there can be used to manipulate 'Element's (however see +-- 'revealPredicate'). +-- +-- The underlying @a@ value can be obtained with 'unrefine'. An @a@ can be +-- validated into an @'Element' s a@ with 'member'. +type Element s = Refined (InSet 'Hashed s) + +unsafeCastElement :: forall s a. Coercion a (Element s a) +unsafeCastElement = reallyUnsafeUnderlyingRefined + +unsafeElement :: a -> Element s a +unsafeElement = coerceWith unsafeCastElement + +-- | An existential wrapper for an as-yet-unknown set. Pattern maching on it +-- gives you a way to refer to the set, e.g. +-- +-- @ +-- case 'fromHashSet' ... of +-- 'SomeHashSet' \@s _ -> doSomethingWith \@s +-- +-- case 'fromHashSet' ... of +-- 'SomeHashSet' (_ :: 'Proxy#' s) -> doSomethingWith \@s +-- @ +data SomeHashSet a where + SomeHashSet :: forall s a. KnownHashSet s a => Proxy# s -> SomeHashSet a + +-- | Apply an unknown set to a continuation that can accept any set. This gives +-- you a way to refer to the set (the parameter @s@), e.g.: +-- +-- @ +-- 'withHashSet' ('fromHashSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s +-- @ +withHashSet + :: forall a r. SomeHashSet a + -> (forall s. KnownHashSet s a => Proxy s -> r) + -> r +withHashSet (SomeHashSet (_ :: Proxy# s)) k = k $ Proxy @s + +-- | Construct a set from a regular 'Data.HashSet.HashSet'. +fromHashSet :: forall a. HashSet.HashSet a -> SomeHashSet a +fromHashSet s = reify s \(_ :: Proxy s) -> SomeHashSet @s proxy# + +-- | An existential wrapper for an as-yet-unknown set, together with a proof of +-- some fact @p@ about the set. Pattern matching on it gives you a way to refer +-- to the set (the parameter @s@). Most functions will return a set in this way, +-- together with a proof that somehow relates the set to the function's inputs. +data SomeHashSetWith p a where + SomeHashSetWith + :: forall s a p. KnownHashSet s a => !(p s) -> SomeHashSetWith p a + +-- | Apply an unknown set with proof to a continuation that can accept any set +-- satisfying the proof. This gives you a way to refer to the set (the parameter +-- @s@). +withHashSetWith + :: forall a r p. SomeHashSetWith p a + -> (forall s. KnownHashSet s a => p s -> r) + -> r +withHashSetWith (SomeHashSetWith p) k = k p + +-- | An existential wrapper for an as-yet-unknown pair of sets, together with +-- a proof of some fact @p@ relating them. +data Some2HashSetWith p a where + Some2HashSetWith + :: forall s t a p. (KnownHashSet s a, KnownHashSet t a) + => !(p s t) -> Some2HashSetWith p a + +-- | Apply a pair of unknown sets with proof to a continuation that can accept +-- any pair of sets satisfying the proof. This gives you a way to refer to the +-- sets (the parameters @s@ and @t@). +with2HashSetWith + :: forall a r p. Some2HashSetWith p a + -> (forall s t. (KnownHashSet s a, KnownHashSet t a) => p s t -> r) + -> r +with2HashSetWith (Some2HashSetWith p) k = k p + +-- | An empty set. +empty :: forall a. SomeHashSetWith (EmptyProof 'Hashed) a +empty = reify HashSet.empty \(_ :: Proxy r) + -> SomeHashSetWith @r $ EmptyProof unsafeSubset + +-- | Create a set with a single element. +singleton + :: forall a. Hashable a => a -> SomeHashSetWith (SingletonProof 'Hashed a) a +singleton x = reify (HashSet.singleton x) \(_ :: Proxy r) + -> SomeHashSetWith @r $ SingletonProof $ unsafeElement x + +-- | Create a set from the elements of an arbitrary traversable. +fromTraversable + :: forall t a. (Traversable t, Hashable a) + => t a -> SomeHashSetWith (FromTraversableProof 'Hashed t a) a +fromTraversable xs = reify set \(_ :: Proxy r) + -> SomeHashSetWith @r $ FromTraversableProof + $ unsafeCoerce @(t (Element _ a)) @(t (Element r a)) proof + where + (set, proof) = mapAccumL + (\s x -> let !s' = HashSet.insert x s in (s', unsafeElement x)) + HashSet.empty + xs + +-- | Insert an element in a set. +insert + :: forall s a. (Hashable a, KnownHashSet s a) + => a -> SomeHashSetWith (InsertProof 'Hashed a s) a +insert x = reify (HashSet.insert x $ reflect $ Proxy @s) \(_ :: Proxy r) + -> SomeHashSetWith @r $ InsertProof (unsafeElement x) unsafeSubset + +-- | Delete an element from a set. +delete + :: forall s a. (Hashable a, KnownHashSet s a) + => a -> SomeHashSetWith (SupersetProof 'Hashed s) a +delete x = reify (HashSet.delete x $ reflect $ Proxy @s) \(_ :: Proxy r) + -> SomeHashSetWith @s $ SupersetProof unsafeSubset + +-- | If an element is in the set, return the proof that it is. +member :: forall s a. (Hashable a, KnownHashSet s a) => a -> Maybe (Element s a) +member x + | x `HashSet.member` reflect (Proxy @s) = Just $ unsafeElement x + | otherwise = Nothing + +-- | If the set is empty, return the proof that it is. +null :: forall s a. KnownHashSet s a => Maybe (EmptyProof 'Hashed s) +null + | HashSet.null $ reflect $ Proxy @s = Just $ EmptyProof unsafeSubset + | otherwise = Nothing + +-- | If @s@ is a subset of @t@ (or is equal to), return a proof of that. +isSubsetOf + :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a) + => Maybe (SubsetProof 'Hashed s t) +isSubsetOf +#if MIN_VERSION_unordered_containers(0, 2, 12) + | reflect (Proxy @s) `HashSet.isSubsetOf` reflect (Proxy @t) +#else + | all (`HashSet.member` reflect (Proxy @t)) (reflect (Proxy @s)) +#endif + = Just $ SubsetProof unsafeSubset + | otherwise = Nothing + +-- | If @s@ and @t@ are disjoint (i.e. their intersection is empty), return a +-- proof of that. +disjoint + :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a) + => Maybe (DisjointProof 'Hashed s t) +disjoint + | HashSet.null + $ HashSet.intersection (reflect $ Proxy @s) (reflect $ Proxy @t) + = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g + | otherwise = Nothing + +-- | The union of two sets. +union + :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a) + => SomeHashSetWith (UnionProof 'Hashed s t) a +union = reify (reflect (Proxy @s) `HashSet.union` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeHashSetWith @r + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- unions :: ? + +-- | HashSet with elements of @s@ that are not in @t@. +difference + :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a) + => SomeHashSetWith (DifferenceProof 'Hashed s t) a +difference = reify (reflect (Proxy @s) `HashSet.difference` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeHashSetWith @r + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Intersection of two sets. +intersection + :: forall s t a. (Hashable a, KnownHashSet s a, KnownHashSet t a) + => SomeHashSetWith (IntersectionProof 'Hashed s t) a +intersection + = reify (reflect (Proxy @s) `HashSet.intersection` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeHashSetWith @r + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Return a subset of elements that satisfy the given predicate. +filter + :: forall s a. KnownHashSet s a + => (Element s a -> Bool) -> SomeHashSetWith (SupersetProof 'Hashed s) a +filter p = reify (HashSet.filter (p . unsafeElement) $ reflect $ Proxy @s) + \(_ :: Proxy r) -> SomeHashSetWith @r $ SupersetProof unsafeSubset + +-- | Partition a set into two disjoint subsets: those that satisfy the +-- predicate, and those that don't. +partition + :: forall s a. KnownHashSet s a + => (Element s a -> Bool) -> Some2HashSetWith (PartitionProof 'Hashed s a) a +partition p = reify (HashSet.filter (p . unsafeElement) $ reflect $ Proxy @s) + \(_ :: Proxy r) + -> reify (HashSet.filter (not . p . unsafeElement) $ reflect $ Proxy @s) + \(_ :: Proxy q) + -> Some2HashSetWith @s @r $ PartitionProof + do \x -> if p x + then Left $ unsafeElement $ unrefine x + else Right $ unsafeElement $ unrefine x + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Apply the given function to each element of the set and collect the +-- results. Note that the resulting set can be smaller. +map + :: forall s a b. (Hashable b, KnownHashSet s a) + => (Element s a -> b) -> SomeHashSetWith (MapProof 'Hashed s a b) b +map f = reify (HashMap.keysSet m) + \(_ :: Proxy r) -> SomeHashSetWith @r + $ MapProof (unsafeElement . f) \y -> case HashMap.lookup (unrefine y) m of + Nothing -> error "map: bug: Data.HashSet.Refined has been subverted" + Just x -> x + where + !m = HashMap.fromList + [ (y, unsafeElement x) + | x <- HashSet.toList $ reflect $ Proxy @s + , let !y = f $ unsafeElement x + ] + +-- | Map each element of @s@ into a monoid (with proof that it was an element), +-- and combine the results using 'Data.Monoid.<>'. +foldMap :: forall s a m. (KnownHashSet s a, Monoid m) => (Element s a -> m) -> m +foldMap f = Foldable.foldMap (f . unsafeElement) $ reflect $ Proxy @s + +-- | Right associative fold with a lazy accumulator. +foldr :: forall s a b. KnownHashSet s a => (Element s a -> b -> b) -> b -> b +foldr f z = HashSet.foldr (f . unsafeElement) z $ reflect $ Proxy @s + +-- | Left associative fold with a lazy accumulator. +foldl :: forall s a b. KnownHashSet s a => (b -> Element s a -> b) -> b -> b +foldl f z = Foldable.foldl ((. unsafeElement) . f) z $ reflect $ Proxy @s + +-- | Right associative fold with a strict accumulator. +foldr' :: forall s a b. KnownHashSet s a => (Element s a -> b -> b) -> b -> b +foldr' f z = Foldable.foldr' (f . unsafeElement) z $ reflect $ Proxy @s + +-- | Left associative fold with a strict accumulator. +foldl' :: forall s a b. KnownHashSet s a => (b -> Element s a -> b) -> b -> b +foldl' f z = HashSet.foldl' ((. unsafeElement) . f) z $ reflect $ Proxy @s + +-- | List of elements in the set. +toList :: forall s a. KnownHashSet s a => [Element s a] +toList = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ HashSet.toList $ reflect $ Proxy @s + + +-- | Convert an 'IntSet' into a 'Set', retaining its set of elements, which can +-- be converted with 'castFlavor'. +asSet :: forall s a. (Ord a, KnownHashSet s a) => Set s a +asSet = hashSet2Set + +-- | Convert an 'IntSet' into a 'HashSet', retaining its set of elements, which +-- can be converted with 'castFlavor'. +asIntSet :: forall s. KnownHashSet s Int => IntSet s +asIntSet = hashSet2IntSet + +-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then +-- @s@ and @t@ actually stand for the same set and @'Element' s@ can be safely +-- interconverted with @'Element' t@. +-- +-- The requirement that the weakenings are natural transformations ensures that +-- they don't actually alter the elements. To build these you can compose +-- ':->''s from proofs returned by functions in this module, or "Refined" +-- functions like 'andLeft' or 'leftOr'. +castElement + :: forall s t a. (forall x. Element s x -> Element t x) + -> (forall x. Element t x -> Element s x) + -> Coercion (Element s a) (Element t a) +castElement = castRefined + +-- | If elements can be interconverted (e.g. as proved by 'castElement'), then +-- the sets can be interconverted too. For example we can establish that the +-- intersection of a set with itself is interconvertible with that set: +-- +-- @ +-- castIntersection +-- :: t'IntersectionProof' ''Data.HashSet.Refined.Hashed' s s r +-- -> 'Coercion' ('HashSet' r a) ('HashSet' s a) +-- castIntersection ( v'IntersectionProof' p1 p2) +-- = 'cast' $ 'castElement' ('andLeft' . p1) (p2 'id' 'id') +-- @ +cast + :: forall s t a. (forall x. Coercion (Element s x) (Element t x)) + -> Coercion (HashSet s a) (HashSet t a) +cast Coercion +#if MIN_VERSION_base(4, 15, 0) + = case unsafeEqualityProof @s @t of UnsafeRefl -> Coercion +#else + = repr $ unsafeCoerce Refl +#endif diff --git a/src/Data/IntMap/Common/Refined.hs b/src/Data/IntMap/Common/Refined.hs new file mode 100644 index 0000000..5f2a692 --- /dev/null +++ b/src/Data/IntMap/Common/Refined.hs @@ -0,0 +1,523 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.IntMap.Common.Refined where + +import Control.Monad.Reader +import Control.DeepSeq +import Data.Coerce +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Distributive +import Data.Foldable.WithIndex +import Data.Functor.Rep +import Data.Functor.WithIndex +import qualified Data.Hashable as Hashable +import qualified Data.IntMap as IntMap +import Data.Proxy +import Data.Reflection +import Data.Traversable.WithIndex +import Data.Type.Coercion +import Data.Type.Equality ((:~:)(..)) +import Refined +import Refined.Unsafe +import Unsafe.Coerce + +#if MIN_VERSION_containers(0, 6, 7) +#elif MIN_VERSION_containers(0, 6, 2) +import qualified Data.List as List +#elif MIN_VERSION_containers(0, 5, 8) +import Data.Functor.Const (Const(..)) +import qualified Data.List as List +import Data.Monoid (Any(..)) +import qualified Data.IntMap.Merge.Lazy as IntMap +#else +import qualified Data.IntMap.Strict as IntMapStrict +import qualified Data.List as List +#endif + + +-- | A wrapper around a regular 'Data.IntMap.IntMap' with a type parameter @s@ +-- identifying the set of keys present in the map. +-- +-- An t'Int' key may not be present in the map, but a @'Key' s@ is guaranteed to +-- be present (if the @s@ parameters match). Thus the map is isomorphic to a +-- (total) function @'Key' s -> a@, which motivates many of the instances below. +-- +-- An 'IntMap' always knows its set of keys, so given @'IntMap' s a@ we can +-- always derive @'KnownIntSet' s@ by pattern matching on the 'Dict' returned by +-- 'keysSet'. +newtype IntMap s a = IntMap (IntMap.IntMap a) + deriving newtype (Eq, Ord, Show, Functor, Foldable, NFData) +#if MIN_VERSION_hashable(1, 3, 4) + deriving newtype (Hashable.Hashable) +#endif + deriving stock (Traversable) +type role IntMap nominal representational + +-- | Convert to a regular 'Data.IntMap.IntMap', forgetting its set of keys. +toIntMap :: forall s a. IntMap s a -> IntMap.IntMap a +toIntMap (IntMap m) = m + +-- | @'Key' s@ is a key of type t'Int' that has been verified to be an element +-- of the set @s@, and thus verified to be present in all @'IntMap' s k@ maps. +-- +-- Thus, @'Key' s@ is a \"refinement\" type of t'Int', and this library +-- integrates with an implementation of refimenement types in "Refined", so +-- the machinery from there can be used to manipulate 'Key's (however see +-- 'Data.IntSet.Refined.revealPredicate'). +-- +-- The underlying t'Int' value can be obtained with 'unrefine'. An t'Int' can be +-- validated into an @'Key' s@ with 'member'. +type Key s = Refined (InSet 'Int s) Int + +unsafeCastKey :: forall s. Coercion Int (Key s) +unsafeCastKey = reallyUnsafeUnderlyingRefined + +unsafeKey :: Int -> Key s +unsafeKey = coerceWith unsafeCastKey + +-- | An existential wrapper for an 'IntMap' with an as-yet-unknown set of keys. +-- Pattern maching on it gives you a way to refer to the set (the parameter +-- @s@), e.g. +-- +-- @ +-- case 'fromIntMap' ... of +-- 'SomeIntMap' \@s m -> doSomethingWith \@s +-- +-- case 'fromIntMap' ... of +-- 'SomeIntMap' (m :: 'IntMap' s a) -> doSomethingWith \@s +-- @ +data SomeIntMap a where + SomeIntMap :: forall s a. !(IntMap s a) -> SomeIntMap a + +-- | Apply a map with an unknown set of keys to a continuation that can accept +-- a map with any set of keys. This gives you a way to refer to the set (the +-- parameter @s@), e.g.: +-- +-- @ +-- 'withIntMap' ('fromIntMap' ...) $ \(m :: 'IntMap' s a) -> doSomethingWith \@s +-- @ +withIntMap :: forall a r. SomeIntMap a -> (forall s. IntMap s a -> r) -> r +withIntMap (SomeIntMap m) k = k m + +-- | Construct a map from a regular 'Data.IntMap.IntMap'. +fromIntMap :: forall a. IntMap.IntMap a -> SomeIntMap a +fromIntMap m = SomeIntMap (IntMap m) + +-- | An existential wrapper for an 'IntMap' with an as-yet-unknown set of keys, +-- together with a proof of some fact @p@ about the set. Pattern matching on it +-- gives you a way to refer to the set (the parameter @s@). Functions that +-- change the set of keys in a map will return the map in this way, together +-- with a proof that somehow relates the keys set to the function's inputs. +data SomeIntMapWith p a where + SomeIntMapWith :: forall s a p. !(IntMap s a) -> !(p s) -> SomeIntMapWith p a + +-- | Apply a map with proof for an unknown set of keys to a continuation that +-- can accept a map with any set of keys satisfying the proof. This gives you a +-- way to refer to the set (the parameter @s@). +withIntMapWith + :: forall a r p. SomeIntMapWith p a -> (forall s. IntMap s a -> p s -> r) -> r +withIntMapWith (SomeIntMapWith m p) k = k m p + +-- | An existential wrapper for a pair of maps with as-yet-unknown sets of keys, +-- together with a proof of some fact @p@ relating them. +data Some2IntMapWith p a b where + Some2IntMapWith + :: forall s t a b p. !(IntMap s a) + -> !(IntMap t b) + -> !(p s t) + -> Some2IntMapWith p a b + +-- | Apply a pair of maps with proof for unknown sets of keys to a continuation +-- that can accept any pair of maps with any sets of keys satisfying the proof. +-- This gives you a way to refer to the sets (the parameters @s@ and @t@). +with2IntMapWith + :: forall a b r p. Some2IntMapWith p a b + -> (forall s t. IntMap s a -> IntMap t b -> p s t -> r) + -> r +with2IntMapWith (Some2IntMapWith m1 m2 p) k = k m1 m2 p + +-- | An empty map. +empty :: forall a. SomeIntMapWith (EmptyProof 'Int) a +empty = SomeIntMapWith (IntMap IntMap.empty) $ EmptyProof unsafeSubset + +-- | Create a map from a set of keys, and a function that for each key computes +-- the corresponding value. +fromSet :: forall s a. KnownIntSet s => (Key s -> a) -> IntMap s a +fromSet f = IntMap $ IntMap.fromSet (f . unsafeKey) (reflect $ Proxy @s) + +-- | Delete a key and its value from the map if present, returning a potentially +-- smaller map. +delete :: forall s a. Int -> IntMap s a -> SomeIntMapWith (SupersetProof 'Int s) a +delete k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.delete k m) + $ SupersetProof unsafeSubset + +-- | If the key is in the map, return the proof of this, and the associated +-- value; otherwise return 'Nothing'. +lookup :: forall s a. Int -> IntMap s a -> Maybe (Key s, a) +lookup k (IntMap m) = (unsafeKey k,) <$> IntMap.lookup k m + +-- | Given a key that is proven to be in the map, return the associated value. +-- +-- Unlike 'Data.IntMap.!' from "Data.IntMap", this function is total, as it is +-- impossible to obtain a @'Key' s@ for a key that is not in the map +-- @'IntMap' s a@. +(!) :: forall s a. IntMap s a -> Key s -> a +(!) (IntMap m) k = case IntMap.lookup (unrefine k) m of + Nothing -> error "(!): bug: Data.IntMap.Refined has been subverted" + Just x -> x + +-- | If a key is in the map, return the proof that it is. +member :: forall s a. Int -> IntMap s a -> Maybe (Key s) +member k (IntMap m) + | k `IntMap.member` m = Just (unsafeKey k) + | otherwise = Nothing + +-- | Find the largest key smaller than the given one, and return the +-- associated key-value pair. +lookupLT :: forall s a. Int -> IntMap s a -> Maybe (Key s, a) +lookupLT = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupLT @a + +-- | Find the smallest key greater than the given one, and return the +-- associated key-value pair. +lookupGT :: forall s a. Int -> IntMap s a -> Maybe (Key s, a) +lookupGT = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupGT @a + +-- | Find the largest key smaller or equal to the given one, and return the +-- associated key-value pair. +lookupLE :: forall s a. Int -> IntMap s a -> Maybe (Key s, a) +lookupLE = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupLE @a + +-- | Find the smallest key greater or equal to the given one, and return the +-- associated key-value pair. +lookupGE :: forall s a. Int -> IntMap s a -> Maybe (Key s, a) +lookupGE = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.lookupGE @a + +-- | If a map is empty, return a proof that it is. +null :: forall s a. IntMap s a -> Maybe (EmptyProof 'Int s) +null (IntMap m) + | IntMap.null m = Just $ EmptyProof unsafeSubset + | otherwise = Nothing + +-- | If all keys of the first map are also present in the second map, and the +-- given function returns 'True' for their associated values, return a proof +-- that the keys form a subset. +isSubmapOfBy + :: forall s t a b. (a -> b -> Bool) + -> IntMap s a + -> IntMap t b + -> Maybe (SubsetProof 'Int s t) +isSubmapOfBy f (IntMap m1) (IntMap m2) + | IntMap.isSubmapOfBy f m1 m2 = Just $ SubsetProof unsafeSubset + | otherwise = Nothing + +-- | If two maps are disjoint (i.e. their intersection is empty), return a proof +-- of that. +disjoint + :: forall s t a b. IntMap s a -> IntMap t b -> Maybe (DisjointProof 'Int s t) +disjoint (IntMap m1) (IntMap m2) +#if MIN_VERSION_containers(0, 6, 2) + | IntMap.disjoint m1 m2 +#elif MIN_VERSION_containers(0, 5, 8) + | Const (Any False) <- IntMap.mergeA + (IntMap.traverseMissing \_ _ -> Const mempty) + (IntMap.traverseMissing \_ _ -> Const mempty) + (IntMap.zipWithAMatched \_ _ _ -> Const $ Any True) + m1 + m2 +#else + | IntMap.null $ IntMapStrict.intersectionWith (\_ _ -> ()) m1 m2 +#endif + = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g + | otherwise = Nothing + +-- | Given two maps proven to have the same keys, for each key apply the +-- function to the associated values, to obtain a new map with the same keys. +zipWithKey + :: forall s a b c. (Key s -> a -> b -> c) + -> IntMap s a + -> IntMap s b + -> IntMap s c +zipWithKey f (IntMap m1) (IntMap m2) = IntMap + $ IntMap.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y) + (error "zipWithKey: bug: Data.IntMap.Refined has been subverted") + (error "zipWithKey: bug: Data.IntMap.Refined has been subverted") + m1 + m2 + +-- | Remove the keys that appear in the second map from the first map. +difference + :: forall s t a b. IntMap s a + -> IntMap t b + -> SomeIntMapWith (DifferenceProof 'Int s t) a +difference (IntMap m1) (IntMap m2) = SomeIntMapWith + (IntMap $ IntMap.difference m1 m2) + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, that are proven to be in the map. The set of keys remains the same. +mapWithKey :: forall s a b. (Key s -> a -> b) -> IntMap s a -> IntMap s b +mapWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.mapWithKey @a @b + +-- | Map an 'Applicative' transformation in ascending order of keys, with access +-- to each value's corresponding key and a proof that it is in the map. The set +-- of keys remains unchanged. +traverseWithKey + :: forall s f a b. Applicative f + => (Key s -> a -> f b) -> IntMap s a -> f (IntMap s b) +traverseWithKey f (IntMap m) = IntMap <$> IntMap.traverseWithKey (f . unsafeKey) m + +-- | Map each key-value pair of a map into a monoid (with proof that the key was +-- in the map), and combine the results using '<>'. +foldMapWithKey + :: forall s a m. Monoid m => (Key s -> a -> m) -> IntMap s a -> m +foldMapWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.foldMapWithKey @m @a + +-- | Right associative fold with a lazy accumulator. +foldrWithKey :: forall s a b. (Key s -> a -> b -> b) -> b -> IntMap s a -> b +foldrWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.foldrWithKey @a @b + +-- | Left associative fold with a lazy accumulator. +foldlWithKey :: forall s a b. (b -> Key s -> a -> b) -> b -> IntMap s a -> b +foldlWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.foldlWithKey @b @a + +-- | Right associative fold with a strict accumulator. +foldrWithKey' :: forall s a b. (Key s -> a -> b -> b) -> b -> IntMap s a -> b +foldrWithKey' = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.foldrWithKey' @a @b + +-- | Left associative fold with a strict accumulator. +foldlWithKey' :: forall s a b. (b -> Key s -> a -> b) -> b -> IntMap s a -> b +foldlWithKey' = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.foldlWithKey' @b @a + +-- | Return the set of keys in the map, with the contents of the set still +-- tracked by the @s@ parameter. See "Data.IntSet.Refined". +keysSet :: forall s a. IntMap s a -> IntSet s +keysSet (IntMap m) = reify (IntMap.keysSet m) + \(_ :: Proxy s') -> case unsafeCoerce Refl :: s :~: s' of + Refl -> Dict + +-- | Convert to a list of key-value pairs in ascending order of keys. +toList :: forall s a. IntMap s a -> [(Key s, a)] +toList = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.toAscList @a + +-- | Convert to a list of key-value pairs in descending order of keys. +toDescList :: forall s a. IntMap s a -> [(Key s, a)] +toDescList = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.toDescList @a + +-- | Retain only the key-value pairs that satisfy the predicate, returning a +-- potentially smaller map. +filterWithKey + :: forall s a. (Key s -> a -> Bool) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +filterWithKey p (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.filterWithKey (p . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Restrict a map to only those keys that are elements of @t@. +restrictKeys + :: forall s t a. KnownIntSet t + => IntMap s a -> SomeIntMapWith (IntersectionProof 'Int s t) a +restrictKeys (IntMap m) = SomeIntMapWith +#if MIN_VERSION_containers(0, 5, 8) + (IntMap $ IntMap.restrictKeys m $ reflect $ Proxy @t) +#else + (IntMap $ IntMap.intersectionWith const m + $ IntMap.fromSet id $ reflect $ Proxy @t) +#endif + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Remove all keys that are elements of @t@ from the map. +withoutKeys + :: forall s t a. KnownIntSet t + => IntMap s a -> SomeIntMapWith (DifferenceProof 'Int s t) a +withoutKeys (IntMap m) = SomeIntMapWith +#if MIN_VERSION_containers(0, 5, 8) + (IntMap $ IntMap.withoutKeys m $ reflect $ Proxy @t) +#else + (IntMap $ IntMap.difference m $ IntMap.fromSet id $ reflect $ Proxy @t) +#endif + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Partition a map into two disjoint submaps: those whose key-value pairs +-- satisfy the predicate, and those whose don't. +partitionWithKey + :: forall s a. (Key s -> a -> Bool) + -> IntMap s a + -> Some2IntMapWith (PartitionProof 'Int s Int) a a +partitionWithKey p (IntMap m) + = case IntMap.partitionWithKey (p . unsafeKey) m of + (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof + do \k -> case IntMap.lookup (unrefine k) m of + Nothing -> error + "partitionWithKey: bug: Data.IntMap.Refined has been subverted" + Just x -> if p k x + then Left $ unsafeKey $ unrefine k + else Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Divide a map into two disjoint submaps at a point where the predicate on +-- the keys stops holding. +-- +-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then +-- this point is uniquely defined. If @p@ is not antitone, a splitting point is +-- chosen in an unspecified way. +spanAntitone + :: forall s a. (Key s -> Bool) + -> IntMap s a + -> Some2IntMapWith (PartialPartitionProof 'Int s) a a +spanAntitone p (IntMap m) = +#if MIN_VERSION_containers(0, 6, 7) + case IntMap.spanAntitone (p . unsafeKey) m of + (m1, m2) +#else + case List.span (p . unsafeKey . fst) $ IntMap.toAscList m of + (xs1, xs2) + | let m1 = IntMap.fromDistinctAscList xs1 + , let m2 = IntMap.fromDistinctAscList xs2 +#endif + -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartialPartitionProof + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Return two disjoint submaps: those whose keys are less than the given key, +-- and those whose keys are greater than the given key. If the key was in the +-- map, also return the associated value and the proof that it was in the map. +splitLookup + :: forall s a. Int + -> IntMap s a + -> Some2IntMapWith (SplitProof 'Int s (Key s, a)) a a +splitLookup k (IntMap m) = case IntMap.splitLookup k m of + (m1, v, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ SplitProof + ((unsafeKey k,) <$> v) unsafeSubset \f g -> unsafeSubsetWith2 f g + +-- | Retrieves the key-value pair corresponding to the smallest key of the map, +-- and the map with that pair removed; or a proof that the map was empty. +minViewWithKey + :: forall s a. IntMap s a + -> Either + (EmptyProof 'Int s) + ((Key s, a), SomeIntMapWith (SupersetProof 'Int s) a) +minViewWithKey (IntMap m) = case IntMap.minViewWithKey m of + Nothing -> Left $ EmptyProof unsafeSubset + Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s) $ coerce kv,) + $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset + +-- | Retrieves the key-value pair corresponding to the greatest key of the map, +-- and the map with that pair removed; or a proof that the map was empty. +maxViewWithKey + :: forall s a. IntMap s a + -> Either + (EmptyProof 'Int s) + ((Key s, a), SomeIntMapWith (SupersetProof 'Int s) a) +maxViewWithKey (IntMap m) = case IntMap.maxViewWithKey m of + Nothing -> Left $ EmptyProof unsafeSubset + Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s) $ coerce kv,) + $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset + +-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then +-- @s@ and @t@ actually stand for the same set and @'Key' s@ can be safely +-- interconverted with @'Key' t@. +-- +-- The requirement that the weakenings are natural transformations ensures that +-- they don't actually alter the keys. To build these you can compose ':->''s +-- from proofs returned by functions in this module, or "Refined" functions like +-- 'andLeft' or 'leftOr'. +castKey + :: forall s t k. + (forall x. Refined (InSet 'Int s) x -> Refined (InSet 'Int t) x) + -> (forall x. Refined (InSet 'Int t) x -> Refined (InSet 'Int s) x) + -> Coercion (Refined (InSet 'Int s) k) (Refined (InSet 'Int t) k) +castKey = castRefined + +-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps +-- can be interconverted too. For example, 'zipWithKey' can be implemented via +-- 'Data.IntMap.Refined.intersectionWithKey' by proving that the set of keys +-- remains unchanged: +-- +-- @ +-- 'zipWithKey' +-- :: forall s a b c. ('Key' s -> a -> b -> c) +-- -> 'IntMap' s a +-- -> 'IntMap' s b +-- -> 'IntMap' s c +-- 'zipWithKey' f m1 m2 +-- | v'SomeIntMapWith' @r m proof <- 'Data.IntMap.Refined.intersectionWithKey' (f . 'andLeft') m1 m2 +-- , v'IntersectionProof' p1 p2 <- proof +-- , ( v'Coercion' :: t'Coercion' ('IntMap' r c) ('IntMap' s c)) +-- <- app $ 'cast' $ 'castKey' ('andLeft' . p1) (p2 'id' 'id') +-- = 'coerce' m +-- where +-- app :: t'Coercion' f g -> t'Coercion' (f x) (g x) +-- app v'Coercion' = v'Coercion' +-- @ +cast + :: forall s t k. (forall x. Coercion + (Refined (InSet 'Int s) x) + (Refined (InSet 'Int t) x)) + -> Coercion (IntMap s k) (IntMap t k) +cast Coercion = Coercion + +instance FunctorWithIndex (Key s) (IntMap s) where + imap = mapWithKey + +instance FoldableWithIndex (Key s) (IntMap s) where + ifoldMap = foldMapWithKey + +instance TraversableWithIndex (Key s) (IntMap s) where + itraverse = traverseWithKey + +-- | Similar to the instance for functions -- zip corresponding keys. To use +-- '<*>'/'liftA2' without 'KnownIntSet' see 'zipWithKey'. +instance KnownIntSet s => Applicative (IntMap s) where + pure x = fromSet \_ -> x + (<*>) = zipWithKey (const id) + +-- | @'bind' m f@ is a map that for each key @k :: 'Key' s@, contains the +-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions. +bind :: forall s a b. IntMap s a -> (a -> IntMap s b) -> IntMap s b +bind m f = mapWithKey (\k x -> f x ! k) m + +-- | Similar to the instance for functions. To use '>>=' without 'KnownIntSet' +-- see 'bind'. +instance KnownIntSet s => Monad (IntMap s) where + (>>=) = bind + +-- | Similar to the instance for functions. See also +-- 'Data.IntMap.Refined.backpermuteKeys'. +instance KnownIntSet s => MonadReader (Key s) (IntMap s) where + ask = fromSet id + local f m = mapWithKey (\k _ -> m ! f k) m + +-- | Append the values at the corresponding keys +instance Semigroup a => Semigroup (IntMap s a) where + (<>) = zipWithKey (const (<>)) + +instance (KnownIntSet s, Monoid a) => Monoid (IntMap s a) where + mempty = fromSet \_ -> mempty + +-- | Similar to the instance for functions +instance KnownIntSet s => Distributive (IntMap s) where + collect = collectRep + distribute = distributeRep + +-- | Witness isomorphism with functions from @'Key' s@ +instance KnownIntSet s => Representable (IntMap s) where + type Rep (IntMap s) = Key s + index = (!) + tabulate = fromSet + +#if MIN_VERSION_hashable(1, 3, 4) +#else +instance Hashable.Hashable a => Hashable.Hashable (IntMap s a) where + hashWithSalt s (IntMap m) = IntMap.foldlWithKey' + (\s' k v -> Hashable.hashWithSalt (Hashable.hashWithSalt s' k) v) + (Hashable.hashWithSalt s (IntMap.size m)) + m +#endif diff --git a/src/Data/IntMap/Refined.hs b/src/Data/IntMap/Refined.hs new file mode 100644 index 0000000..93aa4c9 --- /dev/null +++ b/src/Data/IntMap/Refined.hs @@ -0,0 +1,379 @@ +-- | This module defines a way to prove that a key exists in a map, so that the +-- key can be used to index into the map without using a 'Maybe', or manually +-- handling the \"impossible\" case with 'error' or other partial functions. +-- +-- To do this, @'IntMap' s v@ has a type parameter @s@ that identifies its set +-- ofvkeys, so that if another map has the same type parameter, you know that +-- map has the same set of keys. There is @'Key' s@, a type of keys that have +-- been validated to belong to the set identified by @s@, and for which the +-- operation of indexing into a @'IntMap' s v@ (only for the same @s@) can +-- proceed without failure (see '!'). The type @s@ itself has no internal +-- structure, rather it is merely a skolem type variable (rank-2 polymorphism +-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection". +-- +-- Like "Data.IntMap", functions in this module are strict in the keys but lazy +-- in the values. The "Data.IntMap.Strict.Refined" module reuses the same +-- 'IntMap' type but provides functions that operate strictly on the values. +module Data.IntMap.Refined + ( + -- * Map type + Common.IntMap + , Common.Key + -- * Existentials and common proofs + , Common.SomeIntMap(..) + , Common.withIntMap + , Common.SomeIntMapWith(..) + , Common.withIntMapWith + , Common.Some2IntMapWith(..) + , Common.with2IntMapWith + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , Common.empty + , singleton + , SingletonProof(..) + , fromSet + , Common.fromIntMap + , fromTraversableWithKey + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + , reinsert + , insertLookupWithKey + -- * Deletion/Update + , Common.delete + , adjust + , adjustWithKey + , update + , updateLookupWithKey + -- * Query + , Common.lookup + , (Common.!) + , Common.member + , Common.lookupLT + , Common.lookupGT + , Common.lookupLE + , Common.lookupGE + , Common.null + , Common.isSubmapOfBy + , SubsetProof(..) + , Common.disjoint + , DisjointProof(..) + -- * Combine + , zipWithKey + , bind + , unionWithKey + , UnionProof(..) + , Common.difference + , DifferenceProof(..) + , differenceWithKey + , PartialDifferenceProof(..) + , intersectionWithKey + , IntersectionProof(..) + -- * Traversal + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , MapProof(..) + , backpermuteKeys + -- * Folds + , Common.foldMapWithKey + , Common.foldrWithKey + , Common.foldlWithKey + , Common.foldrWithKey' + , Common.foldlWithKey' + -- * Conversion + , Common.toIntMap + , Common.keysSet + , Common.toList + , Common.toDescList + -- * Filter + , Common.restrictKeys + , Common.withoutKeys + , Common.filterWithKey + , Common.partitionWithKey + , PartitionProof(..) + , Common.spanAntitone + , PartialPartitionProof(..) + , mapMaybeWithKey + , mapEitherWithKey + , Common.splitLookup + , SplitProof(..) + -- * Min/Max + , updateMinWithKey + , updateMaxWithKey + , adjustMinWithKey + , adjustMaxWithKey + , Common.minViewWithKey + , Common.maxViewWithKey + -- * Casts + , Common.castKey + , Common.cast + , castFlavor + ) where + +import Data.Coerce +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Functor +import qualified Data.IntMap as IntMap +import Data.IntMap.Common.Refined + ( IntMap(..), Key, unsafeCastKey, unsafeKey, SomeIntMapWith(..) + , Some2IntMapWith(..), fromSet, (!), zipWithKey, mapWithKey, traverseWithKey + , bind + ) +import qualified Data.IntMap.Common.Refined as Common +import Data.Traversable +import Data.Type.Coercion +import Prelude hiding (lookup, null) +import Refined +import Refined.Unsafe + + +-- | Create a map with a single key-value pair, and return a proof that the +-- key is in the resulting map. +singleton :: forall a. Int -> a -> SomeIntMapWith (SingletonProof 'Int Int) a +singleton k v = SomeIntMapWith (IntMap $ IntMap.singleton k v) + $ SingletonProof (unsafeKey k) + +-- | Create a map from an arbitrary traversable of key-value pairs. +fromTraversableWithKey + :: forall t a. Traversable t + => (Int -> a -> a -> a) + -> t (Int, a) + -> SomeIntMapWith (FromTraversableProof 'Int t Int) a +fromTraversableWithKey f xs + = SomeIntMapWith (IntMap m) $ FromTraversableProof proof + where + (m, proof) = mapAccumL + (\s (k, v) -> let !s' = IntMap.insertWithKey f k v s in (s', unsafeKey k)) + IntMap.empty + xs + +-- | Insert a key-value pair into the map to obtain a potentially larger map, +-- guaranteed to contain the given key. If the key was already present, the +-- associated value is replaced with the supplied value. +insert + :: forall s a. Int + -> a + -> IntMap s a + -> SomeIntMapWith (InsertProof 'Int Int s) a +insert k v (IntMap m) = SomeIntMapWith (IntMap $ IntMap.insert k v m) + $ InsertProof (unsafeKey k) unsafeSubset + +-- | Overwrite a key-value pair that is known to already be in the map. The set +-- of keys remains the same. +reinsert + :: forall s a. Key s -> a -> IntMap s a -> IntMap s a +reinsert = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.insert @a + +-- | Insert a key-value pair into the map using a combining function, and if +-- the key was already present, the old value is returned along with the proof +-- that the key was present. +insertLookupWithKey + :: forall s a. (Key s -> a -> a -> a) + -> Int + -> a + -> IntMap s a + -> (Maybe (Key s, a), SomeIntMapWith (InsertProof 'Int Int s) a) +insertLookupWithKey f k v (IntMap m) + = case IntMap.insertLookupWithKey (f . unsafeKey) k v m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeIntMapWith (IntMap m') $ InsertProof (unsafeKey k) unsafeSubset + +-- | Update the value at a specific key known the be in the map using the given +-- function. The set of keys remains the same. +adjust :: forall s a. (a -> a) -> Key s -> IntMap s a -> IntMap s a +adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.adjust @a + +-- | If the given key is in the map, update the associated value using the given +-- function with a proof that the key was in the map; otherwise return the map +-- unchanged. In any case the set of keys remains the same. +adjustWithKey + :: forall s a. (Key s -> a -> a) -> Int -> IntMap s a -> IntMap s a +adjustWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.adjustWithKey @a + +-- | Update or delete a key known to be in the map using the given function, +-- returning a potentially smaller map. +update + :: forall s a. (a -> Maybe a) + -> Key s + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +update f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f (unrefine k) m) + $ SupersetProof unsafeSubset + +-- | If the given key is in the map, update or delete it using the given +-- function with a proof that the key was in the map; otherwise the map is +-- unchanged. Alongside return the new value if it was updated, or the old value +-- if it was deleted, and a proof that the key was in the map. +updateLookupWithKey + :: forall s a. (Key s -> a -> Maybe a) + -> Int + -> IntMap s a + -> (Maybe (Key s, a), SomeIntMapWith (SupersetProof 'Int s) a) +updateLookupWithKey f k (IntMap m) + = case IntMap.updateLookupWithKey (f . unsafeKey) k m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset + +-- | Return the union of two maps, with a given combining function for keys that +-- exist in both maps simultaneously. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@ +-- respectively. +unionWithKey + :: forall s t a. (Refined (InSet 'Int s && InSet 'Int t) Int -> a -> a -> a) + -> IntMap s a + -> IntMap t a + -> SomeIntMapWith (UnionProof 'Int s t) a +unionWithKey f (IntMap m1) (IntMap m2) + = SomeIntMapWith (IntMap $ IntMap.unionWithKey (f . reallyUnsafeRefine) m1 m2) + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- | For keys that appear in both maps, the given function decides whether the +-- key is removed from the first map. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@ +-- respectively. +differenceWithKey + :: forall s t a b. (Refined (InSet 'Int s && InSet 'Int t) Int + -> a + -> b + -> Maybe a) + -> IntMap s a + -> IntMap t b + -> SomeIntMapWith (PartialDifferenceProof 'Int s t) a +differenceWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith + (IntMap $ IntMap.differenceWithKey (f . reallyUnsafeRefine) m1 m2) + $ PartialDifferenceProof unsafeSubset unsafeSubset + +-- | Return the intersection of two maps with the given combining function. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@ +-- respectively. +intersectionWithKey + :: forall s t a b c. (Refined (InSet 'Int s && InSet 'Int t) Int + -> a + -> b + -> c) + -> IntMap s a + -> IntMap t b + -> SomeIntMapWith (IntersectionProof 'Int s t) c +intersectionWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith + (IntMap $ IntMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Thread an accumularing argument through the map in ascending order of keys. +mapAccumLWithKey + :: forall s a b c. (a -> Key s -> b -> (a, c)) + -> a + -> IntMap s b + -> (a, IntMap s c) +mapAccumLWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.mapAccumWithKey @a @b @c + +-- | Thread an accumularing argument through the map in descending order of +-- keys. +mapAccumRWithKey + :: forall s a b c. (a -> Key s -> b -> (a, c)) + -> a + -> IntMap s b + -> (a, IntMap s c) +mapAccumRWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.mapAccumRWithKey @a @b @c + +-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the +-- results into a new map. For keys that were mapped to the same new key, @c@ +-- acts as the combining function for corresponding values. +mapKeysWith + :: forall s a. (a -> a -> a) + -> (Key s -> Int) + -> IntMap s a + -> SomeIntMapWith (MapProof 'Int s Int Int) a +mapKeysWith f g (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.mapKeysWith f (g . unsafeKey) m) + $ MapProof (unsafeKey . g) \k2 -> + case IntMap.lookup (unrefine k2) backMap of + Nothing -> error + "mapKeysWith: bug: Data.IntMap.Refined has been subverted" + Just k1 -> k1 + where + ~backMap = IntMap.fromList + [ (k2, unsafeKey k1) + | k1 <- IntMap.keys m + , let !k2 = g $ unsafeKey k1 + ] + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect only the 'Just' results, returning a potentially smaller +-- map. +mapMaybeWithKey + :: forall s a b. (Key s -> a -> Maybe b) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) b +mapMaybeWithKey f (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.mapMaybeWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect the 'Left' and 'Right' results into separate (disjoint) +-- maps. +mapEitherWithKey + :: forall s a b c. (Key s -> a -> Either b c) + -> IntMap s a + -> Some2IntMapWith (PartitionProof 'Int s Int) b c +mapEitherWithKey p (IntMap m) + = case IntMap.mapEitherWithKey (p . unsafeKey) m of + (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof + do \k -> case IntMap.lookup (unrefine k) m of + Nothing -> error + "mapEitherWithKey: bug: Data.IntMap.Refined has been subverted" + Just x -> case p k x of + Left _ -> Left $ unsafeKey $ unrefine k + Right _ -> Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Update or delete the value at the smallest key, returning a potentially +-- smaller map. +updateMinWithKey + :: forall s a. (Key s -> a -> Maybe a) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +updateMinWithKey f (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.updateMinWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Update or delete the value at the largest key, returning a potentially +-- smaller map. +updateMaxWithKey + :: forall s a. (Key s -> a -> Maybe a) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +updateMaxWithKey f (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.updateMaxWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Adjust the value at the smallest key. The set of keys remains unchanged. +adjustMinWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a +adjustMinWithKey f (IntMap m) + = IntMap $ IntMap.updateMinWithKey ((Just .) . f . unsafeKey) m + +-- | Adjust the value at the greatest key. The set of keys remains unchanged. +adjustMaxWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a +adjustMaxWithKey f (IntMap m) + = IntMap $ IntMap.updateMaxWithKey ((Just .) . f . unsafeKey) m + +-- | Apply the inverse image of the given function to the keys of the given map, +-- so that for all @k :: 'Key' s2@, +-- @'backpermuteKeys' f m '!' k = m '!' f k@. +-- +-- If maps are identified with functions, this computes the composition. +backpermuteKeys + :: forall s1 s2 a. KnownIntSet s2 + => (Key s2 -> Key s1) -> IntMap s1 a -> IntMap s2 a +backpermuteKeys f m = fromSet \k -> m ! f k diff --git a/src/Data/IntMap/Strict/Refined.hs b/src/Data/IntMap/Strict/Refined.hs new file mode 100644 index 0000000..dcd2f9e --- /dev/null +++ b/src/Data/IntMap/Strict/Refined.hs @@ -0,0 +1,418 @@ +-- | This module defines a way to prove that a key exists in a map, so that the +-- key can be used to index into the map without using a 'Maybe', or manually +-- handling the \"impossible\" case with 'error' or other partial functions. +-- +-- To do this, @'IntMap' s v@ has a type parameter @s@ that identifies its set +-- of keys, so that if another map has the same type parameter, you know that +-- map has the same set of keys. There is @'Key' s@, a type of keys that have +-- been validated to belong to the set identified by @s@, and for which the +-- operation of indexing into a @'IntMap' s v@ (only for the same @s@) can +-- proceed without failure (see '!'). The type @s@ itself has no internal +-- structure, rather it is merely a skolem type variable (rank-2 polymorphism +-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection". +-- +-- Like "Data.IntMap.Strict", functions in this module are strict in the keys +-- and values. The "Data.IntMap.Refined" module reuses the same 'IntMap' type +-- but provides functions that operate lazily on the values. +module Data.IntMap.Strict.Refined + ( + -- * Map type + Common.IntMap + , Common.Key + -- * Existentials and common proofs + , Common.SomeIntMap(..) + , Common.withIntMap + , Common.SomeIntMapWith(..) + , Common.withIntMapWith + , Common.Some2IntMapWith(..) + , Common.with2IntMapWith + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , Common.empty + , singleton + , SingletonProof(..) + , fromSet + , Common.fromIntMap + , fromTraversableWithKey + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + , reinsert + , insertLookupWithKey + -- * Deletion/Update + , Common.delete + , adjust + , adjustWithKey + , update + , updateLookupWithKey + -- * Query + , Common.lookup + , (Common.!) + , Common.member + , Common.lookupLT + , Common.lookupGT + , Common.lookupLE + , Common.lookupGE + , Common.null + , Common.isSubmapOfBy + , SubsetProof(..) + , Common.disjoint + , DisjointProof(..) + -- * Combine + , zipWithKey + , bind + , unionWithKey + , UnionProof(..) + , Common.difference + , DifferenceProof(..) + , differenceWithKey + , PartialDifferenceProof(..) + , intersectionWithKey + , IntersectionProof(..) + -- * Traversal + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , MapProof(..) + , backpermuteKeys + -- * Folds + , Common.foldMapWithKey + , Common.foldrWithKey + , Common.foldlWithKey + , Common.foldrWithKey' + , Common.foldlWithKey' + -- * Conversion + , Common.toIntMap + , Common.keysSet + , Common.toList + , Common.toDescList + -- * Filter + , Common.restrictKeys + , Common.withoutKeys + , Common.filterWithKey + , Common.partitionWithKey + , PartitionProof(..) + , Common.spanAntitone + , PartialPartitionProof(..) + , mapMaybeWithKey + , mapEitherWithKey + , Common.splitLookup + , SplitProof(..) + -- * Min/Max + , updateMinWithKey + , updateMaxWithKey + , adjustMinWithKey + , adjustMaxWithKey + , Common.minViewWithKey + , Common.maxViewWithKey + -- * Casts + , Common.castKey + , Common.cast + , castFlavor + ) where + +import Data.Coerce +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Functor +import qualified Data.IntMap.Strict as IntMap +import Data.IntMap.Common.Refined + ( IntMap(..), Key, unsafeCastKey, unsafeKey, SomeIntMapWith(..) + , Some2IntMapWith(..), (!) + ) +import qualified Data.IntMap.Common.Refined as Common +import Data.Proxy +import Data.Reflection +import Data.Traversable +import Data.Type.Coercion +import Prelude hiding (lookup, null) +import Refined +import Refined.Unsafe + + +-- | Create a map with a single key-value pair, and return a proof that the +-- key is in the resulting map. +singleton :: forall a. Int -> a -> SomeIntMapWith (SingletonProof 'Int Int) a +singleton k v = SomeIntMapWith (IntMap $ IntMap.singleton k v) + $ SingletonProof (unsafeKey k) + +-- | Create a map from a set of keys, and a function that for each key computes +-- the corresponding value. +fromSet :: forall s a. KnownIntSet s => (Key s -> a) -> IntMap s a +fromSet f = IntMap $ IntMap.fromSet (f . unsafeKey) (reflect $ Proxy @s) + +-- | Create a map from an arbitrary traversable of key-value pairs. +fromTraversableWithKey + :: forall t a. Traversable t + => (Int -> a -> a -> a) + -> t (Int, a) + -> SomeIntMapWith (FromTraversableProof 'Int t Int) a +fromTraversableWithKey f xs = SomeIntMapWith (IntMap m) + $ FromTraversableProof proof + where + (m, proof) = mapAccumL + (\s (k, v) -> (IntMap.insertWithKey f k v s, unsafeKey k)) + IntMap.empty + xs + +-- | Insert a key-value pair into the map to obtain a potentially larger map, +-- guaranteed to contain the given key. If the key was already present, the +-- associated value is replaced with the supplied value. +insert + :: forall s a. Int + -> a + -> IntMap s a + -> SomeIntMapWith (InsertProof 'Int Int s) a +insert k v (IntMap m) = SomeIntMapWith (IntMap $ IntMap.insert k v m) + $ InsertProof (unsafeKey k) unsafeSubset + +-- | Overwrite a key-value pair that is known to already be in the map. The set +-- of keys remains the same. +reinsert + :: forall s a. Key s -> a -> IntMap s a -> IntMap s a +reinsert = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.insert @a + +-- | Insert a key-value pair into the map using a combining function, and if +-- the key was already present, the old value is returned along with the proof +-- that the key was present. +insertLookupWithKey + :: forall s a. (Key s -> a -> a -> a) + -> Int + -> a + -> IntMap s a + -> (Maybe (Key s, a), SomeIntMapWith (InsertProof 'Int Int s) a) +insertLookupWithKey f k v (IntMap m) + = case IntMap.insertLookupWithKey (f . unsafeKey) k v m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeIntMapWith (IntMap m') $ InsertProof (unsafeKey k) unsafeSubset + +-- | Update the value at a specific key known the be in the map using the given +-- function. The set of keys remains the same. +adjust :: forall s a. (a -> a) -> Key s -> IntMap s a -> IntMap s a +adjust = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.adjust @a + +-- | If the given key is in the map, update the associated value using the given +-- function with a proof that the key was in the map; otherwise return the map +-- unchanged. In any case the set of keys remains the same. +adjustWithKey + :: forall s a. (Key s -> a -> a) -> Int -> IntMap s a -> IntMap s a +adjustWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.adjustWithKey @a + +-- | Update or delete a key known to be in the map using the given function, +-- returning a potentially smaller map. +update + :: forall s a. (a -> Maybe a) + -> Key s + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +update f k (IntMap m) = SomeIntMapWith (IntMap $ IntMap.update f (unrefine k) m) + $ SupersetProof unsafeSubset + +-- | If the given key is in the map, update or delete it using the given +-- function with a proof that the key was in the map; otherwise the map is +-- unchanged. Alongside return the new value if it was updated, or the old value +-- if it was deleted, and a proof that the key was in the map. +updateLookupWithKey + :: forall s a. (Key s -> a -> Maybe a) + -> Int + -> IntMap s a + -> (Maybe (Key s, a), SomeIntMapWith (SupersetProof 'Int s) a) +updateLookupWithKey f k (IntMap m) + = case IntMap.updateLookupWithKey (f . unsafeKey) k m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeIntMapWith (IntMap m') $ SupersetProof unsafeSubset + +-- | Given two maps proven to have the same keys, for each key apply the +-- function to the associated values, to obtain a new map with the same keys. +zipWithKey + :: forall s a b c. (Key s -> a -> b -> c) + -> IntMap s a + -> IntMap s b + -> IntMap s c +zipWithKey f (IntMap m1) (IntMap m2) = IntMap + $ IntMap.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y) + (error "zipWithKey: bug: Data.IntMap.Strict.Refined has been subverted") + (error "zipWithKey: bug: Data.IntMap.Strict.Refined has been subverted") + m1 + m2 + +-- | Return the union of two maps, with a given combining function for keys that +-- exist in both maps simultaneously. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@ +-- respectively. +unionWithKey + :: forall s t a. (Refined (InSet 'Int s && InSet 'Int t) Int -> a -> a -> a) + -> IntMap s a + -> IntMap t a + -> SomeIntMapWith (UnionProof 'Int s t) a +unionWithKey f (IntMap m1) (IntMap m2) + = SomeIntMapWith (IntMap $ IntMap.unionWithKey (f . reallyUnsafeRefine) m1 m2) + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- | For keys that appear in both maps, the given function decides whether the +-- key is removed from the first map. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@ +-- respectively. +differenceWithKey + :: forall s t a b. (Refined (InSet 'Int s && InSet 'Int t) Int + -> a + -> b + -> Maybe a) + -> IntMap s a + -> IntMap t b + -> SomeIntMapWith (PartialDifferenceProof 'Int s t) a +differenceWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith + (IntMap $ IntMap.differenceWithKey (f . reallyUnsafeRefine) m1 m2) + $ PartialDifferenceProof unsafeSubset unsafeSubset + +-- | Return the intersection of two maps with the given combining function. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s@ and @'Key' t@ +-- respectively. +intersectionWithKey + :: forall s t a b c. (Refined (InSet 'Int s && InSet 'Int t) Int + -> a + -> b + -> c) + -> IntMap s a + -> IntMap t b + -> SomeIntMapWith (IntersectionProof 'Int s t) c +intersectionWithKey f (IntMap m1) (IntMap m2) = SomeIntMapWith + (IntMap $ IntMap.intersectionWithKey (f . reallyUnsafeRefine) m1 m2) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, that are proven to be in the map. The set of keys remains the same. +mapWithKey :: forall s a b. (Key s -> a -> b) -> IntMap s a -> IntMap s b +mapWithKey = gcoerceWith (unsafeCastKey @s) $ coerce $ IntMap.mapWithKey @a @b + +-- | Map an 'Applicative' transformation in ascending order of keys, with access +-- to each value's corresponding key and a proof that it is in the map. The set +-- of keys remains unchanged. +traverseWithKey + :: forall s f a b. Applicative f + => (Key s -> a -> f b) -> IntMap s a -> f (IntMap s b) +traverseWithKey f (IntMap m) + = IntMap <$> IntMap.traverseWithKey (f . unsafeKey) m + +-- | Thread an accumularing argument through the map in ascending order of keys. +mapAccumLWithKey + :: forall s a b c. (a -> Key s -> b -> (a, c)) + -> a + -> IntMap s b + -> (a, IntMap s c) +mapAccumLWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.mapAccumWithKey @a @b @c + +-- | Thread an accumularing argument through the map in descending order of +-- keys. +mapAccumRWithKey + :: forall s a b c. (a -> Key s -> b -> (a, c)) + -> a + -> IntMap s b + -> (a, IntMap s c) +mapAccumRWithKey = gcoerceWith (unsafeCastKey @s) $ coerce + $ IntMap.mapAccumRWithKey @a @b @c + +-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the +-- results into a new map. For keys that were mapped to the same new key, @c@ +-- acts as the combining function for corresponding values. +mapKeysWith + :: forall s a. (a -> a -> a) + -> (Key s -> Int) + -> IntMap s a + -> SomeIntMapWith (MapProof 'Int s Int Int) a +mapKeysWith f g (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.mapKeysWith f (g . unsafeKey) m) + $ MapProof (unsafeKey . g) \k2 -> + case IntMap.lookup (unrefine k2) backMap of + Nothing -> error + "mapKeysWith: bug: Data.IntMap.Strict.Refined has been subverted" + Just k1 -> k1 + where + ~backMap = IntMap.fromList + [ (k2, unsafeKey k1) + | k1 <- IntMap.keys m + , let !k2 = g $ unsafeKey k1 + ] + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect only the 'Just' results, returning a potentially smaller +-- map. +mapMaybeWithKey + :: forall s a b. (Key s -> a -> Maybe b) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) b +mapMaybeWithKey f (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.mapMaybeWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect the 'Left' and 'Right' results into separate (disjoint) +-- maps. +mapEitherWithKey + :: forall s a b c. (Key s -> a -> Either b c) + -> IntMap s a + -> Some2IntMapWith (PartitionProof 'Int s Int) b c +mapEitherWithKey p (IntMap m) + = case IntMap.mapEitherWithKey (p . unsafeKey) m of + (m1, m2) -> Some2IntMapWith (IntMap m1) (IntMap m2) $ PartitionProof + do \k -> case IntMap.lookup (unrefine k) m of + Nothing -> error + "mapEitherWithKey: bug: Data.IntMap.Strict.Refined has been subverted" + Just x -> case p k x of + Left _ -> Left $ unsafeKey $ unrefine k + Right _ -> Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Update or delete the value at the smallest key, returning a potentially +-- smaller map. +updateMinWithKey + :: forall s a. (Key s -> a -> Maybe a) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +updateMinWithKey f (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.updateMinWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Update or delete the value at the largest key, returning a potentially +-- smaller map. +updateMaxWithKey + :: forall s a. (Key s -> a -> Maybe a) + -> IntMap s a + -> SomeIntMapWith (SupersetProof 'Int s) a +updateMaxWithKey f (IntMap m) + = SomeIntMapWith (IntMap $ IntMap.updateMaxWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Adjust the value at the smallest key. The set of keys remains unchanged. +adjustMinWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a +adjustMinWithKey f (IntMap m) + = IntMap $ IntMap.updateMinWithKey ((Just .) . f . unsafeKey) m + +-- | Adjust the value at the greatest key. The set of keys remains unchanged. +adjustMaxWithKey :: forall s a. (Key s -> a -> a) -> IntMap s a -> IntMap s a +adjustMaxWithKey f (IntMap m) + = IntMap $ IntMap.updateMaxWithKey ((Just .) . f . unsafeKey) m + +-- | @'bind' m f@ is a map that for each key @k :: 'Key' s@, contains the +-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions. +bind :: forall s a b. IntMap s a -> (a -> IntMap s b) -> IntMap s b +bind m f = mapWithKey (\k x -> f x ! k) m + +-- | Apply the inverse image of the given function to the keys of the given map, +-- so that for all @k :: 'Key' s2@, +-- @'backpermuteKeys' f m '!' k = m '!' f k@. +-- +-- If maps are identified with functions, this computes the composition. +backpermuteKeys + :: forall s1 s2 a. KnownIntSet s2 + => (Key s2 -> Key s1) -> IntMap s1 a -> IntMap s2 a +backpermuteKeys f m = fromSet \k -> m ! f k diff --git a/src/Data/IntSet/Refined.hs b/src/Data/IntSet/Refined.hs new file mode 100644 index 0000000..69f5ad2 --- /dev/null +++ b/src/Data/IntSet/Refined.hs @@ -0,0 +1,502 @@ +{-# LANGUAGE CPP #-} +-- | This module implements a way of tracking the contents of an +-- 'Data.IntSet.IntSet' at the type level, and provides utilities for +-- manipulating such sets. +-- +-- The contents of a set are associated with a type parameter, e.g. @s@, so that +-- whenever you see the same type parameter, you know you are working with the +-- same set. The type @s@ itself has no internal structure, rather it is merely +-- a skolem type variable (rank-2 polymorphism 'Control.Monad.ST.runST' trick) +-- introduced by "Data.Reflection". +module Data.IntSet.Refined + ( + -- * Set type + KnownIntSet + , IntSet + -- * Refinement type + , InSet(..) + , Flavor(Int) + , Element + , revealPredicate + -- * Existentials and common proofs + , SomeIntSet(..) + , withIntSet + , SomeIntSetWith(..) + , withIntSetWith + , Some2IntSetWith(..) + , with2IntSetWith + , (:->) + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , empty + , singleton + , SingletonProof(..) + , fromIntSet + , fromTraversable + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + -- * Deletion + , delete + -- * Query + , member + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , null + , isSubsetOf + , SubsetProof(..) + , disjoint + , DisjointProof(..) + -- * Combine + , union + , UnionProof(..) + , difference + , DifferenceProof(..) + , intersection + , IntersectionProof(..) + -- * Filter + , filter + , partition + , PartitionProof(..) + , spanAntitone + , PartialPartitionProof(..) + , splitMember + , SplitProof(..) + -- * Map + , map + , MapProof(..) + -- * Folds + , foldMap + , foldr + , foldl + , foldr' + , foldl' + -- * Min/Max + , minView + , maxView + -- * Conversion + , toList + , toDescList + , asSet + , asHashSet + -- * Casts + , castElement + , cast + , castFlavor + ) where + +import Data.Coerce +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Conversion +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import qualified Data.Foldable as Foldable +import qualified Data.IntMap as IntMap +import qualified Data.IntSet as IntSet +import Data.Proxy +import Data.Reflection +import Data.Traversable +import Data.Type.Coercion +import Data.Type.Equality ((:~:)(..)) +import GHC.Exts (Proxy#, proxy#) +import Prelude hiding (filter, foldl, foldMap, foldr, map, null) +import Refined +import Refined.Unsafe +import Unsafe.Coerce + +#if MIN_VERSION_containers(0, 6, 7) +#else +import qualified Data.List as List +#endif + + +-- | To use "Refined" machinery that uses the 'Predicate' typeclass you will +-- need to pattern match on this 'Dict'. +-- +-- The reason is that in the default /fast/ implementation of reflection, we +-- don't have @'Data.Typeable.Typeable' s@, which "Refined" wants for +-- pretty-printing exceptions. We /can/ provide @'Data.TypeableTypeable' s@, but +-- at the cost of using the /slow/ implementation of reflection. +revealPredicate + :: forall s. KnownIntSet s => Dict (Predicate (InSet 'Int s) Int) +revealPredicate = reifyTypeable (reflect (Proxy @s)) + \(_ :: Proxy s') -> + reflect (Proxy @s') `seq` + -- ^ Work around https://github.com/ekmett/reflection/issues/54 + case unsafeCoerce Refl :: s :~: s' of + Refl -> Dict + +-- | @'Element' s@ is an t'Int' that has been verified to be an element of @s@. +-- +-- Thus, @'Element' s@ is a \"refinement\" type of t'Int', and this library +-- integrates with an implementation of refimenement types in "Refined", so +-- the machinery from there can be used to manipulate 'Element's (however see +-- 'revealPredicate'). +-- +-- The underlying t'Int' can be obtained with 'unrefine'. An t'Int' can be +-- validated into an @'Element' s@ with 'member'. +type Element s = Refined (InSet 'Int s) Int + +unsafeCastElement :: forall s. Coercion Int (Element s) +unsafeCastElement = reallyUnsafeUnderlyingRefined + +unsafeElement :: Int -> Element s +unsafeElement = coerceWith unsafeCastElement + +-- | An existential wrapper for an as-yet-unknown set. Pattern maching on it +-- gives you a way to refer to the set, e.g. +-- +-- @ +-- case 'fromIntSet' ... of +-- 'SomeIntSet' \@s _ -> doSomethingWith \@s +-- +-- case 'fromIntSet' ... of +-- 'SomeIntSet' (_ :: 'Proxy#' s) -> doSomethingWith \@s +-- @ +data SomeIntSet where + SomeIntSet :: forall s. KnownIntSet s => Proxy# s -> SomeIntSet + +-- | Apply an unknown set to a continuation that can accept any set. This gives +-- you a way to refer to the set (the parameter @s@), e.g.: +-- +-- @ +-- 'withIntSet' ('fromIntSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s +-- @ +withIntSet + :: forall r. SomeIntSet -> (forall s. KnownIntSet s => Proxy s -> r) -> r +withIntSet (SomeIntSet (_ :: Proxy# s)) k = k $ Proxy @s + +-- | Construct a set from a regular 'Data.IntSet.IntSet'. +fromIntSet :: IntSet.IntSet -> SomeIntSet +fromIntSet s = reify s \(_ :: Proxy s) -> SomeIntSet @s proxy# + +-- | An existential wrapper for an as-yet-unknown set, together with a proof of +-- some fact @p@ about the set. Pattern matching on it gives you a way to refer +-- to the set (the parameter @s@). Most functions will return a set in this way, +-- together with a proof that somehow relates the set to the function's inputs. +data SomeIntSetWith p where + SomeIntSetWith :: forall s p. KnownIntSet s => !(p s) -> SomeIntSetWith p + +-- | Apply an unknown set with proof to a continuation that can accept any set +-- satisfying the proof. This gives you a way to refer to the set (the parameter +-- @s@). +withIntSetWith + :: forall r p. SomeIntSetWith p -> (forall s. KnownIntSet s => p s -> r) -> r +withIntSetWith (SomeIntSetWith p) k = k p + +-- | An existential wrapper for an as-yet-unknown pair of sets, together with +-- a proof of some fact @p@ relating them. +data Some2IntSetWith p where + Some2IntSetWith + :: forall s t p. (KnownIntSet s, KnownIntSet t) + => !(p s t) -> Some2IntSetWith p + +-- | Apply a pair of unknown sets with proof to a continuation that can accept +-- any pair of sets satisfying the proof. This gives you a way to refer to the +-- sets (the parameters @s@ and @t@). +with2IntSetWith + :: forall r p. Some2IntSetWith p + -> (forall s t. (KnownIntSet s, KnownIntSet t) => p s t -> r) + -> r +with2IntSetWith (Some2IntSetWith p) k = k p + +-- | An empty set. +empty :: SomeIntSetWith (EmptyProof 'Int) +empty = reify IntSet.empty \(_ :: Proxy r) + -> SomeIntSetWith @r $ EmptyProof unsafeSubset + + +-- | Create a set with a single element. +singleton :: Int -> SomeIntSetWith (SingletonProof 'Int Int) +singleton x = reify (IntSet.singleton x) \(_ :: Proxy r) + -> SomeIntSetWith @r $ SingletonProof $ unsafeElement x + +-- | Create a set from the elements of an arbitrary traversable. +fromTraversable + :: forall t. Traversable t + => t Int -> SomeIntSetWith (FromTraversableProof 'Int t Int) +fromTraversable xs = reify set \(_ :: Proxy r) + -> SomeIntSetWith @r $ FromTraversableProof + $ unsafeCoerce @(t (Element _)) @(t (Element r)) proof + where + (set, proof) = mapAccumL + (\s x -> let !s' = IntSet.insert x s in (s', unsafeElement x)) + IntSet.empty + xs + +-- | Insert an element in a set. +insert :: forall s. KnownIntSet s + => Int -> SomeIntSetWith (InsertProof 'Int Int s) +insert x = reify (IntSet.insert x $ reflect $ Proxy @s) \(_ :: Proxy r) + -> SomeIntSetWith @r $ InsertProof (unsafeElement x) unsafeSubset + +-- | Delete an element from a set. +delete :: forall s. KnownIntSet s + => Int -> SomeIntSetWith (SupersetProof 'Int s) +delete x = reify (IntSet.delete x $ reflect $ Proxy @s) \(_ :: Proxy r) + -> SomeIntSetWith @s $ SupersetProof unsafeSubset + +-- | If an element is in the set, return the proof that it is. +member :: forall s. KnownIntSet s => Int -> Maybe (Element s) +member x + | x `IntSet.member` reflect (Proxy @s) = Just $ unsafeElement x + | otherwise = Nothing + +-- | Find the largest element smaller than the given one. +lookupLT :: forall s. KnownIntSet s => Int -> Maybe (Element s) +lookupLT x = gcoerceWith (unsafeCastElement @s) $ coerce + $ IntSet.lookupLT x (reflect $ Proxy @s) + +-- | Find the smallest element greater than the given one. +lookupGT :: forall s. KnownIntSet s => Int -> Maybe (Element s) +lookupGT x = gcoerceWith (unsafeCastElement @s) $ coerce + $ IntSet.lookupGT x (reflect $ Proxy @s) + +-- | Find the largest element smaller or equal to the given one. +lookupLE :: forall s. KnownIntSet s => Int -> Maybe (Element s) +lookupLE x = gcoerceWith (unsafeCastElement @s) $ coerce + $ IntSet.lookupLE x (reflect $ Proxy @s) + +-- | Find the smallest element greater or equal to the given one. +lookupGE :: forall s. KnownIntSet s => Int -> Maybe (Element s) +lookupGE x = gcoerceWith (unsafeCastElement @s) $ coerce + $ IntSet.lookupGE x (reflect $ Proxy @s) + +-- | If the set is empty, return the proof that it is. +null :: forall s. KnownIntSet s => Maybe (EmptyProof 'Int s) +null + | IntSet.null $ reflect $ Proxy @s = Just $ EmptyProof unsafeSubset + | otherwise = Nothing + +-- | If @s@ is a subset of @t@ (or is equal to), return a proof of that. +isSubsetOf + :: forall s t. (KnownIntSet s, KnownIntSet t) => Maybe (SubsetProof 'Int s t) +isSubsetOf + | reflect (Proxy @s) `IntSet.isSubsetOf` reflect (Proxy @t) + = Just $ SubsetProof unsafeSubset + | otherwise = Nothing + +-- | If @s@ and @t@ are disjoint (i.e. their intersection is empty), return a +-- proof of that. +disjoint + :: forall s t. (KnownIntSet s, KnownIntSet t) + => Maybe (DisjointProof 'Int s t) +disjoint +#if MIN_VERSION_containers(0, 5, 11) + | IntSet.disjoint (reflect $ Proxy @s) (reflect $ Proxy @t) +#else + | IntSet.null $ IntSet.intersection (reflect $ Proxy @s) (reflect $ Proxy @t) +#endif + = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g + | otherwise = Nothing + +-- | The union of two sets. +union + :: forall s t. (KnownIntSet s, KnownIntSet t) + => SomeIntSetWith (UnionProof 'Int s t) +union = reify (reflect (Proxy @s) `IntSet.union` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeIntSetWith @r + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- unions :: ? + +-- | Set with elements of @s@ that are not in @t@. +difference + :: forall s t. (KnownIntSet s, KnownIntSet t) + => SomeIntSetWith (DifferenceProof 'Int s t) +difference = reify (reflect (Proxy @s) `IntSet.difference` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeIntSetWith @r + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Intersection of two sets. +intersection + :: forall s t. (KnownIntSet s, KnownIntSet t) + => SomeIntSetWith (IntersectionProof 'Int s t) +intersection + = reify (reflect (Proxy @s) `IntSet.intersection` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeIntSetWith @r + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Return a subset of elements that satisfy the given predicate. +filter + :: forall s. KnownIntSet s + => (Element s -> Bool) -> SomeIntSetWith (SupersetProof 'Int s) +filter p = reify (IntSet.filter (p . unsafeElement) $ reflect $ Proxy @s) + \(_ :: Proxy r) -> SomeIntSetWith @r $ SupersetProof unsafeSubset + +-- | Partition a set into two disjoint subsets: those that satisfy the +-- predicate, and those that don't. +partition + :: forall s. KnownIntSet s + => (Element s -> Bool) -> Some2IntSetWith (PartitionProof 'Int s Int) +partition p = case IntSet.partition (p . unsafeElement) $ reflect $ Proxy @s of + (r, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q) + -> Some2IntSetWith @s @r $ PartitionProof + do \x -> if p x + then Left $ unsafeElement $ unrefine x + else Right $ unsafeElement $ unrefine x + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Divide a set into two disjoint subsets at a point where the predicate stops +-- holding. +-- +-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then +-- this point is uniquely defined. If @p@ is not antitone, a splitting point is +-- chosen in an unspecified way. +spanAntitone + :: forall s. KnownIntSet s + => (Element s -> Bool) -> Some2IntSetWith (PartialPartitionProof 'Int s) +spanAntitone p = +#if MIN_VERSION_containers(0, 6, 7) + case IntSet.spanAntitone (p . unsafeElement) $ reflect $ Proxy @s of + (r, q) +#else + case List.span (p . unsafeElement) + $ IntSet.toAscList $ reflect $ Proxy @s of + (rs, qs) + | let r = IntSet.fromDistinctAscList rs + , let q = IntSet.fromDistinctAscList qs +#endif + -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q) + -> Some2IntSetWith @r @q $ PartialPartitionProof + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Return two disjoint subsets: those less than the given element, and those +-- greater than the given element; along with the proof that the given element +-- was in the set, if it was. +splitMember + :: forall s. KnownIntSet s + => Int -> Some2IntSetWith (SplitProof 'Int s (Element s)) +splitMember x = case IntSet.splitMember x $ reflect $ Proxy @s of + (r, m, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q) + -> Some2IntSetWith @r @q $ SplitProof + (if m then Just (unsafeElement x) else Nothing) + unsafeSubset \f g -> unsafeSubsetWith2 f g + +-- | Apply the given function to each element of the set and collect the +-- results. Note that the resulting set can be smaller. +map + :: forall s. KnownIntSet s + => (Element s -> Int) -> SomeIntSetWith (MapProof 'Int s Int Int) +map f = reify (IntMap.keysSet m) \(_ :: Proxy r) -> SomeIntSetWith @r + $ MapProof (unsafeElement . f) \y -> case IntMap.lookup (unrefine y) m of + Nothing -> error "map: bug: Data.IntSet.Refined has been subverted" + Just x -> x + where + !m = IntMap.fromList + [ (y, unsafeElement x) + | x <- IntSet.toList $ reflect $ Proxy @s + , let !y = f $ unsafeElement x + ] + + +-- | Map each element of @s@ into a monoid (with proof that it was an element), +-- and combine the results using 'Data.Monoid.<>'. +foldMap :: forall s m. (KnownIntSet s, Monoid m) => (Element s -> m) -> m +foldMap f = go $ reflect $ Proxy @s + where + go s = case IntSet.splitRoot s of + [s'] -> Foldable.foldMap (f . unsafeElement) $ IntSet.toAscList s' + xs -> Foldable.foldMap go xs + +-- | Right associative fold with a lazy accumulator. +foldr :: forall s a. KnownIntSet s => (Element s -> a -> a) -> a -> a +foldr f z = IntSet.foldr (f . unsafeElement) z $ reflect $ Proxy @s + +-- | Left associative fold with a lazy accumulator. +foldl :: forall s a. KnownIntSet s => (a -> Element s -> a) -> a -> a +foldl f z = IntSet.foldl ((. unsafeElement) . f) z $ reflect $ Proxy @s + +-- | Right associative fold with a strict accumulator. +foldr' :: forall s a. KnownIntSet s => (Element s -> a -> a) -> a -> a +foldr' f z = IntSet.foldr' (f . unsafeElement) z $ reflect $ Proxy @s + +-- | Left associative fold with a strict accumulator. +foldl' :: forall s a. KnownIntSet s => (a -> Element s -> a) -> a -> a +foldl' f z = IntSet.foldl' ((. unsafeElement) . f) z $ reflect $ Proxy @s + +-- | Retrieves the smallest element of the set, and the set with that element +-- removed; or a proof that the set was empty. +minView + :: forall s. KnownIntSet s + => Either + (EmptyProof 'Int s) + (Element s, SomeIntSetWith (SupersetProof 'Int s)) +minView = case IntSet.minView $ reflect $ Proxy @s of + Nothing -> Left $ EmptyProof unsafeSubset + Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r) + -> SomeIntSetWith @r $ SupersetProof unsafeSubset + +-- | Retrieves the greatest element of the set, and the set with that element +-- removed; or a proof that the set was empty. +maxView + :: forall s. KnownIntSet s + => Either + (EmptyProof 'Int s) + (Element s, SomeIntSetWith (SupersetProof 'Int s)) +maxView = case IntSet.maxView $ reflect $ Proxy @s of + Nothing -> Left $ EmptyProof unsafeSubset + Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r) + -> SomeIntSetWith @r $ SupersetProof unsafeSubset + +-- | List of elements in the set in ascending order. +toList :: forall s. KnownIntSet s => [Element s] +toList = gcoerceWith (unsafeCastElement @s) $ coerce + $ IntSet.toAscList $ reflect $ Proxy @s + +-- | List of elements in the set in descending order. +toDescList :: forall s. KnownIntSet s => [Element s] +toDescList = gcoerceWith (unsafeCastElement @s) $ coerce + $ IntSet.toDescList $ reflect $ Proxy @s + +-- | Convert an 'IntSet' into a 'Set', retaining its set of elements, which can +-- be converted with 'castFlavor'. +asSet :: forall s. KnownIntSet s => Set s Int +asSet = intSet2Set + +-- | Convert an 'IntSet' into a 'HashSet', retaining its set of elements, which +-- can be converted with 'castFlavor'. +asHashSet :: forall s. KnownIntSet s => HashSet s Int +asHashSet = intSet2HashSet + +-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then +-- @s@ and @t@ actually stand for the same set and @'Element' s@ can be safely +-- interconverted with @'Element' t@. +-- +-- The requirement that the weakenings are natural transformations ensures that +-- they don't actually alter the elements. To build these you can compose +-- ':->''s from proofs returned by functions in this module, or "Refined" +-- functions like 'andLeft' or 'leftOr'. +castElement + :: forall s t a. + (forall x. Refined (InSet 'Int s) x -> Refined (InSet 'Int t) x) + -> (forall x. Refined (InSet 'Int t) x -> Refined (InSet 'Int s) x) + -> Coercion (Refined (InSet 'Int s) a) (Refined (InSet 'Int t) a) +castElement = castRefined + +-- | If elements can be interconverted (e.g. as proved by 'castElement'), then +-- the sets can be interconverted too. For example we can establish that the +-- intersection of a set with itself is interconvertible with that set: +-- +-- @ +-- castIntersection +-- :: t'IntersectionProof' ''Data.IntSet.Refined.Int' s s r +-- -> 'Coercion' ('IntSet' r) ('IntSet' s) +-- castIntersection ( v'IntersectionProof' p1 p2) +-- = 'cast' $ 'castElement' ('andLeft' . p1) (p2 'id' 'id') +-- @ +cast + :: forall s t. (forall x. Coercion + (Refined (InSet 'Int s) x) + (Refined (InSet 'Int t) x)) + -> Coercion (IntSet s) (IntSet t) +cast Coercion +#if MIN_VERSION_base(4, 15, 0) + = case unsafeEqualityProof @s @t of UnsafeRefl -> Coercion +#else + = repr $ unsafeCoerce Refl +#endif diff --git a/src/Data/Map/Common/Refined.hs b/src/Data/Map/Common/Refined.hs new file mode 100644 index 0000000..14fdd5a --- /dev/null +++ b/src/Data/Map/Common/Refined.hs @@ -0,0 +1,515 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Map.Common.Refined where + +import Control.Monad.Reader +import Control.DeepSeq +import Data.Coerce +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Distributive +import Data.Foldable.WithIndex +import Data.Functor.Rep +import Data.Functor.WithIndex +import qualified Data.Hashable as Hashable +import qualified Data.Map as Map +import Data.Proxy +import Data.Reflection +import Data.Traversable.WithIndex +import Data.Type.Coercion +import Data.Type.Equality ((:~:)(..)) +import Refined +import Refined.Unsafe +import Unsafe.Coerce + +#if MIN_VERSION_containers(0, 6, 2) +#elif MIN_VERSION_containers(0, 5, 8) +import Data.Functor.Const (Const(..)) +import Data.Monoid (Any(..)) +import qualified Data.Map.Merge.Lazy as Map +#else +import qualified Data.List as List +import qualified Data.Map.Strict as MapStrict +#endif + + +-- | A wrapper around a regular 'Data.Map.Map' with a type parameter @s@ +-- identifying the set of keys present in the map. +-- +-- A key of type @k@ may not be present in the map, but a @'Key' s k@ is +-- guaranteed to be present (if the @s@ parameters match). Thus the map is +-- isomorphic to a (total) function @'Key' s k -> a@, which motivates many of +-- the instances below. +-- +-- A 'Map' always knows its set of keys, so given @'Map' s k a@ we can always +-- derive @'KnownSet' s k@ by pattern matching on the 'Dict' returned by +-- 'keysSet'. +newtype Map s k a = Map (Map.Map k a) + deriving newtype (Eq, Ord, Show, Functor, Foldable, NFData) +#if MIN_VERSION_hashable(1, 3, 4) + deriving newtype (Hashable.Hashable) +#endif + deriving stock (Traversable) +type role Map nominal nominal representational + +-- | Convert to a regular 'Data.Map.Map', forgetting its set of keys. +toMap :: forall s k a. Map s k a -> Map.Map k a +toMap (Map m) = m + +-- | @'Key' s k@ is a key of type @k@ that has been verified to be an element +-- of the set @s@, and thus verified to be present in all @'Map' s k@ maps. +-- +-- Thus, @'Key' s k@ is a \"refinement\" type of @k@, and this library +-- integrates with an implementation of refimenement types in "Refined", so +-- the machinery from there can be used to manipulate 'Key's (however see +-- 'Data.Set.Refined.revealPredicate'). +-- +-- The underlying @k@ value can be obtained with 'unrefine'. A @k@ can be +-- validated into an @'Key' s k@ with 'member'. +type Key s = Refined (InSet 'Regular s) + +unsafeCastKey :: forall s k. Coercion k (Key s k) +unsafeCastKey = reallyUnsafeUnderlyingRefined + +unsafeKey :: k -> Key s k +unsafeKey = coerceWith unsafeCastKey + +-- | An existential wrapper for a 'Map' with an as-yet-unknown set of keys. +-- Pattern maching on it gives you a way to refer to the set (the parameter +-- @s@), e.g. +-- +-- @ +-- case 'fromMap' ... of +-- 'SomeMap' \@s m -> doSomethingWith \@s +-- +-- case 'fromMap' ... of +-- 'SomeMap' (m :: 'Map' s k a) -> doSomethingWith \@s +-- @ +data SomeMap k a where + SomeMap :: forall s k a. !(Map s k a) -> SomeMap k a + +-- | Apply a map with an unknown set of keys to a continuation that can accept +-- a map with any set of keys. This gives you a way to refer to the set (the +-- parameter @s@), e.g.: +-- +-- @ +-- 'withMap' ('fromMap' ...) $ \(m :: 'Map' s k a) -> doSomethingWith \@s +-- @ +withMap :: forall k a r. SomeMap k a -> (forall s. Map s k a -> r) -> r +withMap (SomeMap m) k = k m + +-- | Construct a map from a regular 'Data.Map.Map'. +fromMap :: forall k a. Map.Map k a -> SomeMap k a +fromMap m = SomeMap (Map m) + +-- | An existential wrapper for a 'Map' with an as-yet-unknown set of keys, +-- together with a proof of some fact @p@ about the set. Pattern matching on it +-- gives you a way to refer to the set (the parameter @s@). Functions that +-- change the set of keys in a map will return the map in this way, together +-- with a proof that somehow relates the keys set to the function's inputs. +data SomeMapWith p k a where + SomeMapWith :: forall s k a p. !(Map s k a) -> !(p s) -> SomeMapWith p k a + +-- | Apply a map with proof for an unknown set of keys to a continuation that +-- can accept a map with any set of keys satisfying the proof. This gives you a +-- way to refer to the set (the parameter @s@). +withMapWith + :: forall k a r p. SomeMapWith p k a -> (forall s. Map s k a -> p s -> r) -> r +withMapWith (SomeMapWith m p) k = k m p + +-- | An existential wrapper for a pair of maps with as-yet-unknown sets of keys, +-- together with a proof of some fact @p@ relating them. +data Some2MapWith p k a b where + Some2MapWith + :: forall s t k a b p. !(Map s k a) + -> !(Map t k b) + -> !(p s t) + -> Some2MapWith p k a b + +-- | Apply a pair of maps with proof for unknown sets of keys to a continuation +-- that can accept any pair of maps with any sets of keys satisfying the proof. +-- This gives you a way to refer to the sets (the parameters @s@ and @t@). +with2MapWith + :: forall k a b r p. Some2MapWith p k a b + -> (forall s t. Map s k a -> Map t k b -> p s t -> r) + -> r +with2MapWith (Some2MapWith m1 m2 p) k = k m1 m2 p + +-- | An empty map. +empty :: forall k a. SomeMapWith (EmptyProof 'Regular) k a +empty = SomeMapWith (Map Map.empty) $ EmptyProof unsafeSubset + +-- | Create a map from a set of keys, and a function that for each key computes +-- the corresponding value. +fromSet :: forall s k a. KnownSet s k => (Key s k -> a) -> Map s k a +fromSet f = Map $ Map.fromSet (f . unsafeKey) (reflect $ Proxy @s) + +-- | Delete a key and its value from the map if present, returning a potentially +-- smaller map. +delete + :: forall s k a. Ord k + => k -> Map s k a -> SomeMapWith (SupersetProof 'Regular s) k a +delete k (Map m) = SomeMapWith (Map $ Map.delete k m) + $ SupersetProof unsafeSubset + +-- | If the key is in the map, return the proof of this, and the associated +-- value; otherwise return 'Nothing'. +lookup :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a) +lookup k (Map m) = (unsafeKey k,) <$> Map.lookup k m + +-- | Given a key that is proven to be in the map, return the associated value. +-- +-- Unlike 'Data.Map.!' from "Data.Map", this function is total, as it is +-- impossible to obtain a @'Key' s k@ for a key that is not in the map +-- @'Map' s k a@. +(!) :: forall s k a. Ord k => Map s k a -> Key s k -> a +(!) (Map m) k = case Map.lookup (unrefine k) m of + Nothing -> error "(!): bug: Data.Map.Refined has been subverted" + Just x -> x + +-- | If a key is in the map, return the proof that it is. +member :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k) +member k (Map m) + | k `Map.member` m = Just (unsafeKey k) + | otherwise = Nothing + +-- | Find the largest key smaller than the given one, and return the +-- associated key-value pair. +lookupLT :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a) +lookupLT = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupLT @k @a + +-- | Find the smallest key greater than the given one, and return the +-- associated key-value pair. +lookupGT :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a) +lookupGT = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupGT @k @a + +-- | Find the largest key smaller or equal to the given one, and return the +-- associated key-value pair. +lookupLE :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a) +lookupLE = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupLE @k @a + +-- | Find the smallest key greater or equal to the given one, and return the +-- associated key-value pair. +lookupGE :: forall s k a. Ord k => k -> Map s k a -> Maybe (Key s k, a) +lookupGE = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.lookupGE @k @a + +-- | If a map is empty, return a proof that it is. +null :: forall s k a. Map s k a -> Maybe (EmptyProof 'Regular s) +null (Map m) + | Map.null m = Just $ EmptyProof unsafeSubset + | otherwise = Nothing + +-- | If all keys of the first map are also present in the second map, and the +-- given function returns 'True' for their associated values, return a proof +-- that the keys form a subset. +isSubmapOfBy + :: forall s t k a b. Ord k + => (a -> b -> Bool) + -> Map s k a + -> Map t k b + -> Maybe (SubsetProof 'Regular s t) +isSubmapOfBy f (Map m1) (Map m2) + | Map.isSubmapOfBy f m1 m2 = Just $ SubsetProof unsafeSubset + | otherwise = Nothing + +-- | If two maps are disjoint (i.e. their intersection is empty), return a proof +-- of that. +disjoint + :: forall s t k a b. Ord k + => Map s k a -> Map t k b -> Maybe (DisjointProof 'Regular s t) +disjoint (Map m1) (Map m2) +#if MIN_VERSION_containers(0, 6, 2) + | Map.disjoint m1 m2 +#elif MIN_VERSION_containers(0, 5, 8) + | Const (Any False) <- Map.mergeA + (Map.traverseMissing \_ _ -> Const mempty) + (Map.traverseMissing \_ _ -> Const mempty) + (Map.zipWithAMatched \_ _ _ -> Const $ Any True) + m1 + m2 +#else + | Map.null $ MapStrict.intersectionWith (\_ _ -> ()) m1 m2 +#endif + = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g + | otherwise = Nothing + +-- | Given two maps proven to have the same keys, for each key apply the +-- function to the associated values, to obtain a new map with the same keys. +zipWithKey + :: forall s k a b c. Ord k + => (Key s k -> a -> b -> c) -> Map s k a -> Map s k b -> Map s k c +zipWithKey f (Map m1) (Map m2) = Map + $ Map.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y) + (error "zipWithKey: bug: Data.Map.Refined has been subverted") + (error "zipWithKey: bug: Data.Map.Refined has been subverted") + m1 + m2 + +-- | Remove the keys that appear in the second map from the first map. +difference + :: forall s t k a b. Ord k + => Map s k a -> Map t k b -> SomeMapWith (DifferenceProof 'Regular s t) k a +difference (Map m1) (Map m2) = SomeMapWith (Map $ Map.difference m1 m2) + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, that are proven to be in the map. The set of keys remains the same. +mapWithKey :: forall s k a b. (Key s k -> a -> b) -> Map s k a -> Map s k b +mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.mapWithKey @k @a @b + +-- | Map an 'Applicative' transformation in ascending order of keys, with access +-- to each value's corresponding key and a proof that it is in the map. The set +-- of keys remains unchanged. +traverseWithKey + :: forall s f k a b. Applicative f + => (Key s k -> a -> f b) -> Map s k a -> f (Map s k b) +traverseWithKey f (Map m) = Map <$> Map.traverseWithKey (f . unsafeKey) m + +-- | Map each key-value pair of a map into a monoid (with proof that the key was +-- in the map), and combine the results using '<>'. +foldMapWithKey + :: forall s k a m. Monoid m => (Key s k -> a -> m) -> Map s k a -> m +foldMapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.foldMapWithKey @m @k @a + +-- | Right associative fold with a lazy accumulator. +foldrWithKey :: forall s k a b. (Key s k -> a -> b -> b) -> b -> Map s k a -> b +foldrWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.foldrWithKey @k @a @b + +-- | Left associative fold with a lazy accumulator. +foldlWithKey :: forall s k a b. (b -> Key s k -> a -> b) -> b -> Map s k a -> b +foldlWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.foldlWithKey @b @k @a + +-- | Right associative fold with a strict accumulator. +foldrWithKey' :: forall s k a b. (Key s k -> a -> b -> b) -> b -> Map s k a -> b +foldrWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.foldrWithKey' @k @a @b + +-- | Left associative fold with a strict accumulator. +foldlWithKey' :: forall s k a b. (b -> Key s k -> a -> b) -> b -> Map s k a -> b +foldlWithKey' = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.foldlWithKey' @b @k @a + +-- | Return the set of keys in the map, with the contents of the set still +-- tracked by the @s@ parameter. See "Data.Set.Refined". +keysSet :: forall s k a. Map s k a -> Set s k +keysSet (Map m) = reify (Map.keysSet m) + \(_ :: Proxy s') -> case unsafeCoerce Refl :: s :~: s' of + Refl -> Dict + +-- | Convert to a list of key-value pairs in ascending order of keys. +toList :: forall s k a. Map s k a -> [(Key s k, a)] +toList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.toAscList @k @a + +-- | Convert to a list of key-value pairs in descending order of keys. +toDescList :: forall s k a. Map s k a -> [(Key s k, a)] +toDescList = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.toDescList @k @a + +-- | Retain only the key-value pairs that satisfy the predicate, returning a +-- potentially smaller map. +filterWithKey + :: forall s k a. (Key s k -> a -> Bool) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +filterWithKey p (Map m) + = SomeMapWith (Map $ Map.filterWithKey (p . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Restrict a map to only those keys that are elements of @t@. +restrictKeys + :: forall s t k a. (Ord k, KnownSet t k) + => Map s k a -> SomeMapWith (IntersectionProof 'Regular s t) k a +restrictKeys (Map m) = SomeMapWith +#if MIN_VERSION_containers(0, 5, 8) + (Map $ Map.restrictKeys m $ reflect $ Proxy @t) +#else + (Map $ Map.intersectionWith const m $ Map.fromSet id $ reflect $ Proxy @t) +#endif + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Remove all keys that are elements of @t@ from the map. +withoutKeys + :: forall s t k a. (Ord k, KnownSet t k) + => Map s k a -> SomeMapWith (DifferenceProof 'Regular s t) k a +withoutKeys (Map m) = SomeMapWith +#if MIN_VERSION_containers(0, 5, 8) + (Map $ Map.withoutKeys m $ reflect $ Proxy @t) +#else + (Map $ Map.difference m $ Map.fromSet id $ reflect $ Proxy @t) +#endif + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Partition a map into two disjoint submaps: those whose key-value pairs +-- satisfy the predicate, and those whose don't. +partitionWithKey + :: forall s k a. Ord k -- TODO: this is only used in the proof + => (Key s k -> a -> Bool) + -> Map s k a + -> Some2MapWith (PartitionProof 'Regular s k) k a a +partitionWithKey p (Map m) = case Map.partitionWithKey (p . unsafeKey) m of + (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof + do \k -> case Map.lookup (unrefine k) m of + Nothing + -> error "partitionWithKey: bug: Data.Map.Refined has been subverted" + Just x -> if p k x + then Left $ unsafeKey $ unrefine k + else Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Divide a map into two disjoint submaps at a point where the predicate on +-- the keys stops holding. +-- +-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then +-- this point is uniquely defined. If @p@ is not antitone, a splitting point is +-- chosen in an unspecified way. +spanAntitone + :: forall s k a. (Key s k -> Bool) + -> Map s k a + -> Some2MapWith (PartialPartitionProof 'Regular s) k a a +spanAntitone p (Map m) = +#if MIN_VERSION_containers(0, 5, 8) + case Map.spanAntitone (p . unsafeKey) m of + (m1, m2) +#else + case List.span (p . unsafeKey . fst) $ Map.toAscList m of + (xs1, xs2) + | let m1 = Map.fromDistinctAscList xs1 + , let m2 = Map.fromDistinctAscList xs2 +#endif + -> Some2MapWith (Map m1) (Map m2) $ PartialPartitionProof + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Return two disjoint submaps: those whose keys are less than the given key, +-- and those whose keys are greater than the given key. If the key was in the +-- map, also return the associated value and the proof that it was in the map. +splitLookup + :: forall s k a. Ord k + => k -> Map s k a -> Some2MapWith (SplitProof 'Regular s (Key s k, a)) k a a +splitLookup k (Map m) = case Map.splitLookup k m of + (!m1, v, !m2) -> Some2MapWith (Map m1) (Map m2) $ SplitProof + ((unsafeKey k,) <$> v) unsafeSubset \f g -> unsafeSubsetWith2 f g + +-- | Retrieves the key-value pair corresponding to the smallest key of the map, +-- and the map with that pair removed; or a proof that the map was empty. +minViewWithKey + :: forall s k a. Map s k a + -> Either + (EmptyProof 'Regular s) + ((Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a) +minViewWithKey (Map m) = case Map.minViewWithKey m of + Nothing -> Left $ EmptyProof unsafeSubset + Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s @k) $ coerce kv,) + $ SomeMapWith (Map m') $ SupersetProof unsafeSubset + +-- | Retrieves the key-value pair corresponding to the greatest key of the map, +-- and the map with that pair removed; or a proof that the map was empty. +maxViewWithKey + :: forall s k a. Map s k a + -> Either + (EmptyProof 'Regular s) + ((Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a) +maxViewWithKey (Map m) = case Map.maxViewWithKey m of + Nothing -> Left $ EmptyProof unsafeSubset + Just (kv, m') -> Right $ (gcoerceWith (unsafeCastKey @s @k) $ coerce kv,) + $ SomeMapWith (Map m') $ SupersetProof unsafeSubset + +-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then +-- @s@ and @t@ actually stand for the same set and @'Key' s@ can be safely +-- interconverted with @'Key' t@. +-- +-- The requirement that the weakenings are natural transformations ensures that +-- they don't actually alter the keys. To build these you can compose ':->''s +-- from proofs returned by functions in this module, or "Refined" functions like +-- 'andLeft' or 'leftOr'. +castKey + :: forall s t k. (forall x. Key s x -> Key t x) + -> (forall x. Key t x -> Key s x) + -> Coercion (Key s k) (Key t k) +castKey = castRefined + +-- | If keys can be interconverted (e.g. as proved by 'castKey'), then the maps +-- can be interconverted too. For example, 'zipWithKey' can be implemented via +-- 'Data.Map.Refined.intersectionWithKey' by proving that the set of keys +-- remains unchanged: +-- +-- @ +-- 'zipWithKey' +-- :: forall s k a b c. 'Ord' k +-- => ('Key' s k -> a -> b -> c) -> 'Map' s k a -> 'Map' s k b -> 'Map' s k c +-- 'zipWithKey' f m1 m2 +-- | v'SomeMapWith' @r m proof <- 'Data.Map.Refined.intersectionWithKey' (f . 'andLeft') m1 m2 +-- , v'IntersectionProof' p1 p2 <- proof +-- , ( v'Coercion' :: t'Coercion' ('Map' r k c) ('Map' s k c)) +-- <- app $ 'cast' $ 'castKey' ('andLeft' . p1) (p2 'id' 'id') +-- = 'coerce' m +-- where +-- app :: t'Coercion' f g -> t'Coercion' (f x) (g x) +-- app v'Coercion' = v'Coercion' +-- @ +cast + :: forall s t k. (forall x. Coercion (Key s x) (Key t x)) + -> Coercion (Map s k) (Map t k) +cast Coercion = Coercion + +instance FunctorWithIndex (Key s k) (Map s k) where + imap = mapWithKey + +instance FoldableWithIndex (Key s k) (Map s k) where + ifoldMap = foldMapWithKey + +instance TraversableWithIndex (Key s k) (Map s k) where + itraverse = traverseWithKey + +-- | Similar to the instance for functions -- zip corresponding keys. To use +-- '<*>'/'liftA2' without 'KnownSet' see 'zipWithKey'. +instance (Ord k, KnownSet s k) => Applicative (Map s k) where + pure x = fromSet \_ -> x + (<*>) = zipWithKey (const id) + +-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the +-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions. +bind :: forall s k a b. Ord k => Map s k a -> (a -> Map s k b) -> Map s k b +bind m f = mapWithKey (\k x -> f x ! k) m + +-- | Similar to the instance for functions. To use '>>=' without 'KnownSet' see +-- 'bind'. +instance (Ord k, KnownSet s k) => Monad (Map s k) where + (>>=) = bind + +-- | Similar to the instance for functions. See also +-- 'Data.Map.Refined.backpermuteKeys'. +instance (Ord k, KnownSet s k) => MonadReader (Key s k) (Map s k) where + ask = fromSet id + local f m = mapWithKey (\k _ -> m ! f k) m + +-- | Append the values at the corresponding keys +instance (Ord k, Semigroup a) => Semigroup (Map s k a) where + (<>) = zipWithKey (const (<>)) + +instance (Ord k, KnownSet s k, Monoid a) => Monoid (Map s k a) where + mempty = fromSet \_ -> mempty + +-- | Similar to the instance for functions +instance (Ord k, KnownSet s k) => Distributive (Map s k) where + collect = collectRep + distribute = distributeRep + +-- | Witness isomorphism with functions from @'Key' s k@ +instance (Ord k, KnownSet s k) => Representable (Map s k) where + type Rep (Map s k) = Key s k + index = (!) + tabulate = fromSet + +#if MIN_VERSION_hashable(1, 3, 4) +#else +instance (Hashable.Hashable a, Hashable.Hashable k) + => Hashable.Hashable (Map s k a) where + hashWithSalt s (Map m) = Map.foldlWithKey' + (\s' k v -> Hashable.hashWithSalt (Hashable.hashWithSalt s' k) v) + (Hashable.hashWithSalt s (Map.size m)) + m +#endif diff --git a/src/Data/Map/Refined.hs b/src/Data/Map/Refined.hs new file mode 100644 index 0000000..a9cfe66 --- /dev/null +++ b/src/Data/Map/Refined.hs @@ -0,0 +1,384 @@ +-- | This module defines a way to prove that a key exists in a map, so that the +-- key can be used to index into the map without using a 'Maybe', or manually +-- handling the \"impossible\" case with 'error' or other partial functions. +-- +-- To do this, @'Map' s k v@ has a type parameter @s@ that identifies its set of +-- keys, so that if another map has the same type parameter, you know that map +-- has the same set of keys. There is @'Key' s k@, a type of keys that have been +-- validated to belong to the set identified by @s@, and for which the operation +-- of indexing into a @'Map' s k v@ (only for the same @s@) can proceed without +-- failure (see '!'). The type @s@ itself has no internal structure, rather it +-- is merely a skolem type variable (rank-2 polymorphism +-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection". +-- +-- Like "Data.Map", functions in this module are strict in the keys but lazy in +-- the values. The "Data.Map.Strict.Refined" module reuses the same 'Map' type +-- but provides functions that operate strictly on the values. +-- +-- = Warning +-- This module together with "Data.Map" rely on 'Eq' and 'Ord' instances being +-- lawful: that '==' is an equivalence relation, and that the 'Ord' operations +-- define a total order on the quotient defined by this equivalence relation; at +-- least for the subset of keys that are actually encountered at runtime. If +-- this assumption is violated, this module may not be able to uphold its +-- invariants and may throw errors. In particular beware of NaN in 'Float' and +-- 'Double'. +module Data.Map.Refined + ( + -- * Map type + Common.Map + , Common.Key + -- * Existentials and common proofs + , Common.SomeMap(..) + , Common.withMap + , Common.SomeMapWith(..) + , Common.withMapWith + , Common.Some2MapWith(..) + , Common.with2MapWith + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , Common.empty + , singleton + , SingletonProof(..) + , fromSet + , Common.fromMap + , fromTraversableWithKey + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + , reinsert + , insertLookupWithKey + -- * Deletion/Update + , Common.delete + , adjust + , adjustWithKey + , update + , updateLookupWithKey + -- * Query + , Common.lookup + , (Common.!) + , Common.member + , Common.lookupLT + , Common.lookupGT + , Common.lookupLE + , Common.lookupGE + , Common.null + , Common.isSubmapOfBy + , SubsetProof(..) + , Common.disjoint + , DisjointProof(..) + -- * Combine + , zipWithKey + , bind + , unionWithKey + , UnionProof(..) + , Common.difference + , DifferenceProof(..) + , differenceWithKey + , PartialDifferenceProof(..) + , intersectionWithKey + , IntersectionProof(..) + -- * Traversal + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , MapProof(..) + , backpermuteKeys + -- * Folds + , Common.foldMapWithKey + , Common.foldrWithKey + , Common.foldlWithKey + , Common.foldrWithKey' + , Common.foldlWithKey' + -- * Conversion + , Common.toMap + , Common.keysSet + , Common.toList + , Common.toDescList + -- * Filter + , Common.restrictKeys + , Common.withoutKeys + , Common.filterWithKey + , Common.partitionWithKey + , PartitionProof(..) + , Common.spanAntitone + , PartialPartitionProof(..) + , mapMaybeWithKey + , mapEitherWithKey + , Common.splitLookup + , SplitProof(..) + -- * Min/Max + , updateMinWithKey + , updateMaxWithKey + , adjustMinWithKey + , adjustMaxWithKey + , Common.minViewWithKey + , Common.maxViewWithKey + -- * Casts + , Common.castKey + , Common.cast + , castFlavor + ) where + +import Data.Coerce +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Functor +import qualified Data.Map as Map +import Data.Map.Common.Refined + ( Map(..), Key, unsafeCastKey, unsafeKey, SomeMapWith(..), Some2MapWith(..) + , fromSet, (!), zipWithKey, mapWithKey, traverseWithKey, bind + ) +import qualified Data.Map.Common.Refined as Common +import Data.Traversable +import Data.Type.Coercion +import Prelude hiding (lookup, null) +import Refined +import Refined.Unsafe + + +-- | Create a map with a single key-value pair, and return a proof that the +-- key is in the resulting map. +singleton :: forall k a. k -> a -> SomeMapWith (SingletonProof 'Regular k) k a +singleton k v = SomeMapWith (Map $ Map.singleton k v) + $ SingletonProof (unsafeKey k) + +-- | Create a map from an arbitrary traversable of key-value pairs. +fromTraversableWithKey + :: forall t k a. (Traversable t, Ord k) + => (k -> a -> a -> a) + -> t (k, a) + -> SomeMapWith (FromTraversableProof 'Regular t k) k a +fromTraversableWithKey f xs = SomeMapWith (Map m) $ FromTraversableProof proof + where + (m, proof) = mapAccumL + (\s (k, v) -> let !s' = Map.insertWithKey f k v s in (s', unsafeKey k)) + Map.empty + xs + +-- | Insert a key-value pair into the map to obtain a potentially larger map, +-- guaranteed to contain the given key. If the key was already present, the +-- associated value is replaced with the supplied value. +insert + :: forall s k a. Ord k + => k -> a -> Map s k a -> SomeMapWith (InsertProof 'Regular k s) k a +insert k v (Map m) = SomeMapWith (Map $ Map.insert k v m) + $ InsertProof (unsafeKey k) unsafeSubset + +-- | Overwrite a key-value pair that is known to already be in the map. The set +-- of keys remains the same. +reinsert + :: forall s k a. Ord k + => Key s k -> a -> Map s k a -> Map s k a +reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.insert @k @a + +-- | Insert a key-value pair into the map using a combining function, and if +-- the key was already present, the old value is returned along with the proof +-- that the key was present. +insertLookupWithKey + :: forall s k a. Ord k + => (Key s k -> a -> a -> a) + -> k + -> a + -> Map s k a + -> (Maybe (Key s k, a), SomeMapWith (InsertProof 'Regular k s) k a) +insertLookupWithKey f k v (Map m) + = case Map.insertLookupWithKey (f . unsafeKey) k v m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeMapWith (Map m') $ InsertProof (unsafeKey k) unsafeSubset + +-- | Update the value at a specific key known the be in the map using the given +-- function. The set of keys remains the same. +adjust :: forall s k a. Ord k => (a -> a) -> Key s k -> Map s k a -> Map s k a +adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.adjust @k @a + +-- | If the given key is in the map, update the associated value using the given +-- function with a proof that the key was in the map; otherwise return the map +-- unchanged. In any case the set of keys remains the same. +adjustWithKey + :: forall s k a. Ord k => (Key s k -> a -> a) -> k -> Map s k a -> Map s k a +adjustWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.adjustWithKey @k @a + +-- | Update or delete a key known to be in the map using the given function, +-- returning a potentially smaller map. +update + :: forall s k a. Ord k + => (a -> Maybe a) + -> Key s k + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +update f k (Map m) = SomeMapWith (Map $ Map.update f (unrefine k) m) + $ SupersetProof unsafeSubset + +-- | If the given key is in the map, update or delete it using the given +-- function with a proof that the key was in the map; otherwise the map is +-- unchanged. Alongside return the new value if it was updated, or the old value +-- if it was deleted, and a proof that the key was in the map. +updateLookupWithKey + :: forall s k a. Ord k + => (Key s k -> a -> Maybe a) + -> k + -> Map s k a + -> (Maybe (Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a) +updateLookupWithKey f k (Map m) + = case Map.updateLookupWithKey (f . unsafeKey) k m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeMapWith (Map m') $ SupersetProof unsafeSubset + +-- | Return the union of two maps, with a given combining function for keys that +-- exist in both maps simultaneously. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +unionWithKey + :: forall s t k a. Ord k + => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> a -> a) + -> Map s k a + -> Map t k a + -> SomeMapWith (UnionProof 'Regular s t) k a +unionWithKey f (Map m1) (Map m2) + = SomeMapWith (Map $ Map.unionWithKey (f . reallyUnsafeRefine) m1 m2) + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- | For keys that appear in both maps, the given function decides whether the +-- key is removed from the first map. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +differenceWithKey + :: forall s t k a b. Ord k + => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> Maybe a) + -> Map s k a + -> Map t k b + -> SomeMapWith (PartialDifferenceProof 'Regular s t) k a +differenceWithKey f (Map m1) (Map m2) + = SomeMapWith (Map $ Map.differenceWithKey (f . reallyUnsafeRefine) m1 m2) + $ PartialDifferenceProof unsafeSubset unsafeSubset + +-- | Return the intersection of two maps with the given combining function. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +intersectionWithKey + :: forall s t k a b c. Ord k + => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> c) + -> Map s k a + -> Map t k b + -> SomeMapWith (IntersectionProof 'Regular s t) k c +intersectionWithKey f (Map m1) (Map m2) + = SomeMapWith (Map $ Map.intersectionWithKey (f . reallyUnsafeRefine) m1 m2) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Thread an accumularing argument through the map in ascending order of keys. +mapAccumLWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> Map s k b + -> (a, Map s k c) +mapAccumLWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.mapAccumWithKey @a @k @b @c + +-- | Thread an accumularing argument through the map in descending order of +-- keys. +mapAccumRWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> Map s k b + -> (a, Map s k c) +mapAccumRWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.mapAccumRWithKey @a @k @b @c + +-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the +-- results into a new map. For keys that were mapped to the same new key, @c@ +-- acts as the combining function for corresponding values. +mapKeysWith + :: forall s k1 k2 a. Ord k2 + => (a -> a -> a) + -> (Key s k1 -> k2) + -> Map s k1 a + -> SomeMapWith (MapProof 'Regular s k1 k2) k2 a +mapKeysWith f g (Map m) + = SomeMapWith (Map $ Map.mapKeysWith f (g . unsafeKey) m) + $ MapProof (unsafeKey . g) \k2 -> case Map.lookup (unrefine k2) backMap of + Nothing -> error "mapKeysWith: bug: Data.Map.Refined has been subverted" + Just k1 -> k1 + where + ~backMap = Map.fromList + [ (k2, unsafeKey k1) + | k1 <- Map.keys m + , let !k2 = g $ unsafeKey k1 + ] + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect only the 'Just' results, returning a potentially smaller +-- map. +mapMaybeWithKey + :: forall s k a b. (Key s k -> a -> Maybe b) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k b +mapMaybeWithKey f (Map m) + = SomeMapWith (Map $ Map.mapMaybeWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect the 'Left' and 'Right' results into separate (disjoint) +-- maps. +mapEitherWithKey + :: forall s k a b c. Ord k -- TODO: this is only used in the proof + => (Key s k -> a -> Either b c) + -> Map s k a + -> Some2MapWith (PartitionProof 'Regular s k) k b c +mapEitherWithKey p (Map m) = case Map.mapEitherWithKey (p . unsafeKey) m of + (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof + do \k -> case Map.lookup (unrefine k) m of + Nothing + -> error "mapEitherWithKey: bug: Data.Map.Refined has been subverted" + Just x -> case p k x of + Left _ -> Left $ unsafeKey $ unrefine k + Right _ -> Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Update or delete the value at the smallest key, returning a potentially +-- smaller map. +updateMinWithKey + :: forall s k a. (Key s k -> a -> Maybe a) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +updateMinWithKey f (Map m) + = SomeMapWith (Map $ Map.updateMinWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Update or delete the value at the largest key, returning a potentially +-- smaller map. +updateMaxWithKey + :: forall s k a. (Key s k -> a -> Maybe a) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +updateMaxWithKey f (Map m) + = SomeMapWith (Map $ Map.updateMaxWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Adjust the value at the smallest key. The set of keys remains unchanged. +adjustMinWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a +adjustMinWithKey f (Map m) + = Map $ Map.updateMinWithKey ((Just .) . f . unsafeKey) m + +-- | Adjust the value at the greatest key. The set of keys remains unchanged. +adjustMaxWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a +adjustMaxWithKey f (Map m) + = Map $ Map.updateMaxWithKey ((Just .) . f . unsafeKey) m + +-- | Apply the inverse image of the given function to the keys of the given map, +-- so that for all @k :: 'Key' s2 k2@, +-- @'backpermuteKeys' f m '!' k = m '!' f k@. +-- +-- If maps are identified with functions, this computes the composition. +backpermuteKeys + :: forall s1 s2 k1 k2 a. (Ord k1, KnownSet s2 k2) + => (Key s2 k2 -> Key s1 k1) -> Map s1 k1 a -> Map s2 k2 a +backpermuteKeys f m = fromSet \k -> m ! f k diff --git a/src/Data/Map/Strict/Refined.hs b/src/Data/Map/Strict/Refined.hs new file mode 100644 index 0000000..2bac38d --- /dev/null +++ b/src/Data/Map/Strict/Refined.hs @@ -0,0 +1,423 @@ +-- | This module defines a way to prove that a key exists in a map, so that the +-- key can be used to index into the map without using a 'Maybe', or manually +-- handling the \"impossible\" case with 'error' or other partial functions. +-- +-- To do this, @'Map' s k v@ has a type parameter @s@ that identifies its set of +-- keys, so that if another map has the same type parameter, you know that map +-- has the same set of keys. There is @'Key' s k@, a type of keys that have been +-- validated to belong to the set identified by @s@, and for which the operation +-- of indexing into a @'Map' s k v@ (only for the same @s@) can proceed without +-- failure (see '!'). The type @s@ itself has no internal structure, rather it +-- is merely a skolem type variable (rank-2 polymorphism +-- 'Control.Monad.ST.runST' trick) introduced by "Data.Reflection". +-- +-- Like "Data.Map.Strict", functions in this module are strict in the keys and +-- values. The "Data.Map.Refined" module reuses the same 'Map' type but provides +-- functions that operate lazily on the values. +-- +-- = Warning +-- This module together with "Data.Map.Strict" rely on 'Eq' and 'Ord' instances +-- being lawful: that '==' is an equivalence relation, and that the 'Ord' +-- operations define a total order on the quotient defined by this equivalence +-- relation; at least for the subset of keys that are actually encountered at +-- runtime. If this assumption is violated, this module may not be able to +-- uphold its invariants and may throw errors. In particular beware of NaN in +-- 'Float' and 'Double'. +module Data.Map.Strict.Refined + ( + -- * Map type + Common.Map + , Common.Key + -- * Existentials and common proofs + , Common.SomeMap(..) + , Common.withMap + , Common.SomeMapWith(..) + , Common.withMapWith + , Common.Some2MapWith(..) + , Common.with2MapWith + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , Common.empty + , singleton + , SingletonProof(..) + , fromSet + , Common.fromMap + , fromTraversableWithKey + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + , reinsert + , insertLookupWithKey + -- * Deletion/Update + , Common.delete + , adjust + , adjustWithKey + , update + , updateLookupWithKey + -- * Query + , Common.lookup + , (Common.!) + , Common.member + , Common.lookupLT + , Common.lookupGT + , Common.lookupLE + , Common.lookupGE + , Common.null + , Common.isSubmapOfBy + , SubsetProof(..) + , Common.disjoint + , DisjointProof(..) + -- * Combine + , zipWithKey + , bind + , unionWithKey + , UnionProof(..) + , Common.difference + , DifferenceProof(..) + , differenceWithKey + , PartialDifferenceProof(..) + , intersectionWithKey + , IntersectionProof(..) + -- * Traversal + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , MapProof(..) + , backpermuteKeys + -- * Folds + , Common.foldMapWithKey + , Common.foldrWithKey + , Common.foldlWithKey + , Common.foldrWithKey' + , Common.foldlWithKey' + -- * Conversion + , Common.toMap + , Common.keysSet + , Common.toList + , Common.toDescList + -- * Filter + , Common.restrictKeys + , Common.withoutKeys + , Common.filterWithKey + , Common.partitionWithKey + , PartitionProof(..) + , Common.spanAntitone + , PartialPartitionProof(..) + , mapMaybeWithKey + , mapEitherWithKey + , Common.splitLookup + , SplitProof(..) + -- * Min/Max + , updateMinWithKey + , updateMaxWithKey + , adjustMinWithKey + , adjustMaxWithKey + , Common.minViewWithKey + , Common.maxViewWithKey + -- * Casts + , Common.castKey + , Common.cast + , castFlavor + ) where + +import Data.Coerce +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import Data.Functor +import qualified Data.Map.Strict as Map +import Data.Map.Common.Refined + ( Map(..), Key, unsafeCastKey, unsafeKey, SomeMapWith(..), Some2MapWith(..) + , (!) + ) +import qualified Data.Map.Common.Refined as Common +import Data.Proxy +import Data.Reflection +import Data.Traversable +import Data.Type.Coercion +import Prelude hiding (lookup, null) +import Refined +import Refined.Unsafe + + +-- | Create a map with a single key-value pair, and return a proof that the +-- key is in the resulting map. +singleton :: forall k a. k -> a -> SomeMapWith (SingletonProof 'Regular k) k a +singleton k v = SomeMapWith (Map $ Map.singleton k v) + $ SingletonProof (unsafeKey k) + +-- | Create a map from a set of keys, and a function that for each key computes +-- the corresponding value. +fromSet :: forall s k a. KnownSet s k => (Key s k -> a) -> Map s k a +fromSet f = Map $ Map.fromSet (f . unsafeKey) (reflect $ Proxy @s) + +-- | Create a map from an arbitrary traversable of key-value pairs. +fromTraversableWithKey + :: forall t k a. (Traversable t, Ord k) + => (k -> a -> a -> a) + -> t (k, a) + -> SomeMapWith (FromTraversableProof 'Regular t k) k a +fromTraversableWithKey f xs = SomeMapWith (Map m) $ FromTraversableProof proof + where + (m, proof) = mapAccumL + (\s (k, v) -> let !s' = Map.insertWithKey f k v s in (s', unsafeKey k)) + Map.empty + xs + +-- | Insert a key-value pair into the map to obtain a potentially larger map, +-- guaranteed to contain the given key. If the key was already present, the +-- associated value is replaced with the supplied value. +insert + :: forall s k a. Ord k + => k -> a -> Map s k a -> SomeMapWith (InsertProof 'Regular k s) k a +insert k v (Map m) = SomeMapWith (Map $ Map.insert k v m) + $ InsertProof (unsafeKey k) unsafeSubset + +-- | Overwrite a key-value pair that is known to already be in the map. The set +-- of keys remains the same. +reinsert + :: forall s k a. Ord k + => Key s k -> a -> Map s k a -> Map s k a +reinsert = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.insert @k @a + +-- | Insert a key-value pair into the map using a combining function, and if +-- the key was already present, the old value is returned along with the proof +-- that the key was present. +insertLookupWithKey + :: forall s k a. Ord k + => (Key s k -> a -> a -> a) + -> k + -> a + -> Map s k a + -> (Maybe (Key s k, a), SomeMapWith (InsertProof 'Regular k s) k a) +insertLookupWithKey f k v (Map m) + = case Map.insertLookupWithKey (f . unsafeKey) k v m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeMapWith (Map m') $ InsertProof (unsafeKey k) unsafeSubset + +-- | Update the value at a specific key known the be in the map using the given +-- function. The set of keys remains the same. +adjust :: forall s k a. Ord k => (a -> a) -> Key s k -> Map s k a -> Map s k a +adjust = gcoerceWith (unsafeCastKey @s @k) $ coerce $ Map.adjust @k @a + +-- | If the given key is in the map, update the associated value using the given +-- function with a proof that the key was in the map; otherwise return the map +-- unchanged. In any case the set of keys remains the same. +adjustWithKey + :: forall s k a. Ord k => (Key s k -> a -> a) -> k -> Map s k a -> Map s k a +adjustWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.adjustWithKey @k @a + +-- | Update or delete a key known to be in the map using the given function, +-- returning a potentially smaller map. +update + :: forall s k a. Ord k + => (a -> Maybe a) + -> Key s k + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +update f k (Map m) = SomeMapWith (Map $ Map.update f (unrefine k) m) + $ SupersetProof unsafeSubset + +-- | If the given key is in the map, update or delete it using the given +-- function with a proof that the key was in the map; otherwise the map is +-- unchanged. Alongside return the new value if it was updated, or the old value +-- if it was deleted, and a proof that the key was in the map. +updateLookupWithKey + :: forall s k a. Ord k + => (Key s k -> a -> Maybe a) + -> k + -> Map s k a + -> (Maybe (Key s k, a), SomeMapWith (SupersetProof 'Regular s) k a) +updateLookupWithKey f k (Map m) + = case Map.updateLookupWithKey (f . unsafeKey) k m of + (v', !m') -> ((unsafeKey k,) <$> v',) + $ SomeMapWith (Map m') $ SupersetProof unsafeSubset + +-- | Given two maps proven to have the same keys, for each key apply the +-- function to the associated values, to obtain a new map with the same keys. +zipWithKey + :: forall s k a b c. Ord k + => (Key s k -> a -> b -> c) -> Map s k a -> Map s k b -> Map s k c +zipWithKey f (Map m1) (Map m2) = Map + $ Map.mergeWithKey (\k x y -> Just $ f (unsafeKey k) x y) + (error "zipWithKey: bug: Data.Map.Strict.Refined has been subverted") + (error "zipWithKey: bug: Data.Map.Strict.Refined has been subverted") + m1 + m2 + +-- | Return the union of two maps, with a given combining function for keys that +-- exist in both maps simultaneously. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +unionWithKey + :: forall s t k a. Ord k + => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> a -> a) + -> Map s k a + -> Map t k a + -> SomeMapWith (UnionProof 'Regular s t) k a +unionWithKey f (Map m1) (Map m2) + = SomeMapWith (Map $ Map.unionWithKey (f . reallyUnsafeRefine) m1 m2) + $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- | For keys that appear in both maps, the given function decides whether the +-- key is removed from the first map. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +differenceWithKey + :: forall s t k a b. Ord k + => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> Maybe a) + -> Map s k a + -> Map t k b + -> SomeMapWith (PartialDifferenceProof 'Regular s t) k a +differenceWithKey f (Map m1) (Map m2) + = SomeMapWith (Map $ Map.differenceWithKey (f . reallyUnsafeRefine) m1 m2) + $ PartialDifferenceProof unsafeSubset unsafeSubset + +-- | Return the intersection of two maps with the given combining function. +-- +-- You can use 'andLeft' and 'andRight' to obtain @'Key' s k@ and @'Key' t k@ +-- respectively. +intersectionWithKey + :: forall s t k a b c. Ord k + => (Refined (InSet 'Regular s && InSet 'Regular t) k -> a -> b -> c) + -> Map s k a + -> Map t k b + -> SomeMapWith (IntersectionProof 'Regular s t) k c +intersectionWithKey f (Map m1) (Map m2) + = SomeMapWith (Map $ Map.intersectionWithKey (f . reallyUnsafeRefine) m1 m2) + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, that are proven to be in the map. The set of keys remains the same. +mapWithKey :: forall s k a b. (Key s k -> a -> b) -> Map s k a -> Map s k b +mapWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.mapWithKey @k @a @b + +-- | Map an 'Applicative' transformation in ascending order of keys, with access +-- to each value's corresponding key and a proof that it is in the map. The set +-- of keys remains unchanged. +traverseWithKey + :: forall s f k a b. Applicative f + => (Key s k -> a -> f b) -> Map s k a -> f (Map s k b) +traverseWithKey f (Map m) = Map <$> Map.traverseWithKey (f . unsafeKey) m + +-- | Thread an accumularing argument through the map in ascending order of keys. +mapAccumLWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> Map s k b + -> (a, Map s k c) +mapAccumLWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.mapAccumWithKey @a @k @b @c + +-- | Thread an accumularing argument through the map in descending order of +-- keys. +mapAccumRWithKey + :: forall s k a b c. (a -> Key s k -> b -> (a, c)) + -> a + -> Map s k b + -> (a, Map s k c) +mapAccumRWithKey = gcoerceWith (unsafeCastKey @s @k) $ coerce + $ Map.mapAccumRWithKey @a @k @b @c + +-- | @'mapKeysWith' c f m@ applies @f@ to each key of @m@ and collects the +-- results into a new map. For keys that were mapped to the same new key, @c@ +-- acts as the combining function for corresponding values. +mapKeysWith + :: forall s k1 k2 a. Ord k2 + => (a -> a -> a) + -> (Key s k1 -> k2) + -> Map s k1 a + -> SomeMapWith (MapProof 'Regular s k1 k2) k2 a +mapKeysWith f g (Map m) + = SomeMapWith (Map $ Map.mapKeysWith f (g . unsafeKey) m) + $ MapProof (unsafeKey . g) \k2 -> case Map.lookup (unrefine k2) backMap of + Nothing -> error + "mapKeysWith: bug: Data.Map.Strict.Refined has been subverted" + Just k1 -> k1 + where + ~backMap = Map.fromList + [ (k2, unsafeKey k1) + | k1 <- Map.keys m + , let !k2 = g $ unsafeKey k1 + ] + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect only the 'Just' results, returning a potentially smaller +-- map. +mapMaybeWithKey + :: forall s k a b. (Key s k -> a -> Maybe b) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k b +mapMaybeWithKey f (Map m) + = SomeMapWith (Map $ Map.mapMaybeWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Apply a function to all values in a map, together with their corresponding +-- keys, and collect the 'Left' and 'Right' results into separate (disjoint) +-- maps. +mapEitherWithKey + :: forall s k a b c. Ord k -- TODO: this is only used in the proof + => (Key s k -> a -> Either b c) + -> Map s k a + -> Some2MapWith (PartitionProof 'Regular s k) k b c +mapEitherWithKey p (Map m) = case Map.mapEitherWithKey (p . unsafeKey) m of + (m1, m2) -> Some2MapWith (Map m1) (Map m2) $ PartitionProof + do \k -> case Map.lookup (unrefine k) m of + Nothing -> error + "mapEitherWithKey: bug: Data.Map.Strict.Refined has been subverted" + Just x -> case p k x of + Left _ -> Left $ unsafeKey $ unrefine k + Right _ -> Right $ unsafeKey $ unrefine k + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Update or delete the value at the smallest key, returning a potentially +-- smaller map. +updateMinWithKey + :: forall s k a. (Key s k -> a -> Maybe a) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +updateMinWithKey f (Map m) + = SomeMapWith (Map $ Map.updateMinWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Update or delete the value at the largest key, returning a potentially +-- smaller map. +updateMaxWithKey + :: forall s k a. (Key s k -> a -> Maybe a) + -> Map s k a + -> SomeMapWith (SupersetProof 'Regular s) k a +updateMaxWithKey f (Map m) + = SomeMapWith (Map $ Map.updateMaxWithKey (f . unsafeKey) m) + $ SupersetProof unsafeSubset + +-- | Adjust the value at the smallest key. The set of keys remains unchanged. +adjustMinWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a +adjustMinWithKey f (Map m) + = Map $ Map.updateMinWithKey ((Just .) . f . unsafeKey) m + +-- | Adjust the value at the greatest key. The set of keys remains unchanged. +adjustMaxWithKey :: forall s k a. (Key s k -> a -> a) -> Map s k a -> Map s k a +adjustMaxWithKey f (Map m) + = Map $ Map.updateMaxWithKey ((Just .) . f . unsafeKey) m + +-- | @'bind' m f@ is a map that for each key @k :: 'Key' s k@, contains the +-- value @f (m '!' k) '!' k@, similar to @'>>='@ for functions. +bind :: forall s k a b. Ord k => Map s k a -> (a -> Map s k b) -> Map s k b +bind m f = mapWithKey (\k x -> f x ! k) m + +-- | Apply the inverse image of the given function to the keys of the given map, +-- so that for all @k :: 'Key' s2 k2@, +-- @'backpermuteKeys' f m '!' k = m '!' f k@. +-- +-- If maps are identified with functions, this computes the composition. +backpermuteKeys + :: forall s1 s2 k1 k2 a. (Ord k1, KnownSet s2 k2) + => (Key s2 k2 -> Key s1 k1) -> Map s1 k1 a -> Map s2 k2 a +backpermuteKeys f m = fromSet \k -> m ! f k diff --git a/src/Data/Set/Refined.hs b/src/Data/Set/Refined.hs new file mode 100644 index 0000000..2e52a01 --- /dev/null +++ b/src/Data/Set/Refined.hs @@ -0,0 +1,556 @@ +{-# LANGUAGE CPP #-} +-- | This module implements a way of tracking the contents of a 'Data.Set.Set' +-- at the type level, and provides utilities for manipulating such sets. +-- +-- The contents of a set are associated with a type parameter, e.g. @s@, so that +-- whenever you see the same type parameter, you know you are working with the +-- same set. The type @s@ itself has no internal structure, rather it is merely +-- a skolem type variable (rank-2 polymorphism 'Control.Monad.ST.runST' trick) +-- introduced by "Data.Reflection". +-- +-- = Warning +-- This module together with "Data.Set" rely on 'Eq' and 'Ord' instances being +-- lawful: that '==' is an equivalence relation, and that the 'Ord' operations +-- define a total order on the quotient defined by this equivalence relation; at +-- least for the subset of values that are actually encountered at runtime. If +-- this assumption is violated, this module may not be able to uphold its +-- invariants and may throw errors. In particular beware of NaN in 'Float' and +-- 'Double'. +module Data.Set.Refined + ( + -- * Set type + KnownSet + , Set + -- * Refinement type + , InSet(..) + , Flavor(Regular) + , Element + , revealPredicate + -- * Existentials and common proofs + , SomeSet(..) + , withSet + , SomeSetWith(..) + , withSetWith + , Some2SetWith(..) + , with2SetWith + , (:->) + , SupersetProof(..) + , EmptyProof(..) + -- * Construction + , empty + , singleton + , SingletonProof(..) + , fromSet + , fromTraversable + , FromTraversableProof(..) + -- * Insertion + , insert + , InsertProof(..) + -- * Deletion + , delete + -- * Query + , member + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , null + , isSubsetOf + , SubsetProof(..) + , disjoint + , DisjointProof(..) + -- * Combine + , union + , UnionProof(..) + , difference + , DifferenceProof(..) + , intersection + , IntersectionProof(..) + , cartesianProduct + , ProductProof(..) + , disjointUnion + , CoproductProof(..) + -- * Filter + , filter + , partition + , PartitionProof(..) + , spanAntitone + , PartialPartitionProof(..) + , splitMember + , SplitProof(..) + -- * Map + , map + , MapProof(..) + -- * Folds + , foldMap + , foldr + , foldl + , foldr' + , foldl' + -- * Min/Max + , minView + , maxView + -- * Conversion + , toList + , toDescList + , asIntSet + , asHashSet + -- * Casts + , castElement + , cast + , castFlavor + ) where + +import Data.Coerce +import Data.Constraint (Dict(..)) +import Data.Container.Refined.Conversion +import Data.Container.Refined.Hashable +import Data.Container.Refined.Proofs +import Data.Container.Refined.Unsafe +import qualified Data.Foldable as Foldable +import qualified Data.Map as Map +import Data.Proxy +import Data.Reflection +import qualified Data.Set as Set +import Data.Traversable +import Data.Type.Coercion +import Data.Type.Equality ((:~:)(..)) +import Data.Typeable (Typeable) +import GHC.Exts (Proxy#, proxy#) +import Prelude hiding (filter, foldl, foldMap, foldr, map, null) +import Refined +import Refined.Unsafe +import Unsafe.Coerce + +#if MIN_VERSION_containers(0, 5, 8) +#else +import qualified Data.List as List +#endif + + +-- | To use "Refined" machinery that uses the 'Predicate' typeclass you will +-- need to pattern match on this 'Dict'. +-- +-- The reason is that in the default /fast/ implementation of reflection, we +-- don't have @'Typeable' s@, which "Refined" wants for pretty-printing +-- exceptions. We /can/ provide @'Typeable' s@, but at the cost of using the +-- /slow/ implementation of reflection. +revealPredicate + :: forall s a. (Typeable a, Ord a, KnownSet s a) + => Dict (Predicate (InSet 'Regular s) a) +revealPredicate = reifyTypeable (reflect (Proxy @s)) + \(_ :: Proxy s') -> + reflect (Proxy @s') `seq` + -- ^ Work around https://github.com/ekmett/reflection/issues/54 + case unsafeCoerce Refl :: s :~: s' of + Refl -> Dict + +-- | @'Element' s a@ is a value of type @a@ that has been verified to be an +-- element of @s@. +-- +-- Thus, @'Element' s a@ is a \"refinement\" type of @a@, and this library +-- integrates with an implementation of refimenement types in "Refined", so +-- the machinery from there can be used to manipulate 'Element's (however see +-- 'revealPredicate'). +-- +-- The underlying @a@ value can be obtained with 'unrefine'. An @a@ can be +-- validated into an @'Element' s a@ with 'member'. +type Element s = Refined (InSet 'Regular s) + +unsafeCastElement :: forall s a. Coercion a (Element s a) +unsafeCastElement = reallyUnsafeUnderlyingRefined + +unsafeElement :: a -> Element s a +unsafeElement = coerceWith unsafeCastElement + +-- | An existential wrapper for an as-yet-unknown set. Pattern maching on it +-- gives you a way to refer to the set, e.g. +-- +-- @ +-- case 'fromSet' ... of +-- 'SomeSet' \@s _ -> doSomethingWith \@s +-- +-- case 'fromSet' ... of +-- 'SomeSet' (_ :: 'Proxy#' s) -> doSomethingWith \@s +-- @ +data SomeSet a where + SomeSet :: forall s a. KnownSet s a => Proxy# s -> SomeSet a + +-- | Apply an unknown set to a continuation that can accept any set. This gives +-- you a way to refer to the set (the parameter @s@), e.g.: +-- +-- @ +-- 'withSet' ('fromSet' ...) $ \(_ :: 'Proxy' s) -> doSomethingWith \@s +-- @ +withSet + :: forall a r. SomeSet a -> (forall s. KnownSet s a => Proxy s -> r) -> r +withSet (SomeSet (_ :: Proxy# s)) k = k $ Proxy @s + +-- | Construct a set from a regular 'Data.Set.Set'. +fromSet :: forall a. Set.Set a -> SomeSet a +fromSet s = reify s \(_ :: Proxy s) -> SomeSet @s proxy# + +-- | An existential wrapper for an as-yet-unknown set, together with a proof of +-- some fact @p@ about the set. Pattern matching on it gives you a way to refer +-- to the set (the parameter @s@). Most functions will return a set in this way, +-- together with a proof that somehow relates the set to the function's inputs. +data SomeSetWith p a where + SomeSetWith :: forall s a p. KnownSet s a => !(p s) -> SomeSetWith p a + +-- | Apply an unknown set with proof to a continuation that can accept any set +-- satisfying the proof. This gives you a way to refer to the set (the parameter +-- @s@). +withSetWith + :: forall a r p. SomeSetWith p a -> (forall s. KnownSet s a => p s -> r) -> r +withSetWith (SomeSetWith p) k = k p + +-- | An existential wrapper for an as-yet-unknown pair of sets, together with +-- a proof of some fact @p@ relating them. +data Some2SetWith p a where + Some2SetWith + :: forall s t a p. (KnownSet s a, KnownSet t a) + => !(p s t) -> Some2SetWith p a + +-- | Apply a pair of unknown sets with proof to a continuation that can accept +-- any pair of sets satisfying the proof. This gives you a way to refer to the +-- sets (the parameters @s@ and @t@). +with2SetWith + :: forall a r p. Some2SetWith p a + -> (forall s t. (KnownSet s a, KnownSet t a) => p s t -> r) + -> r +with2SetWith (Some2SetWith p) k = k p + +-- | An empty set. +empty :: forall a. SomeSetWith (EmptyProof 'Regular) a +empty = reify Set.empty \(_ :: Proxy r) + -> SomeSetWith @r $ EmptyProof unsafeSubset + +-- | Create a set with a single element. +singleton :: forall a. a -> SomeSetWith (SingletonProof 'Regular a) a +singleton x = reify (Set.singleton x) \(_ :: Proxy r) + -> SomeSetWith @r $ SingletonProof $ unsafeElement x + +-- | Create a set from the elements of an arbitrary traversable. +fromTraversable + :: forall t a. (Traversable t, Ord a) + => t a -> SomeSetWith (FromTraversableProof 'Regular t a) a +fromTraversable xs = reify set \(_ :: Proxy r) + -> SomeSetWith @r $ FromTraversableProof + $ unsafeCoerce @(t (Element _ a)) @(t (Element r a)) proof + where + (set, proof) = mapAccumL + (\s x -> let !s' = Set.insert x s in (s', unsafeElement x)) + Set.empty + xs + +-- | Insert an element in a set. +insert + :: forall s a. (Ord a, KnownSet s a) + => a -> SomeSetWith (InsertProof 'Regular a s) a +insert x = reify (Set.insert x $ reflect $ Proxy @s) \(_ :: Proxy r) + -> SomeSetWith @r $ InsertProof (unsafeElement x) unsafeSubset + +-- | Delete an element from a set. +delete + :: forall s a. (Ord a, KnownSet s a) + => a -> SomeSetWith (SupersetProof 'Regular s) a +delete x = reify (Set.delete x $ reflect $ Proxy @s) \(_ :: Proxy r) + -> SomeSetWith @s $ SupersetProof unsafeSubset + +-- | If an element is in the set, return the proof that it is. +member :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a) +member x + | x `Set.member` reflect (Proxy @s) = Just $ unsafeElement x + | otherwise = Nothing + +-- | Find the largest element smaller than the given one. +lookupLT :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a) +lookupLT x = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ Set.lookupLT x (reflect $ Proxy @s) + +-- | Find the smallest element greater than the given one. +lookupGT :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a) +lookupGT x = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ Set.lookupGT x (reflect $ Proxy @s) + +-- | Find the largest element smaller or equal to the given one. +lookupLE :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a) +lookupLE x = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ Set.lookupLE x (reflect $ Proxy @s) + +-- | Find the smallest element greater or equal to the given one. +lookupGE :: forall s a. (Ord a, KnownSet s a) => a -> Maybe (Element s a) +lookupGE x = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ Set.lookupGE x (reflect $ Proxy @s) + +-- | If the set is empty, return the proof that it is. +null :: forall s a. KnownSet s a => Maybe (EmptyProof 'Regular s) +null + | Set.null $ reflect $ Proxy @s = Just $ EmptyProof unsafeSubset + | otherwise = Nothing + +-- | If @s@ is a subset of @t@ (or is equal to), return a proof of that. +isSubsetOf + :: forall s t a. (Ord a, KnownSet s a, KnownSet t a) + => Maybe (SubsetProof 'Regular s t) +isSubsetOf + | reflect (Proxy @s) `Set.isSubsetOf` reflect (Proxy @t) + = Just $ SubsetProof unsafeSubset + | otherwise = Nothing + +-- | If @s@ and @t@ are disjoint (i.e. their intersection is empty), return a +-- proof of that. +disjoint + :: forall s t a. (Ord a, KnownSet s a, KnownSet t a) + => Maybe (DisjointProof 'Regular s t) +disjoint +#if MIN_VERSION_containers(0, 5, 11) + | Set.disjoint (reflect $ Proxy @s) (reflect $ Proxy @t) +#else + | Set.null $ Set.intersection (reflect $ Proxy @s) (reflect $ Proxy @t) +#endif + = Just $ DisjointProof \f g -> unsafeSubsetWith2 f g + | otherwise = Nothing + +-- | The union of two sets. +union + :: forall s t a. (Ord a, KnownSet s a, KnownSet t a) + => SomeSetWith (UnionProof 'Regular s t) a +union = reify (reflect (Proxy @s) `Set.union` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeSetWith @r $ UnionProof unsafeSubset unsafeSubsetWith2 + +-- unions :: ? + +-- | Set with elements of @s@ that are not in @t@. +difference + :: forall s t a. (Ord a, KnownSet s a, KnownSet t a) + => SomeSetWith (DifferenceProof 'Regular s t) a +difference = reify (reflect (Proxy @s) `Set.difference` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeSetWith @r + $ DifferenceProof unsafeSubset (\f g -> unsafeSubsetWith2 f g) unsafeSubset + +-- | Intersection of two sets. +intersection + :: forall s t a. (Ord a, KnownSet s a, KnownSet t a) + => SomeSetWith (IntersectionProof 'Regular s t) a +intersection = reify (reflect (Proxy @s) `Set.intersection` reflect (Proxy @t)) + \(_ :: Proxy r) -> SomeSetWith @r + $ IntersectionProof unsafeSubset unsafeSubsetWith2 + +-- | Cartesian product of two sets. The elements are all pairs @(x, y)@ for each +-- @x@ from @s@ and for each @y@ from @t@. +cartesianProduct + :: forall s t a b. (KnownSet s a, KnownSet t b) + => SomeSetWith (ProductProof 'Regular s t) (a, b) +cartesianProduct = reify +#if MIN_VERSION_containers(0, 5, 11) + (reflect (Proxy @s) `Set.cartesianProduct` reflect (Proxy @t)) +#else + (Set.fromDistinctAscList $ (,) <$> Set.toAscList (reflect $ Proxy @s) + <*> Set.toAscList (reflect $ Proxy @t)) +#endif + \(_ :: Proxy r) -> SomeSetWith @r $ ProductProof let + proof :: forall x y. Coercion + (Refined (InSet 'Regular s) x, Refined (InSet 'Regular t) y) + (Refined (InSet 'Regular r) (x, y)) + !proof + | Coercion <- reallyUnsafeUnderlyingRefined @x @(InSet 'Regular s) + , Coercion <- reallyUnsafeUnderlyingRefined @y @(InSet 'Regular t) + = Coercion `trans` + reallyUnsafeUnderlyingRefined @(x, y) @(InSet 'Regular r) + in proof + +-- | Disjoint union of two sets. Includes @'Left' x@ for each @x@ from @s@, and +-- @'Right' y@ for each @y@ from @t@. +disjointUnion + :: forall s t a b. (KnownSet s a, KnownSet t b) + => SomeSetWith (CoproductProof 'Regular s t) (Either a b) +disjointUnion = reify +#if MIN_VERSION_containers(0, 5, 11) + (reflect (Proxy @s) `Set.disjointUnion` reflect (Proxy @t)) +#else + (Set.fromDistinctAscList $ (Left <$> Set.toAscList (reflect $ Proxy @s)) + ++ (Right <$> Set.toAscList (reflect $ Proxy @t))) +#endif + \(_ :: Proxy r) -> SomeSetWith @r $ CoproductProof let + proof :: forall x y. Coercion + (Either (Refined (InSet 'Regular s) x) (Refined (InSet 'Regular t) y)) + (Refined (InSet 'Regular r) (Either x y)) + !proof + | Coercion <- reallyUnsafeUnderlyingRefined @x @(InSet 'Regular s) + , Coercion <- reallyUnsafeUnderlyingRefined @y @(InSet 'Regular t) + = Coercion `trans` + reallyUnsafeUnderlyingRefined @(Either x y) @(InSet 'Regular r) + in proof + +-- | Return a subset of elements that satisfy the given predicate. +filter + :: forall s a. KnownSet s a + => (Element s a -> Bool) -> SomeSetWith (SupersetProof 'Regular s) a +filter p = reify (Set.filter (p . unsafeElement) $ reflect $ Proxy @s) + \(_ :: Proxy r) -> SomeSetWith @r $ SupersetProof unsafeSubset + +-- | Partition a set into two disjoint subsets: those that satisfy the +-- predicate, and those that don't. +partition + :: forall s a. KnownSet s a + => (Element s a -> Bool) -> Some2SetWith (PartitionProof 'Regular s a) a +partition p = case Set.partition (p . unsafeElement) $ reflect $ Proxy @s of + (r, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q) + -> Some2SetWith @s @r $ PartitionProof + do \x -> if p x + then Left $ unsafeElement $ unrefine x + else Right $ unsafeElement $ unrefine x + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Divide a set into two disjoint subsets at a point where the predicate stops +-- holding. +-- +-- If @p@ is antitone ( \(\forall x y, x < y \implies p(x) \ge p(y)\) ), then +-- this point is uniquely defined. If @p@ is not antitone, a splitting point is +-- chosen in an unspecified way. +spanAntitone + :: forall s a. KnownSet s a + => (Element s a -> Bool) -> Some2SetWith (PartialPartitionProof 'Regular s) a +spanAntitone p = +#if MIN_VERSION_containers(0, 5, 8) + case Set.spanAntitone (p . unsafeElement) $ reflect $ Proxy @s of + (r, q) +#else + case List.span (p . unsafeElement) + $ Set.toAscList $ reflect $ Proxy @s of + (rs, qs) + | let r = Set.fromDistinctAscList rs + , let q = Set.fromDistinctAscList qs +#endif + -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q) + -> Some2SetWith @r @q $ PartialPartitionProof + unsafeSubset unsafeSubsetWith2 \f g -> unsafeSubsetWith2 f g + +-- | Return two disjoint subsets: those less than the given element, and those +-- greater than the given element; along with the proof that the given element +-- was in the set, if it was. +splitMember + :: forall s a. (Ord a, KnownSet s a) + => a -> Some2SetWith (SplitProof 'Regular s (Element s a)) a +splitMember x = case Set.splitMember x $ reflect $ Proxy @s of + (r, m, q) -> reify r \(_ :: Proxy r) -> reify q \(_ :: Proxy q) + -> Some2SetWith @r @q $ SplitProof + (if m then Just (unsafeElement x) else Nothing) + unsafeSubset \f g -> unsafeSubsetWith2 f g + +-- | Apply the given function to each element of the set and collect the +-- results. Note that the resulting set can be smaller. +map + :: forall s a b. (Ord b, KnownSet s a) + => (Element s a -> b) -> SomeSetWith (MapProof 'Regular s a b) b +map f = reify (Map.keysSet m) \(_ :: Proxy r) -> SomeSetWith @r + $ MapProof (unsafeElement . f) \y -> case Map.lookup (unrefine y) m of + Nothing -> error "map: bug: Data.Set.Refined has been subverted" + Just x -> x + where + !m = Map.fromList + [ (y, unsafeElement x) + | x <- Set.toList $ reflect $ Proxy @s + , let !y = f $ unsafeElement x + ] + +-- | Map each element of @s@ into a monoid (with proof that it was an element), +-- and combine the results using 'Data.Monoid.<>'. +foldMap :: forall s a m. (KnownSet s a, Monoid m) => (Element s a -> m) -> m +foldMap f = Foldable.foldMap (f . unsafeElement) $ reflect $ Proxy @s + +-- | Right associative fold with a lazy accumulator. +foldr :: forall s a b. KnownSet s a => (Element s a -> b -> b) -> b -> b +foldr f z = Set.foldr (f . unsafeElement) z $ reflect $ Proxy @s + +-- | Left associative fold with a lazy accumulator. +foldl :: forall s a b. KnownSet s a => (b -> Element s a -> b) -> b -> b +foldl f z = Set.foldl ((. unsafeElement) . f) z $ reflect $ Proxy @s + +-- | Right associative fold with a strict accumulator. +foldr' :: forall s a b. KnownSet s a => (Element s a -> b -> b) -> b -> b +foldr' f z = Set.foldr' (f . unsafeElement) z $ reflect $ Proxy @s + +-- | Left associative fold with a strict accumulator. +foldl' :: forall s a b. KnownSet s a => (b -> Element s a -> b) -> b -> b +foldl' f z = Set.foldl' ((. unsafeElement) . f) z $ reflect $ Proxy @s + +-- | Retrieves the smallest element of the set, and the set with that element +-- removed; or a proof that the set was empty. +minView + :: forall s a. KnownSet s a + => Either + (EmptyProof 'Regular s) + (Element s a, SomeSetWith (SupersetProof 'Regular s) a) +minView = case Set.minView $ reflect $ Proxy @s of + Nothing -> Left $ EmptyProof unsafeSubset + Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r) + -> SomeSetWith @r $ SupersetProof unsafeSubset + +-- | Retrieves the greatest element of the set, and the set with that element +-- removed; or a proof that the set was empty. +maxView + :: forall s a. KnownSet s a + => Either + (EmptyProof 'Regular s) + (Element s a, SomeSetWith (SupersetProof 'Regular s) a) +maxView = case Set.maxView $ reflect $ Proxy @s of + Nothing -> Left $ EmptyProof unsafeSubset + Just (x, xs) -> Right $ (unsafeElement x,) $ reify xs \(_ :: Proxy r) + -> SomeSetWith @r $ SupersetProof unsafeSubset + +-- | List of elements in the set in ascending order. +toList :: forall s a. KnownSet s a => [Element s a] +toList = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ Set.toAscList $ reflect $ Proxy @s + +-- | List of elements in the set in descending order. +toDescList :: forall s a. KnownSet s a => [Element s a] +toDescList = gcoerceWith (unsafeCastElement @s @a) $ coerce + $ Set.toDescList $ reflect $ Proxy @s + +-- | Convert a 'Set' into an 'IntSet', retaining its set of elements, which can +-- be converted with 'castFlavor'. +asIntSet :: forall s. KnownSet s Int => IntSet s +asIntSet = set2IntSet + +-- | Convert a 'Set' into a 'HashSet', retaining its set of elements, which can +-- be converted with 'castFlavor'. +asHashSet :: forall s a. (Hashable a, KnownSet s a) => HashSet s a +asHashSet = set2HashSet + +-- | If elements of @s@ can be weakened to elements of @t@ and vice versa, then +-- @s@ and @t@ actually stand for the same set and @'Element' s@ can be safely +-- interconverted with @'Element' t@. +-- +-- The requirement that the weakenings are natural transformations ensures that +-- they don't actually alter the elements. To build these you can compose +-- ':->''s from proofs returned by functions in this module, or "Refined" +-- functions like 'andLeft' or 'leftOr'. +castElement + :: forall s t a. (forall x. Element s x -> Element t x) + -> (forall x. Element t x -> Element s x) + -> Coercion (Element s a) (Element t a) +castElement = castRefined + +-- | If elements can be interconverted (e.g. as proved by 'castElement'), then +-- the sets can be interconverted too. For example we can establish that the +-- intersection of a set with itself is interconvertible with that set: +-- +-- @ +-- castIntersection +-- :: t'IntersectionProof' ''Regular' s s r +-- -> 'Coercion' ('Set' r a) ('Set' s a) +-- castIntersection ( v'IntersectionProof' p1 p2) +-- = 'cast' $ 'castElement' ('andLeft' . p1) (p2 'id' 'id') +-- @ +cast + :: forall s t a. (forall x. Coercion (Element s x) (Element t x)) + -> Coercion (Set s a) (Set t a) +cast Coercion +#if MIN_VERSION_base(4, 15, 0) + = case unsafeEqualityProof @s @t of UnsafeRefl -> Coercion +#else + = repr $ unsafeCoerce Refl +#endif