Wasm: Implement PriorityQueue without js.Array in Wasm backend

tanishiking · tanishiking · commit afdda6f804dd · 2025-07-24T11:39:00.000+09:00
Replaced the `js.Array`-based `PriorityQueue` backend with a `scala.Array` implementation when targeting Wasm. The previous js.Array based PriorityQueue implementation was not very performant on Wasm, because of the overhead of JS-interop for each inner array operations. A `linkTimeIf` condition now selects the appropriate array implementation at build time. Benchmarks show a 2-4x speedup for `add`, `peek`, and `poll` operations on the Wasm backend. tanishiking/scalajs-benchmarks#3
diff --git a/javalib/src/main/scala/java/util/PriorityQueue.scala b/javalib/src/main/scala/java/util/PriorityQueue.scala
@@ -12,31 +12,49 @@
 
 package java.util
 
+import scala.language.higherKinds
+
 import scala.annotation.tailrec
 
-import scala.scalajs.js
+import java.lang.Utils.roundUpToPowerOfTwo
+
+import scala.scalajs.LinkingInfo
 
 class PriorityQueue[E] private (
-    private val comp: Comparator[_ >: E], internal: Boolean)
+    private val comp: Comparator[_ >: E], internal: Boolean, initialCapacity: Int)
     extends AbstractQueue[E] with Serializable {
 
+  import PriorityQueue._
+
   def this() =
-    this(NaturalComparator, internal = true)
+    this(NaturalComparator, internal = true, initialCapacity = 16)
 
   def this(initialCapacity: Int) = {
-    this()
-    if (initialCapacity < 1)
-      throw new IllegalArgumentException()
+    this(
+      NaturalComparator,
+      internal = true,
+      {
+        if (initialCapacity < 1)
+          throw new IllegalArgumentException
+        initialCapacity + 1 // index 0 is unused
+      }
+    )
   }
 
   def this(comparator: Comparator[_ >: E]) = {
-    this(NaturalComparator.select(comparator), internal = true)
+    this(NaturalComparator.select(comparator), internal = true, initialCapacity = 16)
   }
 
   def this(initialCapacity: Int, comparator: Comparator[_ >: E]) = {
-    this(comparator)
-    if (initialCapacity < 1)
-      throw new IllegalArgumentException()
+    this(
+      NaturalComparator.select(comparator),
+      internal = true,
+      {
+        if (initialCapacity < 1)
+          throw new IllegalArgumentException()
+        initialCapacity + 1 // index 0 is unused
+      }
+    )
   }
 
   def this(c: Collection[_ <: E]) = {
@@ -47,47 +65,51 @@ class PriorityQueue[E] private (
         NaturalComparator.select(c.comparator().asInstanceOf[Comparator[_ >: E]])
       case _ =>
         NaturalComparator
-    }, internal = true)
+    }, internal = true, roundUpToPowerOfTwo(c.size() + 1)) // index 0 is unused
     addAll(c)
   }
 
   def this(c: PriorityQueue[_ <: E]) = {
-    this(c.comp.asInstanceOf[Comparator[_ >: E]], internal = true)
+    this(c.comp.asInstanceOf[Comparator[_ >: E]], internal = true,
+        roundUpToPowerOfTwo(c.size() + 1)) // index 0 is unused
     addAll(c)
   }
 
   def this(sortedSet: SortedSet[_ <: E]) = {
     this(NaturalComparator.select(
         sortedSet.comparator().asInstanceOf[Comparator[_ >: E]]),
-        internal = true)
+        internal = true,
+        roundUpToPowerOfTwo(sortedSet.size() + 1)) // index 0 is unused
     addAll(sortedSet)
   }
 
   // The index 0 is not used; the root is at index 1.
   // This is standard practice in binary heaps, to simplify arithmetics.
-  private[this] val inner = js.Array[E](null.asInstanceOf[E])
+  private var inner: innerImpl.Repr[E] = innerImpl.make[E](initialCapacity)
 
   override def add(e: E): Boolean = {
     if (e == null)
       throw new NullPointerException()
-    inner.push(e)
-    fixUp(inner.length - 1)
+    val newInner = innerImpl.push(inner, e)
+    if (LinkingInfo.isWebAssembly) // opt: for JS we know it's always the same
+      inner = newInner
+    fixUp(innerImpl.length(inner) - 1)
     true
   }
 
   def offer(e: E): Boolean = add(e)
 
   def peek(): E =
