From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Thibaud Michaud <thibaudm@chromium.org>
Date: Fri, 21 Jun 2024 16:31:15 +0200
Subject: Fix scanning of wasm-to-js params

Wasm-to-js wrappers are sometimes compiled as on-heap Code objects, for
example when tiering-up from a WasmFuncRef call origin. The frames of
these functions are mapped to a subclass of TypedFrame, however
TypedFrame::Iterate() only supports iterating the generic wasm-to-js
wrapper.

Add support for iterating the tagged parameters of optimized wasm-to-js
wrappers in TypedFrame::Iterate. For this we also add two 16-bit fields
in the Code object to encode the incoming tagged parameter region, which
we would normally find in the WasmCode data.

R=jkummerow@chromium.org

Fixed: 346597059
Change-Id: I425619fca86c38f91f1ca9cbeb70e7b5a7b2d6c1
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/5639725
Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#94589}

diff --git a/src/diagnostics/objects-printer.cc b/src/diagnostics/objects-printer.cc
index 4246209b6974d5cb575382c71384fda1c5b6eba0..7759fb1391e6019bf9eaaca0b6e9191f19f64838 100644
--- a/src/diagnostics/objects-printer.cc
+++ b/src/diagnostics/objects-printer.cc
@@ -2048,7 +2048,14 @@ void Code::CodePrint(std::ostream& os, const char* name, Address current_pc) {
   os << "\n - instruction_size: " << instruction_size();
   os << "\n - metadata_size: " << metadata_size();
 
-  os << "\n - inlined_bytecode_size: " << inlined_bytecode_size();
+  if (kind() != CodeKind::WASM_TO_JS_FUNCTION) {
+    os << "\n - inlined_bytecode_size: " << inlined_bytecode_size();
+  } else {
+    os << "\n - wasm_js_tagged_parameter_count: "
+       << wasm_js_tagged_parameter_count();
+    os << "\n - wasm_js_first_tagged_parameter: "
+       << wasm_js_first_tagged_parameter();
+  }
   os << "\n - osr_offset: " << osr_offset();
   os << "\n - handler_table_offset: " << handler_table_offset();
   os << "\n - unwinding_info_offset: " << unwinding_info_offset();
diff --git a/src/execution/frames.cc b/src/execution/frames.cc
index 7b8507384ab5fd8009d2cf2cb2afc764a5332d61..efde8eb70717e785747f335c1b817cdab525ffdc 100644
--- a/src/execution/frames.cc
+++ b/src/execution/frames.cc
@@ -1536,7 +1536,7 @@ void WasmFrame::Iterate(RootVisitor* v) const {
                        frame_header_limit);
 }
 
