@@ -1,10 +1,11 @@
1−{-# LANGUAGE BangPatterns #-}
2−{-# LANGUAGE DataKinds #-}
3−{-# LANGUAGE GADTs #-}
4−{-# LANGUAGE KindSignatures #-}
5−{-# LANGUAGE MagicHash #-}
6−{-# LANGUAGE PolyKinds #-}
7−{-# LANGUAGE TypeFamilies #-}
1+{-# LANGUAGE AllowAmbiguousTypes #-}
2+{-# LANGUAGE BangPatterns #-}
3+{-# LANGUAGE DataKinds #-}
4+{-# LANGUAGE GADTs #-}
5+{-# LANGUAGE KindSignatures #-}
6+{-# LANGUAGE MagicHash #-}
7+{-# LANGUAGE TypeInType #-}
8+{-# LANGUAGE TypeFamilies #-}
89
910-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
1011
@@ -14,8 +15,11 @@ module Data.TypeRep.CacheMap
1415
1516 -- 'TypeRepMap' interface
1617 , empty
18+ , one
1719 , insert
20+ , delete
1821 , lookup
22+ , member
1923 , size
2024
2125 -- * Helpful testing functions
@@ -34,7 +38,6 @@ import Data.List (nubBy)
3438import Data.Maybe (fromJust)
3539import Data.Proxy (Proxy (..))
3640import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
37−import Data.Word (Word64)
3841import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
3942import GHC.Exts (inline, sortWith)
4043import GHC.Fingerprint (Fingerprint (..))
@@ -57,7 +60,7 @@ data TypeRepMap (f :: k -> Type) = TypeRepMap
5760instance Show (TypeRepMap f) where
5861 show = show . toFps
5962
60−-- | Returnes the list of 'Fingerprint's from 'TypeRepMap'.
63+-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
6164toFps :: TypeRepMap f -> [Fingerprint]
6265toFps TypeRepMap{..} = zipWith Fingerprint
6366 (Unboxed.toList fingerprintAs)
@@ -70,20 +73,52 @@ fromAny = unsafeCoerce
7073empty :: TypeRepMap f
7174empty = TypeRepMap mempty mempty mempty
7275
76+-- | Construct a 'TypeRepMap' with a single element.
77+one :: forall a f . Typeable a => f a -> TypeRepMap f
78+one x = insert x empty
79+{-# INLINE one #-}
80+
7381-- | Inserts the value with its type as a key.
7482insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
7583insert x = fromListPairs . addX . toPairList
7684 where
77− toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
78− toPairList trMap = zip (toFps trMap) (V.toList $ anys trMap)
79−
8085 pairX :: (Fingerprint, Any)
8186 pairX@(fpX, _) = (calcFp x, unsafeCoerce x)
8287
8388 addX :: [(Fingerprint, Any)] -> [(Fingerprint, Any)]
84− addX l = pairX : filter ((/= fpX) . fst) l
89+ addX l = pairX : deleteByFst fpX l
8590{-# INLINE insert #-}
8691
92+-- Helper type to workaround some issues with types
93+type KindOf (a :: k) = k
94+
95+{- | Deletes value from list.
96+
97+>>> let trMap = delete @Bool $ insert (Identity True) $ one (Identity 'a')
98+>>> size trMap
99+1
100+>>> member @Bool trMap
101+False
102+>>> member @Char trMap
103+True
104+-}
105+delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
106+delete = fromListPairs . deleteByFst (typeFp @a) . toPairList
107+{-# INLINE delete #-}
108+
109+{- | Returns 'True' if there exist value of given type.
110+
111+>>> member @Char $ one (Identity 'a')
112+True
113+>>> member @Bool $ one (Identity 'a')
114+False
115+-}
116+member :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> Bool
117+member trMap = case lookup @a trMap of
118+ Nothing -> False
119+ Just _ -> True
120+{-# INLINE member #-}
121+
87122{- | Looks up the value at the type.
88123
89124>>> let x = lookup $ insert (Identity (11 :: Int)) empty
@@ -91,11 +126,10 @@ insert x = fromListPairs . addX . toPairList
91126Just (Identity 11)
92127>>> x :: Maybe (Identity ())
93128Nothing
94−
95129-}
96130lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
97131lookup tVect = fromAny . (anys tVect V.!)
98− <$> cachedBinarySearch (typeRepFingerprint $ typeRep $ Proxy @a)
132+ <$> cachedBinarySearch (typeFp @a)
99133 (fingerprintAs tVect)
100134 (fingerprintBs tVect)
101135{-# INLINE lookup #-}
@@ -128,6 +162,23 @@ cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
128162 len = let !(I# l) = Unboxed.length fpAs in l
129163{-# INLINE cachedBinarySearch #-}
130164
165+----------------------------------------------------------------------------
166+-- Internal functions
167+----------------------------------------------------------------------------
168+
169+typeFp :: forall a . Typeable a => Fingerprint
170+typeFp = typeRepFingerprint $ typeRep $ Proxy @a
171+{-# INLINE typeFp #-}
172+
173+toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
174+toPairList trMap = zip (toFps trMap) (V.toList $ anys trMap)
175+
176+deleteByFst :: Eq a => a -> [(a, b)] -> [(a, b)]
177+deleteByFst x = filter ((/= x) . fst)
178+
179+nubByFst :: (Eq a) => [(a, b)] -> [(a, b)]
180+nubByFst = nubBy ((==) `on` fst)
181+
131182----------------------------------------------------------------------------
132183-- Functions for testing and benchmarking
133184----------------------------------------------------------------------------
@@ -165,10 +216,7 @@ fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
165216fromListPairs kvs = TypeRepMap (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
166217 where
167218 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
168− (fps, ans) = unzip $ fromSortedList $ sortWith fst $ nubPairs kvs
169−
170−nubPairs :: (Eq a) => [(a, b)] -> [(a, b)]
171−nubPairs = nubBy ((==) `on` fst)
219+ (fps, ans) = unzip $ fromSortedList $ sortWith fst $ nubByFst kvs
172220
173221----------------------------------------------------------------------------
174222-- Tree-like conversion