[lldb] Change swift unwind heuristic for Q funclets

Q funclets have an ambiguity region where the debugger doesn't know what
the meaning of the x22 register is. Today, it works around this by
sacrificing the ability to unwind recursive async functions. This patch
changes the heuristic to, instead, sacrifice unwinding in a few
instructions after the prologue, and before any branches. Such
instructions should never be hit in the course of stepping or setting
breakpoints, a user would need to go out of their way to stop in one of
them.

The new approach is:
1. If we are inside the prologue, or the first instruction after the
   prologue, assume x22 contains the indirect context.
2. Otherwise, assume it is the direct context.

This approach fails for the few instructions shown below.

This is what the assembly looks like for Q funclets in x86, with the
first non-prologue instruction indicated by the arrow:

```
    0x100004140 <+0>:  orq    0x3ee9(%rip), %rbp ; (void *)0x1000000000000000
    0x100004147 <+7>:  pushq  %rbp
    0x100004148 <+8>:  pushq  %r14
    0x10000414a <+10>: leaq   0x8(%rsp), %rbp
    0x10000414f <+15>: subq   $0x38, %rsp
->  0x100004153 <+19>: movq   (%r14), %rax
    0x100004156 <+22>: movq   %rax, -0x30(%rbp)   << Fail zone start, inclusive
    0x10000415a <+26>: movq   %rbp, %rcx
    0x10000415d <+29>: subq   $0x8, %rcx
    0x100004161 <+33>: movq   %rax, (%rcx)        << Fail zone end, inclusive
    0x100004164 <+36>: movq   0x50(%rax), %rdi
    0x100004168 <+40>: movq   (%r14), %rcx
    0x10000416b <+43>: movq   %rbp, %rdx
    0x10000416e <+46>: subq   $0x8, %rdx
    0x100004172 <+50>: movq   %rcx, (%rdx)
    0x100004175 <+53>: movq   %rcx, 0x30(%rax)
    0x100004179 <+57>: callq  0x100006036    ; symbol stub for: swift_task_dealloc
    ...
```

The same assembly for arm:

```
    0x100000c54 <+0>:   orr    x29, x29, #0x1000000000000000
    0x100000c58 <+4>:   sub    sp, sp, #0x40
    0x100000c5c <+8>:   stp    x29, x30, [sp, #0x30]
    0x100000c60 <+12>:  str    x22, [sp, #0x28]
    0x100000c64 <+16>:  add    x29, sp, #0x30
->  0x100000c68 <+20>:  ldr    x9, [x22]
    0x100000c6c <+24>:  str    x9, [sp]         << Fail zone start, inclusive
    0x100000c70 <+28>:  mov    x8, x29
    0x100000c74 <+32>:  sub    x8, x8, #0x8
    0x100000c78 <+36>:  str    x9, [x8]         << Fail zone end, inclusive
    0x100000c7c <+40>:  ldr    x0, [x9, #0x50]
    0x100000c80 <+44>:  ldr    x8, [x22]
    0x100000c84 <+48>:  mov    x10, x29
    0x100000c88 <+52>:  sub    x10, x10, #0x8
    0x100000c8c <+56>:  str    x8, [x10]
    0x100000c90 <+60>:  str    x8, [x9, #0x30]
    0x100000c94 <+64>:  bl     0x100001cf4    ; symbol stub for: swift_task_dealloc
    ...
```

As a result, TestSwiftAsyncUnwindAllInstructions.py was changed to
accept failure in those regions. A new test is added, showcasing a step
operation that previously failed because of bad unwinding; it sets a
breakpoint in a Q funclet and then step over. This is _not_ an
artificially created situation, as a `step out` operation always takes
the user to a Q funclet; being able to step-over from that point is
crucial.

Author: felipepiovezan

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

