[Macros] Automatically format expanded macros

Rather than requiring macro implementations to add required whitespace
and indentation, basic format all macro expansions. Right now this uses
the default four space indentation, we can consider having that inferred
later. Macros can opt-out of automatic formatting by implementing
`formatMode` and setting it to `.disabled`.

Also moves the extra newlines before/after expansions to a new "Inline
Macro" refactoring.

Resolves rdar://107731047.

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,67 @@ 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
+    if mode == .auto {
+      formatted = self.formatted()
+    } else {
+      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(of: "\n") {
+    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;                                                             \