Merge pull request #2493 from vinu-vanjari/macro-adjustIndentationOfFreestandingMacro

adjust Indentation Of Freestanding Macro

Author: ahoppen

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -93,10 +93,8 @@ private func expandFreestandingMemberDeclList(
     return nil
   }
 
-  let indentedSource =
-    expanded
-    .indented(by: node.indentationOfFirstLine)
-    .wrappingInTrivia(from: node)
+  let indentedSource = adjustIndentationOfFreestandingMacro(expandedCode: expanded, node: node)
+
   return "\(raw: indentedSource)"
 }
 
@@ -118,13 +116,8 @@ private func expandFreestandingCodeItemList(
     return nil
   }
 
-  // The macro expansion just provides an expansion for the content.
-  // We need to make sure that we aren’t dropping the trivia before and after
-  // the expansion.
-  let indentedSource =
-    expanded
-    .indented(by: node.indentationOfFirstLine)
-    .wrappingInTrivia(from: node)
+  let indentedSource = adjustIndentationOfFreestandingMacro(expandedCode: expanded, node: node)
+
   return "\(raw: indentedSource)"
 }
 
@@ -146,13 +139,36 @@ private func expandFreestandingExpr(
     return nil
   }
 
-  let indentedSource =
-    expanded
-    .indented(by: node.indentationOfFirstLine)
-    .wrappingInTrivia(from: node)
+  let indentedSource = adjustIndentationOfFreestandingMacro(expandedCode: expanded, node: node)
+
   return "\(raw: indentedSource)"
 }
 
+/// Adds the appropriate indentation on expanded code even if it's multi line.
+/// Makes sure original macro expression's trivia is maintained by adding it to expanded code.
+private func adjustIndentationOfFreestandingMacro(expandedCode: String, node: some FreestandingMacroExpansionSyntax) -> String {
+
+  if expandedCode.isEmpty {
+    return expandedCode.wrappingInTrivia(from: node)
+  }
+
+  let indentationOfFirstLine = node.indentationOfFirstLine
+  let indentLength = indentationOfFirstLine.sourceLength.utf8Length
+
+  // we are doing 3 step adjustment here
+  // step 1: add indentation to each line of expanded code
+  // step 2: remove indentation from first line of expaned code
+  // step 3: wrap the expanded code into macro expression's trivia. This trivia will contain appropriate existing
+  // indentation. Note that if macro expression occurs in middle of the line then there will be no indentation or extra space.
+  // Hence we are doing step 2
+
+  var indentedSource = expandedCode.indented(by: indentationOfFirstLine)
+  indentedSource.removeFirst(indentLength)
+  indentedSource = indentedSource.wrappingInTrivia(from: node)
+
+  return indentedSource
+}
+
 private func expandMemberMacro(
   definition: MemberMacro.Type,
   attributeNode: AttributeSyntax,
@@ -1284,9 +1300,7 @@ private extension String {
   /// user should think about it as just replacing the `#...` expression without
   /// any trivia.
   func wrappingInTrivia(from node: some SyntaxProtocol) -> String {
-    // We need to remove the indentation from the last line because the macro
-    // expansion buffer already contains the indentation.
-    return node.leadingTrivia.removingIndentationOnLastLine.description
+    return node.leadingTrivia.description
       + self
       + node.trailingTrivia.description
   }

Tests/SwiftSyntaxMacroExpansionTest/LexicalContextTests.swift

@@ -173,6 +173,24 @@ public struct FunctionMacro: ExpressionMacro {
   }
 }
 
+public struct MultilineFunctionMacro: ExpressionMacro {
+  public static func expansion<
+    Node: FreestandingMacroExpansionSyntax,
+    Context: MacroExpansionContext
+  >(
+    of node: Node,
+    in context: Context
+  ) -> ExprSyntax {
+    guard let lexicalContext = context.lexicalContext.first,
+      let name = lexicalContext.functionName(in: context)
+    else {
+      return #""<unknown>""#
+    }
+
+    return ExprSyntax("{\n\(literal: name)\n}")
+  }
+}
+
 public struct AllLexicalContextsMacro: DeclarationMacro {
   public static func expansion(
     of node: some FreestandingMacroExpansionSyntax,
@@ -194,7 +212,7 @@ final class LexicalContextTests: XCTestCase {
       """,
       expandedSource: """
         func f(a: Int, _: Double, c: Int) {
-          print(  "f(a:_:c:)")
+          print("f(a:_:c:)")
         }
         """,
       macros: ["function": FunctionMacro.self],
@@ -263,12 +281,182 @@ final class LexicalContextTests: XCTestCase {
       """,
       expandedSource: """
         extension A {
-          static var staticProp: String =   "staticProp"
+          static var staticProp: String = "staticProp"
         }
         """,
       macros: ["function": FunctionMacro.self],
       indentationWidth: indentationWidth
     )
+
+    assertMacroExpansion(
+      """
+      func f(a: Int, _: Double, c: Int) {
+        print(/*comment*/#function)
+      }
+      """,
+      expandedSource: """
+        func f(a: Int, _: Double, c: Int) {
+          print(/*comment*/"f(a:_:c:)")
+        }
+        """,
+      macros: ["function": FunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      var computed: String {
+        get {
+          /*comment*/#function
+        }
+      }
+      """,
+      expandedSource: """
+        var computed: String {
+          get {
+            /*comment*/"computed"
+          }
+        }
+        """,
+      macros: ["function": FunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testPoundMultilineFunction() {
+    assertMacroExpansion(
+      """
+      func f(a: Int, _: Double, c: Int) {
+        print(#function)
+      }
+      """,
+      expandedSource: """
+        func f(a: Int, _: Double, c: Int) {
+          print({
+            "f(a:_:c:)"
+          })
+        }
+        """,
+      macros: ["function": MultilineFunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      struct X {
+        init(from: String) {
+          #function
+        }
+
+        subscript(a: Int) -> String {
+          #function
+        }
+
+        subscript(a a: Int) -> String {
+          #function
+        }
+      }
+      """,
+      expandedSource: """
+        struct X {
+          init(from: String) {
+            {
+              "init(from:)"
+            }
+          }
+
+          subscript(a: Int) -> String {
+            {
+              "subscript(_:)"
+            }
+          }
+
+          subscript(a a: Int) -> String {
+            {
+              "subscript(a:)"
+            }
+          }
+        }
+        """,
+      macros: ["function": MultilineFunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      var computed: String {
+        get {
+          #function
+        }
+      }
+      """,
+      expandedSource: """
+        var computed: String {
+          get {
+            {
+              "computed"
+            }
+          }
+        }
+        """,
+      macros: ["function": MultilineFunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      extension A {
+        static var staticProp: String = #function
+      }
+      """,
+      expandedSource: """
+        extension A {
+          static var staticProp: String = {
+            "staticProp"
+          }
+        }
+        """,
+      macros: ["function": MultilineFunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      func f(a: Int, _: Double, c: Int) {
+        print(/*comment*/#function)
+      }
+      """,
+      expandedSource: """
+        func f(a: Int, _: Double, c: Int) {
+          print(/*comment*/{
+            "f(a:_:c:)"
+          })
+        }
+        """,
+      macros: ["function": MultilineFunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      var computed: String {
+        get {
+          /*comment*/#function // another comment
+        }
+      }
+      """,
+      expandedSource: """
+        var computed: String {
+          get {
+            /*comment*/{
+              "computed"
+            } // another comment
+          }
+        }
+        """,
+      macros: ["function": MultilineFunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
   }
 
   func testAllLexicalContexts() {