Bug 1929342 - Inline wasm:js-string length. r=rhunt

Differential Revision: https://phabricator.services.mozilla.com/D228071

js/src/jit/CodeGenerator.cpp

@@ -22018,6 +22018,19 @@ void CodeGenerator::visitWasmTrapIfAnyRefIsNotJSString(
   masm.bind(&isJSString);
 }
 
+void CodeGenerator::visitWasmAnyRefJSStringLength(
+    LWasmAnyRefJSStringLength* lir) {
+  Register input = ToRegister(lir->input());
+  Register output = ToRegister(lir->output());
+  Register temp = ToRegister(lir->temp0());
+  Label isJSString;
+  masm.branchWasmAnyRefIsJSString(true, input, temp, &isJSString);
+  masm.wasmTrap(lir->mir()->trap(), lir->mir()->bytecodeOffset());
+  masm.bind(&isJSString);
+  masm.untagWasmAnyRef(input, temp, wasm::AnyRefTag::String);
+  masm.loadStringLength(temp, output);
+}
+
 void CodeGenerator::visitWasmNewI31Ref(LWasmNewI31Ref* lir) {
   if (lir->value()->isConstant()) {
     // i31ref are often created with constants. If that's the case we will

js/src/jit/LIROps.yaml

@@ -4294,6 +4294,13 @@
     input: WordSized
   num_temps: 1
 
+- name: WasmAnyRefJSStringLength
+  mir_op: true
+  result_type: WordSized
+  operands:
+    input: WordSized
+  num_temps: 1
+
 - name: WasmNewI31Ref
   mir_op: true
   result_type: WordSized

js/src/jit/Lowering.cpp

@@ -3720,6 +3720,13 @@ void LIRGenerator::visitWasmTrapIfAnyRefIsNotJSString(
   add(lir, ins);
 }
 
+void LIRGenerator::visitWasmAnyRefJSStringLength(
+    MWasmAnyRefJSStringLength* ins) {
+  LWasmAnyRefJSStringLength* lir = new (alloc())
+      LWasmAnyRefJSStringLength(useRegisterAtStart(ins->input()), temp());
+  define(lir, ins);
+}
+
 void LIRGenerator::visitWasmNewI31Ref(MWasmNewI31Ref* ins) {
   // If it's a constant, it will be put directly into the register.
   LWasmNewI31Ref* lir =

js/src/jit/MIROps.yaml

@@ -679,6 +679,19 @@
   guard: true
   clone: true
 
+- name: WasmAnyRefJSStringLength
+  operands:
+    input: WasmAnyRef
+  arguments:
+    trap: wasm::Trap
+    bytecodeOffset: wasm::BytecodeOffset
+  type_policy: none
+  result_type: Int32
+  congruent_to: if_operands_equal
+  alias_set: none
+  guard: true
+  clone: true
+
 - name: WasmNewI31Ref
   operands:
     input: Int32

js/src/jit/MacroAssembler.cpp

@@ -6973,6 +6973,12 @@ void MacroAssembler::extractWasmAnyRefTag(Register src, Register dest) {
   andPtr(Imm32(int32_t(wasm::AnyRef::TagMask)), dest);
 }
 
+void MacroAssembler::untagWasmAnyRef(Register src, Register dest,
+                                     wasm::AnyRefTag tag) {
+  MOZ_ASSERT(tag != wasm::AnyRefTag::ObjectOrNull, "No untagging needed");
+  computeEffectiveAddress(Address(src, -int32_t(tag)), dest);
+}
+
 void MacroAssembler::branchWasmAnyRefIsNull(bool isNull, Register src,
                                             Label* label) {
   branchTestPtr(isNull ? Assembler::Zero : Assembler::NonZero, src, src, label);
@@ -7579,8 +7585,7 @@ void MacroAssembler::convertWasmAnyRefToValue(Register instance, Register src,
                 &isObjectOrNull);
 
   // If we're not i31, object, or null, we must be a string
-  rshiftPtr(Imm32(wasm::AnyRef::TagShift), src);
-  lshiftPtr(Imm32(wasm::AnyRef::TagShift), src);
+  untagWasmAnyRef(src, src, wasm::AnyRefTag::String);
   moveValue(TypedOrValueRegister(MIRType::String, AnyRegister(src)), dst);
   jump(&done);
 

js/src/jit/MacroAssembler.h

@@ -49,6 +49,7 @@
 #include "vm/FunctionFlags.h"
 #include "vm/Opcodes.h"
 #include "vm/RealmFuses.h"
+#include "wasm/WasmAnyRef.h"
 
 // [SMDOC] MacroAssembler multi-platform overview
 //
@@ -4149,6 +4150,9 @@ class MacroAssembler : public MacroAssemblerSpecific {
   // Extract the tag of wasm anyref `src`.
   void extractWasmAnyRefTag(Register src, Register dest);
 
+  // Remove the known tag of wasm anyref `src`.
+  void untagWasmAnyRef(Register src, Register dest, wasm::AnyRefTag tag);
+
   // Branch if the wasm anyref `src` is or is not the null value.
   void branchWasmAnyRefIsNull(bool isNull, Register src, Label* label);
   // Branch if the wasm anyref `src` is or is not an I31.

js/src/wasm/WasmBuiltinModule.yaml

@@ -338,6 +338,7 @@
     name: StringLength
     type: Args_Int32_GeneralGeneral
   entry: Instance::stringLength
+  inline_op: StringLength
   export: length
   params:
     - 'externref'

js/src/wasm/WasmConstants.h

@@ -985,6 +985,7 @@ enum class BuiltinInlineOp {
 
   StringCast,
   StringTest,
+  StringLength,
 
   // Op limit.
   Limit

js/src/wasm/WasmIonCompile.cpp

@@ -3245,6 +3245,16 @@ class FunctionCompiler {
     return ins;
   }
 
+  [[nodiscard]] MDefinition* stringLength(MDefinition* string) {
+    auto* ins = MWasmAnyRefJSStringLength::New(
+        alloc(), string, wasm::Trap::BadCast, bytecodeOffset());
+    if (!ins) {
+      return nullptr;
+    }
+    curBlock_->add(ins);
+    return ins;
+  }
+
   [[nodiscard]] bool dispatchInlineBuiltinModuleFunc(
       const BuiltinModuleFunc& builtinModuleFunc, const DefVector& params) {
     BuiltinInlineOp inlineOp = builtinModuleFunc.inlineOp();
@@ -3270,6 +3280,16 @@ class FunctionCompiler {
         iter().setResult(test);
         return true;
       }
+      case BuiltinInlineOp::StringLength: {
+        MOZ_ASSERT(params.length() == 1);
+        MDefinition* string = params[0];
+        MDefinition* length = stringLength(string);
+        if (!length) {
+          return false;
+        }
+        iter().setResult(length);
+        return true;
+      }
       case BuiltinInlineOp::None:
       case BuiltinInlineOp::Limit:
         break;