[lldb][swift] Fix StackID comparisons to enable stepping out of async functions

lldb/source/Target/StackID.cpp

@@ -13,18 +13,20 @@
 #include "lldb/Target/MemoryRegionInfo.h"
 #include "lldb/Target/Process.h"
 #include "lldb/Target/Thread.h"
+#include "lldb/Utility/LLDBLog.h"
 #include "lldb/Utility/Stream.h"
 
 using namespace lldb_private;
 
 bool StackID::IsCFAOnStack(Process &process) const {
   if (m_cfa_on_stack == eLazyBoolCalculate) {
-    m_cfa_on_stack = eLazyBoolNo;
+    // Conservatively assume stack memory
+    m_cfa_on_stack = eLazyBoolYes;
     if (m_cfa != LLDB_INVALID_ADDRESS) {
       MemoryRegionInfo mem_info;
       if (process.GetMemoryRegionInfo(m_cfa, mem_info).Success())
-        if (mem_info.IsStackMemory() == MemoryRegionInfo::eYes)
-          m_cfa_on_stack = eLazyBoolYes;
+        if (mem_info.IsStackMemory() == MemoryRegionInfo::eNo)
+          m_cfa_on_stack = eLazyBoolNo;
     }
   }
   return m_cfa_on_stack == eLazyBoolYes;
@@ -84,14 +86,43 @@ IsYoungerHeapCFAs(const StackID &lhs, const StackID &rhs, Process &process) {
   if (lhs_cfa_on_stack && rhs_cfa_on_stack)
     return HeapCFAComparisonResult::NoOpinion;
 
-  // FIXME: rdar://76119439
-  // At the boundary between an async parent frame calling a regular child
-  // frame, the CFA of the parent async function is a heap addresses, and the
-  // CFA of concrete child function is a stack address. Therefore, if lhs is
-  // on stack, and rhs is not, lhs is considered less than rhs, independent of
-  // address values.
+  // If one of the frames has a CFA on the stack and the other doesn't, we are
+  // at the boundary between an asynchronous and a synchronous function.
+  // Synchronous functions cannot call asynchronous functions, therefore the
+  // synchronous frame is always younger.
   if (lhs_cfa_on_stack && !rhs_cfa_on_stack)
     return HeapCFAComparisonResult::Younger;
+  if (!lhs_cfa_on_stack && rhs_cfa_on_stack)
+    return HeapCFAComparisonResult::Older;
+
+  const lldb::addr_t lhs_cfa = lhs.GetCallFrameAddress();
+  const lldb::addr_t rhs_cfa = rhs.GetCallFrameAddress();
+  // If the cfas are the same, fallback to the usual scope comparison.
+  if (lhs_cfa == rhs_cfa)
+    return HeapCFAComparisonResult::NoOpinion;
+
+  // Both CFAs are on the heap and they are distinct.
+  // LHS is younger if and only if its continuation async context is (directly
+  // or indirectly) RHS. Chase continuation pointers to check this case, until
+  // we hit the end of the chain (parent_ctx == 0) or a safety limit in case of
+  // an invalid continuation chain.
+  auto max_num_frames = 512;
+  for (lldb::addr_t parent_ctx = lhs_cfa; parent_ctx && max_num_frames;
+       max_num_frames--) {
+    Status error;
+    lldb::addr_t old_parent_ctx = parent_ctx;
+    // The continuation's context is the first field of an async context.
+    parent_ctx = process.ReadPointerFromMemory(old_parent_ctx, error);
+    if (error.Fail()) {
+      Log *log = GetLog(LLDBLog::Unwind);
+      LLDB_LOGF(log, "Failed to read parent async context of: 0x%8.8" PRIx64,
+                old_parent_ctx);
+      break;
+    }
+    if (parent_ctx == rhs_cfa)
+      return HeapCFAComparisonResult::Younger;
+  }
+
   return HeapCFAComparisonResult::NoOpinion;
 }
 // END SWIFT

lldb/test/API/lang/swift/async/stepping/step_out/TestSteppingOutAsync.py

@@ -30,7 +30,6 @@ def step_out_checks(self, thread, expected_func_names):
 
     @swiftTest
     @skipIf(oslist=["windows", "linux"])
-    @skipIf("rdar://133849022")
     def test(self):
         """Test `frame variable` in async functions"""
         self.build()

lldb/unittests/StackID/StackIDTest.cpp

@@ -39,11 +39,19 @@ class StackIDTest : public ::testing::Test {
 };
 
 struct MockProcess : Process {
+  llvm::DenseMap<addr_t, addr_t> memory_map;
+
   MockProcess(TargetSP target_sp, ListenerSP listener_sp)
       : Process(target_sp, listener_sp) {}
+  MockProcess(TargetSP target_sp, ListenerSP listener_sp,
+              llvm::DenseMap<addr_t, addr_t> &&memory_map)
+      : Process(target_sp, listener_sp), memory_map(memory_map) {}
   size_t DoReadMemory(addr_t vm_addr, void *buf, size_t size,
                       Status &error) override {
-    return 0;
+    assert(memory_map.contains(vm_addr));
+    assert(size == sizeof(addr_t));
+    *reinterpret_cast<addr_t *>(buf) = memory_map[vm_addr];
+    return sizeof(addr_t);
   }
   size_t ReadMemory(addr_t addr, void *buf, size_t size,
                     Status &status) override {
@@ -86,8 +94,32 @@ TEST_F(StackIDTest, StackHeapCFAComparison) {
   MockStackID cfa_on_heap(/*cfa*/ 10, OnStack::No);
 
   EXPECT_TRUE(StackID::IsYounger(cfa_on_stack, cfa_on_heap, process));
+  EXPECT_FALSE(StackID::IsYounger(cfa_on_heap, cfa_on_stack, process));
+}
+
+TEST_F(StackIDTest, HeapHeapCFAComparison) {
+  // Create a mock async continuation chain:
+  // 100 -> 108 -> 116 -> 0
+  // This should be read as:
+  // "Async context whose address is 100 has a continuation context whose
+  // address is 108", etc.
+  llvm::DenseMap<addr_t, addr_t> memory_map;
+  memory_map[100] = 108;
+  memory_map[108] = 116;
+  memory_map[116] = 0;
+  auto process = MockProcess(m_target_sp, Listener::MakeListener("dummy"),
+                             std::move(memory_map));
+
+  MockStackID oldest_cfa(/*cfa*/ 116, OnStack::No);
+  MockStackID middle_cfa(/*cfa*/ 108, OnStack::No);
+  MockStackID youngest_cfa(/*cfa*/ 100, OnStack::No);
+
+  EXPECT_TRUE(StackID::IsYounger(youngest_cfa, oldest_cfa, process));
+  EXPECT_FALSE(StackID::IsYounger(oldest_cfa, youngest_cfa, process));
+
+  EXPECT_TRUE(StackID::IsYounger(youngest_cfa, middle_cfa, process));
+  EXPECT_FALSE(StackID::IsYounger(middle_cfa, youngest_cfa, process));
 
-  // FIXME: if the above returned true, swapping the arguments **must** return
-  // false. And yet it doesn't.
-  // EXPECT_FALSE(StackID::IsYounger(cfa_on_heap, cfa_on_stack, process));
+  EXPECT_TRUE(StackID::IsYounger(middle_cfa, oldest_cfa, process));
+  EXPECT_FALSE(StackID::IsYounger(oldest_cfa, middle_cfa, process));
 }