Initial with functions that horde-ad uses

1 files changed, 284 insertions, 0 deletions

diff --git a/src/Data/Dependent/EnumMap/Strict.hs b/src/Data/Dependent/EnumMap/Strict.hs
new file mode 100644
index 0000000..6c1f11a
--- /dev/null
+++ b/src/Data/Dependent/EnumMap/Strict.hs
@@ -0,0 +1,284 @@
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+module Data.Dependent.EnumMap.Strict where
+
+import Data.Bifunctor (bimap)
+import Data.Dependent.Sum
+import qualified Data.IntMap.Strict as IM
+import Data.Some
+import Unsafe.Coerce (unsafeCoerce)
+
+
+data KV k v = forall a. KV !(Enum1Info k) !(v a)
+
+newtype DEnumMap k v = DEnumMap (IM.IntMap (KV k v))
+
+class Enum1 f where
+  type Enum1Info f
+  fromEnum1 :: f a -> (Int, Enum1Info f)
+  toEnum1 :: Int -> Enum1Info f -> Some f
+
+-- * Construction
+
+empty :: DEnumMap k v
+empty = DEnumMap IM.empty
+
+singleton :: Enum1 k => k a -> v a -> DEnumMap k v
+singleton k v =
+  let (i, inf) = fromEnum1 k
+  in DEnumMap (IM.singleton i (KV inf v))
+
+-- fromSet
+
+-- ** From Unordered Lists
+
+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
+
+-- ** From Ascending Lists
+
+-- fromAscList
+-- fromAscListWith
+-- fromAscListWithKey
+
+fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
+fromDistinctAscList l =
+  DEnumMap (IM.fromDistinctAscList (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))
+
+-- * Insertion
+
+insert :: Enum1 k => k a -> v a -> DEnumMap k v -> DEnumMap k v
+insert k v (DEnumMap m) =
+  let (i, inf) = fromEnum1 k
+  in DEnumMap (IM.insert i (KV inf v) m)
+
+-- insertWith
+
+insertWithKey :: Enum1 k
+              => (forall b. k b -> v b -> v b -> v b)
+              -> k a -> v a -> DEnumMap k v -> DEnumMap k v
+insertWithKey f k v (DEnumMap m) =
+  let (i, inf) = fromEnum1 k
+  in DEnumMap (IM.insertWithKey
+                (\_ (KV _ v1) (KV _ v2) -> KV inf (f k (coe1 v1) (coe1 v2)))
+                i (KV inf v) m)
+
+-- insertLookupWithKey
+
+-- * Deletion\/Update
+
+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 => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
+adjust f k (DEnumMap m) = DEnumMap (IM.adjust (\(KV k' v) -> KV k' (f (coe1 v))) (fst (fromEnum1 k)) m)
+
+-- adjustWithKey
+-- update
+-- updateWithKey
+-- updateLookupWithKey
+
+alter :: forall k v a. Enum1 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
+
+    f' :: Maybe (KV k v) -> Maybe (KV k v)
+    f' Nothing = KV inf <$> f Nothing
+    f' (Just (KV _ v)) = KV inf <$> f (Just (coe1 v))
+
+-- alterF
+
+-- * Query
+-- ** Lookup
+
+lookup :: Enum1 k => k a -> DEnumMap k v -> Maybe (v a)
+lookup k (DEnumMap m) = (\(KV _ v) -> coe1 v) <$> IM.lookup (fst (fromEnum1 k)) m
+
+-- (!?)
+-- (!)
+
+findWithDefault :: Enum1 k => v a -> k a -> DEnumMap k v -> v a
+findWithDefault def k (DEnumMap m) =
+  (\(KV _ v) -> coe1 v) $
+    IM.findWithDefault (KV undefined def) (fst (fromEnum1 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
+
+-- ** Size
+
+null :: DEnumMap k v -> Bool
+null (DEnumMap m) = IM.null m
+
+-- size
+
+-- * Combine
+
+-- ** Union
+
+union :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v
+union (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.union m1 m2)
+
+unionWith :: forall k v. (forall a. v a -> v a -> v a)
+          -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
+unionWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.unionWith f' m1 m2)
+  where
+    f' :: KV k v -> KV k v -> KV k v
+    f' (KV inf v1) (KV _ v2) = KV inf (f v1 (coe1 v2))
+
+-- unionWithKey
+-- unions
+-- unionsWith
+
+-- ** Difference
+
+difference :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v
+difference (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.difference m1 m2)
+
+-- (\\)
+-- differenceWith
+-- differenceWithKey
+
+-- ** Intersection
+
+-- intersection
+-- intersectionWith
+-- intersectionWithKey
+
+-- ** Disjoint
+
+-- disjoint
+
+-- ** Compose
+
+-- compose
+
+-- ** Universal combining function
+
+-- mergeWithKey
+
+-- * Traversal
+-- ** Map
+
+-- map
+-- mapWithKey
+-- traverseWithKey
+-- traverseMaybeWithKey
+-- mapAccum
+-- mapAccumWithKey
+-- mapAccumRWithKey
+-- mapKeys
+-- mapKeysWith
+-- mapKeysMonotonic
+
+-- * Folds
+
+-- foldr
+-- foldl
+-- foldrWithKey
+-- foldlWithKey
+-- foldMapWithKey
+
+-- ** Strict folds
+
+-- foldr'
+-- foldl'
+-- foldrWithKey'
+-- foldlWithKey'
+
+-- * Conversion
+
+elems :: DEnumMap k v -> [Some v]
+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
+-- keysSet
+
+-- ** Lists
+
+-- toList
+
+-- ** Ordered lists
+
+-- toAscList
+
+toDescList :: Enum1 k => DEnumMap k v -> [DSum k v]
+toDescList (DEnumMap m) =
+  map (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v)
+      (IM.toDescList m)
+
+-- * Filter
+
+-- filter
+-- filterWithKey
+-- restrictKeys
+-- withoutKeys
+
+partition :: (forall a. v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v) 
+partition f (DEnumMap m) =
+  bimap DEnumMap DEnumMap (IM.partition (\(KV _ v) -> f v) m)
+
+-- partitionWithKey
+
+-- takeWhileAntitone
+-- dropWhileAntitone
+-- spanAntitone
+
+-- mapMaybe
+-- mapMaybeWithKey
+-- mapEither
+-- mapEitherWithKey
+
+-- split
+-- splitLookup
+-- splitRoot
+
+-- * Submap
+
+-- isSubmapOf, isSubmapOfBy
+-- isProperSubmapOf, isProperSubmapOfBy
+
+-- * Min\/Max
+
+-- lookupMin
+-- lookupMax
+-- findMin
+-- findMax
+-- deleteMin
+-- deleteMax
+-- deleteFindMin
+-- deleteFindMax
+-- updateMin
+-- updateMax
+-- updateMinWithKey
+-- updateMaxWithKey
+-- minView
+-- maxView
+-- minViewWithKey
+
+maxViewWithKey :: Enum1 k => DEnumMap k v -> Maybe (DSum k v, DEnumMap k v)
+maxViewWithKey (DEnumMap m) =
+  bimap (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v) DEnumMap
+    <$> IM.maxViewWithKey m
+
+
+-- * Unsafe helpers
+
+coe1 :: v a -> v b
+coe1 = unsafeCoerce