Opt/Wasm: Optimize reference equality as ref.is_null or ref.eq.

sjrd · sjrd · commit 367d63ab5f12 · 2024-06-24T15:20:00.000+02:00
Profiling showed that using `jsStrictEquals`/`is` is really
prohibitive.
diff --git a/linker/shared/src/main/scala/org/scalajs/linker/backend/wasmemitter/FunctionEmitter.scala b/linker/shared/src/main/scala/org/scalajs/linker/backend/wasmemitter/FunctionEmitter.scala
@@ -1347,19 +1347,53 @@ private class FunctionEmitter private (
     val BinaryOp(op, lhs, rhs) = tree
     assert(op == === || op == !==)
 
-    // TODO Optimize this when the operands have a better type than `any`
-
-    genTree(lhs, AnyType)
-    genTree(rhs, AnyType)
+    def maybeGenInvert(): Unit = {
+      if (op == BinaryOp.!==)
+        genBooleanNot(fb)
+    }
 
-    markPosition(tree)
+    val lhsType = lhs.tpe
+    val rhsType = rhs.tpe
 
-    fb += wa.Call(genFunctionID.is)
+    if (lhsType == NothingType) {
+      genTree(lhs, NothingType)
+      NothingType
+    } else if (rhsType == NothingType) {
+      genTree(lhs, NoType)
+      genTree(rhs, NothingType)
+      NothingType
+    } else if (rhsType == NullType) {
+      /* Note that the optimizer normalizes Literals on the right of `===`,
+       * so testing for the `lhsType == NullType` is not as useful.
+       */
+      genTree(lhs, AnyType)
+      genTree(rhs, NoType) // no-op if it is actually a Null() literal
+      markPosition(tree)
+      fb += wa.RefIsNull
+      maybeGenInvert()
+      BooleanType
+    } else {
+      val lhsWasmType = transformSingleType(lhsType)
+      val rhsWasmType = transformSingleType(rhsType)
+
+      (lhsWasmType, rhsWasmType) match {
+        case (watpe.RefType(_, lhsHeapType), watpe.RefType(_, rhsHeapType))
+            if lhsHeapType != watpe.HeapType.Any && rhsHeapType != watpe.HeapType.Any =>
+          genTree(lhs, lhsType)
+          genTree(rhs, rhsType)
+          markPosition(tree)
+          fb += wa.RefEq
 
-    if (op == !==)
-      genBooleanNot(fb)
+        case _ =>
+          genTree(lhs, AnyType)
+          genTree(rhs, AnyType)
+          markPosition(tree)
+          fb += wa.Call(genFunctionID.is)
+      }
 
-    BooleanType
+      maybeGenInvert()
+      BooleanType
+    }
   }
 
   private def getElementaryBinaryOpInstr(op: BinaryOp.Code): wa.Instr = {
diff --git a/linker/shared/src/main/scala/org/scalajs/linker/frontend/optimizer/OptimizerCore.scala b/linker/shared/src/main/scala/org/scalajs/linker/frontend/optimizer/OptimizerCore.scala
@@ -542,10 +542,14 @@ private[optimizer] abstract class OptimizerCore(
           pretransformExpr(expr) { texpr =>
             val result = {
               if (isSubtype(texpr.tpe.base, testType)) {
-                if (texpr.tpe.isNullable)
-                  JSBinaryOp(JSBinaryOp.!==, finishTransformExpr(texpr), Null())
-                else
+                if (texpr.tpe.isNullable) {
+                  if (isWasm)
+                    BinaryOp(BinaryOp.!==, finishTransformExpr(texpr), Null())
+                  else
+                    JSBinaryOp(JSBinaryOp.!==, finishTransformExpr(texpr), Null())
+                } else {
                   Block(finishTransformStat(texpr), BooleanLiteral(true))
+                }
               } else {
                 if (texpr.tpe.isExact)
                   Block(finishTransformStat(texpr), BooleanLiteral(false))
@@ -3491,6 +3495,13 @@ private[optimizer] abstract class OptimizerCore(
                 if lhsIdent2 == lhsIdent && rhsIdent2 == rhsIdent =>
               elsep
 
+            // Same, but on Wasm, which does not turn BinaryOp.=== into JSBinaryOp.===
+            case (BinaryOp(BinaryOp.===, VarRef(lhsIdent), Null()),
+                BinaryOp(BinaryOp.===, VarRef(rhsIdent), Null()),
+                BinaryOp(BinaryOp.===, VarRef(lhsIdent2), VarRef(rhsIdent2)))
+                if lhsIdent2 == lhsIdent && rhsIdent2 == rhsIdent =>
+              elsep
+
             // Example: (x > y) || (x == y)  ->  (x >= y)
             case (BinaryOp(op1 @ (Int_== | Int_!= | Int_< | Int_<= | Int_> | Int_>=), l1, r1),
                   BooleanLiteral(true),
@@ -4077,7 +4088,9 @@ private[optimizer] abstract class OptimizerCore(
           case (PreTransLit(l), PreTransLit(r)) =>
             val isSame = literal_===(l, r, isJSStrictEq = false)
             PreTransLit(BooleanLiteral(if (op == ===) isSame else !isSame))
-          case _ if canOptimizeAsJSStrictEq(lhs.tpe, rhs.tpe) =>
+          case (PreTransLit(_), _) =>
+            foldBinaryOp(op, rhs, lhs)
+          case _ if !isWasm && canOptimizeAsJSStrictEq(lhs.tpe, rhs.tpe) =>
             foldJSBinaryOp(
                 if (op == ===) JSBinaryOp.=== else JSBinaryOp.!==,
                 lhs, rhs)
@@ -5027,7 +5040,7 @@ private[optimizer] abstract class OptimizerCore(
                 finishTransformStat(lhs),
                 BooleanLiteral(!positive)).toPreTransform
 
-          case (PreTransLit(_), _) => foldBinaryOp(op, rhs, lhs)
+          case (PreTransLit(_), _) => foldJSBinaryOp(op, rhs, lhs)
           case _                   => default
         }
 
