Matching indentation with SyntaxRewriter

Using Swift
1

I've started playing around with a SyntaxRewriter to add @MainActor attributes to types or functions based on various criteria. This is a script that would run over a source tree and modify the source code in place. I'm somewhat close, but I not clear if this is the right approach, and specifically, I don't know how to deal with trivia.

My current version adds @MainActor to any type that inherits ObservableObject.

Given the following input:

import SwiftUI

enum Wrapper {
    /// An observed object
    internal final class Observered: ObservableObject {
    }
}

I'm currently able to generate the following output:

import SwiftUI

enum Wrapper {
@MainActor
    /// An observed object
    internal final class Observered: ObservableObject {
    }
}

So @MainActor is added, but it's at the wrong indentation, and should be after the comment rather than before it. I'm not sure how to approach this correctly, or what docs to read to learn how to do this properly.

class MainActorRewriter: SyntaxRewriter {
    override func visit(_ node: ClassDeclSyntax) -> DeclSyntax {
        
        let conformsToObservableObject = node.inheritanceClause?.inheritedTypes.contains {
            $0.type.as(IdentifierTypeSyntax.self)?.name.text == "ObservableObject"
        } ?? false

        let alreadyHasMainActor = node.attributes.contains(where: { $0.as(AttributeSyntax.self)?.attributeName == "@MainActor" })

        if conformsToObservableObject && !alreadyHasMainActor {
            let mainActorAttribute = AttributeSyntax("@MainActor").with(\.leadingTrivia, .newline)

            let newAttributes = node.attributes + [.attribute(mainActorAttribute)]

            return DeclSyntax(node.with(\.attributes, newAttributes))
        }

        // For classes not conforming to ObservableObject or already having @MainActor, return the node as is
        return DeclSyntax(node)
    }
}
2

Use

AttributeSyntax("@MainActor").with(\.leadingTrivia, node.leadingTrivia)

instead. Trailing trivia would be the last part (newlines and spaces) of the leading trivia of the declaration to get the indentation right.

That should work in your posted example. In case other attributes already exist, this becomes a bit trickier, though, depending on whether you want to have the new attribute first or last. Then you need to think thoroughly about which comments to put where. But the basic idea is the same: Take existing trivia and attach it to other nodes.

1 Like
3

I ended up asking, and answering, this question today: