Add support for basic line and documentation comments.

Author: dabelknap

Sources/SwiftFormatPrettyPrint/Comment.swift

@@ -38,13 +38,15 @@ struct Comment {
     }
   }
   let kind: Kind
-  var text: String
+  var text: [String]
+  public var length: Int
 
   init(kind: Kind, text: String) {
-    self.text = text
+    self.text = [text]
     self.kind = kind
+    self.length = text.count + kind.prefixLength + 1
 
-    self.text.removeFirst(kind.prefixLength)
+    self.text[0].removeFirst(kind.prefixLength)
 
     switch kind {
     case .docBlock:
@@ -55,51 +57,15 @@ struct Comment {
     }
   }
 
-  mutating func addText(_ text: String) {
-    self.text += "\n" + text
+  func print(indent: Int) -> String {
+    let separator = "\n" + String(repeating: " ", count: indent) + kind.prefix
+    return kind.prefix + self.text.joined(separator: separator)
   }
 
-  func wordWrap(lineLength: Int) -> [String] {
-    let maxLength = lineLength - (kind.prefixLength + 1)
-    let scanner = Scanner(string: text)
-    var lines = [String]()
-
-    // FIXME: Word wrapping doesn't work for documentation comments, as we need to preserve all the
-    //        intricacies of Markdown formatting.
-    // FIXME: If there's a totally blank comment line, it doesn't get a prefix for some reason.
-    // TODO: Allow for leading `*` characters for each line of a block comment.
-    if kind == .docLine || kind == .docBlock {
-      lines = text.split(separator: "\n").map { "\(kind.prefix)\($0)" }
-    } else {
-      var currentLine = ""
-      var currentLineLength = 0
-      var buffer: NSString! = ""
-      while scanner.scanUpToCharacters(from: .whitespacesAndNewlines, into: &buffer) {
-        let strBuf = buffer as String
-        if currentLineLength + strBuf.count > maxLength {
-          lines.append(currentLine)
-          currentLine = ""
-          currentLineLength = 0
-        }
-        currentLine += strBuf + " "
-        currentLineLength += strBuf.count + 1
-      }
-      if currentLineLength > 0 {
-        lines.append(currentLine.trimmingCharacters(in: .whitespaces))
-      }
-      for i in 0..<lines.count {
-        lines[i] = "\(kind.prefix) \(lines[i])"
-      }
-    }
-    switch kind {
-    case .block:
-      lines.insert("/*", at: 0)
-      lines.append(" */")
-    case .docBlock:
-      lines.insert("/**", at: 0)
-      lines.append(" */")
-    default: break
+  mutating func addText(_ text: [String]) {
+    for line in text {
+      self.text.append(line)
+      self.length += line.count + self.kind.prefixLength + 1
     }
-    return lines
   }
 }

Sources/SwiftFormatPrettyPrint/PrettyPrint.swift

@@ -225,8 +225,18 @@ public class PrettyPrinter {
       write(syntaxToken.text)
       spaceRemaining -= syntaxToken.text.count
 
-    // TODO(dabelknap): Implement comments
-    case .comment: ()
+    case .comment(let comment):
+      if lastBreakConsecutive {
+        // If the last token created a new line, we need to apply indentation.
+        writeSpaces(lastBreakValue)
+
+        lastBreak = false
+        lastBreakConsecutive = false
+        lastBreakOffset = 0
+        lastBreakValue = 0
+      }
+      write(comment.print(indent: lastBreakValue))
+      spaceRemaining -= comment.length
     }
   }
 
@@ -317,8 +327,9 @@ public class PrettyPrinter {
         lengths.append(syntaxToken.text.count)
         total += syntaxToken.text.count
 
-      // TODO(dabelknap): Implement comments
-      case .comment: ()
+      case .comment(let comment):
+        lengths.append(comment.length)
+        total += comment.length
       }
     }
 

Sources/SwiftFormatPrettyPrint/Token.swift

@@ -33,7 +33,7 @@ enum Token {
   case `break`(size: Int, offset: Int)
   case space(size: Int)
   case newlines(Int, offset: Int)
-  case comment(Comment, hasTrailingSpace: Bool)
+  case comment(Comment)
   case reset
 
   // Convenience overloads for the enum types

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -1205,27 +1205,29 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ token: TokenSyntax) {
-    breakDownTrivia(token.leadingTrivia, before: token)
+    let fileStart = token.previousToken == nil
+    extractTrivia(token.leadingTrivia, fileStart: fileStart)
     if let before = beforeMap[token] {
       tokens += before
     }
     appendToken(.syntax(token))
+    extractTrailingLineComment(token)
     if let afterGroups = afterMap[token] {
       for after in afterGroups.reversed() {
         tokens += after
       }
     }
-    breakDownTrivia(token.trailingTrivia)
+    extractTrivia(token.trailingTrivia)
   }
 
   func appendToken(_ token: Token) {
     if let last = tokens.last {
       switch (last, token) {
-      case (.comment(let c1, _), .comment(let c2, _))
+      case (.comment(let c1), .comment(let c2))
         where c1.kind == .docLine && c2.kind == .docLine:
         var newComment = c1
         newComment.addText(c2.text)
-        tokens[tokens.count - 1] = .comment(newComment, hasTrailingSpace: false)
+        tokens[tokens.count - 1] = .comment(newComment)
         return
       default:
         break
@@ -1242,52 +1244,49 @@ private final class TokenStreamCreator: SyntaxVisitor {
     return true
   }
 
-  private func breakDownTrivia(_ trivia: Trivia, before: TokenSyntax? = nil) {
+  private func extractTrailingLineComment(_ token: TokenSyntax) {
+    guard let nextToken = token.nextToken else {
+      return
+    }
+    let trivia = nextToken.leadingTrivia
     for (offset, piece) in trivia.enumerated() {
       switch piece {
       case .lineComment(let text):
-        appendToken(.comment(Comment(kind: .line, text: text), hasTrailingSpace: false))
-        if case .newlines? = trivia[safe: offset + 1],
-           case .lineComment? = trivia[safe: offset + 2] {
-          /* do nothing */
+        if offset > 0, case .newlines? = trivia[safe: offset - 1] {
+          return
         } else {
+          appendToken(.break(size: 2, offset: 2))
+          appendToken(.comment(Comment(kind: .line, text: text)))
+          if let parent = nextToken.parent?.parent, !(parent is CodeBlockItemSyntax) {
+            appendToken(.newline)
+          }
+          return
+        }
+      default:
+        break
+      }
+    }
+  }
+
+  private func extractTrivia(_ trivia: Trivia, fileStart: Bool = false) {
+    for (offset, piece) in trivia.enumerated() {
+      switch piece {
+      case .lineComment(let text):
+        if fileStart {
+          appendToken(.comment(Comment(kind: .line, text: text)))
+          appendToken(.newline)
+        } else if offset > 0, case .newlines? = trivia[safe: offset - 1] {
+          appendToken(.comment(Comment(kind: .line, text: text)))
           appendToken(.newline)
         }
       case .docLineComment(let text):
-        appendToken(.comment(Comment(kind: .docLine, text: text), hasTrailingSpace: false))
+        appendToken(.comment(Comment(kind: .docLine, text: text)))
         if case .newlines? = trivia[safe: offset + 1],
            case .docLineComment? = trivia[safe: offset + 2] {
           /* do nothing */
         } else {
           appendToken(.newline)
         }
-      case .blockComment(let text), .docBlockComment(let text):
-        var hasTrailingSpace = false
-        var hasTrailingNewline = false
-
-        // Detect if a newline or trailing space comes after this comment and preserve it.
-        if let next = trivia[safe: offset + 1] {
-          switch next {
-          case .newlines, .carriageReturns, .carriageReturnLineFeeds:
-            hasTrailingNewline = true
-          case .spaces, .tabs:
-            hasTrailingSpace = true
-          default:
-            break
-          }
-        }
-
-        let commentKind: Comment.Kind
-        if case .blockComment = piece {
-          commentKind = .block
-        } else {
-          commentKind = .docBlock
-        }
-        let comment = Comment(kind: commentKind, text: text)
-        appendToken(.comment(comment, hasTrailingSpace: hasTrailingSpace))
-        if hasTrailingNewline {
-          appendToken(.newline)
-        }
       case .newlines(let n), .carriageReturns(let n), .carriageReturnLineFeeds(let n):
         if n > 1 {
           appendToken(.newlines(min(n - 1, config.maximumBlankLines), offset: 0))

Tests/SwiftFormatPrettyPrintTests/CommentTests.swift

@@ -0,0 +1,151 @@
+public class CommentTests: PrettyPrintTestCase {
+  public func testDocumentationComments() {
+    let input =
+      """
+      /// This is a documentation comment
+      ///
+      /// - Parameters:
+      ///   - param1: Param1 comment
+      ///   - param2: Param2 comment
+      /// - Returns: The output
+      func myfun(param1: Int, param2: Double) -> String {
+        let out = "123"
+        return out
+      }
+
+      /// A brief doc comment
+      func myfun() {
+        let a = 123
+        let b = "456"
+      }
+
+      public class MyClass {
+        /// Doc comment
+        var myVariable: Int
+
+        /// Method doc comment
+        ///
+        /// - Parameters:
+        ///   - param1: Param1 comment
+        ///   - param2: Param2 comment
+        /// - Returns: The output
+        func myFun(param1: Int, param2: Int) -> String {
+          let a = 123
+          let b = "456"
+          return b
+        }
+      }
+      """
+
+    let expected =
+      """
+      /// This is a documentation comment
+      ///
+      /// - Parameters:
+      ///   - param1: Param1 comment
+      ///   - param2: Param2 comment
+      /// - Returns: The output
+      func myfun(param1: Int, param2: Double) -> String {
+        let out = "123"
+        return out
+      }
+
+      /// A brief doc comment
+      func myfun() {
+        let a = 123
+        let b = "456"
+      }
+
+      public class MyClass {
+        /// Doc comment
+        var myVariable: Int
+
+        /// Method doc comment
+        ///
+        /// - Parameters:
+        ///   - param1: Param1 comment
+        ///   - param2: Param2 comment
+        /// - Returns: The output
+        func myFun(param1: Int, param2: Int) -> String {
+          let a = 123
+          let b = "456"
+          return b
+        }
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 55)
+  }
+
+  public func testLineComments() {
+    let input =
+      """
+      // Line Comment1
+      // Line Comment2
+      let a = 123
+      let b = "456"  // End of line comment
+      let c = "More content"
+
+      // Comment 3
+      // Comment 4
+
+      let reallyLongVariableName = 123 // This comment should wrap
+
+      let d = 123
+      // Trailing Comment
+      """
+
+    let expected =
+      """
+      // Line Comment1
+      // Line Comment2
+      let a = 123
+      let b = "456"  // End of line comment
+      let c = "More content"
+
+      // Comment 3
+      // Comment 4
+
+      let reallyLongVariableName = 123
+        // This comment should wrap
+
+      let d = 123
+      // Trailing Comment
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 45)
+  }
+
+  public func testContainerLineComments() {
+    let input =
+      """
+      // Array comment
+      let a = [456, // small comment
+        789]
+
+      // Dictionary comment
+      let b = ["abc": 456, // small comment
+        "def": 789]
+      """
+
+    let expected =
+      """
+      // Array comment
+      let a = [
+        456,  // small comment
+        789
+      ]
+
+      // Dictionary comment
+      let b = [
+        "abc": 456,  // small comment
+        "def": 789
+      ]
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 80)
+  }
+}