Skip to content

Commit

Permalink
Add AttributeRemover600 and conditional usage
Browse files Browse the repository at this point in the history
- AttributeRemover600 is vendored from SwiftSyntax
- Conditionally use AttributeRemover600 or AttributeRemover509
  • Loading branch information
gohanlon committed Jul 12, 2024
1 parent 398e1b0 commit 8dee94c
Show file tree
Hide file tree
Showing 2 changed files with 128 additions and 9 deletions.
23 changes: 14 additions & 9 deletions Sources/MacroTesting/AssertMacro.swift
Original file line number Diff line number Diff line change
Expand Up @@ -278,17 +278,22 @@ public func assertMacro(
// TODO: write a test where didExpand returns false
// For now, covered in MemberwiseInitTests.testAppliedToEnum_FailsWithDiagnostic
var didExpand: Bool {
let origSourceWithMacroAttributesRemoved = MacroTesting.AttributeRemover509(
removingWhere: {
guard let name = $0.attributeName.as(IdentifierTypeSyntax.self)?.name.text
else { return false }
return macros.keys.contains(name)
}
).rewrite(origSourceFile)
let removingWhere: (AttributeSyntax) -> Bool = {
guard let name = $0.attributeName.as(IdentifierTypeSyntax.self)?.name.text
else { return false }
return macros.keys.contains(name)
}

#if canImport(SwiftSyntax600)
let remover = MacroTesting.AttributeRemover600(removingWhere: removingWhere)
#else
let remover = MacroTesting.AttributeRemover509(removingWhere: removingWhere)
#endif

let removedSource = remover.rewrite(origSourceFile)

return expandedSourceFile.description.trimmingCharacters(in: .newlines)
!= origSourceWithMacroAttributesRemoved.description.trimmingCharacters(
in: .newlines)
!= removedSource.description.trimmingCharacters(in: .newlines)
}

if didExpand {
Expand Down
114 changes: 114 additions & 0 deletions Sources/MacroTesting/Internal/AttributeRemover600.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
import SwiftSyntax

public class AttributeRemover600: SyntaxRewriter {
let predicate: (AttributeSyntax) -> Bool

var triviaToAttachToNextToken: Trivia = Trivia()

/// Initializes an attribute remover with a given predicate to determine which attributes to remove.
///
/// - Parameter predicate: A closure that determines whether a given `AttributeSyntax` should be removed.
/// If this closure returns `true` for an attribute, that attribute will be removed.
public init(removingWhere predicate: @escaping (AttributeSyntax) -> Bool) {
self.predicate = predicate
}

public override func visit(_ node: AttributeListSyntax) -> AttributeListSyntax {
var filteredAttributes: [AttributeListSyntax.Element] = []
for case .attribute(let attribute) in node {
if self.predicate(attribute) {
var leadingTrivia = attribute.leadingTrivia

// Don't leave behind an empty line when the attribute being removed is on its own line,
// based on the following conditions:
// - Leading trivia ends with a newline followed by arbitrary number of spaces or tabs
// - All leading trivia pieces after the last newline are just whitespace, ensuring
// there are no comments or other non-whitespace characters on the same line
// preceding the attribute.
// - There is no trailing trivia and the next token has leading trivia.
if let lastNewline = leadingTrivia.pieces.lastIndex(where: \.isNewline),
leadingTrivia.pieces[lastNewline...].allSatisfy(\.isWhitespace),
attribute.trailingTrivia.isEmpty,
let nextToken = attribute.nextToken(viewMode: .sourceAccurate),
!nextToken.leadingTrivia.isEmpty
{
leadingTrivia = Trivia(pieces: leadingTrivia.pieces[..<lastNewline])
}

// Drop any spaces or tabs from the trailing trivia because there’s no
// more attribute they need to separate.
let trailingTrivia = attribute.trailingTrivia.trimmingPrefix(while: \.isSpaceOrTab)
triviaToAttachToNextToken += leadingTrivia + trailingTrivia

// If the attribute is not separated from the previous attribute by trivia, as in
// `@First@Second var x: Int` (yes, that's valid Swift), removing the `@Second`
// attribute and dropping all its trivia would cause `@First` and `var` to join
// without any trivia in between, which is invalid. In such cases, the trailing trivia
// of the attribute is significant and must be retained.
if triviaToAttachToNextToken.isEmpty,
let previousToken = attribute.previousToken(viewMode: .sourceAccurate),
previousToken.trailingTrivia.isEmpty
{
triviaToAttachToNextToken = attribute.trailingTrivia
}
} else {
filteredAttributes.append(.attribute(prependAndClearAccumulatedTrivia(to: attribute)))
}
}

// Ensure that any horizontal whitespace trailing the attributes list is trimmed if the next
// token starts a new line.
if let nextToken = node.nextToken(viewMode: .sourceAccurate),
nextToken.leadingTrivia.startsWithNewline
{
if !triviaToAttachToNextToken.isEmpty {
triviaToAttachToNextToken = triviaToAttachToNextToken.trimmingSuffix(while: \.isSpaceOrTab)
} else if let lastAttribute = filteredAttributes.last {
filteredAttributes[filteredAttributes.count - 1].trailingTrivia = lastAttribute
.trailingTrivia
.trimmingSuffix(while: \.isSpaceOrTab)
}
}
return AttributeListSyntax(filteredAttributes)
}

public override func visit(_ token: TokenSyntax) -> TokenSyntax {
return prependAndClearAccumulatedTrivia(to: token)
}

/// Prepends the accumulated trivia to the given node's leading trivia.
///
/// To preserve correct formatting after attribute removal, this function reassigns
/// significant trivia accumulated from removed attributes to the provided subsequent node.
/// Once attached, the accumulated trivia is cleared.
///
/// - Parameter node: The syntax node receiving the accumulated trivia.
/// - Returns: The modified syntax node with the prepended trivia.
private func prependAndClearAccumulatedTrivia<T: SyntaxProtocol>(to syntaxNode: T) -> T {
defer { triviaToAttachToNextToken = Trivia() }
return syntaxNode.with(\.leadingTrivia, triviaToAttachToNextToken + syntaxNode.leadingTrivia)
}
}

extension Trivia {
fileprivate func trimmingPrefix(
while predicate: (TriviaPiece) -> Bool
) -> Trivia {
Trivia(pieces: self.drop(while: predicate))
}

fileprivate func trimmingSuffix(
while predicate: (TriviaPiece) -> Bool
) -> Trivia {
Trivia(
pieces: self[...]
.reversed()
.drop(while: predicate)
.reversed()
)
}

fileprivate var startsWithNewline: Bool {
self.first?.isNewline ?? false
}
}

0 comments on commit 8dee94c

Please sign in to comment.