Wasm: Implement PriorityQueue without js.Array in Wasm backend #5168

tanishiking · 2025-05-15T02:36:34Z

Simliar to ju.ArrayList, this commit adds a Wasm version of internal storage implementation to remove JS interop for better performance.

Branched from #5164

This implementation uses the same technique as parseFloat to simplify having the two implementation for different platforms. see the conversation #5164 (comment)
Also, this is a draft PR for now, because without linkTimeIf it seems that the optimizer fails to de-virtualize the function calls of methods on InnerArrayImpl. Waiting for it
5e842d8

sjrd · 2025-05-16T05:53:38Z

javalib/src/main/scala/java/util/PriorityQueue.scala

+
+      // The index 0 is not used; the root is at index 1.
+      // This is standard practice in binary heaps, to simplify arithmetics.
+      private var _size = 1


Something went wrong during the refactoring, because this state variable was moved to the global impl object. 😅

Oops, thanks, moved back to the PriorityQueue

sjrd

Partial review. We need to address the _size issue before going further.

sjrd · 2025-06-06T09:23:18Z

javalib/src/main/scala/java/util/PriorityQueue.scala

+
+  private sealed abstract class InnerArrayImpl {


This should be in the companion object of PriorityQueue, so that instances do not depend on the enclosing PriorityQueue object.

Suggested change

private sealed abstract class InnerArrayImpl {

}

object PriorityQueue {

private sealed abstract class InnerArrayImpl {

Ah! I now understand that they actually refer to the enclosing object for _size. 😱 That's super confusing: the dependencies on inner are made explicit (params), but those on _size are implicit (scope).

We should also pass _size explicitly. But then that poses an issue for methods that need to return a new Repr and a new size. 😖

Here is a crazy idea: we have an unused slot at index 0 of the array. What if we use it to store the size on Wasm? It would be boxed, but unless it exceeds 2^30, it will fit in a (ref i31), so no JS interop should occur. (and if it does exceed 2^30, I don't think that's going to be the bottleneck :p).

This way we still have a single Repr that embodies both the size and the capacity, on Wasm as well as on JS.

Thanks! I think it's reasonable technique to re-use the free space to store size information 👍
Here I make a prototype implementation for moving InnerArrayImpl to companion object, and store size to array(0) 0832ba7

sjrd · 2025-06-06T09:25:12Z

javalib/src/main/scala/java/util/PriorityQueue.scala

+   * the function calls.
+   */
+
+  private val innerImpl: InnerArrayImpl =


It is unfortunate to store this field inside instances of PriorityQueue. I think this can be avoided by putting this private val in the companion object of PriorityQueue.

If my prediction is correct, that companion will entirely disappear after optimizations.

sjrd · 2025-06-06T09:26:15Z

javalib/src/main/scala/java/util/PriorityQueue.scala

+import java.lang.Utils.roundUpToPowerOfTwo
+
 import scala.annotation.tailrec


Imports of scala.annotation.* always come first (after scala.language.* imports, if any):

Suggested change

import java.lang.Utils.roundUpToPowerOfTwo

import scala.annotation.tailrec

import scala.annotation.tailrec

import java.lang.Utils.roundUpToPowerOfTwo

javalib/src/main/scala/java/util/PriorityQueue.scala

sjrd · 2025-06-06T09:35:16Z

javalib/src/main/scala/java/util/PriorityQueue.scala

  // The index 0 is not used; the root is at index 1.
  // This is standard practice in binary heaps, to simplify arithmetics.


We should keep these two lines as comments on private var inner. They apply to JS as well.

Also, that means that the notion of "capacity" is a skewed in this class. The array should have 1 more elements than the capacity, for "capacity" being the external notion of capacity.

Put differently, it is probably easier to think of "capacity" as the internal notion of capacity, which is 1 more than the "public capacity". The public capacity only shows up in def this(initialCapacity) and def this(c: Collection).

sjrd · 2025-06-06T09:35:49Z

javalib/src/main/scala/java/util/PriorityQueue.scala

+      {
+        if (initialCapacity < 1)
+          throw new IllegalArgumentException
+        initialCapacity


Suggested change

initialCapacity

initialCapacity + 1 // index 0 is unused

sjrd · 2025-06-06T09:36:24Z

javalib/src/main/scala/java/util/PriorityQueue.scala

@@ -47,47 +62,74 @@ class PriorityQueue[E] private (
        NaturalComparator.select(c.comparator().asInstanceOf[Comparator[_ >: E]])
      case _ =>
        NaturalComparator
-    }, internal = true)
+    }, internal = true, roundUpToPowerOfTwo(c.size()))


Suggested change

}, internal = true, roundUpToPowerOfTwo(c.size()))

}, internal = true, roundUpToPowerOfTwo(c.size() + 1)) // index 0 is unused

