[wasm] Fix OSR on wasm calls

This fixes issues with replacing the return address of deeper (non-top)
wasm frames, i.e. frames which are at a call position. The replaced
address should also point after the call in the new code, so we don't
execute the same call again.

This is achieved by using slightly different encodings for breakpoint
positions and other (wasm instruction) positions. Breakpoints set
{is_instruction} to {false} in the source position table entry, whereas
usual wasm instruction set it to {true}.
Also, during stack walking for OSR, we remember whether we want to OSR
to the position before the instruction (if it's the top frame), or after
the call instruction (if it's deeper in the stack). We then use the
{is_instruction} predicate to find the right location.

R=thibaudm@chromium.org

Bug: v8:10321
Change-Id: I73212a7532c6ecf4c82bde76fe4059c8203e422c
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2116206Reviewed-by: Thibaud Michaud <thibaudm@chromium.org>
Commit-Queue: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#66850}

src/wasm/baseline/liftoff-compiler.cc

@@ -649,7 +649,7 @@ class LiftoffCompiler {
     if (!ool->regs_to_save.is_empty()) __ PushRegisters(ool->regs_to_save);
 
     source_position_table_builder_.AddPosition(
-        __ pc_offset(), SourcePosition(ool->position), false);
+        __ pc_offset(), SourcePosition(ool->position), true);
     __ CallRuntimeStub(ool->stub);
     DCHECK_EQ(!debug_sidetable_builder_, !ool->debug_sidetable_entry_builder);
     if (V8_UNLIKELY(ool->debug_sidetable_entry_builder)) {
@@ -686,12 +686,11 @@ class LiftoffCompiler {
   }
 
   void NextInstruction(FullDecoder* decoder, WasmOpcode opcode) {
-    bool breakpoint = false;
     if (V8_UNLIKELY(next_breakpoint_ptr_)) {
       if (*next_breakpoint_ptr_ == 0) {
         // A single breakpoint at offset 0 indicates stepping.
         DCHECK_EQ(next_breakpoint_ptr_ + 1, next_breakpoint_end_);
-        breakpoint = true;
+        EmitBreakpoint(decoder);
       } else {
         while (next_breakpoint_ptr_ != next_breakpoint_end_ &&
                *next_breakpoint_ptr_ < decoder->position()) {
@@ -701,17 +700,12 @@ class LiftoffCompiler {
         if (next_breakpoint_ptr_ == next_breakpoint_end_) {
           next_breakpoint_ptr_ = next_breakpoint_end_ = nullptr;
         } else if (*next_breakpoint_ptr_ == decoder->position()) {
-          breakpoint = true;
+          EmitBreakpoint(decoder);
         }
       }
     }
-    if (breakpoint) {
-      EmitBreakpoint(decoder);
-    } else if (opcode != kExprCallFunction) {
-      // There is no breakpoint and no call instruction, but we might still need
-      // a source position to get the return address for a removed breakpoint.
-      MaybeGenerateExtraSourcePos(decoder, false);
-    }
+    // Potentially generate the source position to OSR to this instruction.
+    MaybeGenerateExtraSourcePos(decoder);
     TraceCacheState(decoder);
 #ifdef DEBUG
     SLOW_DCHECK(__ ValidateCacheState());
@@ -731,7 +725,6 @@ class LiftoffCompiler {
         __ pc_offset(), SourcePosition(decoder->position()), false);
     __ CallRuntimeStub(WasmCode::kWasmDebugBreak);
     RegisterDebugSideTableEntry(DebugSideTableBuilder::kAllowRegisters);
-    MaybeGenerateExtraSourcePos(decoder);
     safepoint_table_builder_.DefineSafepoint(&asm_, Safepoint::kNoLazyDeopt);
   }
 
@@ -1923,7 +1916,7 @@ class LiftoffCompiler {
     }
 
     source_position_table_builder_.AddPosition(__ pc_offset(),
-                                               SourcePosition(position), false);
+                                               SourcePosition(position), true);
     __ CallRuntimeStub(WasmCode::kWasmTraceMemory);
     safepoint_table_builder_.DefineSafepoint(&asm_, Safepoint::kNoLazyDeopt);
 
@@ -2089,6 +2082,11 @@ class LiftoffCompiler {
     call_descriptor =
         GetLoweredCallDescriptor(compilation_zone_, call_descriptor);
 
+    // Place the source position before any stack manipulation, since this will
+    // be used for OSR in debugging.
+    source_position_table_builder_.AddPosition(
+        __ pc_offset(), SourcePosition(decoder->position()), true);
+
     if (imm.index < env_->module->num_imported_functions) {
       // A direct call to an imported function.
       LiftoffRegList pinned;
@@ -2111,17 +2109,11 @@ class LiftoffCompiler {
 
       Register* explicit_instance = &imported_function_ref;
       __ PrepareCall(imm.sig, call_descriptor, &target, explicit_instance);
-      source_position_table_builder_.AddPosition(
-          __ pc_offset(), SourcePosition(decoder->position()), false);
-
       __ CallIndirect(imm.sig, call_descriptor, target);
     } else {
       // A direct call within this module just gets the current instance.
       __ PrepareCall(imm.sig, call_descriptor);
 
-      source_position_table_builder_.AddPosition(
-          __ pc_offset(), SourcePosition(decoder->position()), false);
-
       // Just encode the function index. This will be patched at instantiation.
       Address addr = static_cast<Address>(imm.index);
       __ CallNativeWasmCode(addr);
@@ -2150,6 +2142,11 @@ class LiftoffCompiler {
       return;
     }
 
+    // Place the source position before any stack manipulation, since this will
+    // be used for OSR in debugging.
+    source_position_table_builder_.AddPosition(
+        __ pc_offset(), SourcePosition(decoder->position()), true);
+
     // Pop the index.
     Register index = __ PopToRegister().gp();
     // If that register is still being used after popping, we move it to another
@@ -2250,9 +2247,6 @@ class LiftoffCompiler {
     __ Load(LiftoffRegister(scratch), table, index, 0, kPointerLoadType,
             pinned);
 
-    source_position_table_builder_.AddPosition(
-        __ pc_offset(), SourcePosition(decoder->position()), false);
-
     auto call_descriptor =
         compiler::GetWasmCallDescriptor(compilation_zone_, imm.sig);
     call_descriptor =
@@ -2265,6 +2259,8 @@ class LiftoffCompiler {
     RegisterDebugSideTableEntry(DebugSideTableBuilder::kDidSpill);
     safepoint_table_builder_.DefineSafepoint(&asm_, Safepoint::kNoLazyDeopt);
 
+    MaybeGenerateExtraSourcePos(decoder);
+
     __ FinishCall(imm.sig, call_descriptor);
   }
 
@@ -3089,34 +3085,23 @@ class LiftoffCompiler {
  private:
   // Emit additional source positions for return addresses. Used by debugging to
   // OSR frames with different sets of breakpoints.
-  void MaybeGenerateExtraSourcePos(Decoder* decoder, bool after_call = true) {
-    if (V8_UNLIKELY(next_extra_source_pos_ptr_) &&
-        *next_extra_source_pos_ptr_ == static_cast<int>(decoder->position())) {
-      if (*next_extra_source_pos_ptr_ ==
-          static_cast<int>(decoder->position())) {
-        next_extra_source_pos_ptr_++;
-        if (next_extra_source_pos_ptr_ == next_extra_source_pos_end_) {
-          next_extra_source_pos_ptr_ = next_extra_source_pos_end_ = nullptr;
-        }
-        source_position_table_builder_.AddPosition(
-            __ pc_offset(), SourcePosition(decoder->position()), false);
-        // Add a nop after the breakpoint, such that following code has another
-        // PC and does not collide with the source position recorded above.
-        // TODO(thibaudm/clemens): Remove this.
-        __ nop();
-        if (!after_call) {
-          // A frame position is given by the last source position before the
-          // return address. If the return address does not follow a call, which
-          // happens after removing a breakpoint, this would be incorrect.
-          // Generate two source positions in this case. The second one will be
-          // used as the return address and the first one as the actual
-          // position.
-          source_position_table_builder_.AddPosition(
-              __ pc_offset(), SourcePosition(decoder->position()), false);
-          __ nop();
-        }
+  void MaybeGenerateExtraSourcePos(Decoder* decoder) {
+    if (V8_LIKELY(next_extra_source_pos_ptr_ == nullptr)) return;
+    int position = static_cast<int>(decoder->position());
+    while (*next_extra_source_pos_ptr_ < position) {
+      ++next_extra_source_pos_ptr_;
+      if (next_extra_source_pos_ptr_ == next_extra_source_pos_end_) {
+        next_extra_source_pos_ptr_ = next_extra_source_pos_end_ = nullptr;
+        return;
       }
     }
+    if (*next_extra_source_pos_ptr_ != position) return;
+    source_position_table_builder_.AddPosition(
+        __ pc_offset(), SourcePosition(decoder->position()), true);
+    // Add a nop here, such that following code has another
+    // PC and does not collide with the source position recorded above.
+    // TODO(thibaudm/clemens): Remove this.
+    __ nop();
   }
 
   static constexpr WasmOpcode kNoOutstandingOp = kExprUnreachable;

src/wasm/wasm-debug.cc

@@ -476,23 +476,35 @@ std::vector<int> StackFramePositions(int func_index, Isolate* isolate) {
   return byte_offsets;
 }
 
-Address FindNewPC(int byte_offset, WasmCode* wasm_code) {
+enum ReturnLocation { kAfterBreakpoint, kAfterWasmCall };
+
+Address FindNewPC(WasmCode* wasm_code, int byte_offset,
+                  ReturnLocation return_location) {
   Vector<const uint8_t> new_pos_table = wasm_code->source_positions();
 
   DCHECK_LE(0, byte_offset);
 
+  // If {return_location == kAfterBreakpoint} we search for the first code
+  // offset which is marked as instruction (i.e. not the breakpoint).
+  // If {return_location == kAfterWasmCall} we return the last code offset
+  // associated with the byte offset.
   SourcePositionTableIterator it(new_pos_table);
   while (!it.done() && it.source_position().ScriptOffset() != byte_offset) {
     it.Advance();
   }
-  DCHECK(!it.done());
-  DCHECK_EQ(byte_offset, it.source_position().ScriptOffset());
-  // If there are two source positions with this offset, one is for the call and
-  // one for the return address. Get the return address.
-  it.Advance();
-  DCHECK(!it.done());
-  DCHECK_EQ(byte_offset, it.source_position().ScriptOffset());
-  return wasm_code->instruction_start() + it.code_offset();
+  if (return_location == kAfterBreakpoint) {
+    while (!it.is_statement()) it.Advance();
+    DCHECK_EQ(byte_offset, it.source_position().ScriptOffset());
+    return wasm_code->instruction_start() + it.code_offset();
+  }
+
+  DCHECK_EQ(kAfterWasmCall, return_location);
+  int code_offset;
+  do {
+    code_offset = it.code_offset();
+    it.Advance();
+  } while (!it.done() && it.source_position().ScriptOffset() == byte_offset);
+  return wasm_code->instruction_start() + code_offset;
 }
 
 }  // namespace
@@ -640,8 +652,7 @@ class DebugInfoImpl {
     std::vector<int> stack_frame_positions =
         StackFramePositions(func_index, current_isolate);
 
-    WasmCompilationResult result;
-    result = ExecuteLiftoffCompilation(
+    WasmCompilationResult result = ExecuteLiftoffCompilation(
         native_module_->engine()->allocator(), &env, body, func_index, nullptr,
         nullptr, offsets, &debug_sidetable, VectorOf(stack_frame_positions));
     // Liftoff compilation failure is a FATAL error. We rely on complete Liftoff
@@ -710,7 +721,8 @@ class DebugInfoImpl {
   bool IsStepping(WasmCompiledFrame* frame) {
     Isolate* isolate = frame->wasm_instance().GetIsolate();
     StepAction last_step_action = isolate->debug()->last_step_action();
-    return last_step_action == StepIn || stepping_frame_ == frame->id();
+    return stepping_frame_ == frame->id() ||
+           (last_step_action == StepIn && stepping_frame_ != NO_ID);
   }
 
   void RemoveBreakpoint(int func_index, int offset, Isolate* current_isolate) {
@@ -828,7 +840,11 @@ class DebugInfoImpl {
   // TODO(thibaudm): update other threads as well.
   void UpdateReturnAddresses(Isolate* isolate, WasmCode* new_code) {
     DCHECK(new_code->is_liftoff());
-    for (StackTraceFrameIterator it(isolate); !it.done(); it.Advance()) {
+    // The first return location is after the breakpoint, others are after wasm
+    // calls.
+    ReturnLocation return_location = kAfterBreakpoint;
+    for (StackTraceFrameIterator it(isolate); !it.done();
+         it.Advance(), return_location = kAfterWasmCall) {
       // We still need the flooded function for stepping.
       if (it.frame()->id() == stepping_frame_) continue;
       if (!it.is_wasm()) continue;
@@ -840,7 +856,7 @@ class DebugInfoImpl {
       int pc_offset =
           static_cast<int>(frame->pc() - old_code->instruction_start());
       int byte_offset = FindByteOffset(pc_offset, old_code);
-      Address new_pc = FindNewPC(byte_offset, new_code);
+      Address new_pc = FindNewPC(new_code, byte_offset, return_location);
       PointerAuthentication::ReplacePC(frame->pc_address(), new_pc,
                                        kSystemPointerSize);
     }

test/inspector/debugger/wasm-stepping-liftoff-expected.txt

@@ -14,7 +14,7 @@ Debugger.resume called
 Paused at wasm://wasm/42af3c82:0:73
 Debugger.stepOver called
 Paused at wasm://wasm/42af3c82:0:75
-Debugger.stepOver called
+Debugger.stepInto called
 Paused at wasm://wasm/42af3c82:0:59
 Debugger.resume called
 Paused at wasm://wasm/42af3c82:0:73

test/inspector/debugger/wasm-stepping-liftoff.js

@@ -96,7 +96,7 @@ function instantiate(bytes) {
   await waitForPauseAndStep('stepOver');  // over call to wasm_A
   await waitForPauseAndStep('resume');    // stop on breakpoint
   await waitForPauseAndStep('stepOver');  // over call
-  await waitForPauseAndStep('stepOver');  // over br
+  await waitForPauseAndStep('stepInto');  // == stepOver br
   await waitForPauseAndStep('resume');    // to next breakpoint (3rd iteration)
   await waitForPauseAndStep('stepInto');  // into wasm_A
   // Step out of wasm_A, back to wasm_B.