Merge pull request #5211 from tanishiking/copy-on-write-arraylist

Optimize CopyOnWriteArrayList to use scala.Array on Wasm

Author: sjrd

javalib/src/main/scala/java/util/concurrent/CopyOnWriteArrayList.scala

@@ -12,6 +12,8 @@
 
 package java.util.concurrent
 
+import scala.language.higherKinds
+
 import java.lang.Cloneable
 import java.lang.Utils._
 import java.lang.{reflect => jlr}
@@ -23,32 +25,46 @@ import scala.annotation.tailrec
 import ScalaOps._
 
 import scala.scalajs._
+import scala.scalajs.LinkingInfo._
 
-class CopyOnWriteArrayList[E <: AnyRef] private (private var inner: js.Array[E])
+class CopyOnWriteArrayList[E <: AnyRef] private (initialCapacity: Int)
     extends List[E] with RandomAccess with Cloneable with Serializable {
   self =>
 
+  /* This class has two different implementations for the
+   * internal data storage, depending on whether we are on Wasm or JS.
+   * We use `js.Array` on JS, and `scala.Array` on Wasm.
+   * The initialCapacity parameter is effective only in Wasm,
+   * since js.Array doesn't support explicit pre-allocation.
+   *
+   * On Wasm, we store the length at the index 0 of the array.
+   */
+
+  import CopyOnWriteArrayList._
+
+  private var inner: innerImpl.Repr[E] = innerImpl.make(initialCapacity)
+
   // requiresCopyOnWrite is false if and only if no other object
   // (like the iterator) may have a reference to inner
   private var requiresCopyOnWrite = false
 
   def this() = {
-    this(new js.Array[E])
+    this(16)
   }
 
   def this(c: Collection[_ <: E]) = {
-    this()
+    this(c.size())
     addAll(c)
   }
 
   def this(toCopyIn: Array[E]) = {
-    this()
+    this(toCopyIn.length)
     for (i <- 0 until toCopyIn.length)
       add(toCopyIn(i))
   }
 
   def size(): Int =
-    inner.length
+    innerImpl.length(inner)
 
   def isEmpty(): Boolean =
     size() == 0
@@ -175,7 +191,7 @@ class CopyOnWriteArrayList[E <: AnyRef] private (private var inner: js.Array[E])
   }
 
   def clear(): Unit = {
-    inner = new js.Array[E]
+    inner = innerImpl.make(16)
     requiresCopyOnWrite = false
   }
 
@@ -274,43 +290,49 @@ class CopyOnWriteArrayList[E <: AnyRef] private (private var inner: js.Array[E])
   }
 
   protected def innerGet(index: Int): E =
-    inner(index)
+    innerImpl.get(inner, index)
 
   protected def innerSet(index: Int, elem: E): Unit =
-    inner(index) = elem
+    innerImpl.set(inner, index, elem)
 
-  protected def innerPush(elem: E): Unit =
-    inner.push(elem)
+  protected def innerPush(elem: E): Unit = {
+    val newInner = innerImpl.push(inner, elem)
+    if (LinkingInfo.isWebAssembly) // opt: for JS we know it's always the same
+      inner = newInner
+  }
 
-  protected def innerInsert(index: Int, elem: E): Unit =
-    inner.splice(index, 0, elem)
+  protected def innerInsert(index: Int, elem: E): Unit = {
+    val newInner = innerImpl.add(inner, index, elem)
+    if (LinkingInfo.isWebAssembly) // opt: for JS we know it's always the same
+      inner = newInner
+  }
 
   protected def innerInsertMany(index: Int, items: Collection[_ <: E]): Unit = {
-    val itemsArray = js.Array[E]()
-    items.scalaOps.foreach(itemsArray.push(_))
-    inner.splice(index, 0, itemsArray.toSeq: _*)
+    val newInner = innerImpl.add(inner, index, items)
+    if (LinkingInfo.isWebAssembly) // opt: for JS we know it's always the same
+      inner = newInner
   }
 
   protected def innerRemove(index: Int): E =
-    arrayRemoveAndGet(inner, index)
+    innerImpl.remove(inner, index)
 
   protected def innerRemoveMany(index: Int, count: Int): Unit =
-    inner.splice(index, count)
+    innerImpl.remove(inner, index, count)
 
   protected def copyIfNeeded(): Unit = {
     if (requiresCopyOnWrite) {
-      inner = inner.jsSlice()
+      inner = innerImpl.clone(inner)
       requiresCopyOnWrite = false
     }
   }
 
