diff options
Diffstat (limited to 'src/Data/Dependent/EnumMap/Strict')
| -rw-r--r-- | src/Data/Dependent/EnumMap/Strict/Internal.hs | 616 | ||||
| -rw-r--r-- | src/Data/Dependent/EnumMap/Strict/Unsafe.hs | 45 |
2 files changed, 537 insertions, 124 deletions
diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs index 7e54a64..dcad7d1 100644 --- a/src/Data/Dependent/EnumMap/Strict/Internal.hs +++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs @@ -1,34 +1,80 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE PolyKinds #-} -{-# LANGUAGE RankNTypes #-} {-# LANGUAGE QuantifiedConstraints #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneKindSignatures #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} module Data.Dependent.EnumMap.Strict.Internal where -import Data.Bifunctor (bimap) +import Prelude hiding (lookup, map) + +import Control.Exception +import Data.Bifunctor (bimap, second) +import Data.Coerce import Data.Dependent.Sum +import qualified Data.Foldable as Foldable import qualified Data.IntMap.Strict 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) !(v a) +-- Invariant: 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 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 :: kind -> Type) (v :: kind -> Type)) where + => 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. class Enum1 f where type Enum1Info f fromEnum1 :: f a -> (Int, Enum1Info f) toEnum1 :: Int -> Enum1Info f -> Some f +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) = case toEnum1 i inf of Some k -> k :=> coe1 v + -- * Construction empty :: DEnumMap k v @@ -39,26 +85,61 @@ 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 (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l)) - --- fromListWith --- fromListWithKey +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) -> + typeCheck2 (Proxy @k) i inf1 inf2 $ + KV inf1 (f v1 (coe1 v2))) + (dSumToKV <$> l)) + +fromListWithKey :: (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) -> + case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2))) + (dSumToKV <$> l)) -- ** From Ascending Lists --- fromAscList --- fromAscListWith --- fromAscListWithKey +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) -> + typeCheck2 (Proxy @k) i inf1 inf2 $ + KV inf1 (f v1 (coe1 v2))) + (dSumToKV <$> l)) + +fromAscListWithKey :: (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) -> + case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2))) + (dSumToKV <$> l)) fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v -fromDistinctAscList l = - DEnumMap (IM.fromDistinctAscList (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l)) +fromDistinctAscList l = DEnumMap (IM.fromDistinctAscList (dSumToKV <$> l)) -- * Insertion @@ -67,156 +148,345 @@ insert k v (DEnumMap m) = let (i, inf) = fromEnum1 k in DEnumMap (IM.insert i (KV inf v) m) --- insertWith +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 +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.insertWithKey - (\_ (KV _ v1) (KV _ v2) -> KV inf (f k (coe1 v1) (coe1 v2))) + in DEnumMap (IM.insertWith + (\_ (KV inf' v2) -> typeCheck1 k i inf' $ KV inf (f k v (coe1 v2))) i (KV inf v) m) --- insertLookupWithKey +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) -> typeCheck1 k i inf' $ KV inf (f k v (coe1 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 => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v -adjust f k (DEnumMap m) = DEnumMap (IM.adjust (\(KV k' v) -> KV k' (f (coe1 v))) (fst (fromEnum1 k)) m) - --- adjustWithKey --- update --- updateWithKey --- updateLookupWithKey - -alter :: forall k v a. Enum1 k => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v +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) -> typeCheck1 k i inf $ KV inf (f k (coe1 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) -> typeCheck1 k i inf $ KV inf <$> f k (coe1 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) -> typeCheck1 k i inf $ KV inf <$> f k (coe1 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 _ v)) = KV inf <$> f (Just (coe1 v)) + f' (Just (KV inf' v)) = typeCheck1 k i inf' $ KV inf <$> f (Just (coe1 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 --- alterF + f' :: Maybe (KV k v) -> f (Maybe (KV k v)) + f' Nothing = fmap (KV inf) <$> f Nothing + f' (Just (KV inf' v)) = typeCheck1 k i inf' $ fmap (KV inf) <$> f (Just (coe1 v)) -- * Query -- ** Lookup -lookup :: Enum1 k => k a -> DEnumMap k v -> Maybe (v a) -lookup k (DEnumMap m) = (\(KV _ v) -> coe1 v) <$> IM.lookup (fst (fromEnum1 k)) m +lookup :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (v a) +lookup k (DEnumMap m) = + let (i, _) = fromEnum1 k + in (\(KV inf v) -> typeCheck1 k i inf $ coe1 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 => v a -> k a -> DEnumMap k v -> v a +findWithDefault :: (Enum1 k, TestEquality k) => v a -> k a -> DEnumMap k v -> v a findWithDefault def k (DEnumMap m) = - case IM.findWithDefault (KV undefined def) (fst (fromEnum1 k)) m of - KV _ v -> coe1 v + let (i, _) = fromEnum1 k + in case IM.findWithDefault (KV undefined def) i m of + KV inf' v -> typeCheck1 k i inf' $ coe1 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 --- lookupLT --- lookupGT --- lookupLE --- lookupGE +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 +size :: DEnumMap k v -> Int +size (DEnumMap m) = IM.size m -- * Combine -- ** Union -union :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v -union (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.union m1 m2) +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 :: forall k v. (forall a. v a -> v a -> v a) - -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v -unionWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.unionWith f' m1 m2) +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' :: KV k v -> KV k v -> KV k v - f' (KV inf v1) (KV _ v2) = KV inf (f v1 (coe1 v2)) + f' :: Int -> KV k v -> KV k v -> KV k v + f' i (KV inf1 v1) (KV inf2 v2) = typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 (f v1 (coe1 v2)) --- unionWithKey --- unions --- unionsWith +unionWithKey :: (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) = case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 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 v -> DEnumMap k v -> DEnumMap k v +difference :: DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1 difference (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.difference m1 m2) --- (\\) --- differenceWith --- differenceWithKey +(\\) :: 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) = + typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 <$> f (coe1 v1) (coe1 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) = case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 <$> f k1 (coe1 v1) (coe1 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) = case (toEnum1 i inf1, toEnum1 i inf2) of + (Some k1, Some k2) -> KV inf1 <$> f k1 (coe1 v1) k2 (coe1 v2) -- ** Intersection --- intersection --- intersectionWith --- intersectionWithKey +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) = + typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 $ f (coe1 v1) (coe1 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) = case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 $ f k1 (coe1 v1) (coe1 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) = case (toEnum1 i inf1, toEnum1 i inf2) of + (Some k1, Some k2) -> KV inf1 $ f k1 (coe1 v1) k2 (coe1 v2) -- ** Disjoint --- disjoint +disjoint :: DEnumMap k v1 -> DEnumMap k v2 -> Bool +disjoint (DEnumMap m1) (DEnumMap m2) = IM.disjoint m1 m2 -- ** Compose --- 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 +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) = case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 <$> f k1 (coe1 v1) (coe1 v2) -- * Traversal -- ** Map --- map --- mapWithKey --- traverseWithKey --- traverseMaybeWithKey --- mapAccum --- mapAccumWithKey --- mapAccumRWithKey +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) -> case toEnum1 i inf of Some k -> KV inf $ f k (coe1 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) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 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) -> case toEnum1 i inf of Some k -> fmap (KV inf) <$> f k (coe1 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) -> case toEnum1 i inf of Some k -> second (KV inf) $ f acc k (coe1 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) -> case toEnum1 i inf of Some k -> second (KV inf) $ f acc k (coe1 v)) acc0 m + +-- TODO: These are hard. Probably we can't avoid using a fold, analogously as in IntMap. -- mapKeys -- mapKeysWith -- mapKeysMonotonic -- * Folds --- foldr --- foldl --- foldrWithKey --- foldlWithKey --- foldMapWithKey +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 -> case toEnum1 i inf of Some k -> f k (coe1 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) -> case toEnum1 i inf of Some k -> f acc k (coe1 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) -> case toEnum1 i inf of Some k -> f k (coe1 v)) m -- ** Strict folds --- foldr' --- foldl' --- foldrWithKey' --- foldlWithKey' +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 -> case toEnum1 i inf of Some k -> f k (coe1 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) -> case toEnum1 i inf of Some k -> f acc k (coe1 v)) acc0 m -- * Conversion elems :: DEnumMap k v -> [Some v] -elems (DEnumMap m) = map (\(KV _ v) -> Some v) (IM.elems m) +elems (DEnumMap m) = (\(KV _ v) -> Some v) <$> IM.elems m keys :: Enum1 k => DEnumMap k v -> [Some k] -keys (DEnumMap m) = map (\(k, KV inf _) -> toEnum1 k inf) (IM.assocs m) +keys (DEnumMap m) = (\(k, KV inf _) -> toEnum1 k inf) <$> IM.assocs m --- assocs +assocs :: Enum1 k => DEnumMap k v -> [DSum k v] +assocs (DEnumMap m) = kVToDSum <$> IM.assocs m + +-- TODO: Wait for DEnumSet. -- keysSet -- ** Lists @@ -227,71 +497,169 @@ toList = toAscList -- ** Ordered lists toAscList :: Enum1 k => DEnumMap k v -> [DSum k v] -toAscList (DEnumMap m) = - map (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v) - (IM.toAscList m) +toAscList (DEnumMap m) = kVToDSum <$> IM.toAscList m toDescList :: Enum1 k => DEnumMap k v -> [DSum k v] -toDescList (DEnumMap m) = - map (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v) - (IM.toDescList m) +toDescList (DEnumMap m) = kVToDSum <$> IM.toDescList m -- * Filter --- filter --- filterWithKey +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) -> case toEnum1 i inf of Some k -> f k (coe1 v)) m) + +-- TODO: Wait for DEnumSet. -- restrictKeys -- withoutKeys -partition :: (forall a. v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v) +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 +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) -> case toEnum1 i inf of Some k -> f k (coe1 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 'IM.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 _ -> case toEnum1 i inf of Some k -> 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 _ -> case toEnum1 i inf of Some k -> 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 _ -> case toEnum1 i inf of Some k -> 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) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 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) -> case toEnum1 i inf of Some k -> bimap (KV inf) (KV inf) $ f k (coe1 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 --- takeWhileAntitone --- dropWhileAntitone --- spanAntitone +-- * Submap --- mapMaybe --- mapMaybeWithKey --- mapEither --- mapEitherWithKey +-- 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 --- split --- splitLookup --- splitRoot +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 --- * Submap +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 --- isSubmapOf, isSubmapOfBy --- isProperSubmapOf, isProperSubmapOfBy +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 --- lookupMax --- findMin --- findMax --- deleteMin --- deleteMax --- deleteFindMin --- deleteFindMax --- updateMin --- updateMax --- updateMinWithKey --- updateMaxWithKey --- minView --- maxView --- minViewWithKey +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) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 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) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 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 (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v) DEnumMap - <$> IM.maxViewWithKey m +maxViewWithKey (DEnumMap m) = bimap kVToDSum DEnumMap <$> IM.maxViewWithKey m --- * Unsafe helpers +-- * Helpers coe1 :: v a -> v b coe1 = unsafeCoerce + +typeCheck1 :: (Enum1 k, TestEquality k) + => k a -> Int -> Enum1Info k -> r -> r +typeCheck1 k1 i inf2 x = + assert (case toEnum1 i inf2 of { Some k2 -> + case testEquality k1 k2 of + Just Refl -> True + Nothing -> False }) + x + +typeCheck2 :: forall k proxy r. (Enum1 k, TestEquality k) + => proxy k -> Int -> Enum1Info k -> Enum1Info k -> r -> r +typeCheck2 _ i inf1 inf2 x = + assert (case toEnum1 @k i inf1 of { Some k1 -> + case toEnum1 i inf2 of { Some k2 -> + case testEquality k1 k2 of + Just Refl -> True + Nothing -> False }}) + x diff --git a/src/Data/Dependent/EnumMap/Strict/Unsafe.hs b/src/Data/Dependent/EnumMap/Strict/Unsafe.hs new file mode 100644 index 0000000..10dd744 --- /dev/null +++ b/src/Data/Dependent/EnumMap/Strict/Unsafe.hs @@ -0,0 +1,45 @@ +{-| +These are variants of the functions in "Data.Dependent.EnumMap.Strict" that do +not type-check keys: they do not check that you don't create two keys with the +same 'Int' and different types. As a result, these functions do not have a +'Data.Type.Equality.TestEquality' constraint, and are faster. + +Be careful though, because one can easily create @unsafeCoerce@ with this API: + +@ +{-# LANGUAGE ScopedTypeVariables TypeFamilies #-} + +import qualified Data.Dependent.EnumMap.Strict as DE +import qualified Data.Dependent.EnumMap.Strict.Unsafe as DEU + +import Data.Functor.Identity +import Data.Maybe +import Data.Some + +data Foo a = Foo Int + deriving (Show) + +instance DE.Enum1 Foo where + type Enum1Info Foo = () + fromEnum1 (Foo i) = (i, ()) + toEnum1 i () = Some (Foo i) + +unsafe :: forall a b. a -> b +unsafe x = runIdentity $ fromJust $ + DEU.lookupUnsafe (Foo 1 :: Foo b) $ + DE.singleton (Foo 1 :: Foo a) (Identity x) +@ + +-} +module Data.Dependent.EnumMap.Strict.Unsafe ( + adjustUnsafe, + alterUnsafe, + lookupUnsafe, + findWithDefaultUnsafe, + unionUnsafe, + unionWithUnsafe, +) where + +import Prelude () + +import Data.Dependent.EnumMap.Strict.Internal |