Commits

Bryan O'Sullivan committed d46f575

Implement and test lazy takeWhile and dropWhile

Comments (0)

Files changed (2)

Data/Text/Lazy.hs

     -- ** Breaking strings
     , take
     , drop
-    -- , takeWhile
-    -- , dropWhile
+    , takeWhile
+    , dropWhile
     , splitAt
     -- , span
     -- , break
     unstream (S.drop n (stream t)) = drop n t
   #-}
 
+-- | /O(n)/ 'takeWhile', applied to a predicate @p@ and a 'Text', returns
+-- the longest prefix (possibly empty) of elements that satisfy @p@.
+-- This function is subject to array fusion.
+takeWhile :: (Char -> Bool) -> Text -> Text
+takeWhile p t0 = takeWhile' t0
+  where takeWhile' Empty        = Empty
+        takeWhile' (Chunk t ts) =
+          case T.findIndex (not . p) t of
+            Just n | n > 0     -> Chunk (T.take n t) Empty
+                   | otherwise -> Empty
+            Nothing            -> Chunk t (takeWhile' ts)
+{-# INLINE [1] takeWhile #-}
+
+{-# RULES
+"LAZY TEXT takeWhile -> fused" [~1] forall p t.
+    takeWhile p t = unstream (S.takeWhile p (stream t))
+"LAZY TEXT takeWhile -> unfused" [1] forall p t.
+    unstream (S.takeWhile p (stream t)) = takeWhile p t
+  #-}
+
+-- | /O(n)/ 'dropWhile' @p@ @xs@ returns the suffix remaining after
+-- 'takeWhile' @p@ @xs@. This function is subject to array fusion.
+dropWhile :: (Char -> Bool) -> Text -> Text
+dropWhile p t0 = dropWhile' t0
+  where dropWhile' Empty        = Empty
+        dropWhile' (Chunk t ts) =
+          case T.findIndex (not . p) t of
+            Just n  -> Chunk (T.drop n t) ts
+            Nothing -> dropWhile' ts
+{-# INLINE [1] dropWhile #-}
+
+{-# RULES
+"LAZY TEXT dropWhile -> fused" [~1] forall p t.
+    dropWhile p t = unstream (S.dropWhile p (stream t))
+"LAZY TEXT dropWhile -> unfused" [1] forall p t.
+    unstream (S.dropWhile p (stream t)) = dropWhile p t
+  #-}
+
 -- | /O(n)/ 'splitAt' @n t@ returns a pair whose first element is a
 -- prefix of @t@ of length @n@, and whose second is the remainder of
 -- the string. It is equivalent to @('take' n t, 'drop' n t)@.

tests/Properties.hs

 prop_S_drop n          = L.drop n      `eqP` (unpackT . S.drop n)
 prop_T_drop n          = L.drop n      `eqP` (unpackT . T.drop n)
 prop_TL_drop n         = L.drop n      `eqP` (unpackT . TL.drop n)
+prop_S_takeWhile p     = L.takeWhile p `eqP` (unpackT . S.takeWhile p)
 prop_T_takeWhile p     = L.takeWhile p `eqP` (unpackT . T.takeWhile p)
-prop_T_takeWhileS p    = L.takeWhile p `eqP` (unpackT . S.unstream . S.takeWhile p . S.stream)
+prop_TL_takeWhile p    = L.takeWhile p `eqP` (unpackT . TL.takeWhile p)
+prop_S_dropWhile p     = L.dropWhile p `eqP` (unpackT . S.dropWhile p)
 prop_T_dropWhile p     = L.dropWhile p `eqP` (unpackT . T.dropWhile p)
-prop_T_dropWhileS p    = L.dropWhile p `eqP` (unpackT . S.unstream . S.dropWhile p . S.stream)
+prop_TL_dropWhile p    = L.dropWhile p `eqP` (unpackT . S.dropWhile p)
 prop_T_splitAt n       = L.splitAt n   `eqP` (unpack2 . T.splitAt n)
 prop_TL_splitAt n      = L.splitAt n   `eqP` (unpack2 . TL.splitAt n)
 prop_T_span p          = L.span p      `eqP` (unpack2 . T.span p)
   ("prop_S_drop", mytest prop_S_drop),
   ("prop_T_drop", mytest prop_T_drop),
   ("prop_TL_drop", mytest prop_TL_drop),
+  ("prop_S_takeWhile", mytest prop_S_takeWhile),
   ("prop_T_takeWhile", mytest prop_T_takeWhile),
-  ("prop_T_takeWhileS", mytest prop_T_takeWhileS),
+  ("prop_TL_takeWhile", mytest prop_TL_takeWhile),
+  ("prop_S_dropWhile", mytest prop_S_dropWhile),
   ("prop_T_dropWhile", mytest prop_T_dropWhile),
-  ("prop_T_dropWhileS", mytest prop_T_dropWhileS),
+  ("prop_TL_dropWhile", mytest prop_TL_dropWhile),
   ("prop_T_splitAt", mytest prop_T_splitAt),
   ("prop_T_span", mytest prop_T_span),
   ("prop_T_break", mytest prop_T_break),