[SourceKit] Implement macro expansion for (some) attached macros.

DougGregor · DougGregor · commit 96380624db20 · 2023-02-11T11:23:28.000-08:00
Extend the macro-expansion refactoring to work with member and
member-attribute attached macros. These expansions can return several
different changes, e.g., adding new members, sprinkling member
attributes around, and so on.
diff --git a/include/swift/AST/ASTWalker.h b/include/swift/AST/ASTWalker.h
@@ -47,9 +47,17 @@ struct ReferenceMetaData {
   SemaReferenceKind Kind;
   llvm::Optional<AccessKind> AccKind;
   bool isImplicit = false;
-  ReferenceMetaData(SemaReferenceKind Kind, llvm::Optional<AccessKind> AccKind,
-                    bool isImplicit = false)
-      : Kind(Kind), AccKind(AccKind), isImplicit(isImplicit) {}
+
+  /// When non-none, this is a custom attribute reference.
+  Optional<std::pair<const CustomAttr *, Decl *>> CustomAttrRef;
+
+  ReferenceMetaData(
+      SemaReferenceKind Kind, llvm::Optional<AccessKind> AccKind,
+      bool isImplicit = false,
+      Optional<std::pair<const CustomAttr *, Decl *>> customAttrRef
+        = None
+  ) : Kind(Kind), AccKind(AccKind), isImplicit(isImplicit),
+      CustomAttrRef(customAttrRef) {}
 };
 
 /// Specifies how the initialization expression of a \c lazy variable should be
