[IDE] Visit auxiliary declarations if walking expansions

After visiting declarations, also walk into their auxiliary decls if
expansions are being walked.

Resolves rdar://109548265.

Author: bnbarham

Diff for: lib/AST/ASTWalker.cpp

@@ -455,20 +455,20 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       else
         return true;
     }
+
     bool alreadyFailed = false;
     if (shouldWalkExpansion) {
       MED->forEachExpandedNode([&](ASTNode expandedNode) {
         if (alreadyFailed) return;
+
         if (auto *expr = expandedNode.dyn_cast<Expr *>()) {
-          if (!doIt(expr))
-            alreadyFailed = true;
+          alreadyFailed = doIt(expr) == nullptr;
         } else if (auto *stmt = expandedNode.dyn_cast<Stmt *>()) {
-          if (!doIt(stmt))
-            alreadyFailed = true;
+          alreadyFailed = doIt(stmt) == nullptr;
         } else {
           auto decl = expandedNode.get<Decl *>();
           if (!isa<VarDecl>(decl))
-            alreadyFailed = inherited::visit(decl);
+            alreadyFailed = doIt(decl);
         }
       });
     }

Diff for: lib/IDE/SourceEntityWalker.cpp

@@ -236,6 +236,22 @@ ASTWalker::PreWalkAction SemaAnnotator::walkToDeclPreProper(Decl *D) {
 ASTWalker::PostWalkAction SemaAnnotator::walkToDeclPost(Decl *D) {
   auto Action = walkToDeclPostProper(D);
   SEWalker.endBalancedASTOrderDeclVisit(D);
+
+  if (Action.Action == PostWalkAction::Stop)
+    return Action;
+
+  // Walk into peer and conformance expansions if walking expansions
+  if (shouldWalkMacroArgumentsAndExpansion().second) {
+    D->visitAuxiliaryDecls([&](Decl *auxDecl) {
+      if (Action.Action == PostWalkAction::Stop)
+        return;
+
+      if (auxDecl->walk(*this)) {
+        Action = Action::Stop();
+      }
+    }, /*visitFreestandingExpanded=*/false);
+  }
+
   return Action;
 }
 

Diff for: test/Index/index_macros.swift

@@ -1,27 +1,284 @@
+// REQUIRES: swift_swift_parser
+
 // RUN: %empty-directory(%t)
-// RUN: %target-swift-ide-test -print-indexed-symbols -source-filename %s | %FileCheck %s
-// REQUIRES: OS=macosx
+// RUN: split-file --leading-lines %s %t
+
+// Check that we index code expanded from macros, especially nested references
+// (ie. calls within an added function).
+
+// Create the plugin with various macros for testing
+// RUN: %host-build-swift -swift-version 5 -emit-library -o %t/%target-library-name(IndexMacros) -module-name=IndexMacros %t/IndexMacros.swift -g -no-toolchain-stdlib-rpath
+
+// Check indexed symbols
+// RUN: %target-swift-ide-test -print-indexed-symbols -source-filename %t/IndexTest.swift -load-plugin-library %t/%target-library-name(IndexMacros) -parse-as-library 2>&1 | tee %t/test.idx | %FileCheck %s
+
+//--- IndexTest.swift
+@freestanding(expression)
+macro freestandingExpr() = #externalMacro(module: "IndexMacros", type: "FreestandingExprMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | freestandingExpr() |  [[EXPR_USR:.*]] | Def
+
+@freestanding(declaration, names: named(TestFree))
+macro freestandingDecl() = #externalMacro(module: "IndexMacros", type: "FreestandingDeclMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | freestandingDecl() |  [[DECL_USR:.*]] | Def
+
+@attached(accessor)
+macro Accessor() = #externalMacro(module: "IndexMacros", type: "SomeAccessorMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | Accessor() |  [[ACCESSOR_USR:.*]] | Def
+
+@attached(conformance)
+macro Conformance() = #externalMacro(module: "IndexMacros", type: "SomeConformanceMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | Conformance() |  [[CONFORMANCE_USR:.*]] | Def
+
+@attached(member, names: named(memberFunc))
+macro Member() = #externalMacro(module: "IndexMacros", type: "SomeMemberMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | Member() |  [[MEMBER_USR:.*]] | Def
+
+@attached(memberAttribute)
+macro MemberAttribute() = #externalMacro(module: "IndexMacros", type: "SomeMemberAttributeMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | MemberAttribute() |  [[MEMBER_ATTRIBUTE_USR:.*]] | Def
+
+@attached(peer, names: named(TestPeer))
+macro Peer() = #externalMacro(module: "IndexMacros", type: "SomePeerMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | Peer() |  [[PEER_USR:.*]] | Def
+
+@attached(peer, names: named(peerMember))
+macro PeerMember() = #externalMacro(module: "IndexMacros", type: "SomePeerMemberMacro")
+// CHECK: [[@LINE-1]]:7 | macro/Swift | PeerMember() |  [[PEER_MEMBER_USR:.*]] | Def
+
+protocol TestProto {}
+// CHECK: [[@LINE-1]]:10 | protocol/Swift | TestProto | [[PROTO_USR:.*]] | Def
+
+func accessorLog() {}
+// CHECK: [[@LINE-1]]:6 | function/Swift | accessorLog() | [[ACC_LOG_USR:.*]] | Def
+func exprLog() {}
+// CHECK: [[@LINE-1]]:6 | function/Swift | exprLog() | [[EXPR_LOG_USR:.*]] | Def
+func freeLog() {}
+// CHECK: [[@LINE-1]]:6 | function/Swift | freeLog() | [[FREE_LOG_USR:.*]] | Def
+func memberLog() {}
+// CHECK: [[@LINE-1]]:6 | function/Swift | memberLog() | [[MEMBER_LOG_USR:.*]] | Def
+func peerLog() {}
+// CHECK: [[@LINE-1]]:6 | function/Swift | peerLog() | [[PEER_LOG_USR:.*]] | Def
+
+// CHECK: [[@LINE+2]]:8 | struct/Swift | AddOne | [[ADD_ONE_USR:.*]] | Def
+@propertyWrapper
+struct AddOne {
+  var value: Int = 1
+  var wrappedValue: Int {
+    get { value }
+    set { value = newValue + 1 }
+  }
+  init(wrappedValue: Int) {
+    self.wrappedValue = wrappedValue
+  }
+}
+
+// Creates a `TestFree` struct with `freeFunc` calling `freeLog`
+#freestandingDecl
+// CHECK: [[@LINE-1]]:2 | macro/Swift | freestandingDecl() | [[DECL_USR]] | Ref
+// CHECK: [[@LINE-2]]:1 | struct/Swift | TestFree | [[FREE_STRUCT_USR:.*]] | Def,Impl
+// CHECK: [[@LINE-3]]:1 | instance-method/Swift | freeFunc() | [[FREE_FUNC_USR:.*]] | Def,Impl,RelChild
+// CHECK-NEXT: RelChild | struct/Swift | TestFree | [[FREE_STRUCT_USR]]
+// CHECK: [[@LINE-5]]:1 | function/Swift | freeLog() | [[FREE_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+// CHECK-NEXT: RelCall,RelCont | instance-method/Swift | freeFunc() | [[FREE_FUNC_USR]]
+
+// CHECK: [[@LINE+4]]:40 | macro/Swift | Peer() | [[PEER_USR]] | Ref
+// CHECK: [[@LINE+3]]:23 | macro/Swift | MemberAttribute() | [[MEMBER_ATTRIBUTE_USR]] | Ref
+// CHECK: [[@LINE+2]]:15 | macro/Swift | Member() | [[MEMBER_USR]] | Ref
+// CHECK: [[@LINE+1]]:2 | macro/Swift | Conformance() | [[CONFORMANCE_USR]] | Ref
+@Conformance @Member @MemberAttribute @Peer
+struct TestAttached {
+  var attachedMember: Int
+
+  @Accessor
+  var attachedMemberAccessors: Int
+}
+// `MemberAttribute` adds `@AddOne` to attachedMember
+// CHECK: [[@LINE-8]]:22 | struct/Swift | AddOne | [[ADD_ONE_USR]] | Ref,Impl,RelCont
+// CHECK-NEXT: RelCont | instance-property/Swift | attachedMember
+
+// `Accessor` adds getters/setters to `attachedMemberAccessors` that both call `accessorLog`
+// CHECK: [[@LINE-8]]:3 | function/Swift | accessorLog() | [[ACC_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+// CHECK-NEXT: RelCall,RelCont | instance-method/acc-get/Swift | getter:attachedMemberAccessors
 
+// `Member` adds a new member `memberFunc` that calls `memberLog`
+// CHECK: [[@LINE-16]]:14 | instance-method/Swift | memberFunc() | [[MEMBER_FUNC_USR:.*]] | Def,Impl,RelChild
+// CHECK: [[@LINE-17]]:14 | function/Swift | memberLog() | [[MEMBER_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+// CHECK-NEXT: RelCall,RelCont | instance-method/Swift | memberFunc() | [[MEMBER_FUNC_USR]]
 
-@freestanding(expression) macro myLine() -> Int = #externalMacro(module: "MacroDefinition", type: "LineMacro")
-@freestanding(expression) macro myFilename<T: ExpressibleByStringLiteral>() -> T = #externalMacro(module: "MacroDefinition", type: "FileMacro")
-@freestanding(expression) macro myStringify<T>(_: T) -> (T, String) = #externalMacro(module: "MacroDefinition", type: "StringifyMacro")
+// `Peer` adds a new inner type `TestPeer` that contains `peerFunc` with a call to `peerLog`
+// CHECK: [[@LINE-21]]:39 | struct/Swift | TestPeer | [[PEER_STRUCT_USR:.*]] | Def,Impl
+// CHECK: [[@LINE-22]]:39 | instance-method/Swift | peerFunc() | [[PEER_FUNC_USR:.*]] | Def,Impl,RelChild
+// CHECK-NEXT: RelChild | struct/Swift | TestPeer | [[PEER_STRUCT_USR]]
+// CHECK: [[@LINE-24]]:39 | function/Swift | peerLog() | [[PEER_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+// CHECK-NEXT: RelCall,RelCont | instance-method/Swift | peerFunc() | [[PEER_FUNC_USR]]
 
-func test(x: Int) {
-  _ = #myLine
-  let _: String = #myFilename
-  _ = #myStringify(x + x)
+// `Conformance` adds `TestProto` as a conformance on an extension of `TestAttached`
+// CHECK: [[@LINE-28]]:1 | extension/ext-struct/Swift | TestAttached | {{.*}} | Def,Impl
+// CHECK: [[@LINE-29]]:1 | protocol/Swift | TestProto | [[PROTO_USR]] | Ref,Impl,RelBase
+// CHECK-NEXT: RelBase | extension/ext-struct/Swift | TestAttached
+
+// CHECK: [[@LINE+1]]:8 | struct/Swift | Outer | [[OUTER_USR:.*]] | Def
+struct Outer {
+  // CHECK: [[@LINE+1]]:4 | macro/Swift | PeerMember() | [[PEER_MEMBER_USR]] | Ref
+  @PeerMember
+  var anyMember: Int
+  // `PeerMember` adds a new `peerMember`
+  // CHECK: [[@LINE-3]]:3 | instance-property/Swift | peerMember | {{.*}} | Def,Impl,RelChild
+  // CHECK-NEXT: RelChild | struct/Swift | Outer | [[OUTER_USR]]
+
+  // CHECK: [[@LINE+2]]:17 | macro/Swift | Member() | [[MEMBER_USR]] | Ref
+  // CHECK: [[@LINE+1]]:4 | macro/Swift | Conformance() | [[CONFORMANCE_USR]] | Ref
+  @Conformance @Member
+  struct TestInner {}
+}
+// `Member` adds a new member `memberFunc` that calls `memberLog`
+// CHECK: [[@LINE-4]]:16 | instance-method/Swift | memberFunc() | [[INNER_FUNC_USR:.*]] | Def,Impl
+// CHECK-NEXT: RelChild | struct/Swift | TestInner
+// CHECK: [[@LINE-6]]:16 | function/Swift | memberLog() | [[MEMBER_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+// CHECK-NEXT: RelCall,RelCont | instance-method/Swift | memberFunc() | [[INNER_FUNC_USR]]
+
+// `Conformance` adds `TestProto` as a conformance on an extension of `TestInner`
+// CHECK: [[@LINE-10]]:3 | extension/ext-struct/Swift | TestInner | {{.*}} | Def,Impl
+// CHECK: [[@LINE-11]]:3 | protocol/Swift | TestProto | [[PROTO_USR]] | Ref,Impl,RelBase
+// CHECK-NEXT: RelBase | extension/ext-struct/Swift | TestInner
+
+func testExpr() {
+  #freestandingExpr
+  // CHECK: [[@LINE-1]]:3 | function/Swift | exprLog() | [[EXPR_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+  // CHECK-NEXT: RelCall,RelCont | function/Swift | testExpr()
 }
 
-// CHECK: 6:33 | macro/Swift | myLine() | s:14swift_ide_test6myLineSiycfm | Def | rel: 0
-// CHECK: 6:45 | struct/Swift | Int | s:Si | Ref | rel: 0
-// CHECK: 7:33 | macro/Swift | myFilename() | s:14swift_ide_test10myFilenamexycs26ExpressibleByStringLiteralRzlufm | Def | rel: 0
-// CHECK: 7:47 | protocol/Swift | ExpressibleByStringLiteral | s:s26ExpressibleByStringLiteralP | Ref | rel: 0
-// CHECK: 8:33 | macro/Swift | myStringify(_:) | s:14swift_ide_test11myStringifyyx_SStxclufm | Def | rel: 0
+//--- IndexMacros.swift
+import SwiftSyntax
+import SwiftSyntaxBuilder
+import SwiftSyntaxMacros
+
+public struct FreestandingExprMacro: ExpressionMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) -> ExprSyntax {
+    return "exprLog()"
+  }
+}
+
+public struct FreestandingDeclMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["""
+      struct TestFree {
+        func freeFunc() {
+          freeLog()
+        }
+      }
+      """]
+  }
+}
+
+public struct SomeAccessorMacro: AccessorMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingAccessorsOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [AccessorDeclSyntax] {
+    return [
+      """
+        get {
+          accessorLog()
+          return 1
+        }
+      """,
+      """
+        set {
+          accessorLog()
+        }
+      """,
+    ]
+  }
+}
 
