Runtime implementation of try catch inlining

The main differences in the runtime are:
1) We now use a list of dex_pcs to find the correct catch handler
   instead of a single dex pc
2) We now need to restore vregs of the whole frame, which may be
   an inline frame.

Bug: 227283224
Test: art/test/testrunner/testrunner.py --host --64 --optimizing -b
Change-Id: I95d2f32088e1d420c83962a1693be18f3b63f8b4

12 files changed, 301 insertions, 95 deletions

diff --git a/compiler/exception_test.cc b/compiler/exception_test.cc
index f86e0b9b01..f55ba8e0e5 100644
--- a/compiler/exception_test.cc
+++ b/compiler/exception_test.cc
@@ -194,15 +194,13 @@ TEST_F(ExceptionTest, StackTraceElement) {
 
   OatQuickMethodHeader* header = OatQuickMethodHeader::FromEntryPoint(
       method_g_->GetEntryPointFromQuickCompiledCode());
-  fake_stack.push_back(header->ToNativeQuickPc(
-      method_g_, kDexPc, /* is_for_catch_handler= */ false));  // return pc
+  fake_stack.push_back(header->ToNativeQuickPc(method_g_, kDexPc));  // return pc
 
   // Create/push fake 16byte stack frame for method g
   fake_stack.push_back(reinterpret_cast<uintptr_t>(method_g_));
   fake_stack.push_back(0);
   fake_stack.push_back(0);
-  fake_stack.push_back(header->ToNativeQuickPc(
-      method_g_, kDexPc, /* is_for_catch_handler= */ false));  // return pc
+  fake_stack.push_back(header->ToNativeQuickPc(method_g_, kDexPc));  // return pc
 
   // Create/push fake 16byte stack frame for method f
   fake_stack.push_back(reinterpret_cast<uintptr_t>(method_f_));
diff --git a/compiler/optimizing/code_generator.cc b/compiler/optimizing/code_generator.cc
index bd8db30d44..1c5912d87e 100644
--- a/compiler/optimizing/code_generator.cc
+++ b/compiler/optimizing/code_generator.cc
@@ -15,6 +15,7 @@
  */
 
 #include "code_generator.h"
+#include "base/globals.h"
 
 #ifdef ART_ENABLE_CODEGEN_arm
 #include "code_generator_arm_vixl.h"
@@ -1339,8 +1340,11 @@ void CodeGenerator::RecordCatchBlockInfo() {
       continue;
     }
 
-    // Get the outer dex_pc
-    uint32_t outer_dex_pc = block->GetDexPc();
+    // Get the outer dex_pc. We save the full environment list for DCHECK purposes in kIsDebugBuild.
+    std::vector<uint32_t> dex_pc_list_for_verification;
+    if (kIsDebugBuild) {
+      dex_pc_list_for_verification.push_back(block->GetDexPc());
+    }
     DCHECK(block->GetFirstInstruction()->IsNop());
     DCHECK(block->GetFirstInstruction()->AsNop()->NeedsEnvironment());
     HEnvironment* const environment = block->GetFirstInstruction()->GetEnvironment();
@@ -1348,15 +1352,26 @@ void CodeGenerator::RecordCatchBlockInfo() {
     HEnvironment* outer_environment = environment;
     while (outer_environment->GetParent() != nullptr) {
       outer_environment = outer_environment->GetParent();
+      if (kIsDebugBuild) {
+        dex_pc_list_for_verification.push_back(outer_environment->GetDexPc());
+      }
+    }
+
+    if (kIsDebugBuild) {
+      // dex_pc_list_for_verification is set from innnermost to outermost. Let's reverse it
+      // since we are expected to pass from outermost to innermost.
+      std::reverse(dex_pc_list_for_verification.begin(), dex_pc_list_for_verification.end());
+      DCHECK_EQ(dex_pc_list_for_verification.front(), outer_environment->GetDexPc());
     }
-    outer_dex_pc = outer_environment->GetDexPc();
 
     uint32_t native_pc = GetAddressOf(block);
-    stack_map_stream->BeginStackMapEntry(outer_dex_pc,
+    stack_map_stream->BeginStackMapEntry(outer_environment->GetDexPc(),
                                          native_pc,
                                          /* register_mask= */ 0,
                                          /* sp_mask= */ nullptr,
-                                         StackMap::Kind::Catch);
+                                         StackMap::Kind::Catch,
+                                         /* needs_vreg_info= */ true,
+                                         dex_pc_list_for_verification);
 
     EmitEnvironment(environment,
                     /* slow_path= */ nullptr,
diff --git a/compiler/optimizing/stack_map_stream.cc b/compiler/optimizing/stack_map_stream.cc
index c13a35567b..31428d5c2a 100644
--- a/compiler/optimizing/stack_map_stream.cc
+++ b/compiler/optimizing/stack_map_stream.cc
@@ -20,6 +20,7 @@
 #include <vector>
 
 #include "art_method-inl.h"
+#include "base/globals.h"
 #include "base/stl_util.h"
 #include "class_linker.h"
 #include "dex/dex_file.h"
@@ -101,16 +102,21 @@ void StackMapStream::EndMethod(size_t code_size) {
   }
 }
 
-void StackMapStream::BeginStackMapEntry(uint32_t dex_pc,
-                                        uint32_t native_pc_offset,
-                                        uint32_t register_mask,
-                                        BitVector* stack_mask,
-                                        StackMap::Kind kind,
-                                        bool needs_vreg_info) {
+void StackMapStream::BeginStackMapEntry(
+    uint32_t dex_pc,
+    uint32_t native_pc_offset,
+    uint32_t register_mask,
+    BitVector* stack_mask,
+    StackMap::Kind kind,
+    bool needs_vreg_info,
+    const std::vector<uint32_t>& dex_pc_list_for_catch_verification) {
   DCHECK(in_method_) << "Call BeginMethod first";
   DCHECK(!in_stack_map_) << "Mismatched Begin/End calls";
   in_stack_map_ = true;
 
+  DCHECK_IMPLIES(!dex_pc_list_for_catch_verification.empty(), kind == StackMap::Kind::Catch);
+  DCHECK_IMPLIES(!dex_pc_list_for_catch_verification.empty(), kIsDebugBuild);
+
   current_stack_map_ = BitTableBuilder<StackMap>::Entry();
   current_stack_map_[StackMap::kKind] = static_cast<uint32_t>(kind);
   current_stack_map_[StackMap::kPackedNativePc] =
@@ -151,7 +157,8 @@ void StackMapStream::BeginStackMapEntry(uint32_t dex_pc,
                                                                     instruction_set_);
         CHECK_EQ(stack_map.Row(), stack_map_index);
       } else if (kind == StackMap::Kind::Catch) {
-        StackMap stack_map = code_info.GetCatchStackMapForDexPc(dex_pc);
+        StackMap stack_map = code_info.GetCatchStackMapForDexPc(
+            ArrayRef<const uint32_t>(dex_pc_list_for_catch_verification));
         CHECK_EQ(stack_map.Row(), stack_map_index);
       }
       StackMap stack_map = code_info.GetStackMapAt(stack_map_index);
diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index 1aaa6aee9e..2ba60279ff 100644
--- a/compiler/optimizing/stack_map_stream.h
+++ b/compiler/optimizing/stack_map_stream.h
@@ -68,12 +68,15 @@ class StackMapStream : public DeletableArenaObject<kArenaAllocStackMapStream> {
                    bool debuggable);
   void EndMethod(size_t code_size);
 
-  void BeginStackMapEntry(uint32_t dex_pc,
-                          uint32_t native_pc_offset,
-                          uint32_t register_mask = 0,
-                          BitVector* sp_mask = nullptr,
-                          StackMap::Kind kind = StackMap::Kind::Default,
-                          bool needs_vreg_info = true);
+  void BeginStackMapEntry(
+      uint32_t dex_pc,
+      uint32_t native_pc_offset,
+      uint32_t register_mask = 0,
+      BitVector* sp_mask = nullptr,
+      StackMap::Kind kind = StackMap::Kind::Default,
+      bool needs_vreg_info = true,
+      const std::vector<uint32_t>& dex_pc_list_for_catch_verification = std::vector<uint32_t>());
+
   void EndStackMapEntry();
 
   void AddDexRegisterEntry(DexRegisterLocation::Kind kind, int32_t value) {
diff --git a/runtime/oat_quick_method_header.cc b/runtime/oat_quick_method_header.cc
index 8fbf02abea..f0764da703 100644
--- a/runtime/oat_quick_method_header.cc
+++ b/runtime/oat_quick_method_header.cc
@@ -57,33 +57,61 @@ uint32_t OatQuickMethodHeader::ToDexPc(ArtMethod** frame,
 
 uintptr_t OatQuickMethodHeader::ToNativeQuickPc(ArtMethod* method,
                                                 const uint32_t dex_pc,
-                                                bool is_for_catch_handler,
                                                 bool abort_on_failure) const {
   const void* entry_point = GetEntryPoint();
   DCHECK(!method->IsNative());
+  // For catch handlers use the ArrayRef<const uint32_t> version of ToNativeQuickPc.
+  DCHECK(!IsNterpMethodHeader());
+  DCHECK(IsOptimized());
+  // Search for the dex-to-pc mapping in stack maps.
+  CodeInfo code_info = CodeInfo::DecodeInlineInfoOnly(this);
+
+  StackMap stack_map = code_info.GetStackMapForDexPc(dex_pc);
+  if (stack_map.IsValid()) {
+    return reinterpret_cast<uintptr_t>(entry_point) + stack_map.GetNativePcOffset(kRuntimeISA);
+  }
+  if (abort_on_failure) {
+    ScopedObjectAccess soa(Thread::Current());
+    LOG(FATAL) << "Failed to find native offset for dex pc 0x" << std::hex << dex_pc << " in "
+               << method->PrettyMethod();
+  }
+  return UINTPTR_MAX;
+}
+
+uintptr_t OatQuickMethodHeader::ToNativeQuickPcForCatchHandlers(
+    ArtMethod* method,
+    ArrayRef<const uint32_t> dex_pc_list,
+    /* out */ uint32_t* stack_map_row,
+    bool abort_on_failure) const {
+  const void* entry_point = GetEntryPoint();
+  DCHECK(!method->IsNative());
   if (IsNterpMethodHeader()) {
-    // This should only be called on an nterp frame for getting a catch handler.
-    CHECK(is_for_catch_handler);
     return NterpGetCatchHandler();
   }
   DCHECK(IsOptimized());
   // Search for the dex-to-pc mapping in stack maps.
   CodeInfo code_info = CodeInfo::DecodeInlineInfoOnly(this);
 
-  // All stack maps are stored in the same CodeItem section, safepoint stack
-  // maps first, then catch stack maps. We use `is_for_catch_handler` to select
-  // the order of iteration.
-  StackMap stack_map =
-      LIKELY(is_for_catch_handler) ? code_info.GetCatchStackMapForDexPc(dex_pc)
-                                   : code_info.GetStackMapForDexPc(dex_pc);
+  StackMap stack_map = code_info.GetCatchStackMapForDexPc(dex_pc_list);
+  *stack_map_row = stack_map.Row();
   if (stack_map.IsValid()) {
     return reinterpret_cast<uintptr_t>(entry_point) +
            stack_map.GetNativePcOffset(kRuntimeISA);
   }
   if (abort_on_failure) {
+    std::stringstream ss;
+    bool first = true;
+    ss << "Failed to find native offset for dex pcs (from outermost to innermost) " << std::hex;
+    for (auto dex_pc : dex_pc_list) {
+      if (!first) {
+        ss << ", ";
+      }
+      first = false;
+      ss << "0x" << dex_pc;
+    }
     ScopedObjectAccess soa(Thread::Current());
-    LOG(FATAL) << "Failed to find native offset for dex pc 0x" << std::hex << dex_pc
-               << " in " << method->PrettyMethod();
+    ss << " in " << method->PrettyMethod();
+    LOG(FATAL) << ss.str();
   }
   return UINTPTR_MAX;
 }
diff --git a/runtime/oat_quick_method_header.h b/runtime/oat_quick_method_header.h
index 8e6d08e3a1..769d67fb61 100644
--- a/runtime/oat_quick_method_header.h
+++ b/runtime/oat_quick_method_header.h
@@ -161,11 +161,17 @@ class PACKED(4) OatQuickMethodHeader {
     return frame_size - core_spill_size - fpu_spill_size - kShouldDeoptimizeFlagSize;
   }
 
+  // For non-catch handlers. Only used in test code.
   uintptr_t ToNativeQuickPc(ArtMethod* method,
                             const uint32_t dex_pc,
-                            bool is_for_catch_handler,
                             bool abort_on_failure = true) const;
 
+  // For catch handlers.
+  uintptr_t ToNativeQuickPcForCatchHandlers(ArtMethod* method,
+                                            ArrayRef<const uint32_t> dex_pc_list,
+                                            /* out */ uint32_t* stack_map_row,
+                                            bool abort_on_failure = true) const;
+
   uint32_t ToDexPc(ArtMethod** frame,
                    const uintptr_t pc,
                    bool abort_on_failure = true) const
diff --git a/runtime/quick_exception_handler.cc b/runtime/quick_exception_handler.cc
index 82f50346e2..b9dd3e11b6 100644
--- a/runtime/quick_exception_handler.cc
+++ b/runtime/quick_exception_handler.cc
@@ -15,11 +15,14 @@
  */
 
 #include "quick_exception_handler.h"
+
 #include <ios>
 #include <queue>
+#include <sstream>
 
 #include "arch/context.h"
 #include "art_method-inl.h"
+#include "base/array_ref.h"
 #include "base/enums.h"
 #include "base/globals.h"
 #include "base/logging.h"  // For VLOG_IS_ON.
@@ -56,7 +59,6 @@ QuickExceptionHandler::QuickExceptionHandler(Thread* self, bool is_deoptimizatio
       handler_quick_frame_pc_(0),
       handler_method_header_(nullptr),
       handler_quick_arg0_(0),
-      handler_dex_pc_(0),
       clear_exception_(false),
       handler_frame_depth_(kInvalidFrameDepth),
       full_fragment_done_(false) {}
@@ -133,10 +135,12 @@ class CatchBlockStackVisitor final : public StackVisitor {
       uint32_t found_dex_pc = method->FindCatchBlock(to_find, dex_pc, &clear_exception);
       exception_handler_->SetClearException(clear_exception);
       if (found_dex_pc != dex::kDexNoIndex) {
-        exception_handler_->SetHandlerDexPc(found_dex_pc);
+        exception_handler_->SetHandlerDexPcList(ComputeDexPcList(found_dex_pc));
+        uint32_t stack_map_row = -1;
         exception_handler_->SetHandlerQuickFramePc(
-            GetCurrentOatQuickMethodHeader()->ToNativeQuickPc(
-                method, found_dex_pc, /* is_for_catch_handler= */ true));
+            GetCurrentOatQuickMethodHeader()->ToNativeQuickPcForCatchHandlers(
+                method, exception_handler_->GetHandlerDexPcList(), &stack_map_row));
+        exception_handler_->SetCatchStackMapRow(stack_map_row);
         exception_handler_->SetHandlerQuickFrame(GetCurrentQuickFrame());
         exception_handler_->SetHandlerMethodHeader(GetCurrentOatQuickMethodHeader());
         return false;  // End stack walk.
@@ -215,10 +219,24 @@ void QuickExceptionHandler::FindCatch(ObjPtr<mirror::Throwable> exception,
       }
       if (GetHandlerMethod() != nullptr) {
         const DexFile* dex_file = GetHandlerMethod()->GetDexFile();
-        int line_number =
-            annotations::GetLineNumFromPC(dex_file, GetHandlerMethod(), handler_dex_pc_);
+        DCHECK_GE(handler_dex_pc_list_.size(), 1u);
+        int line_number = annotations::GetLineNumFromPC(
+            dex_file, GetHandlerMethod(), handler_dex_pc_list_.front());
+
+        // We may have an inlined method. If so, we can add some extra logging.
+        std::stringstream ss;
+        ArtMethod* maybe_inlined_method = visitor.GetMethod();
+        if (maybe_inlined_method != GetHandlerMethod()) {
+          const DexFile* inlined_dex_file = maybe_inlined_method->GetDexFile();
+          DCHECK_GE(handler_dex_pc_list_.size(), 2u);
+          int inlined_line_number = annotations::GetLineNumFromPC(
+              inlined_dex_file, maybe_inlined_method, handler_dex_pc_list_.back());
+          ss << " which ends up calling inlined method " << maybe_inlined_method->PrettyMethod()
+             << " (line: " << inlined_line_number << ")";
+        }
+
         LOG(INFO) << "Handler: " << GetHandlerMethod()->PrettyMethod() << " (line: "
-                  << line_number << ")";
+                  << line_number << ")" << ss.str();
       }
     }
     // Exception was cleared as part of delivery.
@@ -283,15 +301,18 @@ void QuickExceptionHandler::SetCatchEnvironmentForOptimizedHandler(StackVisitor*
     self_->DumpStack(LOG_STREAM(INFO) << "Setting catch phis: ");
   }
 
-  CodeItemDataAccessor accessor(GetHandlerMethod()->DexInstructionData());
-  const size_t number_of_vregs = accessor.RegistersSize();
   CodeInfo code_info(handler_method_header_);
 
   // Find stack map of the catch block.
-  StackMap catch_stack_map = code_info.GetCatchStackMapForDexPc(GetHandlerDexPc());
+  ArrayRef<const uint32_t> dex_pc_list = GetHandlerDexPcList();
+  DCHECK_GE(dex_pc_list.size(), 1u);
+  StackMap catch_stack_map = code_info.GetStackMapAt(GetCatchStackMapRow());
   DCHECK(catch_stack_map.IsValid());
-  DexRegisterMap catch_vreg_map = code_info.GetDexRegisterMapOf(catch_stack_map);
-  DCHECK_EQ(catch_vreg_map.size(), number_of_vregs);
+  DCHECK_EQ(catch_stack_map.Row(), code_info.GetCatchStackMapForDexPc(dex_pc_list).Row());
+  const uint32_t catch_depth = dex_pc_list.size() - 1;
+  const size_t number_of_registers = stack_visitor->GetNumberOfRegisters(&code_info, catch_depth);
+  DexRegisterMap catch_vreg_map =
+      code_info.GetDexRegisterMapOf(catch_stack_map, /* first= */ 0, number_of_registers);
 
   if (!catch_vreg_map.HasAnyLiveDexRegisters()) {
     return;
@@ -301,26 +322,47 @@ void QuickExceptionHandler::SetCatchEnvironmentForOptimizedHandler(StackVisitor*
   StackMap throw_stack_map =
       code_info.GetStackMapForNativePcOffset(stack_visitor->GetNativePcOffset());
   DCHECK(throw_stack_map.IsValid());
-  DexRegisterMap throw_vreg_map = code_info.GetDexRegisterMapOf(throw_stack_map);
-  DCHECK_EQ(throw_vreg_map.size(), number_of_vregs);
+  const uint32_t throw_depth = stack_visitor->InlineDepth();
+  DCHECK_EQ(throw_depth, catch_depth);
+  DexRegisterMap throw_vreg_map =
+      code_info.GetDexRegisterMapOf(throw_stack_map, /* first= */ 0, number_of_registers);
+  DCHECK_EQ(throw_vreg_map.size(), catch_vreg_map.size());
+
+  // First vreg that it is part of the catch's environment.
+  const size_t catch_vreg_start = catch_depth == 0
+    ? 0
+    : stack_visitor->GetNumberOfRegisters(&code_info, catch_depth - 1);
+
+  // We don't need to copy anything in the parent's environment.
+  for (size_t vreg = 0; vreg < catch_vreg_start; ++vreg) {
+    DexRegisterLocation::Kind catch_location_kind = catch_vreg_map[vreg].GetKind();
+    DCHECK(catch_location_kind == DexRegisterLocation::Kind::kNone ||
+           catch_location_kind == DexRegisterLocation::Kind::kConstant ||
+           catch_location_kind == DexRegisterLocation::Kind::kInStack)
+        << "Unexpected catch_location_kind: " << catch_location_kind;
+  }
 
-  // Copy values between them.
-  for (uint16_t vreg = 0; vreg < number_of_vregs; ++vreg) {
-    DexRegisterLocation::Kind catch_location = catch_vreg_map[vreg].GetKind();
-    if (catch_location == DexRegisterLocation::Kind::kNone) {
+  // Copy values between the throw and the catch.
+  for (size_t vreg = catch_vreg_start; vreg < catch_vreg_map.size(); ++vreg) {
+    DexRegisterLocation::Kind catch_location_kind = catch_vreg_map[vreg].GetKind();
+    if (catch_location_kind == DexRegisterLocation::Kind::kNone) {
       continue;
     }
-    DCHECK(catch_location == DexRegisterLocation::Kind::kInStack);
 
-    // Get vreg value from its current location.
+    // Consistency checks.
+    DCHECK_EQ(catch_location_kind, DexRegisterLocation::Kind::kInStack);
     uint32_t vreg_value;
     VRegKind vreg_kind = ToVRegKind(throw_vreg_map[vreg].GetKind());
-    bool get_vreg_success =
-        stack_visitor->GetVReg(stack_visitor->GetMethod(),
-                               vreg,
-                               vreg_kind,
-                               &vreg_value,
-                               throw_vreg_map[vreg]);
+    DCHECK_NE(vreg_kind, kReferenceVReg)
+        << "The fast path in GetVReg doesn't expect a kReferenceVReg.";
+
+    // Get vreg value from its current location.
+    bool get_vreg_success = stack_visitor->GetVReg(stack_visitor->GetMethod(),
+                                                   vreg,
+                                                   vreg_kind,
+                                                   &vreg_value,
+                                                   throw_vreg_map[vreg],
+                                                   /* need_full_register_list= */ true);
     CHECK(get_vreg_success) << "VReg " << vreg << " was optimized out ("
                             << "method=" << ArtMethod::PrettyMethod(stack_visitor->GetMethod())
                             << ", dex_pc=" << stack_visitor->GetDexPc() << ", "
@@ -701,8 +743,10 @@ void QuickExceptionHandler::DoLongJump(bool smash_caller_saves) {
   if (!is_deoptimization_ &&
       handler_method_header_ != nullptr &&
       handler_method_header_->IsNterpMethodHeader()) {
+    // Interpreter procceses one method at a time i.e. not inlining
+    DCHECK_EQ(handler_dex_pc_list_.size(), 1u) << "We shouldn't have any inlined frames.";
     context_->SetNterpDexPC(reinterpret_cast<uintptr_t>(
-        GetHandlerMethod()->DexInstructions().Insns() + handler_dex_pc_));
+        GetHandlerMethod()->DexInstructions().Insns() + handler_dex_pc_list_.front()));
   }
   context_->DoLongJump();
   UNREACHABLE();
diff --git a/runtime/quick_exception_handler.h b/runtime/quick_exception_handler.h
index 9554f1d7d6..c87e38df3f 100644
--- a/runtime/quick_exception_handler.h
+++ b/runtime/quick_exception_handler.h
@@ -18,10 +18,13 @@
 #define ART_RUNTIME_QUICK_EXCEPTION_HANDLER_H_
 
 #include <android-base/logging.h>
+#include <cstdint>
 
+#include "base/array_ref.h"
 #include "base/macros.h"
 #include "base/mutex.h"
 #include "deoptimization_kind.h"
+#include "stack_map.h"
 #include "stack_reference.h"
 
 namespace art {
@@ -103,12 +106,20 @@ class QuickExceptionHandler {
     return *handler_quick_frame_;
   }
 
-  uint32_t GetHandlerDexPc() const {
-    return handler_dex_pc_;
+  ArrayRef<const uint32_t> GetHandlerDexPcList() const {
+    return ArrayRef<const uint32_t>(handler_dex_pc_list_);
   }
 
-  void SetHandlerDexPc(uint32_t dex_pc) {
-    handler_dex_pc_ = dex_pc;
+  void SetHandlerDexPcList(std::vector<uint32_t>&& handler_dex_pc_list) {
+    handler_dex_pc_list_ = std::move(handler_dex_pc_list);
+  }
+
+  uint32_t GetCatchStackMapRow() const {
+    return catch_stack_map_row_;
+  }
+
+  void SetCatchStackMapRow(uint32_t stack_map_row) {
+    catch_stack_map_row_ = stack_map_row;
   }
 
   bool GetClearException() const {
@@ -151,8 +162,11 @@ class QuickExceptionHandler {
   const OatQuickMethodHeader* handler_method_header_;
   // The value for argument 0.
   uintptr_t handler_quick_arg0_;
-  // The handler's dex PC, zero implies an uncaught exception.
-  uint32_t handler_dex_pc_;
+  // The handler's dex PC list including the inline dex_pcs. The dex_pcs are ordered from outermost
+  // to innermost. An empty list implies an uncaught exception.
+  std::vector<uint32_t> handler_dex_pc_list_;
+  // StackMap row corresponding to the found catch.
+  uint32_t catch_stack_map_row_;
   // Should the exception be cleared as the catch block has no move-exception?
   bool clear_exception_;
   // Frame depth of the catch handler or the upcall.
diff --git a/runtime/stack.cc b/runtime/stack.cc
index 5ee20e5fd3..c83814abd2 100644
--- a/runtime/stack.cc
+++ b/runtime/stack.cc
@@ -140,6 +140,29 @@ uint32_t StackVisitor::GetDexPc(bool abort_on_failure) const {
   }
 }
 
+std::vector<uint32_t> StackVisitor::ComputeDexPcList(uint32_t handler_dex_pc) const {
+  std::vector<uint32_t> result;
+  if (cur_shadow_frame_ == nullptr && cur_quick_frame_ != nullptr && IsInInlinedFrame()) {
+    const BitTableRange<InlineInfo>& infos = current_inline_frames_;
+    DCHECK_NE(infos.size(), 0u);
+
+    // Outermost dex_pc.
+    result.push_back(GetCurrentStackMap()->GetDexPc());
+
+    // The mid dex_pcs. Note that we skip the last one since we want to change that for
+    // `handler_dex_pc`.
+    for (size_t index = 0; index < infos.size() - 1; ++index) {
+      result.push_back(infos[index].GetDexPc());
+    }
+  }
+
+  // The innermost dex_pc has to be the handler dex_pc. In the case of no inline frames, it will be
+  // just the one dex_pc. In the case of inlining we will be replacing the innermost InlineInfo's
+  // dex_pc with this one.
+  result.push_back(handler_dex_pc);
+  return result;
+}
+
 extern "C" mirror::Object* artQuickGetProxyThisObject(ArtMethod** sp)
     REQUIRES_SHARED(Locks::mutator_lock_);
 
@@ -213,7 +236,8 @@ bool StackVisitor::GetVReg(ArtMethod* m,
                            uint16_t vreg,
                            VRegKind kind,
                            uint32_t* val,
-                           std::optional<DexRegisterLocation> location) const {
+                           std::optional<DexRegisterLocation> location,
+                           bool need_full_register_list) const {
   if (cur_quick_frame_ != nullptr) {
     DCHECK(context_ != nullptr);  // You can't reliably read registers without a context.
     DCHECK(m == GetMethod());
@@ -235,10 +259,10 @@ bool StackVisitor::GetVReg(ArtMethod* m,
         // which does not decode the stack maps.
         result = GetVRegFromOptimizedCode(location.value(), val);
         // Compare to the slower overload.
-        DCHECK_EQ(result, GetVRegFromOptimizedCode(m, vreg, kind, &val2));
+        DCHECK_EQ(result, GetVRegFromOptimizedCode(m, vreg, kind, &val2, need_full_register_list));
         DCHECK_EQ(*val, val2);
       } else {
-        result = GetVRegFromOptimizedCode(m, vreg, kind, val);
+        result = GetVRegFromOptimizedCode(m, vreg, kind, val, need_full_register_list);
       }
     }
     if (kind == kReferenceVReg) {
@@ -261,16 +285,20 @@ bool StackVisitor::GetVReg(ArtMethod* m,
   }
 }
 
+size_t StackVisitor::GetNumberOfRegisters(CodeInfo* code_info, int depth) const {
+  return depth == 0
+    ? code_info->GetNumberOfDexRegisters()
+    : current_inline_frames_[depth - 1].GetNumberOfDexRegisters();
+}
+
 bool StackVisitor::GetVRegFromOptimizedCode(ArtMethod* m,
                                             uint16_t vreg,
                                             VRegKind kind,
-                                            uint32_t* val) const {
+                                            uint32_t* val,
+                                            bool need_full_register_list) const {
   DCHECK_EQ(m, GetMethod());
   // Can't be null or how would we compile its instructions?
   DCHECK(m->GetCodeItem() != nullptr) << m->PrettyMethod();
-  CodeItemDataAccessor accessor(m->DexInstructionData());
-  uint16_t number_of_dex_registers = accessor.RegistersSize();
-  DCHECK_LT(vreg, number_of_dex_registers);
   const OatQuickMethodHeader* method_header = GetCurrentOatQuickMethodHeader();
   CodeInfo code_info(method_header);
 
@@ -278,13 +306,18 @@ bool StackVisitor::GetVRegFromOptimizedCode(ArtMethod* m,
   StackMap stack_map = code_info.GetStackMapForNativePcOffset(native_pc_offset);
   DCHECK(stack_map.IsValid());
 
-  DexRegisterMap dex_register_map = IsInInlinedFrame()
-      ? code_info.GetInlineDexRegisterMapOf(stack_map, current_inline_frames_.back())
-      : code_info.GetDexRegisterMapOf(stack_map);
+  DexRegisterMap dex_register_map = (IsInInlinedFrame() && !need_full_register_list)
+    ? code_info.GetInlineDexRegisterMapOf(stack_map, current_inline_frames_.back())
+    : code_info.GetDexRegisterMapOf(stack_map,
+                                    /* first= */ 0,
+                                    GetNumberOfRegisters(&code_info, InlineDepth()));
+
   if (dex_register_map.empty()) {
     return false;
   }
-  DCHECK_EQ(dex_register_map.size(), number_of_dex_registers);
+
+  const size_t number_of_dex_registers = dex_register_map.size();
+  DCHECK_LT(vreg, number_of_dex_registers);
   DexRegisterLocation::Kind location_kind = dex_register_map[vreg].GetKind();
   switch (location_kind) {
     case DexRegisterLocation::Kind::kInStack: {
diff --git a/runtime/stack.h b/runtime/stack.h
index bfda57b136..56f38aecef 100644
--- a/runtime/stack.h
+++ b/runtime/stack.h
@@ -192,6 +192,12 @@ class StackVisitor {
 
   uint32_t GetDexPc(bool abort_on_failure = true) const REQUIRES_SHARED(Locks::mutator_lock_);
 
+  // Returns a vector of the inlined dex pcs, in order from outermost to innermost but it replaces
+  // the innermost one with `handler_dex_pc`. In essence, (outermost dex pc, mid dex pc #1, ..., mid
+  // dex pc #n-1, `handler_dex_pc`).
+  std::vector<uint32_t> ComputeDexPcList(uint32_t handler_dex_pc) const
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
   ObjPtr<mirror::Object> GetThisObject() const REQUIRES_SHARED(Locks::mutator_lock_);
 
   size_t GetNativePcOffset() const REQUIRES_SHARED(Locks::mutator_lock_);
@@ -225,9 +231,8 @@ class StackVisitor {
                uint16_t vreg,
                VRegKind kind,
                uint32_t* val,
-               std::optional<DexRegisterLocation> location =
-                   std::optional<DexRegisterLocation>()) const
-      REQUIRES_SHARED(Locks::mutator_lock_);
+               std::optional<DexRegisterLocation> location = std::optional<DexRegisterLocation>(),
+               bool need_full_register_list = false) const REQUIRES_SHARED(Locks::mutator_lock_);
 
   bool GetVRegPair(ArtMethod* m, uint16_t vreg, VRegKind kind_lo, VRegKind kind_hi,
                    uint64_t* val) const
@@ -263,10 +268,16 @@ class StackVisitor {
     return !current_inline_frames_.empty();
   }
 
+  size_t InlineDepth() const { return current_inline_frames_.size(); }
+
   InlineInfo GetCurrentInlinedFrame() const {
     return current_inline_frames_.back();
   }
 
+  const BitTableRange<InlineInfo>& GetCurrentInlinedFrames() const {
+    return current_inline_frames_;
+  }
+
   uintptr_t GetCurrentQuickFramePc() const {
     return cur_quick_frame_pc_;
   }
@@ -306,6 +317,11 @@ class StackVisitor {
     return (should_deopt_flag & static_cast<uint8_t>(DeoptimizeFlagValue::kDebug)) != 0;
   }
 
+  // Return the number of dex register in the map from the outermost frame to the number of inlined
+  // frames indicated by `depth`. If `depth` is 0, grab just the registers from the outermost level.
+  // If it is greater than 0, grab as many inline frames as `depth` indicates.
+  size_t GetNumberOfRegisters(CodeInfo* code_info, int depth) const;
+
  private:
   // Private constructor known in the case that num_frames_ has already been computed.
   StackVisitor(Thread* thread,
@@ -334,7 +350,8 @@ class StackVisitor {
   bool GetVRegFromOptimizedCode(ArtMethod* m,
                                 uint16_t vreg,
                                 VRegKind kind,
-                                uint32_t* val) const
+                                uint32_t* val,
+                                bool need_full_register_list = false) const
       REQUIRES_SHARED(Locks::mutator_lock_);
 
   bool GetVRegPairFromDebuggerShadowFrame(uint16_t vreg,
diff --git a/runtime/stack_map.h b/runtime/stack_map.h
index 430aaa3204..514d30ed0c 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -20,9 +20,12 @@
 #include <limits>
 
 #include "arch/instruction_set.h"
+#include "base/array_ref.h"
 #include "base/bit_memory_region.h"
 #include "base/bit_table.h"
 #include "base/bit_utils.h"
+#include "base/globals.h"
+#include "base/logging.h"
 #include "base/memory_region.h"
 #include "dex/dex_file_types.h"
 #include "dex_register_location.h"
@@ -360,15 +363,12 @@ class CodeInfo {
     return GetMethodInfoOf(inline_info).GetMethodIndex();
   }
 
+  // Returns the dex registers for `stack_map`, ignoring any inlined dex registers.
   ALWAYS_INLINE DexRegisterMap GetDexRegisterMapOf(StackMap stack_map) const {
-    if (stack_map.HasDexRegisterMap()) {
-      DexRegisterMap map(number_of_dex_registers_, DexRegisterLocation::Invalid());
-      DecodeDexRegisterMap(stack_map.Row(), /* first_dex_register= */ 0, &map);
-      return map;
-    }
-    return DexRegisterMap(0, DexRegisterLocation::None());
+    return GetDexRegisterMapOf(stack_map, /* first= */ 0, number_of_dex_registers_);
   }
 
+  // Returns the dex register map of `inline_info`, and just those registers.
   ALWAYS_INLINE DexRegisterMap GetInlineDexRegisterMapOf(StackMap stack_map,
                                                          InlineInfo inline_info) const {
     if (stack_map.HasDexRegisterMap()) {
@@ -381,6 +381,17 @@ class CodeInfo {
           ? number_of_dex_registers_
           : inline_infos_.GetRow(inline_info.Row() - 1).GetNumberOfDexRegisters();
       uint32_t last = inline_info.GetNumberOfDexRegisters();
+      return GetDexRegisterMapOf(stack_map, first, last);
+    }
+    return DexRegisterMap(0, DexRegisterLocation::None());
+  }
+
+  // Returns the dex register map of `stack_map` in the range the range [first, last).
+  ALWAYS_INLINE DexRegisterMap GetDexRegisterMapOf(StackMap stack_map,
+                                                   uint32_t first,
+                                                   uint32_t last) const {
+    if (stack_map.HasDexRegisterMap()) {
+      DCHECK_LE(first, last);
       DexRegisterMap map(last - first, DexRegisterLocation::Invalid());
       DecodeDexRegisterMap(stack_map.Row(), first, &map);
       return map;
@@ -409,12 +420,39 @@ class CodeInfo {
     return stack_maps_.GetInvalidRow();
   }
 
-  // Searches the stack map list backwards because catch stack maps are stored at the end.
-  StackMap GetCatchStackMapForDexPc(uint32_t dex_pc) const {
+  StackMap GetCatchStackMapForDexPc(ArrayRef<const uint32_t> dex_pcs) const {
+    // Searches the stack map list backwards because catch stack maps are stored at the end.
     for (size_t i = GetNumberOfStackMaps(); i > 0; --i) {
       StackMap stack_map = GetStackMapAt(i - 1);
-      if (stack_map.GetDexPc() == dex_pc && stack_map.GetKind() == StackMap::Kind::Catch) {
-        return stack_map;
+      if (UNLIKELY(stack_map.GetKind() != StackMap::Kind::Catch)) {
+        // Early break since we should have catch stack maps only at the end.
+        if (kIsDebugBuild) {
+          for (size_t j = i - 1; j > 0; --j) {
+            DCHECK(GetStackMapAt(j - 1).GetKind() != StackMap::Kind::Catch);
+          }
+        }
+        break;
+      }
+
+      // Both the handler dex_pc and all of the inline dex_pcs have to match i.e. we want dex_pcs to
+      // be [stack_map_dex_pc, inline_dex_pc_1, ..., inline_dex_pc_n].
+      if (stack_map.GetDexPc() != dex_pcs.front()) {
+        continue;
+      }
+
+      const BitTableRange<InlineInfo>& inline_infos = GetInlineInfosOf(stack_map);
+      if (inline_infos.size() == dex_pcs.size() - 1) {
+        bool matching_dex_pcs = true;
+        for (size_t inline_info_index = 0; inline_info_index < inline_infos.size();
+             ++inline_info_index) {
+          if (inline_infos[inline_info_index].GetDexPc() != dex_pcs[inline_info_index + 1]) {
+            matching_dex_pcs = false;
+            break;
+          }
+        }
+        if (matching_dex_pcs) {
+          return stack_map;
+        }
       }
     }
     return stack_maps_.GetInvalidRow();
@@ -453,6 +491,10 @@ class CodeInfo {
     return (*code_info_data & kIsDebuggable) != 0;
   }
 
+  uint32_t GetNumberOfDexRegisters() {
+    return number_of_dex_registers_;
+  }
+
  private:
   // Scan backward to determine dex register locations at given stack map.
   void DecodeDexRegisterMap(uint32_t stack_map_index,
diff --git a/test/004-ReferenceMap/stack_walk_refmap_jni.cc b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
index 0659c0beb5..c84bd88db5 100644
--- a/test/004-ReferenceMap/stack_walk_refmap_jni.cc
+++ b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
@@ -25,9 +25,8 @@ namespace art {
   int t_size = sizeof(t) / sizeof(*t);                                                        \
   const OatQuickMethodHeader* method_header = GetCurrentOatQuickMethodHeader();               \
   uintptr_t native_quick_pc = method_header->ToNativeQuickPc(GetMethod(),                     \
-                                                 dex_pc,                                      \
-                                                 /* is_catch_handler */ false,                \
-                                                 abort_if_not_found);                         \
+                                                             dex_pc,                          \
+                                                             abort_if_not_found);             \
   if (native_quick_pc != UINTPTR_MAX) {                                                       \
     CheckReferences(t,                                                                        \
                     t_size,                                                                   \