[lldb][swift] Consider ProtocolWitness when demangling async functions

Author: felipepiovezan

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

@@ -82,10 +82,15 @@ static bool IsSwiftAsyncFunctionSymbol(swift::Demangle::NodePointer node) {
     return false;
   if (hasChild(node, Node::Kind::AsyncSuspendResumePartialFunction))
     return false;
-  // Peel off a Static node. If it exists, there will be a single instance and a
-  // top level node.
-  if (node->getFirstChild()->getKind() == Node::Kind::Static)
-    node = node->getFirstChild();
+
+  // Peel off a ProtocolWitness node. If it exists, there will be a single
+  // instance, before Static nodes.
+  if (NodePointer witness = childAtPath(node, Node::Kind::ProtocolWitness))
+    node = witness;
+
+  // Peel off a Static node. If it exists, there will be a single instance.
+  if (NodePointer static_node = childAtPath(node, Node::Kind::Static))
+    node = static_node;
 
   // Get the {Implicit, Explicit} Closure or Function node.
   // For nested closures in Swift, the demangle tree is inverted: the
@@ -203,14 +208,16 @@ SwiftLanguageRuntime::AreFuncletsOfSameAsyncFunction(
       !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;
+  // Peel off ProtocolWitnes/Static nodes, in this order.
+  for (auto wrapper : {Node::Kind::ProtocolWitness, Node::Kind::Static}) {
+    NodePointer wrapper1 = childAtPath(node1, wrapper);
+    NodePointer wrapper2 = childAtPath(node2, wrapper);
+    if (wrapper1 || wrapper2) {
+      if (!wrapper1 | !wrapper2)
+        return FuncletComparisonResult::DifferentAsyncFunctions;
+      node1 = wrapper1;
+      node2 = wrapper2;
+    }
   }
 
   // If there are closures involved, do the closure-specific comparison.
@@ -230,6 +237,11 @@ SwiftLanguageRuntime::AreFuncletsOfSameAsyncFunction(
   // Otherwise, find the corresponding function and compare the two.
   NodePointer function1 = childAtPath(node1, Node::Kind::Function);
   NodePointer function2 = childAtPath(node2, Node::Kind::Function);
+
+  // If we fail to find a function node, conservatively fail.
+  if (!function1 || !function2)
+    return FuncletComparisonResult::NotBothFunclets;
+
   return Node::deepEquals(function1, function2)
              ? FuncletComparisonResult::SameAsyncFunction
              : FuncletComparisonResult::DifferentAsyncFunctions;

lldb/unittests/Symbol/TestSwiftDemangler.cpp

@@ -22,6 +22,22 @@ static constexpr auto AreFuncletsOfSameAsyncFunction = [](StringRef name1,
 
 using FuncletComparisonResult = SwiftLanguageRuntime::FuncletComparisonResult;
 
+/// Helpful printer for the tests.
+namespace lldb_private {
+std::ostream &operator<<(std::ostream &os,
+                         const FuncletComparisonResult &result) {
+  switch (result) {
+  case FuncletComparisonResult::NotBothFunclets:
+    return os << "NotBothFunclets";
+  case FuncletComparisonResult::SameAsyncFunction:
+    return os << "SameAsyncFunction";
+  case FuncletComparisonResult::DifferentAsyncFunctions:
+    return os << "DifferentAsyncFunctions";
+  }
+  return os << "FuncletComparisonResult invalid enumeration";
+}
+} // namespace lldb_private
+
 /// Checks that all names in \c funclets belong to the same function.
 static void CheckGroupOfFuncletsFromSameFunction(ArrayRef<StringRef> funclets) {
   for (StringRef funclet1 : funclets)
@@ -140,6 +156,30 @@ TEST(TestSwiftDemangleAsyncNames, BasicAsync) {
 //    }(10)
 //  print(await another_explicit_inside_implicit_closure)
 //}
+// protocol RandomNumberGenerator {
+//     func random(in range: ClosedRange<Int>) async -> Int
+//     static func static_random() async -> Int
+// }
+// class Generator: RandomNumberGenerator {
+//     func random(in range: ClosedRange<Int>) async -> Int {
+//         try? await Task.sleep(for: .milliseconds(500))
+//         return Int.random(in: range)
+//     }
+//     static func static_random() async -> Int {
+//       print("hello")
+//       try? await Task.sleep(for: .milliseconds(500))
+//       print("hello")
+//       return 42
+//     }
+// }
+// func doMath<RNG: RandomNumberGenerator>(with rng: RNG) async {
+//     let x = await rng.random(in: 0...100)
+//     print("X is \(x)")
+//     let y = await rng.random(in: 101...200)
+//     print("Y is \(y)")
+//     print("The magic number is \(x + y)")
+// }
+
 
 TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
   // These are all async closures
@@ -193,6 +233,28 @@ TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
       "$s1a9sayHello2yyYaFSiyYaYbcfu_S2iYaXEfU0_TQ1_",
       "$s1a9sayHello2yyYaFSiyYaYbcfu_S2iYaXEfU0_TY2_",
   };
+  SmallVector<StringRef> witness_funclets = {
+      "$s1a9GeneratorCAA012RandomNumberA0A2aDP6random2inSiSNySiG_tYaFTW",
+      "$s1a9GeneratorCAA012RandomNumberA0A2aDP6random2inSiSNySiG_tYaFTWTQ0_",
+  };
+  SmallVector<StringRef> actual_funclets = {
+      "$s1a9GeneratorC6random2inSiSNySiG_tYaF",
+      "$s1a9GeneratorC6random2inSiSNySiG_tYaFTY0_",
+      "$s1a9GeneratorC6random2inSiSNySiG_tYaFTQ1_",
+      "$s1a9GeneratorC6random2inSiSNySiG_tYaFTY2_",
+      "$s1a9GeneratorC6random2inSiSNySiG_tYaFTY3_",
+  };
+  SmallVector<StringRef> static_witness_funclets = {
+      "$s1a9GeneratorCAA012RandomNumberA0A2aDP13static_randomSiyYaFZTW",
+      "$s1a9GeneratorCAA012RandomNumberA0A2aDP13static_randomSiyYaFZTWTQ0_",
+  };
+  SmallVector<StringRef> static_actual_funclets = {
+      "$s1a9GeneratorC13static_randomSiyYaFZ",
+      "$s1a9GeneratorC13static_randomSiyYaFZTY0_",
+      "$s1a9GeneratorC13static_randomSiyYaFZTQ1_",
+      "$s1a9GeneratorC13static_randomSiyYaFZTY2_",
+      "$s1a9GeneratorC13static_randomSiyYaFZTY3_",
+  };
 
   SmallVector<ArrayRef<StringRef>, 0> funclet_groups = {
       nested1_funclets,
@@ -203,6 +265,10 @@ TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
       implicit_closure_inside_explicit_closure,
       explicit_closure_inside_implicit_closure,
       another_explicit_closure_inside_implicit_closure,
+      witness_funclets,
+      actual_funclets,
+      static_witness_funclets,
+      static_actual_funclets,
   };
 
   for (ArrayRef<StringRef> funclet_group : funclet_groups)