Better Chain Arbitrary (typelevel#2420)

* Better Chain Arbitrary * Fix bug in uncons
catostrophe · Sep 15, 2018 · bedd44e · bedd44e
1 parent 3e8320e
commit bedd44e
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 21 deletions.
diff --git a/core/src/main/scala/cats/data/Chain.scala b/core/src/main/scala/cats/data/Chain.scala
@@ -31,8 +31,10 @@ sealed abstract class Chain[+A] {
           result = Some(a -> next)
         case Append(l, r) => c = l; rights += r
         case Wrap(seq) =>
-          val tail = seq.tail
-          val next = fromSeq(tail)
+          val tail = fromSeq(seq.tail)
+          val next =
+            if (rights.isEmpty) tail
+            else tail ++ rights.reduceLeft((x, y) => Append(y, x))
           result = Some((seq.head, next))
         case Empty =>
           if (rights.isEmpty) {
@@ -424,22 +426,7 @@ object Chain extends ChainInstances {
 
   /** Creates a Chain from the specified elements. */
   def apply[A](as: A*): Chain[A] =
-    as match {
-      case w: collection.mutable.WrappedArray[A] =>
-        if (w.isEmpty) nil
-        else if (w.size == 1) one(w.head)
-        else {
-          val arr: Array[A] = w.array
-          var c: Chain[A] = one(arr.last)
-          var idx = arr.size - 2
-          while (idx >= 0) {
-            c = Append(one(arr(idx)), c)
-            idx -= 1
-          }
-          c
-        }
-      case _ => fromSeq(as)
-    }
+    fromSeq(as)
 
   // scalastyle:off null
   class ChainIterator[A](self: Chain[A]) extends Iterator[A] {

diff --git a/laws/src/main/scala/cats/laws/discipline/Arbitrary.scala b/laws/src/main/scala/cats/laws/discipline/Arbitrary.scala
@@ -293,9 +293,20 @@ object arbitrary extends ArbitraryInstances0 {
       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 => Chain.fromSeq(Range.apply(0, n)).foldLeft(Gen.const(Chain.empty[A])) { (gen, _) =>
-        gen.flatMap(cat => A.arbitrary.map(a => cat :+ a))
-      }
+      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 catsLawsCogenForChain[A](implicit A: Cogen[A]): Cogen[Chain[A]] =