[lldb] Accept address in task backtrace and select commands (#10191)

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

@@ -2164,58 +2164,72 @@ class CommandObjectSwift_RefCount : public CommandObjectRaw {
 /// Construct a `ThreadTask` instance for a Task variable contained in the first
 /// argument.
 static llvm::Expected<ThreadSP>
-ThreadForTaskVariable(Args &command, ExecutionContext &exe_ctx) {
+ThreadForTaskArgument(Args &command, ExecutionContext &exe_ctx) {
   if (!exe_ctx.GetFramePtr())
     return llvm::createStringError("no active frame selected");
 
   if (command.empty() || command[0].ref().empty())
     return llvm::createStringError("missing task variable argument");
 
+  StringRef arg = command[0].ref();
+
   StackFrame &frame = exe_ctx.GetFrameRef();
   uint32_t path_options =
       StackFrame::eExpressionPathOptionsAllowDirectIVarAccess;
   VariableSP var_sp;
   Status status;
   ValueObjectSP valobj_sp = frame.GetValueForVariableExpressionPath(
-      command[0].c_str(), eDynamicDontRunTarget, path_options, var_sp, status);
-  if (!valobj_sp)
-    return status.takeError();
+      arg, eDynamicDontRunTarget, path_options, var_sp, status);
 
   addr_t task_ptr = LLDB_INVALID_ADDRESS;
-  ThreadSafeReflectionContext reflection_ctx;
-  if (ValueObjectSP task_obj_sp = valobj_sp->GetChildMemberWithName("_task")) {
-    task_ptr = task_obj_sp->GetValueAsUnsigned(LLDB_INVALID_ADDRESS);
-    if (task_ptr != LLDB_INVALID_ADDRESS)
-      if (auto *runtime = SwiftLanguageRuntime::Get(exe_ctx.GetProcessSP()))
-        reflection_ctx = runtime->GetReflectionContext();
+  if (status.Success() && valobj_sp) {
+    if (auto task_obj_sp = valobj_sp->GetChildMemberWithName("_task"))
+      task_ptr = task_obj_sp->GetValueAsUnsigned(LLDB_INVALID_ADDRESS);
+    if (task_ptr == LLDB_INVALID_ADDRESS)
+      return llvm::createStringError("failed to access Task pointer");
+  } else {
+    // The argument is not a valid variable expression, try parsing it as a
+    // (task) address.
+    if (arg.getAsInteger(0, task_ptr))
+      return status.takeError();
   }
-  if (task_ptr == LLDB_INVALID_ADDRESS || !reflection_ctx)
-    return llvm::createStringError("failed to access Task data from runtime");
 
-  llvm::Expected<ReflectionContextInterface::AsyncTaskInfo> task_info =
-      reflection_ctx->asyncTaskInfo(task_ptr);
-  if (auto error = task_info.takeError())
-    return error;
+  if (auto *runtime = SwiftLanguageRuntime::Get(exe_ctx.GetProcessSP()))
+    if (auto reflection_ctx = runtime->GetReflectionContext()) {
+      if (auto task_info = reflection_ctx->asyncTaskInfo(task_ptr))
+        return ThreadTask::Create(task_info->id, task_info->resumeAsyncContext,
+                                  exe_ctx);
+      else
+        return task_info.takeError();
+    }
 
-  return ThreadTask::Create(task_info->id, task_info->resumeAsyncContext,
-                            exe_ctx);
+  return llvm::createStringError("failed to access Task data from runtime");
 }
 
 class CommandObjectLanguageSwiftTaskBacktrace final
     : public CommandObjectParsed {
 public:
   CommandObjectLanguageSwiftTaskBacktrace(CommandInterpreter &interpreter)
-      : CommandObjectParsed(interpreter, "backtrace",
-                            "Show the backtrace of Swift tasks. See `thread "
-                            "backtrace` for customizing backtrace output.",
-                            "language swift task backtrace <variable-name>") {
-    AddSimpleArgumentList(eArgTypeVarName);
+      : CommandObjectParsed(
+            interpreter, "backtrace",
+            "Show the backtrace of Swift tasks. See `thread "
+            "backtrace` for customizing backtrace output.",
+            "language swift task backtrace <variable-name | address>") {
+    CommandArgumentEntry arg_entry;
+    arg_entry.emplace_back(eArgTypeVarName, eArgRepeatPlain, LLDB_OPT_SET_1);
+    arg_entry.emplace_back(eArgTypeAddress, eArgRepeatPlain, LLDB_OPT_SET_2);
+    m_arguments.push_back(arg_entry);
   }
 
 private:
   void DoExecute(Args &command, CommandReturnObject &result) override {
+    if (command.GetArgumentCount() != 1) {
+      result.AppendError("missing <variable-name> or <address> argument");
+      return;
+    }
+
     llvm::Expected<ThreadSP> thread_task =
-        ThreadForTaskVariable(command, m_exe_ctx);
+        ThreadForTaskArgument(command, m_exe_ctx);
     if (auto error = thread_task.takeError()) {
       result.AppendError(toString(std::move(error)));
       return;
@@ -2235,14 +2249,22 @@ class CommandObjectLanguageSwiftTaskSelect final : public CommandObjectParsed {
             interpreter, "select",
             "Change the currently selected thread to thread representation of "
             "the given Swift Task. See `thread select`.",
-            "language swift task select <variable-name>") {
-    AddSimpleArgumentList(eArgTypeVarName);
+            "language swift task select <variable-name | address>") {
+    CommandArgumentEntry arg_entry;
+    arg_entry.emplace_back(eArgTypeVarName, eArgRepeatPlain, LLDB_OPT_SET_1);
+    arg_entry.emplace_back(eArgTypeAddress, eArgRepeatPlain, LLDB_OPT_SET_2);
+    m_arguments.push_back(arg_entry);
   }
 
 private:
   void DoExecute(Args &command, CommandReturnObject &result) override {
+    if (command.GetArgumentCount() != 1) {
+      result.AppendError("missing <variable-name> or <address> argument");
+      return;
+    }
+
     llvm::Expected<ThreadSP> thread_task =
-        ThreadForTaskVariable(command, m_exe_ctx);
+        ThreadForTaskArgument(command, m_exe_ctx);
     if (auto error = thread_task.takeError()) {
       result.AppendError(toString(std::move(error)));
       return;

lldb/test/API/lang/swift/async/tasks/TestSwiftTaskBacktrace.py

@@ -5,13 +5,27 @@
 
 
 class TestCase(TestBase):
-    def test(self):
+
+    def test_backtrace_task_variable(self):
         self.build()
         lldbutil.run_to_source_breakpoint(
             self, "break here", lldb.SBFileSpec("main.swift")
         )
+        self.do_backtrace("task")
+
+    def test_backtrace_task_address(self):
+        self.build()
+        _, _, thread, _ = lldbutil.run_to_source_breakpoint(
+            self, "break here", lldb.SBFileSpec("main.swift")
+        )
+        frame = thread.frames[0]
+        task = frame.FindVariable("task")
+        task_addr = task.GetChildMemberWithName("address").unsigned
+        self.do_backtrace(task_addr)
+
+    def do_backtrace(self, arg):
         self.expect(
-            "language swift task backtrace task",
+            f"language swift task backtrace {arg}",
             substrs=[
                 ".sleep(",
                 "`second() at main.swift:6",

lldb/test/API/lang/swift/async/tasks/TestSwiftTaskSelect.py

@@ -6,12 +6,25 @@
 
 class TestCase(TestBase):
 
-    def test_backtrace_selected_task(self):
+    def test_backtrace_selected_task_variable(self):
         self.build()
         lldbutil.run_to_source_breakpoint(
             self, "break here", lldb.SBFileSpec("main.swift")
         )
-        self.runCmd("language swift task select task")
+        self.do_backtrace_selected_task("task")
+
+    def test_backtrace_selected_task_address(self):
+        self.build()
+        _, _, thread, _ = lldbutil.run_to_source_breakpoint(
+            self, "break here", lldb.SBFileSpec("main.swift")
+        )
+        frame = thread.frames[0]
+        task = frame.FindVariable("task")
+        task_addr = task.GetChildMemberWithName("address").unsigned
+        self.do_backtrace_selected_task(task_addr)
+
+    def do_backtrace_selected_task(self, arg):
+        self.runCmd(f"language swift task select {arg}")
         self.expect(
             "thread backtrace",
             substrs=[
@@ -22,16 +35,27 @@ def test_backtrace_selected_task(self):
             ],
         )
 
-    def test_navigate_selected_task_stack(self):
+    def test_navigate_stack_of_selected_task_variable(self):
         self.build()
-        _, process, _, _ = lldbutil.run_to_source_breakpoint(
+        _, process, thread, _ = lldbutil.run_to_source_breakpoint(
             self, "break here", lldb.SBFileSpec("main.swift")
         )
-        self.runCmd("language swift task select task")
+        self.do_test_navigate_selected_task_stack(process, "task")
+
+    def test_navigate_stack_of_selected_task_address(self):
+        self.build()
+        _, process, thread, _ = lldbutil.run_to_source_breakpoint(
+            self, "break here", lldb.SBFileSpec("main.swift")
+        )
+        frame = thread.frames[0]
+        task = frame.FindVariable("task")
+        task_addr = task.GetChildMemberWithName("address").unsigned
+        self.do_test_navigate_selected_task_stack(process, task_addr)
+
+    def do_test_navigate_selected_task_stack(self, process, arg):
+        self.runCmd(f"language swift task select {arg}")
+        thread = process.GetSelectedThread()
 
-        thread = process.selected_thread
-        self.assertEqual(thread.id, 2)
-        self.assertEqual(thread.idx, 0xFFFFFFFF)
         self.assertIn(
             "libswift_Concurrency.", thread.GetSelectedFrame().module.file.basename
         )