-void TypedFrame::IterateParamsOfWasmToJSWrapper(RootVisitor* v) const {
+void TypedFrame::IterateParamsOfGenericWasmToJSWrapper(RootVisitor* v) const {
   Tagged<Object> maybe_signature = Tagged<Object>(
       Memory<Address>(fp() + WasmToJSWrapperConstants::kSignatureOffset));
   if (IsSmi(maybe_signature)) {
@@ -1652,6 +1652,18 @@ void TypedFrame::IterateParamsOfWasmToJSWrapper(RootVisitor* v) const {
     }
   }
 }
+
+void TypedFrame::IterateParamsOfOptimizedWasmToJSWrapper(RootVisitor* v) const {
+  Tagged<GcSafeCode> code = GcSafeLookupCode();
+  if (code->wasm_js_tagged_parameter_count() > 0) {
+    FullObjectSlot tagged_parameter_base(&Memory<Address>(caller_sp()));
+    tagged_parameter_base += code->wasm_js_first_tagged_parameter();
+    FullObjectSlot tagged_parameter_limit =
+        tagged_parameter_base + code->wasm_js_tagged_parameter_count();
+    v->VisitRootPointers(Root::kStackRoots, nullptr, tagged_parameter_base,
+                         tagged_parameter_limit);
+  }
+}
 #endif  // V8_ENABLE_WEBASSEMBLY
 
 void TypedFrame::Iterate(RootVisitor* v) const {
@@ -1683,10 +1695,13 @@ void TypedFrame::Iterate(RootVisitor* v) const {
   CHECK(entry->code.has_value());
   Tagged<GcSafeCode> code = entry->code.value();
 #if V8_ENABLE_WEBASSEMBLY
-  bool is_wasm_to_js =
+  bool is_generic_wasm_to_js =
       code->is_builtin() && code->builtin_id() == Builtin::kWasmToJsWrapperCSA;
-  if (is_wasm_to_js) {
-    IterateParamsOfWasmToJSWrapper(v);
+  bool is_optimized_wasm_to_js = this->type() == WASM_TO_JS_FUNCTION;
+  if (is_generic_wasm_to_js) {
+    IterateParamsOfGenericWasmToJSWrapper(v);
+  } else if (is_optimized_wasm_to_js) {
+    IterateParamsOfOptimizedWasmToJSWrapper(v);
   }
 #endif  // V8_ENABLE_WEBASSEMBLY
   DCHECK(code->is_turbofanned());
@@ -1719,10 +1734,13 @@ void TypedFrame::Iterate(RootVisitor* v) const {
   // wrapper switched to before pushing the outgoing stack parameters and
   // calling the target. It marks the limit of the stack param area, and is
   // distinct from the beginning of the spill area.
-  Address central_stack_sp =
-      Memory<Address>(fp() + WasmToJSWrapperConstants::kCentralStackSPOffset);
+  int central_stack_sp_offset =
+      is_generic_wasm_to_js
+          ? WasmToJSWrapperConstants::kCentralStackSPOffset
+          : WasmImportWrapperFrameConstants::kCentralStackSPOffset;
+  Address central_stack_sp = Memory<Address>(fp() + central_stack_sp_offset);
   FullObjectSlot parameters_limit(
-      v8_flags.experimental_wasm_stack_switching && is_wasm_to_js &&
+      v8_flags.experimental_wasm_stack_switching && is_generic_wasm_to_js &&
               central_stack_sp != kNullAddress
           ? central_stack_sp
           : frame_header_base.address() - spill_slots_size);
diff --git a/src/execution/frames.h b/src/execution/frames.h
index 428e50c32a7eca9ad579b2c634607f4be5d1734e..a578a87d67905b19f045465d04532bbb1706e585 100644
--- a/src/execution/frames.h
+++ b/src/execution/frames.h
@@ -626,7 +626,8 @@ class TypedFrame : public CommonFrame {
   Tagged<HeapObject> unchecked_code() const override { return {}; }
   void Iterate(RootVisitor* v) const override;
 
-  void IterateParamsOfWasmToJSWrapper(RootVisitor* v) const;
+  void IterateParamsOfGenericWasmToJSWrapper(RootVisitor* v) const;
+  void IterateParamsOfOptimizedWasmToJSWrapper(RootVisitor* v) const;
 
  protected:
   inline explicit TypedFrame(StackFrameIteratorBase* iterator);
diff --git a/src/objects/code-inl.h b/src/objects/code-inl.h
index da3c0bd4fe8b9c593ecc25d9a3c277160b9df95b..4f6082f171948385c0fe9feeba51cc51b6fd9674 100644
--- a/src/objects/code-inl.h
+++ b/src/objects/code-inl.h
@@ -48,6 +48,8 @@ GCSAFE_CODE_FWD_ACCESSOR(bool, has_tagged_outgoing_params)
 GCSAFE_CODE_FWD_ACCESSOR(bool, marked_for_deoptimization)
 GCSAFE_CODE_FWD_ACCESSOR(Tagged<Object>, raw_instruction_stream)
 GCSAFE_CODE_FWD_ACCESSOR(int, stack_slots)
+GCSAFE_CODE_FWD_ACCESSOR(uint16_t, wasm_js_tagged_parameter_count)
+GCSAFE_CODE_FWD_ACCESSOR(uint16_t, wasm_js_first_tagged_parameter)
 GCSAFE_CODE_FWD_ACCESSOR(Address, constant_pool)
 GCSAFE_CODE_FWD_ACCESSOR(Address, safepoint_table_address)
 #undef GCSAFE_CODE_FWD_ACCESSOR
@@ -371,6 +373,31 @@ void Code::set_inlined_bytecode_size(unsigned size) {
   RELAXED_WRITE_UINT_FIELD(*this, kInlinedBytecodeSizeOffset, size);
 }
 
+// For optimized on-heap wasm-js wrappers, we repurpose the (otherwise unused)
+// 32-bit InlinedBytecodeSize field to encode two 16 values needed for scanning
+// the frame: the count and starting offset of incoming tagged parameters.
+// TODO(wasm): Eventually the wrappers should be managed off-heap by the wasm
+// engine. Remove these accessors when that is the case.
+void Code::set_wasm_js_tagged_parameter_count(uint16_t count) {
+  DCHECK_EQ(kind(), CodeKind::WASM_TO_JS_FUNCTION);
+  RELAXED_WRITE_UINT16_FIELD(*this, kInlinedBytecodeSizeOffset, count);
+}
+
+uint16_t Code::wasm_js_tagged_parameter_count() const {
+  DCHECK_EQ(kind(), CodeKind::WASM_TO_JS_FUNCTION);
+  return RELAXED_READ_UINT16_FIELD(*this, kInlinedBytecodeSizeOffset);
+}
+
+void Code::set_wasm_js_first_tagged_parameter(uint16_t count) {
+  DCHECK_EQ(kind(), CodeKind::WASM_TO_JS_FUNCTION);
+  RELAXED_WRITE_UINT16_FIELD(*this, kInlinedBytecodeSizeOffset + 2, count);
+}
+
+uint16_t Code::wasm_js_first_tagged_parameter() const {
+  DCHECK_EQ(kind(), CodeKind::WASM_TO_JS_FUNCTION);
+  return RELAXED_READ_UINT16_FIELD(*this, kInlinedBytecodeSizeOffset + 2);
+}
+
 BytecodeOffset Code::osr_offset() const {
   return BytecodeOffset(RELAXED_READ_INT32_FIELD(*this, kOsrOffsetOffset));
 }
diff --git a/src/objects/code.h b/src/objects/code.h
index c6b0eda42849ac46498bb292566b6f74a6716a25..b500da60feb9c60ef77add1a8e8102bdb0b3829a 100644
--- a/src/objects/code.h
+++ b/src/objects/code.h
@@ -124,6 +124,25 @@ class Code : public ExposedTrustedObject {
   // [deoptimization_data]: Array containing data for deopt for non-baseline
   // code.
   DECL_ACCESSORS(deoptimization_data, Tagged<FixedArray>)
+  // [parameter_count]: The number of formal parameters, including the
+  // receiver. Currently only available for optimized functions.
+  // TODO(saelo): make this always available. This is just a matter of figuring
+  // out how to obtain the parameter count during code generation when no
+  // BytecodeArray is available from which it can be copied.
+  DECL_PRIMITIVE_ACCESSORS(parameter_count, uint16_t)
+  inline uint16_t parameter_count_without_receiver() const;
+  DECL_PRIMITIVE_ACCESSORS(wasm_js_tagged_parameter_count, uint16_t)
+  DECL_PRIMITIVE_ACCESSORS(wasm_js_first_tagged_parameter, uint16_t)
+
+  // Whether this type of Code uses deoptimization data, in which case the
+  // deoptimization_data field will be populated.
+  inline bool uses_deoptimization_data() const;
+
+  // If neither deoptimization data nor bytecode/interpreter data are used
+  // (e.g. for builtin code), the respective field will contain Smi::zero().
+  inline void clear_deoptimization_data_and_interpreter_data();
+  inline bool has_deoptimization_data_or_interpreter_data() const;
+
   // [bytecode_or_interpreter_data]: BytecodeArray or InterpreterData for
   // baseline code.
   // As BytecodeArrays are located in trusted space, but InterpreterData
@@ -482,6 +501,10 @@ class GcSafeCode : public HeapObject {
   inline bool CanDeoptAt(Isolate* isolate, Address pc) const;
   inline Tagged<Object> raw_instruction_stream(
       PtrComprCageBase code_cage_base) const;
+  // The two following accessors repurpose the InlinedBytecodeSize field, see
+  // comment in code-inl.h.
+  inline uint16_t wasm_js_tagged_parameter_count() const;
+  inline uint16_t wasm_js_first_tagged_parameter() const;
 
  private:
   OBJECT_CONSTRUCTORS(GcSafeCode, HeapObject);
