Allow Lst of Array

Author: odersky

tests/run/lst/Lst.scala

@@ -5,13 +5,13 @@ import printing.{Printer, Texts}
 import Texts.Text
 import collection.mutable.{ListBuffer, StringBuilder}
 
-
 /** A lightweight class for lists, optimized for short and medium lengths.
  *  A list is represented at runtime as
  *
- *    If it is empty:               the value `Lst.Empty`
- *    If it contains one element:   the element itself
- *    If it contains more elements: an Array[Any] containing the elements
+ *    If it is empty:                    the value `Lst.Empty`
+ *    If it contains one element,
+ *    and the element is not an array:   the element itself
+ *    Otherwise:                         an Array[Any] containing the elements
  */
 class Lst[+T](val elems: Any) extends AnyVal {
   import Lst._
@@ -68,8 +68,8 @@ class Lst[+T](val elems: Any) extends AnyVal {
         val newElems = new Arr(elems.length)
         var i = 0
         while (i < elems.length) { newElems(i) = op(elem(i)); i += 1 }
-        new Lst[U](newElems)
-      case elem: T @ unchecked => new Lst[U](op(elem))
+        multi[U](newElems)
+      case elem: T @ unchecked => single[U](op(elem))
     }
   }
 
@@ -89,8 +89,8 @@ class Lst[+T](val elems: Any) extends AnyVal {
         }
         i += 1
       }
-      if (newElems == null) this.asInstanceOf[Lst[U]] else new Lst[U](newElems)
-    case elem: T @ unchecked => new Lst[U](f(elem))
+      if (newElems == null) this.asInstanceOf[Lst[U]] else multi[U](newElems)
+    case elem: T @ unchecked => single[U](f(elem))
   }
 
   def flatMap[U](f: T => Lst[U]): Lst[U] = elems match {
@@ -121,9 +121,9 @@ class Lst[+T](val elems: Any) extends AnyVal {
           j += ys.length
           i += 1
         }
-        new Lst[U](newElems)
+        multi[U](newElems)
       }
-    case elem: T @ unchecked => new Lst[U](f(elem).elems)
+    case elem: T @ unchecked => f(elem)
   }
 
   def filter(p: T => Boolean): Lst[T] = elems match {
@@ -214,7 +214,7 @@ class Lst[+T](val elems: Any) extends AnyVal {
         newElems(elems.length - 1 - i) = elem(i)
         i += 1
       }
-      new Lst[T](newElems)
+      multi[T](newElems)
     case _ => this
   }
 
@@ -252,7 +252,7 @@ class Lst[+T](val elems: Any) extends AnyVal {
       val newElems = new Arr(len1 + len2)
       this.copyToArray(newElems, 0)
       that.copyToArray(newElems, len1)
-      new Lst[U](newElems)
+      multi[U](newElems)
     }
 
   def zipWith[U, V](that: Lst[U])(op: (T, U) => V): Lst[V] =
@@ -264,7 +264,7 @@ class Lst[+T](val elems: Any) extends AnyVal {
           case elems2: Arr => def elem2(i: Int) = elems2(i).asInstanceOf[U]
             val len = elems1.length min elems2.length
             if (len == 0) Empty
-            else if (len == 1) new Lst[V](op(elem1(0), elem2(0)))
+            else if (len == 1) single[V](op(elem1(0), elem2(0)))
             else {
               var newElems: Arr = null
               var i = 0
@@ -279,17 +279,17 @@ class Lst[+T](val elems: Any) extends AnyVal {
                 }
                 i += 1
               }
-              new Lst[V](newElems)
+              multi[V](newElems)
             }
           case elem2: U @unchecked =>
-            new Lst[V](op(elem1(0), elem2))
+            single[V](op(elem1(0), elem2))
         }
       case elem1: T @unchecked =>
         that.elems match {
           case null => Empty
           case elems2: Arr => def elem2(i: Int) = elems2(i).asInstanceOf[U]
-            new Lst[V](op(elem1, elem2(0)))
-          case elem2: U @unchecked => new Lst[V](op(elem1, elem2))
+            single[V](op(elem1, elem2(0)))
+          case elem2: U @unchecked => single[V](op(elem1, elem2))
         }
     }
 
@@ -301,9 +301,9 @@ class Lst[+T](val elems: Any) extends AnyVal {
       val newElems = new Arr(elems.length)
       var i = 0
       while (i < elems.length) { newElems(i) = (elem(i), i); i += 1 }
-      new Lst[(T, Int)](newElems)
+      multi[(T, Int)](newElems)
     case elem: T @unchecked =>
-      new Lst[(T, Int)]((elem, 0))
+      single[(T, Int)]((elem, 0))
   }
 
   def corresponds[U](that: Lst[U])(p: (T, U) => Boolean): Boolean =
@@ -381,7 +381,7 @@ object Lst {
 
   def apply[T](): Lst[T] = Empty
 
-  def apply[T](x0: T): Lst[T] = new Lst[T](x0)
+  def apply[T](x0: T): Lst[T] = single[T](x0)
 
   def apply[T](x0: T, x1: T): Lst[T] = {
     val elems = new Arr(2)
@@ -427,12 +427,12 @@ object Lst {
 
   class Buffer[T] {
     private var len = 0
-    private var elem: T = _
+    private var elem: Any = _
     private var elems: Arr = _
 
     def size = len
 
-    /** pre: len > 0, n > 1 */
+    /** pre: len > 0, n >= 1 */
     private def ensureSize(n: Int) =
       if (len == 1) {
         elems = new Arr(n `max` 16)
@@ -459,7 +459,9 @@ object Lst {
       xs.elems match {
         case null => this
         case elems2: Arr =>
-          if (len == 0) elems = elems2
+          if (len == 0 && elems2.length != 1)
+            // if elems2.length == 1, elems2 is a wrapped single element list, which has to be unpacked
+            elems = elems2
           else {
             ensureSize(len + elems2.length)
             System.arraycopy(elems2, 0, elems, len, elems2.length)
@@ -472,17 +474,30 @@ object Lst {
 
     def toLst: Lst[T] =
       if (len == 0) Empty
-      else if (len == 1) new Lst[T](elem)
+      else if (len == 1) single(elem)
       else _fromArray(elems, 0, len)
 
     def clear() =
       len = 0
   }
 
+  private def single[T](elem: Any): Lst[T] = elem match {
+    case elem: Arr =>
+      val wrapped = new Arr(1)
+      wrapped(0) = elem
+      new Lst[T](wrapped)
+    case _ =>
+      new Lst[T](elem)
+  }
+
+  private def multi[T](elems: Array[Any]) =
+    if (elems.length == 1) new Lst[T](elems(0))
+    else new Lst[T](elems)
+
   private def _fromArray[T](elems: Arr, start: Int, end: Int): Lst[T] = {
     val len = end - start
     if (len <= 0) Empty
-    else if (len == 1) new Lst[T](elems(start))
+    else if (len == 1) single[T](elems(start))
     else if (start == 0 && end == elems.length) new Lst[T](elems)
     else {
       val newElems = new Arr(len)

tests/run/lst/LstTest.check

@@ -0,0 +1 @@
+testing