[Refactor] Allow expanding nested macros

This adds a new `primary_file` key, which defaults to `sourcefile`. For
nested expansions, `primary_file` should be set to the containing file
and `sourcefile` to the name of the macro expansion buffer.

Author: bnbarham

include/swift/AST/Decl.h

@@ -892,8 +892,8 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
 
   /// Retrieve the discriminator for the given custom attribute that names
   /// an attached macro.
-  unsigned getAttachedMacroDiscriminator(
-      MacroRole role, const CustomAttr *attr) const;
+  unsigned getAttachedMacroDiscriminator(DeclBaseName macroName, MacroRole role,
+                                         const CustomAttr *attr) const;
 
   /// Returns the innermost enclosing decl with an availability annotation.
   const Decl *getInnermostDeclWithAvailability() const;
@@ -8318,24 +8318,23 @@ class PostfixOperatorDecl : public OperatorDecl {
 /// is used for parser recovery, e.g. when parsing a floating
 /// attribute list.
 class MissingDecl: public Decl {
-  MissingDecl(DeclContext *DC) : Decl(DeclKind::Missing, DC) {
+  /// The location that the decl would be if it wasn't missing.
+  SourceLoc Loc;
+
+  MissingDecl(DeclContext *DC, SourceLoc loc)
+      : Decl(DeclKind::Missing, DC), Loc(loc) {
     setImplicit();
   }
 
   friend class Decl;
-  SourceLoc getLocFromSource() const {
-    return SourceLoc();
-  }
+  SourceLoc getLocFromSource() const { return Loc; }
 
 public:
-  static MissingDecl *
-  create(ASTContext &ctx, DeclContext *DC) {
-    return new (ctx) MissingDecl(DC);
+  static MissingDecl *create(ASTContext &ctx, DeclContext *DC, SourceLoc loc) {
+    return new (ctx) MissingDecl(DC, loc);
   }
 
-  SourceRange getSourceRange() const {
-    return SourceRange();
-  }
+  SourceRange getSourceRange() const { return SourceRange(Loc); }
 
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::Missing;

include/swift/Refactoring/Refactoring.h

@@ -33,7 +33,7 @@ enum class RefactoringKind : int8_t {
 };
 
 struct RangeConfig {
-  unsigned BufferId;
+  unsigned BufferID;
   unsigned Line;
   unsigned Column;
   unsigned Length;

lib/AST/ASTMangler.cpp

@@ -3810,13 +3810,14 @@ void ASTMangler::appendMacroExpansionContext(
 
     outerExpansionLoc = decl->getLoc();
     outerExpansionDC = decl->getDeclContext();
-    discriminator = decl->getAttachedMacroDiscriminator(role, attr);
 
     if (auto *macroDecl = decl->getResolvedMacro(attr))
       baseName = macroDecl->getBaseName();
     else
       baseName = ctx.getIdentifier("__unknown_macro__");
 
+    discriminator = decl->getAttachedMacroDiscriminator(baseName, role, attr);
+
     break;
   }
 
@@ -3896,24 +3897,30 @@ std::string ASTMangler::mangleAttachedMacroExpansion(
   beginMangling();
 
   const Decl *attachedTo = decl;
-  if (role == MacroRole::MemberAttribute) {
-    appendContextOf(cast<ValueDecl>(decl));
-    attachedTo = decl->getDeclContext()->getAsDecl();
-  } else if (auto valueDecl = dyn_cast<ValueDecl>(decl)) {
-    appendAnyDecl(valueDecl);
+  if (auto valueDecl = dyn_cast<ValueDecl>(decl)) {
+    if (role != MacroRole::MemberAttribute) {
+      appendAnyDecl(valueDecl);
+    } else {
+      // Appending the member would result in a cycle since `VarDecl` appends
+      // its type, which would then loop back around to getting the attributes
+      // again. We'll instead add a discriminator for each member.
+      appendContextOf(valueDecl);
+      attachedTo = decl->getDeclContext()->getAsDecl();
+    }
   } else {
     appendContext(decl->getDeclContext(), "");
   }
 
-  StringRef macroName;
-  if (auto *macroDecl = attachedTo->getResolvedMacro(attr))
-    macroName = macroDecl->getName().getBaseName().userFacingName();
-  else
-    macroName = "__unknown_macro__";
+  DeclBaseName macroName;
+  if (auto *macroDecl = attachedTo->getResolvedMacro(attr)) {
+    macroName = macroDecl->getName().getBaseName();
+  } else {
+    macroName = decl->getASTContext().getIdentifier("__unknown_macro__");
+  }
 
   appendMacroExpansionOperator(
-      macroName, role,
-      decl->getAttachedMacroDiscriminator(role, attr));
+      macroName.userFacingName(), role,
+      decl->getAttachedMacroDiscriminator(macroName, role, attr));
   return finalize();
 }
 

lib/AST/Decl.cpp

@@ -441,39 +441,34 @@ MacroDecl *Decl::getResolvedMacro(CustomAttr *customAttr) const {
   return dyn_cast_or_null<MacroDecl>(declRef.getDecl());
 }
 
-unsigned Decl::getAttachedMacroDiscriminator(
-    MacroRole role, const CustomAttr *attr
-) const {
+unsigned Decl::getAttachedMacroDiscriminator(DeclBaseName macroName,
+                                             MacroRole role,
+                                             const CustomAttr *attr) const {
   assert(isAttachedMacro(role) && "Not an attached macro role");
 
-  // Member-attribute macros are written on the enclosing declaration,
-  // but apply to the member itself. Adjust the owning declaration accordingly.
-  const Decl *owningDecl = this;
-  if (role == MacroRole::MemberAttribute) {
-    owningDecl = getDeclContext()->getAsDecl();
-  }
-
-  llvm::SmallDenseMap<Identifier, unsigned> nextDiscriminator;
-  Optional<unsigned> foundDiscriminator;
+  if (role != MacroRole::MemberAttribute) {
+    llvm::SmallDenseMap<Identifier, unsigned> nextDiscriminator;
+    Optional<unsigned> foundDiscriminator;
 
-  owningDecl->forEachAttachedMacro(
-    role,
-    [&](CustomAttr *foundAttr, MacroDecl *foundMacro) {
-      unsigned discriminator =
-        nextDiscriminator[foundMacro->getBaseIdentifier()]++;
-      if (attr == foundAttr)
-        foundDiscriminator = discriminator;
-    });
+    forEachAttachedMacro(
+        role, [&](CustomAttr *foundAttr, MacroDecl *foundMacro) {
+          unsigned discriminator =
+              nextDiscriminator[foundMacro->getBaseIdentifier()]++;
+          if (attr == foundAttr)
+            foundDiscriminator = discriminator;
+        });
 
-  if (foundDiscriminator)
-    return *foundDiscriminator;
+    if (foundDiscriminator)
+      return *foundDiscriminator;
+  }
 
   // If that failed, conjure up a discriminator.
+  // FIXME: Better discriminator for member attributes - add the member name?
   ASTContext &ctx = getASTContext();
-  assert(ctx.Diags.hadAnyError());
+  assert(role == MacroRole::MemberAttribute || ctx.Diags.hadAnyError());
   return ctx.getNextMacroDiscriminator(
       MacroDiscriminatorContext::getParentOf(getLoc(), getDeclContext()),
-      DeclBaseName());
+      macroName);
 }
 
 const Decl *Decl::getInnermostDeclWithAvailability() const {

lib/AST/Module.cpp

@@ -640,7 +640,9 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
 
   auto foundSourceFile = *found;
   auto foundRange = sourceMgr.getRangeForBuffer(*foundSourceFile->getBufferID());
-  if (!foundRange.contains(adjustedLoc) && adjustedLoc != foundRange.getStart())
+  // Positions inside an empty file or at EOF should still be considered within
+  // this file.
+  if (!foundRange.contains(adjustedLoc) && adjustedLoc != foundRange.getEnd())
     return nullptr;
 
   // Update the last source file.

lib/IDE/SourceEntityWalker.cpp

@@ -684,11 +684,13 @@ bool SemaAnnotator::handleCustomAttributes(Decl *D) {
     }
   }
 
-  // FIXME: This should be getSemanticAttrs if we want to walk macro
-  // expansions. We've just already typechecked and this list is mutable so...
-  for (auto *customAttr : D->getAttrs().getAttributes<CustomAttr, true>()) {
-    if (!shouldWalkMacroArgumentsAndExpansion().second &&
-        D->getModuleContext()->isInGeneratedBuffer(customAttr->getLocation()))
+  ModuleDecl *MD = D->getModuleContext();
+  for (auto *customAttr :
+       D->getSemanticAttrs().getAttributes<CustomAttr, true>()) {
+    SourceFile *SF =
+        MD->getSourceFileContainingLocation(customAttr->getLocation());
+    ASTNode expansion = SF ? SF->getMacroExpansion() : nullptr;
+    if (!shouldWalkMacroArgumentsAndExpansion().second && expansion)
       continue;
 
     if (auto *Repr = customAttr->getTypeRepr()) {
@@ -700,10 +702,12 @@ bool SemaAnnotator::handleCustomAttributes(Decl *D) {
       if (auto macroDecl = D->getResolvedMacro(mutableAttr)) {
         Type macroRefType = macroDecl->getDeclaredInterfaceType();
         if (!passReference(
-              macroDecl, macroRefType, DeclNameLoc(Repr->getStartLoc()),
-              ReferenceMetaData(SemaReferenceKind::DeclRef, None,
-                                /*isImplicit=*/false,
-                                std::make_pair(customAttr, D))))
+                macroDecl, macroRefType, DeclNameLoc(Repr->getStartLoc()),
+                ReferenceMetaData(
+                    SemaReferenceKind::DeclRef, None,
+                    /*isImplicit=*/false,
+                    std::make_pair(customAttr,
+                                   expansion ? expansion.get<Decl *>() : D))))
           return false;
       } else if (!Repr->walk(*this)) {
         return false;
@@ -865,9 +869,23 @@ bool SemaAnnotator::passCallArgNames(Expr *Fn, ArgumentList *ArgList) {
 }
 
 bool SemaAnnotator::shouldIgnore(Decl *D) {
+  if (!D->isImplicit())
+    return false;
+
   // TODO: There should really be a separate field controlling whether
   //       constructors are visited or not
-  return D->isImplicit() && !isa<ConstructorDecl>(D);
+  if (isa<ConstructorDecl>(D))
+    return false;
+
+  // Walk into missing decls to visit their attributes if they were generated
+  // by a member attribute expansion. Note that we would have already skipped
+  // this decl if we were ignoring expansions, so no need to check that.
+  if (auto *missing = dyn_cast<MissingDecl>(D)) {
+    if (D->isInGeneratedBuffer())
+      return false;
+  }
+
+  return true;
 }
 
 bool SourceEntityWalker::walk(SourceFile &SrcFile) {

lib/Parse/ParseDecl.cpp

@@ -7136,7 +7136,8 @@ void Parser::parseExpandedAttributeList(SmallVectorImpl<ASTNode> &items) {
 
   // Create a `MissingDecl` as a placeholder for the declaration the
   // macro will attach the attribute list to.
-  MissingDecl *missing = MissingDecl::create(Context, CurDeclContext);
+  MissingDecl *missing =
+      MissingDecl::create(Context, CurDeclContext, Tok.getLoc());
   missing->getAttrs() = attributes;
 
   items.push_back(ASTNode(missing));

lib/Refactoring/Refactoring.cpp

@@ -730,27 +730,12 @@ RenameRangeCollector::indexSymbolToRenameLoc(const index::IndexSymbol &symbol,
                    isFunctionLike, isNonProtocolType};
 }
 
-ArrayRef<SourceFile*>
-collectSourceFiles(ModuleDecl *MD, SmallVectorImpl<SourceFile *> &Scratch) {
-  for (auto Unit : MD->getFiles()) {
-    if (auto SF = dyn_cast<SourceFile>(Unit)) {
-      Scratch.push_back(SF);
-    }
-  }
-  return llvm::makeArrayRef(Scratch);
-}
-
-/// Get the source file that contains the given range and belongs to the module.
+/// Get the source file that corresponds to the given buffer.
 SourceFile *getContainingFile(ModuleDecl *M, RangeConfig Range) {
-  SmallVector<SourceFile*, 4> Files;
-  for (auto File : collectSourceFiles(M, Files)) {
-    if (File->getBufferID()) {
-      if (File->getBufferID().value() == Range.BufferId) {
-        return File;
-      }
-    }
-  }
-  return nullptr;
+  auto &SM = M->getASTContext().SourceMgr;
+  // TODO: We should add an ID -> SourceFile mapping.
+  return M->getSourceFileContainingLocation(
+      SM.getRangeForBuffer(Range.BufferID).getStart());
 }
 
 class RefactoringAction {
@@ -8802,7 +8787,7 @@ getDescriptiveRefactoringKindName(RefactoringKind Kind) {
   }
 
 SourceLoc swift::ide::RangeConfig::getStart(SourceManager &SM) {
-  return SM.getLocForLineCol(BufferId, Line, Column);
+  return SM.getLocForLineCol(BufferID, Line, Column);
 }
 
 SourceLoc swift::ide::RangeConfig::getEnd(SourceManager &SM) {

test/SourceKit/Macros/macro_basic.swift

@@ -138,14 +138,25 @@ struct S4 { }
 // ATTACHED_EXPAND: source.edit.kind.active:
 // ATTACHED_EXPAND-NEXT: 23:3-23:3 (@__swiftmacro_9MacroUser1SV13myTypeWrapperfMA_.swift) "@accessViaStorage "
 // ATTACHED_EXPAND-NEXT: source.edit.kind.active:
-// ATTACHED_EXPAND-NEXT: 24:3-24:3 (@__swiftmacro_9MacroUser1SV13myTypeWrapperfMA_.swift) "@accessViaStorage "
+// ATTACHED_EXPAND-NEXT: 24:3-24:3 (@__swiftmacro_9MacroUser1SV13myTypeWrapperfMA0_.swift) "@accessViaStorage "
 // ATTACHED_EXPAND-NEXT: source.edit.kind.active:
 // ATTACHED_EXPAND-NEXT: 22:11-22:11 (@__swiftmacro_9MacroUser1SV13myTypeWrapperfMm_.swift) "
 // ATTACHED_EXPAND-NEXT: private var _storage = _Storage()
 // ATTACHED_EXPAND-NEXT: "
 // ATTACHED_EXPAND-NEXT: source.edit.kind.active:
 // ATTACHED_EXPAND-NEXT: 21:1-21:15 ""
 
+//##-- Refactoring expanding the attribute expanded by @myTypeWrapper
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=1:2 @__swiftmacro_9MacroUser1SV13myTypeWrapperfMA_.swift -primary-file %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=NESTED_ATTACHED_EXPAND %s
+// NESTED_ATTACHED_EXPAND: source.edit.kind.active:
+// NESTED_ATTACHED_EXPAND-NEXT: Macros/macro_basic.swift 23:13-23:13 (@__swiftmacro_9MacroUser1SV1xSivp16accessViaStoragefMa_.swift) "{
+// NESTED_ATTACHED_EXPAND-NEXT:  get { _storage.x }
+// NESTED_ATTACHED_EXPAND-EMPTY:
+// NESTED_ATTACHED_EXPAND-NEXT:  set { _storage.x = newValue }
+// NESTED_ATTACHED_EXPAND-NEXT: }"
+// NESTED_ATTACHED_EXPAND-NEXT: source.edit.kind.active:
+// NESTED_ATTACHED_EXPAND-NEXT: 1:1-1:18 ""
+
 //##-- Refactoring expanding the first accessor macro
 // RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=30:4 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=ACCESSOR1_EXPAND %s
 // ACCESSOR1_EXPAND: source.edit.kind.active:

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -637,6 +637,10 @@ enum class SemanticRefactoringKind {
 
 struct SemanticRefactoringInfo {
   SemanticRefactoringKind Kind;
+  // The name of the source file to start the refactoring in. Empty if it is
+  // the primary file (in which case the primary file from the AST is used).
+  // This must match the buffer identifier stored in the source manager.
+  StringRef SourceFile;
   unsigned Line;
   unsigned Column;
   unsigned Length;
@@ -1045,7 +1049,7 @@ class LangSupport {
                         SourceKitCancellationToken CancellationToken,
                         CategorizedRenameRangesReceiver Receiver) = 0;
 
-  virtual void semanticRefactoring(StringRef Filename,
+  virtual void semanticRefactoring(StringRef PrimaryFile,
                                    SemanticRefactoringInfo Info,
                                    ArrayRef<const char *> Args,
                                    SourceKitCancellationToken CancellationToken,