[6.2] [IDE] Better handle macro trailing closures

include/swift/AST/ASTWalker.h

@@ -603,6 +603,13 @@ class ASTWalker {
   /// params in AbstractFunctionDecl and NominalTypeDecl.
   virtual bool shouldWalkIntoGenericParams() { return false; }
 
+  /// Whether the walker should walk into any attached CustomAttrs.
+  virtual bool shouldWalkIntoCustomAttrs() const {
+    // Default to false currently since some walkers don't handle this case
+    // well.
+    return false;
+  }
+
   /// This method configures how the walker should walk the initializers of
   /// lazy variables. These initializers are semantically different from other
   /// initializers in their context and so sometimes should be visited as part

include/swift/IDE/CompletionLookup.h

@@ -458,6 +458,8 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   void addEnumElementRef(const EnumElementDecl *EED, DeclVisibilityKind Reason,
                          DynamicLookupInfo dynamicLookupInfo,
                          bool HasTypeContext);
+  void addMacroCallArguments(const MacroDecl *MD, DeclVisibilityKind Reason,
+                             bool forTrivialTrailingClosure = false);
   void addMacroExpansion(const MacroDecl *MD, DeclVisibilityKind Reason);
 
   void addKeyword(

lib/AST/ASTWalker.cpp

@@ -120,6 +120,8 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   [[nodiscard]]
   bool visit(Decl *D) {
     SetParentRAII SetParent(Walker, D);
+    if (visitDeclCommon(D))
+      return true;
     return inherited::visit(D);
   }
 
@@ -138,6 +140,40 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   //                                 Decls
   //===--------------------------------------------------------------------===//
 
+  [[nodiscard]]
+  bool visitCustomAttr(CustomAttr *CA) {
+    auto *newTypeExpr = doIt(CA->getTypeExpr());
+    if (!newTypeExpr)
+      return true;
+
+    ASSERT(newTypeExpr == CA->getTypeExpr() &&
+           "Cannot change CustomAttr TypeExpr");
+
+    if (auto *args = CA->getArgs()) {
+      auto *newArgs = doIt(args);
+      if (!newArgs)
+        return true;
+
+      CA->setArgs(newArgs);
+    }
+    return false;
+  }
+
+  [[nodiscard]]
+  bool visitDeclCommon(Decl *D) {
+    if (Walker.shouldWalkIntoCustomAttrs()) {
+      for (auto *attr : D->getAttrs()) {
+        auto *CA = dyn_cast<CustomAttr>(attr);
+        if (!CA)
+          continue;
+
+        if (visitCustomAttr(CA))
+          return true;
+      }
+    }
+    return false;
+  }
+
   [[nodiscard]]
   bool visitGenericParamListIfNeeded(GenericContext *GC) {
     // Must check this first in case extensions have not been bound yet
@@ -2218,6 +2254,15 @@ bool Traversal::visitErrorTypeRepr(ErrorTypeRepr *T) {
 }
 
 bool Traversal::visitAttributedTypeRepr(AttributedTypeRepr *T) {
+  if (Walker.shouldWalkIntoCustomAttrs()) {
+    for (auto attr : T->getAttrs()) {
+      auto *CA = attr.dyn_cast<CustomAttr *>();
+      if (!CA)
+        continue;
+      if (visitCustomAttr(CA))
+        return true;
+    }
+  }
   return doIt(T->getTypeRepr());
 }
 

lib/IDE/CompletionLookup.cpp

@@ -104,12 +104,18 @@ static Type defaultTypeLiteralKind(CodeCompletionLiteralKind kind,
   llvm_unreachable("Unhandled CodeCompletionLiteralKind in switch.");
 }
 
-/// Whether funcType has a single argument (not including defaulted arguments)
-/// that is of type () -> ().
-static bool hasTrivialTrailingClosure(const FuncDecl *FD,
-                                      AnyFunctionType *funcType) {
-  ParameterListInfo paramInfo(funcType->getParams(), FD,
-                              /*skipCurriedSelf*/ FD->hasCurriedSelf());
+/// Whether the provided type has a single argument (not including defaulted
+/// arguments) that is of type () -> ().
+static bool hasTrivialTrailingClosure(const ValueDecl *VD, Type type) {
+  if (!VD->hasParameterList())
+    return false;
+
+  auto *funcType = type->getAs<AnyFunctionType>();
+  if (!funcType)
+    return false;
+
+  ParameterListInfo paramInfo(funcType->getParams(), VD,
+                              /*skipCurriedSelf*/ VD->hasCurriedSelf());
 
   if (paramInfo.size() - paramInfo.numNonDefaultedParameters() == 1) {
     auto param = funcType->getParams().back();
@@ -1946,34 +1952,22 @@ static StringRef getTypeAnnotationString(const MacroDecl *MD,
   return {stash.data(), stash.size()};
 }
 
-void CompletionLookup::addMacroExpansion(const MacroDecl *MD,
-                                         DeclVisibilityKind Reason) {
-  if (!MD->hasName() || !MD->isAccessibleFrom(CurrDeclContext) ||
-      MD->shouldHideFromEditor())
-    return;
-
-  OptionSet<CustomAttributeKind> expectedKinds =
-      expectedTypeContext.getExpectedCustomAttributeKinds();
-  if (expectedKinds) {
-    CodeCompletionMacroRoles expectedRoles =
-        getCompletionMacroRoles(expectedKinds);
-    CodeCompletionMacroRoles roles = getCompletionMacroRoles(MD);
-    if (!(roles & expectedRoles))
-      return;
-  }
-
+void CompletionLookup::addMacroCallArguments(const MacroDecl *MD,
+                                             DeclVisibilityKind Reason,
+                                             bool forTrivialTrailingClosure) {
   CodeCompletionResultBuilder Builder =
       makeResultBuilder(CodeCompletionResultKind::Declaration,
                         getSemanticContext(MD, Reason, DynamicLookupInfo()));
   Builder.setAssociatedDecl(MD);
 
   addValueBaseName(Builder, MD->getBaseIdentifier());
 
-  Type macroType = MD->getInterfaceType();
-  if (MD->parameterList && MD->parameterList->size() > 0) {
+  if (forTrivialTrailingClosure) {
+    Builder.addBraceStmtWithCursor(" { code }");
+  } else if (MD->parameterList && MD->parameterList->size() > 0) {
+    auto *macroTy = MD->getInterfaceType()->castTo<AnyFunctionType>();
     Builder.addLeftParen();
-    addCallArgumentPatterns(Builder, macroType->castTo<AnyFunctionType>(),
-                            MD->parameterList,
+    addCallArgumentPatterns(Builder, macroTy, MD->parameterList,
                             MD->getGenericSignature());
     Builder.addRightParen();
   }
@@ -1988,6 +1982,28 @@ void CompletionLookup::addMacroExpansion(const MacroDecl *MD,
   }
 }
 
+void CompletionLookup::addMacroExpansion(const MacroDecl *MD,
+                                         DeclVisibilityKind Reason) {
+  if (!MD->hasName() || !MD->isAccessibleFrom(CurrDeclContext) ||
+      MD->shouldHideFromEditor())
+    return;
+
+  OptionSet<CustomAttributeKind> expectedKinds =
+      expectedTypeContext.getExpectedCustomAttributeKinds();
+  if (expectedKinds) {
+    CodeCompletionMacroRoles expectedRoles =
+        getCompletionMacroRoles(expectedKinds);
+    CodeCompletionMacroRoles roles = getCompletionMacroRoles(MD);
+    if (!(roles & expectedRoles))
+      return;
+  }
+
+  if (hasTrivialTrailingClosure(MD, MD->getInterfaceType()))
+    addMacroCallArguments(MD, Reason, /*forTrivialTrailingClosure*/ true);
+
+  addMacroCallArguments(MD, Reason);
+}
+
 void CompletionLookup::addKeyword(StringRef Name, Type TypeAnnotation,
                                   SemanticContextKind SK,
                                   CodeCompletionKeywordKind KeyKind,

test/IDE/complete_macros.swift

@@ -21,6 +21,15 @@ public macro freestandingExprStringMacro() -> String
 @freestanding(expression)
 public macro freestandingExprTMacro<T>(_ value: T) -> T
 
+@freestanding(expression)
+public macro freestandingExprVoidClosureMacro(fn: () -> Void)
+
+@freestanding(expression)
+public macro freestandingExprIntClosureMacro(fn: () -> Int)
+
+@freestanding(declaration)
+public macro freestandingDeclVoidClosureMacro(fn: () -> Void)
+
 @freestanding(declaration)
 public macro freestandingDeclMacro()
 
@@ -158,8 +167,21 @@ func nestedFreestanding() {
 // ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingExprTMacro({#(value): T#})[#T#]; name=freestandingExprTMacro(:)
 // ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: EverythingMacro[#Expression Macro, Declaration Macro, Accessor Macro, Member Attribute Macro, Member Macro, Peer Macro, Extension Macro#]; name=EverythingMacro
 //
+// We offer a trailing closure completion for the Void case, but not the
+// Int case currently. Placeholder expansion will turn the latter into the
+// former though.
+//
+// ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingExprVoidClosureMacro {|}[#Void#]; name=freestandingExprVoidClosureMacro
+// ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingExprVoidClosureMacro({#fn: () -> Void##() -> Void#})[#Void#]; name=freestandingExprVoidClosureMacro(fn:)
+//
+// ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingExprIntClosureMacro({#fn: () -> Int##() -> Int#})[#Void#]; name=freestandingExprIntClosureMacro(fn:)
+//
+// ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingDeclVoidClosureMacro {|}[#Declaration Macro#]; name=freestandingDeclVoidClosureMacro
+// ALL_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingDeclVoidClosureMacro({#fn: () -> Void##() -> Void#})[#Declaration Macro#]; name=freestandingDeclVoidClosureMacro(fn:)
+//
 // ALL_FREESTANDING_NOT-NOT: Attached
 // ALL_FREESTANDING_NOT-NOT: BodyMacro
+// ALL_FREESTANDING_NOT-NOT: freestandingExprIntClosureMacro {|}
 
 func exprFreestanding(arg: Int) {
   _ = arg + ##^EXPR_FREESTANDING?check=EXPR_FREESTANDING;check=EXPR_FREESTANDING_NOT^#
@@ -181,6 +203,9 @@ struct NestedFreestanding {
 // ITEM_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingDeclMacro[#Declaration Macro#]; name=freestandingDeclMacro
 // ITEM_FREESTANDING-DAG: Decl[Macro]/{{.*}}: EverythingMacro[#Expression Macro, Declaration Macro, Accessor Macro, Member Attribute Macro, Member Macro, Peer Macro, Extension Macro#]; name=EverythingMacro
 //
+// ITEM_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingDeclVoidClosureMacro {|}[#Declaration Macro#]; name=freestandingDeclVoidClosureMacro
+// ITEM_FREESTANDING-DAG: Decl[Macro]/{{.*}}: freestandingDeclVoidClosureMacro({#fn: () -> Void##() -> Void#})[#Declaration Macro#]; name=freestandingDeclVoidClosureMacro(fn:)
+//
 // ITEM_FREESTANDING_NOT-NOT: Attached
 // ITEM_FREESTANDING_NOT-NOT: freestandingExpr
 // ITEM_FREESTANDING_NOT-NOT: freestandingCodeItemMacro

test/SourceKit/CodeExpand/code-expand.swift

@@ -25,6 +25,10 @@ await foo(x: <#T##() -> Void#>)
 // CHECK-NEXT: <#code#>
 // CHECK-NEXT: }
 
+foo(bar(<#T##() -> Void#>))
+// CHECK:      foo(bar({
+// CHECK-NEXT: <#code#>
+// CHECK-NEXT: }))
 
 anArr.indexOfObjectPassingTest(<#T##predicate: ((AnyObject!, Int, UnsafePointer<ObjCBool>) -> Bool)?##((AnyObject!, Int, UnsafePointer<ObjCBool>) -> Bool)?#>)
 // CHECK:      anArr.indexOfObjectPassingTest { <#AnyObject!#>, <#Int#>, <#UnsafePointer<ObjCBool>#> in
@@ -273,3 +277,61 @@ expandClosureWithInternalParameterNames {
 // CHECK: withtrail { a, b in
 // CHECK-NEXT: <#code#>
 }
+
+// CHECK-LABEL: func expandMacro()
+func expandMacro() {
+  #foo(<#T##() -> Int#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+
+  #foo(bar: <#T##() -> ()#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+
+  #foo(bar: <#T##() -> Int#>, baz: <#T##() -> ()#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: } baz: {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: struct ExpandDeclMacro
+struct ExpandDeclMacro {
+  #foo(<#T##() -> ()#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+
+  #foo(bar(<#T##() -> ()#>))
+  // CHECK:      #foo(bar({
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }))
+
+  #foo(#bar(<#T##() -> ()#>))
+  // CHECK:      #foo(#bar({
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }))
+}
+
+@Foo(<#Int#>)
+func testDeclAttr1() {}
+// CHECK: @Foo(<#Int#>)
+// CHECK-NEXT: func testDeclAttr1() {}
+
+@Foo(<#T##() -> ()#>)
+func testDeclAttr2() {}
+// CHECK:      @Foo({
+// CHECK-NEXT:   <#code#>
+// CHECK-NEXT: })
+// CHECK-NEXT: func testDeclAttr2() {}
+
+func testTypeAttr1(x: @Foo(<#Int#>) String) {}
+// CHECK: func testTypeAttr1(x: @Foo(<#Int#>) String) {}
+
+func testTypeAttr2(x: @Foo(<#T##() -> ()#>) Int) {}
+// CHECK: func testTypeAttr2(x: @Foo({
+// CHECK-NEXT:   <#code#>
+// CHECK-NEXT: }) Int) {}