Implement many filtering operations

1 files changed, 38 insertions, 8 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs
index d2fdcb6..369cd48 100644
--- a/src/Data/Dependent/EnumMap/Strict/Internal.hs
+++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs
@@ -407,7 +407,10 @@ elems (DEnumMap m) = map (\(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)
 
--- assocs
+assocs :: Enum1 k => DEnumMap k v -> [DSum k v]
+assocs (DEnumMap m) = map kVToDSum $ IM.assocs m
+
+-- TODO: probably doesn't make much sense until we have DEnumSet?
 -- keysSet
 
 -- ** Lists
@@ -426,7 +429,15 @@ toDescList (DEnumMap m) = map 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: these use IntSet. Do we use a list instead of wait for DEnumSet?
 -- restrictKeys
 -- withoutKeys
 
@@ -438,18 +449,37 @@ partitionWithKey :: Enum1 k => (forall a. k a -> v a -> Bool) -> 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)
 
--- takeWhileAntitone
--- dropWhileAntitone
--- spanAntitone
+-- To make this more efficient, we'd need to define takeWhileAntitoneWithValue
+-- for IntMap and use it here.
+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)
+
+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)
+
+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
 -- mapMaybeWithKey
 -- mapEither
 -- mapEitherWithKey
 
--- split
--- splitLookup
--- splitRoot
+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)
+
+-- TODO: is this coe1 fine or can we readably check that IM doesn't cheat
+-- and give us a value with a wrong type?
+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
+  in (DEnumMap m1, (\(KV _ v) -> coe1 v) <$> mkv, DEnumMap m2)
+
+splitRoot :: DEnumMap k v -> [DEnumMap k v]
+splitRoot (DEnumMap m) = DEnumMap <$> IM.splitRoot m
 
 -- * Submap