[lldb][swift] Use assembly unwinders to recover async registers

The SwiftLanguageRuntime unwind plan for async functions relies on the extended
frame setup to recover the asynchronous contexts of these functions. This
extended frame consists of:

1. Tagging the saved fp register by setting its 60 bit, indicating the existence
of the extended frame.
2. An additional stack slot where the asynchronous context is saved.

The current unwind plans recover the asynchronous context by reading the stack
slot described in (2). This creates a number of challenges:

* The extended frame is not always created (e.g. leaf funclets).
* The extended frame is not necessarily valid inside the prologue or epilogue.
* The code to handle all possible combinations of funclet x {is_prologue,
is_epilogue} is very complex.
* Entry funclets don't have the extended frame.

This patch uses a new approach, exploiting the property that the asynchronous
context is passed in a callee-saved register. We use existing unwind plans that
inspect assembly and track spill slots, and recover the asynchronous register
from those.

Unwind Plans are not setup to query other Unwind Plans, more specifically they
are not setup to materialize the abstract location provided by an Unwind Plan
into a concrete value. Doing so would require a lot of plumbing. Instead, this
patch reimplements a very targeted conversion between the abstract locations and
concrete register values. We argue this is worth it, as the new unwind procedure
is much simpler, relies on existing plans, and works in many cases where the
previous implementation doesn't.

Testing-wise, this commit creates a fairly complex async function, with multiple
split points, and sets breakpoints on all instructions of all funclets. It then
asserts that unwinding works correctly in all of them. This has revealed an
issue with Q funclets, which clobber the async context location prior to freeing
the context of the funclet that has just finished executing (see the FIXME in
the test). This is not a problem with the new approach, but rather it was
discovered because of the extensive testing done here.

Author: felipepiovezan

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.cpp

@@ -37,9 +37,11 @@
 #include "lldb/Interpreter/CommandObject.h"
 #include "lldb/Interpreter/CommandObjectMultiword.h"
 #include "lldb/Interpreter/CommandReturnObject.h"
+#include "lldb/Symbol/FuncUnwinders.h"
 #include "lldb/Symbol/Function.h"
 #include "lldb/Symbol/VariableList.h"
 #include "lldb/Target/RegisterContext.h"
+#include "lldb/Target/UnwindLLDB.h"
 #include "lldb/Utility/LLDBLog.h"
 #include "lldb/Utility/Log.h"
 #include "lldb/Utility/OptionParsing.h"
@@ -2584,37 +2586,111 @@ lldb::addr_t SwiftLanguageRuntime::GetAsyncContext(RegisterContext *regctx) {
   return LLDB_INVALID_ADDRESS;
 }
 
