Commits

Bryan O'Sullivan committed c4397da

Rewrite mapAccumL to be fast

  • Participants
  • Parent commits 65fbb95

Comments (0)

Files changed (3)

File Data/Text.hs

            | otherwise = scanr f (last t) (init t)
 {-# INLINE scanr1 #-}
 
--- | /O(n)/ Like a combination of 'map' and 'foldl'. Applies a
+-- | /O(n)/ Like a combination of 'map' and 'foldl''. Applies a
 -- function to each element of a 'Text', passing an accumulating
 -- parameter from left to right, and returns a final 'Text'.
 mapAccumL :: (a -> Char -> (a,Char)) -> a -> Text -> (a, Text)
-mapAccumL f = go
-  where go s t = case uncons t of
-                   Nothing -> (s, empty)
-                   Just (x, xs) -> (s'', cons y ys)
-                       where (s', y ) = f s x
-                             (s'',ys) = go s' xs
+mapAccumL f z0 = S.mapAccumL f z0 . stream
 {-# INLINE mapAccumL #-}
 
 -- | The 'mapAccumR' function behaves like a combination of 'map' and

File Data/Text/Fusion.hs

     -- ** Scans
     , reverseScanr
 
+    -- ** Accumulating maps
+    , mapAccumL
+
     -- ** Generation and unfolding
     , unfoldrN
 
 countChar :: Char -> Stream Char -> Int
 countChar = S.countCharI
 {-# INLINE [0] countChar #-}
+
+-- | /O(n)/ Like a combination of 'map' and 'foldl''. Applies a
+-- function to each element of a 'Text', passing an accumulating
+-- parameter from left to right, and returns a final 'Text'.
+mapAccumL :: (a -> Char -> (a,Char)) -> a -> Stream Char -> (a, Text)
+mapAccumL f z0 (Stream next0 s0 len) = (nz,I.textP na 0 nl)
+  where
+    (na,(nz,nl)) = A.run2 (A.unsafeNew mlen >>= \arr -> outer arr mlen z0 s0 0)
+      where mlen = upperBound 4 len
+    outer arr top = loop
+      where
+        loop !z !s !i =
+            case next0 s of
+              Done          -> return (arr, (z,i))
+              Skip s'       -> loop z s' i
+              Yield x s'
+                | j >= top  -> {-# SCC "mapAccumL/resize" #-} do
+                               let top' = (top + 1) `shiftL` 1
+                               arr' <- A.unsafeNew top'
+                               A.copyM arr' 0 arr 0 top
+                               outer arr' top' z s i
+                | otherwise -> do let (z',c) = f z x
+                                  d <- unsafeWrite arr i c
+                                  loop z' s' (i+d)
+                where j | ord x < 0x10000 = i
+                        | otherwise       = i + 1
+{-# INLINE [0] mapAccumL #-}

File tests/benchmarks/StripBrackets.hs

+-- From Petr Prokhorenkov.
+
+import Data.Text as T
+import Data.Text.IO as T
+
+stripBrackets :: T.Text -> T.Text
+stripBrackets text = snd $ T.mapAccumL f 0 text where
+   f depth c = let
+       depth' = depth + d' c
+       c' | depth > 0 || depth' > 0 = ' '
+          | otherwise = c
+       in
+       (depth', c')
+
+   d' '{' = 1
+   d' '[' = 1
+   d' '}' = -1
+   d' ']' = -1
+   d' _   = 0
+
+main = T.interact stripBrackets