sjrd · 2025-06-06T09:36:48Z

javalib/src/main/scala/java/util/PriorityQueue.scala

    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()))


Suggested change

this(c.comp.asInstanceOf[Comparator[_ >: E]], internal = true, roundUpToPowerOfTwo(c.size()))

this(c.comp.asInstanceOf[Comparator[_ >: E]], internal = true,

roundUpToPowerOfTwo(c.size() + 1)) // index 0 is unused

sjrd · 2025-06-06T09:37:42Z

javalib/src/main/scala/java/util/PriorityQueue.scala

    addAll(c)
  }

  def this(sortedSet: SortedSet[_ <: E]) = {
    this(NaturalComparator.select(
        sortedSet.comparator().asInstanceOf[Comparator[_ >: E]]),
-        internal = true)
+        internal = true,
+        roundUpToPowerOfTwo(sortedSet.size()))


Suggested change

roundUpToPowerOfTwo(sortedSet.size()))

roundUpToPowerOfTwo(sortedSet.size() + 1)) // index 0 is unused

sjrd · 2025-06-06T09:46:29Z

javalib/src/main/scala/java/util/PriorityQueue.scala

+    if (LinkingInfo.isWebAssembly) {
+      val minCapacity = innerImpl.length(inner) + 1
+      if (innerImpl.capacity(inner) < minCapacity)
+        inner = innerImpl.resized(inner, minCapacity)
+    }


This can be moved to JArrayImpl.push.

The use-case of this function is for ensuring the size of internal buffers (e.g. internal Array of `ju.ArrayList`).

Current js.Array-based PriorityQueue implementation on Wasm requires JS-interop for every operation, and JS-interop is very slow. We use a `linkTimeIf` to select a `scala.Array`-based implementation of internal data structure for PriorityQueue, based on wether it is on JS or WebAssembly.

tanishiking · 2025-06-12T04:32:28Z

I measured performance of PriorityQueue using a simple priority-queue micro benchmark with fastLinkJS. It does add/peek/poll 10000 numbers (in order from 0-10000, maybe inappropriate for priorityqueue.)

The performance change on wasm (red and purple bars) was okay: add/peek gets a bit slower (probably because of array expansion), but the overall performance with poll is much faster. 🎉

The problem is the execution performance on JS.
As you can see from the blue (main) and orange (this change) graphs, the performance on add/peek is about 10 times slower, and poll is also degraded by 2 times.

null checks on innnerImpl

One problem is that, with linktimeIf, the null check on innerImpl has inserted everywhere $n($m_ju_PriorityQueue$().ju_PriorityQueue$__f_java$util$PriorityQueue$$innerImpl); The diff between main vs this branch, generated from the Scala code below: https://gist.github.com/tanishiking/9d0e24de18e994cae5e8742aa9749d98#file-foo-diff

Interestingly, without linkeTimeIf on innerImpl private val innerImpl: InnerArrayImpl = InnerArrayImpl.JSArrayImpl (green bar), innerImpl gets inlined and there's no null check on it. diff: https://gist.github.com/tanishiking/9d0e24de18e994cae5e8742aa9749d98#file-foo2-diff

linktimeIf should generate the same code as the innerImpl = InnerArrayImpl.JSArrayImpl.

package helloworld

import java.util.PriorityQueue

