From fc17ab793e271f947697d42c3f9300424ae6cf4c Mon Sep 17 00:00:00 2001 From: Tom Smeding Date: Mon, 22 Jul 2024 15:14:55 +0200 Subject: Initial with functions that horde-ad uses --- src/Data/Dependent/EnumMap/Strict.hs | 284 +++++++++++++++++++++++++++++++++++ 1 file changed, 284 insertions(+) create mode 100644 src/Data/Dependent/EnumMap/Strict.hs (limited to 'src/Data') 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 -- cgit v1.2.3-70-g09d2