[lldb][NFC] Prepare SwiftLanguageRuntime for vectorized memory reads

Once the proposal [1] is implemented, the Process class will have the
ability to read multiple memory addresses at once. This patch is an NFC
refactor to make the code more amenable to that future: methods are
rewritten work with vectors of threads and addresses, instead of
individual values.

[1]: https://discourse.llvm.org/t/rfc-a-new-vectorized-memory-read-packet/

Author: felipepiovezan

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

@@ -2915,19 +2915,6 @@ std::optional<lldb::addr_t> SwiftLanguageRuntime::TrySkipVirtualParentProlog(
   return pc_value;
 }
 
-/// Attempts to read the memory location at `task_addr_location`, producing
-/// the Task pointer if possible.
-static llvm::Expected<lldb::addr_t>
-ReadTaskAddr(lldb::addr_t task_addr_location, Process &process) {
-  Status status;
-  addr_t task_addr = process.ReadPointerFromMemory(task_addr_location, status);
-  if (status.Fail())
-    return llvm::joinErrors(
-        llvm::createStringError("could not get current task from thread"),
-        status.takeError());
-  return task_addr;
-}
-
 /// Compute the location where the Task pointer for `real_thread` is stored by
 /// the runtime.
 static llvm::Expected<lldb::addr_t>
@@ -2955,48 +2942,133 @@ ComputeTaskAddrLocationFromThreadLocalStorage(Thread &real_thread) {
 #endif
 }
 
