Improve replicateChar

src/Data/Text.hs

@@ -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'

src/Data/Text/Array.hs

@@ -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