diff --git a/core/src/main/scala/cats/data/Chain.scala b/core/src/main/scala/cats/data/Chain.scala index fb1594a588..0333ac404a 100644 --- a/core/src/main/scala/cats/data/Chain.scala +++ b/core/src/main/scala/cats/data/Chain.scala @@ -47,11 +47,47 @@ sealed abstract class Chain[+A] { result } + /** + * Returns the init and last of this Chain if non empty, none otherwise. Amortized O(1). + */ + final def initLast: Option[(Chain[A], A)] = { + var c: Chain[A] = this + val lefts = new collection.mutable.ArrayBuffer[Chain[A]] + // scalastyle:off null + var result: Option[(Chain[A], A)] = null + while (result eq null) { + c match { + case Singleton(a) => + val pre = + if (lefts.isEmpty) nil + else lefts.reduceLeft((x, y) => Append(x, y)) + result = Some(pre -> a) + case Append(l, r) => c = r; lefts += l + case Wrap(seq) => + val init = fromSeq(seq.init) + val pre = + if (lefts.isEmpty) init + else lefts.reduceLeft((x, y) => Append(x, y)) ++ init + result = Some((pre, seq.last)) + case Empty => + // Empty is only top level, it is never internal to an Append + result = None + } + } + // scalastyle:on null + result + } + /** * Returns the head of this Chain if non empty, none otherwise. Amortized O(1). */ def headOption: Option[A] = uncons.map(_._1) + /** + * Returns the last of this Chain if non empty, none otherwise. Amortized O(1). + */ + final def lastOption: Option[A] = initLast.map(_._2) + /** * Returns true if there are no elements in this collection. */ diff --git a/core/src/main/scala/cats/data/NonEmptyChain.scala b/core/src/main/scala/cats/data/NonEmptyChain.scala index bce726b7b4..85f66f8d60 100644 --- a/core/src/main/scala/cats/data/NonEmptyChain.scala +++ b/core/src/main/scala/cats/data/NonEmptyChain.scala @@ -191,15 +191,30 @@ class NonEmptyChainOps[A](private val value: NonEmptyChain[A]) extends AnyVal { final def uncons: (A, Chain[A]) = toChain.uncons.get /** - * Returns the first element of this chain. + * Returns the init and last of this NonEmptyChain. Amortized O(1). + */ + final def initLast: (Chain[A], A) = toChain.initLast.get + + /** + * Returns the first element of this NonEmptyChain. Amortized O(1). */ final def head: A = uncons._1 /** - * Returns all but the first element of this chain. + * Returns all but the first element of this NonEmptyChain. Amortized O(1). */ final def tail: Chain[A] = uncons._2 + /** + * Returns all but the last element of this NonEmptyChain. Amortized O(1). + */ + final def init: Chain[A] = initLast._1 + + /** + * Returns the last element of this NonEmptyChain. Amortized O(1). + */ + final def last: A = initLast._2 + /** * Tests if some element is contained in this chain. * {{{ diff --git a/tests/src/test/scala/cats/tests/ChainSuite.scala b/tests/src/test/scala/cats/tests/ChainSuite.scala index c6069efd8e..2b870fa18c 100644 --- a/tests/src/test/scala/cats/tests/ChainSuite.scala +++ b/tests/src/test/scala/cats/tests/ChainSuite.scala @@ -68,6 +68,12 @@ class ChainSuite extends CatsSuite { } } + test("lastOption") { + forAll { (c: Chain[Int]) => + c.lastOption should ===(c.toList.lastOption) + } + } + test("size is consistent with toList.size") { forAll { (ci: Chain[Int]) => ci.size.toInt should ===(ci.toList.size) diff --git a/tests/src/test/scala/cats/tests/NonEmptyChainSuite.scala b/tests/src/test/scala/cats/tests/NonEmptyChainSuite.scala index b9939ee25c..5a2ea55e28 100644 --- a/tests/src/test/scala/cats/tests/NonEmptyChainSuite.scala +++ b/tests/src/test/scala/cats/tests/NonEmptyChainSuite.scala @@ -139,4 +139,16 @@ class NonEmptyChainSuite extends CatsSuite { ci.distinct.toList should ===(ci.toList.distinct) } } + + test("init") { + forAll { ci: NonEmptyChain[Int] => + ci.init.toList should ===(ci.toList.init) + } + } + + test("last") { + forAll { ci: NonEmptyChain[Int] => + ci.last should ===(ci.toList.last) + } + } }