Initial implementation of the CommentWhitespace rule and tests (swiftlang#52)

Author: atamez31

tools/swift-format/Sources/Rules/CommentWhitespace.swift

@@ -11,5 +11,96 @@ import SwiftSyntax
 ///
 /// - SeeAlso: https://google.github.io/swift#horizontal-whitespace
 public final class CommentWhitespace: SyntaxFormatRule {
+  public override func visit (_ token: TokenSyntax) -> Syntax {
+    var pieces = [TriviaPiece]()
+    var validToken = token
+    var needsWhitespaceFix = false
 
+    guard let nextToken = token.nextToken else {
+      // In the case there is a line comment at the end of the file, it ensures
+      // that the line comment has a single space after the `//`.
+      pieces = checksSpacesAfterLineComment(isInvalid: &needsWhitespaceFix, token: token)
+      return needsWhitespaceFix ? token.withLeadingTrivia(Trivia.init(pieces: pieces)) : token
+    }
+
+    // Ensures the line comment has at least 2 spaces before the `//`.
+    if hasInlineLineComment(trivia: nextToken.leadingTrivia) {
+      let numSpaces = token.trailingTrivia.numberOfSpaces
+      if numSpaces < 2 {
+        needsWhitespaceFix = true
+        let addSpaces = 2 - numSpaces
+        diagnose(.addSpacesBeforeLineComment(count: addSpaces), on:token)
+        validToken = token.withTrailingTrivia(token.trailingTrivia.appending(.spaces(addSpaces)))
+      }
+    }
+
+    pieces = checksSpacesAfterLineComment(isInvalid: &needsWhitespaceFix, token: token)
+    return needsWhitespaceFix ? validToken.withLeadingTrivia(Trivia.init(pieces: pieces)) : token
+  }
+
+  /// Returns a boolean indicating if the given trivia contains
+  /// a line comment inline with code.
+  private func hasInlineLineComment (trivia: Trivia) -> Bool {
+    // Comments are inline unless the trivia begins with a
+    // with a newline.
+    if let firstPiece = trivia.reversed().last {
+      if case .newlines(_) = firstPiece {
+        return false
+      }
+    }
+    for piece in trivia {
+      if case .lineComment(_) = piece {
+        return true
+      }
+    }
+    return false
+  }
+
+  /// Ensures the line comment has exactly one space after the `//`.
+  private func checksSpacesAfterLineComment(isInvalid: inout Bool, token: TokenSyntax) -> [TriviaPiece] {
+    var pieces = [TriviaPiece]()
+
+    for piece in token.leadingTrivia {
+      // Checks if the line comment has exactly one space after the `//`,
+      // if it doesn't it removes or add an space, depending on what the
+      // comment needs in order to follow the right format.
+      if case .lineComment(let text) = piece,
+         let formatText = formatLineComment(textLineComment: text, token: token) {
+        isInvalid = true
+        pieces.append(TriviaPiece.lineComment(formatText))
+      }
+      else {
+        pieces.append(piece)
+      }
+    }
+    return pieces
+  }
+
+  /// Given a string with the text of a line comment, it ensures there
+  /// is exactly one space after the `//`. If the string doesn't follow
+  /// this rule a new string is returned with the right format.
+  private func formatLineComment (textLineComment: String, token: TokenSyntax) -> String? {
+    let text = textLineComment.dropFirst(2)
+    if text.first != " " {
+      diagnose(.addSpaceAfterLineComment, on: token)
+      return "// " + text.trimmingCharacters(in: .whitespaces)
+    }
+    else if text.dropFirst(1).first == " " {
+      diagnose(.removeSpacesAfterLineComment, on: token)
+      return "// " + text.trimmingCharacters(in: .whitespaces)
+    }
+    return nil
+  }
+}
+
+extension Diagnostic.Message {
+  static func addSpacesBeforeLineComment(count: Int) -> Diagnostic.Message {
+    let ending = count == 1 ? "" : "s"
+    return Diagnostic.Message(.warning, "add \(count) space\(ending) before the //")
+  }
+
+  static let addSpaceAfterLineComment =
+    Diagnostic.Message(.warning, "add one space after `//`")
+  static let removeSpacesAfterLineComment =
+    Diagnostic.Message(.warning, "remove excess of spaces after the `//`")
 }

tools/swift-format/Tests/SwiftFormatTests/CommentWhitespaceTest.swift

@@ -0,0 +1,63 @@
+import SwiftSyntax
+import XCTest
+
+@testable import Rules
+
+public class CommentWhitespaceTest: DiagnosingTestCase {
+  public func testInvalidCommentWhiteSpace() {
+    XCTAssertFormatting(
+    CommentWhitespace.self,
+    input: """
+           let initialFactor = 2 //        Line comment of initialFactor.
+           let finalFactor = 2//Line comment of finalFactor.
+
+           //Lorem ipsum dolor sit amet, at nonumes adipisci sea, natum
+           //       offendit vis ex. Audiam legendos expetenda ei quo, nonumes
+           // msensibus eloquentiam ex vix.
+           let fin = 3//        End of file.
+           """,
+    expected: """
+              let initialFactor = 2  // Line comment of initialFactor.
+              let finalFactor = 2  // Line comment of finalFactor.
+
+              // Lorem ipsum dolor sit amet, at nonumes adipisci sea, natum
+              // offendit vis ex. Audiam legendos expetenda ei quo, nonumes
+              // msensibus eloquentiam ex vix.
+              let fin = 3  // End of file.
+              """)
+  }
+  
+  public func testInvalidFuncCommentWhiteSpace() {
+    XCTAssertFormatting(
+    CommentWhitespace.self,
+    input: """
+           func testLineComment(paramA: Int) -> Int {
+             //LineComment.
+             if paramA < 50 {
+               //LineComment.
+               return paramA - 100
+             }
+             //LineComment.
+             return paramA + 100
+           }
+           """,
+    expected: """
+              func testLineComment(paramA: Int) -> Int {
+                // LineComment.
+                if paramA < 50 {
+                  // LineComment.
+                  return paramA - 100
+                }
+                // LineComment.
+                return paramA + 100
+              }
+              """)
+  }
+  
+  #if !os(macOS)
+  static let allTests = [
+    CommentWhitespaceTest.testInvalidCommentWhiteSpace,
+    CommentWhitespaceTest.testInvalidFuncCommentWhiteSpace,
+    ]
+  #endif
+}