diff --git a/include/swift/IDE/Utils.h b/include/swift/IDE/Utils.h
@@ -151,6 +151,8 @@ struct ResolvedCursorInfo {
     bool IsRef = true;
     Type Ty;
     Type ContainerType;
+    Optional<std::pair<const CustomAttr *, Decl *>> CustomAttrRef = None;
+
     bool IsKeywordArgument = false;
     /// It this is a ref, whether it is "dynamic". See \c ide::isDynamicRef.
     bool IsDynamic = false;
@@ -262,6 +264,13 @@ struct ResolvedValueRefCursorInfo : public ResolvedCursorInfo {
                                   : ValueRefInfo.ValueD;
   }
 
+  Optional<std::pair<const CustomAttr *, Decl *>> getCustomAttrRef() const {
+    return ValueRefInfo.CustomAttrRef;
+  }
+  void setCustomAttrRef(Optional<std::pair<const CustomAttr *, Decl *>> ref) {
+    ValueRefInfo.CustomAttrRef = ref;
+  }
+
   static bool classof(const ResolvedCursorInfo *Info) {
     return Info->getKind() == CursorInfoKind::ValueRef;
   }
diff --git a/lib/IDE/IDERequests.cpp b/lib/IDE/IDERequests.cpp
@@ -94,7 +94,8 @@ class CursorInfoResolver : public SourceEntityWalker {
   bool rangeContainsLoc(CharSourceRange Range) const;
   bool isDone() const { return CursorInfo.isValid(); }
   bool tryResolve(ValueDecl *D, TypeDecl *CtorTyRef, ExtensionDecl *ExtTyRef,
-                  SourceLoc Loc, bool IsRef, Type Ty = Type());
+                  SourceLoc Loc, bool IsRef, Type Ty = Type(),
+                  Optional<ReferenceMetaData> Data = None);
   bool tryResolve(ModuleEntity Mod, SourceLoc Loc);
   bool tryResolve(Stmt *St);
   bool visitSubscriptReference(ValueDecl *D, CharSourceRange Range,
@@ -109,7 +110,8 @@ SourceManager &CursorInfoResolver::getSourceMgr() const
 
 bool CursorInfoResolver::tryResolve(ValueDecl *D, TypeDecl *CtorTyRef,
                                     ExtensionDecl *ExtTyRef, SourceLoc Loc,
-                                    bool IsRef, Type Ty) {
+                                    bool IsRef, Type Ty,
+                                    Optional<ReferenceMetaData> Data) {
   if (!D->hasName())
     return false;
 
@@ -137,6 +139,9 @@ bool CursorInfoResolver::tryResolve(ValueDecl *D, TypeDecl *CtorTyRef,
     }
   }
 
+  if (Data)
+    ValueRefInfo.setCustomAttrRef(Data->CustomAttrRef);
+
   CursorInfo = ValueRefInfo;
   return true;
 }
@@ -260,7 +265,7 @@ bool CursorInfoResolver::visitDeclReference(ValueDecl *D,
     return false;
   if (Data.isImplicit || !Range.isValid())
     return true;
-  return !tryResolve(D, CtorTyRef, ExtTyRef, Range.getStart(), /*IsRef=*/true, T);
+  return !tryResolve(D, CtorTyRef, ExtTyRef, Range.getStart(), /*IsRef=*/true, T, Data);
 }
 
 static bool isCursorOn(Expr *E, SourceLoc Loc) {
diff --git a/lib/IDE/SourceEntityWalker.cpp b/lib/IDE/SourceEntityWalker.cpp
@@ -19,6 +19,7 @@
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/Stmt.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/Defer.h"
@@ -691,9 +692,23 @@ bool SemaAnnotator::handleCustomAttributes(Decl *D) {
       return true;
     }
   }
+
   for (auto *customAttr : D->getAttrs().getAttributes<CustomAttr, true>()) {
     if (auto *Repr = customAttr->getTypeRepr()) {
-      if (!Repr->walk(*this))
+      // If this attribute resolves to a macro, index that.
+      ASTContext &ctx = D->getASTContext();
+      ResolveMacroRequest req{const_cast<CustomAttr *>(customAttr),
+                              getAttachedMacroRoles(),
+                              D->getInnermostDeclContext()};
+      if (auto macroDecl = evaluateOrDefault(ctx.evaluator, req, nullptr)) {
+        Type macroRefType = macroDecl->getDeclaredInterfaceType();
+        if (!passReference(
+              macroDecl, macroRefType, DeclNameLoc(Repr->getStartLoc()),
+              ReferenceMetaData(SemaReferenceKind::DeclRef, None,
+                                /*isImplicit=*/false,
+                                std::make_pair(customAttr, D))))
+          return false;
+      } else if (!Repr->walk(*this))
         return false;
     }
     if (auto *SemaInit = customAttr->getSemanticInit()) {
diff --git a/lib/Refactoring/Refactoring.cpp b/lib/Refactoring/Refactoring.cpp
@@ -8470,21 +8470,91 @@ bool RefactoringActionAddAsyncWrapper::performChange() {
   return false;
 }
 
-static MacroExpansionExpr *
-findMacroExpansionTargetExpr(const ResolvedCursorInfo &Info) {
+/// Retrieve the macro expansion buffer for the given macro expansion
+/// expression.
+static Optional<unsigned> getMacroExpansionBuffer(
+    SourceManager &sourceMgr, MacroExpansionExpr *expansion) {
+  if (auto rewritten = expansion->getRewritten()) {
+    return sourceMgr.findBufferContainingLoc(rewritten->getStartLoc());
+  }
+
+  return None;
+}
+
+/// Retrieve the macro expansion buffers for the given attached macro reference.
+static llvm::SmallVector<unsigned, 2>
+getMacroExpansionBuffers(MacroDecl *macro, const CustomAttr *attr, Decl *decl) {
+  auto roles = macro->getMacroRoles() & getAttachedMacroRoles();
+  if (!roles)
+    return { };
+
+  ASTContext &ctx = macro->getASTContext();
+  llvm::SmallVector<unsigned, 2> allBufferIDs;
+  if (roles.contains(MacroRole::Accessor)) {
+    // FIXME: Need to requestify.
+  }
+
+  if (roles.contains(MacroRole::MemberAttribute)) {
+    if (auto idc = dyn_cast<IterableDeclContext>(decl)) {
+      for (auto memberDecl : idc->getAllMembers()) {
+        auto bufferIDs = evaluateOrDefault(
+            ctx.evaluator, ExpandMemberAttributeMacros{memberDecl}, { });
+        allBufferIDs.append(bufferIDs.begin(), bufferIDs.end());
+      }
+    }
+  }
 
+  if (roles.contains(MacroRole::Member)) {
+    auto bufferIDs = evaluateOrDefault(
+        ctx.evaluator, ExpandSynthesizedMemberMacroRequest{decl}, { });
+    allBufferIDs.append(bufferIDs.begin(), bufferIDs.end());
+  }
+
+  // Drop any buffers that come from other macros. We could eliminate this
+  // step by adding more fine-grained requests above, which only expand for a
+  // single custom attribute.
+  SourceManager &sourceMgr = ctx.SourceMgr;
+  auto removedAt = std::remove_if(
+      allBufferIDs.begin(), allBufferIDs.end(),
+      [&](unsigned bufferID) {
+        auto generatedInfo = sourceMgr.getGeneratedSourceInfo(bufferID);
+        if (!generatedInfo)
+          return true;
+
+        return generatedInfo->attachedMacroCustomAttr != attr;
+      });
+  allBufferIDs.erase(removedAt, allBufferIDs.end());
+  return allBufferIDs;
+}
+
+/// Given a resolved cursor, determine whether it is for a macro expansion and
+/// return the list of macro expansion buffer IDs that are associated with the
+/// macro reference here.
+static llvm::SmallVector<unsigned, 2>
+getMacroExpansionBuffers(
+    SourceManager &sourceMgr, const ResolvedCursorInfo &Info
+) {
   // Handle '#' position in '#macroName(...)'.
   if (auto exprInfo = dyn_cast<ResolvedExprStartCursorInfo>(&Info)) {
     if (auto target =
-            dyn_cast_or_null<MacroExpansionExpr>(exprInfo->getTrailingExpr()))
-      if (target->getRewritten())
-        return target;
-    return nullptr;
+          dyn_cast_or_null<MacroExpansionExpr>(exprInfo->getTrailingExpr())) {
+      if (auto bufferID = getMacroExpansionBuffer(sourceMgr, target))
+        return { *bufferID };
+    }
+
+    return { };
   }
 
-  // Handle 'macroName' position in '#macroName(...)'.
   if (auto refInfo = dyn_cast<ResolvedValueRefCursorInfo>(&Info)) {
     if (refInfo->isRef() && isa_and_nonnull<MacroDecl>(refInfo->getValueD())) {
+      // Handle 'macroName' position in '@macroName(...)'.
+      if (auto customAttrRef = refInfo->getCustomAttrRef()) {
+        auto macro = cast<MacroDecl>(refInfo->getValueD());
+        return getMacroExpansionBuffers(
+            macro, customAttrRef->first, customAttrRef->second);
+      }
+
+      // Handle 'macroName' position in '#macroName(...)'.
       ContextFinder Finder(
           *Info.getSourceFile(), Info.getLoc(), [&](ASTNode N) {
             auto *expr =
@@ -8496,33 +8566,82 @@ findMacroExpansionTargetExpr(const ResolvedCursorInfo &Info) {
       if (!Finder.getContexts().empty()) {
         auto *target =
             dyn_cast<MacroExpansionExpr>(Finder.getContexts()[0].get<Expr *>());
-        if (target->getRewritten())
-          return target;
+        if (target) {
+          if (auto bufferID = getMacroExpansionBuffer(sourceMgr, target))
+            return { *bufferID };
+        }
       }
     }
-    return nullptr;
+
+    return { };
   }
 
   // TODO: handle MacroExpansionDecl.
-  return nullptr;
+  return { };
 }
 
 bool RefactoringActionExpandMacro::isApplicable(const ResolvedCursorInfo &Info,
                                                 DiagnosticEngine &Diag) {
-  return findMacroExpansionTargetExpr(Info) != nullptr;
+  return !getMacroExpansionBuffers(Diag.SourceMgr, Info).empty();
 }
 
 bool RefactoringActionExpandMacro::performChange() {
-  auto target = findMacroExpansionTargetExpr(CursorInfo);
-  if (!target)
+  auto bufferIDs = getMacroExpansionBuffers(SM, CursorInfo);
+  if (bufferIDs.empty())
     return true;
 
-  auto exprRange =
-      Lexer::getCharSourceRangeFromSourceRange(SM, target->getSourceRange());
-  auto rewrittenRange = Lexer::getCharSourceRangeFromSourceRange(
-      SM, target->getRewritten()->getSourceRange());
-  auto rewrittenBuffer = SM.extractText(rewrittenRange);
-  EditConsumer.accept(SM, exprRange, rewrittenBuffer);
+  // Send all of the rewritten buffer snippets.
+  CustomAttr *attachedMacroAttr = nullptr;
+  for (auto bufferID: bufferIDs) {
+    auto generatedInfo = SM.getGeneratedSourceInfo(bufferID);
+    if (!generatedInfo || generatedInfo->originalSourceRange.isInvalid())
+      continue;
+
+    auto rewrittenBuffer = SM.extractText(generatedInfo->generatedSourceRange);
+
+    // If there's no change, drop the edit entirely.
+    if (generatedInfo->originalSourceRange.getStart() ==
+          generatedInfo->originalSourceRange.getEnd() &&
+        rewrittenBuffer.empty())
+      continue;
+
+    auto originalSourceRange = generatedInfo->originalSourceRange;
+
+    // For member macros, adjust the source range from before-the-close-brace
+    // to after-the-open-brace.
+    if (generatedInfo->kind == GeneratedSourceInfo::MemberMacroExpansion) {
+      ASTNode node = ASTNode::getFromOpaqueValue(generatedInfo->astNode);
+      auto decl = node.dyn_cast<Decl *>();
+      if (!decl)
+        continue;
+
+      SourceLoc leftBraceLoc;
+      if (auto nominal = dyn_cast<NominalTypeDecl>(decl)) {
+        leftBraceLoc = nominal->getBraces().Start;
+      } else if (auto ext = dyn_cast<ExtensionDecl>(decl)) {
+        leftBraceLoc = ext->getBraces().Start;
+      }
+      if (leftBraceLoc.isInvalid())
+        continue;
+
+      auto afterLeftBraceLoc = Lexer::getLocForEndOfToken(SM, leftBraceLoc);
+      originalSourceRange = CharSourceRange(afterLeftBraceLoc, 0);
+    }
+
+    EditConsumer.accept(SM, originalSourceRange, rewrittenBuffer);
+
+    if (generatedInfo->attachedMacroCustomAttr && !attachedMacroAttr)
+      attachedMacroAttr = generatedInfo->attachedMacroCustomAttr;
+  }
+
+  // For an attached macro, remove the custom attribute; it's been fully
+  // subsumed by its expansions.
+  if (attachedMacroAttr) {
+    SourceRange range = attachedMacroAttr->getRangeWithAt();
+    auto charRange = Lexer::getCharSourceRangeFromSourceRange(SM, range);
+    EditConsumer.accept(SM, charRange, StringRef());
+  }
+
   return false;
 }
 
diff --git a/test/SourceKit/Macros/macro_basic.swift b/test/SourceKit/Macros/macro_basic.swift
@@ -4,6 +4,26 @@ func testStringify(a: Int, b: Int) {
   _ = #stringify(a + b)
 }
 
+@attached(memberAttribute)
+@attached(member)
+macro myTypeWrapper() = #externalMacro(module: "MacroDefinition", type: "TypeWrapperMacro")
+@attached(accessor) macro accessViaStorage() = #externalMacro(module: "MacroDefinition", type: "AccessViaStorageMacro")
+
+struct _Storage {
+  var x: Int = 0 {
+    willSet { print("setting \(newValue)") }
+  }
+  var y: Int = 0 {
+    willSet { print("setting \(newValue)") }
+  }
+}
+
+@myTypeWrapper
+struct S {
+  var x: Int
+  var y: Int
+}
+
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
@@ -54,3 +74,21 @@ func testStringify(a: Int, b: Int) {
 // EXPAND: source.edit.kind.active:
 // EXPAND-NEXT: 4:7-4:24 "(a + b, "a + b")"
 
+//##-- cursor-info at 'macro name' position following @.
+// RUN: %sourcekitd-test -req=cursor -pos=21:2 -cursor-action %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=CURSOR_ATTACHED_MACRO %s
+
+// CURSOR_ATTACHED_MACRO-LABEL: ACTIONS BEGIN
+// CURSOR_ATTACHED_MACRO: source.refactoring.kind.expand.macro
+// CURSOR_ATTACHED_MACRO-NEXT: Expand Macro
+// CURSOR_ATTACHED_MACRO: ACTIONS END
+
+//##-- Refactoring expanding the attached macro
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=21:2 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=ATTACHED_EXPAND %s
+// ATTACHED_EXPAND: source.edit.kind.active:
+// ATTACHED_EXPAND:   23:3-23:3 "@accessViaStorage"
+// ATTACHED_EXPAND: source.edit.kind.active:
+// ATTACHED_EXPAND:   24:3-24:3 "@accessViaStorage"
+// ATTACHED_EXPAND: source.edit.kind.active:
+// ATTACHED_EXPAND:   22:11-22:11 "private var _storage = _Storage()"
+// ATTACHED_EXPAND: source.edit.kind.active:
+// ATTACHED_EXPAND:   21:1-21:15 ""
