Fill in all ops up to Query; not tested!

1 files changed, 42 insertions, 9 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs
index 6cc46a3..f87c3cd 100644
--- a/src/Data/Dependent/EnumMap/Strict/Internal.hs
+++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs
@@ -151,7 +151,16 @@ insertWithKey f k v (DEnumMap m) =
                 (\_ (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, dmap) =
+        IM.insertLookupWithKey
+          (\_ _ (KV inf' v2) -> typeCheck1 k i inf' $ KV inf (f k v (coe1 v2)))
+          i (KV inf v) m
+  in ((\(KV inf' v2) -> typeCheck1 k i inf' $ coe1 v2) <$> mx, DEnumMap dmap)
 
 -- * Deletion\/Update
 
@@ -159,16 +168,32 @@ 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 f k (DEnumMap m) =
+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 (coe1 v))) i m)
+  in DEnumMap (IM.adjust (\(KV inf v) -> typeCheck1 k i inf $ KV inf (f k (coe1 v))) i m)
 
--- adjustWithKey
--- update
--- updateWithKey
--- updateLookupWithKey
+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)
 
-alter :: forall k v a. (Enum1 k, TestEquality k) => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
+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, dmap) =
+        IM.updateLookupWithKey
+          (\_ (KV inf v) -> typeCheck1 k i inf . KV inf <$> f k (coe1 v))
+          i m
+  in ((\(KV inf' v2) -> typeCheck1 k i inf' $ coe1 v2) <$> mx, DEnumMap dmap)
+
+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
@@ -177,7 +202,15 @@ alter f k (DEnumMap m) = DEnumMap (IM.alter f' i m)
     f' Nothing = KV inf <$> f Nothing
     f' (Just (KV inf' v)) = typeCheck1 k i inf' $ KV inf <$> f (Just (coe1 v))
 
--- alterF
+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)) = typeCheck1 k i inf' $ fmap (KV inf) <$> f (Just (coe1 v))
 
 -- * Query
 -- ** Lookup