on 4b66779 help sign in
typerep-map/src/Data/TypeRep/OptimalVector.hs
177 lines 6.3 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 ScopedTypeVariables #-}
9{-# LANGUAGE TypeFamilies #-}
12module Data.TypeRep.OptimalVector
13 ( TypeRepVector (..)
14 , TF (..)
15 , empty
16 , insert
17 , lookup
18 , size
19 , fromList
20 ) where
22import Prelude hiding (lookup)
24import Control.Arrow ((&&&))
25import Data.Foldable (for_)
26import Data.Proxy (Proxy (..))
27import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
28import Data.Word (Word64)
29import GHC.Base (Any, liftM)
30import GHC.Exts (sortWith)
31import GHC.Fingerprint (Fingerprint (..))
32import Unsafe.Coerce (unsafeCoerce)
34import qualified Data.Vector as V
35import qualified Data.Vector.Generic as G
36import qualified Data.Vector.Generic.Mutable as M
37import qualified Data.Vector.Unboxed as Unboxed
39data instance Unboxed.MVector s Fingerprint = MFingerprintVectorOpt (Unboxed.MVector s Word64)
40data instance Unboxed.Vector Fingerprint = FingerprintVectorOpt (Unboxed.Vector Word64)
42instance Unboxed.Unbox Fingerprint
44instance M.MVector Unboxed.MVector Fingerprint where
45 {-# INLINE basicLength #-}
46 basicLength (MFingerprintVectorOpt v) = M.basicLength v `div` 2
47 {-# INLINE basicUnsafeSlice #-}
48 basicUnsafeSlice i n (MFingerprintVectorOpt v) =
49 MFingerprintVectorOpt $ M.basicUnsafeSlice (2 * i) (2 * n) v
50 {-# INLINE basicOverlaps #-}
51 basicOverlaps (MFingerprintVectorOpt v1) (MFingerprintVectorOpt v2) =
52 M.basicOverlaps v1 v2
53 {-# INLINE basicUnsafeNew #-}
54 basicUnsafeNew n = MFingerprintVectorOpt `liftM` M.basicUnsafeNew (2 * n)
55 {-# INLINE basicInitialize #-}
56 basicInitialize (MFingerprintVectorOpt v) = M.basicInitialize v
57 {-# INLINE basicUnsafeReplicate #-}
58 basicUnsafeReplicate n_ (Fingerprint a b) = do
59 v <- M.basicUnsafeNew $ 2 * n_
60 for_ [0, 2 .. 2 * n_ - 1] $ \i -> do
61 M.basicUnsafeWrite v i a
62 M.basicUnsafeWrite v (i + 1) b
63 return $ MFingerprintVectorOpt v
64 {-# INLINE basicUnsafeRead #-}
65 basicUnsafeRead (MFingerprintVectorOpt v) i = do
66 a <- M.basicUnsafeRead v (2 * i)
67 b <- M.basicUnsafeRead v (2 * i + 1)
68 return (Fingerprint a b)
69 {-# INLINE basicUnsafeWrite #-}
70 basicUnsafeWrite (MFingerprintVectorOpt v) i (Fingerprint a b) = do
71 M.basicUnsafeWrite v (2 * i) a
72 M.basicUnsafeWrite v (2 * i + 1) b
73 {-# INLINE basicClear #-}
74 basicClear (MFingerprintVectorOpt v) = M.basicClear v
75 {-# INLINE basicSet #-}
76 basicSet (MFingerprintVectorOpt v) (Fingerprint a b) = do
77 let n = M.basicLength v
78 for_ [0, 2 .. n - 1] $ \i -> do
79 M.basicUnsafeWrite v i a
80 M.basicUnsafeWrite v (i + 1) b
81 {-# INLINE basicUnsafeCopy #-}
82 basicUnsafeCopy (MFingerprintVectorOpt v1) (MFingerprintVectorOpt v2) =
83 M.basicUnsafeCopy v1 v2
84 {-# INLINE basicUnsafeMove #-}
85 basicUnsafeMove (MFingerprintVectorOpt v1) (MFingerprintVectorOpt v2) =
86 M.basicUnsafeMove v1 v2
87 {-# INLINE basicUnsafeGrow #-}
88 basicUnsafeGrow (MFingerprintVectorOpt v) m_ =
89 MFingerprintVectorOpt <$> M.basicUnsafeGrow v (2 * m_)
91instance G.Vector Unboxed.Vector Fingerprint where
92 {-# INLINE basicUnsafeFreeze #-}
93 basicUnsafeFreeze (MFingerprintVectorOpt v) =
94 FingerprintVectorOpt <$> G.basicUnsafeFreeze v
95 {-# INLINE basicUnsafeThaw #-}
96 basicUnsafeThaw (FingerprintVectorOpt v) = do
97 MFingerprintVectorOpt <$> G.basicUnsafeThaw v
98 {-# INLINE basicLength #-}
99 basicLength (FingerprintVectorOpt v) = G.basicLength v `div` 2
100 {-# INLINE basicUnsafeSlice #-}
101 basicUnsafeSlice i_ m_ (FingerprintVectorOpt v) =
102 FingerprintVectorOpt (G.basicUnsafeSlice (2 * i_) (2 * m_) v)
103 {-# INLINE basicUnsafeIndexM #-}
104 basicUnsafeIndexM (FingerprintVectorOpt v) i_ = do
105 a <- G.basicUnsafeIndexM v (2 * i_)
106 b <- G.basicUnsafeIndexM v (2 * i_ + 1)
107 return (Fingerprint a b)
108 {-# INLINE basicUnsafeCopy #-}
109 basicUnsafeCopy (MFingerprintVectorOpt v1) (FingerprintVectorOpt v2) =
110 G.basicUnsafeCopy v1 v2
111 {-# INLINE elemseq #-}
112 elemseq _ (Fingerprint a b)
113 = G.elemseq (undefined :: Unboxed.Vector a) a
114 . G.elemseq (undefined :: Unboxed.Vector b) b
116data TypeRepVector f = TypeRepVect
117 { fingerprints :: Unboxed.Vector Fingerprint
118 , anys :: V.Vector Any
119 }
121fromAny :: Any -> f a
122fromAny = unsafeCoerce
124-- | Empty structure.
125empty :: TypeRepVector f
126empty = TypeRepVect mempty mempty
128-- | Inserts the value with its type as a key.
129insert :: forall a f . Typeable a => a -> TypeRepVector f -> TypeRepVector f
130insert = undefined
132-- | Looks up the value at the type.
133-- >>> let x = lookup $ insert (11 :: Int) empty
134-- >>> x :: Maybe Int
135-- Just 11
136-- >>> x :: Maybe ()
137-- Nothing
138lookup :: forall a f . Typeable a => TypeRepVector f -> Maybe (f a)
139lookup tVect = fromAny . (anys tVect V.!)
140 <$> binarySearch (typeRepFingerprint (typeRep (Proxy :: Proxy a))) (fingerprints tVect)
142-- | Returns the size of the 'TypeRepVect'.
143size :: TypeRepVector f -> Int
144size = Unboxed.length . fingerprints
146data TF f where
147 TF :: Typeable a => f a -> TF f
149fromF :: Typeable a => f a -> Proxy a
150fromF _ = Proxy
152fromList :: forall f . [TF f] -> TypeRepVector f
153fromList tfs = TypeRepVect (Unboxed.fromList fps) (V.fromList ans)
154 where
155 (fps, ans) = unzip $ sortWith fst $ map (fp &&& an) tfs
157 fp :: TF f -> Fingerprint
158 fp (TF x) = typeRepFingerprint $ typeRep $ fromF x
160 an :: TF f -> Any
161 an (TF x) = unsafeCoerce x
164-- | Returns the index is found.
165binarySearch :: Fingerprint -> Unboxed.Vector Fingerprint -> Maybe Int
166binarySearch fp fpVect =
167 let ind = binSearchHelp (-1) (Unboxed.length fpVect) in
168 if fp == (fpVect Unboxed.! ind) then Just ind else Nothing
169 where
170 binSearchHelp :: Int -> Int -> Int
171 binSearchHelp l r = if l < r - 1
172 then (
173 let m = (l + r) `div` 2 in
174 if fpVect Unboxed.! m < fp
175 then binSearchHelp m r
176 else binSearchHelp l m )
177 else r