dependent-enummap-0.2.0.0: src/Data/Dependent/EnumMap/Strict/Internal.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Dependent.EnumMap.Strict.Internal where
import Prelude hiding (lookup, map)
import Control.Exception
import Control.Monad ((<$!>))
import Data.Bifunctor (bimap, second)
import Data.Coerce
import Data.Dependent.Sum
import Data.Foldable qualified as Foldable
import Data.IntMap.Strict qualified as IM
import Data.Kind (Type)
import Data.Proxy
import Data.Some
import Data.Type.Equality
import Text.Show (showListWith)
import Unsafe.Coerce (unsafeCoerce)
type KV :: forall kind. (kind -> Type) -> (kind -> Type) -> Type
data KV k v = forall a. KV !(Enum1Info k a) !(v a)
-- Invariant ensured by types: the key-value pairs in a DEnumMap
-- are type-consistent. That is to say: they have the same type-index.
-- Any other type equalities, like between the key argument to 'lookup'
-- and the key-value pairs in the map argument to 'lookup',
-- may /not/ hold, and should be runtime-type-checked as much as we're able.
newtype DEnumMap k v = DEnumMap (IM.IntMap (KV k v))
instance (Enum1 k, forall a. Show (k a), forall a. Show (v a))
=> Show (DEnumMap k v) where
showsPrec d mp = showParen (d > 10) $
showString "fromList " . showListWith (\(k :=> v) -> showsPrec 2 k . showString " :=> " . showsPrec 1 v) (toList mp)
-- | This class attempts to generalise 'Enum' to indexed data types: data types
-- with a GADT-like type parameter. Conversion to an 'Int' naturally loses type
-- information, and furthermore it is common to actually need some additional
-- data alongside the 'Int' to be able to reconstruct the original (in
-- 'toEnum1'). This additional data lives in 'Enum1Info'. The laws are:
--
-- [Unique IDs]
-- If @'fst' ('fromEnum1' x) == 'fst' ('fromEnum1' y)@ then @'testEquality' x y == 'Just' 'Refl' && x '==' y@
-- [Persistent IDs]
-- @'fst' ('fromEnum1' ('uncurry' 'toEnum1' ('fromEnum1' x))) == 'fst' ('fromEnum1' x)@
--
-- The "Unique IDs" law states that if the IDs of two values are equal, then
-- the values themselves must have the same type index, and furthermore be
-- equal. If @f@ does not implement 'TestEquality' or 'Eq', the law should
-- morally hold (but most of the API will be unusable).
--
-- The "Persistent IDs" law states that reconstructing a value using 'toEnum1'
-- does not change its ID.
--
-- __Note__: The methods on 'DEnumMap' attempt to check these laws using
-- 'assert' assertions (which are by default __disabled__ when optimisations
-- are on!), but full consistency cannot always be checked;
-- __if you break these laws in a sufficiently clever way, the internals of 'DEnumMap' may 'unsafeCoerce' unequal things and engage nasal demons, including crashes and worse.__
--
-- To enable assertions in optimized builds, e.g., for test runs
-- using optimized code, add the following to the cabal.project.local
-- file of your project:
--
-- > package dependent-enummap
-- > ghc-options: -fno-ignore-asserts
class Enum1 f where
type Enum1Info f :: kind -> Type
fromEnum1 :: f a -> (Int, Enum1Info f a)
toEnum1 :: Int -> Enum1Info f a -> f a
dSumToKV :: Enum1 k => DSum k v -> (Int, KV k v)
dSumToKV (k :=> v) = let (i, inf) = fromEnum1 k in (i, KV inf v)
-- | Assumes that the input was obtained via 'fromEnum1'.
kVToDSum :: Enum1 k => (Int, KV k v) -> DSum k v
kVToDSum (i, KV inf v) = toEnum1 i inf :=> v
-- * Construction
empty :: DEnumMap k v
empty = DEnumMap IM.empty
singleton :: Enum1 k => k a -> v a -> DEnumMap k v
singleton k v =
let (i, inf) = fromEnum1 k
in DEnumMap (IM.singleton i (KV inf v))
-- TODO: Wait for DEnumSet.
-- fromSet
-- ** From Unordered Lists
fromList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromList l = DEnumMap (IM.fromList (dSumToKV <$> l))
fromListWith :: (Enum1 k, TestEquality k)
=> (forall a. v a -> v a -> v a)
-> [DSum k v] -> DEnumMap k v
fromListWith f (l :: [DSum k v]) =
DEnumMap (IM.fromListWithKey
(\i (KV inf1 v1) (KV inf2 v2) ->
case typeCheck2 (Proxy @k) i inf1 inf2 of
Refl -> KV inf1 (f v1 v2))
(dSumToKV <$> l))
fromListWithKey :: forall k v. (Enum1 k, TestEquality k)
=> (forall a. k a -> v a -> v a -> v a)
-> [DSum k v] -> DEnumMap k v
fromListWithKey f l =
DEnumMap (IM.fromListWithKey
(\i (KV inf1 v1) (KV inf2 v2) ->
let k1 = toEnum1 @k i inf1
in case typeCheck1 k1 i inf2 of
Refl -> KV inf1 (f k1 v1 v2))
(dSumToKV <$> l))
-- ** From Ascending Lists
fromAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromAscList l = DEnumMap (IM.fromAscList (dSumToKV <$> l))
fromAscListWith :: (Enum1 k, TestEquality k)
=> (forall a. v a -> v a -> v a)
-> [DSum k v] -> DEnumMap k v
fromAscListWith f (l :: [DSum k v]) =
DEnumMap (IM.fromAscListWithKey
(\i (KV inf1 v1) (KV inf2 v2) ->
case typeCheck2 (Proxy @k) i inf1 inf2 of
Refl -> KV inf1 (f v1 v2))
(dSumToKV <$> l))
fromAscListWithKey :: forall k v. (Enum1 k, TestEquality k)
=> (forall a. k a -> v a -> v a -> v a)
-> [DSum k v] -> DEnumMap k v
fromAscListWithKey f l =
DEnumMap (IM.fromAscListWithKey
(\i (KV inf1 v1) (KV inf2 v2) ->
let k1 = toEnum1 @k i inf1
in case typeCheck1 k1 i inf2 of
Refl -> KV inf1 (f k1 v1 v2))
(dSumToKV <$> l))
fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromDistinctAscList l = DEnumMap (IM.fromDistinctAscList (dSumToKV <$> l))
-- * Insertion
insert :: Enum1 k => k a -> v a -> DEnumMap k v -> DEnumMap k v
insert k v (DEnumMap m) =
let (i, inf) = fromEnum1 k
in DEnumMap (IM.insert i (KV inf v) m)
insertWith :: (Enum1 k, TestEquality k)
=> (v a -> v a -> v a)
-> k a -> v a -> DEnumMap k v -> DEnumMap k v
insertWith = insertWithKey . const
insertWithKey :: (Enum1 k, TestEquality k)
=> (k a -> v a -> v a -> v a)
-> k a -> v a -> DEnumMap k v -> DEnumMap k v
insertWithKey f k v (DEnumMap m) =
let (i, inf) = fromEnum1 k
in DEnumMap (IM.insertWith
(\_ (KV inf' v2) -> case typeCheck1 k i inf' of
Refl -> KV inf (f k v v2))
i (KV inf v) m)
insertLookupWithKey :: (Enum1 k, TestEquality k)
=> (k a -> v a -> v a -> v a)
-> k a -> v a -> DEnumMap k v -> (Maybe (v a), DEnumMap k v)
insertLookupWithKey f k v (DEnumMap m) =
let (i, inf) = fromEnum1 k
(!mx, !m') =
IM.insertLookupWithKey
(\_ _ (KV inf' v2) -> case typeCheck1 k i inf' of
Refl -> KV inf (f k v v2))
i (KV inf v) m
-- Note: type checking unnecessary here, because by the BangPatterns,
-- evaluating mx evaluates dmap, and the IntMap is strict, so the lambda
-- will have run and typechecked the old value already.
-- Second note: the BangPatterns don't do anything operationally because
-- with the current implementation of IM.insertLookupWithKey, the pair
-- components are already strict.
in ((\(KV _ v2) -> coe1 v2) <$!> mx, DEnumMap m')
-- * Deletion\/Update
delete :: Enum1 k => k a -> DEnumMap k v -> DEnumMap k v
delete k (DEnumMap m) = DEnumMap (IM.delete (fst (fromEnum1 k)) m)
adjust :: (Enum1 k, TestEquality k) => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjust = adjustWithKey . const
adjustWithKey :: (Enum1 k, TestEquality k) => (k a -> v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjustWithKey f k (DEnumMap m) =
let (i, _) = fromEnum1 k
in DEnumMap (IM.adjust (\(KV inf v) -> case typeCheck1 k i inf of
Refl -> KV inf (f k v))
i m)
update :: (Enum1 k, TestEquality k) => (v a -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
update = updateWithKey . const
updateWithKey :: (Enum1 k, TestEquality k)
=> (k a -> v a -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
updateWithKey f k (DEnumMap m) =
let (i, _) = fromEnum1 k
in DEnumMap (IM.update (\(KV inf v) -> case typeCheck1 k i inf of
Refl -> KV inf <$> f k v)
i m)
updateLookupWithKey :: (Enum1 k, TestEquality k)
=> (k a -> v a -> Maybe (v a)) -> k a -> DEnumMap k v -> (Maybe (v a), DEnumMap k v)
updateLookupWithKey f k (DEnumMap m) =
let (i, _) = fromEnum1 k
(!mx, !m') =
IM.updateLookupWithKey
(\_ (KV inf v) -> case typeCheck1 k i inf of
Refl -> KV inf <$> f k v)
i m
-- Note: type checking unnecessary here for the same reason as insertLookupWithKey
in ((\(KV _ v2) -> coe1 v2) <$!> mx, DEnumMap m')
alter :: forall k v a. (Enum1 k, TestEquality k)
=> (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
alter f k (DEnumMap m) = DEnumMap (IM.alter f' i m)
where
(i, inf) = fromEnum1 k
f' :: Maybe (KV k v) -> Maybe (KV k v)
f' Nothing = KV inf <$> f Nothing
f' (Just (KV inf' v)) = case typeCheck1 k i inf' of
Refl -> KV inf <$> f (Just v)
alterF :: forall k v a f. (Functor f, Enum1 k, TestEquality k)
=> (Maybe (v a) -> f (Maybe (v a))) -> k a -> DEnumMap k v -> f (DEnumMap k v)
alterF f k (DEnumMap m) = DEnumMap <$> IM.alterF f' i m
where
(i, inf) = fromEnum1 k
f' :: Maybe (KV k v) -> f (Maybe (KV k v))
f' Nothing = fmap (KV inf) <$> f Nothing
f' (Just (KV inf' v)) = case typeCheck1 k i inf' of
Refl -> fmap (KV inf) <$> f (Just v)
-- * Query
-- ** Lookup
lookup :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (v a)
{-# INLINEABLE lookup #-}
lookup k (DEnumMap m) =
let (i, _) = fromEnum1 k
in (\(KV inf v) -> case typeCheck1 k i inf of
Refl -> v)
<$!> IM.lookup i m
(!?) :: (Enum1 k, TestEquality k) => DEnumMap k v -> k a -> Maybe (v a)
(!?) m k = lookup k m
findWithDefault :: (Enum1 k, TestEquality k) => v a -> k a -> DEnumMap k v -> v a
findWithDefault def k (DEnumMap m) =
let (i, _) = fromEnum1 k
in case IM.findWithDefault (KV undefined def) i m of
KV inf' v -> case typeCheck1 k i inf' of
Refl -> v
find :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> v a
find k = findWithDefault (error ("Data.Dependent.EnumMap.!: key " ++ show (fst (fromEnum1 k)) ++ " is not an element of the map")) k
(!) :: (Enum1 k, TestEquality k) => DEnumMap k v -> k a -> v a
(!) m k = find k m
member :: Enum1 k => k a -> DEnumMap k v -> Bool
member k (DEnumMap m) = IM.member (fst (fromEnum1 k)) m
notMember :: Enum1 k => k a -> DEnumMap k v -> Bool
notMember k m = not $ member k m
lookupLT, lookupGT, lookupLE, lookupGE :: Enum1 k => k a -> DEnumMap k v -> Maybe (DSum k v)
lookupLT k (DEnumMap m) = let (i, _) = fromEnum1 k in kVToDSum <$!> IM.lookupLT i m
lookupGT k (DEnumMap m) = let (i, _) = fromEnum1 k in kVToDSum <$!> IM.lookupGT i m
lookupLE k (DEnumMap m) = let (i, _) = fromEnum1 k in kVToDSum <$!> IM.lookupLE i m
lookupGE k (DEnumMap m) = let (i, _) = fromEnum1 k in kVToDSum <$!> IM.lookupGE i m
-- ** Size
null :: DEnumMap k v -> Bool
null (DEnumMap m) = IM.null m
size :: DEnumMap k v -> Int
size (DEnumMap m) = IM.size m
-- * Combine
-- ** Union
union :: (Enum1 k, TestEquality k) => DEnumMap k v -> DEnumMap k v -> DEnumMap k v
union = unionWith const -- if we're type checking, we need unionWith anyway, so might as well just delegate here already
unionWith :: (Enum1 k, TestEquality k)
=> (forall a. v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWith f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWithKey f' m1 m2)
where
f' :: Int -> KV k v -> KV k v -> KV k v
f' i (KV inf1 v1) (KV inf2 v2) = case typeCheck2 (Proxy @k) i inf1 inf2 of
Refl -> KV inf1 (f v1 v2)
unionWithKey :: forall k v. (Enum1 k, TestEquality k)
=> (forall a. k a -> v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWithKey f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWithKey f' m1 m2)
where
f' :: Int -> KV k v -> KV k v -> KV k v
f' i (KV inf1 v1) (KV inf2 v2) =
let k1 = toEnum1 @k i inf1
in case typeCheck1 k1 i inf2 of
Refl -> KV inf1 (f k1 v1 v2)
unions :: (Foldable f, Enum1 k, TestEquality k) => f (DEnumMap k v) -> DEnumMap k v
unions = Foldable.foldl' union empty
unionsWith :: (Foldable f, Enum1 k, TestEquality k)
=> (forall a. v a -> v a -> v a) -> f (DEnumMap k v) -> DEnumMap k v
unionsWith f = Foldable.foldl' (unionWith f) empty
-- ** Difference
difference :: DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
difference (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.difference m1 m2)
(\\) :: DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
m1 \\ m2 = difference m1 m2
differenceWith :: forall k v1 v2. (Enum1 k, TestEquality k)
=> (forall a. v1 a -> v2 a -> Maybe (v1 a)) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
differenceWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
where
f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v1)
f' i (KV inf1 v1) (KV inf2 v2) =
case typeCheck2 (Proxy @k) i inf1 inf2 of
Refl -> KV inf1 <$> f v1 v2
differenceWithKey :: forall k v1 v2. (Enum1 k, TestEquality k)
=> (forall a. k a -> v1 a -> v2 a -> Maybe (v1 a)) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
differenceWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
where
f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v1)
f' i (KV inf1 v1) (KV inf2 v2) =
let k1 = toEnum1 @k i inf1
in case typeCheck1 k1 i inf2 of
Refl -> KV inf1 <$> f k1 v1 v2
-- | Because the underlying maps are keyed on integers, it is possible to
-- subtract a map from another even if the key types differ. This function
-- assumes that the @Int@ identifiers of @k1@ and @k2@ are compatible, i.e.
-- that "2" in @k1@ somehow means the same thing as "2" in @k2@.
--
-- Because the key types are different, there is no guarantee whatsoever (even
-- not by 'Enum1' laws) that equal key IDs in @k1@ and @k2@ actually have the
-- same type index (@a@). Hence, the combining function gets key-value pairs
-- with potentially distinct type indices.
differenceWithKey' :: forall k1 k2 v1 v2. (Enum1 k1, Enum1 k2)
=> (forall a b. k1 a -> v1 a -> k2 b -> v2 b -> Maybe (v1 a))
-> DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
differenceWithKey' f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
where
f' :: Int -> KV k1 v1 -> KV k2 v2 -> Maybe (KV k1 v1)
f' i (KV inf1 v1) (KV inf2 v2) =
let k1 = toEnum1 i inf1
k2 = toEnum1 i inf2
in KV inf1 <$> f k1 v1 k2 v2
-- ** Intersection
intersection :: DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
intersection (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersection m1 m2)
intersectionWith :: forall k v1 v2 v3. (Enum1 k, TestEquality k)
=> (forall a. v1 a -> v2 a -> v3 a) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3
intersectionWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2)
where
f' :: Int -> KV k v1 -> KV k v2 -> KV k v3
f' i (KV inf1 v1) (KV inf2 v2) =
case typeCheck2 (Proxy @k) i inf1 inf2 of
Refl -> KV inf1 $ f v1 v2
intersectionWithKey :: forall k v1 v2 v3. (Enum1 k, TestEquality k)
=> (forall a. k a -> v1 a -> v2 a -> v3 a) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3
intersectionWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2)
where
f' :: Int -> KV k v1 -> KV k v2 -> KV k v3
f' i (KV inf1 v1) (KV inf2 v2) =
let k1 = toEnum1 @k i inf1
in case typeCheck1 k1 i inf2 of
Refl -> KV inf1 $ f k1 v1 v2
-- | Generalises 'intersectionWithKey' in the same way as 'differenceWithKey''
-- generalises 'differenceWithKey'.
intersectionWithKey' :: forall k1 k2 v1 v2 v3. (Enum1 k1, Enum1 k2)
=> (forall a b. k1 a -> v1 a -> k2 b -> v2 b -> v3 a)
-> DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v3
intersectionWithKey' f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2)
where
f' :: Int -> KV k1 v1 -> KV k2 v2 -> KV k1 v3
f' i (KV inf1 v1) (KV inf2 v2) =
let k1 = toEnum1 i inf1
k2 = toEnum1 i inf2
in KV inf1 $ f k1 v1 k2 v2
-- ** Disjoint
disjoint :: DEnumMap k v1 -> DEnumMap k v2 -> Bool
disjoint (DEnumMap m1) (DEnumMap m2) = IM.disjoint m1 m2
-- ** Compose
compose :: (Enum1 k2, TestEquality k2) => DEnumMap k2 v -> DEnumMap k1 k2 -> DEnumMap k1 v
compose m2v (DEnumMap m12) =
DEnumMap (IM.mapMaybe (\(KV inf1 k2) -> KV inf1 <$> m2v !? k2) m12)
-- ** Universal combining function
mergeWithKey :: forall k v1 v2 v3. (Enum1 k, TestEquality k)
=> (forall a. k a -> v1 a -> v2 a -> Maybe (v3 a))
-> (DEnumMap k v1 -> DEnumMap k v3)
-> (DEnumMap k v2 -> DEnumMap k v3)
-> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3
mergeWithKey f g1 g2 (DEnumMap m1) (DEnumMap m2) =
DEnumMap (IM.mergeWithKey f' (coerce g1) (coerce g2) m1 m2)
where
f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v3)
f' i (KV inf1 v1) (KV inf2 v2) =
let k1 = toEnum1 @k i inf1
in case typeCheck1 k1 i inf2 of
Refl -> KV inf1 <$> f k1 v1 v2
-- * Traversal
-- ** Map
map :: Enum1 k => (forall a. v1 a -> v2 a) -> DEnumMap k v1 -> DEnumMap k v2
map f = mapWithKey (const f)
mapWithKey :: Enum1 k => (forall a. k a -> v1 a -> v2 a) -> DEnumMap k v1 -> DEnumMap k v2
mapWithKey f (DEnumMap m) =
DEnumMap (IM.mapWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in KV inf $ f k v)
m)
traverseWithKey :: (Applicative f, Enum1 k)
=> (forall a. k a -> v1 a -> f (v2 a)) -> DEnumMap k v1 -> f (DEnumMap k v2)
traverseWithKey f (DEnumMap m) =
DEnumMap <$> IM.traverseWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in KV inf <$> f k v)
m
traverseMaybeWithKey :: (Applicative f, Enum1 k)
=> (forall a. k a -> v1 a -> f (Maybe (v2 a))) -> DEnumMap k v1 -> f (DEnumMap k v2)
traverseMaybeWithKey f (DEnumMap m) =
DEnumMap <$> IM.traverseMaybeWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in fmap (KV inf) <$> f k v)
m
mapAccum :: Enum1 k => (forall a. acc -> v1 a -> (acc, v2 a)) -> acc -> DEnumMap k v1 -> (acc, DEnumMap k v2)
mapAccum f = mapAccumWithKey (\x _ y -> f x y)
mapAccumWithKey :: Enum1 k => (forall a. acc -> k a -> v1 a -> (acc, v2 a)) -> acc -> DEnumMap k v1 -> (acc, DEnumMap k v2)
mapAccumWithKey f acc0 (DEnumMap m) =
second DEnumMap $ IM.mapAccumWithKey (\acc i (KV inf v) ->
let k = toEnum1 i inf
in second (KV inf) $ f acc k v)
acc0 m
mapAccumRWithKey :: Enum1 k => (forall a. acc -> k a -> v1 a -> (acc, v2 a)) -> acc -> DEnumMap k v1 -> (acc, DEnumMap k v2)
mapAccumRWithKey f acc0 (DEnumMap m) =
second DEnumMap $ IM.mapAccumRWithKey (\acc i (KV inf v) ->
let k = toEnum1 i inf
in second (KV inf) $ f acc k v)
acc0 m
-- TODO: These are hard. Probably we can't avoid using a fold, analogously as in IntMap.
-- mapKeys
-- mapKeysWith
-- mapKeysMonotonic
-- * Folds
foldr :: (forall a. v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldr f acc0 (DEnumMap m) = IM.foldr (\(KV _ v) acc -> f v acc) acc0 m
foldl :: (forall a. acc -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldl f acc0 (DEnumMap m) = IM.foldl (\acc (KV _ v) -> f acc v) acc0 m
foldrWithKey :: Enum1 k => (forall a. k a -> v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldrWithKey f acc0 (DEnumMap m) =
IM.foldrWithKey (\i (KV inf v) acc ->
let k = toEnum1 i inf
in f k v acc)
acc0 m
foldlWithKey :: Enum1 k => (forall a. acc -> k a -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldlWithKey f acc0 (DEnumMap m) =
IM.foldlWithKey (\acc i (KV inf v) ->
let k = toEnum1 i inf
in f acc k v)
acc0 m
foldMapWithKey :: (Monoid m, Enum1 k) => (forall a. k a -> v a -> m) -> DEnumMap k v -> m
foldMapWithKey f (DEnumMap m) =
IM.foldMapWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in f k v)
m
-- ** Strict folds
foldr' :: (forall a. v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldr' f acc0 (DEnumMap m) = IM.foldr' (\(KV _ v) acc -> f v acc) acc0 m
foldl' :: (forall a. acc -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldl' f acc0 (DEnumMap m) = IM.foldl' (\acc (KV _ v) -> f acc v) acc0 m
foldrWithKey' :: Enum1 k => (forall a. k a -> v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldrWithKey' f acc0 (DEnumMap m) =
IM.foldrWithKey' (\i (KV inf v) acc ->
let k = toEnum1 i inf
in f k v acc)
acc0 m
foldlWithKey' :: Enum1 k => (forall a. acc -> k a -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldlWithKey' f acc0 (DEnumMap m) =
IM.foldlWithKey' (\acc i (KV inf v) ->
let k = toEnum1 i inf
in f acc k v)
acc0 m
-- * Conversion
elems :: DEnumMap k v -> [Some v]
elems (DEnumMap m) = (\(KV _ v) -> Some v) <$> IM.elems m
keys :: Enum1 k => DEnumMap k v -> [Some k]
keys (DEnumMap m) = (\(k, KV inf _) -> Some (toEnum1 k inf)) <$> IM.assocs m
assocs :: Enum1 k => DEnumMap k v -> [DSum k v]
assocs (DEnumMap m) = kVToDSum <$> IM.assocs m
-- TODO: Wait for DEnumSet.
-- keysSet
-- ** Lists
toList :: Enum1 k => DEnumMap k v -> [DSum k v]
toList = toAscList
-- ** Ordered lists
toAscList :: Enum1 k => DEnumMap k v -> [DSum k v]
toAscList (DEnumMap m) = kVToDSum <$> IM.toAscList m
toDescList :: Enum1 k => DEnumMap k v -> [DSum k v]
toDescList (DEnumMap m) = kVToDSum <$> IM.toDescList m
-- * Filter
filter :: (forall a. v a -> Bool) -> DEnumMap k v -> DEnumMap k v
filter f (DEnumMap m) = DEnumMap (IM.filter (\(KV _ v) -> f v) m)
filterWithKey :: Enum1 k => (forall a. k a -> v a -> Bool) -> DEnumMap k v -> DEnumMap k v
filterWithKey f (DEnumMap m) =
DEnumMap (IM.filterWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in f k v)
m)
-- TODO: Wait for DEnumSet.
-- restrictKeys
-- withoutKeys
partition :: (forall a. v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
partition f (DEnumMap m) =
bimap DEnumMap DEnumMap (IM.partition (\(KV _ v) -> f v) m)
partitionWithKey :: Enum1 k => (forall a. k a -> v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
partitionWithKey f (DEnumMap m) =
bimap DEnumMap DEnumMap (IM.partitionWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in f k v)
m)
-- | \(O(\min(n,W)^2)\). Because of the lack of a @takeWhileAntitoneWithValue@
-- operation on 'Data.IntMap.Strict.IntMap', this function performs additional lookups to
-- reconstruct the full keys to pass to the predicate, resulting in a somewhat
-- worse complexity than 'Data.IntMap.Strict.takeWhileAntitone'.
takeWhileAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> DEnumMap k v
takeWhileAntitone f (DEnumMap m) =
DEnumMap (IM.takeWhileAntitone (\i -> case m IM.! i of
KV inf _ -> let k = toEnum1 i inf
in f k)
m)
-- | \(O(\min(n,W)^2)\). See 'takeWhileAntitone'.
dropWhileAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> DEnumMap k v
dropWhileAntitone f (DEnumMap m) =
DEnumMap (IM.dropWhileAntitone (\i -> case m IM.! i of
KV inf _ -> let k = toEnum1 i inf
in f k)
m)
-- | \(O(\min(n,W)^2)\). See 'takeWhileAntitone'.
spanAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
spanAntitone f (DEnumMap m) =
bimap DEnumMap DEnumMap
(IM.spanAntitone (\i -> case m IM.! i of
KV inf _ -> let k = toEnum1 i inf
in f k)
m)
mapMaybe :: Enum1 k => (forall a. v1 a -> Maybe (v2 a)) -> DEnumMap k v1 -> DEnumMap k v2
mapMaybe f = mapMaybeWithKey (const f)
mapMaybeWithKey :: Enum1 k
=> (forall a. k a -> v1 a -> Maybe (v2 a)) -> DEnumMap k v1 -> DEnumMap k v2
mapMaybeWithKey f (DEnumMap m) =
DEnumMap (IM.mapMaybeWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in KV inf <$> f k v)
m)
mapEither :: Enum1 k
=> (forall a. v1 a -> Either (v2 a) (v3 a)) -> DEnumMap k v1 -> (DEnumMap k v2, DEnumMap k v3)
mapEither f = mapEitherWithKey (const f)
mapEitherWithKey :: Enum1 k
=> (forall a. k a -> v1 a -> Either (v2 a) (v3 a)) -> DEnumMap k v1 -> (DEnumMap k v2, DEnumMap k v3)
mapEitherWithKey f (DEnumMap m) =
bimap DEnumMap DEnumMap (IM.mapEitherWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in bimap (KV inf) (KV inf) $ f k v)
m)
split :: Enum1 k => k a -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
split k (DEnumMap m) = bimap DEnumMap DEnumMap (IM.split (fst $ fromEnum1 k) m)
splitLookup :: Enum1 k => k a -> DEnumMap k v -> (DEnumMap k v, Maybe (v a), DEnumMap k v)
splitLookup k (DEnumMap m) =
let (!m1, !mkv, !m2) = IM.splitLookup (fst $ fromEnum1 k) m
-- Note: this coe1 is fine because of the invariant on DEnumMap.
in (DEnumMap m1, (\(KV _ v) -> coe1 v) <$!> mkv, DEnumMap m2)
splitRoot :: DEnumMap k v -> [DEnumMap k v]
splitRoot (DEnumMap m) = DEnumMap <$> IM.splitRoot m
-- * Submap
-- TODO: the submap operations can't check any laws because there is no IM.isSubmapOfByKey.
isSubmapOf :: (forall a. Eq (v a)) => DEnumMap k v -> DEnumMap k v -> Bool
isSubmapOf (DEnumMap m1) (DEnumMap m2) = IM.isSubmapOfBy (\(KV _ v1) (KV _ v2) -> v1 == coe1 v2) m1 m2
isSubmapOfBy :: (forall a. v1 a -> v2 a -> Bool) -> DEnumMap k v1 -> DEnumMap k v2 -> Bool
isSubmapOfBy f (DEnumMap m1) (DEnumMap m2) =
IM.isSubmapOfBy (\(KV _ v1) (KV _ v2) -> f v1 (coe1 v2)) m1 m2
isProperSubmapOf :: (forall a. Eq (v a)) => DEnumMap k v -> DEnumMap k v -> Bool
isProperSubmapOf (DEnumMap m1) (DEnumMap m2) = IM.isProperSubmapOfBy (\(KV _ v1) (KV _ v2) -> v1 == coe1 v2) m1 m2
isProperSubmapOfBy :: (forall a. v1 a -> v2 a -> Bool) -> DEnumMap k v1 -> DEnumMap k v2 -> Bool
isProperSubmapOfBy f (DEnumMap m1) (DEnumMap m2) =
IM.isProperSubmapOfBy (\(KV _ v1) (KV _ v2) -> f v1 (coe1 v2)) m1 m2
-- * Min\/Max
lookupMin :: Enum1 k => DEnumMap k v -> Maybe (DSum k v)
lookupMin (DEnumMap m) = kVToDSum <$!> IM.lookupMin m
lookupMax :: Enum1 k => DEnumMap k v -> Maybe (DSum k v)
lookupMax (DEnumMap m) = kVToDSum <$!> IM.lookupMax m
findMin :: Enum1 k => DEnumMap k v -> DSum k v
findMin (DEnumMap m) = kVToDSum $ IM.findMin m
findMax :: Enum1 k => DEnumMap k v -> DSum k v
findMax (DEnumMap m) = kVToDSum $ IM.findMax m
deleteMin :: DEnumMap k v -> DEnumMap k v
deleteMin (DEnumMap m) = DEnumMap $ IM.deleteMin m
deleteMax :: DEnumMap k v -> DEnumMap k v
deleteMax (DEnumMap m) = DEnumMap $ IM.deleteMax m
deleteFindMin :: Enum1 k => DEnumMap k v -> (DSum k v, DEnumMap k v)
deleteFindMin (DEnumMap m) = bimap' kVToDSum DEnumMap $ IM.deleteFindMin m
deleteFindMax :: Enum1 k => DEnumMap k v -> (DSum k v, DEnumMap k v)
deleteFindMax (DEnumMap m) = bimap' kVToDSum DEnumMap $ IM.deleteFindMax m
updateMin :: Enum1 k => (forall a. v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMin f = updateMinWithKey (const f)
updateMinWithKey :: Enum1 k => (forall a. k a -> v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMinWithKey f (DEnumMap m) =
DEnumMap (IM.updateMinWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in KV inf <$> f k v)
m)
updateMax :: Enum1 k => (forall a. v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMax f = updateMaxWithKey (const f)
updateMaxWithKey :: Enum1 k => (forall a. k a -> v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMaxWithKey f (DEnumMap m) =
DEnumMap (IM.updateMaxWithKey (\i (KV inf v) ->
let k = toEnum1 i inf
in KV inf <$> f k v)
m)
minView :: DEnumMap k v -> Maybe (v a, DEnumMap k v)
minView (DEnumMap m) = bimap' (\(KV _ v) -> coe1 v) DEnumMap <$!> IM.minView m
maxView :: DEnumMap k v -> Maybe (v a, DEnumMap k v)
maxView (DEnumMap m) = bimap' (\(KV _ v) -> coe1 v) DEnumMap <$!> IM.maxView m
minViewWithKey :: Enum1 k => DEnumMap k v -> Maybe (DSum k v, DEnumMap k v)
minViewWithKey (DEnumMap m) = bimap' kVToDSum DEnumMap <$!> IM.minViewWithKey m
maxViewWithKey :: Enum1 k => DEnumMap k v -> Maybe (DSum k v, DEnumMap k v)
maxViewWithKey (DEnumMap m) = bimap' kVToDSum DEnumMap <$!> IM.maxViewWithKey m
-- * Helpers
coe1 :: v a -> v b
coe1 = unsafeCoerce
typeCheck1 :: forall k a b. (Enum1 k, TestEquality k)
=> k a -> Int -> Enum1Info k b -> k a :~: k b
typeCheck1 k1 i inf2 =
let ret :: k a :~: k b
ret = unsafeCoerce Refl
in assert (let k2 = toEnum1 i inf2
in case testEquality k1 k2 of
Just r -> justifies r ret True
Nothing -> False)
ret
where
justifies :: a :~: b -> k a :~: k b -> r -> r
justifies Refl _ x = x
typeCheck2 :: forall k proxy a b. (Enum1 k, TestEquality k)
=> proxy k -> Int -> Enum1Info k a -> Enum1Info k b -> k a :~: k b
typeCheck2 _ i inf1 inf2 =
let ret :: k a :~: k b
ret = unsafeCoerce Refl
in assert (let k1 = toEnum1 @k i inf1
k2 = toEnum1 i inf2
in case testEquality k1 k2 of
Just r -> justifies r ret True
Nothing -> False)
ret
where
justifies :: a :~: b -> k a :~: k b -> r -> r
justifies Refl _ x = x
bimap' :: (a -> b) -> (c -> d) -> (a, c) -> (b, d)
bimap' f g (a, c) =
let !b = f a
!d = g c
in (b, d)