-/// Creates an expression accessing *(fp - 8) or **(fp - 8) if
-/// `with_double_deref` is true. This is only valid for x86_64 or aarch64.
-llvm::ArrayRef<uint8_t>
-GetAsyncRegFromFramePointerDWARFExpr(llvm::Triple::ArchType triple,
-                                     bool with_double_deref) {
-  assert(triple == llvm::Triple::x86_64 || triple == llvm::Triple::aarch64);
-
-  // These expressions must have static storage, due to how UnwindPlan::Row
-  // works.
-  static const uint8_t g_cfa_dwarf_expression_x86_64[] = {
-      llvm::dwarf::DW_OP_breg6, // DW_OP_breg6, register 6 == rbp
-      0x78,                     //    sleb128 -8 (ptrsize)
-      llvm::dwarf::DW_OP_deref,
-      llvm::dwarf::DW_OP_deref,
-  };
-  static const uint8_t g_cfa_dwarf_expression_arm64[] = {
-      llvm::dwarf::DW_OP_breg29, // DW_OP_breg29, register 29 == fp
-      0x78,                      //    sleb128 -8 (ptrsize)
-      llvm::dwarf::DW_OP_deref,
-      llvm::dwarf::DW_OP_deref,
-  };
+/// Functional wrapper to read a register as an address.
+static std::optional<addr_t> ReadRegisterAsAddress(RegisterContext &regctx,
+                                                   unsigned regnum) {
+  auto reg = regctx.ReadRegisterAsUnsigned(regnum, LLDB_INVALID_ADDRESS);
+  if (reg != LLDB_INVALID_ADDRESS)
+    return reg;
+  return {};
+}
+
+/// Functional wrapper to read a pointer from process memory at `addr +
+/// offset`.
+static std::optional<addr_t>
+ReadPtrFromAddr(Process &process, std::optional<addr_t> addr, int offset = 0) {
+  if (!addr)
+    return {};
+  Status error;
+  addr_t ptr = process.ReadPointerFromMemory(*addr + offset, error);
+  if (ptr != LLDB_INVALID_ADDRESS)
+    return ptr;
+  return {};
+}
 
-  const uint8_t *expr = triple == llvm::Triple::x86_64
-                            ? g_cfa_dwarf_expression_x86_64
-                            : g_cfa_dwarf_expression_arm64;
-  auto size = triple == llvm::Triple::x86_64
-                  ? sizeof(g_cfa_dwarf_expression_x86_64)
-                  : sizeof(g_cfa_dwarf_expression_arm64);
-  if (with_double_deref)
-    return llvm::ArrayRef<uint8_t>(expr, size);
-  return llvm::ArrayRef<uint8_t>(expr, size - 1);
+/// Computes the Canonical Frame Address (CFA) by converting the abstract
+/// location of UnwindPlan::Row::FAValue into a concrete address. This is a
+/// simplified version of the methods in RegisterContextUnwind, since plumbing
+/// access to those here would be challenging.
+static std::optional<addr_t> GetCFA(Process &process, RegisterContext &regctx,
+                                    UnwindPlan::Row::FAValue cfa_loc) {
+  using ValueType = UnwindPlan::Row::FAValue::ValueType;
+  switch (cfa_loc.GetValueType()) {
+  case ValueType::isRegisterPlusOffset: {
+    unsigned regnum = cfa_loc.GetRegisterNumber();
+    if (std::optional<addr_t> regvalue = ReadRegisterAsAddress(regctx, regnum))
+      return *regvalue + cfa_loc.GetOffset();
+    break;
+  }
+  case ValueType::isConstant:
+  case ValueType::isDWARFExpression:
+  case ValueType::isRaSearch:
+  case ValueType::isRegisterDereferenced:
+  case ValueType::unspecified:
+    break;
+  }
+  return {};
+}
+
+/// Attempts to use UnwindPlans that inspect assembly to recover the entry value
+/// of the async context register. This is a simplified version of the methods
+/// in RegisterContextUnwind, since plumbing access to those here would be
+/// challenging.
+static std::optional<addr_t> ReadAsyncContextRegisterFromUnwind(
+    SymbolContext &sc, Process &process, Address pc, Address func_start_addr,
+    RegisterContext &regctx, AsyncUnwindRegisterNumbers regnums) {
+  FuncUnwindersSP unwinders =
+      pc.GetModule()->GetUnwindTable().GetFuncUnwindersContainingAddress(pc,
+                                                                         sc);
+  if (!unwinders)
+    return {};
+
+  Target &target = process.GetTarget();
+  UnwindPlanSP unwind_plan =
+      unwinders->GetUnwindPlanAtNonCallSite(target, regctx.GetThread());
+  if (!unwind_plan)
+    return {};
+
+  UnwindPlan::RowSP row = unwind_plan->GetRowForFunctionOffset(
+      pc.GetFileAddress() - func_start_addr.GetFileAddress());
+  UnwindPlan::Row::AbstractRegisterLocation regloc;
+
+  // If the plan doesn't have information about the async register, we can use
+  // its current value, as this is a callee saved register.
+  if (!row->GetRegisterInfo(regnums.async_ctx_regnum, regloc))
+    return ReadRegisterAsAddress(regctx, regnums.async_ctx_regnum);
+
+  // Handle the few abstract locations we are likely to encounter.
+  using RestoreType = UnwindPlan::Row::AbstractRegisterLocation::RestoreType;
+  RestoreType loctype = regloc.GetLocationType();
+  switch (loctype) {
+  case RestoreType::same:
+  case RestoreType::inOtherRegister: {
+    unsigned regnum = loctype == RestoreType::same ? regnums.async_ctx_regnum
+                                                   : regloc.GetRegisterNumber();
+    return ReadRegisterAsAddress(regctx, regnum);
+  }
+  case RestoreType::atCFAPlusOffset: {
+    std::optional<addr_t> cfa = GetCFA(process, regctx, row->GetCFAValue());
+    return ReadPtrFromAddr(process, cfa, regloc.GetOffset());
+  }
+  case RestoreType::isCFAPlusOffset: {
+    std::optional<addr_t> cfa = GetCFA(process, regctx, row->GetCFAValue());
+    if (!cfa)
+      return {};
+    return *cfa + regloc.GetOffset();
+  }
+  case RestoreType::isConstant:
+    return regloc.GetConstant();
+  case RestoreType::unspecified:
+  case RestoreType::undefined:
+  case RestoreType::atAFAPlusOffset:
+  case RestoreType::isAFAPlusOffset:
+  case RestoreType::isDWARFExpression:
+  case RestoreType::atDWARFExpression:
+    return {};
+  }
+  return {};
 }
 
 // Examine the register state and detect the transition from a real
@@ -2644,12 +2720,6 @@ SwiftLanguageRuntime::GetRuntimeUnwindPlan(ProcessSP process_sp,
     return UnwindPlanSP();
   }
 
-  // If we're in the prologue of a function, don't provide a Swift async
-  // unwind plan.  We can be tricked by unmodified caller-registers that
-  // make this look like an async frame when this is a standard ABI function
-  // call, and the parent is the async frame.
-  // This assumes that the frame pointer register will be modified in the
-  // prologue.
   Address pc;
   pc.SetLoadAddress(regctx->GetPC(), &target);
   SymbolContext sc;
@@ -2659,111 +2729,56 @@ SwiftLanguageRuntime::GetRuntimeUnwindPlan(ProcessSP process_sp,
       return UnwindPlanSP();
 
   Address func_start_addr;
-  uint32_t prologue_size;
   ConstString mangled_name;
   if (sc.function) {
     func_start_addr = sc.function->GetAddressRange().GetBaseAddress();
-    prologue_size = sc.function->GetPrologueByteSize();
     mangled_name = sc.function->GetMangled().GetMangledName();
   } else if (sc.symbol) {
     func_start_addr = sc.symbol->GetAddress();
-    prologue_size = sc.symbol->GetPrologueByteSize();
     mangled_name = sc.symbol->GetMangled().GetMangledName();
   } else {
     return UnwindPlanSP();
   }
 
-  AddressRange prologue_range(func_start_addr, prologue_size);
-  bool in_prologue = (func_start_addr == pc ||
-                      prologue_range.ContainsLoadAddress(pc, &target));
-
-  if (in_prologue) {
-    if (!IsAnySwiftAsyncFunctionSymbol(mangled_name.GetStringRef()))
-      return UnwindPlanSP();
-  } else {
-    addr_t saved_fp = LLDB_INVALID_ADDRESS;
-    Status error;
-    if (!process_sp->ReadMemory(fp, &saved_fp, 8, error))
-      return UnwindPlanSP();
-
-    // Get the high nibble of the dreferenced fp; if the 60th bit is set,
-    // this is the transition to a swift async AsyncContext chain.
-    if ((saved_fp & (0xfULL << 60)) >> 60 != 1)
-      return UnwindPlanSP();
-  }
+  if (!IsAnySwiftAsyncFunctionSymbol(mangled_name.GetStringRef()))
+    return UnwindPlanSP();
 
-  // The coroutine funclets split from an async function have 2 different ABIs:
-  //  - Async suspend partial functions and the first funclet get their async
-  //    context directly in the async register.
-  //  - Async await resume partial functions take their context indirectly, it
-  //    needs to be dereferenced to get the actual function's context.
-  // The debug info for locals reflects this difference, so our unwinding of the
-  // context register needs to reflect it too.
+  // The async register contains, at the start of the funclet:
+  // 1. The async context of the async function that just finished executing,
+  // for await resume ("Q") funclets ("indirect context").
+  // 2. The async context for the currently executing async function, for all
+  // other funclets ("Y" and "Yx" funclets, where "x" is a number).
   bool indirect_context =
       IsSwiftAsyncAwaitResumePartialFunctionSymbol(mangled_name.GetStringRef());
 
+  std::optional<addr_t> async_reg = ReadAsyncContextRegisterFromUnwind(
+      sc, *process_sp, pc, func_start_addr, *regctx, *regnums);
+  std::optional<addr_t> async_ctx =
+      indirect_context ? ReadPtrFromAddr(*m_process, async_reg) : async_reg;
+  if (!async_reg || !async_ctx)
+    return UnwindPlanSP();
+
   UnwindPlan::RowSP row(new UnwindPlan::Row);
   const int32_t ptr_size = 8;
   row->SetOffset(0);
 
-  if (in_prologue) {
-    if (indirect_context)
-      row->GetCFAValue().SetIsRegisterDereferenced(regnums->async_ctx_regnum);
-    else
-      row->GetCFAValue().SetIsRegisterPlusOffset(regnums->async_ctx_regnum, 0);
-  } else {
-    // In indirect funclets, dereferencing (fp-8) once produces the CFA of the
-    // frame above. Dereferencing twice will produce the current frame's CFA.
-    bool with_double_deref = indirect_context;
-    llvm::ArrayRef<uint8_t> expr = GetAsyncRegFromFramePointerDWARFExpr(
-        arch.GetMachine(), with_double_deref);
-    row->GetCFAValue().SetIsDWARFExpression(expr.data(), expr.size());
-  }
-
-  if (indirect_context) {
-    if (in_prologue) {
-      row->SetRegisterLocationToSame(regnums->async_ctx_regnum, false);
-    } else {
-      llvm::ArrayRef<uint8_t> expr = GetAsyncRegFromFramePointerDWARFExpr(
-          arch.GetMachine(), false /*with_double_deref*/);
-      row->SetRegisterLocationToIsDWARFExpression(
-          regnums->async_ctx_regnum, expr.data(), expr.size(), false);
-    }
-  } else {
-    // In the first part of a split async function, the context is passed
-    // directly, so we can use the CFA value directly.
-    row->SetRegisterLocationToIsCFAPlusOffset(regnums->async_ctx_regnum, 0,
-                                              false);
-    // The fact that we are in this case needs to be communicated to the frames
-    // below us as they need to react differently. There is no good way to
-    // expose this, so we set another dummy register to communicate this state.
-    static const uint8_t g_dummy_dwarf_expression[] = {
-        llvm::dwarf::DW_OP_const1u, 0
-    };
-    row->SetRegisterLocationToIsDWARFExpression(
-        regnums->dummy_regnum, g_dummy_dwarf_expression,
-        sizeof(g_dummy_dwarf_expression), false);
-  }
-
-  std::optional<addr_t> pc_after_prologue = [&]() -> std::optional<addr_t> {
-    // In the prologue, use the async_reg as is, it has not been clobbered.
-    if (in_prologue)
-      return TrySkipVirtualParentProlog(GetAsyncContext(regctx), *process_sp,
-                                        indirect_context);
-
-    // Both ABIs (x86_64 and aarch64) guarantee the async reg is saved at:
-    // *(fp - 8).
-    Status error;
-    addr_t async_reg_entry_value = LLDB_INVALID_ADDRESS;
-    process_sp->ReadMemory(fp - ptr_size, &async_reg_entry_value, ptr_size,
-                           error);
-    if (error.Fail())
-      return {};
-    return TrySkipVirtualParentProlog(async_reg_entry_value, *process_sp,
-                                      indirect_context);
-  }();
+  // The CFA of a funclet is its own async context.
+  row->GetCFAValue().SetIsConstant(*async_ctx);
+
+  // The value of the async register in the parent frame is the entry value of
+  // the async register in the current frame. This mimics a function call, as
+  // if the parent funclet had called the current funclet.
+  row->SetRegisterLocationToIsConstant(regnums->async_ctx_regnum, *async_reg,
+                                       /*can_replace=*/false);
+
+  // The parent frame needs to know how to interpret the value it is given for
+  // its own async register. A dummy register is used to communicate that.
+  if (!indirect_context)
+    row->SetRegisterLocationToIsConstant(regnums->dummy_regnum, 0,
+                                         /*can_replace=*/false);
 
-  if (pc_after_prologue)
+  if (std::optional<addr_t> pc_after_prologue =
+          TrySkipVirtualParentProlog(*async_ctx, *process_sp))
     row->SetRegisterLocationToIsConstant(regnums->pc_regnum, *pc_after_prologue,
                                          false);
   else

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.h

@@ -496,7 +496,7 @@ class SwiftLanguageRuntime : public LanguageRuntime {
   /// address of the first non-prologue instruction.
   std::optional<lldb::addr_t>
   TrySkipVirtualParentProlog(lldb::addr_t async_reg_val, Process &process,
-                             unsigned num_indirections);
+                             unsigned num_indirections = 0);
 };
 
 } // namespace lldb_private

lldb/test/API/lang/swift/async/stepping/step-in/task-switch/TestSwiftTaskSwitch.py

@@ -7,7 +7,6 @@
 class TestCase(lldbtest.TestBase):
     @swiftTest
     @skipIf(oslist=["windows", "linux"])
-    @skipIf(bugnumber="rdar://136083188")
     def test(self):
         """Test conditions for async step-in."""
         self.build()
@@ -28,7 +27,7 @@ def test(self):
         # Using the line table, build a set of the non-zero line numbers for
         # this this function - and verify that there is exactly one line.
         lines = {inst.addr.line_entry.line for inst in instructions}
-        lines.remove(0)
+        lines.discard(0)
         self.assertEqual(lines, {3})
 
         # Required for builds that have debug info.

lldb/test/API/lang/swift/async/unwind/unwind_in_all_instructions/Makefile

@@ -0,0 +1,3 @@
+SWIFT_SOURCES := main.swift
+SWIFTFLAGS_EXTRAS := -parse-as-library
+include Makefile.rules