+/// Helper function to read all `pointers` from process memory at once.
+/// Consumes any errors from the input by propagating them to the output.
+static llvm::SmallVector<llvm::Expected<addr_t>> MultiReadPointers(
+    Process &process,
+    llvm::MutableArrayRef<llvm::Expected<addr_t>> maybe_pointers) {
+  llvm::SmallVector<llvm::Expected<addr_t>> final_results;
+  llvm::SmallVector<addr_t> to_read;
+
+  for (llvm::Expected<addr_t> &maybe_ptr : maybe_pointers) {
+    if (!maybe_ptr)
+      final_results.emplace_back(maybe_ptr.takeError());
+    else {
+      final_results.push_back(LLDB_INVALID_ADDRESS);
+      to_read.push_back(*maybe_ptr);
+    }
+  }
+
+  /// TODO: convert this loop into a call to the vectorized memory read, once
+  /// that is available in Process.
+  llvm::SmallVector<llvm::Expected<addr_t>> read_results;
+  for (addr_t pointer : to_read) {
+    Status status;
+    addr_t result = process.ReadPointerFromMemory(pointer, status);
+    if (status.Fail())
+      read_results.push_back(status.takeError());
+    else
+      read_results.push_back(result);
+  }
+
+  llvm::MutableArrayRef<llvm::Expected<addr_t>> results_ref = read_results;
+  for (llvm::Expected<addr_t> &maybe_result : final_results) {
+    // Skip over errors propagated from the input.
+    if (!maybe_result)
+      continue;
+    maybe_result = std::move(results_ref.consume_front());
+  }
+
+  return final_results;
+}
+
+/// Helper function to read `addr` from process memory. Errors in the input are
+/// propagated to to the output.
+static llvm::Expected<addr_t> ReadPointer(Process &process,
+                                          llvm::Expected<addr_t> addr) {
+  auto read_result = MultiReadPointers(process, addr);
+  return std::move(read_result[0]);
+}
+
 llvm::Expected<lldb::addr_t>
 TaskInspector::GetTaskAddrFromThreadLocalStorage(Thread &thread) {
-  // Look through backing threads when inspecting TLS.
-  Thread &real_thread =
-      thread.GetBackingThread() ? *thread.GetBackingThread() : thread;
+  return std::move(GetTaskAddrFromThreadLocalStorage(&thread)[0]);
+}
+
+llvm::SmallVector<llvm::Expected<lldb::addr_t>>
+TaskInspector::GetTaskAddrLocations(llvm::ArrayRef<Thread *> threads) {
+  llvm::SmallVector<llvm::Expected<addr_t>> addr_locations;
+  addr_locations.reserve(threads.size());
 
-  if (auto it = m_tid_to_task_addr_location.find(real_thread.GetID());
-      it != m_tid_to_task_addr_location.end()) {
+  for (auto [idx, thread] : llvm::enumerate(threads)) {
+    Thread &real_thread =
+        thread->GetBackingThread() ? *thread->GetBackingThread() : *thread;
+
+    auto it = m_tid_to_task_addr_location.find(real_thread.GetID());
+    if (it != m_tid_to_task_addr_location.end()) {
+      addr_locations.push_back(it->second);
 #ifndef NDEBUG
-    // In assert builds, check that caching did not produce incorrect results.
-    llvm::Expected<lldb::addr_t> task_addr_location =
-        ComputeTaskAddrLocationFromThreadLocalStorage(real_thread);
-    assert(task_addr_location);
-    assert(it->second == *task_addr_location);
+      // In assert builds, check that caching did not produce incorrect results.
+      llvm::Expected<lldb::addr_t> task_addr_location =
+          ComputeTaskAddrLocationFromThreadLocalStorage(real_thread);
+      assert(task_addr_location);
+      assert(it->second == *task_addr_location);
 #endif
-    llvm::Expected<lldb::addr_t> task_addr =
-        ReadTaskAddr(it->second, *thread.GetProcess());
-    if (task_addr)
-      return task_addr;
-    // If the cached task addr location became invalid, invalidate the cache.
-    m_tid_to_task_addr_location.erase(it);
-    LLDB_LOG_ERROR(GetLog(LLDBLog::OS), task_addr.takeError(),
+      continue;
+    }
+    addr_locations.push_back(
+        ComputeTaskAddrLocationFromThreadLocalStorage(real_thread));
+  }
+  return addr_locations;
+}
+
+llvm::SmallVector<llvm::Expected<addr_t>>
+TaskInspector::GetTaskAddrFromThreadLocalStorage(
+    llvm::ArrayRef<Thread *> threads) {
+  if (threads.empty())
+    return {};
+
+  llvm::SmallVector<llvm::Expected<addr_t>> addr_locations =
+      GetTaskAddrLocations(threads);
+
+  Process &process = *threads[0]->GetProcess();
+  llvm::SmallVector<llvm::Expected<addr_t>> mem_read_results =
+      MultiReadPointers(process, addr_locations);
+
+  for (auto [idx, thread] : llvm::enumerate(threads)) {
+    Thread &real_thread =
+        thread->GetBackingThread() ? *thread->GetBackingThread() : *thread;
+    user_id_t tid = real_thread.GetID();
+
+    // If the read was successful, cache the address.
+    if (mem_read_results[idx]) {
+      m_tid_to_task_addr_location[tid] = *addr_locations[idx];
+      continue;
+    }
+
+    // For unsuccessful reads whose address was not cached, don't try again.
+    if (!m_tid_to_task_addr_location.erase(tid))
+      continue;
+
+    LLDB_LOG_ERROR(GetLog(LLDBLog::OS), mem_read_results[idx].takeError(),
                    "TaskInspector: evicted task location address due to "
                    "invalid memory read: {0}");
-  }
 
-  llvm::Expected<lldb::addr_t> task_addr_location =
-      ComputeTaskAddrLocationFromThreadLocalStorage(real_thread);
-  if (!task_addr_location)
-    return task_addr_location;
+    // The cached address could not be loaded. "This should never happen", but
+    // recompute the address and try again for completeness.
+    llvm::Expected<addr_t> task_addr_loc =
+        ComputeTaskAddrLocationFromThreadLocalStorage(real_thread);
+    if (!task_addr_loc) {
+      mem_read_results[idx] = std::move(task_addr_loc);
+      continue;
+    }
 
-  llvm::Expected<lldb::addr_t> task_addr =
-      ReadTaskAddr(*task_addr_location, *thread.GetProcess());
+    mem_read_results[idx] = ReadPointer(process, *task_addr_loc);
+    if (mem_read_results[idx])
+      m_tid_to_task_addr_location[tid] = *task_addr_loc;
+  }
 
-  // If the read from this TLS address is successful, cache the TLS address.
-  // Caching without a valid read is dangerous: earlier in the thread
-  // lifetime, the result of GetExtendedInfo can be invalid.
-  if (task_addr &&
-      real_thread.GetProcess()->GetTarget().GetSwiftCacheTaskPointerLocation())
-    m_tid_to_task_addr_location.try_emplace(real_thread.GetID(),
-                                            *task_addr_location);
-  return task_addr;
+  return mem_read_results;
 }
 
 namespace {

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

@@ -921,7 +921,17 @@ class TaskInspector {
   llvm::Expected<lldb::addr_t>
   GetTaskAddrFromThreadLocalStorage(Thread &thread);
 
+  /// Inspects thread local storage to find the address of the currently
+  /// executing task, if any.
+  llvm::SmallVector<llvm::Expected<lldb::addr_t>>
+  GetTaskAddrFromThreadLocalStorage(llvm::ArrayRef<Thread *> threads);
+
 private:
+  /// For each thread in `threads`, return the location of the its task
+  /// pointer, if it exists.
+  llvm::SmallVector<llvm::Expected<lldb::addr_t>>
+  GetTaskAddrLocations(llvm::ArrayRef<Thread *> threads);
+
   llvm::DenseMap<uint64_t, lldb::addr_t> m_tid_to_task_addr_location;
 };
 

lldb/source/Plugins/OperatingSystem/SwiftTasks/OperatingSystemSwiftTasks.cpp

@@ -151,19 +151,34 @@ OperatingSystemSwiftTasks::FindOrCreateSwiftThread(ThreadList &old_thread_list,
                                                  /*register_data_addr*/ 0);
 }
 
-static std::optional<addr_t> FindTaskAddress(TaskInspector &task_inspector,
-                                             Thread &thread) {
-  llvm::Expected<addr_t> task_addr =
-      task_inspector.GetTaskAddrFromThreadLocalStorage(thread);
-  if (!task_addr) {
-    LLDB_LOG_ERROR(GetLog(LLDBLog::OS), task_addr.takeError(),
-                   "OperatingSystemSwiftTasks: failed to find task address in "
-                   "thread local storage: {0}");
-    return {};
+/// For each thread in `threads_it`, computes the task address that is being run
+/// by the thread, if any.
+static llvm::SmallVector<std::optional<addr_t>>
+FindTaskAddresses(TaskInspector &task_inspector,
+                  ThreadCollection::ThreadIterable &threads_it) {
+  llvm::SmallVector<Thread *> threads;
+  for (const ThreadSP &thread : threads_it)
+    threads.push_back(thread.get());
+
+  llvm::SmallVector<std::optional<addr_t>> task_addrs;
+  task_addrs.reserve(threads.size());
+
+  for (llvm::Expected<addr_t> &task_addr :
+       task_inspector.GetTaskAddrFromThreadLocalStorage(threads)) {
+    if (!task_addr) {
+      LLDB_LOG_ERROR(
+          GetLog(LLDBLog::OS), task_addr.takeError(),
+          "OperatingSystemSwiftTasks: failed to find task address in "
+          "thread local storage: {0}");
+      task_addrs.push_back(std::nullopt);
+      continue;
+    }
+    if (*task_addr == 0 || *task_addr == LLDB_INVALID_ADDRESS)
+      task_addrs.push_back(std::nullopt);
+    else
+      task_addrs.push_back(*task_addr);
   }
-  if (*task_addr == 0 || *task_addr == LLDB_INVALID_ADDRESS)
-    return std::nullopt;
-  return *task_addr;
+  return task_addrs;
 }
 
 static std::optional<uint64_t> FindTaskId(addr_t task_addr, Process &process) {
@@ -189,10 +204,13 @@ bool OperatingSystemSwiftTasks::UpdateThreadList(ThreadList &old_thread_list,
   Log *log = GetLog(LLDBLog::OS);
   LLDB_LOG(log, "OperatingSystemSwiftTasks: Updating thread list");
 
-  for (const ThreadSP &real_thread : core_thread_list.Threads()) {
-    std::optional<addr_t> task_addr =
-        FindTaskAddress(m_task_inspector, *real_thread);
+  ThreadCollection::ThreadIterable locked_core_threads =
+      core_thread_list.Threads();
+  llvm::SmallVector<std::optional<addr_t>> task_addrs =
+      FindTaskAddresses(m_task_inspector, locked_core_threads);
 
+  for (const auto &[real_thread, task_addr] :
+       llvm::zip(locked_core_threads, task_addrs)) {
     // If this is not a thread running a Task, add it to the list as is.
     if (!task_addr) {
       new_thread_list.AddThread(real_thread);