-// CHECK: 11:8 | macro/Swift | myLine() | s:14swift_ide_test6myLineSiycfm | Ref,RelCont | rel: 1
-// CHECK: 12:20 | macro/Swift | myFilename() | s:14swift_ide_test10myFilenamexycs26ExpressibleByStringLiteralRzlufm | Ref,RelCont | rel: 1
-// CHECK: 13:8 | macro/Swift | myStringify(_:) | s:14swift_ide_test11myStringifyyx_SStxclufm | Ref,RelCont | rel: 1
-// CHECK: 13:20 | param/Swift | x | s:14swift_ide_test0C01xySi_tFACL_Sivp | Ref,Read,RelCont | rel: 1
-// CHECK: 13:22 | static-method/infix-operator/Swift | +(_:_:) | s:Si1poiyS2i_SitFZ | Ref,Call,RelCall,RelCont | rel: 1
-// CHECK: 13:24 | param/Swift | x | s:14swift_ide_test0C01xySi_tFACL_Sivp | Ref,Read,RelCont | rel: 1
+public struct SomeConformanceMacro: ConformanceMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingConformancesOf decl: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [(TypeSyntax, GenericWhereClauseSyntax?)] {
+    let protocolName: TypeSyntax = "TestProto"
+    return [(protocolName, nil)]
+  }
+}
+
+public struct SomeMemberMacro: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf declaration: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    let newFunc: DeclSyntax =
+      """
+      func memberFunc() {
+        memberLog()
+      }
+      """
+    return [
+      newFunc,
+    ]
+  }
+}
+
+public struct SomeMemberAttributeMacro: MemberAttributeMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    attachedTo parent: some DeclGroupSyntax,
+    providingAttributesFor member: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [AttributeSyntax] {
+    guard let varDecl = member.as(VariableDeclSyntax.self),
+      let binding = varDecl.bindings.first,
+      let identifier = binding.pattern.as(IdentifierPatternSyntax.self)?.identifier.text,
+      identifier == "attachedMember"
+    else {
+      return []
+    }
+
+    return [AttributeSyntax(
+      attributeName: SimpleTypeIdentifierSyntax(
+        name: .identifier("AddOne")
+      )
+    )]
+  }
+}
+
+public struct SomePeerMacro: PeerMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return [
+      """
+      struct TestPeer {
+        func peerFunc() {
+          peerLog()
+        }
+      }
+      """
+    ]
+  }
+}
+
+public struct SomePeerMemberMacro: PeerMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return [
+      """
+      var peerMember: Int
+      """
+    ]
+  }
+}

Diff for: test/Macros/Inputs/syntax_macro_definitions.swift

@@ -1449,23 +1449,22 @@ public struct SimpleCodeItemMacro: CodeItemMacro {
   ) throws -> [CodeBlockItemSyntax] {
     [
       .init(item: .decl("""
-
       struct \(context.makeUniqueName("foo")) {
         var x: Int
       }
       """)),
       .init(item: .stmt("""
-
       if true {
         print("from stmt")
         usedInExpandedStmt()
       }
+      """)),
+      .init(item: .stmt("""
       if false {
         print("impossible")
       }
       """)),
       .init(item: .expr("""
-
       print("from expr")
       """)),
     ]