Skip to content

Commit

Permalink
recurse on NonEmpty
Browse files Browse the repository at this point in the history
  • Loading branch information
@johnynek
johnynek committed Jan 3, 2025
1 parent 49625fe commit f56cad6
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 72 deletions.
141 changes: 73 additions & 68 deletions core/src/main/scala/cats/data/Chain.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
final def take(count: Long): Chain[A] = {
// invariant count >= 1
@tailrec
def go(lhs: Chain[A], count: Long, arg: Chain[A], rhs: Chain[A]): Chain[A] =
def go(lhs: Chain[A], count: Long, arg: NonEmpty[A], rhs: Chain[A]): Chain[A] =
arg match {
case Wrap(seq) =>
if (count == 1) {
Expand All @@ -277,31 +277,31 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
val wrapped = Wrap(taken)
// this is more efficient than using concat
val newLhs = if (lhs.isEmpty) wrapped else Append(lhs, wrapped)
if (newCount > 0) {
// we have to keep taking on the rhs
go(newLhs, newCount, rhs, Chain.nil)
} else {
// newCount == 0, we have taken enough
newLhs
rhs match {
case rhsNE: NonEmpty[A] if newCount > 0L =>
// we have to keep taking on the rhs
go(newLhs, newCount, rhsNE, Empty)
case _ =>
newLhs
}
}
case Append(l, r) =>
go(lhs, count, l, if (rhs.isEmpty) r else Append(r, rhs))
case s @ Singleton(_) =>
// due to the invariant count >= 1
val newLhs = if (lhs.isEmpty) s else Append(lhs, s)
if (count > 1L) {
go(newLhs, count - 1L, rhs, Chain.nil)
} else newLhs
case Empty =>
// this empty check isn't an optimization but to ensure
// the recursion terminates.
if (rhs.isEmpty) lhs
else go(lhs, count, rhs, Empty)
rhs match {
case rhsNE: NonEmpty[A] if count > 1L =>
go(newLhs, count - 1L, rhsNE, Empty)
case _ => newLhs
}
}

if (count <= 0L) Empty
else go(Empty, count, this, Empty)
this match {
case ne: NonEmpty[A] if count > 0L =>
go(Empty, count, ne, Empty)
case _ => Empty
}
}

/**
Expand All @@ -310,7 +310,7 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
final def takeRight(count: Long): Chain[A] = {
// invariant count >= 1
@tailrec
def go(lhs: Chain[A], count: Long, arg: Chain[A], rhs: Chain[A]): Chain[A] =
def go(lhs: Chain[A], count: Long, arg: NonEmpty[A], rhs: Chain[A]): Chain[A] =
arg match {
case Wrap(seq) =>
if (count == 1L) {
Expand All @@ -324,31 +324,29 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
val newCount = count - taken.length
val wrapped = Wrap(taken)
val newRhs = if (rhs.isEmpty) wrapped else Append(wrapped, rhs)
if (newCount > 0) {
// we have to keep taking on the rhs
go(Chain.nil, newCount, lhs, newRhs)
} else {
// newCount == 0, we have taken enough
newRhs
lhs match {
case lhsNE: NonEmpty[A] if newCount > 0 =>
go(Empty, newCount, lhsNE, newRhs)
case _ => newRhs
}
}
case Append(l, r) =>
go(if (lhs.isEmpty) l else Append(lhs, l), count, r, rhs)
case s @ Singleton(_) =>
// due to the invariant count >= 1
val newRhs = if (rhs.isEmpty) s else Append(s, rhs)
if (count > 1) {
go(Empty, count - 1, lhs, newRhs)
} else newRhs
case Empty =>
// this empty check isn't an optimization but to ensure
// the recursion terminates.
if (lhs.isEmpty) rhs
else go(Empty, count, lhs, rhs)
lhs match {
case lhsNE: NonEmpty[A] if count > 1 =>
go(Empty, count - 1, lhsNE, newRhs)
case _ => newRhs
}
}

if (count <= 0) Empty
else go(Empty, count, this, Empty)
this match {
case ne: NonEmpty[A] if count > 0L =>
go(Empty, count, ne, Empty)
case _ => Empty
}
}

/**
Expand Down Expand Up @@ -376,20 +374,21 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
final def drop(count: Long): Chain[A] = {
// invariant count >= 1
@tailrec
def go(count: Long, arg: Chain[A], rhs: Chain[A]): Chain[A] =
def go(count: Long, arg: NonEmpty[A], rhs: Chain[A]): Chain[A] =
arg match {
case Wrap(seq) =>
val dropped = if (count < Int.MaxValue) seq.drop(count.toInt) else seq.drop(Int.MaxValue)
if (dropped.isEmpty) {
// we may have not dropped all of count
val newCount = count - seq.length
if (newCount > 0) {
// we have to keep dropping on the rhs
go(newCount, rhs, Chain.nil)
} else {
// we know that count >= seq.length else we wouldn't be empty
// so in this case, it is exactly count == seq.length
rhs
rhs match {
case rhsNE: NonEmpty[A] if newCount > 0 =>
// we have to keep dropping on the rhs
go(newCount, rhsNE, Empty)
case _ =>
// we know that count >= seq.length else we wouldn't be empty
// so in this case, it is exactly count == seq.length
rhs
}
} else {
// dropped is not empty
Expand All @@ -401,17 +400,19 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
go(count, l, if (rhs.isEmpty) r else Append(r, rhs))
case Singleton(_) =>
// due to the invariant count >= 1
if (count > 1L) go(count - 1L, rhs, Chain.nil)
else rhs
case Empty =>
// this empty check isn't an optimization but to ensure
// the recursion terminates.
if (rhs.isEmpty) Empty
else go(count, rhs, Empty)
rhs match {
case rhsNE: NonEmpty[A] if count > 1L =>
go(count - 1L, rhsNE, Empty)
case _ =>
rhs
}
}

if (count <= 0L) this
else go(count, this, Empty)
this match {
case ne: NonEmpty[A] if count > 0L =>
go(count, ne, Empty)
case _ => this
}
}

/**
Expand All @@ -420,20 +421,21 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
final def dropRight(count: Long): Chain[A] = {
// invariant count >= 1
@tailrec
def go(lhs: Chain[A], count: Long, arg: Chain[A]): Chain[A] =
def go(lhs: Chain[A], count: Long, arg: NonEmpty[A]): Chain[A] =
arg match {
case Wrap(seq) =>
val dropped = if (count < Int.MaxValue) seq.dropRight(count.toInt) else seq.dropRight(Int.MaxValue)
if (dropped.isEmpty) {
// we may have not dropped all of count
val newCount = count - seq.length
if (newCount > 0L) {
// we have to keep dropping on the rhs
go(Chain.nil, newCount, lhs)
} else {
// we know that count >= seq.length else we wouldn't be empty
// so in this case, it is exactly count == seq.length
lhs
lhs match {
case lhsNE: NonEmpty[A] if newCount > 0L =>
// we have to keep dropping on the lhs
go(Empty, newCount, lhsNE)
case _ =>
// we know that count >= seq.length else we wouldn't be empty
// so in this case, it is exactly count == seq.length
lhs
}
} else {
// we must be done
Expand All @@ -445,17 +447,20 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
go(if (lhs.isEmpty) l else Append(lhs, l), count, r)
case Singleton(_) =>
// due to the invariant count >= 1
if (count > 1L) go(Chain.nil, count - 1L, lhs)
else lhs
case Empty =>
// this empty check isn't an optimization but to ensure
// the recursion terminates.
if (lhs.isEmpty) Empty
else go(Empty, count, lhs)
lhs match {
case lhsNE: NonEmpty[A] if count > 1L =>
go(Empty, count - 1L, lhsNE)
case _ =>
lhs
}
}

if (count <= 0) this
else go(Empty, count, this)
this match {
case ne: NonEmpty[A] if count > 0L =>
go(Empty, count, ne)
case _ =>
this
}
}

/**
Expand Down
8 changes: 4 additions & 4 deletions tests/shared/src/test/scala/cats/tests/ChainSuite.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -462,24 +462,24 @@ class ChainSuite extends CatsSuite {

test("drop(cnt).toList == toList.drop(cnt)") {
forAll(genChainDropTakeArgs) { case (chain: Chain[Int], count: Int) =>
assert(chain.drop(count).toList == chain.toList.drop(count))
assertEquals(chain.drop(count).toList, chain.toList.drop(count))
}
}

test("dropRight(cnt).toList == toList.dropRight(cnt)") {
forAll(genChainDropTakeArgs) { case (chain: Chain[Int], count: Int) =>
assert(chain.dropRight(count).toList == chain.toList.dropRight(count))
assertEquals(chain.dropRight(count).toList, chain.toList.dropRight(count))
}
}
test("take(cnt).toList == toList.take(cnt)") {
forAll(genChainDropTakeArgs) { case (chain: Chain[Int], count: Int) =>
assert(chain.take(count).toList == chain.toList.take(count))
assertEquals(chain.take(count).toList, chain.toList.take(count))
}
}

test("takeRight(cnt).toList == toList.takeRight(cnt)") {
forAll(genChainDropTakeArgs) { case (chain: Chain[Int], count: Int) =>
assert(chain.takeRight(count).toList == chain.toList.takeRight(count))
assertEquals(chain.takeRight(count).toList, chain.toList.takeRight(count))
}
}
}

0 comments on commit f56cad6

Please sign in to comment.