[lldb] Introduce command to select task "threads"

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

@@ -2086,6 +2086,46 @@ 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) {
+  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");
+
+  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();
+
+  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 (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;
+
+  return ThreadTask::Create(task_info->id, task_info->resumeAsyncContext,
+                            exe_ctx);
+}
+
 class CommandObjectLanguageSwiftTaskBacktrace final
     : public CommandObjectParsed {
 public:
@@ -2099,60 +2139,44 @@ class CommandObjectLanguageSwiftTaskBacktrace final
 
 private:
   void DoExecute(Args &command, CommandReturnObject &result) override {
-    if (!m_exe_ctx.GetFramePtr()) {
-      result.AppendError("no active frame selected");
-      return;
-    }
-
-    if (command.empty() || command[0].ref().empty()) {
-      result.AppendError("no task variable");
-      return;
-    }
-
-    StackFrame &frame = m_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) {
-      result.AppendError(status.AsCString());
+    llvm::Expected<ThreadSP> thread_task =
+        ThreadForTaskVariable(command, m_exe_ctx);
+    if (auto error = thread_task.takeError()) {
+      result.AppendError(toString(std::move(error)));
       return;
     }
 
-    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(m_exe_ctx.GetProcessSP()))
-          reflection_ctx = runtime->GetReflectionContext();
-    }
-    if (task_ptr == LLDB_INVALID_ADDRESS || !reflection_ctx) {
-      result.AppendError("failed to access Task data from runtime");
-      return;
-    }
+    // GetStatus prints the backtrace.
+    thread_task.get()->GetStatus(result.GetOutputStream(), 0, UINT32_MAX, 0,
+                                 false, true);
+    result.SetStatus(lldb::eReturnStatusSuccessFinishResult);
+  }
+};
 
-    llvm::Expected<ReflectionContextInterface::AsyncTaskInfo> task_info =
-        reflection_ctx->asyncTaskInfo(task_ptr);
-    if (auto error = task_info.takeError()) {
-      result.AppendError(toString(std::move(error)));
-      return;
-    }
+class CommandObjectLanguageSwiftTaskSelect final : public CommandObjectParsed {
+public:
+  CommandObjectLanguageSwiftTaskSelect(CommandInterpreter &interpreter)
+      : 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);
+  }
 
-    auto thread_task = ThreadTask::Create(
-        task_info->id, task_info->resumeAsyncContext, m_exe_ctx);
+private:
+  void DoExecute(Args &command, CommandReturnObject &result) override {
+    llvm::Expected<ThreadSP> thread_task =
+        ThreadForTaskVariable(command, m_exe_ctx);
     if (auto error = thread_task.takeError()) {
       result.AppendError(toString(std::move(error)));
       return;
     }
 
-    // GetStatus prints the backtrace.
-    thread_task.get()->GetStatus(result.GetOutputStream(), 0, UINT32_MAX, 0,
-                                 false, true);
+    auto &thread_list = m_exe_ctx.GetProcessRef().GetThreadList();
+    thread_list.AddThread(thread_task.get());
+    thread_list.SetSelectedThreadByID(thread_task.get()->GetID());
+
     result.SetStatus(lldb::eReturnStatusSuccessFinishResult);
   }
 };
@@ -2166,6 +2190,9 @@ class CommandObjectLanguageSwiftTask final : public CommandObjectMultiword {
     LoadSubCommand("backtrace",
                    CommandObjectSP(new CommandObjectLanguageSwiftTaskBacktrace(
                        interpreter)));
+    LoadSubCommand(
+        "select",
+        CommandObjectSP(new CommandObjectLanguageSwiftTaskSelect(interpreter)));
   }
 };
 

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

@@ -0,0 +1,75 @@
+import lldb
+from lldbsuite.test.decorators import *
+from lldbsuite.test.lldbtest import TestBase
+import lldbsuite.test.lldbutil as lldbutil
+
+
+class TestCase(TestBase):
+
+    def test_backtrace_selected_task(self):
+        self.build()
+        lldbutil.run_to_source_breakpoint(
+            self, "break here", lldb.SBFileSpec("main.swift")
+        )
+        self.runCmd("language swift task select task")
+        self.expect(
+            "thread backtrace",
+            substrs=[
+                ".sleep(",
+                "`second() at main.swift:6:",
+                "`first() at main.swift:2:",
+                "`closure #1 in static Main.main() at main.swift:12:",
+            ],
+        )
+
+    def test_navigate_selected_task_stack(self):
+        self.build()
+        # target, process, thread, bkpt
+        _, process, _, _ = lldbutil.run_to_source_breakpoint(
+            self, "break here", lldb.SBFileSpec("main.swift")
+        )
+        self.runCmd("language swift task select task")
+
+        self.expect(
+            "thread list",
+            substrs=[
+                "* thread #4294967295: tid = 0x0002,",
+                "libswift_Concurrency.dylib",
+                "._sleep(",
+            ],
+        )
+
+        frame_index = 0
+        for frame in process.selected_thread:
+            if "`second()" in str(frame):
+                frame_index = frame.idx
+        self.assertNotEqual(frame_index, -1)
+
+        self.expect(
+            f"frame select {frame_index}",
+            substrs=[
+                f"frame #{frame_index}:",
+                "`second() at main.swift:6:",
+                "   5   \tfunc second() async {",
+                "-> 6   \t    try? await Task.sleep(for: .seconds(10))",
+            ],
+        )
+
+        self.expect(
+            "up",
+            substrs=[
+                f"frame #{frame_index + 1}:",
+                "`first() at main.swift:2:",
+                "   1   \tfunc first() async {",
+                "-> 2   \t    await second()",
+            ],
+        )
+
+        self.expect(
+            "up",
+            substrs=[
+                f"frame #{frame_index + 2}:",
+                "`closure #1 in static Main.main() at main.swift:12:",
+                "-> 12  \t            await first()",
+            ],
+        )