Indentation issue with BodyMacro expansion

Using Swift
1

Hi everyone,

I'm working on a Swift macro called @AutoGuardSelf that automatically inserts a guard let self else { return } statement in closures capturing self weakly. The macro works as expected, but I'm running into an indentation issue when expanding the function body.

Here’s the relevant implementation:

import SwiftSyntax
import SwiftSyntaxMacros

public struct AutoGuardSelfMacro: BodyMacro {
	public static func expansion(
		of node: AttributeSyntax,
		providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
		in context: some MacroExpansionContext
	) throws -> [CodeBlockItemSyntax] {
		guard
			let functionDecl = declaration.as(FunctionDeclSyntax.self),
			let body = functionDecl.body
		else {
			throw MacroError.message("Expected a function declaration with a body.")
		}

		if let newBody = AutoGuardSelfRewriter().rewrite(body).as(CodeBlockSyntax.self) {
			return Array(newBody.statements)
		} else {
			throw MacroError.message("Failed to rewrite function body.")
		}
	}
}

final class AutoGuardSelfRewriter: SyntaxRewriter {
	var shouldInsertGuardSelf: Bool? = nil

	override func visit(_ node: ClosureExprSyntax) -> ExprSyntax {
		guard
			let signature = node.signature,
			let capture = signature.capture
		else {
			return super.visit(node)
		}

		let containsWeakSelfCapture = capture.items.contains {
			$0.specifier?.specifier.text == "weak" &&
			$0.expression.as(DeclReferenceExprSyntax.self)?.baseName.text == "self"
		}

		if containsWeakSelfCapture {
			shouldInsertGuardSelf = true
		}

		return super.visit(node)
	}

	override func visit(_ node: CodeBlockItemListSyntax) -> CodeBlockItemListSyntax {
		guard shouldInsertGuardSelf == true else {
			return super.visit(node)
		}

		shouldInsertGuardSelf = nil

		let guardStmt = GuardStmtSyntax(
			conditions: [
				ConditionElementSyntax(
					condition: .optionalBinding(
						OptionalBindingConditionSyntax(
							bindingSpecifier: .keyword(.let),
							pattern: IdentifierPatternSyntax(identifier: .keyword(.`self`))
						)
					)
				)
			],
			body: CodeBlockSyntax(statements: [
				CodeBlockItemSyntax(item: .stmt(StmtSyntax(ReturnStmtSyntax())))
			])
		)

		return [CodeBlockItemSyntax(item: .stmt(StmtSyntax(guardStmt)))] + node
	}
}

enum MacroError: Error, CustomStringConvertible {
	case message(String)

	var description: String {
		switch self {
		case .message(let msg):
			return msg
		}
	}
}

The Issue

When expanding the function body in a test case, the indentation gets messed up. Here's an example:

Input:

@AutoGuardSelf
func myFunction() {
	foo.bar { [weak self] in
		self.baz()
	}
}

Expected Output:

func myFunction() {
	foo.bar { [weak self] in
		guard let self else {
			return
		}
		self.baz()
	}
}

Actual Output:

func myFunction() {
	foo.bar { [weak self] in
		guard let self else {
			return
		}
				self.baz()
		}
}

As you can see, the indentation of self.baz() and the closing } is incorrect.

I assume SwiftSyntax handles formatting, but it doesn’t seem to be working here. Is there a best practice for maintaining indentation when modifying the syntax tree?

Any guidance would be greatly appreciated.

2

Found a solution:

import SwiftSyntax

extension SyntaxProtocol {
	var recursivelyTrimmed: Self {
		let trimmedSyntax = TriviaTrimmer().rewrite(Syntax(self))
		return trimmedSyntax.as(Self.self) ?? self
	}
}

private class TriviaTrimmer: SyntaxRewriter {
	override func visit(_ node: TokenSyntax) -> TokenSyntax {
		node.trimmed
	}
}

then...

if let newBody = AutoGuardSelfRewriter().rewrite(body.recursivelyTrimmed).as(CodeBlockSyntax.self) {