-  protected def innerSnapshot(): js.Array[E] = {
+  protected def innerSnapshot(): innerImpl.Repr[E] = {
     requiresCopyOnWrite = true
     inner
   }
 
   private class CopyOnWriteArrayListView(fromIndex: Int, private var toIndex: Int)
-      extends CopyOnWriteArrayList[E](null: js.Array[E]) {
+      extends CopyOnWriteArrayList[E](initialCapacity) {
     viewSelf =>
 
     override def size(): Int =
@@ -381,7 +403,7 @@ class CopyOnWriteArrayList[E <: AnyRef] private (private var inner: js.Array[E])
     override protected def copyIfNeeded(): Unit =
       self.copyIfNeeded()
 
-    override protected def innerSnapshot(): js.Array[E] =
+    override protected def innerSnapshot(): innerImpl.Repr[E] =
       self.innerSnapshot()
 
     protected def changeSize(delta: Int): Unit =
@@ -399,7 +421,8 @@ class CopyOnWriteArrayList[E <: AnyRef] private (private var inner: js.Array[E])
   }
 }
 
-private class CopyOnWriteArrayListIterator[E](arraySnapshot: js.Array[E], i: Int, start: Int, end: Int)
+private class CopyOnWriteArrayListIterator[E](
+    arraySnapshot: CopyOnWriteArrayList.innerImpl.Repr[E], i: Int, start: Int, end: Int)
     extends AbstractRandomAccessListIterator[E](i, start, end) {
   override def remove(): Unit =
     throw new UnsupportedOperationException
@@ -411,7 +434,7 @@ private class CopyOnWriteArrayListIterator[E](arraySnapshot: js.Array[E], i: Int
     throw new UnsupportedOperationException
 
   protected def get(index: Int): E =
-    arraySnapshot(index)
+    CopyOnWriteArrayList.innerImpl.get(arraySnapshot, index)
 
   protected def remove(index: Int): Unit =
     throw new UnsupportedOperationException
@@ -422,3 +445,162 @@ private class CopyOnWriteArrayListIterator[E](arraySnapshot: js.Array[E], i: Int
   protected def add(index: Int, e: E): Unit =
     throw new UnsupportedOperationException
 }
+
+object CopyOnWriteArrayList {
+
+  /* Get the best implementation of inner array for the given platform.
+   *
+   * Use Array[AnyRef] in WebAssembly to avoid JS-interop. In JS, use js.Array.
+   * It is resizable by nature, so manual resizing is not needed.
+   *
+   * `linkTimeIf` is needed here to ensure the optimizer knows
+   * there is only one implementation of `InnerArrayImpl`, and de-virtualize/inline
+   * the function calls.
+   */
+
+  // package private so that `protected def innerSnapshot` can access this field.
+  private[concurrent] val innerImpl: InnerArrayImpl = {
+    LinkingInfo.linkTimeIf[InnerArrayImpl](LinkingInfo.isWebAssembly) {
+      InnerArrayImpl.JArrayImpl
+    } {
+      InnerArrayImpl.JSArrayImpl
+    }
+  }
+
+  private[concurrent] sealed abstract class InnerArrayImpl {
+    type Repr[E] <: AnyRef
+
+    def make[E](initialCapacity: Int): Repr[E]
+    def length(v: Repr[_]): Int
+    def get[E](v: Repr[E], index: Int): E
+    def set[E](v: Repr[E], index: Int, e: E): Unit
+    def push[E](v: Repr[E], e: E): Repr[E]
+    def add[E](v: Repr[E], index: Int, e: E): Repr[E]
+    def add[E](v: Repr[E], index: Int, items: Collection[_ <: E]): Repr[E]
+    def remove[E](v: Repr[E], index: Int): E
+    def remove(v: Repr[_], index: Int, count: Int): Unit
+    def clone[E](v: Repr[E]): Repr[E]
+  }
+
+  private object InnerArrayImpl {
+    object JSArrayImpl extends InnerArrayImpl {
+      import scala.scalajs.js
+
+      type Repr[E] = js.Array[AnyRef]
+
+      @inline def make[E](_initialCapacity: Int): Repr[E] = js.Array[AnyRef]()
+
+      @inline def length(v: Repr[_]): Int = v.length
+
+      @inline def get[E](v: Repr[E], index: Int): E = v(index).asInstanceOf[E]
+
+      @inline def set[E](v: Repr[E], index: Int, e: E): Unit =
+        v(index) = e.asInstanceOf[AnyRef]
+
+      @inline def push[E](v: Repr[E], e: E): Repr[E] = {
+        v.push(e.asInstanceOf[AnyRef])
+        v
+      }
+
+      @inline def add[E](v: Repr[E], index: Int, e: E): Repr[E] = {
+        v.splice(index, 0, e.asInstanceOf[AnyRef])
+        v
+      }
+
+      @inline def add[E](v: Repr[E], index: Int, items: Collection[_ <: E]): Repr[E] = {
+        val itemsArray = js.Array[AnyRef]()
+        items.scalaOps.foreach(e => itemsArray.push(e.asInstanceOf[AnyRef]))
+        v.splice(index, 0, itemsArray.toSeq: _*)
+        v
+      }
+
+      @inline def remove[E](v: Repr[E], index: Int): E =
+        arrayRemoveAndGet(v, index).asInstanceOf[E]
+
+      @inline def remove(v: Repr[_], index: Int, count: Int): Unit =
+        v.splice(index, count)
+
+      @inline def clone[E](v: Repr[E]): Repr[E] =
+        v.jsSlice(0)
+    }
+
+    object JArrayImpl extends InnerArrayImpl {
+      type Repr[E] = Array[AnyRef]
+
+      @inline def make[E](initialCapacity: Int): Repr[E] = {
+        val v = new Array[AnyRef](roundUpToPowerOfTwo(initialCapacity + 1))
+        v(0) = 0.asInstanceOf[AnyRef]
+        v
+      }
+
+      @inline def length(v: Repr[_]): Int = v(0).asInstanceOf[Int]
+
+      @inline def get[E](v: Repr[E], index: Int): E = v(index + 1).asInstanceOf[E] // Index 0 stores the length
+
+      @inline def set[E](v: Repr[E], index: Int, e: E): Unit =
+        v(index + 1) = e.asInstanceOf[AnyRef]
+
+      @inline def push[E](v: Repr[E], e: E): Repr[E] = {
+        val size = length(v)
+        val newArr = ensureCapacity(v, size + 1)
+        newArr(size + 1) = e.asInstanceOf[AnyRef]
+        newArr(0) = (size + 1).asInstanceOf[AnyRef]
+        newArr
+      }
+
+      @inline def add[E](v: Repr[E], index: Int, e: E): Repr[E] = {
+        val innerIdx = index + 1
+        val size = length(v)
+        val newArr = ensureCapacity(v, size + 1)
+        System.arraycopy(newArr, innerIdx, newArr, innerIdx + 1, size + 1 - innerIdx)
+        newArr(innerIdx) = e.asInstanceOf[AnyRef]
+        newArr(0) = (size + 1).asInstanceOf[AnyRef]
+        newArr
+      }
+
+      @inline def add[E](v: Repr[E], index: Int, items: Collection[_ <: E]): Repr[E] = {
+        val innerIdx = index + 1
+        val size = length(v)
+        val itemsSize = items.size()
+        val newArr = ensureCapacity(v, size + itemsSize)
+        System.arraycopy(newArr, innerIdx, newArr, innerIdx + itemsSize, size + 1 - innerIdx)
+        System.arraycopy(items.toArray(), 0, newArr, innerIdx, itemsSize)
+        newArr(0) = (size + itemsSize).asInstanceOf[AnyRef]
+        newArr
+      }
+
+      @inline def remove[E](v: Repr[E], index: Int): E = {
+        val innerIdx = index + 1
+        val size = length(v)
+        val removed = v(innerIdx)
+        System.arraycopy(v, innerIdx + 1, v, innerIdx, size - innerIdx)
+        v(size) = null // free reference for GC
+        v(0) = (size - 1).asInstanceOf[AnyRef]
+        removed.asInstanceOf[E]
+      }
+
+      @inline def remove(v: Repr[_], index: Int, count: Int): Unit = {
+        val innerIdx = index + 1
+        val size = length(v)
+        val toIndex = innerIdx + count
+        System.arraycopy(v, toIndex, v, innerIdx, size + 1 - toIndex)
+        val newSize = size - count
+        Arrays.fill(v, newSize + 1, newSize + 1 + count, null) // free references for GC
+        v(0) = newSize.asInstanceOf[AnyRef]
+      }
+
+      @inline def clone[E](v: Repr[E]): Repr[E] =
+        v.clone()
+
+      @inline private def ensureCapacity[E](v: Repr[E], minCapacity: Int): Repr[E] = {
+        val capacity = v.length - 1
+        if (capacity < minCapacity) {
+          val newCapacity = roundUpToPowerOfTwo(minCapacity + 1) // Index 0 stores the length
+          Arrays.copyOf(v, newCapacity)
+        } else {
+          v
+        }
+      }
+    }
+  }
+}