[lldb][swift] Properly detect async functions with ImplicitClosure demangle nodes

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

@@ -82,26 +82,39 @@ 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();
 
-  // Get the ExplicitClosure or Function node.
+  // 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
-  // inner-most closure is the top-most ExplicitClosure node.
+  // inner-most closure is the top-most closure node.
   NodePointer func_node = [&] {
     if (NodePointer func = childAtPath(node, Node::Kind::Function))
       return func;
+    if (NodePointer implicit = childAtPath(node, Node::Kind::ImplicitClosure))
+      return implicit;
     return childAtPath(node, Node::Kind::ExplicitClosure);
   }();
 
-  return childAtPath(func_node, {Node::Kind::Type, Node::Kind::FunctionType,
-                                 Node::Kind::AsyncAnnotation}) ||
-         childAtPath(func_node,
-                     {Node::Kind::Type, Node::Kind::DependentGenericType,
-                      Node::Kind::Type, Node::Kind::FunctionType,
-                      Node::Kind::AsyncAnnotation});
+  using Kind = Node::Kind;
+  static const llvm::SmallVector<llvm::SmallVector<Kind>>
+      async_annotation_paths = {
+          {Kind::Type, Kind::FunctionType, Kind::AsyncAnnotation},
+          {Kind::Type, Kind::NoEscapeFunctionType, Kind::AsyncAnnotation},
+          {Kind::Type, Kind::DependentGenericType, Kind::Type,
+           Kind::FunctionType, Kind::AsyncAnnotation},
+      };
+  return llvm::any_of(async_annotation_paths,
+                      [func_node](llvm::ArrayRef<Kind> path) {
+                        return childAtPath(func_node, path);
+                      });
 }
 
 /// Returns true if closure1 and closure2 have the same number, type, and
@@ -122,11 +135,16 @@ AreFuncletsOfSameAsyncClosure(swift::Demangle::NodePointer closure1,
 
   // 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 =
+  NodePointer explicit_closure1_parent =
       childAtPath(closure1, Node::Kind::ExplicitClosure);
-  NodePointer closure2_parent =
+  NodePointer explicit_closure2_parent =
       childAtPath(closure2, Node::Kind::ExplicitClosure);
-  if (!Node::deepEquals(closure1_parent, closure2_parent))
+  NodePointer implicit_closure1_parent =
+      childAtPath(closure1, Node::Kind::ImplicitClosure);
+  NodePointer implicit_closure2_parent =
+      childAtPath(closure2, Node::Kind::ImplicitClosure);
+  if (!Node::deepEquals(explicit_closure1_parent, explicit_closure2_parent) ||
+      !Node::deepEquals(implicit_closure1_parent, implicit_closure2_parent))
     return false;
 
   // If there are no ExplicitClosure as parents, there may still be a
@@ -190,30 +208,40 @@ 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.
-  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;
+  for (auto closure_kind :
+       {Node::Kind::ImplicitClosure, Node::Kind::ExplicitClosure}) {
+    NodePointer closure1 = childAtPath(node1, closure_kind);
+    NodePointer closure2 = childAtPath(node2, closure_kind);
+    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);
+
+  // 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/test/API/lang/swift/async_breakpoints/TestSwiftAsyncBreakpoints.py

@@ -24,9 +24,7 @@ def test(self):
         self.assertEquals(breakpoint1.GetNumLocations(), 1)
         self.assertEquals(breakpoint2.GetNumLocations(), 1)
         self.assertEquals(breakpoint3.GetNumLocations(), 1)
-        # FIXME: there should be two breakpoints here, but the "entry" funclet of the
-        # implicit closure is mangled slightly differently. rdar://147035260
-        self.assertEquals(breakpoint4.GetNumLocations(), 3)
+        self.assertEquals(breakpoint4.GetNumLocations(), 2)
         self.assertEquals(breakpoint5.GetNumLocations(), 1)
 
         location11 = breakpoint1.GetLocationAtIndex(0)

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)
@@ -86,6 +102,85 @@ TEST(TestSwiftDemangleAsyncNames, BasicAsync) {
   CheckFuncletNumbersAreARange(generic_funclets);
 }
 
