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