[SourceKit] Handle CustomAttrs arguments for placeholder expansion

Introduce a new ASTWalker option for walking CustomAttrs and use it
for the placeholder scanner to ensure we can expand placeholders in
attribute arguments.

Author: hamishknight

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

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());
 }
 

test/SourceKit/CodeExpand/code-expand.swift

@@ -315,3 +315,23 @@ struct ExpandDeclMacro {
   // 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) {}

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -1592,6 +1592,10 @@ class PlaceholderExpansionScanner {
       return MacroWalking::Arguments;
     }
 
+    bool shouldWalkIntoCustomAttrs() const override {
+      return true;
+    }
+
     PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
       if (isa<EditorPlaceholderExpr>(E) && E->getStartLoc() == PlaceholderLoc) {
         Found = cast<EditorPlaceholderExpr>(E);