[lldb][swift] Implement async funclet comparison based on mangling

lldb/source/Plugins/Language/Swift/SwiftLanguage.cpp

@@ -1834,6 +1834,23 @@ bool SwiftLanguage::IgnoreForLineBreakpoints(const SymbolContext &sc) const {
       name);
 }
 
+std::optional<bool>
+SwiftLanguage::AreEqualForFrameComparison(const SymbolContext &sc1,
+                                          const SymbolContext &sc2) const {
+  auto result = SwiftLanguageRuntime::AreFuncletsOfSameAsyncFunction(
+      sc1.GetFunctionName(Mangled::ePreferMangled),
+      sc2.GetFunctionName(Mangled::ePreferMangled));
+  switch (result) {
+  case SwiftLanguageRuntime::FuncletComparisonResult::NotBothFunclets:
+    return {};
+  case SwiftLanguageRuntime::FuncletComparisonResult::SameAsyncFunction:
+    return true;
+  case SwiftLanguageRuntime::FuncletComparisonResult::DifferentAsyncFunctions:
+    return false;
+  }
+  llvm_unreachable("unhandled enumeration in AreEquivalentFunctions");
+}
+
 //------------------------------------------------------------------
 // Static Functions
 //------------------------------------------------------------------

lldb/source/Plugins/Language/Swift/SwiftLanguage.h

@@ -75,6 +75,12 @@ class SwiftLanguage : public Language {
   ConstString
   GetDemangledFunctionNameWithoutArguments(Mangled mangled) const override;
 
+  /// Returns whether two SymbolContexts correspond to funclets of the same
+  /// async function.
+  /// If either SymbolContext is not a funclet, nullopt is returned.
+  std::optional<bool>
+  AreEqualForFrameComparison(const SymbolContext &sc1,
+                             const SymbolContext &sc2) const override;
   //------------------------------------------------------------------
   // Static Functions
   //------------------------------------------------------------------

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

@@ -132,13 +132,31 @@ class SwiftLanguageRuntime : public LanguageRuntime {
   /// since some day we may want to support more than one swift variant.
   static bool IsSwiftMangledName(llvm::StringRef name);
 
+  enum class FuncletComparisonResult {
+    NotBothFunclets,
+    DifferentAsyncFunctions,
+    SameAsyncFunction
+  };
+
+  /// Compares name1 and name2 to decide whether they are both async funclets.
+  /// If either is not an async funclet, returns NotBothFunclets.
+  /// If they are both funclets but of different async functions, returns
+  /// DifferentAsyncFunctions.
+  /// Otherwise, returns SameAsyncFunction.
+  static FuncletComparisonResult
+  AreFuncletsOfSameAsyncFunction(StringRef name1, StringRef name2);
+
   /// Return true if name is a Swift async function symbol.
   static bool IsSwiftAsyncFunctionSymbol(llvm::StringRef name);
 
   /// Return true if name is a Swift async function, await resume partial
   /// function, or suspend resume partial function symbol.
   static bool IsAnySwiftAsyncFunctionSymbol(llvm::StringRef name);
 
+  /// Return true if node is a Swift async function, await resume partial
+  /// function, or suspend resume partial function symbol.
+  static bool IsAnySwiftAsyncFunctionSymbol(NodePointer node);
+
   /// 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

@@ -105,6 +105,77 @@ static bool IsSwiftAsyncFunctionSymbol(swift::Demangle::NodePointer node) {
                       Node::Kind::AsyncAnnotation});
 }
 