@@ -2559,31 +2559,29 @@ static llvm::Expected<addr_t> ReadAsyncContextRegisterFromUnwind(
 /// CFA of the currently executing function. This is the case at the start of
 /// "Q" funclets, before the low level code changes the meaning of the async
 /// register to not require the indirection.
-/// This implementation detects the transition point by comparing the
-/// continuation pointer in the async context with the currently executing
-/// funclet given by the SymbolContext sc. If they are the same, the PC is
-/// before the transition point.
-/// FIXME: this fails in some recursive async functions. See: rdar://139676623
+/// The end of the prologue approximates the transition point.
+/// FIXME: In the few instructions between the end of the prologue and the
+/// transition point, this approximation fails. rdar://139676623
 static llvm::Expected<bool> IsIndirectContext(Process &process,
                                               StringRef mangled_name,
-                                              addr_t async_reg,
-                                              SymbolContext &sc) {
+                                              Address pc, SymbolContext &sc) {
   if (!SwiftLanguageRuntime::IsSwiftAsyncAwaitResumePartialFunctionSymbol(
           mangled_name))
     return false;
 
-  llvm::Expected<addr_t> continuation_ptr = ReadPtrFromAddr(
-      process, async_reg, /*offset*/ process.GetAddressByteSize());
-  if (!continuation_ptr)
-    return continuation_ptr.takeError();
-
-  if (sc.function)
-    return sc.function->GetAddressRange().ContainsLoadAddress(
-        *continuation_ptr, &process.GetTarget());
-  assert(sc.symbol);
-  Address continuation_addr;
-  continuation_addr.SetLoadAddress(*continuation_ptr, &process.GetTarget());
-  return sc.symbol->ContainsFileAddress(continuation_addr.GetFileAddress());
+  // This is checked prior to calling this function.
+  assert(sc.function || sc.symbol);
+  uint32_t prologue_size = sc.function ? sc.function->GetPrologueByteSize()
+                                       : sc.symbol->GetPrologueByteSize();
+  Address func_start_addr =
+      sc.function ? sc.function->GetAddressRange().GetBaseAddress()
+                  : sc.symbol->GetAddress();
+  // Include one instruction after the prologue. This is where breakpoints
+  // by function name are set, so it's important to get this point right. This
+  // instruction is exactly at address "base + prologue", so adding 1
+  // in the range will do.
+  AddressRange prologue_range(func_start_addr, prologue_size + 1);
+  return prologue_range.ContainsLoadAddress(pc, &process.GetTarget());
 }
 
 // Examine the register state and detect the transition from a real
@@ -2653,7 +2651,7 @@ SwiftLanguageRuntime::GetRuntimeUnwindPlan(ProcessSP process_sp,
     return log_expected(async_reg.takeError());
 
   llvm::Expected<bool> maybe_indirect_context =
-      IsIndirectContext(*process_sp, mangled_name, *async_reg, sc);
+      IsIndirectContext(*process_sp, mangled_name, pc, sc);
   if (!maybe_indirect_context)
     return log_expected(maybe_indirect_context.takeError());
 

lldb/test/API/lang/swift/async/unwind/unwind_in_all_instructions/TestSwiftAsyncUnwindAllInstructions.py

@@ -62,7 +62,80 @@ def set_breakpoints_all_funclets(self, target):
                 breakpoints.add(bp.GetID())
         return breakpoints
 
+    unwind_fail_range_cache = dict()
+
+    # There are challenges when unwinding Q funclets ("await resume"): LLDB cannot
+    # detect the transition point where x22 stops containing the indirect context,
+    # and instead contains the direct context.
+    # Up to and including the first non-prologue instruction, LLDB correctly assumes
+    # it is the indirect context.
+    # After that, it assume x22 contains the direct context. There are a few
+    # instructions where this is not true; this function computes a range that
+    # includes such instructions, so the test may skip checks while stopped in them.
+    def compute_unwind_fail_range(self, function, target):
+        name = function.GetName()
+        if name in TestCase.unwind_fail_range_cache:
+            return TestCase.unwind_fail_range_cache[name]
+
+        if "await resume" not in function.GetName():
+            TestCase.unwind_fail_range_cache[name] = range(0)
+            return range(0)
+
+        first_pc_after_prologue = function.GetStartAddress()
+        first_pc_after_prologue.OffsetAddress(function.GetPrologueByteSize())
+        first_bad_instr = None
+        first_good_instr = None
+        for instr in function.GetInstructions(target):
+            instr_addr = instr.GetAddress()
+
+            # The first bad instruction is approximately the second instruction after the prologue
+            # In actuality, it is at some point after that.
+            if first_bad_instr is None and (
+                instr_addr.GetFileAddress() > first_pc_after_prologue.GetFileAddress()
+            ):
+                first_bad_instr = instr
+                continue
+
+            # The first good instr is approximately the branch to swift_task_dealloc.
+            # In actuality, it is at some point before that.
+            if "swift_task_dealloc" in instr.GetComment(target):
+                first_good_instr = instr
+                break
+
+            # If inside the bad range, no branches can be found.
+            # If this happens, this test must fail so we know unwinding will be broken during stepping.
+            if first_bad_instr is not None:
+                # GetControlFlowKind is only implemented for x86.
+                if "x86" in target.GetTriple():
+                    self.assertEqual(
+                        instr.GetControlFlowKind(target),
+                        lldb.eInstructionControlFlowKindOther,
+                        str(instr),
+                    )
+
+        self.assertNotEqual(first_bad_instr, None)
+        self.assertNotEqual(first_good_instr, None)
+
+        fail_range = range(
+            first_bad_instr.GetAddress().GetFileAddress(),
+            first_good_instr.GetAddress().GetFileAddress(),
+        )
+        TestCase.unwind_fail_range_cache[name] = fail_range
+        return fail_range
+
+    def should_skip_Q_funclet(self, thread):
+        current_frame = thread.frames[0]
+        function = current_frame.GetFunction()
+        fail_range = self.compute_unwind_fail_range(
+            function, thread.GetProcess().GetTarget()
+        )
+
+        current_pc = current_frame.GetPCAddress()
+        return current_pc.GetFileAddress() in fail_range
+
     def check_unwind_ok(self, thread, bpid):
+        if self.should_skip_Q_funclet(thread):
+            return
         # Check that we see the virtual backtrace:
         expected_funcnames = [
             "ASYNC___1___",

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

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

lldb/test/API/lang/swift/async/unwind/unwind_recursive_q_funclets/TestSwiftAsyncUnwindRecursiveQFunclets.py

@@ -0,0 +1,44 @@
+import lldb
+from lldbsuite.test.decorators import *
+import lldbsuite.test.lldbtest as lldbtest
+import lldbsuite.test.lldbutil as lldbutil
+
+
+class TestCase(lldbtest.TestBase):
+
+    mydir = lldbtest.TestBase.compute_mydir(__file__)
+
+    unwind_fail_range_cache = dict()
+
+    @swiftTest
+    @skipIf(oslist=["windows", "linux"])
+    def test(self):
+        """Test that the debugger can unwind at all instructions of all funclets"""
+        self.build()
+
+        source_file = lldb.SBFileSpec("main.swift")
+        target, process, thread, bkpt = lldbutil.run_to_name_breakpoint(
+            self, "$s1a9factorialyS2iYaFTQ1_"
+        )
+
+        # Ensure we are on the last factorial call which recurses (n == 1).
+        frame = thread.frames[0]
+        result = frame.EvaluateExpression("n == 1")
+        self.assertSuccess(result.GetError())
+        self.assertEqual(result.GetSummary(), "true")
+
+        # Disable the breakpoint and step over the call.
+        bkpt.SetEnabled(False)
+
+        # Make sure we are still in the frame of n == 1.
+        thread.StepOver()
+        frame = thread.frames[0]
+        result = frame.EvaluateExpression("n == 1")
+        self.assertSuccess(result.GetError())
+        self.assertEqual(result.GetSummary(), "true")
+
+        thread.StepOver()
+        frame = thread.frames[0]
+        result = frame.EvaluateExpression("n == 1")
+        self.assertSuccess(result.GetError())
+        self.assertEqual(result.GetSummary(), "true")

lldb/test/API/lang/swift/async/unwind/unwind_recursive_q_funclets/main.swift

@@ -0,0 +1,14 @@
+func factorial(_ n: Int) async -> Int {
+  if (n == 0) {
+    return 1;
+  }
+  let n1 = await factorial(n - 1)
+  return n * n1
+}
+
+@main struct Main {
+  static func main() async {
+    let result = await factorial(10)
+    print(result)
+  }
+}