git.iko.soy/typerep-map
on 71aa479 help sign in
typerep-map/internal/Data/TypeRep/Vector.hs
189 lines 6.7 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE BangPatterns #-}
4{-# LANGUAGE GADTs #-}
5{-# LANGUAGE KindSignatures #-}
6{-# LANGUAGE MagicHash #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PolyKinds #-}
9{-# LANGUAGE TypeFamilies #-}
11module Data.TypeRep.Vector
12 ( TypeRepVector (..)
13 , TF (..)
14 , empty
15 , insert
16 , lookup
17 , size
18 , fromList
19 ) where
21import Prelude hiding (lookup)
23import Control.Arrow ((&&&))
24import Data.Proxy (Proxy (..))
25import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
26import Data.Word (Word64)
27import GHC.Base hiding (empty)
28import GHC.Exts (sortWith)
29import GHC.Fingerprint (Fingerprint (..))
30import Unsafe.Coerce (unsafeCoerce)
32import qualified Data.Vector as V
33import qualified Data.Vector.Generic as G
34import qualified Data.Vector.Generic.Mutable as M
35import qualified Data.Vector.Unboxed as Unboxed
37data instance Unboxed.MVector s Fingerprint = MFingerprintVector (Unboxed.MVector s Word64) (Unboxed.MVector s Word64)
38data instance Unboxed.Vector Fingerprint = FingerprintVector (Unboxed.Vector Word64) (Unboxed.Vector Word64)
40instance Unboxed.Unbox Fingerprint
42instance M.MVector Unboxed.MVector Fingerprint where
43 {-# INLINE basicLength #-}
44 basicLength (MFingerprintVector x _) = M.basicLength x
45 {-# INLINE basicUnsafeSlice #-}
46 basicUnsafeSlice i m (MFingerprintVector a b) =
47 MFingerprintVector (M.basicUnsafeSlice i m a) (M.basicUnsafeSlice i m b)
48 {-# INLINE basicOverlaps #-}
49 basicOverlaps (MFingerprintVector as1 bs1) (MFingerprintVector as2 bs2) =
50 M.basicOverlaps as1 as2 || M.basicOverlaps bs1 bs2
51 {-# INLINE basicUnsafeNew #-}
52 basicUnsafeNew n_ = do
53 as <- M.basicUnsafeNew n_
54 bs <- M.basicUnsafeNew n_
55 return $ MFingerprintVector as bs
56 {-# INLINE basicInitialize #-}
57 basicInitialize (MFingerprintVector as bs) = do
58 M.basicInitialize as
59 M.basicInitialize bs
60 {-# INLINE basicUnsafeReplicate #-}
61 basicUnsafeReplicate n_ (Fingerprint a b) = do
62 as <- M.basicUnsafeReplicate n_ a
63 bs <- M.basicUnsafeReplicate n_ b
64 return $ MFingerprintVector as bs
65 {-# INLINE basicUnsafeRead #-}
66 basicUnsafeRead (MFingerprintVector as bs) i_ = do
67 a <- M.basicUnsafeRead as i_
68 b <- M.basicUnsafeRead bs i_
69 return (Fingerprint a b)
70 {-# INLINE basicUnsafeWrite #-}
71 basicUnsafeWrite (MFingerprintVector as bs) i_ (Fingerprint a b) = do
72 M.basicUnsafeWrite as i_ a
73 M.basicUnsafeWrite bs i_ b
74 {-# INLINE basicClear #-}
75 basicClear (MFingerprintVector as bs) = do
76 M.basicClear as
77 M.basicClear bs
78 {-# INLINE basicSet #-}
79 basicSet (MFingerprintVector as bs) (Fingerprint a b) = do
80 M.basicSet as a
81 M.basicSet bs b
82 {-# INLINE basicUnsafeCopy #-}
83 basicUnsafeCopy (MFingerprintVector as1 bs1) (MFingerprintVector as2 bs2) = do
84 M.basicUnsafeCopy as1 as2
85 M.basicUnsafeCopy bs1 bs2
86 {-# INLINE basicUnsafeMove #-}
87 basicUnsafeMove (MFingerprintVector as1 bs1) (MFingerprintVector as2 bs2) = do
88 M.basicUnsafeMove as1 as2
89 M.basicUnsafeMove bs1 bs2
90 {-# INLINE basicUnsafeGrow #-}
91 basicUnsafeGrow (MFingerprintVector as bs) m_ = do
92 as' <- M.basicUnsafeGrow as m_
93 bs' <- M.basicUnsafeGrow bs m_
94 return $ MFingerprintVector as' bs'
96instance G.Vector Unboxed.Vector Fingerprint where
97 {-# INLINE basicUnsafeFreeze #-}
98 basicUnsafeFreeze (MFingerprintVector as bs) = do
99 as' <- G.basicUnsafeFreeze as
100 bs' <- G.basicUnsafeFreeze bs
101 return $ FingerprintVector as' bs'
102 {-# INLINE basicUnsafeThaw #-}
103 basicUnsafeThaw (FingerprintVector as bs) = do
104 as' <- G.basicUnsafeThaw as
105 bs' <- G.basicUnsafeThaw bs
106 return $ MFingerprintVector as' bs'
107 {-# INLINE basicLength #-}
108 basicLength (FingerprintVector x _) = G.basicLength x
109 {-# INLINE basicUnsafeSlice #-}
110 basicUnsafeSlice i_ m_ (FingerprintVector as bs) =
111 FingerprintVector (G.basicUnsafeSlice i_ m_ as) (G.basicUnsafeSlice i_ m_ bs)
112 {-# INLINE basicUnsafeIndexM #-}
113 basicUnsafeIndexM (FingerprintVector as bs) i_ = do
114 a <- G.basicUnsafeIndexM as i_
115 b <- G.basicUnsafeIndexM bs i_
116 return (Fingerprint a b)
117 {-# INLINE basicUnsafeCopy #-}
118 basicUnsafeCopy (MFingerprintVector as1 bs1) (FingerprintVector as2 bs2) = do
119 G.basicUnsafeCopy as1 as2
120 G.basicUnsafeCopy bs1 bs2
121 {-# INLINE elemseq #-}
122 elemseq _ (Fingerprint a b)
123 = G.elemseq (undefined :: Unboxed.Vector a) a
124 . G.elemseq (undefined :: Unboxed.Vector b) b
126data TypeRepVector f = TypeRepVect
127 { fingerprints :: Unboxed.Vector Fingerprint
128 , anys :: V.Vector Any
129 }
131fromAny :: Any -> f a
132fromAny = unsafeCoerce
134-- | Empty structure.
135empty :: TypeRepVector f
136empty = TypeRepVect mempty mempty
138-- | Inserts the value with its type as a key.
139insert :: forall a f . Typeable a => a -> TypeRepVector f -> TypeRepVector f
140insert = undefined
142-- | Looks up the value at the type.
143-- >>> let x = lookup $ insert (11 :: Int) empty
144-- >>> x :: Maybe Int
145-- Just 11
146-- >>> x :: Maybe ()
147-- Nothing
148lookup :: forall a f . Typeable a => TypeRepVector f -> Maybe (f a)
149lookup tVect = fromAny . (anys tVect V.!)
150 <$> binarySearch (typeRepFingerprint (typeRep (Proxy :: Proxy a))) (fingerprints tVect)
152-- | Returns the size of the 'TypeRepVect'.
153size :: TypeRepVector f -> Int
154size = Unboxed.length . fingerprints
156data TF f where
157 TF :: Typeable a => f a -> TF f
159fromF :: Typeable a => f a -> Proxy a
160fromF _ = Proxy
162fromList :: forall f . [TF f] -> TypeRepVector f
163fromList tfs = TypeRepVect (Unboxed.fromList fps) (V.fromList ans)
164 where
165 (fps, ans) = unzip $ sortWith fst $ map (fp &&& an) tfs
167 fp :: TF f -> Fingerprint
168 fp (TF x) = typeRepFingerprint $ typeRep $ fromF x
170 an :: TF f -> Any
171 an (TF x) = unsafeCoerce x
173-- | Returns the index is found.
174binarySearch :: Fingerprint -> Unboxed.Vector Fingerprint -> Maybe Int
175binarySearch fp fpVect =
176 let
177 !(I# len) = Unboxed.length fpVect
178 ind = I# (binSearchHelp (-1#) len)
179 in
180 if fp == (fpVect Unboxed.! ind) then Just ind else Nothing
181 where
182 binSearchHelp :: Int# -> Int# -> Int#
183 binSearchHelp l r = case l <# (r -# 1#) of
184 0# -> r
185 _ ->
186 let m = uncheckedIShiftRA# (l +# r) 1# in
187 if Unboxed.unsafeIndex fpVect (I# m) < fp
188 then binSearchHelp m r
189 else binSearchHelp l m