diff options
| -rw-r--r-- | src/Data/Dependent/EnumMap/Strict/Internal.hs | 51 |
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 -alter :: forall k v a. (Enum1 k, TestEquality k) => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v +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, 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 |