git.iko.soy/typerep-map
on from-internal help sign in
typerep-map/src/Data/TypeRep/OptimalVector.hs
118 lines 3.6 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE GADTs #-}
3{-# LANGUAGE KindSignatures #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE PolyKinds #-}
6{-# LANGUAGE TypeFamilies #-}
8module Data.TypeRep.OptimalVector
9 ( -- * Map type
10 TypeRepMap (..)
12 -- 'TypeRepMap' interface
13 , empty
14 , insert
15 , lookup
16 , size
18 -- * Helpful testing functions
19 , TF (..)
20 , fromList
21 ) where
23import Prelude hiding (lookup)
25import Control.Arrow ((&&&))
26import Data.Kind (Type)
27import Data.Proxy (Proxy (..))
28import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
29import Data.Word (Word64)
30import GHC.Base (Any, Int (..), Int#, uncheckedIShiftRA#, (+#), (-#), (<#))
31import GHC.Exts (inline, sortWith)
32import GHC.Fingerprint (Fingerprint (..))
33import Unsafe.Coerce (unsafeCoerce)
35import qualified Data.Vector as V
36import qualified Data.Vector.Unboxed as Unboxed
38data TypeRepMap (f :: k -> Type) = TypeRepMap
39 { fingerprintAs :: Unboxed.Vector Word64
40 , fingerprintBs :: Unboxed.Vector Word64
41 , anys :: V.Vector Any
42 }
44fromAny :: Any -> f a
45fromAny = unsafeCoerce
47-- | Empty structure.
48empty :: TypeRepMap f
49empty = TypeRepMap mempty mempty mempty
51-- | Inserts the value with its type as a key.
52insert :: forall a f . Typeable a => a -> TypeRepMap f -> TypeRepMap f
53insert = undefined
55-- | Looks up the value at the type.
56-- >>> let x = lookup $ insert (11 :: Int) empty
57-- >>> x :: Maybe Int
58-- Just 11
59-- >>> x :: Maybe ()
60-- Nothing
61lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
62lookup tVect = fromAny . (anys tVect V.!)
63 <$> binarySearch (typeRepFingerprint $ typeRep $ Proxy @a)
64 (fingerprintAs tVect)
65 (fingerprintBs tVect)
67-- | Returns the size of the 'TypeRepMap'.
68size :: TypeRepMap f -> Int
69size = Unboxed.length . fingerprintAs
71-- | Returns the index is found.
72binarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
73binarySearch (Fingerprint a b) fpAs fpBs =
74 let
75 !(I# len) = Unboxed.length fpAs
76 checkfpBs :: Int# -> Maybe Int
77 checkfpBs i =
78 case i <# len of
79 0# -> Nothing
80 _ ->
81 if a /= Unboxed.unsafeIndex fpAs (I# i)
82 then Nothing
83 else if b == Unboxed.unsafeIndex fpBs (I# i)
84 then Just (I# i)
85 else checkfpBs (i +# 1#)
86 in
87 inline (checkfpBs (binSearchHelp (-1#) len))
88 where
89 binSearchHelp :: Int# -> Int# -> Int#
90 binSearchHelp l r = case l <# (r -# 1#) of
91 0# -> r
92 _ ->
93 let m = uncheckedIShiftRA# (l +# r) 1# in
94 if Unboxed.unsafeIndex fpAs (I# m) < a
95 then binSearchHelp m r
96 else binSearchHelp l m
98----------------------------------------------------------------------------
99-- Functions for testing and benchmarking
100----------------------------------------------------------------------------
102data TF f where
103 TF :: Typeable a => f a -> TF f
105fromF :: Typeable a => f a -> Proxy a
106fromF _ = Proxy
108fromList :: forall f . [TF f] -> TypeRepMap f
109fromList tfs = TypeRepMap (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
110 where
111 (fpAs, fpBs) = unzip $ fmap (\(Fingerprint a b) -> (a, b)) fps
112 (fps, ans) = unzip $ sortWith fst $ map (fp &&& an) tfs
114 fp :: TF f -> Fingerprint
115 fp (TF x) = typeRepFingerprint $ typeRep $ fromF x
117 an :: TF f -> Any
118 an (TF x) = unsafeCoerce x