{-# LANGUAGE BangPatterns, CPP #-}
-- |
-- Module      : Data.Text.Lazy.Fusion
-- Copyright   : (c) 2009, 2010 Bryan O'Sullivan
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : GHC
--
-- /Warning/: this is an internal module, and does not have a stable
-- API or name. Functions in this module may not check or enforce
-- preconditions expected by public modules. Use at your own risk!
--
-- Core stream fusion functionality for text.

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)

-- | /O(n)/ Convert a 'Text' into a 'Stream Char'.
stream ::
#if defined(ASSERTS)
  HasCallStack =>
#endif
  Text -> Stream Char
stream :: Text -> Stream Char
stream Text
text = (PairS Text Int -> Step (PairS Text Int) Char)
-> PairS Text Int -> Size -> Stream Char
forall a s. (s -> Step s a) -> s -> Size -> Stream a
Stream PairS Text Int -> Step (PairS Text Int) Char
next (Text
text Text -> Int -> PairS Text Int
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
_) = Step (PairS Text Int) Char
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len  = PairS Text Int -> Step (PairS Text Int) Char
next (Text
ts Text -> Int -> PairS Text Int
forall a b. a -> b -> PairS a b
:*: Int
0)
        | Bool
otherwise = Char -> PairS Text Int -> Step (PairS Text Int) Char
forall s a. a -> s -> Step s a
Yield Char
c (Text
txt Text -> Int -> PairS Text Int
forall a b. a -> b -> PairS a b
:*: Int
iInt -> Int -> Int
forall 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 #-}

-- | /O(n)/ Convert a 'Stream Char' into a 'Text', using the given
-- chunk size.
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 s. ST s Text) -> Text
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Text) -> Text) -> (forall s. ST s Text) -> Text
forall a b. (a -> b) -> a -> b
$ do
                                a <- Int -> ST s (MArray s)
forall s. Int -> ST s (MArray s)
A.new Int
unknownLength
                                unsafeWrite a 0 x >>= inner a unknownLength 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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
chunkSize = MArray s -> Int -> s -> ST s Text
forall {s}. MArray s -> Int -> s -> ST s Text
finish MArray s
marr Int
i s
s
            | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len       = {-# SCC "unstreamChunks/resize" #-} do
                let newLen :: Int
newLen = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
len Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Int
chunkSize
                marr' <- Int -> ST s (MArray s)
forall s. Int -> ST s (MArray s)
A.new Int
newLen
                A.copyM marr' 0 marr 0 len
                inner marr' newLen s i
            | Bool
otherwise =
                {-# SCC "unstreamChunks/inner" #-}
                case s -> Step s Char
next s
s of
                  Step s Char
Done        -> MArray s -> Int -> s -> ST s Text
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 d <- MArray s -> Int -> Char -> ST s Int
forall s. MArray s -> Int -> Char -> ST s Int
unsafeWrite MArray s
marr Int
i Char
x
                                    inner marr len s' (i+d)
        finish :: MArray s -> Int -> s -> ST s Text
finish MArray s
marr Int
len s
s' = do
          MArray s -> Int -> ST s ()
forall s. MArray s -> Int -> ST s ()
A.shrinkM MArray s
marr Int
len
          arr <- MArray s -> ST s Array
forall s. MArray s -> ST s Array
A.unsafeFreeze MArray s
marr
          return (I.Text arr 0 len `Chunk` outer s')
{-# INLINE [0] unstreamChunks #-}

-- | /O(n)/ Convert a 'Stream Char' into a 'Text', using
-- 'defaultChunkSize'.
unstream ::
#if defined(ASSERTS)
  HasCallStack =>
#endif
  Stream Char -> Text
unstream :: Stream Char -> Text
unstream = Int -> Stream Char -> Text
unstreamChunks Int
defaultChunkSize
{-# INLINE [0] unstream #-}

-- | /O(n)/ Returns the number of characters in a text.
length :: Stream Char -> Int64
length :: Stream Char -> Int64
length = Stream Char -> Int64
forall a. Integral a => Stream Char -> a
S.lengthI
{-# INLINE[0] length #-}

{-# RULES "LAZY STREAM stream/unstream fusion" forall s.
    stream (unstream s) = s #-}

-- | /O(n)/ Like 'unfoldr', 'unfoldrN' builds a stream from a seed
-- value. However, the length of the result is limited by the
-- first argument to 'unfoldrN'. This function is more efficient than
-- 'unfoldr' when the length of the result is known.
unfoldrN :: Int64 -> (a -> Maybe (Char,a)) -> a -> Stream Char
unfoldrN :: forall a. Int64 -> (a -> Maybe (Char, a)) -> a -> Stream Char
unfoldrN Int64
n = Int64 -> (a -> Maybe (Char, a)) -> a -> Stream Char
forall a b.
Integral a =>
a -> (b -> Maybe (Char, b)) -> b -> Stream Char
S.unfoldrNI Int64
n
{-# INLINE [0] unfoldrN #-}

-- | /O(n)/ stream index (subscript) operator, starting from 0.
index :: HasCallStack => Stream Char -> Int64 -> Char
index :: HasCallStack => Stream Char -> Int64 -> Char
index = Stream Char -> Int64 -> Char
forall a. (HasCallStack, Integral a) => Stream Char -> a -> Char
S.indexI
{-# INLINE [0] index #-}

-- | /O(n)/ The 'count' function returns the number of times the query
-- element appears in the given stream.
countChar :: Char -> Stream Char -> Int64
countChar :: Char -> Stream Char -> Int64
countChar = Char -> Stream Char -> Int64
forall a. Integral a => Char -> Stream Char -> a
S.countCharI
{-# INLINE [0] countChar #-}