+// The funclets below are created from this program:
+// swiftc -g -Onone test.swift -o - -emit-ir -module-name a \
+// | grep "define.*sayHello"
+// func work() async {}
+// func async_int() async -> Int { return 42; }
+// func sayHello() async {
+//   let closure: (Any) async -> () = { _ in
+//     print("hello")
+//     await work()
+//     print("hello")
+//
+//     let inner_closure: (Any) async -> () = { _ in
+//       print("hello")
+//       await work()
+//       print("hello")
+//     }
+//     await inner_closure(10)
+//     print("hello")
+//
+//     let inner_closure2: (Any) async -> () = { _ in
+//       print("hello")
+//       await work()
+//       print("hello")
+//     }
+//
+//     await inner_closure2(10)
+//     print("hello")
+//     async let x = await async_int();
+//     print(await x);
+//   }
+//   async let x = await async_int();
+//   print(await x);
+//   await closure(10)
+//   async let explicit_inside_implicit_closure =
+//     { _ in
+//       print("hello")
+//       await work()
+//       print("hello")
+//       return 42
+//     }(10)
+//   print(await explicit_inside_implicit_closure)
+// }
+// func sayHello2() async {
+//  {_ in
+//  }(10)
+//  async let another_explicit_inside_implicit_closure =
+//    { _ in
+//      print("hello")
+//      await work()
+//      print("hello")
+//      return 42
+//    }(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
   SmallVector<StringRef> nested1_funclets = {
@@ -116,34 +211,83 @@ TEST(TestSwiftDemangleAsyncNames, ClosureAsync) {
       "$s1a18myNonAsyncFunctionyyFyyYacfU_SiypYacfU_SSypYacfU0_TQ1_",
       "$s1a18myNonAsyncFunctionyyFyyYacfU_SiypYacfU_SSypYacfU0_TY2_"};
 
-  for (StringRef async_name : llvm::concat<StringRef>(
-           nested1_funclets, nested2_funclets1, nested2_funclets2,
-           nested2_funclets_top_not_async)) {
-    EXPECT_TRUE(IsSwiftMangledName(async_name)) << async_name;
-    EXPECT_TRUE(IsAnySwiftAsyncFunctionSymbol(async_name)) << async_name;
-  }
+  SmallVector<StringRef> implicit_closure_inside_function = {
+      "$s1a8sayHelloyyYaFSiyYaYbcfu_",
+      "$s1a8sayHelloyyYaFSiyYaYbcfu_TQ0_",
+      "$s1a8sayHelloyyYaFSiyYaYbcfu_TY1_",
+  };
+  SmallVector<StringRef> implicit_closure_inside_explicit_closure = {
+      "$s1a8sayHelloyyYaFyypYacfU_SiyYaYbcfu_",
+      "$s1a8sayHelloyyYaFyypYacfU_SiyYaYbcfu_TQ0_",
+      "$s1a8sayHelloyyYaFyypYacfU_SiyYaYbcfu_TY1_",
+  };
+  SmallVector<StringRef> explicit_closure_inside_implicit_closure = {
+      "$s1a8sayHelloyyYaFSiyYaYbcfu0_S2iYaXEfU0_",
+      "$s1a8sayHelloyyYaFSiyYaYbcfu0_S2iYaXEfU0_TY0_",
+      "$s1a8sayHelloyyYaFSiyYaYbcfu0_S2iYaXEfU0_TQ1_",
+      "$s1a8sayHelloyyYaFSiyYaYbcfu0_S2iYaXEfU0_TY2_",
+  };
+  SmallVector<StringRef> another_explicit_closure_inside_implicit_closure = {
+      "$s1a9sayHello2yyYaFSiyYaYbcfu_S2iYaXEfU0_",
+      "$s1a9sayHello2yyYaFSiyYaYbcfu_S2iYaXEfU0_TY0_",
+      "$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,
+      nested2_funclets1,
+      nested2_funclets2,
+      nested2_funclets_top_not_async,
+      implicit_closure_inside_function,
+      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)
+    for (StringRef async_name : funclet_group) {
+      EXPECT_TRUE(IsSwiftMangledName(async_name)) << async_name;
+      EXPECT_TRUE(IsAnySwiftAsyncFunctionSymbol(async_name)) << async_name;
+    }
+
+  for (ArrayRef<StringRef> funclet_group : funclet_groups)
+    CheckGroupOfFuncletsFromSameFunction(funclet_group);
+
+  for (ArrayRef<StringRef> funclet_group1 : funclet_groups)
+    for (ArrayRef<StringRef> funclet_group2 : funclet_groups)
+      if (funclet_group1.data() != funclet_group2.data())
+        CheckGroupOfFuncletsFromDifferentFunctions(funclet_group1,
+                                                   funclet_group2);
 
-  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);
-  CheckFuncletNumbersAreARange(nested1_funclets);
-  CheckFuncletNumbersAreARange(nested2_funclets1);
-  CheckFuncletNumbersAreARange(nested2_funclets2);
-  CheckFuncletNumbersAreARange(nested2_funclets_top_not_async);
+  for (ArrayRef<StringRef> funclet_group : funclet_groups)
+    CheckFuncletNumbersAreARange(funclet_group);
 }
 
 TEST(TestSwiftDemangleAsyncNames, StaticAsync) {