Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feature: initsNE and tailsNE. #558

Merged
merged 14 commits into from
Feb 20, 2024
29 changes: 24 additions & 5 deletions src/Data/Text.hs
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,9 @@ module Data.Text
, group
, groupBy
, inits
, initsNE
, tails
, tailsNE

-- ** Breaking into many substrings
-- $split
Expand Down Expand Up @@ -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(..))
Expand Down Expand Up @@ -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
--
Expand Down
28 changes: 22 additions & 6 deletions src/Data/Text/Lazy.hs
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,9 @@ module Data.Text.Lazy
, group
, groupBy
, inits
, initsNE
, tails
, tailsNE

-- ** Breaking into many substrings
-- $split
Expand Down Expand Up @@ -1427,18 +1429,32 @@ groupBy eq (Chunk t ts) = cons x ys : groupBy eq zs
-- | /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',
Lysxia marked this conversation as resolved.
Show resolved Hide resolved
-- 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
--
Expand Down
9 changes: 9 additions & 0 deletions tests/Tests/Properties/Substrings.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
6 changes: 5 additions & 1 deletion tests/Tests/ShareEmpty.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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 "" ""
Expand Down
Loading