1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Vector.Hybrid.Internal (Vector) where
import Control.Monad (liftM2)
import Data.Functor.Identity (Identity(..))
import GHC.ST (ST, runST)
import Text.Read (ReadPrec, readPrec)
import Data.Kind (Type)
-----
class Monad m => PrimMonad m where
type PrimState m
instance PrimMonad (ST s) where
type PrimState (ST s) = s
class GMVector v a where
gmbasicLength :: v s a -> Int
gmbasicUnsafeSlice :: Int -> Int -> v s a -> v s a
gmbasicUnsafeNew :: PrimMonad m => Int -> m (v (PrimState m) a)
gmbasicUnsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()
type family GMutable (v :: Type -> Type) :: Type -> Type -> Type
class GMVector (GMutable v) a => GVector v a where
gbasicUnsafeFreeze :: PrimMonad m => GMutable v (PrimState m) a -> m (v a)
data Step s a where
Yield :: a -> s -> Step s a
instance Functor (Step s) where
{-# INLINE fmap #-}
fmap f (Yield x s) = Yield (f x) s
data Stream m a = forall s. Stream (s -> m (Step s a)) s
data Chunk v a = Chunk Int (forall m. (PrimMonad m, GVector v a) => GMutable v (PrimState m) a -> m ())
data New v a = New { newrun :: forall s. ST s (GMutable v s a) }
type MBundle m v a = Stream m (Chunk v a)
type Bundle v a = MBundle Identity v a
mbfromStream :: Monad m => Stream m a -> MBundle m v a
{-# INLINE mbfromStream #-}
mbfromStream (Stream step t) = Stream step' t
where
step' s = do r <- step s
return $ fmap (\x -> Chunk 1 (\v -> gmbasicUnsafeWrite v 0 x)) r
mbunsafeFromList :: Monad m => [a] -> MBundle m v a
{-# INLINE [1] mbunsafeFromList #-}
mbunsafeFromList xs = mbfromStream (sfromList xs)
blift :: Monad m => Bundle v a -> MBundle m v a
{-# INLINE [1] blift #-}
blift (Stream vstep t) = Stream (return . runIdentity . vstep) t
sfromList :: Monad m => [a] -> Stream m a
{-# INLINE sfromList #-}
sfromList zs = Stream step zs
where
step (x:xs) = return (Yield x xs)
step _ = undefined
sfoldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
{-# INLINE [1] sfoldlM #-}
sfoldlM m w (Stream step t) = foldlM_loop w t
where
foldlM_loop z s
= do
r <- step s
case r of
Yield x s' -> do { z' <- m z x; foldlM_loop z' s' }
gmvunstream :: (PrimMonad m, GVector v a)
=> Bundle v a -> m (GMutable v (PrimState m) a)
{-# INLINE [1] gmvunstream #-}
gmvunstream s = gmvmunstreamUnknown (blift s)
gmvmunstreamUnknown :: (PrimMonad m, GVector v a)
=> MBundle m v a -> m (GMutable v (PrimState m) a)
{-# INLINE gmvmunstreamUnknown #-}
gmvmunstreamUnknown s
= do
v <- gmbasicUnsafeNew 0
(_, _) <- sfoldlM copyChunk (v,0) s
return undefined
where
{-# INLINE [0] copyChunk #-}
copyChunk (v,i) (Chunk n f)
= do
let j = i+n
v' <- if gmbasicLength v < j
then gmbasicUnsafeNew undefined
else return v
f (gmbasicUnsafeSlice i n v')
return (v',j)
newunstream :: GVector v a => Bundle v a -> New v a
{-# INLINE [1] newunstream #-}
newunstream s = s `seq` New (gmvunstream s)
gnew :: GVector v a => New v a -> v a
{-# INLINE [1] gnew #-}
gnew m = m `seq` runST (gbasicUnsafeFreeze =<< newrun m)
gunstream :: GVector v a => Bundle v a -> v a
{-# INLINE gunstream #-}
gunstream s = gnew (newunstream s)
gfromList :: GVector v a => [a] -> v a
{-# INLINE gfromList #-}
gfromList = gunstream . mbunsafeFromList
greadPrec :: (GVector v a, Read a) => ReadPrec (v a)
{-# INLINE greadPrec #-}
greadPrec = do
xs <- readPrec
return (gfromList xs)
-----
data MVector :: (Type -> Type -> Type) ->
(Type -> Type -> Type) ->
(Type -> Type -> Type) where
MV :: !(u s a) -> !(v s b) -> MVector u v s (a, b)
instance (GMVector u a, GMVector v b) => GMVector (MVector u v) (a, b) where
gmbasicLength (MV ks _) = gmbasicLength ks
gmbasicUnsafeSlice s e (MV ks vs) = MV (gmbasicUnsafeSlice s e ks) (gmbasicUnsafeSlice s e vs)
gmbasicUnsafeNew n = liftM2 MV (gmbasicUnsafeNew n) (gmbasicUnsafeNew n)
-- Removing this INLINE pragma makes it compile
{-# INLINE gmbasicUnsafeNew #-}
gmbasicUnsafeWrite (MV ks vs) n (k,v) = do
gmbasicUnsafeWrite ks n k
gmbasicUnsafeWrite vs n v
data Vector :: (Type -> Type) -> (Type -> Type) -> Type -> Type
type instance GMutable (Vector u v) = MVector (GMutable u) (GMutable v)
instance (GVector u a, GVector v b) => GVector (Vector u v) (a, b) where
gbasicUnsafeFreeze = undefined
instance (GVector u a, GVector v b, Read a, Read b, c ~ (a, b)) => Read (Vector u v c) where
readPrec = greadPrec
|