+/// Returns true if closure1 and closure2 have the same number, type, and
+/// parent closures / function.
+static bool AreFuncletsOfSameAsyncClosure(NodePointer closure1,
+                                          NodePointer closure2) {
+  NodePointer closure1_number = childAtPath(closure1, Node::Kind::Number);
+  NodePointer closure2_number = childAtPath(closure2, Node::Kind::Number);
+  if (!Node::deepEquals(closure1_number, closure2_number))
+    return false;
+
+  NodePointer closure1_type = childAtPath(closure1, Node::Kind::Type);
+  NodePointer closure2_type = childAtPath(closure2, Node::Kind::Type);
+  if (!Node::deepEquals(closure1_type, closure2_type))
+    return false;
+
+  // Because the tree is inverted, a parent closure (in swift code) is a child
+  // *node* (in the demangle tree). Check that any such parents are identical.
+  NodePointer closure1_parent =
+      childAtPath(closure1, Node::Kind::ExplicitClosure);
+  NodePointer closure2_parent =
+      childAtPath(closure2, Node::Kind::ExplicitClosure);
+  if (!Node::deepEquals(closure1_parent, closure2_parent))
+    return false;
+
+  // If there are no ExplicitClosure as parents, there may still be a
+  // Function. Also check that they are identical.
+  NodePointer closure1_function = childAtPath(closure1, Node::Kind::Function);
+  NodePointer closure2_function = childAtPath(closure2, Node::Kind::Function);
+  return Node::deepEquals(closure1_function, closure2_function);
+}
+
+SwiftLanguageRuntime::FuncletComparisonResult
+SwiftLanguageRuntime::AreFuncletsOfSameAsyncFunction(StringRef name1,
+                                                     StringRef name2) {
+  using namespace swift::Demangle;
+  Context ctx;
+  NodePointer node1 = DemangleSymbolAsNode(name1, ctx);
+  NodePointer node2 = DemangleSymbolAsNode(name2, ctx);
+
+  if (!IsAnySwiftAsyncFunctionSymbol(node1) ||
+      !IsAnySwiftAsyncFunctionSymbol(node2))
+    return FuncletComparisonResult::NotBothFunclets;
+
+  // Peel off Static nodes.
+  NodePointer static_wrapper1 = childAtPath(node1, Node::Kind::Static);
+  NodePointer static_wrapper2 = childAtPath(node2, Node::Kind::Static);
+  if (static_wrapper1 || static_wrapper2) {
+    if (!static_wrapper1 | !static_wrapper2)
+      return FuncletComparisonResult::DifferentAsyncFunctions;
+    node1 = static_wrapper1;
+    node2 = static_wrapper2;
+  }
+
+  // If there are closures involved, do the closure-specific comparison.
+  NodePointer closure1 = childAtPath(node1, Node::Kind::ExplicitClosure);
+  NodePointer closure2 = childAtPath(node2, Node::Kind::ExplicitClosure);
+  if (closure1 || closure2) {
+    if (!closure1 || !closure2)
+      return FuncletComparisonResult::DifferentAsyncFunctions;
+    return AreFuncletsOfSameAsyncClosure(closure1, closure2)
+               ? FuncletComparisonResult::SameAsyncFunction
+               : FuncletComparisonResult::DifferentAsyncFunctions;
+  }
+
+  // Otherwise, find the corresponding function and compare the two.
+  NodePointer function1 = childAtPath(node1, Node::Kind::Function);
+  NodePointer function2 = childAtPath(node2, Node::Kind::Function);
+  return Node::deepEquals(function1, function2)
+             ? FuncletComparisonResult::SameAsyncFunction
+             : FuncletComparisonResult::DifferentAsyncFunctions;
+}
+
 bool SwiftLanguageRuntime::IsSwiftAsyncFunctionSymbol(StringRef name) {
   if (!IsSwiftMangledName(name))
     return false;
@@ -130,6 +201,10 @@ bool SwiftLanguageRuntime::IsAnySwiftAsyncFunctionSymbol(StringRef name) {
   using namespace swift::Demangle;
   Context ctx;
   NodePointer node = SwiftLanguageRuntime::DemangleSymbolAsNode(name, ctx);
+  return IsAnySwiftAsyncFunctionSymbol(node);
+}
+
+bool SwiftLanguageRuntime::IsAnySwiftAsyncFunctionSymbol(NodePointer node) {
   if (!node || node->getKind() != Node::Kind::Global || !node->getNumChildren())
     return false;
   auto marker = node->getFirstChild()->getKind();

lldb/unittests/Symbol/TestSwiftDemangler.cpp

@@ -11,6 +11,39 @@ static constexpr auto IsAnySwiftAsyncFunctionSymbol = [](StringRef name) {
   return SwiftLanguageRuntime::IsAnySwiftAsyncFunctionSymbol(name);
 };
 
+using FuncletComparisonResult = SwiftLanguageRuntime::FuncletComparisonResult;
+static constexpr auto AreFuncletsOfSameAsyncFunction =
+    SwiftLanguageRuntime::AreFuncletsOfSameAsyncFunction;
+
+/// Checks that all names in \c funclets belong to the same function.
+static void CheckGroupOfFuncletsFromSameFunction(ArrayRef<StringRef> funclets) {
+  for (StringRef funclet1 : funclets)
+    for (StringRef funclet2 : funclets) {
+      EXPECT_EQ(FuncletComparisonResult::SameAsyncFunction,
+                AreFuncletsOfSameAsyncFunction(funclet1, funclet2))
+          << funclet1 << " -- " << funclet2;
+      EXPECT_EQ(FuncletComparisonResult::SameAsyncFunction,
+                AreFuncletsOfSameAsyncFunction(funclet2, funclet1))
+          << funclet1 << " -- " << funclet2;
+    }
+}
+
+/// Checks that all pairs of combinations of names from \c funclets1 and \c
+/// funclets2 belong to different functions.
+static void
+CheckGroupOfFuncletsFromDifferentFunctions(ArrayRef<StringRef> funclets1,
+                                           ArrayRef<StringRef> funclets2) {
+  for (StringRef funclet1 : funclets1)
+    for (StringRef funclet2 : funclets2) {
+      EXPECT_EQ(FuncletComparisonResult::DifferentAsyncFunctions,
+                AreFuncletsOfSameAsyncFunction(funclet1, funclet2))
+          << funclet1 << " -- " << funclet2;
+      EXPECT_EQ(FuncletComparisonResult::DifferentAsyncFunctions,
+                AreFuncletsOfSameAsyncFunction(funclet2, funclet1))
+          << funclet1 << " -- " << funclet2;
+    }
+}
+
 TEST(TestSwiftDemangleAsyncNames, BasicAsync) {
   // "sayBasic" == a basic async function
   // "sayGeneric" == a generic async function
@@ -31,6 +64,10 @@ TEST(TestSwiftDemangleAsyncNames, BasicAsync) {
     EXPECT_TRUE(IsSwiftMangledName(async_name)) << async_name;
     EXPECT_TRUE(IsAnySwiftAsyncFunctionSymbol(async_name)) << async_name;
   }
+
+  CheckGroupOfFuncletsFromSameFunction(basic_funclets);
+  CheckGroupOfFuncletsFromSameFunction(generic_funclets);
+  CheckGroupOfFuncletsFromDifferentFunctions(basic_funclets, generic_funclets);
 }
 
 TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
@@ -69,19 +106,48 @@ TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
     EXPECT_TRUE(IsSwiftMangledName(async_name)) << async_name;
     EXPECT_TRUE(IsAnySwiftAsyncFunctionSymbol(async_name)) << async_name;
   }
+
+  CheckGroupOfFuncletsFromSameFunction(nested1_funclets);
+  CheckGroupOfFuncletsFromSameFunction(nested2_funclets1);
+  CheckGroupOfFuncletsFromSameFunction(nested2_funclets2);
+  CheckGroupOfFuncletsFromSameFunction(nested2_funclets_top_not_async);
+
+  CheckGroupOfFuncletsFromDifferentFunctions(nested1_funclets,
+                                             nested2_funclets1);
+  CheckGroupOfFuncletsFromDifferentFunctions(nested1_funclets,
+                                             nested2_funclets2);
+  CheckGroupOfFuncletsFromDifferentFunctions(nested1_funclets,
+                                             nested2_funclets_top_not_async);
+  CheckGroupOfFuncletsFromDifferentFunctions(nested2_funclets1,
+                                             nested2_funclets2);
+  CheckGroupOfFuncletsFromDifferentFunctions(nested2_funclets1,
+                                             nested2_funclets_top_not_async);
+  CheckGroupOfFuncletsFromDifferentFunctions(nested2_funclets2,
+                                             nested2_funclets_top_not_async);
 }
 
 TEST(TestSwiftDemangleAsyncNames, StaticAsync) {
   // static async functions
-  SmallVector<StringRef> async_names = {
-      "$s1a6StructV9sayStaticyySSYaFZ"
+  SmallVector<StringRef> static_async_funclets = {
+      "$s1a6StructV9sayStaticyySSYaFZ",
       "$s1a6StructV9sayStaticyySSYaFZTY0_",
       "$s1a6StructV9sayStaticyySSYaFZTQ1_",
       "$s1a6StructV9sayStaticyySSYaFZTY2_",
   };
 
-  for (StringRef async_name : async_names) {
+  for (StringRef async_name : static_async_funclets) {
     EXPECT_TRUE(IsSwiftMangledName(async_name)) << async_name;
     EXPECT_TRUE(IsAnySwiftAsyncFunctionSymbol(async_name)) << async_name;
   }
+
+  CheckGroupOfFuncletsFromSameFunction(static_async_funclets);
+
+  // Make sure we can compare static funclets to other kinds of funclets
+  SmallVector<StringRef> other_funclets = {
+      // Nested funclets:
+      "$s1a8sayHelloyyYaFyypYacfU_", "$s1a8sayHelloyyYaFyypYacfU_TY0_",
+      // "Normal" funclets:
+      "$s1a8sayBasicyySSYaF", "$s1a8sayBasicyySSYaFTY0_"};
+  CheckGroupOfFuncletsFromDifferentFunctions(static_async_funclets,
+                                             other_funclets);
 }