[lldb] Use unique line count in swift_task_switch logic

lldb/include/lldb/Target/StackID.h

@@ -19,12 +19,7 @@ class StackID {
   // Constructors and Destructors
   StackID() = default;
 
-  explicit StackID(lldb::addr_t pc, lldb::addr_t cfa,
-                   SymbolContextScope *symbol_scope)
-      : m_pc(pc), m_cfa(cfa), m_symbol_scope(symbol_scope) {}
-
-  StackID(const StackID &rhs)
-      : m_pc(rhs.m_pc), m_cfa(rhs.m_cfa), m_symbol_scope(rhs.m_symbol_scope) {}
+  StackID(const StackID &rhs) = default;
 
   ~StackID() = default;
 
@@ -41,6 +36,7 @@ class StackID {
   void Clear() {
     m_pc = LLDB_INVALID_ADDRESS;
     m_cfa = LLDB_INVALID_ADDRESS;
+    m_cfa_on_stack = true;
     m_symbol_scope = nullptr;
   }
 
@@ -55,14 +51,22 @@ class StackID {
     if (this != &rhs) {
       m_pc = rhs.m_pc;
       m_cfa = rhs.m_cfa;
+      m_cfa_on_stack = rhs.m_cfa_on_stack;
       m_symbol_scope = rhs.m_symbol_scope;
     }
     return *this;
   }
 
+  bool IsCFAOnStack() const { return m_cfa_on_stack; }
+
 protected:
   friend class StackFrame;
 
+  explicit StackID(lldb::addr_t pc, lldb::addr_t cfa, lldb::ThreadSP thread_sp)
+      : m_pc(pc), m_cfa(cfa), m_cfa_on_stack(IsStackAddress(cfa, thread_sp)) {}
+
+  bool IsStackAddress(lldb::addr_t addr, lldb::ThreadSP thread_sp) const;
+
   void SetPC(lldb::addr_t pc) { m_pc = pc; }
 
   void SetCFA(lldb::addr_t cfa) { m_cfa = cfa; }
@@ -78,6 +82,9 @@ class StackID {
                             // at the beginning of the function that uniquely
                             // identifies this frame (along with m_symbol_scope
                             // below)
+  // True if the CFA is an address on the stack, false if it's an address
+  // elsewhere (ie heap).
+  bool m_cfa_on_stack = true;
   SymbolContextScope *m_symbol_scope =
       nullptr; // If nullptr, there is no block or symbol for this frame.
                // If not nullptr, this will either be the scope for the

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

@@ -28,6 +28,7 @@
 
 #include "Plugins/Process/Utility/RegisterContext_x86.h"
 #include "Utility/ARM64_DWARF_Registers.h"
+#include "llvm/ADT/SmallSet.h"
 
 using namespace lldb;
 using namespace lldb_private;
@@ -234,7 +235,7 @@ class ThreadPlanStepInAsync : public ThreadPlan {
       return false;
     auto fn_start = sc.symbol->GetFileAddress();
     auto fn_end = sc.symbol->GetFileAddress() + sc.symbol->GetByteSize();
-    int line_entry_count = 0;
+    llvm::SmallSet<uint32_t, 2> unique_debug_lines;
     if (auto *line_table = sc.comp_unit->GetLineTable()) {
       for (uint32_t i = 0; i < line_table->GetSize(); ++i) {
         LineEntry line_entry;
@@ -243,11 +244,23 @@ class ThreadPlanStepInAsync : public ThreadPlan {
             continue;
 
           auto line_start = line_entry.range.GetBaseAddress().GetFileAddress();
-          if (line_start >= fn_start && line_start < fn_end)
-            if (++line_entry_count > 1)
-              // This is an async function with a proper body of code, no step
-              // into `swift_task_switch` required.
+          if (fn_start <= line_start && line_start < fn_end) {
+            unique_debug_lines.insert(line_entry.line);
+            // This logic is to distinguish between async functions that only
+            // call `swift_task_switch` (which, from the perspective of the
+            // user, has no meaningful function body), vs async functions that
+            // do have a function body. In the first case, lldb should step
+            // further to find the function body, in the second case lldb has
+            // found a body and should stop.
+            //
+            // Currently, async functions that go through `swift_task_switch`
+            // are generated with a reference to a single line. If this function
+            // has more than one unique debug line, then it is a function that
+            // has a body, and execution can stop here.
+            if (unique_debug_lines.size() >= 2)
+              // No step into `swift_task_switch` required.
               return false;
+          }
         }
       }
     }

lldb/source/Target/StackFrame.cpp

@@ -57,7 +57,7 @@ StackFrame::StackFrame(const ThreadSP &thread_sp, user_id_t frame_idx,
                        const SymbolContext *sc_ptr)
     : m_thread_wp(thread_sp), m_frame_index(frame_idx),
       m_concrete_frame_index(unwind_frame_index), m_reg_context_sp(),
-      m_id(pc, cfa, nullptr), m_frame_code_addr(pc), m_sc(), m_flags(),
+      m_id(pc, cfa, thread_sp), m_frame_code_addr(pc), m_sc(), m_flags(),
       m_frame_base(), m_frame_base_error(), m_cfa_is_valid(cfa_is_valid),
       m_stack_frame_kind(kind),
       m_behaves_like_zeroth_frame(behaves_like_zeroth_frame),
@@ -83,7 +83,7 @@ StackFrame::StackFrame(const ThreadSP &thread_sp, user_id_t frame_idx,
                        const SymbolContext *sc_ptr)
     : m_thread_wp(thread_sp), m_frame_index(frame_idx),
       m_concrete_frame_index(unwind_frame_index),
-      m_reg_context_sp(reg_context_sp), m_id(pc, cfa, nullptr),
+      m_reg_context_sp(reg_context_sp), m_id(pc, cfa, thread_sp),
       m_frame_code_addr(pc), m_sc(), m_flags(), m_frame_base(),
       m_frame_base_error(), m_cfa_is_valid(true),
       m_stack_frame_kind(StackFrame::Kind::Regular),
@@ -111,7 +111,7 @@ StackFrame::StackFrame(const ThreadSP &thread_sp, user_id_t frame_idx,
       m_concrete_frame_index(unwind_frame_index),
       m_reg_context_sp(reg_context_sp),
       m_id(pc_addr.GetLoadAddress(thread_sp->CalculateTarget().get()), cfa,
-           nullptr),
+           thread_sp),
       m_frame_code_addr(pc_addr), m_sc(), m_flags(), m_frame_base(),
       m_frame_base_error(), m_cfa_is_valid(true),
       m_stack_frame_kind(StackFrame::Kind::Regular),

lldb/source/Target/StackID.cpp

@@ -10,14 +10,31 @@
 #include "lldb/Symbol/Block.h"
 #include "lldb/Symbol/Symbol.h"
 #include "lldb/Symbol/SymbolContext.h"
+#include "lldb/Target/MemoryRegionInfo.h"
+#include "lldb/Target/Process.h"
+#include "lldb/Target/Thread.h"
 #include "lldb/Utility/Stream.h"
 
 using namespace lldb_private;
 
+bool StackID::IsStackAddress(lldb::addr_t addr,
+                             lldb::ThreadSP thread_sp) const {
+  if (addr == LLDB_INVALID_ADDRESS)
+    return false;
+
+  if (thread_sp)
+    if (auto process_sp = thread_sp->GetProcess()) {
+      MemoryRegionInfo mem_info;
+      if (process_sp->GetMemoryRegionInfo(addr, mem_info).Success())
+        return mem_info.IsStackMemory() == MemoryRegionInfo::eYes;
+    }
+  return true; // assumed
+}
+
 void StackID::Dump(Stream *s) {
   s->Printf("StackID (pc = 0x%16.16" PRIx64 ", cfa = 0x%16.16" PRIx64
-            ", symbol_scope = %p",
-            m_pc, m_cfa, static_cast<void *>(m_symbol_scope));
+            ", cfa_on_stack = %d, symbol_scope = %p",
+            m_pc, m_cfa, m_cfa_on_stack, static_cast<void *>(m_symbol_scope));
   if (m_symbol_scope) {
     SymbolContext sc;
 
@@ -62,10 +79,12 @@ bool lldb_private::operator<(const StackID &lhs, const StackID &rhs) {
   const lldb::addr_t rhs_cfa = rhs.GetCallFrameAddress();
 
   // FIXME: rdar://76119439
-  // This heuristic is a *temporary* fallback while proper fixes are
-  // determined. The heuristic assumes the CFA of async functions is a low
-  // (heap) address, and for normal functions it's a high (stack) address.
-  if (lhs_cfa - rhs_cfa >= 0x00007ff000000000ULL)
+  // Heuristic: 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 addresses. Therefore, if lhs
+  // is on stack, and rhs is not, lhs is considered less than rhs (regardless of
+  // address values).
+  if (lhs.IsCFAOnStack() && !rhs.IsCFAOnStack())
     return true;
 
   // FIXME: We are assuming that the stacks grow downward in memory.  That's not

lldb/test/API/lang/swift/async/stepping/step-in/task-switch/Makefile

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

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

@@ -0,0 +1,35 @@
+import lldb
+from lldbsuite.test.decorators import *
+import lldbsuite.test.lldbtest as lldbtest
+import lldbsuite.test.lldbutil as lldbutil
+
+
+class TestCase(lldbtest.TestBase):
+    @swiftTest
+    @skipIf(oslist=["windows", "linux"])
+    def test(self):
+        """Test conditions for async step-in."""
+        self.build()
+
+        src = lldb.SBFileSpec("main.swift")
+        target, _, thread, _ = lldbutil.run_to_source_breakpoint(self, "await f()", src)
+        self.assertEqual(thread.frame[0].function.mangled, "$s1a5entryO4mainyyYaFZ")
+
+        function = target.FindFunctions("$s1a5entryO4mainyyYaFZTQ0_")[0].function
+        instructions = list(function.GetInstructions(target))
+
+        # Expected to be a trampoline that tail calls `swift_task_switch`.
+        self.assertIn("swift_task_switch", instructions[-1].GetComment(target))
+
+        # 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)
+        self.assertEqual(lines, {3})
+
+        # Required for builds that have debug info.
+        self.runCmd("settings set target.process.thread.step-avoid-libraries libswift_Concurrency.dylib")
+        thread.StepInto()
+        frame = thread.frame[0]
+        # Step in from `main` should progress through to `f`.
+        self.assertEqual(frame.name, "a.f() async -> Swift.Int")

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

@@ -0,0 +1,10 @@
+@main enum entry {
+  static func main() async {
+    let x = await f()
+    print(x)
+  }
+}
+
+func f() async -> Int {
+  return 30
+}