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.

< 10000 div data-view-component="true" class="comment-reactions js-reactions-container js-reaction-buttons-container social-reactions reactions-container d-flex">

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

10000
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.

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

sjrd · 2025-06-16T13:36:45Z

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

I managed to get this to happen by changing genApplyTypeApply in GenJSCode.scala to

    private def genApplyTypeApply(tree: Apply, isStat: Boolean): js.Tree = {
      implicit val pos = tree.pos
      val Apply(TypeApply(fun @ Select(obj, _), targs), args) = tree
      val sym = fun.symbol

      sym match {
        case Object_isInstanceOf =>
          genIsAsInstanceOf(obj, targs, cast = false)
        case Object_asInstanceOf =>
          val targetTpe = targs.head.tpe
          obj match {
            case Apply(fun, List(cond, thenp, elsep))
                if fun.symbol == jsDefinitions.LinkingInfo_linkTimeIf &&
                thenp.tpe <:< targetTpe && elsep.tpe <:< targetTpe =>
              val genObj = genExpr(obj).asInstanceOf[js.LinkTimeIf]
              js.LinkTimeIf(genObj.cond, genObj.thenp, genObj.elsep)(toIRType(targetTpe))(genObj.pos)
            case _ =>
              genIsAsInstanceOf(obj, targs, cast = true)
          }
        case Object_synchronized =>
          genSynchronized(obj, args.head, isStat)
        case _ =>
          abort("Unexpected type application " + fun +
              "[sym: " + sym.fullName + "]" + " in: " + tree)
      }
    }

That removes a spurious js.AsInstanceOf around the js.LinkTimeIf, which prevents the module aliasing analysis done in IncOptimizer.scala.

(It's a rough draft; it should be cleaned up and tested in order to accept that change into main.)

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

tanishiking · 2025-07-15T13:08:54Z

I'm getting back on this. Sorry for the delayed response, and thank you very much for addressing the issue about spurious js.AsInstanceOfs!

I found that, another significant performance degradation in add is caused by fixUp implementation. It always traversed all the way up to the root node even after finding a suitable location. fixed in dbfd7a7
(Interestingly, it wasn't a problem without innerImpl, because JS will be optimized to stop traversing when fixUp found the location).

Fixing the fixUp and applying your change, we could optimize Wasm without degrade the JS performance 🎉
Now there's a small diff between main and HEAD in JS https://gist.github.com/tanishiking/5f7b24a1e4e3e724b85d5e60415bf1e4

Benchmark(ms)	JS (HEAD)	JS (main)	Wasm (HEAD)	Wasm (main)
add_N	65.82	70.81	232.06	705.57
poll_N_FromFull	1557.52	1537.54	4025.59	8748.01
peek_Repeatedly	83.25	81.29	289.83	1077.45
build_FromCollection_N	77.42	82.55	285.58	721.59
pipelined_AddPoll	1724.26	1529.60	3539.56	8522.60

Your change in #5168 (comment) looks good to me, maybe we can make another PR with tests?

sjrd · 2025-07-15T15:44:39Z

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

        }
+        loop(parent)


This is where the loop(parent) was moved to the wrong place. So the "fix fixUp" commit should be squashed with this one. It was correct in the original code. ;)

Whoa..., I didn't realize I had made this change myself 🤣

When B<:A and C<:A, `linkTimeIf[A](cond){ B }{ C }` would generate an `.asInstanceOf[A]` that casts the resulting value to a supertype. However, this cast is redundant, as the linker will prune one branch at linktime, and the remaining expression is already known to be a compatible type. This commit eliminates the surrounding `.asInstanceOf` around the `linkTimeIf[T]` when both `thenp` and `elsep` are surely subtypes of `T`. context: scala-js#5168 Co-authored-by: Sébastien Doeraene <sjrdoeraene@gmail.com>

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.

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

Co-authored-by: Sébastien Doeraene <sjrdoeraene@gmail.com>

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 force-pushed the priorityqueue-wasm branch from da5c453 to 83aa955 Compare June 12, 2025 04:16

tanishiking force-pushed the priorityqueue-wasm branch 4 times, most recently from 0832ba7 to 144bca0 Compare June 16, 2025 12:33

Add java.lang.Utils.roundUpToPowerOfTwo utility function

04efdb2

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 144bca0 to a219596 Compare July 15, 2025 13:04

sjrd reviewed Jul 15, 2025

View reviewed changes

tanishiking mentioned this pull request Jul 16, 2025

Removes spurious js.AsInstanceOf around the js.LinkTimeIf #5210

Open

tanishiking and others added 3 commits July 16, 2025 15:45

Move innerArrayImpl to companion object

281cbde

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

Removes spurious js.AsInstanceOf around the js.LinkTimeIf

3f504eb

Co-authored-by: Sébastien Doeraene <sjrdoeraene@gmail.com>

tanishiking force-pushed the priorityqueue-wasm branch from a219596 to 3f504eb Compare July 16, 2025 06:46

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

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!

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!

Uh oh!

null checks on innnerImpl

It's still slower

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!