Implement the easy mapping operations

1 files changed, 45 insertions, 12 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs
index db083a6..add7fa5 100644
--- a/src/Data/Dependent/EnumMap/Strict/Internal.hs
+++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs
@@ -8,7 +8,7 @@
 module Data.Dependent.EnumMap.Strict.Internal where
 
 import Control.Exception
-import Data.Bifunctor (bimap)
+import Data.Bifunctor (bimap, second)
 import Data.Dependent.Sum
 import qualified Data.Foldable as Foldable
 import qualified Data.IntMap.Strict as IM
@@ -373,13 +373,36 @@ mergeWithKey f g1 g2 (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.mergeWithKey f'
 -- * Traversal
 -- ** Map
 
--- map
--- mapWithKey
--- traverseWithKey
--- traverseMaybeWithKey
--- mapAccum
--- mapAccumWithKey
--- mapAccumRWithKey
+-- TODO: can the kind in the typing be avoided? Or moved to DEnumMap or Enum1 definition or somewhere?
+map :: forall c (k :: c -> Type) (v1 :: c -> Type) v2. (Enum1 k, TestEquality k) => (forall a. v1 a -> v2 a) -> DEnumMap k v1 -> DEnumMap k v2
+map f = mapWithKey (\_ x -> f x)
+
+mapWithKey :: (Enum1 k, TestEquality 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 t, Enum1 k, TestEquality k)
+                => (forall a. k a -> v1 a -> t (v2 a)) -> DEnumMap k v1 -> t (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 t, Enum1 k, TestEquality k)
+                     => (forall a. k a -> v1 a -> t (Maybe (v2 a))) -> DEnumMap k v1 -> t (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 :: forall c (k :: c -> Type) (v1 :: c -> Type) v2 acc. (Enum1 k, TestEquality 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, TestEquality 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, TestEquality 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
@@ -463,10 +486,20 @@ spanAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> (DEnumMap
 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
--- mapMaybeWithKey
--- mapEither
--- mapEitherWithKey
+-- TODO: can the kind in the typing be avoided? Or moved to DEnumMap or Enum1 definition or somewhere?
+mapMaybe :: forall c (k :: c -> Type) (v1 :: c -> Type) v2. (Enum1 k, TestEquality k) => (forall a. v1 a -> Maybe (v2 a)) -> DEnumMap k v1 -> DEnumMap k v2
+mapMaybe f = mapMaybeWithKey (\_ x -> f x)
+
+mapMaybeWithKey :: (Enum1 k, TestEquality 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 :: forall c (k :: c -> Type) (v1 :: c -> Type) v2 v3. (Enum1 k, TestEquality k) => (forall a. v1 a -> Either (v2 a) (v3 a)) -> DEnumMap k v1 -> (DEnumMap k v2, DEnumMap k v3)
+mapEither f = mapEitherWithKey (\_ x -> f x)
+
+mapEitherWithKey :: (Enum1 k, TestEquality 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)