make Bufferable collections safer

bijection-core/src/main/scala/com/twitter/bijection/Bufferable.scala

@@ -22,6 +22,7 @@ import java.nio.channels.Channel
 import scala.annotation.implicitNotFound
 import scala.annotation.tailrec
 import scala.collection.mutable.{Builder, Map => MMap, Set => MSet, Buffer => MBuffer}
+import scala.collection.generic.CanBuildFrom
 import scala.util.{ Failure, Success, Try }
 import com.twitter.bijection.Inversion.attempt
 
@@ -222,50 +223,41 @@ object Bufferable extends GeneratedTupleBufferable with Serializable {
     val nextBb = reallocatingPut(bb){ _.putInt(size) }
     l.foldLeft(nextBb) { (oldbb, t) => reallocatingPut(oldbb) { buf.put(_, t) } }
   }
+  // Assumes the builder passed in was just created from a CanBuildFrom
+  // If the same builder is used concurrently twice anywhere, you will have bugs.
+  // TODO: remove this at the next binary change
   def getCollection[T,C](initbb: ByteBuffer, builder: Builder[T,C])(implicit buf: Bufferable[T]):
-    Try[(ByteBuffer, C)] = {
+    Try[(ByteBuffer, C)] = Try {
 
-    val bbOpt: Try[ByteBuffer] = Try(initbb.duplicate)
-    val size = bbOpt.get.getInt
-    // We can't mutate the builder while calling other functions (not safe)
-    // so we write into this array:
-    val ary = new Array[Any](size)
-    (0 until size).foldLeft(bbOpt) { (oldBb, idx) =>
-      oldBb.flatMap { bb =>
-        buf.get(bb) match {
-          case Success((newbb, t)) =>
-            //Side-effect! scary!!!
-            ary(idx) = t
-            Success(newbb)
-          case Failure(t) => Failure(t)
-        }
-      }
-    }
-    .map { bb =>
-      // Now use the builder:
-      builder.clear()
-      builder.sizeHint(size)
-      ary.foreach { item => builder += item.asInstanceOf[T] }
-      (bb, builder.result())
+    var bb: ByteBuffer = initbb.duplicate
+    val size = bb.getInt
+    var idx = 0
+    builder.clear()
+    builder.sizeHint(size)
+    while(idx < size) {
+      val tup = buf.get(bb).get
+      bb = tup._1
+      builder += tup._2
+      idx += 1
     }
+    (bb, builder.result)
   }
-  def collection[C<:Traversable[T],T](builder: Builder[T,C])(implicit buf: Bufferable[T]):
+
+  def collection[C<:Traversable[T],T](implicit buf: Bufferable[T], cbf: CanBuildFrom[Nothing,T,C]):
     Bufferable[C] = build[C] { (bb, l) => putCollection(bb, l) }
-      { bb => getCollection(bb, builder) }
+      { bb => getCollection(bb, cbf()) }
 
-  implicit def list[T](implicit buf: Bufferable[T]) = collection[List[T], T](List.newBuilder[T])
-  implicit def set[T](implicit buf: Bufferable[T]) = collection[Set[T], T](Set.newBuilder[T])
-  implicit def indexedSeq[T](implicit buf: Bufferable[T]) =
-    collection[IndexedSeq[T], T](IndexedSeq.newBuilder[T])
-  implicit def vector[T](implicit buf: Bufferable[T]) =
-    collection[Vector[T], T](Vector.newBuilder[T])
+  implicit def list[T](implicit buf: Bufferable[T]) = collection[List[T], T]
+  implicit def set[T](implicit buf: Bufferable[T]) = collection[Set[T], T]
+  implicit def indexedSeq[T](implicit buf: Bufferable[T]) = collection[IndexedSeq[T], T]
+  implicit def vector[T](implicit buf: Bufferable[T]) = collection[Vector[T], T]
   implicit def map[K,V](implicit bufk: Bufferable[K], bufv: Bufferable[V]) =
-    collection[Map[K,V], (K,V)](Map.newBuilder[K,V])
+    collection[Map[K,V], (K,V)]
   // Mutable collections
   implicit def mmap[K,V](implicit bufk: Bufferable[K], bufv: Bufferable[V]) =
-    collection[MMap[K,V], (K,V)](MMap.newBuilder[K,V])
-  implicit def buffer[T](implicit buf: Bufferable[T]) = collection[MBuffer[T], T](MBuffer.newBuilder[T])
-  implicit def mset[T](implicit buf: Bufferable[T]) = collection[MSet[T], T](MSet.newBuilder[T])
+    collection[MMap[K,V], (K,V)]
+  implicit def buffer[T](implicit buf: Bufferable[T]) = collection[MBuffer[T], T]
+  implicit def mset[T](implicit buf: Bufferable[T]) = collection[MSet[T], T]
 
   // TODO we could add IntBuffer/FloatBuffer etc.. to have faster implementations Array[Int]
   implicit def array[T](implicit buf: Bufferable[T], cm: ClassManifest[T]): Bufferable[Array[T]] =

bijection-core/src/test/scala/com/twitter/bijection/BufferableLaws.scala

@@ -44,6 +44,9 @@ trait BaseBufferable {
   implicit protected def aeq[T:Equiv]: Equiv[Array[T]] = Equiv.fromFunction { (a1, a2) =>
     a1.zip(a2).forall { tup => Equiv[T].equiv(tup._1, tup._2) }
   }
+  def itereq[C <: Iterable[T], T:Equiv]: Equiv[C] = Equiv.fromFunction { (a1, a2) =>
+    a1.zip(a2).forall { tup => Equiv[T].equiv(tup._1, tup._2) }
+  }
 }
 
 object BufferableLaws extends Properties("Bufferable") with BaseBufferable {
@@ -61,6 +64,7 @@ object BufferableLaws extends Properties("Bufferable") with BaseBufferable {
     Bufferable.getBytes(bb).toList == bytes.toList
   }
 
+  implicit val laeq = itereq[List[Array[Byte]], Array[Byte]]
   property("Ints roundtrip") = roundTrips[Int]
   property("Doubles roundtrip") = roundTrips[Double]
   property("Floats roundtrip") = roundTrips[Float]
@@ -72,6 +76,7 @@ object BufferableLaws extends Properties("Bufferable") with BaseBufferable {
   property("Array[Byte] roundtrip") = roundTrips[Array[Byte]]
   property("Array[Int] roundtrip") = roundTrips[Array[Int]]
   property("List[Double] roundtrip") = roundTrips[List[Double]]
+  property("List[Array[Byte]] roundtrip") = roundTrips[List[Array[Byte]]]
   property("Set[Double] roundtrip") = roundTrips[Set[Double]]
   property("Map[String, Int] roundtrip") = roundTrips[Map[String, Int]]
   property("Option[(Long,Long)] roundtrip") = roundTrips[Option[(Long,Long)]]