[lldb][swift] Add function to extract funclet numbers

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

@@ -181,6 +181,14 @@ class SwiftLanguageRuntime : public LanguageRuntime {
   /// function, or suspend resume partial function symbol.
   static bool IsAnySwiftAsyncFunctionSymbol(swift::Demangle::NodePointer node);
 
+  /// If node is a Swift async funclet, return its funclet number.
+  static std::optional<uint64_t>
+  GetFuncletNumber(swift::Demangle::NodePointer node);
+
+  /// If name is a Swift async funclet, return its funclet number.
+  static std::optional<uint64_t>
+  GetFuncletNumber(llvm::StringRef name);
+
   /// Return the async context address using the target's specific register.
   static lldb::addr_t GetAsyncContext(RegisterContext *regctx);
 

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

@@ -136,6 +136,43 @@ AreFuncletsOfSameAsyncClosure(swift::Demangle::NodePointer closure1,
   return Node::deepEquals(closure1_function, closure2_function);
 }
 
+std::optional<uint64_t>
+SwiftLanguageRuntime::GetFuncletNumber(swift::Demangle::NodePointer node) {
+  if (!node)
+    return {};
+  NodePointer node_id =
+      childAtPath(node, {Node::Kind::AsyncSuspendResumePartialFunction,
+                         Node::Kind::Number});
+  if (!node_id)
+    node_id = childAtPath(node, {Node::Kind::AsyncAwaitResumePartialFunction,
+                                 Node::Kind::Number});
+  if (node_id) {
+    if (node_id->hasIndex())
+      return node_id->getIndex();
+
+    // This should never happen, log it if it does:
+    std::string node_dump = getNodeTreeAsString(node);
+    LLDB_LOGF(GetLog(LLDBLog::Demangle),
+              "%s::Found a number node without a number:\n%s", __FUNCTION__,
+              node_dump.c_str());
+    return {};
+  }
+
+  // If this is an async function and there was no number, then this is an entry
+  // funclet. Assign it number 0.
+  if (SwiftLanguageRuntime::IsAnySwiftAsyncFunctionSymbol(node))
+    return 0;
+  return {};
+}
+
+std::optional<uint64_t>
+SwiftLanguageRuntime::GetFuncletNumber(llvm::StringRef name) {
+  using namespace swift::Demangle;
+  Context ctx;
+  NodePointer node = SwiftLanguageRuntime::DemangleSymbolAsNode(name, ctx);
+  return GetFuncletNumber(node);
+}
+
 SwiftLanguageRuntime::FuncletComparisonResult
 SwiftLanguageRuntime::AreFuncletsOfSameAsyncFunction(llvm::StringRef name1,
                                                      llvm::StringRef name2) {

lldb/unittests/Symbol/TestSwiftDemangler.cpp

@@ -1,12 +1,16 @@
 #include "gtest/gtest.h"
 
 #include "Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.h"
+#include "llvm/ADT/Sequence.h"
 
 using namespace lldb;
 using namespace lldb_private;
 using namespace llvm;
 static constexpr auto IsSwiftMangledName =
     SwiftLanguageRuntime::IsSwiftMangledName;
+static constexpr auto GetFuncletNumber = [](StringRef name) {
+  return SwiftLanguageRuntime::GetFuncletNumber(name);
+};
 static constexpr auto IsAnySwiftAsyncFunctionSymbol = [](StringRef name) {
   return SwiftLanguageRuntime::IsAnySwiftAsyncFunctionSymbol(name);
 };
@@ -47,6 +51,13 @@ CheckGroupOfFuncletsFromDifferentFunctions(ArrayRef<StringRef> funclets1,
     }
 }
 
+/// Check that funclets contain a sequence of funclet names whose "async
+/// numbers" go from 0 to size(funclets).
+static void CheckFuncletNumbersAreARange(ArrayRef<StringRef> funclets) {
+  for (auto idx : llvm::seq<int>(0, funclets.size()))
+    EXPECT_EQ(idx, GetFuncletNumber(funclets[idx]));
+}
+
 TEST(TestSwiftDemangleAsyncNames, BasicAsync) {
   // "sayBasic" == a basic async function
   // "sayGeneric" == a generic async function
@@ -71,6 +82,8 @@ TEST(TestSwiftDemangleAsyncNames, BasicAsync) {
   CheckGroupOfFuncletsFromSameFunction(basic_funclets);
   CheckGroupOfFuncletsFromSameFunction(generic_funclets);
   CheckGroupOfFuncletsFromDifferentFunctions(basic_funclets, generic_funclets);
+  CheckFuncletNumbersAreARange(basic_funclets);
+  CheckFuncletNumbersAreARange(generic_funclets);
 }
 
 TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
@@ -127,6 +140,10 @@ TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
                                              nested2_funclets_top_not_async);
   CheckGroupOfFuncletsFromDifferentFunctions(nested2_funclets2,
                                              nested2_funclets_top_not_async);
+  CheckFuncletNumbersAreARange(nested1_funclets);
+  CheckFuncletNumbersAreARange(nested2_funclets1);
+  CheckFuncletNumbersAreARange(nested2_funclets2);
+  CheckFuncletNumbersAreARange(nested2_funclets_top_not_async);
 }
 
 TEST(TestSwiftDemangleAsyncNames, StaticAsync) {
@@ -153,4 +170,20 @@ TEST(TestSwiftDemangleAsyncNames, StaticAsync) {
       "$s1a8sayBasicyySSYaF", "$s1a8sayBasicyySSYaFTY0_"};
   CheckGroupOfFuncletsFromDifferentFunctions(static_async_funclets,
                                              other_funclets);
+  CheckFuncletNumbersAreARange(static_async_funclets);
+}
+
+TEST(TestSwiftDemangleAsyncNames, NonAsync) {
+  // func factorial(_ n:Int) -> Int {
+  {
+    StringRef func = "$s4test9factorialyS2iF";
+    EXPECT_EQ(GetFuncletNumber(func), std::nullopt);
+  }
+
+  // func factorial(_ n:Int) async -> Int {
+  //   func inner_factorial(_ n:Int) -> Int {
+  {
+    StringRef func = "$s4test9factorialyS2iYaF06inner_B0L_yS2iF";
+    EXPECT_EQ(GetFuncletNumber(func), std::nullopt);
+  }
 }