module Streamly.Internal.Data.Array.Stream.Fold.Foreign
(
Fold (..)
, fromFold
, fromParser
, fromParserD
, fromArrayFold
, rmapM
, fromPure
, fromEffect
, serialWith
, concatMap
, take
)
where
#include "ArrayMacros.h"
import Control.Applicative (liftA2)
import Control.Exception (assert)
import Control.Monad.Catch (MonadThrow)
import Control.Monad.IO.Class (MonadIO(..))
import Foreign.Ptr (minusPtr, plusPtr)
import Foreign.Storable (Storable(..))
import GHC.Types (SPEC(..))
import Streamly.Internal.Data.Array.Foreign.Mut.Type (touch)
import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
import Streamly.Internal.Data.Parser.ParserD (Initial(..), Step(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import qualified Streamly.Internal.Data.Array.Foreign as Array
import qualified Streamly.Internal.Data.Fold as Fold
import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
import qualified Streamly.Internal.Data.Parser.ParserD.Type as ParserD
import qualified Streamly.Internal.Data.Parser.ParserK.Type as ParserK
import qualified Streamly.Internal.Data.Parser as Parser
import Prelude hiding (concatMap, take)
newtype Fold m a b = Fold (ParserD.Parser m (Array a) b)
{-# INLINE fromFold #-}
fromFold :: forall m a b. (MonadIO m, Storable a) =>
Fold.Fold m a b -> Fold m a b
fromFold :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Fold m a b -> Fold m a b
fromFold (Fold.Fold s -> a -> m (Step s b)
fstep m (Step s b)
finitial s -> m b
fextract) =
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold (forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m a b
ParserD.Parser s -> Array a -> m (Step s b)
step m (Initial s b)
initial s -> m b
fextract)
where
initial :: m (Initial s b)
initial = do
Step s b
res <- m (Step s b)
finitial
forall (m :: * -> *) a. Monad m => a -> m a
return
forall a b. (a -> b) -> a -> b
$ case Step s b
res of
Fold.Partial s
s1 -> forall s b. s -> Initial s b
IPartial s
s1
Fold.Done b
b -> forall s b. b -> Initial s b
IDone b
b
step :: s -> Array a -> m (Step s b)
step s
s (Array ArrayContents
contents Ptr a
start Ptr a
end) = do
SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC Ptr a
start s
s
where
goArray :: SPEC -> Ptr a -> s -> m (Step s b)
goArray !SPEC
_ !Ptr a
cur !s
fs | Ptr a
cur forall a. Ord a => a -> a -> Bool
>= Ptr a
end = do
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Ptr a
cur forall a. Eq a => a -> a -> Bool
== Ptr a
end) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> s -> Step s b
Partial Int
0 s
fs
goArray !SPEC
_ !Ptr a
cur !s
fs = do
a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
cur
Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
let elemSize :: Int
elemSize = SIZE_OF(a)
next :: Ptr b
next = PTR_NEXT(cur,a)
case Step s b
res of
Fold.Done b
b ->
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> b -> Step s b
Done ((Ptr a
end forall a b. Ptr a -> Ptr b -> Int
`minusPtr` forall {b}. Ptr b
next) forall a. Integral a => a -> a -> a
`div` Int
elemSize) b
b
Fold.Partial s
fs1 ->
SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC forall {b}. Ptr b
next s
fs1
{-# INLINE fromParserD #-}
fromParserD :: forall m a b. (MonadIO m, Storable a) =>
ParserD.Parser m a b -> Fold m a b
fromParserD :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Parser m a b -> Fold m a b
fromParserD (ParserD.Parser s -> a -> m (Step s b)
step1 m (Initial s b)
initial1 s -> m b
extract1) =
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold (forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m a b
ParserD.Parser s -> Array a -> m (Step s b)
step m (Initial s b)
initial1 s -> m b
extract1)
where
step :: s -> Array a -> m (Step s b)
step s
s (Array ArrayContents
contents Ptr a
start Ptr a
end) = do
if Ptr a
start forall a. Ord a => a -> a -> Bool
>= Ptr a
end
then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> s -> Step s b
Continue Int
0 s
s
else SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC Ptr a
start s
s
where
{-# INLINE partial #-}
partial :: Int
-> Ptr a
-> Ptr a
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Ptr a
cur Ptr a
next Int
elemSize Int -> s -> Step s b
st Int
n s
fs1 = do
let next1 :: Ptr b
next1 = Ptr a
next forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. Num a => a -> a
negate (Int
n forall a. Num a => a -> a -> a
* Int
elemSize)
if forall {b}. Ptr b
next1 forall a. Ord a => a -> a -> Bool
>= Ptr a
start Bool -> Bool -> Bool
&& Ptr a
cur forall a. Ord a => a -> a -> Bool
< Ptr a
end
then SPEC -> Ptr a -> s -> m (Step s b)
goArray SPEC
SPEC forall {b}. Ptr b
next1 s
fs1
else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Int -> s -> Step s b
st (Int
arrRem forall a. Num a => a -> a -> a
+ Int
n) s
fs1
goArray :: SPEC -> Ptr a -> s -> m (Step s b)
goArray !SPEC
_ !Ptr a
cur !s
fs = do
a
x <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
cur
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ArrayContents -> IO ()
touch ArrayContents
contents
Step s b
res <- s -> a -> m (Step s b)
step1 s
fs a
x
let elemSize :: Int
elemSize = SIZE_OF(a)
next :: Ptr b
next = PTR_NEXT(cur,a)
arrRem :: Int
arrRem = (Ptr a
end forall a b. Ptr a -> Ptr b -> Int
`minusPtr` forall {b}. Ptr b
next) forall a. Integral a => a -> a -> a
`div` Int
elemSize
case Step s b
res of
ParserD.Done Int
n b
b -> do
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> b -> Step s b
Done (Int
arrRem forall a. Num a => a -> a -> a
+ Int
n) b
b
ParserD.Partial Int
n s
fs1 ->
Int
-> Ptr a
-> Ptr a
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Ptr a
cur forall {b}. Ptr b
next Int
elemSize forall s b. Int -> s -> Step s b
Partial Int
n s
fs1
ParserD.Continue Int
n s
fs1 -> do
Int
-> Ptr a
-> Ptr a
-> Int
-> (Int -> s -> Step s b)
-> Int
-> s
-> m (Step s b)
partial Int
arrRem Ptr a
cur forall {b}. Ptr b
next Int
elemSize forall s b. Int -> s -> Step s b
Continue Int
n s
fs1
Error String
err -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. String -> Step s b
Error String
err
{-# INLINE fromParser #-}
fromParser :: forall m a b. (MonadThrow m, MonadIO m, Storable a) =>
Parser.Parser m a b -> Fold m a b
fromParser :: forall (m :: * -> *) a b.
(MonadThrow m, MonadIO m, Storable a) =>
Parser m a b -> Fold m a b
fromParser = forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Parser m a b -> Fold m a b
fromParserD forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
MonadThrow m =>
Parser m a b -> Parser m a b
ParserK.fromParserK
{-# INLINE fromArrayFold #-}
fromArrayFold :: forall m a b. (MonadIO m) =>
Fold.Fold m (Array a) b -> Fold m a b
fromArrayFold :: forall (m :: * -> *) a b.
MonadIO m =>
Fold m (Array a) b -> Fold m a b
fromArrayFold Fold m (Array a) b
f = forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b. Monad m => Fold m a b -> Parser m a b
ParserD.fromFold Fold m (Array a) b
f
instance Functor m => Functor (Fold m a) where
{-# INLINE fmap #-}
fmap :: forall a b. (a -> b) -> Fold m a a -> Fold m a b
fmap a -> b
f (Fold Parser m (Array a) a
p) = forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Parser m (Array a) a
p
{-# INLINE rmapM #-}
rmapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
rmapM :: forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Fold m a b -> Fold m a c
rmapM b -> m c
f (Fold Parser m (Array a) b
p) = forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Parser m a b -> Parser m a c
ParserD.rmapM b -> m c
f Parser m (Array a) b
p
{-# INLINE fromPure #-}
fromPure :: Monad m => b -> Fold m a b
fromPure :: forall (m :: * -> *) b a. Monad m => b -> Fold m a b
fromPure = forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) b a. Monad m => b -> Parser m a b
ParserD.fromPure
{-# INLINE fromEffect #-}
fromEffect :: Monad m => m b -> Fold m a b
fromEffect :: forall (m :: * -> *) b a. Monad m => m b -> Fold m a b
fromEffect = forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) b a. Monad m => m b -> Parser m a b
ParserD.fromEffect
{-# INLINE serial_ #-}
serial_ :: MonadThrow m => Fold m x a -> Fold m x b -> Fold m x b
serial_ :: forall (m :: * -> *) x a b.
MonadThrow m =>
Fold m x a -> Fold m x b -> Fold m x b
serial_ (Fold Parser m (Array x) a
p1) (Fold Parser m (Array x) b
p2) = forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) x a b.
MonadThrow m =>
Parser m x a -> Parser m x b -> Parser m x b
ParserD.noErrorUnsafeSplit_ Parser m (Array x) a
p1 Parser m (Array x) b
p2
{-# INLINE serialWith #-}
serialWith :: MonadThrow m
=> (a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
serialWith :: forall (m :: * -> *) a b c x.
MonadThrow m =>
(a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
serialWith a -> b -> c
f (Fold Parser m (Array x) a
p1) (Fold Parser m (Array x) b
p2) =
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Parser m x a -> Parser m x b -> Parser m x c
ParserD.noErrorUnsafeSplitWith a -> b -> c
f Parser m (Array x) a
p1 Parser m (Array x) b
p2
instance MonadThrow m => Applicative (Fold m a) where
{-# INLINE pure #-}
pure :: forall a. a -> Fold m a a
pure = forall (m :: * -> *) b a. Monad m => b -> Fold m a b
fromPure
{-# INLINE (<*>) #-}
<*> :: forall a b. Fold m a (a -> b) -> Fold m a a -> Fold m a b
(<*>) = forall (m :: * -> *) a b c x.
MonadThrow m =>
(a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
serialWith forall a. a -> a
id
{-# INLINE (*>) #-}
*> :: forall a b. Fold m a a -> Fold m a b -> Fold m a b
(*>) = forall (m :: * -> *) x a b.
MonadThrow m =>
Fold m x a -> Fold m x b -> Fold m x b
serial_
#if MIN_VERSION_base(4,10,0)
{-# INLINE liftA2 #-}
liftA2 :: forall a b c.
(a -> b -> c) -> Fold m a a -> Fold m a b -> Fold m a c
liftA2 a -> b -> c
f Fold m a a
x = forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b -> c
f Fold m a a
x)
#endif
{-# INLINE concatMap #-}
concatMap :: MonadThrow m =>
(b -> Fold m a c) -> Fold m a b -> Fold m a c
concatMap :: forall (m :: * -> *) b a c.
MonadThrow m =>
(b -> Fold m a c) -> Fold m a b -> Fold m a c
concatMap b -> Fold m a c
func (Fold Parser m (Array a) b
p) =
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b a c.
MonadThrow m =>
(b -> Parser m a c) -> Parser m a b -> Parser m a c
ParserD.noErrorUnsafeConcatMap (\b
x -> let Fold Parser m (Array a) c
y = b -> Fold m a c
func b
x in Parser m (Array a) c
y) Parser m (Array a) b
p
instance MonadThrow m => Monad (Fold m a) where
{-# INLINE return #-}
return :: forall a. a -> Fold m a a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure
{-# INLINE (>>=) #-}
>>= :: forall a b. Fold m a a -> (a -> Fold m a b) -> Fold m a b
(>>=) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) b a c.
MonadThrow m =>
(b -> Fold m a c) -> Fold m a b -> Fold m a c
concatMap
{-# INLINE (>>) #-}
>> :: forall a b. Fold m a a -> Fold m a b -> Fold m a b
(>>) = forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)
{-# INLINE take #-}
take :: forall m a b. (Monad m, Storable a) => Int -> Fold m a b -> Fold m a b
take :: forall (m :: * -> *) a b.
(Monad m, Storable a) =>
Int -> Fold m a b -> Fold m a b
take Int
n (Fold (ParserD.Parser s -> Array a -> m (Step s b)
step1 m (Initial s b)
initial1 s -> m b
extract1)) =
forall (m :: * -> *) a b. Parser m (Array a) b -> Fold m a b
Fold forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Initial s b) -> (s -> m b) -> Parser m a b
ParserD.Parser Tuple' Int s -> Array a -> m (Step (Tuple' Int s) b)
step m (Initial (Tuple' Int s) b)
initial forall {a}. Tuple' a s -> m b
extract
where
initial :: m (Initial (Tuple' Int s) b)
initial = do
Initial s b
res <- m (Initial s b)
initial1
case Initial s b
res of
IPartial s
s ->
if Int
n forall a. Ord a => a -> a -> Bool
> Int
0
then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Initial s b
IPartial forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> Tuple' a b
Tuple' Int
n s
s
else forall s b. b -> Initial s b
IDone forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
extract1 s
s
IDone b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. b -> Initial s b
IDone b
b
IError String
err -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. String -> Initial s b
IError String
err
{-# INLINE partial #-}
partial :: a -> (a -> Tuple' a s -> Step s b) -> a -> s -> m (Step s b)
partial a
i1 a -> Tuple' a s -> Step s b
st a
j s
s =
let i2 :: a
i2 = a
i1 forall a. Num a => a -> a -> a
+ a
j
in if a
i2 forall a. Ord a => a -> a -> Bool
> a
0
then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ a -> Tuple' a s -> Step s b
st a
j (forall a b. a -> b -> Tuple' a b
Tuple' a
i2 s
s)
else forall s b. Int -> b -> Step s b
Done Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
extract1 s
s
step :: Tuple' Int s -> Array a -> m (Step (Tuple' Int s) b)
step (Tuple' Int
i s
r) Array a
arr = do
let len :: Int
len = forall a. Storable a => Array a -> Int
Array.length Array a
arr
i1 :: Int
i1 = Int
i forall a. Num a => a -> a -> a
- Int
len
if Int
i1 forall a. Ord a => a -> a -> Bool
>= Int
0
then do
Step s b
res <- s -> Array a -> m (Step s b)
step1 s
r Array a
arr
case Step s b
res of
Partial Int
j s
s -> forall {a} {s}.
(Ord a, Num a) =>
a -> (a -> Tuple' a s -> Step s b) -> a -> s -> m (Step s b)
partial Int
i1 forall s b. Int -> s -> Step s b
Partial Int
j s
s
Continue Int
j s
s -> forall {a} {s}.
(Ord a, Num a) =>
a -> (a -> Tuple' a s -> Step s b) -> a -> s -> m (Step s b)
partial Int
i1 forall s b. Int -> s -> Step s b
Continue Int
j s
s
Done Int
j b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> b -> Step s b
Done Int
j b
b
Error String
err -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. String -> Step s b
Error String
err
else do
let !(Array ArrayContents
contents Ptr a
start Ptr a
_) = Array a
arr
end :: Ptr b
end = PTR_INDEX(start,i,a)
arr1 :: Array a
arr1 = forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array ArrayContents
contents Ptr a
start forall {b}. Ptr b
end
remaining :: Int
remaining = forall a. Num a => a -> a
negate Int
i1
Step s b
res <- s -> Array a -> m (Step s b)
step1 s
r Array a
arr1
case Step s b
res of
Partial Int
0 s
s -> forall s b. Int -> b -> Step s b
Done Int
remaining forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
extract1 s
s
Partial Int
j s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> s -> Step s b
Partial (Int
remaining forall a. Num a => a -> a -> a
+ Int
j) (forall a b. a -> b -> Tuple' a b
Tuple' Int
j s
s)
Continue Int
0 s
s -> forall s b. Int -> b -> Step s b
Done Int
remaining forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m b
extract1 s
s
Continue Int
j s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> s -> Step s b
Continue (Int
remaining forall a. Num a => a -> a -> a
+ Int
j) (forall a b. a -> b -> Tuple' a b
Tuple' Int
j s
s)
Done Int
j b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. Int -> b -> Step s b
Done (Int
remaining forall a. Num a => a -> a -> a
+ Int
j) b
b
Error String
err -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. String -> Step s b
Error String
err
extract :: Tuple' a s -> m b
extract (Tuple' a
_ s
r) = s -> m b
extract1 s
r