{-# LANGUAGE BangPatterns, CPP #-}
module Data.Text.Internal.Lazy.Fusion
(
stream
, unstream
, unstreamChunks
, length
, unfoldrN
, index
, countChar
) where
import Prelude hiding (length)
import Data.Bits (shiftL)
import qualified Data.Text.Internal.Fusion.Common as S
import Control.Monad.ST (runST)
import Data.Text.Internal.Fusion.Types
import Data.Text.Internal.Fusion.Size (isEmpty, unknownSize)
import Data.Text.Internal.Lazy
import qualified Data.Text.Internal as I
import qualified Data.Text.Array as A
import Data.Text.Internal.Unsafe.Char (unsafeWrite)
import Data.Text.Unsafe (Iter(..), iter)
import Data.Int (Int64)
import GHC.Stack (HasCallStack)
default(Int64)
stream ::
#if defined(ASSERTS)
HasCallStack =>
#endif
Text -> Stream Char
stream :: Text -> Stream Char
stream Text
text = forall a s. (s -> Step s a) -> s -> Size -> Stream a
Stream PairS Text Int -> Step (PairS Text Int) Char
next (Text
text forall a b. a -> b -> PairS a b
:*: Int
0) Size
unknownSize
where
next :: PairS Text Int -> Step (PairS Text Int) Char
next (Text
Empty :*: Int
_) = forall s a. Step s a
Done
next (txt :: Text
txt@(Chunk t :: Text
t@(I.Text Array
_ Int
_ Int
len) Text
ts) :*: Int
i)
| Int
i forall a. Ord a => a -> a -> Bool
>= Int
len = PairS Text Int -> Step (PairS Text Int) Char
next (Text
ts forall a b. a -> b -> PairS a b
:*: Int
0)
| Bool
otherwise = forall s a. a -> s -> Step s a
Yield Char
c (Text
txt forall a b. a -> b -> PairS a b
:*: Int
iforall a. Num a => a -> a -> a
+Int
d)
where Iter Char
c Int
d = Text -> Int -> Iter
iter Text
t Int
i
{-# INLINE [0] stream #-}
unstreamChunks ::
#if defined(ASSERTS)
HasCallStack =>
#endif
Int -> Stream Char -> Text
unstreamChunks :: Int -> Stream Char -> Text
unstreamChunks !Int
chunkSize (Stream s -> Step s Char
next s
s0 Size
len0)
| Size -> Bool
isEmpty Size
len0 = Text
Empty
| Bool
otherwise = s -> Text
outer s
s0
where
outer :: s -> Text
outer s
so = {-# SCC "unstreamChunks/outer" #-}
case s -> Step s Char
next s
so of
Step s Char
Done -> Text
Empty
Skip s
s' -> s -> Text
outer s
s'
Yield Char
x s
s' -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
MArray s
a <- forall s. Int -> ST s (MArray s)
A.new Int
unknownLength
forall s. MArray s -> Int -> Char -> ST s Int
unsafeWrite MArray s
a Int
0 Char
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall {s}. MArray s -> Int -> s -> Int -> ST s Text
inner MArray s
a Int
unknownLength s
s'
where unknownLength :: Int
unknownLength = Int
4
where
inner :: MArray s -> Int -> s -> Int -> ST s Text
inner MArray s
marr !Int
len s
s !Int
i
| Int
i forall a. Num a => a -> a -> a
+ Int
3 forall a. Ord a => a -> a -> Bool
>= Int
chunkSize = forall {s}. MArray s -> Int -> s -> ST s Text
finish MArray s
marr Int
i s
s
| Int
i forall a. Num a => a -> a -> a
+ Int
3 forall a. Ord a => a -> a -> Bool
>= Int
len = {-# SCC "unstreamChunks/resize" #-} do
let newLen :: Int
newLen = forall a. Ord a => a -> a -> a
min (Int
len forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Int
chunkSize
MArray s
marr' <- forall s. Int -> ST s (MArray s)
A.new Int
newLen
forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
A.copyM MArray s
marr' Int
0 MArray s
marr Int
0 Int
len
MArray s -> Int -> s -> Int -> ST s Text
inner MArray s
marr' Int
newLen s
s Int
i
| Bool
otherwise =
{-# SCC "unstreamChunks/inner" #-}
case s -> Step s Char
next s
s of
Step s Char
Done -> forall {s}. MArray s -> Int -> s -> ST s Text
finish MArray s
marr Int
i s
s
Skip s
s' -> MArray s -> Int -> s -> Int -> ST s Text
inner MArray s
marr Int
len s
s' Int
i
Yield Char
x s
s' -> do Int
d <- forall s. MArray s -> Int -> Char -> ST s Int
unsafeWrite MArray s
marr Int
i Char
x
MArray s -> Int -> s -> Int -> ST s Text
inner MArray s
marr Int
len s
s' (Int
iforall a. Num a => a -> a -> a
+Int
d)
finish :: MArray s -> Int -> s -> ST s Text
finish MArray s
marr Int
len s
s' = do
forall s. MArray s -> Int -> ST s ()
A.shrinkM MArray s
marr Int
len
Array
arr <- forall s. MArray s -> ST s Array
A.unsafeFreeze MArray s
marr
forall (m :: * -> *) a. Monad m => a -> m a
return (Array -> Int -> Int -> Text
I.Text Array
arr Int
0 Int
len Text -> Text -> Text
`Chunk` s -> Text
outer s
s')
{-# INLINE [0] unstreamChunks #-}
unstream ::
#if defined(ASSERTS)
HasCallStack =>
#endif
Stream Char -> Text
unstream :: Stream Char -> Text
unstream = Int -> Stream Char -> Text
unstreamChunks Int
defaultChunkSize
{-# INLINE [0] unstream #-}
length :: Stream Char -> Int64
length :: Stream Char -> Int64
length = forall a. Integral a => Stream Char -> a
S.lengthI
{-# INLINE[0] length #-}
{-# RULES "LAZY STREAM stream/unstream fusion" forall s.
stream (unstream s) = s #-}
unfoldrN :: Int64 -> (a -> Maybe (Char,a)) -> a -> Stream Char
unfoldrN :: forall a. Int64 -> (a -> Maybe (Char, a)) -> a -> Stream Char
unfoldrN Int64
n = forall a b.
Integral a =>
a -> (b -> Maybe (Char, b)) -> b -> Stream Char
S.unfoldrNI Int64
n
{-# INLINE [0] unfoldrN #-}
index :: HasCallStack => Stream Char -> Int64 -> Char
index :: HasCallStack => Stream Char -> Int64 -> Char
index = forall a. (HasCallStack, Integral a) => Stream Char -> a -> Char
S.indexI
{-# INLINE [0] index #-}
countChar :: Char -> Stream Char -> Int64
countChar :: Char -> Stream Char -> Int64
countChar = forall a. Integral a => Char -> Stream Char -> a
S.countCharI
{-# INLINE [0] countChar #-}