Commit 10e270e1 authored by Bodigrim's avatar Bodigrim
Browse files

Improve replicateChar

parent 30e1d175
......@@ -211,8 +211,8 @@ import Control.DeepSeq (NFData(rnf))
import Control.Exception (assert)
import GHC.Stack (HasCallStack)
#endif
import Data.Bits (shiftL, (.&.))
import Data.Char (isSpace)
import Data.Bits ((.&.))
import Data.Char (isSpace, isAscii, ord)
import Data.Data (Data(gfoldl, toConstr, gunfold, dataTypeOf), constrIndex,
Constr, mkConstr, DataType, mkDataType, Fixity(Prefix))
import Control.Monad (foldM)
......@@ -224,13 +224,14 @@ import Data.Binary (Binary(get, put))
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
import Data.String (IsString(..))
import Data.Text.Internal.Encoding.Utf8 (chr3, utf8LengthByLeader)
import Data.Text.Internal.Encoding.Utf8 (chr3, utf8Length, utf8LengthByLeader)
import qualified Data.Text.Internal.Fusion as S
import qualified Data.Text.Internal.Fusion.Common as S
import Data.Text.Encoding (decodeUtf8', encodeUtf8)
import Data.Text.Internal.Fusion (stream, reverseStream, unstream)
import Data.Text.Internal.Private (span_)
import Data.Text.Internal (Text(..), empty, firstf, mul, safe, text)
import Data.Text.Internal.Unsafe.Char (unsafeWrite)
import Data.Text.Show (singleton, unpack, unpackCString#)
import qualified Prelude as P
import Data.Text.Unsafe (Iter(..), iter, iter_, lengthWord8, reverseIter,
......@@ -1053,21 +1054,15 @@ replicate n t@(Text a o l)
| n <= 0 || l <= 0 = empty
| n == 1 = t
| isSingleton t = replicateChar n (unsafeHead t)
| otherwise = Text (A.run x) 0 len
where
len = l `mul` n -- TODO: detect overflows
x :: ST s (A.MArray s)
x = do
arr <- A.new len
A.copyI l arr 0 a o
let loop !l1 =
let rest = len - l1 in
if rest <= l1 then A.copyM arr l1 arr 0 rest >> return arr
else A.copyM arr l1 arr 0 l1 >> loop (l1 `shiftL` 1)
loop l
| otherwise = runST $ do
let totalLen = n `mul` l
marr <- A.new totalLen
A.copyI l marr 0 a o
A.tile marr l
arr <- A.unsafeFreeze marr
return $ Text arr 0 totalLen
{-# INLINE [1] replicate #-}
{-# RULES
"TEXT replicate/singleton -> replicateChar" [~1] forall n c.
replicate n (singleton c) = replicateChar n c
......@@ -1076,7 +1071,22 @@ replicate n t@(Text a o l)
-- | /O(n)/ 'replicateChar' @n@ @c@ is a 'Text' of length @n@ with @c@ the
-- value of every element.
replicateChar :: Int -> Char -> Text
replicateChar n c = unstream (S.replicateCharI n (safe c))
replicateChar !len !c'
| len <= 0 = empty
| isAscii c = runST $ do
marr <- A.newFilled len (ord c)
arr <- A.unsafeFreeze marr
return $ Text arr 0 len
| otherwise = runST $ do
let cLen = utf8Length c
totalLen = cLen P.* len
marr <- A.new totalLen
_ <- unsafeWrite marr 0 c
A.tile marr cLen
arr <- A.unsafeFreeze marr
return $ Text arr 0 totalLen
where
c = safe c'
{-# INLINE replicateChar #-}
-- | /O(n)/, where @n@ is the length of the result. The 'unfoldr'
......
......@@ -40,7 +40,9 @@ module Data.Text.Array
, unsafeIndex
, new
, newPinned
, newFilled
, unsafeWrite
, tile
) where
#if defined(ASSERTS)
......@@ -83,6 +85,22 @@ newPinned (I# len#)
(# s2#, marr# #) -> (# s2#, MutableByteArray marr# #)
{-# INLINE newPinned #-}
newFilled :: Int -> Int -> ST s (MArray s)
newFilled (I# len#) (I# c#) = ST $ \s1# ->
case newByteArray# len# s1# of
(# s2#, marr# #) -> case setByteArray# marr# 0# len# c# s2# of
s3# -> (# s3#, MutableByteArray marr# #)
{-# INLINE newFilled #-}
tile :: MArray s -> Int -> ST s ()
tile marr tileLen = do
totalLen <- getSizeofMArray marr
let go l
| 2 * l > totalLen = copyM marr l marr 0 (totalLen - l)
| otherwise = copyM marr l marr 0 l >> go (2 * l)
go tileLen
{-# INLINE tile #-}
-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
unsafeFreeze :: MArray s -> ST s Array
unsafeFreeze (MutableByteArray marr) = ST $ \s1# ->
......@@ -107,7 +125,6 @@ unsafeIndex (ByteArray arr) i@(I# i#) =
case indexWord8Array# arr i# of r# -> (W8# r#)
{-# INLINE unsafeIndex #-}
#if defined(ASSERTS)
-- sizeofMutableByteArray# is deprecated, because it is unsafe in the presence of
-- shrinkMutableByteArray# and resizeMutableByteArray#.
getSizeofMArray :: MArray s -> ST s Int
......@@ -115,6 +132,7 @@ getSizeofMArray (MutableByteArray marr) = ST $ \s0# ->
case getSizeofMutableByteArray# marr s0# of
(# s1#, word8len# #) -> (# s1#, I# word8len# #)
#if defined(ASSERTS)
checkBoundsM :: HasCallStack => MArray s -> Int -> Int -> ST s ()
checkBoundsM ma i elSize = do
len <- getSizeofMArray ma
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment