Commits

Bryan O'Sullivan committed 0b09e9b

Implement and test lazy take

  • Participants
  • Parent commits 6b47c79

Comments (0)

Files changed (3)

File Data/Text/Fusion/Common.hs

 -- | /O(n)/ take n, applied to a stream, returns the prefix of the
 -- stream of length @n@, or the stream itself if @n@ is greater than the
 -- length of the stream.
-take :: Int -> Stream Char -> Stream Char
+take :: Integral a => a -> Stream Char -> Stream Char
 take n0 (Stream next0 s0 len) = Stream next (n0 :!: s0) len
     where
       {-# INLINE next #-}

File Data/Text/Lazy.hs

     -- * Substrings
 
     -- ** Breaking strings
-    -- , take
+    , take
     -- , drop
     -- , takeWhile
     -- , dropWhile
 unfoldrN n f s = unstream (S.unfoldrN n f s)
 {-# INLINE unfoldrN #-}
 
+-- | /O(n)/ 'take' @n@, applied to a 'Text', returns the prefix of the
+-- 'Text' of length @n@, or the 'Text' itself if @n@ is greater than
+-- the length of the Text. Subject to fusion.
+take :: Int64 -> Text -> Text
+take i _ | i <= 0 = Empty
+take i t0         = take' i t0
+  where take' 0 _            = Empty
+        take' _ Empty        = Empty
+        take' n (Chunk t ts)
+            | n < len   = Chunk (T.take (fromIntegral n) t) Empty
+            | otherwise = Chunk t (take' (n - len) ts)
+            where len = fromIntegral (T.length t)
+{-# INLINE [1] take #-}
+
+{-# RULES
+"LAZY TEXT take -> fused" [~1] forall n t.
+    take n t = unstream (S.take n (stream t))
+"LAZY TEXT take -> unfused" [1] forall n t.
+    unstream (S.take n (stream t)) = take n 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)@.

File tests/Properties.hs

 unpack2 :: (Target t) => (t,t) -> (String,String)
 unpack2 = unpackT *** unpackT
 
+prop_S_take n          = L.take n      `eqP` (unpackT . S.take n)
 prop_T_take n          = L.take n      `eqP` (unpackT . T.take n)
+prop_TL_take n         = L.take n      `eqP` (unpackT . TL.take (fromIntegral n))
 prop_T_drop n          = L.drop n      `eqP` (unpackT . T.drop n)
 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_T_unfoldrN", mytest prop_T_unfoldrN),
   ("prop_TL_unfoldrN", mytest prop_TL_unfoldrN),
 
+  ("prop_S_take", mytest prop_S_take),
   ("prop_T_take", mytest prop_T_take),
+  ("prop_TL_take", mytest prop_TL_take),
   ("prop_T_drop", mytest prop_T_drop),
   ("prop_T_takeWhile", mytest prop_T_takeWhile),
   ("prop_T_takeWhileS", mytest prop_T_takeWhileS),