scanl/scanr should replace invalid Char in the initial value

Author: Bodigrim

src/Data/Text.hs

@@ -254,7 +254,7 @@ import Data.Text.Internal.Encoding.Utf8 (utf8Length, utf8LengthByLeader, chr3, o
 import qualified Data.Text.Internal.Fusion as S
 import qualified Data.Text.Internal.Fusion.Common as S
 import Data.Text.Encoding (decodeUtf8', encodeUtf8Builder)
-import Data.Text.Internal.Fusion (stream, reverseStream, unstream)
+import Data.Text.Internal.Fusion (stream, unstream)
 import Data.Text.Internal.Private (span_)
 import Data.Text.Internal (Text(..), StrictText, empty, firstf, mul, safe, text, append, pack)
 import Data.Text.Internal.Unsafe.Char (unsafeWrite)
@@ -1210,8 +1210,32 @@ minimum = foldl1' min
 --
 -- @'last' ('scanl' f z xs) = 'foldl' f z xs@
 scanl :: (Char -> Char -> Char) -> Char -> Text -> Text
-scanl f z t = unstream (S.scanl g z (stream t))
-    where g a b = safe (f a b)
+scanl f c0 = go
+  where
+    go (Text src o l) = runST $ do
+      let l' = l + 4
+          c0' = safe c0
+      marr <- A.new l'
+      d' <- unsafeWrite marr 0 c0'
+      outer marr l' o d' c0'
+      where
+        outer :: forall s. A.MArray s -> Int -> Int -> Int -> Char -> ST s Text
+        outer !dst !dstLen = inner
+          where
+            inner !srcOff !dstOff !c
+              | srcOff >= l + o = do
+                A.shrinkM dst dstOff
+                arr <- A.unsafeFreeze dst
+                pure $ Text arr 0 dstOff
+              | dstOff + 4 > dstLen = do
+                let !dstLen' = dstLen + (l + o) - srcOff + 4
+                dst' <- A.resizeM dst dstLen'
+                outer dst' dstLen' srcOff dstOff c
+              | otherwise = do
+                let !(Iter c' d) = iterArray src srcOff
+                    c'' = safe $ f c c'
+                d' <- unsafeWrite dst dstOff c''
+                inner (srcOff + d) (dstOff + d') c''
 {-# INLINE scanl #-}
 
 -- | /O(n)/ 'scanl1' is a variant of 'scanl' that has no starting
@@ -1228,8 +1252,37 @@ scanl1 f t | null t    = empty
 --
 -- > scanr f v == reverse . scanl (flip f) v . reverse
 scanr :: (Char -> Char -> Char) -> Char -> Text -> Text
-scanr f z = S.reverse . S.reverseScanr g z . reverseStream
-    where g a b = safe (f a b)
+scanr f c0 = go
+  where
+    go (Text src o l) = runST $ do
+      let l' = l + 4
+          c0' = safe c0
+          !d' = utf8Length c0'
+      marr <- A.new l'
+      _ <- unsafeWrite marr (l' - d') c0'
+      outer marr (l + o - 1) (l' - d' - 1) c0'
+      where
+        outer :: forall s. A.MArray s -> Int -> Int -> Char -> ST s Text
+        outer !dst = inner
+          where
+            inner !srcOff !dstOff !c
+              | srcOff < o = do
+                dstLen <- A.getSizeofMArray dst
+                arr <- A.unsafeFreeze dst
+                pure $ Text arr (dstOff + 1) (dstLen - dstOff - 1)
+              | dstOff < 3 = do
+                dstLen <- A.getSizeofMArray dst
+                let !dstLen' = dstLen + (srcOff - o) + 4
+                dst' <- A.new dstLen'
+                A.copyM dst' (dstLen' - dstLen) dst 0 dstLen
+                outer dst' srcOff (dstOff + dstLen' - dstLen) c
+              | otherwise = do
+                let !(Iter c' d) = reverseIterArray src srcOff
+                    c'' = safe $ f c' c
+                    !d' = utf8Length c''
+                    dstOff' = dstOff - d'
+                _ <- unsafeWrite dst (dstOff' + 1) c''
+                inner (srcOff + d) dstOff' c''
 {-# INLINE scanr #-}
 
 -- | /O(n)/ 'scanr1' is a variant of 'scanr' that has no starting

src/Data/Text/Lazy.hs

@@ -972,8 +972,12 @@ isAscii = foldrChunks (\chnk acc -> T.isAscii chnk && acc) True
 --
 -- > last (scanl f z xs) == foldl f z xs.
 scanl :: (Char -> Char -> Char) -> Char -> Text -> Text
-scanl f z t = unstream (S.scanl g z (stream t))
-    where g a b = safe (f a b)
+scanl f z t = cons z $ P.snd $
+  mapAccumL (\acc c -> let c' = f acc c in (c', c')) (safe z) t
+-- This is a bit suboptimal: we could have used
+-- Data.Text.scanl for the first chunk and mapAccumL
+-- for subsequent ones, but but I doubt anyone cares
+-- about the performance of 'scanl' much.
 {-# INLINE scanl #-}
 
 -- | /O(n)/ 'scanl1' is a variant of 'scanl' that has no starting
@@ -991,8 +995,10 @@ scanl1 f t0 = case uncons t0 of
 --
 -- > scanr f v == reverse . scanl (flip f) v . reverse
 scanr :: (Char -> Char -> Char) -> Char -> Text -> Text
-scanr f v = reverse . scanl g v . reverse
-    where g a b = safe (f b a)
+scanr f z t = (`snoc` z) $ P.snd $
+  mapAccumR (\acc c -> let c' = f c acc in (c', c')) (safe z) t
+-- See the comment for 'scanl' above.
+{-# INLINE scanr #-}
 
 -- | /O(n)/ 'scanr1' is a variant of 'scanr' that has no starting
 -- value argument.  Performs replacement on invalid scalar values.

tests/Tests/Properties/Folds.hs

@@ -140,6 +140,15 @@ tl_scanr (applyFun2 -> f) z      = L.scanr f z   `eqP` (unpackS . TL.scanr f z)
 t_scanr1 (applyFun2 -> f)        = L.scanr1 f    `eqP` (unpackS . T.scanr1 f)
 tl_scanr1 (applyFun2 -> f)       = L.scanr1 f    `eqP` (unpackS . TL.scanr1 f)
 
+t_scanl_is_safe = let c = '\55296' in
+    T.scanl undefined c mempty === T.singleton c
+tl_scanl_is_safe = let c = '\55296' in
+    TL.scanl undefined c mempty === TL.singleton c
+t_scanr_is_safe = let c = '\55296' in
+    T.scanr undefined c mempty === T.singleton c
+tl_scanr_is_safe = let c = '\55296' in
+    TL.scanr undefined c mempty === TL.singleton c
+
 t_mapAccumL_char c t =
     snd (T.mapAccumL (const (const (0 :: Int, c))) 0 t) === T.replicate (T.length t) (T.singleton c)
 t_mapAccumL (applyFun2 -> f) z   = L.mapAccumL f z `eqP` (second unpackS . T.mapAccumL f z)
@@ -303,7 +312,12 @@ testFolds =
       testProperty "t_scanr" t_scanr,
       testProperty "tl_scanr" tl_scanr,
       testProperty "t_scanr1" t_scanr1,
-      testProperty "tl_scanr1" tl_scanr1
+      testProperty "tl_scanr1" tl_scanr1,
+
+      testProperty "t_scanl_is_safe" t_scanl_is_safe,
+      testProperty "tl_scanl_is_safe" tl_scanl_is_safe,
+      testProperty "t_scanr_is_safe" t_scanr_is_safe,
+      testProperty "tl_scanr_is_safe" tl_scanr_is_safe
     ],
 
     testGroup "mapAccum" [