Replacing a method body in a SyntaxRewriter

Using Swift
1

I've built a SyntaxRewriter as part of transitioning to 5.10 concurrency. It finds makeBody methods that need a MainActor.assumeIsolated wrapper, and rewrites them with the wrapper (for backward compatibility to 15.2, I've called it runUnsafely here). This code works well, except that I don't yet check if the change has already been applied. However, I'm wondering if it's working too hard. In particular, do I need to be reassembling the statements, then the body, then the decl, then MemberBlockItemSyntax, and finally re-inserting it into members? Or is there a more straightforward way? Is the double-reverse-prefix algorithm an appropriate way to get the indentation?

(inherits(from) is a custom extension)

class RewriteMakeBodyRewriter: SyntaxRewriter {
    override func visit(_ node: StructDeclSyntax) -> DeclSyntax {
        guard
            node.inherits(from: "ButtonStyle"),
            let makeBody = node.memberBlock.members
                .first(where: { $0.decl.as(FunctionDeclSyntax.self)?.name.text == "makeBody" }),
            let decl = makeBody.decl.as(FunctionDeclSyntax.self),
            let statements = decl.body?.statements
        else { return DeclSyntax(node) }

        let leadingSpaces = Trivia(pieces: statements.leadingTrivia.reversed().prefix { $0.isSpaceOrTab }.reversed())

        let newStatements = statements.with(\.leadingTrivia, leadingSpaces).with(\.trailingTrivia, [])
        let newBody = ExprSyntax("""
            MainActor.runUnsafely {
                \(newStatements)
            \(leadingSpaces)}
            """)
            .with(\.leadingTrivia, statements.leadingTrivia)
            .with(\.trailingTrivia, statements.trailingTrivia)

        let newCodeBlock = decl.body?.with(\.statements, [CodeBlockItemSyntax(item: .expr(newBody))])

        let newDecl = decl.with(\.body, newCodeBlock)

        let newMakeBody = makeBody.with(\.decl, DeclSyntax(newDecl))

        let index = node.memberBlock.members.index(of: makeBody)!
        var newMembers = node.memberBlock
        newMembers.members[index] = newMakeBody

        return DeclSyntax(node.with(\.memberBlock, newMembers))
    }
}

The following test case shows the function:

func testMakeBody() throws {
    let source = """
        public struct LinkButtonStyle: ButtonStyle {

            public func makeBody(configuration: Configuration) -> some View {
                // Comment inside body
                VStack {
                    EmptyView()
                }
            }
        }
        """

    let expected = """
        public struct LinkButtonStyle: ButtonStyle {

            public func makeBody(configuration: Configuration) -> some View {
                // Comment inside body
                MainActor.runUnsafely {
                    VStack {
                        EmptyView()
                    }
                }
            }
        }
        """

    let result = automain().process(source: source)
    XCTAssertEqual(result.description, expected)
}
2

If you start from the StructDeclSyntax, yes this is necessary.

My advice is to implement the visit methods only for the nodes that really require an update and that's CodeBlockSyntax or CodeBlockItemListSyntax in your case.

In there, you can check if the node's parent is a FunctionDeclSyntax with the right name. To make sure that you are in a type conforming to ButtonStyle, implement a little stack that tracks this detail and check it out when you visit the function body.

The same question I ask myself regularly, too. :sweat_smile: I think, that's fine, though.

1 Like
3

I originally tried writing it this way, but it created some corner cases I'm not certain how to address. In particular, how to detect nested types (including functions inside of functions). I think that's the point of your stack, to keep track of the current enclosing scope, but it wasn't clear how to build it and be confident it covered the cases. Do you create sub-Rewriter on the struct members? Are there any examples of that kind of setup?

This is my first stab at what I think you mean. I haven't checked corner cases, yet, though:

class RewriteMakeBodyRewriter: SyntaxRewriter {

    var inButtonStyle = false
    override func visit(_ node: StructDeclSyntax) -> DeclSyntax {
        let previousInButtonStyle = inButtonStyle

        inButtonStyle = node.inherits(from: "ButtonStyle")
        defer { inButtonStyle = previousInButtonStyle }

        return super.visit(node)
    }

    override func visit(_ node: FunctionDeclSyntax) -> DeclSyntax {
        guard inButtonStyle,
              node.name.text == "makeBody",
              let body = node.body,
              !body.statements.contains(where: { element in
                  element.tokens(viewMode: .fixedUp).contains { token in token.text == "MainActor" }})
        else { return super.visit(node) }

        let statements = body.statements

        let leadingSpaces = Trivia(pieces: statements.leadingTrivia.reversed().prefix { $0.isSpaceOrTab }.reversed())

        let newStatements = statements.with(\.leadingTrivia, leadingSpaces).with(\.trailingTrivia, [])
        let newBody = ExprSyntax("""
                MainActor.runUnsafely {
                    \(newStatements)
                \(leadingSpaces)}
                """)
            .with(\.leadingTrivia, statements.leadingTrivia)
            .with(\.trailingTrivia, statements.trailingTrivia)

        return DeclSyntax(
            node.with(\.body,
                       body.with(\.statements,
                                  [CodeBlockItemSyntax(item: .expr(newBody))])))
    }
}
4

Your example implementation looks quite good already. Similar to what I'd have written, too. I don't see a need for a sub-rewriter here, but, depending on the use case and complexity, this might also come in handy.

You might want to re-think your choice with viewMode: .fixedUp, as viewMode: .sourceAccurate is the recommended kind for a source transformation.

Regarding functions in functions: In case you only want to rewrite functions at the first level, think about checking the parent of FunctionDeclSyntax or implementing visitAny to skip traversing code blocks entirely (for example).

1 Like