Feature: initsNE and tailsNE. (#558)

* Write test cases and stub functions.

* fixup! Write test cases and stub functions.

* Offer high level definitions.

* Offer definitions in primitive terms.

* fixup! Offer definitions in primitive terms.

* Define `inits` in terms of `initsNE`.

* Define `tails` in terms of `tailsNE`.

* Be compatible with older `base`.

* `NonEmptyList.toList` is too lazy.

* fixup! `NonEmptyList.toList` is too lazy.

* fixup! Write test cases and stub functions.

* Add `since` metadata.

* Make sure everything inlines.

L.drop does not inline because it is recursive.

Co-authored-by: Xia Li-yao <Lysxia@users.noreply.github.com>

* Fix false complexity statements.

---------

Co-authored-by: Xia Li-yao <Lysxia@users.noreply.github.com>

src/Data/Text.hs

@@ -155,7 +155,9 @@ module Data.Text
     , group
     , groupBy
     , inits
+    , initsNE
     , tails
+    , tailsNE
 
     -- ** Breaking into many substrings
     -- $split
@@ -227,6 +229,7 @@ import Control.Monad (foldM)
 import Control.Monad.ST (ST, runST)
 import qualified Data.Text.Array as A
 import qualified Data.List as L hiding (head, tail)
+import qualified Data.List.NonEmpty as NonEmptyList
 import Data.Binary (Binary(get, put))
 import Data.Monoid (Monoid(..))
 import Data.Semigroup (Semigroup(..))
@@ -1522,17 +1525,33 @@ group = groupBy (==)
 -- | /O(n)/ Return all initial segments of the given 'Text', shortest
 -- first.
 inits :: Text -> [Text]
-inits t = empty : case t of
+inits = (NonEmptyList.toList $!) . initsNE
+
+-- | /O(n)/ Return all initial segments of the given 'Text', shortest
+-- first.
+--
+-- @since 2.1.2
+initsNE :: Text -> NonEmptyList.NonEmpty Text
+initsNE t = empty NonEmptyList.:| case t of
   Text arr off len ->
-    let loop i | i >= len = []
-               | otherwise = let !j = i + iter_ t i in Text arr off j : loop j
+    let loop i
+          | i >= len = []
+          | otherwise = let !j = i + iter_ t i in Text arr off j : loop j
     in loop 0
 
 -- | /O(n)/ Return all final segments of the given 'Text', longest
 -- first.
 tails :: Text -> [Text]
-tails t | null t    = [empty]
-        | otherwise = t : tails (unsafeTail t)
+tails = (NonEmptyList.toList $!) . tailsNE
+
+-- | /O(n)/ Return all final segments of the given 'Text', longest
+-- first.
+--
+-- @since 2.1.2
+tailsNE :: Text -> NonEmptyList.NonEmpty Text
+tailsNE t
+  | null t = empty NonEmptyList.:| []
+  | otherwise = t NonEmptyList.:| tails (unsafeTail t)
 
 -- $split
 --

src/Data/Text/Lazy.hs

@@ -156,7 +156,9 @@ module Data.Text.Lazy
     , group
     , groupBy
     , inits
+    , initsNE
     , tails
+    , tailsNE
 
     -- ** Breaking into many substrings
     -- $split
@@ -1424,21 +1426,35 @@ groupBy eq (Chunk t ts) = cons x ys : groupBy eq zs
                                 x  = T.unsafeHead t
                                 xs = chunk (T.unsafeTail t) ts
 
--- | /O(n)/ Return all initial segments of the given 'Text',
+-- | /O(n²)/ Return all initial segments of the given 'Text',
 -- shortest first.
 inits :: Text -> [Text]
-inits = (Empty :) . inits'
+inits = (NE.toList P.$!) . initsNE
+
+-- | /O(n²)/ Return all initial segments of the given 'Text',
+-- shortest first.
+--
+-- @since 2.1.2
+initsNE :: Text -> NonEmpty Text
+initsNE = (Empty NE.:|) . inits'
   where inits' Empty        = []
-        inits' (Chunk t ts) = L.map (\t' -> Chunk t' Empty) (L.drop 1 (T.inits t))
+        inits' (Chunk t ts) = L.map (\t' -> Chunk t' Empty) (NE.tail (T.initsNE t))
                            ++ L.map (Chunk t) (inits' ts)
 
 -- | /O(n)/ Return all final segments of the given 'Text', longest
 -- first.
 tails :: Text -> [Text]
-tails Empty         = Empty : []
-tails ts@(Chunk t ts')
-  | T.length t == 1 = ts : tails ts'
-  | otherwise       = ts : tails (Chunk (T.unsafeTail t) ts')
+tails = (NE.toList P.$!) . tailsNE
+
+-- | /O(n)/ Return all final segments of the given 'Text', longest
+-- first.
+--
+-- @since 2.1.2
+tailsNE :: Text -> NonEmpty Text
+tailsNE Empty = Empty :| []
+tailsNE ts@(Chunk t ts')
+  | T.length t == 1 = ts :| tails ts'
+  | otherwise       = ts :| tails (Chunk (T.unsafeTail t) ts')
 
 -- $split
 --

tests/Tests/Properties/Substrings.hs

@@ -15,6 +15,7 @@ import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.QuickCheck (testProperty)
 import Tests.QuickCheckUtils
 import qualified Data.List as L
+import qualified Data.List.NonEmpty as NonEmptyList
 import qualified Data.Text as T
 import qualified Data.Text.Internal.Fusion as S
 import qualified Data.Text.Internal.Fusion.Common as S
@@ -156,9 +157,13 @@ t_groupBy (applyFun2 -> p)
 tl_groupBy (applyFun2 -> p)
                   = L.groupBy p   `eqP` (map unpackS . TL.groupBy p)
 t_inits           = L.inits       `eqP` (map unpackS . T.inits)
+t_initsNE         = NonEmptyList.inits `eqP` (fmap unpackS . T.initsNE)
 tl_inits          = L.inits       `eqP` (map unpackS . TL.inits)
+tl_initsNE        = NonEmptyList.inits `eqP` (fmap unpackS . TL.initsNE)
 t_tails           = L.tails       `eqP` (map unpackS . T.tails)
+t_tailsNE         = NonEmptyList.tails `eqP` (fmap unpackS . T.tailsNE)
 tl_tails          = L.tails       `eqPSqrt` (map unpackS . TL.tails)
+tl_tailsNE        = NonEmptyList.tails `eqP` (fmap unpackS . TL.tailsNE)
 
 spanML :: Monad m => (b -> m Bool) -> [b] -> m ([b], [b])
 spanML p s = go [] s
@@ -361,9 +366,13 @@ testSubstrings =
       testProperty "t_groupBy" t_groupBy,
       testProperty "tl_groupBy" tl_groupBy,
       testProperty "t_inits" t_inits,
+      testProperty "t_initsNE" t_initsNE,
       testProperty "tl_inits" tl_inits,
+      testProperty "tl_initsNE" tl_initsNE,
       testProperty "t_tails" t_tails,
+      testProperty "t_tailsNE" t_tailsNE,
       testProperty "tl_tails" tl_tails,
+      testProperty "tl_tailsNE" tl_tailsNE,
       testProperty "t_spanM" t_spanM,
       testProperty "t_spanEndM" t_spanEndM,
       testProperty "tl_spanM" tl_spanM,

tests/Tests/ShareEmpty.hs

@@ -15,6 +15,7 @@ import Test.Tasty (TestTree, testGroup)
 import GHC.Exts
 import GHC.Stack
 import qualified Data.List as L
+import qualified Data.List.NonEmpty as NonEmptyList
 import qualified Data.Text as T
 
 
@@ -105,10 +106,13 @@ tests = testGroup "empty Text values are shared"
       assertPtrEqEmpty . snd =<< T.spanEndM (const $ pure False) "123"
   , testCase "groupBy _ empty = [empty]" $ mapM_ assertPtrEqEmpty $ T.groupBy (==) empty
   , testCase "inits empty = [empty]" $ mapM_ assertPtrEqEmpty $ T.inits empty
+  , testCase "initsNE empty = singleton empty" $ mapM_ assertPtrEqEmpty $ T.initsNE empty
   , testCase "inits _ = [empty, ...]" $ assertPtrEqEmpty $ L.head $ T.inits "123"
+  , testCase "initsNE _ = empty :| ..." $ assertPtrEqEmpty $ NonEmptyList.head $ T.initsNE "123"
   , testCase "tails empty = [empty]" $ mapM_ assertPtrEqEmpty $ T.tails empty
+  , testCase "tailsNE empty = singleton empty" $ mapM_ assertPtrEqEmpty $ T.tailsNE empty
   , testCase "tails _ = [..., empty]" $ assertPtrEqEmpty $ L.last $ T.tails "123"
-  , testCase "tails _ = [..., empty]" $ assertPtrEqEmpty $ L.last $ T.tails "123"
+  , testCase "tailsNE _ = reverse (empty :| ...)" $ assertPtrEqEmpty $ NonEmptyList.last $ T.tailsNE "123"
   , testCase "split _ empty = [empty]" $ mapM_ assertPtrEqEmpty $ T.split (== 'a') ""
   , testCase "filter (const False) _ = empty" $ assertPtrEqEmpty $ T.filter (const False) "1234"
   , testCase "zipWith const empty empty = empty" $ assertPtrEqEmpty $ T.zipWith const "" ""