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