object HelloWorld {
  def main(args: Array[String]): Unit = {
    val pq = new PriorityQueue[Int]()
    for (element <- List(1,2,3,4,5)) {
      pq.add(element)
    }
  }
}

It's still slower

With private val innerImpl: InnerArrayImpl = InnerArrayImpl.JSArrayImpl, the performance gets better (green bar), but it's still much slower than main.

I'm not sure why it's much slower with this diff 🤔 https://gist.github.com/tanishiking/9d0e24de18e994cae5e8742aa9749d98#file-foo2-diff

Also, store the size of array in the first element of Array

tanishiking mentioned this pull request May 15, 2025

Implement ArrayDeque using scala.Array. #5164

Merged

sjrd reviewed May 16, 2025

View reviewed changes

tanishiking force-pushed the priorityqueue-wasm branch 4 times, most recently from 6f8093e to ea2171f Compare May 17, 2025 04:36

tanishiking marked this pull request as ready for review May 17, 2025 05:09

tanishiking force-pushed the priorityqueue-wasm branch from ea2171f to bca2c21 Compare May 23, 2025 13:31

tanishiking force-pushed the priorityqueue-wasm branch from bca2c21 to da5c453 Compare June 6, 2025 05:49

tanishiking requested a review from sjrd June 6, 2025 05:50

sjrd requested changes Jun 6, 2025

View reviewed changes

tanishiking added 2 commits June 11, 2025 13:34

Add java.lang.Utils.roundUpToPowerOfTwo utility function

808a130

The use-case of this function is for ensuring the size of internal buffers (e.g. internal Array of `ju.ArrayList`).

tanishiking force-pushed the priorityqueue-wasm branch from da5c453 to 83aa955 Compare June 12, 2025 04:16

tanishiking force-pushed the priorityqueue-wasm branch 2 times, most recently from 471d8b6 to 2ee44af Compare June 12, 2025 07:55

Move innerArrayImpl to companion object

0832ba7

Also, store the size of array in the first element of Array

tanishiking force-pushed the priorityqueue-wasm branch from 2ee44af to 0832ba7 Compare June 12, 2025 09:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Wasm: Implement PriorityQueue without js.Array in Wasm backend #5168

Wasm: Implement PriorityQueue without js.Array in Wasm backend #5168

tanishiking commented May 15, 2025 •

edited

Loading

Uh oh!

sjrd May 16, 2025

Uh oh!

tanishiking May 17, 2025

Uh oh!

sjrd left a comment

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

tanishiking Jun 12, 2025 •

edited

Loading

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

sjrd Jun 6, 2025

Uh oh!

tanishiking commented Jun 12, 2025 •

edited

Loading

Uh oh!

Uh oh!

		import java.lang.Utils.roundUpToPowerOfTwo

		import scala.annotation.tailrec

		// The index 0 is not used; the root is at index 1.
		// This is standard practice in binary heaps, to simplify arithmetics.

	}, internal = true, roundUpToPowerOfTwo(c.size()))
	}, internal = true, roundUpToPowerOfTwo(c.size() + 1)) // index 0 is unused

	this(c.comp.asInstanceOf[Comparator[_ >: E]], internal = true, roundUpToPowerOfTwo(c.size()))
	this(c.comp.asInstanceOf[Comparator[_ >: E]], internal = true,
	roundUpToPowerOfTwo(c.size() + 1)) // index 0 is unused

	roundUpToPowerOfTwo(sortedSet.size()))
	roundUpToPowerOfTwo(sortedSet.size() + 1)) // index 0 is unused

Wasm: Implement PriorityQueue without js.Array in Wasm backend #5168

Are you sure you want to change the base?

Wasm: Implement PriorityQueue without js.Array in Wasm backend #5168

Conversation

tanishiking commented May 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sjrd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tanishiking Jun 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tanishiking commented Jun 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

null checks on innnerImpl

It's still slower

Uh oh!

Uh oh!

tanishiking commented May 15, 2025 •

edited

Loading

tanishiking Jun 12, 2025 •

edited

Loading

tanishiking commented Jun 12, 2025 •

edited

Loading