|
| 1 | +/* |
| 2 | + This source file is part of the Swift.org open source project |
| 3 | + |
| 4 | + Copyright (c) 2021 Apple Inc. and the Swift project authors |
| 5 | + Licensed under Apache License v2.0 with Runtime Library Exception |
| 6 | + |
| 7 | + See http://swift.org/LICENSE.txt for license information |
| 8 | + See http://swift.org/CONTRIBUTORS.txt for Swift project authors |
| 9 | +*/ |
| 10 | + |
| 11 | +import SwiftSyntax |
| 12 | + |
| 13 | +extension ArrayExprSyntax { |
| 14 | + public func withAdditionalElementExpr(_ expr: ExprSyntax) -> ArrayExprSyntax { |
| 15 | + if self.elements.count >= 2 { |
| 16 | + // If the array expression has >=2 elements, use the trivia between |
| 17 | + // the last and second-to-last elements to determine how we insert |
| 18 | + // the new one. |
| 19 | + let lastElement = self.elements.last! |
| 20 | + let secondToLastElement = self.elements[self.elements.index(self.elements.endIndex, offsetBy: -2)] |
| 21 | + |
| 22 | + let newElements = self.elements |
| 23 | + .removingLast() |
| 24 | + .appending( |
| 25 | + lastElement.withTrailingComma( |
| 26 | + TokenSyntax.commaToken( |
| 27 | + trailingTrivia: (lastElement.trailingTrivia ?? []) + |
| 28 | + rightSquare.leadingTrivia.droppingPiecesAfterLastComment() + |
| 29 | + (secondToLastElement.trailingTrivia ?? []) |
| 30 | + ) |
| 31 | + ) |
| 32 | + ) |
| 33 | + .appending( |
| 34 | + ArrayElementSyntax( |
| 35 | + expression: expr, |
| 36 | + trailingComma: TokenSyntax.commaToken() |
| 37 | + ).withLeadingTrivia(lastElement.leadingTrivia?.droppingPiecesUpToAndIncludingLastComment() ?? []) |
| 38 | + ) |
| 39 | + |
| 40 | + return self.withElements(newElements) |
| 41 | + .withRightSquare( |
| 42 | + self.rightSquare.withLeadingTrivia( |
| 43 | + self.rightSquare.leadingTrivia.droppingPiecesUpToAndIncludingLastComment() |
| 44 | + ) |
| 45 | + ) |
| 46 | + } else { |
| 47 | + // For empty and single-element array exprs, we determine the indent |
| 48 | + // of the line the opening square bracket appears on, and then use |
| 49 | + // that to indent the added element and closing brace onto newlines. |
| 50 | + let (indentTrivia, unitIndent) = self.leftSquare.determineIndentOfStartingLine() |
| 51 | + var newElements: [ArrayElementSyntax] = [] |
| 52 | + if !self.elements.isEmpty { |
| 53 | + let existingElement = self.elements.first! |
| 54 | + newElements.append( |
| 55 | + ArrayElementSyntax(expression: existingElement.expression, |
| 56 | + trailingComma: TokenSyntax.commaToken()) |
| 57 | + .withLeadingTrivia(indentTrivia + unitIndent) |
| 58 | + .withTrailingTrivia((existingElement.trailingTrivia ?? []) + .newlines(1)) |
| 59 | + ) |
| 60 | + } |
| 61 | + |
| 62 | + newElements.append( |
| 63 | + ArrayElementSyntax(expression: expr, trailingComma: TokenSyntax.commaToken()) |
| 64 | + .withLeadingTrivia(indentTrivia + unitIndent) |
| 65 | + ) |
| 66 | + |
| 67 | + return self.withLeftSquare(self.leftSquare.withTrailingTrivia(.newlines(1))) |
| 68 | + .withElements(ArrayElementListSyntax(newElements)) |
| 69 | + .withRightSquare(self.rightSquare.withLeadingTrivia(.newlines(1) + indentTrivia)) |
| 70 | + } |
| 71 | + } |
| 72 | +} |
| 73 | + |
| 74 | +extension ArrayExprSyntax { |
| 75 | + func reindentingLastCallExprElement() -> ArrayExprSyntax { |
| 76 | + let lastElement = elements.last! |
| 77 | + let (indent, unitIndent) = lastElement.determineIndentOfStartingLine() |
| 78 | + let formattingVisitor = MultilineArgumentListRewriter(indent: indent, unitIndent: unitIndent) |
| 79 | + let formattedLastElement = formattingVisitor.visit(lastElement).as(ArrayElementSyntax.self)! |
| 80 | + return self.withElements(elements.replacing(childAt: elements.count - 1, with: formattedLastElement)) |
| 81 | + } |
| 82 | +} |
| 83 | + |
| 84 | +fileprivate extension TriviaPiece { |
| 85 | + var isComment: Bool { |
| 86 | + switch self { |
| 87 | + case .spaces, .tabs, .verticalTabs, .formfeeds, .newlines, |
| 88 | + .carriageReturns, .carriageReturnLineFeeds, .unexpectedText, .shebang: |
| 89 | + return false |
| 90 | + case .lineComment, .blockComment, .docLineComment, .docBlockComment: |
| 91 | + return true |
| 92 | + } |
| 93 | + } |
| 94 | + |
| 95 | + var isHorizontalWhitespace: Bool { |
| 96 | + switch self { |
| 97 | + case .spaces, .tabs: |
| 98 | + return true |
| 99 | + default: |
| 100 | + return false |
| 101 | + } |
| 102 | + } |
| 103 | + |
| 104 | + var isSpaces: Bool { |
| 105 | + guard case .spaces = self else { return false } |
| 106 | + return true |
| 107 | + } |
| 108 | + |
| 109 | + var isTabs: Bool { |
| 110 | + guard case .tabs = self else { return false } |
| 111 | + return true |
| 112 | + } |
| 113 | +} |
| 114 | + |
| 115 | +fileprivate extension Trivia { |
| 116 | + func droppingPiecesAfterLastComment() -> Trivia { |
| 117 | + Trivia(pieces: .init(self.lazy.reversed().drop(while: { !$0.isComment }).reversed())) |
| 118 | + } |
| 119 | + |
| 120 | + func droppingPiecesUpToAndIncludingLastComment() -> Trivia { |
| 121 | + Trivia(pieces: .init(self.lazy.reversed().prefix(while: { !$0.isComment }).reversed())) |
| 122 | + } |
| 123 | +} |
| 124 | + |
| 125 | +extension SyntaxProtocol { |
| 126 | + func determineIndentOfStartingLine() -> (indent: Trivia, unitIndent: Trivia) { |
| 127 | + let sourceLocationConverter = SourceLocationConverter(file: "", tree: self.root.as(SourceFileSyntax.self)!) |
| 128 | + let line = startLocation(converter: sourceLocationConverter).line ?? 0 |
| 129 | + let visitor = DetermineLineIndentVisitor(lineNumber: line, sourceLocationConverter: sourceLocationConverter) |
| 130 | + visitor.walk(self.root) |
| 131 | + return (indent: visitor.lineIndent, unitIndent: visitor.lineUnitIndent) |
| 132 | + } |
| 133 | +} |
| 134 | + |
| 135 | +public final class DetermineLineIndentVisitor: SyntaxVisitor { |
| 136 | + |
| 137 | + let lineNumber: Int |
| 138 | + let locationConverter: SourceLocationConverter |
| 139 | + private var bestMatch: TokenSyntax? |
| 140 | + |
| 141 | + public var lineIndent: Trivia { |
| 142 | + guard let pieces = bestMatch?.leadingTrivia |
| 143 | + .lazy |
| 144 | + .reversed() |
| 145 | + .prefix(while: \.isHorizontalWhitespace) |
| 146 | + .reversed() else { return .spaces(4) } |
| 147 | + return Trivia(pieces: Array(pieces)) |
| 148 | + } |
| 149 | + |
| 150 | + public var lineUnitIndent: Trivia { |
| 151 | + if lineIndent.allSatisfy(\.isSpaces) { |
| 152 | + let addedSpaces = lineIndent.reduce(0, { |
| 153 | + guard case .spaces(let count) = $1 else { fatalError() } |
| 154 | + return $0 + count |
| 155 | + }) % 4 == 0 ? 4 : 2 |
| 156 | + return .spaces(addedSpaces) |
| 157 | + } else if lineIndent.allSatisfy(\.isTabs) { |
| 158 | + return .tabs(1) |
| 159 | + } else { |
| 160 | + // If we can't determine the indent, default to 4 spaces. |
| 161 | + return .spaces(4) |
| 162 | + } |
| 163 | + } |
| 164 | + |
| 165 | + public init(lineNumber: Int, sourceLocationConverter: SourceLocationConverter) { |
| 166 | + self.lineNumber = lineNumber |
| 167 | + self.locationConverter = sourceLocationConverter |
| 168 | + super.init(viewMode: .sourceAccurate) |
| 169 | + } |
| 170 | + |
| 171 | + public override func visit(_ tokenSyntax: TokenSyntax) -> SyntaxVisitorContinueKind { |
| 172 | + let range = tokenSyntax.sourceRange(converter: locationConverter, |
| 173 | + afterLeadingTrivia: false, |
| 174 | + afterTrailingTrivia: true) |
| 175 | + guard let startLine = range.start.line, |
| 176 | + let endLine = range.end.line, |
| 177 | + let startColumn = range.start.column, |
| 178 | + let endColumn = range.end.column else { |
| 179 | + return .skipChildren |
| 180 | + } |
| 181 | + |
| 182 | + if (startLine, startColumn) <= (lineNumber, 1), |
| 183 | + (lineNumber, 1) <= (endLine, endColumn) { |
| 184 | + bestMatch = tokenSyntax |
| 185 | + return .visitChildren |
| 186 | + } else { |
| 187 | + return .skipChildren |
| 188 | + } |
| 189 | + } |
| 190 | +} |
| 191 | + |
| 192 | +/// Moves each argument to a function call expression onto a new line and indents them appropriately. |
| 193 | +final class MultilineArgumentListRewriter: SyntaxRewriter { |
| 194 | + let indent: Trivia |
| 195 | + let unitIndent: Trivia |
| 196 | + |
| 197 | + init(indent: Trivia, unitIndent: Trivia) { |
| 198 | + self.indent = indent |
| 199 | + self.unitIndent = unitIndent |
| 200 | + } |
| 201 | + |
| 202 | + override func visit(_ token: TokenSyntax) -> Syntax { |
| 203 | + guard token.tokenKind == .rightParen else { return Syntax(token) } |
| 204 | + return Syntax(token.withLeadingTrivia(.newlines(1) + indent)) |
| 205 | + } |
| 206 | + |
| 207 | + override func visit(_ node: TupleExprElementSyntax) -> Syntax { |
| 208 | + return Syntax(node.withLeadingTrivia(.newlines(1) + indent + unitIndent)) |
| 209 | + } |
| 210 | +} |
0 commit comments