{-# LANGUAGE ExistentialQuantification #-}
module Data.StorableVector.Cursor (
Vector,
create,
unfoldrNTerm,
unfoldrN,
pack,
cons,
zipWith,
zipNWith,
whileM,
switchL,
viewL,
foldr,
unpack,
null,
drop,
take,
filter,
) where
import Control.Exception (assert, )
import Control.Monad.Trans.State (StateT(StateT), runStateT, )
import Data.IORef (IORef, newIORef, readIORef, writeIORef, )
import Foreign.Storable (Storable(peekElemOff, pokeElemOff))
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, )
import Data.StorableVector.Memory (mallocForeignPtrArray, )
import Control.Monad (when)
import Data.Maybe (isNothing)
import qualified System.Unsafe as Unsafe
import qualified Data.List.HT as ListHT
import Data.Tuple.HT (mapSnd)
import Data.Maybe (Maybe(Just, Nothing), maybe)
import Data.Function (const, ($))
import Data.Tuple (uncurry)
import Data.Bool (Bool(True,False), (||))
import Data.Ord ((<), (<=), (>=), min, max)
import Data.Eq ((==))
import Control.Monad (Monad, return, fmap, (>>))
import Text.Show (Show, showsPrec)
import Prelude (IO, pred, succ, ($!), Int, (+), (-))
data Generator a =
forall s.
Generator
!(StateT s Maybe a)
{-# UNPACK #-}
!(IORef (Maybe s))
data Buffer a =
Buffer {
forall a. Buffer a -> ForeignPtr a
memory :: {-# UNPACK #-} !(ForeignPtr a),
forall a. Buffer a -> Int
_size :: {-# UNPACK #-} !Int,
forall a. Buffer a -> Generator a
_gen :: !(Generator a),
forall a. Buffer a -> IORef Int
cursor :: {-# UNPACK #-} !(IORef Int)
}
data Vector a =
Vector {
forall a. Vector a -> Buffer a
buffer :: {-# UNPACK #-} !(Buffer a),
forall a. Vector a -> Int
start :: {-# UNPACK #-} !Int,
forall a. Vector a -> Int
maxLen :: {-# UNPACK #-} !Int
}
{-# INLINE create #-}
create :: (Storable a) => Int -> Generator a -> Buffer a
create :: forall a. Storable a => Int -> Generator a -> Buffer a
create Int
l Generator a
g = forall a. IO a -> a
Unsafe.performIO (forall a. Storable a => Int -> Generator a -> IO (Buffer a)
createIO Int
l Generator a
g)
createIO :: (Storable a) => Int -> Generator a -> IO (Buffer a)
createIO :: forall a. Storable a => Int -> Generator a -> IO (Buffer a)
createIO Int
l Generator a
g = do
ForeignPtr a
fp <- forall a. Storable a => Int -> IO (ForeignPtr a)
mallocForeignPtrArray Int
l
IORef Int
cur <- forall a. a -> IO (IORef a)
newIORef Int
0
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a.
ForeignPtr a -> Int -> Generator a -> IORef Int -> Buffer a
Buffer ForeignPtr a
fp Int
l Generator a
g IORef Int
cur
unfoldrNTerm :: (Storable b) =>
Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm :: forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm Int
l a -> Maybe (b, a)
f a
x0 =
forall a. IO a -> a
Unsafe.performIO (forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b)
unfoldrNTermIO Int
l a -> Maybe (b, a)
f a
x0)
unfoldrNTermIO :: (Storable b) =>
Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b)
unfoldrNTermIO :: forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b)
unfoldrNTermIO Int
l a -> Maybe (b, a)
f a
x0 =
do IORef (Maybe a)
ref <- forall a. a -> IO (IORef a)
newIORef (forall a. a -> Maybe a
Just a
x0)
Buffer b
buf <- forall a. Storable a => Int -> Generator a -> IO (Buffer a)
createIO Int
l (forall a s. StateT s Maybe a -> IORef (Maybe s) -> Generator a
Generator (forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT a -> Maybe (b, a)
f) IORef (Maybe a)
ref)
forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Buffer a -> Int -> Int -> Vector a
Vector Buffer b
buf Int
0 Int
l)
unfoldrN :: (Storable b) =>
Int -> (a -> Maybe (b, a)) -> a -> (Vector b, Maybe a)
unfoldrN :: forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> (Vector b, Maybe a)
unfoldrN Int
l a -> Maybe (b, a)
f a
x0 =
forall a. IO a -> a
Unsafe.performIO (forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b, Maybe a)
unfoldrNIO Int
l a -> Maybe (b, a)
f a
x0)
unfoldrNIO :: (Storable b) =>
Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b, Maybe a)
unfoldrNIO :: forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b, Maybe a)
unfoldrNIO Int
l a -> Maybe (b, a)
f a
x0 =
do IORef (Maybe a)
ref <- forall a. a -> IO (IORef a)
newIORef (forall a. a -> Maybe a
Just a
x0)
Buffer b
buf <- forall a. Storable a => Int -> Generator a -> IO (Buffer a)
createIO Int
l (forall a s. StateT s Maybe a -> IORef (Maybe s) -> Generator a
Generator (forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT a -> Maybe (b, a)
f) IORef (Maybe a)
ref)
Maybe a
s <- forall a. IO a -> IO a
Unsafe.interleaveIO forall a b. (a -> b) -> a -> b
$
do forall a. Storable a => Int -> Buffer a -> IO ()
evaluateToIO Int
l Buffer b
buf
forall a. IORef a -> IO a
readIORef IORef (Maybe a)
ref
forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Buffer a -> Int -> Int -> Vector a
Vector Buffer b
buf Int
0 Int
l, Maybe a
s)
{-# INLINE pack #-}
pack :: (Storable a) => Int -> [a] -> Vector a
pack :: forall a. Storable a => Int -> [a] -> Vector a
pack Int
n = forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm Int
n forall a. [a] -> Maybe (a, [a])
ListHT.viewL
{-# INLINE cons #-}
cons :: Storable a =>
a -> Vector a -> Vector a
cons :: forall a. Storable a => a -> Vector a -> Vector a
cons a
x Vector a
xs =
forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm (forall a. Enum a => a -> a
succ (forall a. Vector a -> Int
maxLen Vector a
xs))
(\(Maybe a
mx0,Vector a
xs0) ->
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd ((,) forall a. Maybe a
Nothing)) forall a b. (a -> b) -> a -> b
$
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
(forall a. Storable a => Vector a -> Maybe (a, Vector a)
viewL Vector a
xs0)
(\a
x0 -> forall a. a -> Maybe a
Just (a
x0, Vector a
xs0))
Maybe a
mx0) forall a b. (a -> b) -> a -> b
$
(forall a. a -> Maybe a
Just a
x, Vector a
xs)
{-# INLINE zipWith #-}
zipWith :: (Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith :: forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith a -> b -> c
f Vector a
ps0 Vector b
qs0 =
forall a b c.
(Storable a, Storable b, Storable c) =>
Int -> (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipNWith (forall a. Ord a => a -> a -> a
min (forall a. Vector a -> Int
maxLen Vector a
ps0) (forall a. Vector a -> Int
maxLen Vector b
qs0)) a -> b -> c
f Vector a
ps0 Vector b
qs0
{-# INLINE zipNWith #-}
zipNWith :: (Storable a, Storable b, Storable c) =>
Int -> (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipNWith :: forall a b c.
(Storable a, Storable b, Storable c) =>
Int -> (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipNWith Int
n a -> b -> c
f Vector a
ps0 Vector b
qs0 =
forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm Int
n
(\(Vector a
ps,Vector b
qs) ->
do (a
ph,Vector a
pt) <- forall a. Storable a => Vector a -> Maybe (a, Vector a)
viewL Vector a
ps
(b
qh,Vector b
qt) <- forall a. Storable a => Vector a -> Maybe (a, Vector a)
viewL Vector b
qs
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b -> c
f a
ph b
qh, (Vector a
pt,Vector b
qt)))
(Vector a
ps0,Vector b
qs0)
advanceIO :: Storable a =>
Buffer a -> IO (Maybe a)
advanceIO :: forall a. Storable a => Buffer a -> IO (Maybe a)
advanceIO (Buffer ForeignPtr a
p Int
sz (Generator StateT s Maybe a
n IORef (Maybe s)
s) IORef Int
cr) =
do Int
c <- forall a. IORef a -> IO a
readIORef IORef Int
cr
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
c forall a. Ord a => a -> a -> Bool
< Int
sz) forall a b. (a -> b) -> a -> b
$
do forall a. IORef a -> a -> IO ()
writeIORef IORef Int
cr (forall a. Enum a => a -> a
succ Int
c)
Maybe s
ms <- forall a. IORef a -> IO a
readIORef IORef (Maybe s)
s
case Maybe s
ms of
Maybe s
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
Just s
s0 ->
case forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT s Maybe a
n s
s0 of
Maybe (a, s)
Nothing ->
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe s)
s forall a. Maybe a
Nothing forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
Just (a
a,s
s1) ->
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe s)
s (forall a. a -> Maybe a
Just s
s1) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
p (\Ptr a
q -> forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr a
q Int
c a
a) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
a)
evaluateToIO :: Storable a =>
Int -> Buffer a -> IO ()
evaluateToIO :: forall a. Storable a => Int -> Buffer a -> IO ()
evaluateToIO Int
l buf :: Buffer a
buf@(Buffer ForeignPtr a
_p Int
_sz Generator a
_g IORef Int
cr) =
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m ()
whileM
(forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Ord a => a -> a -> Bool
<Int
l) (forall a. IORef a -> IO a
readIORef IORef Int
cr))
(forall a. Storable a => Buffer a -> IO (Maybe a)
advanceIO Buffer a
buf)
whileM :: Monad m => m Bool -> m a -> m ()
whileM :: forall (m :: * -> *) a. Monad m => m Bool -> m a -> m ()
whileM m Bool
p m a
f =
let recourse :: m ()
recourse =
do Bool
b <- m Bool
p
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (m a
f forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m ()
recourse)
in m ()
recourse
{-# INLINE switchL #-}
switchL :: Storable a => b -> (a -> Vector a -> b) -> Vector a -> b
switchL :: forall a b.
Storable a =>
b -> (a -> Vector a -> b) -> Vector a -> b
switchL b
n a -> Vector a -> b
j Vector a
v = forall b a. b -> (a -> b) -> Maybe a -> b
maybe b
n (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Vector a -> b
j) (forall a. Storable a => Vector a -> Maybe (a, Vector a)
viewL Vector a
v)
_obviousNullIO :: Vector a -> IO Bool
_obviousNullIO :: forall a. Vector a -> IO Bool
_obviousNullIO (Vector (Buffer ForeignPtr a
_ Int
_ (Generator StateT s Maybe a
_ IORef (Maybe s)
s) IORef Int
_) Int
_ Int
ml) =
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
ml forall a. Ord a => a -> a -> Bool
>= Int
0) forall a b. (a -> b) -> a -> b
$
do Maybe s
b <- forall a. IORef a -> IO a
readIORef IORef (Maybe s)
s
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
ml forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| forall a. Maybe a -> Bool
isNothing Maybe s
b)
{-# INLINE viewL #-}
viewL :: Storable a => Vector a -> Maybe (a, Vector a)
viewL :: forall a. Storable a => Vector a -> Maybe (a, Vector a)
viewL Vector a
v = forall a. IO a -> a
Unsafe.performIO (forall a. Storable a => Vector a -> IO (Maybe (a, Vector a))
viewLIO Vector a
v)
{-# INLINE viewLIO #-}
viewLIO :: Storable a => Vector a -> IO (Maybe (a, Vector a))
viewLIO :: forall a. Storable a => Vector a -> IO (Maybe (a, Vector a))
viewLIO (Vector Buffer a
buf Int
st Int
ml) =
do Int
c <- forall a. IORef a -> IO a
readIORef (forall a. Buffer a -> IORef Int
cursor Buffer a
buf)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
a -> (a
a, forall a. Buffer a -> Int -> Int -> Vector a
Vector Buffer a
buf (forall a. Enum a => a -> a
succ Int
st) (forall a. Enum a => a -> a
pred Int
ml)))) forall a b. (a -> b) -> a -> b
$
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
st forall a. Ord a => a -> a -> Bool
<= Int
c) forall a b. (a -> b) -> a -> b
$
if Int
st forall a. Eq a => a -> a -> Bool
== Int
c
then forall a. Storable a => Buffer a -> IO (Maybe a)
advanceIO Buffer a
buf
else forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (forall a. Buffer a -> ForeignPtr a
memory Buffer a
buf) (\Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
st)
{-# INLINE foldr #-}
foldr :: (Storable a) => (a -> b -> b) -> b -> Vector a -> b
foldr :: forall a b. Storable a => (a -> b -> b) -> b -> Vector a -> b
foldr a -> b -> b
k b
z =
let recourse :: Vector a -> b
recourse = forall a b.
Storable a =>
b -> (a -> Vector a -> b) -> Vector a -> b
switchL b
z (\a
h Vector a
t -> a -> b -> b
k a
h (Vector a -> b
recourse Vector a
t))
in Vector a -> b
recourse
{-# INLINE unpack #-}
unpack :: (Storable a) => Vector a -> [a]
unpack :: forall a. Storable a => Vector a -> [a]
unpack = forall a b. Storable a => (a -> b -> b) -> b -> Vector a -> b
foldr (:) []
instance (Show a, Storable a) => Show (Vector a) where
showsPrec :: Int -> Vector a -> ShowS
showsPrec Int
p Vector a
x = forall a. Show a => Int -> a -> ShowS
showsPrec Int
p (forall a. Storable a => Vector a -> [a]
unpack Vector a
x)
{-# INLINE null #-}
null :: Storable a => Vector a -> Bool
null :: forall a. Storable a => Vector a -> Bool
null = forall a b.
Storable a =>
b -> (a -> Vector a -> b) -> Vector a -> b
switchL Bool
True (forall a b. a -> b -> a
const (forall a b. a -> b -> a
const Bool
False))
drop :: (Storable a) => Int -> Vector a -> Vector a
drop :: forall a. Storable a => Int -> Vector a -> Vector a
drop Int
n Vector a
v = forall a. IO a -> a
Unsafe.performIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> Vector a -> IO (Vector a)
dropIO Int
n Vector a
v
dropIO :: (Storable a) => Int -> Vector a -> IO (Vector a)
dropIO :: forall a. Storable a => Int -> Vector a -> IO (Vector a)
dropIO Int
n Vector a
v =
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
nforall a. Ord a => a -> a -> Bool
>=Int
0) forall a b. (a -> b) -> a -> b
$
let pos :: Int
pos = forall a. Ord a => a -> a -> a
min (forall a. Vector a -> Int
maxLen Vector a
v) (forall a. Vector a -> Int
start Vector a
v forall a. Num a => a -> a -> a
+ Int
n)
in do forall a. Storable a => Int -> Buffer a -> IO ()
evaluateToIO Int
pos (forall a. Vector a -> Buffer a
buffer Vector a
v)
forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Buffer a -> Int -> Int -> Vector a
Vector (forall a. Vector a -> Buffer a
buffer Vector a
v) Int
pos (forall a. Ord a => a -> a -> a
max Int
0 (forall a. Vector a -> Int
maxLen Vector a
v forall a. Num a => a -> a -> a
- Int
n)))
take :: (Storable a) => Int -> Vector a -> Vector a
take :: forall a. Storable a => Int -> Vector a -> Vector a
take Int
n Vector a
v =
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
nforall a. Ord a => a -> a -> Bool
>=Int
0) forall a b. (a -> b) -> a -> b
$
Vector a
v{maxLen :: Int
maxLen = forall a. Ord a => a -> a -> a
min Int
n (forall a. Vector a -> Int
maxLen Vector a
v)}
{-# INLINE filter #-}
filter :: (Storable a) => (a -> Bool) -> Vector a -> Vector a
filter :: forall a. Storable a => (a -> Bool) -> Vector a -> Vector a
filter a -> Bool
p Vector a
xs0 =
forall b a.
Storable b =>
Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm (forall a. Vector a -> Int
maxLen Vector a
xs0)
(let recourse :: Vector a -> Maybe (a, Vector a)
recourse = forall a b.
Storable a =>
b -> (a -> Vector a -> b) -> Vector a -> b
switchL forall a. Maybe a
Nothing (\a
x Vector a
xs -> if a -> Bool
p a
x then forall a. a -> Maybe a
Just (a
x,Vector a
xs) else Vector a -> Maybe (a, Vector a)
recourse Vector a
xs)
in Vector a -> Maybe (a, Vector a)
recourse)
Vector a
xs0