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