🍒 [lldb] Skip prologue when stepping through swift_task_switch

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

@@ -2924,31 +2924,15 @@ std::optional<lldb::addr_t> SwiftLanguageRuntime::TrySkipVirtualParentProlog(
   addr_t pc_value = process.ReadPointerFromMemory(pc_location, error);
   if (error.Fail())
     return {};
-  // Clear any high order bits of this code address so that SetLoadAddress works
-  // properly.
-  pc_value = process.FixCodeAddress(pc_value);
 
-  Address pc;
-  Target &target = process.GetTarget();
-  pc.SetLoadAddress(pc_value, &target);
-  if (!pc.IsValid())
-    return {};
-
-  SymbolContext sc;
-  bool sc_ok = pc.CalculateSymbolContext(&sc, eSymbolContextFunction |
-                                                  eSymbolContextSymbol);
-  if (!sc_ok || (!sc.symbol && !sc.function)) {
-    Log *log = GetLog(LLDBLog::Unwind);
-    LLDB_LOGF(log,
-              "SwiftLanguageRuntime::%s Failed to find a symbol context for "
-              "address 0x%" PRIx64,
-              __FUNCTION__, pc_value);
-    return {};
-  }
+  llvm::Expected<uint64_t> maybe_prologue_size =
+      FindPrologueSize(process, pc_value);
+  if (maybe_prologue_size)
+    return pc_value + *maybe_prologue_size;
 
-  auto prologue_size = sc.symbol ? sc.symbol->GetPrologueByteSize()
-                                 : sc.function->GetPrologueByteSize();
-  return pc_value + prologue_size;
+  LLDB_LOG_ERROR(GetLog(LLDBLog::Unwind), maybe_prologue_size.takeError(),
+                 "{1}::{0}", __FUNCTION__);
+  return pc_value;
 }
 
 /// Attempts to read the memory location at `task_addr_location`, producing
@@ -3139,4 +3123,30 @@ llvm::Expected<std::optional<std::string>> GetTaskName(lldb::addr_t task_addr,
   return status.takeError();
 }
 
+llvm::Expected<uint64_t> FindPrologueSize(Process &process,
+                                          uint64_t load_address) {
+  Address addr;
+  Target &target = process.GetTarget();
+  addr.SetLoadAddress(process.FixCodeAddress(load_address), &target);
+  if (!addr.IsValid())
+    return llvm::createStringError(
+        llvm::formatv("Invalid load address for {0:x}", load_address));
+
+  SymbolContext sc;
+  bool sc_ok = addr.CalculateSymbolContext(&sc, eSymbolContextFunction |
+                                                    eSymbolContextSymbol);
+  if (!sc_ok || (!sc.symbol && !sc.function))
+    return llvm::createStringError(llvm::formatv(
+        "Failed to find a symbol context for address {1:x}", load_address));
+
+  uint64_t prologue_size = sc.symbol ? sc.symbol->GetPrologueByteSize()
+                                     : sc.function->GetPrologueByteSize();
+
+  if (prologue_size == 0)
+    return llvm::createStringError(llvm::formatv(
+        "Prologue size is 0 for function {0}",
+        sc.GetFunctionName(Mangled::NamePreference::ePreferMangled)));
+
+  return prologue_size;
+}
 } // namespace lldb_private

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

@@ -944,6 +944,10 @@ class TaskInspector {
 llvm::Expected<std::optional<std::string>> GetTaskName(lldb::addr_t task,
                                                        Process &process);
 
+/// Finds the function for which `load_address` belongs, and returns its
+/// prologue size.
+llvm::Expected<uint64_t> FindPrologueSize(Process &process,
+                                          uint64_t load_address);
 } // namespace lldb_private
 
 #endif // liblldb_SwiftLanguageRuntime_h_

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

@@ -394,7 +394,18 @@ CreateRunThroughTaskSwitchThreadPlan(Thread &thread,
     return {};
 
   resume_fn_ptr = thread.GetProcess()->FixCodeAddress(resume_fn_ptr);
-  return std::make_shared<ThreadPlanRunToAddress>(thread, resume_fn_ptr,
+
+  llvm::Expected<uint64_t> maybe_prologue_size =
+      FindPrologueSize(*thread.GetProcess(), resume_fn_ptr);
+
+  uint64_t dest_address =
+      resume_fn_ptr + (maybe_prologue_size ? *maybe_prologue_size : 0);
+
+  if (!maybe_prologue_size)
+    LLDB_LOG_ERROR(GetLog(LLDBLog::Step), maybe_prologue_size.takeError(),
+                   "{1}::{0}", __FUNCTION__);
+
+  return std::make_shared<ThreadPlanRunToAddress>(thread, dest_address,
                                                   /*stop_others*/ false);
 }
 

lldb/test/API/lang/swift/async/stepping/step_over/TestSwiftAsyncStepOver.py

@@ -7,6 +7,11 @@
 @skipIfAsan  # rdar://138777205
 class TestCase(lldbtest.TestBase):
 
+    def check_x_is_available(self, frame):
+        x_var = frame.FindVariable("x")
+        self.assertTrue(x_var.IsValid(), f"Failed to find x in {frame}")
+        self.assertEqual(x_var.GetValueAsUnsigned(), 30)
+
     def check_is_in_line(self, thread, linenum):
         frame = thread.frames[0]
         line_entry = frame.GetLineEntry()
@@ -25,10 +30,11 @@ def test(self):
         bkpt.SetEnabled(False) # avoid hitting multiple locations in async breakpoints
 
         expected_line_nums = [4]  # print(x)
-        expected_line_nums += [5, 6, 7, 5, 6, 7, 5]  # two runs over the loop
-        expected_line_nums += [8, 9]  # if line + if block
+        expected_line_nums += [5, 6, 7, 8, 5, 6, 7, 8, 5]  # two runs over the loop
+        expected_line_nums += [9, 10]  # if line + if block
         for expected_line_num in expected_line_nums:
             thread.StepOver()
             stop_reason = thread.GetStopReason()
             self.assertStopReason(stop_reason, lldb.eStopReasonPlanComplete)
             self.check_is_in_line(thread, expected_line_num)
+            self.check_x_is_available(thread.frames[0])

lldb/test/API/lang/swift/async/stepping/step_over/main.swift

@@ -4,6 +4,7 @@
     print(x)
     for i in 1...2 {
       await f()
+      print("hello!")
     }
     if (await f() == 30) {
       print("here!")