git.iko.soy/typerep-map
on ff8e4de help sign in
typerep-map/src/Data/TypeRep/Map.hs
170 lines 5.7 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE GADTs #-}
3{-# LANGUAGE KindSignatures #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE PolyKinds #-}
6{-# LANGUAGE TypeFamilies #-}
8{-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
10module Data.TypeRep.Map
11 ( -- * Map type
12 TypeRepMap (..)
14 -- 'TypeRepMap' interface
15 , empty
16 , insert
17 , lookup
18 , size
20 -- * Helpful testing functions
21 , TF (..)
22 , fromList
23 ) where
25import Prelude hiding (lookup)
27import Control.Arrow ((&&&))
28import Data.Kind (Type)
29import Data.Proxy (Proxy (..))
30import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
31import Data.Word (Word64)
32import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
33import GHC.Exts (inline, sortWith)
34import GHC.Fingerprint (Fingerprint (..))
35import GHC.Prim (eqWord#, ltWord#)
36import GHC.Word (Word64 (..))
37import Unsafe.Coerce (unsafeCoerce)
39import qualified Data.Vector as V
40import qualified Data.Vector.Unboxed as Unboxed
42data TypeRepMap (f :: k -> Type) = TypeRepMap
43 { fingerprintAs :: Unboxed.Vector Word64
44 , fingerprintBs :: Unboxed.Vector Word64
45 , anys :: V.Vector Any
46 }
48fromAny :: Any -> f a
49fromAny = unsafeCoerce
51-- | Empty structure.
52empty :: TypeRepMap f
53empty = TypeRepMap mempty mempty mempty
55-- | Inserts the value with its type as a key.
56insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
57insert x = fromListPairs . addX . toPairList
58 where
59 toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
60 toPairList (TypeRepMap as bs ans) = zip (zipWith Fingerprint
61 (Unboxed.toList as)
62 (Unboxed.toList bs))
63 (V.toList ans)
65 pairX :: (Fingerprint, Any)
66 pairX@(fpX, _) = (calcFp x, unsafeCoerce x)
68 addX :: [(Fingerprint, Any)] -> [(Fingerprint, Any)]
69 addX l = pairX : filter ((/= fpX) . fst) l
70{-# INLINE insert #-}
72-- | Looks up the value at the type.
73-- >>> let x = lookup $ insert (11 :: Int) empty
74-- >>> x :: Maybe Int
75-- Just 11
76-- >>> x :: Maybe ()
77-- Nothing
78lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
79lookup tVect = fromAny . (anys tVect V.!)
80 <$> cachedBinarySearch (typeRepFingerprint $ typeRep $ Proxy @a)
81 (fingerprintAs tVect)
82 (fingerprintBs tVect)
83-- {-# INLINE lookup #-}
85-- | Returns the size of the 'TypeRepMap'.
86size :: TypeRepMap f -> Int
87size = Unboxed.length . fingerprintAs
89-- | Binary searched based on this article
90-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
91-- with modification for our two-vector search case.
92cachedBinarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
93cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
94 where
95 go :: Int# -> Maybe Int
96 go i = case i <# len of
97 0# -> Nothing
98 _ -> let !(W64# valA) = Unboxed.unsafeIndex fpAs (I# i) in case a `ltWord#` valA of
99 0# -> case a `eqWord#` valA of
100 0# -> go (2# *# i +# 2#)
101 _ -> let !(W64# valB) = Unboxed.unsafeIndex fpBs (I# i) in case b `eqWord#` valB of
102 0# -> case b `ltWord#` valB of
103 0# -> go (2# *# i +# 2#)
104 _ -> go (2# *# i +# 1#)
105 _ -> Just (I# i)
106 _ -> go (2# *# i +# 1#)
108 len :: Int#
109 len = let !(I# l) = Unboxed.length fpAs in l
110-- {-# INLINE cachedBinarySearch #-}
112----------------------------------------------------------------------------
113-- Functions for testing and benchmarking
114----------------------------------------------------------------------------
116data TF f where
117 TF :: Typeable a => f a -> TF f
119fromF :: Typeable a => f a -> Proxy a
120fromF _ = Proxy
122calcFp :: Typeable a => f a -> Fingerprint
123calcFp = typeRepFingerprint . typeRep . fromF
125fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
126fromListPairs kvs = TypeRepMap (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
127 where
128 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
129 (fps, ans) = unzip $ breadthFirst $ fromListToTree $ sortWith fst kvs
131fromList :: forall f . [TF f] -> TypeRepMap f
132fromList = fromListPairs . map (fp &&& an)
133 where
134 fp :: TF f -> Fingerprint
135 fp (TF x) = calcFp x
137 an :: TF f -> Any
138 an (TF x) = unsafeCoerce x
140----------------------------------------------------------------------------
141-- Tree
142----------------------------------------------------------------------------
144data Tree a = Leaf | Node a (Tree a) (Tree a)
145 deriving (Show)
147fromListToTree :: [a] -> Tree a
148fromListToTree [] = Leaf
149fromListToTree xs =
150 let len = length xs
151 in case splitAt (len `div` 2) xs of
152 ([], []) -> Leaf
153 (ls, []) -> Node (last ls) (fromListToTree $ init ls) Leaf
154 ([], (r:rs)) -> Node r Leaf $ fromListToTree rs
155 (ls, r:rs) -> Node r (fromListToTree ls) (fromListToTree rs)
157breadthFirst :: Tree a -> [a]
158breadthFirst tree = bf [tree]
159 where
160 bf :: [Tree a] -> [a]
161 bf [] = []
162 bf xs = mapValue xs ++ bf (concatMap children xs)
164 mapValue :: [Tree a] -> [a]
165 mapValue [] = []
166 mapValue (Leaf : xs) = mapValue xs
167 mapValue (Node a _ _ : xs) = a : mapValue xs
169 children Leaf = []
170 children (Node _ l r) = [l, r]