Merge pull request #65491 from bnbarham/cherry-format-macro-expansions

[5.9][Macros] Automatically format expanded macros

Author: bnbarham

include/swift/Refactoring/RefactoringKinds.def

@@ -66,6 +66,8 @@ CURSOR_REFACTORING(AddAsyncWrapper, "Add Async Wrapper", add.async-wrapper)
 
 CURSOR_REFACTORING(ExpandMacro, "Expand Macro", expand.macro)
 
+CURSOR_REFACTORING(InlineMacro, "Inline Macro", inline.macro)
+
 RANGE_REFACTORING(ExtractExpr, "Extract Expression", extract.expr)
 
 RANGE_REFACTORING(ExtractFunction, "Extract Method", extract.function)

lib/ASTGen/CMakeLists.txt

@@ -25,6 +25,7 @@ add_pure_swift_host_library(swiftASTGen STATIC
   DEPENDENCIES
     swiftAST
   SWIFT_DEPENDENCIES
+    SwiftSyntax::SwiftBasicFormat
     SwiftSyntax::SwiftDiagnostics
     SwiftSyntax::SwiftOperators
     SwiftSyntax::SwiftParser

lib/ASTGen/Sources/ASTGen/Macros.swift

@@ -566,10 +566,11 @@ func expandFreestandingMacroInProcess(
   }
 
   let macroPtr = macroPtr.bindMemory(to: ExportedMacro.self, capacity: 1)
+  let macro = macroPtr.pointee.macro
 
   let evaluatedSyntax: Syntax
   do {
-    switch macroPtr.pointee.macro {
+    switch macro {
     // Handle expression macro.
     case let exprMacro as ExpressionMacro.Type:
       func expandExpressionMacro<Node: FreestandingMacroExpansionSyntax>(
@@ -630,7 +631,7 @@ func expandFreestandingMacroInProcess(
     return nil
   }
 
-  return evaluatedSyntax.trimmedDescription
+  return evaluatedSyntax.formattedExpansion(macro.formatMode)
 }
 
 /// Retrieve a syntax node in the given source file, with the given type.
@@ -936,7 +937,7 @@ func expandAttachedMacroInProcess(
 
       // Form a buffer of accessor declarations to return to the caller.
       expandedSources = accessors.map {
-        $0.trimmedDescription
+        $0.formattedExpansion(macro.formatMode)
       }
 
     case (let attachedMacro as MemberAttributeMacro.Type, .MemberAttribute):
@@ -970,7 +971,7 @@ func expandAttachedMacroInProcess(
 
       // Form a buffer containing an attribute list to return to the caller.
       expandedSources = attributes.map {
-        $0.trimmedDescription
+        $0.formattedExpansion(macro.formatMode)
       }
 
     case (let attachedMacro as MemberMacro.Type, .Member):
@@ -998,7 +999,7 @@ func expandAttachedMacroInProcess(
 
       // Form a buffer of member declarations to return to the caller.
       expandedSources = members.map {
-        $0.trimmedDescription
+        $0.formattedExpansion(macro.formatMode)
       }
 
     case (let attachedMacro as PeerMacro.Type, .Peer):
@@ -1016,7 +1017,7 @@ func expandAttachedMacroInProcess(
 
       // Form a buffer of peer declarations to return to the caller.
       expandedSources = peers.map {
-        $0.trimmedDescription
+        $0.formattedExpansion(macro.formatMode)
       }
 
     case (let attachedMacro as ConformanceMacro.Type, .Conformance):
@@ -1064,43 +1065,68 @@ func expandAttachedMacroInProcess(
   return expandedSources
 }
 
+fileprivate extension SyntaxProtocol {
+  /// Perform a format if required and then trim any leading/trailing
+  /// whitespace.
+  func formattedExpansion(_ mode: FormatMode) -> String {
+    let formatted: Syntax
+    switch mode {
+    case .auto:
+      formatted = self.formatted()
+    case .disabled:
+      formatted = Syntax(self)
+    }
+    return formatted.trimmedDescription(matching: { $0.isWhitespace })
+  }
+}
+
 fileprivate func collapse<Node: SyntaxProtocol>(
   expansions: [String],
   for role: MacroRole,
   attachedTo declarationNode: Node
 ) -> String {
+  if expansions.isEmpty {
+    return ""
+  }
+
+  var expansions = expansions
   var separator: String = "\n\n"
-  var prefix: String = ""
-  var suffix: String = ""
-
-  switch role {
-  case .Accessor:
-    if let varDecl = declarationNode.as(VariableDeclSyntax.self),
-       let binding = varDecl.bindings.first,
-       binding.accessor == nil {
-      prefix = "{\n"
-      suffix = "\n}"
-    }
-  case .Member:
-    prefix = "\n\n"
-    suffix = "\n"
-  case .MemberAttribute:
+
+  if role == .Accessor,
+     let varDecl = declarationNode.as(VariableDeclSyntax.self),
+     let binding = varDecl.bindings.first,
+     binding.accessor == nil {
+    let indentation = String(repeating: " ", count: 4)
+
+    expansions = expansions.map({ indent($0, with: indentation) })
+    expansions[0] = "{\n" + expansions[0]
+    expansions[expansions.count - 1] += "\n}"
+  } else if role == .MemberAttribute {
     separator = " "
-    suffix = " "
-  case .Peer:
-    prefix = "\n\n"
-    suffix = "\n"
-  case .Conformance:
-    prefix = "\n\n"
-    suffix = "\n"
-  case .Expression,
-      .FreestandingDeclaration:
-    fatalError("unreachable")
   }
 
-  let separated = expansions.joined(separator: separator)
-  if separated.isEmpty {
-    return separated
+  return expansions.joined(separator: separator)
+}
+
+fileprivate func indent(_ source: String, with indentation: String) -> String {
+  if source.isEmpty || indentation.isEmpty {
+    return source
   }
-  return prefix + separated + suffix
+
+  var indented = ""
+  var remaining = source[...]
+  while let nextNewline = remaining.firstIndex(where: { $0.isNewline }) {
+    if nextNewline != remaining.startIndex {
+      indented += indentation
+    }
+    indented += remaining[...nextNewline]
+    remaining = remaining[remaining.index(after: nextNewline)...]
+  }
+
+  if !remaining.isEmpty {
+    indented += indentation
+    indented += remaining
+  }
+
+  return indented
 }

lib/Refactoring/Refactoring.cpp

@@ -8756,18 +8756,53 @@ getMacroExpansionBuffers(SourceManager &sourceMgr, ResolvedCursorInfoPtr Info) {
   return {};
 }
 
-bool RefactoringActionExpandMacro::isApplicable(ResolvedCursorInfoPtr Info,
-                                                DiagnosticEngine &Diag) {
-  return !getMacroExpansionBuffers(Diag.SourceMgr, Info).empty();
-}
-
-bool RefactoringActionExpandMacro::performChange() {
-  auto bufferIDs = getMacroExpansionBuffers(SM, CursorInfo);
+/// Given the expanded code for a particular macro, perform whitespace
+/// adjustments to make the refactoring more suitable for inline insertion.
+static StringRef adjustMacroExpansionWhitespace(
+    GeneratedSourceInfo::Kind kind, StringRef expandedCode,
+    llvm::SmallString<64> &scratch) {
+  scratch.clear();
+
+  switch (kind) {
+  case GeneratedSourceInfo::MemberAttributeMacroExpansion:
+    // Attributes are added to the beginning, add a space to separate from
+    // any existing.
+    scratch += expandedCode;
+    scratch += " ";
+    return scratch;
+
+  case GeneratedSourceInfo::MemberMacroExpansion:
+  case GeneratedSourceInfo::PeerMacroExpansion:
+  case GeneratedSourceInfo::ConformanceMacroExpansion:
+    // All added to the end. Note that conformances are always expanded as
+    // extensions, hence treating them the same as peer.
+    scratch += "\n\n";
+    scratch += expandedCode;
+    scratch += "\n";
+    return scratch;
+
+  case GeneratedSourceInfo::ExpressionMacroExpansion:
+  case GeneratedSourceInfo::FreestandingDeclMacroExpansion:
+  case GeneratedSourceInfo::AccessorMacroExpansion:
+  case GeneratedSourceInfo::ReplacedFunctionBody:
+  case GeneratedSourceInfo::PrettyPrinted:
+    return expandedCode;
+  }
+}
+
+static bool expandMacro(SourceManager &SM, ResolvedCursorInfoPtr cursorInfo,
+                        SourceEditConsumer &editConsumer, bool adjustExpansion) {
+  auto bufferIDs = getMacroExpansionBuffers(SM, cursorInfo);
   if (bufferIDs.empty())
     return true;
 
+  SourceFile *containingSF = cursorInfo->getSourceFile();
+  if (!containingSF)
+    return true;
+
   // Send all of the rewritten buffer snippets.
   CustomAttr *attachedMacroAttr = nullptr;
+  SmallString<64> scratchBuffer;
   for (auto bufferID: bufferIDs) {
     auto generatedInfo = SM.getGeneratedSourceInfo(bufferID);
     if (!generatedInfo || generatedInfo->originalSourceRange.isInvalid())
@@ -8781,8 +8816,12 @@ bool RefactoringActionExpandMacro::performChange() {
         rewrittenBuffer.empty())
       continue;
 
-    // `TheFile` is the file of the actual expansion site, where as
-    // `OriginalFile` is the possibly enclosing buffer. Concretely:
+    if (adjustExpansion) {
+      rewrittenBuffer = adjustMacroExpansionWhitespace(generatedInfo->kind, rewrittenBuffer, scratchBuffer);
+    }
+
+    // `containingFile` is the file of the actual expansion site, where as
+    // `originalFile` is the possibly enclosing buffer. Concretely:
     // ```
     // // m.swift
     // @AddMemberAttributes
@@ -8801,14 +8840,14 @@ bool RefactoringActionExpandMacro::performChange() {
     // expansion.
     auto originalSourceRange = generatedInfo->originalSourceRange;
     SourceFile *originalFile =
-        MD->getSourceFileContainingLocation(originalSourceRange.getStart());
+        containingSF->getParentModule()->getSourceFileContainingLocation(originalSourceRange.getStart());
     StringRef originalPath;
     if (originalFile->getBufferID().hasValue() &&
-        TheFile->getBufferID() != originalFile->getBufferID()) {
+        containingSF->getBufferID() != originalFile->getBufferID()) {
       originalPath = SM.getIdentifierForBuffer(*originalFile->getBufferID());
     }
 
-    EditConsumer.accept(SM, {originalPath,
+    editConsumer.accept(SM, {originalPath,
                              originalSourceRange,
                              SM.getIdentifierForBuffer(bufferID),
                              rewrittenBuffer,
@@ -8823,12 +8862,31 @@ bool RefactoringActionExpandMacro::performChange() {
   if (attachedMacroAttr) {
     SourceRange range = attachedMacroAttr->getRangeWithAt();
     auto charRange = Lexer::getCharSourceRangeFromSourceRange(SM, range);
-    EditConsumer.accept(SM, charRange, StringRef());
+    editConsumer.accept(SM, charRange, StringRef());
   }
 
   return false;
 }
 
+bool RefactoringActionExpandMacro::isApplicable(ResolvedCursorInfoPtr Info,
+                                                DiagnosticEngine &Diag) {
+  // Never list in available refactorings. Only allow requesting directly.
+  return false;
+}
+
+bool RefactoringActionExpandMacro::performChange() {
+  return expandMacro(SM, CursorInfo, EditConsumer, /*adjustExpansion=*/false);
+}
+
+bool RefactoringActionInlineMacro::isApplicable(ResolvedCursorInfoPtr Info,
+                                                DiagnosticEngine &Diag) {
+  return !getMacroExpansionBuffers(Diag.SourceMgr, Info).empty();
+}
+
+bool RefactoringActionInlineMacro::performChange() {
+  return expandMacro(SM, CursorInfo, EditConsumer, /*adjustExpansion=*/true);
+}
+
 } // end of anonymous namespace
 
 StringRef swift::ide::
@@ -8940,8 +8998,8 @@ swift::ide::collectRefactorings(ResolvedCursorInfoPtr CursorInfo,
 
   // Only macro expansion is available within generated buffers
   if (CursorInfo->getSourceFile()->Kind == SourceFileKind::MacroExpansion) {
-    if (RefactoringActionExpandMacro::isApplicable(CursorInfo, DiagEngine)) {
-      Infos.emplace_back(RefactoringKind::ExpandMacro,
+    if (RefactoringActionInlineMacro::isApplicable(CursorInfo, DiagEngine)) {
+      Infos.emplace_back(RefactoringKind::InlineMacro,
                          RefactorAvailableKind::Available);
     }
     return Infos;
@@ -9019,6 +9077,7 @@ refactorSwiftModule(ModuleDecl *M, RefactoringOptions Opts,
 case RefactoringKind::KIND: {                                                  \
       RefactoringAction##KIND Action(M, Opts, EditConsumer, DiagConsumer);     \
       if (RefactoringKind::KIND == RefactoringKind::LocalRename ||             \
+          RefactoringKind::KIND == RefactoringKind::ExpandMacro ||             \
           Action.isApplicable())                                               \
         return Action.performChange();                                         \
       return true;                                                             \