adjust Indentation Of Freestanding Macro

Author: vinu-vanjari

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -78,10 +78,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)"
 }
 
@@ -103,13 +101,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)"
 }
 
@@ -131,13 +124,29 @@ 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 {
+  let indentationOfFirstLine = node.indentationOfFirstLine
+
+  var indentedSource =
+  expandedCode
+    .indented(by: indentationOfFirstLine)
+    .wrappingInTrivia(from: node)
+
+  // if the experssion started in middle of the line, then remove indentation of the first line
+  if !node.leadingTrivia.contains(where: \.isNewline) {
+    indentedSource.removeFirst(indentationOfFirstLine.sourceLength.utf8Length)
+  }
+  
+  return indentedSource
+}
+
 private func expandMemberMacro(
   definition: MemberMacro.Type,
   attributeNode: AttributeSyntax,

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,14 +281,113 @@ final class LexicalContextTests: XCTestCase {
       """,
       expandedSource: """
         extension A {
-          static var staticProp: String =   "staticProp"
+          static var staticProp: String = "staticProp"
         }
         """,
       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
+    )
+  }
+
   func testAllLexicalContexts() {
     assertMacroExpansion(
       """