Make Chain Arbitraries recursively build concatenations

core/src/main/scala/cats/data/Chain.scala

@@ -37,12 +37,8 @@ sealed abstract class Chain[+A] {
             else tail ++ rights.reduceLeft((x, y) => Append(y, x))
           result = Some((seq.head, next))
         case Empty =>
-          if (rights.isEmpty) {
-            result = None
-          } else {
-            c = rights.last
-            rights.trimEnd(1)
-          }
+          // Empty is only top level, it is never internal to an Append
+          result = None
       }
     }
     // scalastyle:on null
@@ -297,12 +293,8 @@ sealed abstract class Chain[+A] {
             else rights.reduceLeft((x, y) => Append(y, x))
           rights.clear()
         case Empty =>
-          if (rights.isEmpty) {
-            c = null
-          } else {
-            c = rights.last
-            rights.trimEnd(1)
-          }
+          // Empty is only top level, it is never internal to an Append
+          c = null
       }
     }
   }
@@ -429,10 +421,10 @@ object Chain extends ChainInstances {
     fromSeq(as)
 
   // scalastyle:off null
-  class ChainIterator[A](self: Chain[A]) extends Iterator[A] {
-    var c: Chain[A] = if (self.isEmpty) null else self
-    val rights = new collection.mutable.ArrayBuffer[Chain[A]]
-    var currentIterator: Iterator[A] = null
+  private class ChainIterator[A](self: Chain[A]) extends Iterator[A] {
+    private[this] var c: Chain[A] = if (self.isEmpty) null else self
+    private[this] val rights = new collection.mutable.ArrayBuffer[Chain[A]]
+    private[this] var currentIterator: Iterator[A] = null
 
     override def hasNext: Boolean = (c ne null) || ((currentIterator ne null) && currentIterator.hasNext)
 
@@ -461,8 +453,8 @@ object Chain extends ChainInstances {
               rights.clear()
               currentIterator = seq.iterator
               currentIterator.next
-            case Empty =>
-              go // This shouldn't happen
+            case null | Empty =>
+              throw new java.util.NoSuchElementException("next called on empty iterator")
           }
         }
 
@@ -473,10 +465,10 @@ object Chain extends ChainInstances {
 
 
   // scalastyle:off null
-  class ChainReverseIterator[A](self: Chain[A]) extends Iterator[A] {
-    var c: Chain[A] = if (self.isEmpty) null else self
-    val lefts = new collection.mutable.ArrayBuffer[Chain[A]]
-    var currentIterator: Iterator[A] = null
+  private class ChainReverseIterator[A](self: Chain[A]) extends Iterator[A] {
+    private[this] var c: Chain[A] = if (self.isEmpty) null else self
+    private[this] val lefts = new collection.mutable.ArrayBuffer[Chain[A]]
+    private[this] var currentIterator: Iterator[A] = null
 
     override def hasNext: Boolean = (c ne null) || ((currentIterator ne null) && currentIterator.hasNext)
 
@@ -505,8 +497,8 @@ object Chain extends ChainInstances {
               lefts.clear()
               currentIterator = seq.reverseIterator
               currentIterator.next
-            case Empty =>
-              go // This shouldn't happen
+            case null | Empty =>
+              throw new java.util.NoSuchElementException("next called on empty iterator")
           }
         }
 

laws/src/main/scala/cats/laws/discipline/Arbitrary.scala

@@ -77,12 +77,12 @@ object arbitrary extends ArbitraryInstances0 {
     Cogen[List[A]].contramap(_.toList)
 
   implicit def catsLawsArbitraryForNonEmptyChain[A](implicit A: Arbitrary[A]): Arbitrary[NonEmptyChain[A]] =
-    Arbitrary(implicitly[Arbitrary[Chain[A]]].arbitrary.flatMap(fa =>
-      Gen.oneOf(
-        A.arbitrary.map(a => NonEmptyChain.fromChainPrepend(a, fa)),
-        A.arbitrary.map(a => NonEmptyChain.fromChainAppend(fa, a)),
-        A.arbitrary.map(NonEmptyChain.one)
-    )))
+    Arbitrary(implicitly[Arbitrary[Chain[A]]].arbitrary.flatMap { chain =>
+      NonEmptyChain.fromChain(chain) match {
+        case None => A.arbitrary.map(NonEmptyChain.one)
+        case Some(ne) => Gen.const(ne)
+      }
+    })
 
   implicit def catsLawsCogenForNonEmptyChain[A](implicit A: Cogen[A]): Cogen[NonEmptyChain[A]] =
     Cogen[Chain[A]].contramap(_.toChain)
@@ -287,27 +287,32 @@ object arbitrary extends ArbitraryInstances0 {
   implicit def catsLawsCogenForAndThen[A, B](implicit F: Cogen[A => B]): Cogen[AndThen[A, B]] =
     Cogen((seed, x) => F.perturb(seed, x))
 
-  implicit def catsLawsArbitraryForChain[A](implicit A: Arbitrary[A]): Arbitrary[Chain[A]] =
-    Arbitrary(Gen.sized {
-      case 0 => Gen.const(Chain.nil)
-      case 1 => A.arbitrary.map(Chain.one)
-      case 2 => A.arbitrary.flatMap(a1 => A.arbitrary.flatMap(a2 =>
-        Chain.concat(Chain.one(a1), Chain.one(a2))))
-      case n =>
-        def fromList(m: Int) = Range.apply(0, m).toList.foldLeft(Gen.const(List.empty[A])) { (gen, _) =>
-          gen.flatMap(list => A.arbitrary.map(_ :: list))
-        }.map(Chain.fromSeq)
-        val split = fromList(n / 2).flatMap(as => fromList(n / 2).map(_ ++ as))
-        val appended = fromList(n - 1).flatMap(as => A.arbitrary.map(as :+ _))
-        val prepended = fromList(n - 1).flatMap(as => A.arbitrary.map(_ +: as))
-        val startEnd = fromList(n - 2).flatMap(as => A.arbitrary.flatMap(a => A.arbitrary.map(_ +: (as :+ a))))
-        val betweenListsAndEnd = fromList((n - 1) / 2).flatMap(as => A.arbitrary.flatMap(a =>
-          fromList((n - 1) / 2).flatMap(as2 => A.arbitrary.map(a2 => (as2 ++ (a +: as)) :+ a2))))
-        val betweenListsAndFront = fromList((n - 1) / 2).flatMap(as => A.arbitrary.flatMap(a =>
-          fromList((n - 1) / 2).flatMap(as2 => A.arbitrary.map(a2 => a2 +: (as2 ++ (a +: as))))))
-        Gen.oneOf(fromList(n), split, appended, prepended, startEnd, betweenListsAndEnd, betweenListsAndFront)
+  implicit def catsLawsArbitraryForChain[A](implicit A: Arbitrary[A]): Arbitrary[Chain[A]] = {
+    val genA = A.arbitrary
+
+    def genSize(sz: Int): Gen[Chain[A]] = {
+      val fromSeq = Gen.listOfN(sz, genA).map(Chain.fromSeq)
+      val recursive =
+        sz match {
+          case 0 => Gen.const(Chain.nil)
+          case 1 => genA.map(Chain.one)
+          case n =>
+            // Here we concat two chains
+            for {
+              n0 <- Gen.choose(1, n-1)
+              n1 = n - n0
+              left <- genSize(n0)
+              right <- genSize(n1)
+            } yield left ++ right
+        }
+
+      // prefer to generate recursively built Chains
+      // but sometimes create fromSeq
+      Gen.frequency((5, recursive), (1, fromSeq))
+    }
 
-    })
+    Arbitrary(Gen.sized(genSize))
+  }
 
   implicit def catsLawsCogenForChain[A](implicit A: Cogen[A]): Cogen[Chain[A]] =
     Cogen[List[A]].contramap(_.toList)

tests/src/test/scala/cats/tests/ChainSuite.scala

@@ -126,4 +126,45 @@ class ChainSuite extends CatsSuite {
       ci.reverse.toList should === (ci.toList.reverse)
     }
   }
+
+  test("(a ++ b).isEmpty ==> a.isEmpty and b.isEmpty") {
+    forAll { (a: Chain[Int], b: Chain[Int]) =>
+      assert((a ++ b).nonEmpty || (a.isEmpty && b.isEmpty))
+    }
+  }
+
+  test("a.isEmpty == (a eq Chain.nil)") {
+    assert(Chain.fromSeq(Nil) eq Chain.nil)
+
+    forAll { (a: Chain[Int]) =>
+      assert(a.isEmpty == (a eq Chain.nil))
+    }
+  }
+
+  test("(nil ++ a) eq a") {
+    forAll { (a: Chain[Int]) =>
+      assert((Chain.nil ++ a) eq a)
+      assert((a ++ Chain.nil) eq a)
+    }
+  }
+
+  test("Chain.iterator.next should throw NoSuchElementException") {
+    forAll { (a: Chain[Int]) =>
+      val it = a.iterator
+
+      while(it.hasNext) it.next
+
+      intercept[java.util.NoSuchElementException] {
+        it.next
+      }
+
+      val rit = a.reverseIterator
+
+      while(rit.hasNext) rit.next
+
+      intercept[java.util.NoSuchElementException] {
+        rit.next
+      }
+    }
+  }
 }