diff options
| -rw-r--r-- | src/Data/Dependent/EnumMap/Strict/Internal.hs | 37 |
1 files changed, 31 insertions, 6 deletions
diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs index ae56051..d2fdcb6 100644 --- a/src/Data/Dependent/EnumMap/Strict/Internal.hs +++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs @@ -329,21 +329,46 @@ differenceWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey -- ** Intersection --- intersection --- intersectionWith --- intersectionWithKey +intersection :: DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1 +intersection (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersection m1 m2) + +intersectionWith :: forall k v1 v2 v3. (Enum1 k, TestEquality k) + => (forall a. v1 a -> v2 a -> v3 a) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3 +intersectionWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2) + where + f' :: Int -> KV k v1 -> KV k v2 -> KV k v3 + f' i (KV inf1 v1) (KV inf2 v2) = + typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 $ f (coe1 v1) (coe1 v2) + +intersectionWithKey :: forall k v1 v2 v3. (Enum1 k, TestEquality k) + => (forall a. k a -> v1 a -> v2 a -> v3 a) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3 +intersectionWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2) + where + f' :: Int -> KV k v1 -> KV k v2 -> KV k v3 + f' i (KV inf1 v1) (KV inf2 v2) = case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 $ f k1 (coe1 v1) (coe1 v2) -- ** Disjoint --- disjoint +disjoint :: DEnumMap k v1 -> DEnumMap k v2 -> Bool +disjoint (DEnumMap m1) (DEnumMap m2) = IM.disjoint m1 m2 -- ** Compose --- compose +compose :: Enum1 k => DEnumMap k v -> DEnumMap k k -> DEnumMap k v +compose (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.compose m1 (IM.map (\(KV _ v) -> fst $ fromEnum1 v) m2)) -- ** Universal combining function --- mergeWithKey +mergeWithKey :: forall k v1 v2 v3. (Enum1 k, TestEquality k) + => (forall a. k a -> v1 a -> v2 a -> Maybe (v3 a)) -> (DEnumMap k v1 -> DEnumMap k v3) -> (DEnumMap k v2 -> DEnumMap k v3) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3 +mergeWithKey f g1 g2 (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.mergeWithKey f' g1' g2' m1 m2) + where + f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v3) + f' i (KV inf1 v1) (KV inf2 v2) = case toEnum1 i inf1 of + Some k1 -> typeCheck1 k1 i inf2 . KV inf1 <$> f k1 (coe1 v1) (coe1 v2) + g1' m = let DEnumMap m' = g1 (DEnumMap m) in m' + g2' m = let DEnumMap m' = g2 (DEnumMap m) in m' -- * Traversal -- ** Map |