Fill in all ops up to Intersection; not tested!

1 files changed, 68 insertions, 14 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict/Internal.hs b/src/Data/Dependent/EnumMap/Strict/Internal.hs
index f87c3cd..ae56051 100644
--- a/src/Data/Dependent/EnumMap/Strict/Internal.hs
+++ b/src/Data/Dependent/EnumMap/Strict/Internal.hs
@@ -10,6 +10,7 @@ module Data.Dependent.EnumMap.Strict.Internal where
 import Control.Exception
 import Data.Bifunctor (bimap)
 import Data.Dependent.Sum
+import qualified Data.Foldable as Foldable
 import qualified Data.IntMap.Strict as IM
 import Data.Kind (Type)
 import Data.Proxy
@@ -18,6 +19,7 @@ import Data.Type.Equality
 import Text.Show (showListWith)
 import Unsafe.Coerce (unsafeCoerce)
 
+import Prelude hiding (lookup)
 
 data KV k v = forall a. KV !(Enum1Info k) !(v a)
 
@@ -220,8 +222,8 @@ lookup k (DEnumMap m) =
   let (i, _) = fromEnum1 k
   in (\(KV inf v) -> typeCheck1 k i inf $ coe1 v) <$> IM.lookup i m
 
--- (!?)
--- (!)
+(!?) :: (Enum1 k, TestEquality k) => DEnumMap k v -> k a -> Maybe (v a)
+(!?) m k = lookup k m
 
 findWithDefault :: (Enum1 k, TestEquality k) => v a -> k a -> DEnumMap k v -> v a
 findWithDefault def k (DEnumMap m) =
@@ -229,14 +231,37 @@ findWithDefault def k (DEnumMap m) =
   in case IM.findWithDefault (KV undefined def) i m of
        KV inf' v -> typeCheck1 k i inf' $ coe1 v
 
+find :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> v a
+find k = findWithDefault (error ("Dependent.EnumMap.!: key " ++ show (fst $ fromEnum1 k) ++ " is not an element of the map")) k
+
+(!) :: (Enum1 k, TestEquality k) => DEnumMap k v -> k a -> v a
+(!) m k = find k m
+
 member :: Enum1 k => k a -> DEnumMap k v -> Bool
 member k (DEnumMap m) = IM.member (fst (fromEnum1 k)) m
 
--- notMember
--- lookupLT
--- lookupGT
--- lookupLE
--- lookupGE
+notMember :: Enum1 k => k a -> DEnumMap k v -> Bool
+notMember k m = not $ member k m
+
+lookupLT :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
+lookupLT k (DEnumMap m) =
+  let (i, _) = fromEnum1 k
+  in kVToDSum <$> IM.lookupLT i m
+
+lookupGT :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
+lookupGT k (DEnumMap m) =
+  let (i, _) = fromEnum1 k
+  in kVToDSum <$> IM.lookupGT i m
+
+lookupLE :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
+lookupLE k (DEnumMap m) =
+  let (i, _) = fromEnum1 k
+  in kVToDSum <$> IM.lookupLE i m
+
+lookupGE :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
+lookupGE k (DEnumMap m) =
+  let (i, _) = fromEnum1 k
+  in kVToDSum <$> IM.lookupGE i m
 
 -- ** Size
 
@@ -260,18 +285,47 @@ unionWith f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWith
     f' :: Int -> KV k v -> KV k v -> KV k v
     f' i (KV inf1 v1) (KV inf2 v2) = typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 (f v1 (coe1 v2))
 
--- unionWithKey
--- unions
--- unionsWith
+unionWithKey :: (Enum1 k, TestEquality k)
+             => (forall a. k a -> v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
+unionWithKey f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWithKey f' m1 m2)
+  where
+    f' :: Int -> KV k v -> KV k v -> KV k v
+    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))
+        -- TODO: are the coe1 correct? is the use of typeCheck1 fine?
+
+unions :: (Foldable f, Enum1 k, TestEquality k) => f (DEnumMap k v) -> DEnumMap k v
+unions xs = Foldable.foldl' union empty xs
+
+unionsWith :: (Foldable f, Enum1 k, TestEquality k) => (forall a. v a -> v a -> v a) -> f (DEnumMap k v) -> DEnumMap k v
+unionsWith f xs = Foldable.foldl' (unionWith f) empty xs
 
 -- ** Difference
 
-difference :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v
+-- TODO: should this be v1, v2 or both v? what about k1 and k2?
+difference :: DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
 difference (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.difference m1 m2)
 
--- (\\)
--- differenceWith
--- differenceWithKey
+(\\) :: DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
+m1 \\ m2 = difference m1 m2
+
+-- TODO: what about k1 and k2 here?
+differenceWith :: forall k v1 v2. (Enum1 k, TestEquality k)
+               => (forall a. v1 a -> v2 a -> Maybe (v1 a)) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
+differenceWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
+  where
+    f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v1)
+    f' i (KV inf1 v1) (KV inf2 v2) =
+      typeCheck2 (Proxy @k) i inf1 inf2 . KV inf1 <$> f (coe1 v1) (coe1 v2)
+
+-- TODO: what about k1 and k2 here?
+differenceWithKey :: forall k v1 v2. (Enum1 k, TestEquality k)
+                  => (forall a. k a -> v1 a -> v2 a -> Maybe (v1 a)) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
+differenceWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
+  where
+    f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v1)
+    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)
 
 -- ** Intersection