@@ -9,6 +9,7 @@
99{-# LANGUAGE Rank2Types #-}
1010{-# LANGUAGE TypeFamilies #-}
1111{-# LANGUAGE TypeInType #-}
12+{-# LANGUAGE ViewPatterns #-}
1213
1314-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
1415
@@ -22,7 +23,6 @@ module Data.TypeRepMap.Internal where
2223
2324import Prelude hiding (lookup)
2425
25−import Control.Arrow ((&&&))
2626import Data.Function (on)
2727import Data.IntMap.Strict (IntMap)
2828import Data.Kind (Type)
@@ -30,14 +30,14 @@ import Data.List (nubBy)
3030import Data.Maybe (fromJust)
3131import Data.Primitive.Array (Array, indexArray, mapArray')
3232import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
33−import Data.Proxy (Proxy (..))
3433import Data.Semigroup (Semigroup (..))
35−import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
3634import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
37−import GHC.Exts (inline, sortWith)
35+import GHC.Exts (IsList (..), inline, sortWith)
3836import GHC.Fingerprint (Fingerprint (..))
3937import GHC.Prim (eqWord#, ltWord#)
4038import GHC.Word (Word64 (..))
39+import Type.Reflection (Typeable, typeRep, withTypeable)
40+import Type.Reflection.Unsafe (typeRepFingerprint)
4141import Unsafe.Coerce (unsafeCoerce)
4242
4343import qualified Data.IntMap.Strict as IM
@@ -70,6 +70,7 @@ data TypeRepMap (f :: k -> Type) =
7070 { fingerprintAs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ first components of key fingerprints
7171 , fingerprintBs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ second components of key fingerprints
7272 , anys :: {-# UNPACK #-} !(Array Any) -- ^ values stored in the map
73+ , keys :: {-# UNPACK #-} !(Array Any) -- ^ typerep keys
7374 }
7475 -- ^ an unsafe constructor for 'TypeRepMap'
7576
@@ -84,7 +85,7 @@ instance Semigroup (TypeRepMap f) where
8485 {-# INLINE (<>) #-}
8586
8687instance Monoid (TypeRepMap f) where
87− mempty = TypeRepMap mempty mempty mempty
88+ mempty = TypeRepMap mempty mempty mempty mempty
8889 mappend = (<>)
8990 {-# INLINE mempty #-}
9091 {-# INLINE mappend #-}
@@ -127,13 +128,13 @@ prop> member @a (insert (x :: f a) tm) == True
127128
128129-}
129130insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
130−insert x = fromListPairs . addX . toPairList
131+insert x = fromTriples . addX . toTriples
131132 where
132− pairX :: (Fingerprint, Any)
133− pairX@(fpX, _) = (calcFp x, toAny x)
133+ tripleX :: (Fingerprint, Any, Any)
134+ tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
134135
135− addX :: [(Fingerprint, Any)] -> [(Fingerprint, Any)]
136− addX l = pairX : deleteByFst fpX l
136+ addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
137+ addX l = tripleX : deleteByFst fpX l
137138{-# INLINE insert #-}
138139
139140-- Extract the kind of a type. We use it to work around lack of syntax for
@@ -154,7 +155,7 @@ False
154155True
155156-}
156157delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
157−delete = fromListPairs . deleteByFst (typeFp @a) . toPairList
158+delete = fromTriples . deleteByFst (typeFp @a) . toTriples
158159{-# INLINE delete #-}
159160
160161{- | Map over the elements of a 'TypeRepMap'.
@@ -171,19 +172,30 @@ Just [True]
171172Just "a"
172173-}
173174hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
174−hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ mapArray' (toAny . f . fromAny) ans
175+hoist f (TypeRepMap as bs ans ks) = TypeRepMap as bs (mapArray' (toAny . f . fromAny) ans) ks
175176{-# INLINE hoist #-}
176177
178+hoistA :: (Applicative t) => (forall x. f x -> t (g x)) -> TypeRepMap f -> t (TypeRepMap g)
179+hoistA f (TypeRepMap as bs (toList -> ans) ks) = (\l -> TypeRepMap as bs (fromList $ map toAny l) ks)
180+ <$> traverse (f . fromAny) ans
181+{-# INLINE hoistA #-}
182+
177183-- | The union of two 'TypeRepMap's using a combining function.
178184unionWith :: (forall x. f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
179−unionWith f m1 m2 = fromListPairs
180− $ Map.toList
185+unionWith f m1 m2 = fromTriples
186+ $ toTripleList
181187 $ Map.unionWith combine
182− (Map.fromList $ toPairList m1)
183− (Map.fromList $ toPairList m2)
188+ (fromTripleList $ toTriples m1)
189+ (fromTripleList $ toTriples m2)
184190 where
185− combine :: Any -> Any -> Any
186− combine a b = toAny $ f (fromAny a) (fromAny b)
191+ combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
192+ combine (av, ak) (bv, _) = (toAny $ f (fromAny av) (fromAny bv), ak)
193+
194+ fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
195+ fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
196+
197+ toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
198+ toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
187199{-# INLINE unionWith #-}
188200
189201-- | The (left-biased) union of two 'TypeRepMap's. It prefers the first map when
@@ -259,56 +271,73 @@ fromAny :: Any -> f a
259271fromAny = unsafeCoerce
260272
261273typeFp :: forall a . Typeable a => Fingerprint
262−typeFp = typeRepFingerprint $ typeRep $ Proxy @a
274+typeFp = typeRepFingerprint $ typeRep @a
263275{-# INLINE typeFp #-}
264276
265−toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
266−toPairList tm = zip (toFingerprints tm) (GHC.toList $ anys tm)
277+toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
278+toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ anys tm) (GHC.toList $ keys tm)
267279
268−deleteByFst :: Eq a => a -> [(a, b)] -> [(a, b)]
269−deleteByFst x = filter ((/= x) . fst)
280+deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
281+deleteByFst x = filter ((/= x) . fst3)
270282
271−nubByFst :: (Eq a) => [(a, b)] -> [(a, b)]
272−nubByFst = nubBy ((==) `on` fst)
283+nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
284+nubByFst = nubBy ((==) `on` fst3)
285+
286+fst3 :: (a, b, c) -> a
287+fst3 (a, _, _) = a
273288
274289----------------------------------------------------------------------------
275290-- Functions for testing and benchmarking
276291----------------------------------------------------------------------------
277292
278293-- | Existential wrapper around 'Typeable' indexed by @f@ type parameter.
279−-- Useful for 'TypeRepMap' structure creation form list of 'TF's.
280−data TF f where
281− TF :: Typeable a => f a -> TF f
294+-- Useful for 'TypeRepMap' structure creation form list of 'WrapTypeable's.
295+data WrapTypeable f where
296+ WrapTypeable :: Typeable a => f a -> WrapTypeable f
297+
298+instance Show (WrapTypeable f) where
299+ show (WrapTypeable (_ :: f a)) = show $ calcFp @a
282300
283−instance Show (TF f) where
284− show (TF tf) = show $ calcFp tf
301+{- |
285302
286−{- | Creates 'TypeRepMap' from a list of 'TF's.
303+prop> fromList . toList == 'id'
287304
288−>>> size $ fromList [TF $ Identity True, TF $ Identity 'a']
305+Creates 'TypeRepMap' from a list of 'WrapTypeable's.
306+
307+>>> size $ fromList [WrapTypeable $ Identity True, WrapTypeable $ Identity 'a']
2893082
290309
310+
291311-}
292−fromList :: forall f . [TF f] -> TypeRepMap f
293−fromList = fromListPairs . map (fp &&& an)
294− where
295− fp :: TF f -> Fingerprint
296− fp (TF x) = calcFp x
312+instance IsList (TypeRepMap f) where
313+ type Item (TypeRepMap f) = WrapTypeable f
314+
315+ fromList :: [WrapTypeable f] -> TypeRepMap f
316+ fromList = fromTriples . map (\x -> (fp x, an x, k x))
317+ where
318+ fp :: WrapTypeable f -> Fingerprint
319+ fp (WrapTypeable (_ :: f a)) = calcFp @a
320+
321+ an :: WrapTypeable f -> Any
322+ an (WrapTypeable x) = toAny x
297323
298− an :: TF f -> Any
299− an (TF x) = toAny x
324+ k :: WrapTypeable f -> Any
325+ k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
300326
301−fromF :: Typeable a => f a -> Proxy a
302−fromF _ = Proxy
327+ toList :: TypeRepMap f -> [WrapTypeable f]
328+ toList = map toWrapTypeable . toTriples
329+ where
330+ toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
331+ toWrapTypeable (_, an, k) = withTypeable (unsafeCoerce k) $ fromAny an
303332
304−calcFp :: Typeable a => f a -> Fingerprint
305−calcFp = typeRepFingerprint . typeRep . fromF
333+calcFp :: forall a . Typeable a => Fingerprint
334+calcFp = typeRepFingerprint $ typeRep @a
306335
307−fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
308−fromListPairs kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans)
336+fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
337+fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
309338 where
310339 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
311− (fps, ans) = unzip $ fromSortedList $ sortWith fst $ nubByFst kvs
340+ (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
312341
313342----------------------------------------------------------------------------
314343-- Tree-like conversion