Inserting syntax nodes in sensible places with respect to trivia

KeithBauerANZ · 2024-03-18T04:19:57.731Z

I'm trying to write a SyntaxRewriter that adds and removes @MainActor from various places to do a multi-thousand-line edit required by Swift 5.10.

The problem I'm having is, when I want to add @MainActor, how to get it in the right place with respect to existing "leading trivia" of the declaration I'm attaching it to.

For example, if I just do node.attributes.append(mainActor), then this code:

func xxx() {
     // unimportant
}

// MARK: - Stuff needing main actor

/// maybe a doc comment
func addMainActorHere() {
    // unimportant
}

Is transformed into this:

func xxx() {
     // unimportant
}@MainActor

// MARK: - Stuff needing main actor

/// maybe a doc comment
func addMainActorHere() {
    // unimportant
}

... which I mean, might be syntactically fine, but isn't an acceptable output of my tool

Obviously I need to move trivia around to achieve the desired effect. But what I can't figure out is, how to do that robustly. In particular

how to know if my added @MainActor will be the first node in the FunctionDeclSyntax (in which case I need to move all the existing leading trivia onto it)
how to reliably get the indentation of the FunctionDeclSyntax (since @MainActor should go on its own line, it generally needs its own copy of the indentation)
how to reliably distinguish "comments that are truly attached to the declaration" (doc comments, swiftlint:disable, etc.) from "comments that happen to be before the declaration" (// MARK:, above)

It's entirely possible that the "right" answer here is "just ignore this problem and run swift-format on the result", but our code has never been autoformatted before, and doing so now isn't a realistic possibility.

Any help or suggestions gratefully received!

SimplyDanny · 2024-03-18T08:59:45.995Z

A similar question came up recently. My answer to it might help you as well.

KeithBauerANZ · 2024-03-18T20:08:31.846Z

Thanks. I did try that, but duplicating the leading trivia in this case will copy everything after the } of the previous function, including the // MARK: - comment and the doc comment, leading to:

func xxx() {
     // unimportant
}

// MARK: - Stuff needing main actor

/// maybe a doc comment
@MainActor

// MARK: - Stuff needing main actor

/// maybe a doc comment
func addMainActorHere() {
    // unimportant
}

Which is much worse!

I've also tried

remove the leading trivia from the FunctionDeclSyntax (which takes it off the first token)
add @MainActor
add the leading trivia back to the FunctionDeclSyntax (which adds it back to whatever is the new first token)

That's slightly better, providing something like this:

    func xxx() {
         // unimportant
    }

    // MARK: - Stuff needing main actor

    /// maybe a doc comment
    @MainActor
func addMainActorHere() {
        // unimportant
    }

Which obviously now requires me to compute the indentation of the function to improve. Perhaps functionDecl.leadingTrivia.pieces.last is reliably enough the indent?

SimplyDanny · 2024-03-18T20:16:12.673Z

Yes, moving (not copying) the trivia to the new first token is the way to go.

Which obviously now requires me to compute the indentation of the function to improve. Perhaps functionDecl.leadingTrivia.pieces.last is reliably enough the indent?

The last piece should do. But you also want to have a line break. So you might want to use the last pieces of the trivia array that are not comments.

KeithBauerANZ · 2024-03-19T02:32:04.970Z

This seems more or less to work:

extension TriviaPiece {
    var isSwiftLintDirective: Bool {
        switch self {
        case let .lineComment(text):
            return text.hasPrefix("// swiftlint:")
        default:
            return false
        }
    }
}

extension Trivia {
    /// Splits the trivia into two parts: everything before the last SwiftLint directive,
    /// and the directive itself + everything after.
    func splitSwiftLintDirective() -> (leadingTrivia: Trivia, swiftLintDirective: Trivia?) {
        guard let swiftLintIndex = pieces.lastIndex(where: { $0.isSwiftLintDirective }) else {
            return (self, nil)
        }
        return (Trivia(pieces: pieces[..<swiftLintIndex]), Trivia(pieces: pieces[swiftLintIndex...]))
    }
}

extension SyntaxProtocol {
    mutating func addMainActor(_ attributes: WritableKeyPath<Self, AttributeListSyntax>) {
        guard !self[keyPath: attributes].hasMainActor else {
            return
        }

        let savedTrivia = leadingTrivia.splitSwiftLintDirective()
        let indent: Trivia = {
            if let last = savedTrivia.0.pieces.last,
               case .spaces = last
            {
                return [.newlines(1), last]
            } else {
                return [.newlines(1)]
            }
        }()

        leadingTrivia = []
        self[keyPath: attributes].append(
            .attribute(
                AttributeSyntax(attributeName: IdentifierTypeSyntax(name: .identifier("MainActor")))
                    .with(self[keyPath: attributes].isEmpty ? \.trailingTrivia : \.leadingTrivia, indent)
            )
        )
        if let swiftLintDirective = savedTrivia.swiftLintDirective {
            self[keyPath: attributes].trailingTrivia = self[keyPath: attributes]
                .trailingTrivia
                .appending(swiftLintDirective)
        }
        leadingTrivia = savedTrivia.leadingTrivia
    }
}

extension AttributeListSyntax.Element {
    var isMainActor: Bool {
        "MainActor" == self.as(AttributeSyntax.self)?
            .attributeName
            .as(IdentifierTypeSyntax.self)?
            .name
            .text
    }
}

extension AttributeListSyntax {

    var hasMainActor: Bool {
        contains(where: \.isMainActor)
    }

}

Probably still some edge cases to handle, but it's working pretty well for a lot of cases.