Implement fromListWith and similar

1 files changed, 36 insertions, 4 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs
index f9679df..49fef13 100644
--- a/src/Data/Dependent/EnumMap/Strict/Internal.hs
+++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs
@@ -50,8 +50,24 @@ 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
+-- TODO: check that inf1 and inf2 match somehow? Comment why not needed or not possible, if so
+fromListWith :: Enum1 k
+             => (forall a. v a -> v a -> v a)
+             -> [DSum k v] -> DEnumMap k v
+fromListWith f l =
+  DEnumMap (IM.fromListWith
+             (\(KV inf1 v1) (KV _ v2) -> KV inf1 (f (coe1 v1) (coe1 v2)))
+             (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))
+
+-- TODO: check that inf1 and inf2 match somehow? Comment why not needed or not possible, if so
+fromListWithKey :: Enum1 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 _ v2) -> case toEnum1 i inf1 of
+                                             Some k -> KV inf1 (f k (coe1 v1) (coe1 v2)))
+             (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))
 
 -- ** From Ascending Lists
 
@@ -59,8 +75,24 @@ fromAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
 fromAscList l =
   DEnumMap (IM.fromAscList (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))
 
--- fromAscListWith
--- fromAscListWithKey
+-- TODO: check that inf1 and inf2 match somehow? Comment why not needed or not possible, if so
+fromAscListWith :: Enum1 k
+                => (forall a. v a -> v a -> v a)
+                -> [DSum k v] -> DEnumMap k v
+fromAscListWith f l =
+  DEnumMap (IM.fromAscListWith
+             (\(KV inf1 v1) (KV _ v2) -> KV inf1 (f (coe1 v1) (coe1 v2)))
+             (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))
+
+-- TODO: check that inf1 and inf2 match somehow? Comment why not needed or not possible, if so
+fromAscListWithKey :: Enum1 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 _ v2) -> case toEnum1 i inf1 of
+                                             Some k -> KV inf1 (f k (coe1 v1) (coe1 v2)))
+             (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))
 
 fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
 fromDistinctAscList l =