-    if (inner.length > 1) inner(1)
+    if (innerImpl.length(inner) > 1) innerImpl.get(inner, 1)
     else null.asInstanceOf[E]
 
   override def remove(o: Any): Boolean = {
     if (o == null) {
       false
     } else {
-      val len = inner.length
+      val len = innerImpl.length(inner)
       var i = 1
-      while (i != len && !o.equals(inner(i))) {
+      while (i != len && !o.equals(innerImpl.get(inner, i))) {
         i += 1
       }
 
@@ -101,9 +123,9 @@ class PriorityQueue[E] private (
   }
 
   private def removeExact(o: Any): Unit = {
-    val len = inner.length
+    val len = innerImpl.length(inner)
     var i = 1
-    while (i != len && (o.asInstanceOf[AnyRef] ne inner(i).asInstanceOf[AnyRef])) {
+    while (i != len && (o.asInstanceOf[AnyRef] ne innerImpl.get(inner, i).asInstanceOf[AnyRef])) {
       i += 1
     }
     if (i == len)
@@ -112,12 +134,12 @@ class PriorityQueue[E] private (
   }
 
   private def removeAt(i: Int): Unit = {
-    val newLength = inner.length - 1
+    val newLength = innerImpl.length(inner) - 1
     if (i == newLength) {
-      inner.length = newLength
+      innerImpl.decLength(inner)
     } else {
-      inner(i) = inner(newLength)
-      inner.length = newLength
+      innerImpl.set(inner, i, innerImpl.get(inner, newLength))
+      innerImpl.decLength(inner)
       fixUpOrDown(i)
     }
   }
@@ -126,9 +148,9 @@ class PriorityQueue[E] private (
     if (o == null) {
       false
     } else {
-      val len = inner.length
+      val len = innerImpl.length(inner)
       var i = 1
-      while (i != len && !o.equals(inner(i))) {
+      while (i != len && !o.equals(innerImpl.get(inner, i))) {
         i += 1
       }
       i != len
@@ -137,16 +159,16 @@ class PriorityQueue[E] private (
 
   def iterator(): Iterator[E] = {
     new Iterator[E] {
-      private[this] var inner: js.Array[E] = PriorityQueue.this.inner
+      private[this] var inner: innerImpl.Repr[E] = PriorityQueue.this.inner
       private[this] var nextIdx: Int = 1
       private[this] var last: E = _ // null
 
-      def hasNext(): Boolean = nextIdx < inner.length
+      def hasNext(): Boolean = nextIdx < innerImpl.length(inner)
 
       def next(): E = {
         if (!hasNext())
           throw new NoSuchElementException("empty iterator")
-        last = inner(nextIdx)
+        last = innerImpl.get(inner, nextIdx)
         nextIdx += 1
         last
       }
@@ -173,27 +195,27 @@ class PriorityQueue[E] private (
         if (last == null)
           throw new IllegalStateException()
         if (inner eq PriorityQueue.this.inner) {
-          inner = inner.jsSlice(nextIdx)
-          nextIdx = 0
+          inner = innerImpl.copyFrom(inner, nextIdx)
+          nextIdx = 1
         }
         removeExact(last)
         last = null.asInstanceOf[E]
       }
     }
   }
 
-  def size(): Int = inner.length - 1
+  def size(): Int = innerImpl.length(inner) - 1
 
   override def clear(): Unit =
-    inner.length = 1
+    innerImpl.clear(inner)
 
   def poll(): E = {
     val inner = this.inner // local copy
-    if (inner.length > 1) {
-      val newSize = inner.length - 1
-      val result = inner(1)
-      inner(1) = inner(newSize)
-      inner.length = newSize
+    if (innerImpl.length(inner) > 1) {
+      val newSize = innerImpl.length(inner) - 1
+      val result = innerImpl.get(inner, 1)
+      innerImpl.set(inner, 1, innerImpl.get(inner, newSize))
+      innerImpl.decLength(inner)
       fixDown(1)
       result
     } else {
@@ -212,7 +234,7 @@ class PriorityQueue[E] private (
    */
   private[this] def fixUpOrDown(m: Int): Unit = {
     val inner = this.inner // local copy
-    if (m > 1 && comp.compare(inner(m >> 1), inner(m)) > 0)
+    if (m > 1 && comp.compare(innerImpl.get(inner, m >> 1), innerImpl.get(inner, m)) > 0)
       fixUp(m)
     else
       fixDown(m)
@@ -227,16 +249,16 @@ class PriorityQueue[E] private (
     /* At each step, even though `m` changes, the element moves with it, and
      * hence inner(m) is always the same initial `innerAtM`.
      */
-    val innerAtM = inner(m)
+    val innerAtM = innerImpl.get(inner, m)
 
     @inline @tailrec
     def loop(m: Int): Unit = {
       if (m > 1) {
         val parent = m >> 1
-        val innerAtParent = inner(parent)
+        val innerAtParent = innerImpl.get(inner, parent)
         if (comp.compare(innerAtParent, innerAtM) > 0) {
-          inner(parent) = innerAtM
-          inner(m) = innerAtParent
+          innerImpl.set(inner, parent, innerAtM)
+          innerImpl.set(inner, m, innerAtParent)
           loop(parent)
         }
       }
@@ -250,22 +272,22 @@ class PriorityQueue[E] private (
    */
   private[this] def fixDown(m: Int): Unit = {
     val inner = this.inner // local copy
-    val size = inner.length - 1
+    val size = innerImpl.length(inner) - 1
 
     /* At each step, even though `m` changes, the element moves with it, and
      * hence inner(m) is always the same initial `innerAtM`.
      */
-    val innerAtM = inner(m)
+    val innerAtM = innerImpl.get(inner, m)
 
     @inline @tailrec
     def loop(m: Int): Unit = {
       var j = 2 * m // left child of `m`
       if (j <= size) {
-        var innerAtJ = inner(j)
+        var innerAtJ = innerImpl.get(inner, j)
 
         // if the left child is greater than the right child, switch to the right child
         if (j < size) {
-          val innerAtJPlus1 = inner(j + 1)
+          val innerAtJPlus1 = innerImpl.get(inner, j + 1)
           if (comp.compare(innerAtJ, innerAtJPlus1) > 0) {
             j += 1
             innerAtJ = innerAtJPlus1
@@ -274,13 +296,116 @@ class PriorityQueue[E] private (
 
         // if the node `m` is greater than the selected child, swap and recurse
         if (comp.compare(innerAtM, innerAtJ) > 0) {
-          inner(m) = innerAtJ
-          inner(j) = innerAtM
+          innerImpl.set(inner, m, innerAtJ)
+          innerImpl.set(inner, j, innerAtM)
           loop(j)
         }
       }
     }
 
     loop(m)
   }
+
+}
+
+object PriorityQueue {
+
+  /* Get the best available 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.
+   */
+
+  private val innerImpl: InnerArrayImpl = {
+    LinkingInfo.linkTimeIf[InnerArrayImpl](LinkingInfo.isWebAssembly) {
+      InnerArrayImpl.JArrayImpl
+    } {
+      InnerArrayImpl.JSArrayImpl
+    }
+  }
+
+  private sealed abstract class InnerArrayImpl {
+    type Repr[E] <: AnyRef
+
+    def make[E](initialCapacity: Int): Repr[E]
+    def length(v: Repr[_]): Int
+    def decLength(v: Repr[_]): Unit
+    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 copyFrom[E](v: Repr[E], from: Int): Repr[E]
+    def clear(v: Repr[_]): Unit
+  }
+
+  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](null)
+      @inline def length(v: Repr[_]): Int = v.length
+      @inline def decLength(v: Repr[_]): Unit =
+        v.length = v.length - 1
+      @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 copyFrom[E](v: Repr[E], from: Int): Repr[E] = v.jsSlice(from - 1)
+      @inline def clear(v: Repr[_]): Unit =
+        v.length = 1
+    }
+
+    /* We store the effective length in the index 0 of the array,
+     * which is unused both in JSArrayImpl and in this impl.
+     */
+    object JArrayImpl extends InnerArrayImpl {
+      type Repr[E] = Array[AnyRef]
+
+      @inline def make[E](initialCapacity: Int): Repr[E] = {
+        val v = new Array[AnyRef](initialCapacity)
+        v(0) = 1.asInstanceOf[AnyRef]
+        v
+      }
+      @inline def length(v: Repr[_]): Int = v(0).asInstanceOf[Int]
+      @inline def decLength(v: Repr[_]): Unit = {
+        val newLength = length(v) - 1
+        v(0) = newLength.asInstanceOf[AnyRef]
+        v(newLength) = null // free reference for GC
+      }
+      @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] = {
+        val l = length(v)
+        val minCapacity = l + 1
+        val newArr =
+          if (v.length < minCapacity)
+            Arrays.copyOf(v, roundUpToPowerOfTwo(minCapacity))
+          else v
+        newArr(l) = e.asInstanceOf[AnyRef]
+        newArr(0) = (l + 1).asInstanceOf[AnyRef]
+        newArr
+      }
+      @inline def copyFrom[E](v: Repr[E], from: Int): Repr[E] = {
+        val elemLength = length(v) - from
+        val newArr = new Array[AnyRef](elemLength + 1)
+        newArr(0) = (elemLength + 1).asInstanceOf[AnyRef]
+        System.arraycopy(v, from, newArr, 1, elemLength)
+        newArr
+      }
+      @inline def clear(v: Repr[_]): Unit = {
+        Arrays.fill(v, null)
+        v(0) = 1.asInstanceOf[AnyRef]
+      }
+    }
+  }
+
 }
