Merge pull request swiftlang#141 from dabelknap/comment-fixes

Fixes to comment behavior

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -398,8 +398,10 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: CodeBlockSyntax) {
-    for i in 0..<(node.statements.count - 1) {
-      after(node.statements[i].lastToken, tokens: .newline)
+    if node.statements.count > 0 {
+      for i in 0..<(node.statements.count - 1) {
+        after(node.statements[i].lastToken, tokens: .newline)
+      }
     }
     super.visit(node)
   }
@@ -1233,123 +1235,145 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ token: TokenSyntax) {
-    let fileStart = token.previousToken == nil
-    extractTrivia(token.leadingTrivia, fileStart: fileStart)
+    extractLeadingTrivia(token)
     if let before = beforeMap[token] {
       tokens += before
     }
     appendToken(.syntax(token))
-    extractTrailingLineComment(token)
+    extractTrailingComment(token)
     if let afterGroups = afterMap[token] {
       for after in afterGroups.reversed() {
         tokens += after
       }
     }
-    extractTrivia(token.trailingTrivia)
   }
 
-  func appendToken(_ token: Token) {
-    if let last = tokens.last {
-      switch (last, token) {
-      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)
-        return
-      default:
-        break
+  private func extractTrailingComment(_ token: TokenSyntax) {
+    let nextToken = token.nextToken
+    guard let trivia = nextToken?.leadingTrivia,
+          let firstPiece = trivia[safe: 0] else {
+      return
+    }
+
+    switch firstPiece {
+    case .lineComment(let text):
+      appendToken(.break(size: 2, offset: 2))
+      appendToken(.comment(Comment(kind: .line, text: text)))
+      if isInContainer(token) {
+        appendToken(.newline)
       }
+
+    case .blockComment(let text):
+      appendToken(.break(size: 2, offset: 2))
+      appendToken(.comment(Comment(kind: .block, text: text)))
+      if isInContainer(token) {
+        appendToken(.newline)
+      }
+
+    default:
+      return
     }
-    tokens.append(token)
   }
 
-  private func shouldAddNewlineBefore(_ token: TokenSyntax?) -> Bool {
-    guard let token = token, let before = beforeMap[token] else { return false }
-    for item in before {
-      if case .newlines = item { return false }
+  private func isInContainer(_ token: TokenSyntax) -> Bool {
+    if token.parent is ArrayElementSyntax || token.parent is DictionaryElementSyntax || token.parent is TupleElementSyntax {
+      return true
     }
-    return true
+    return false
   }
 
-  private func extractTrailingLineComment(_ token: TokenSyntax) {
-    guard let nextToken = token.nextToken else {
-      return
-    }
-    let trivia = nextToken.leadingTrivia
-    for (offset, piece) in trivia.enumerated() {
+  private func extractLeadingTrivia(_ token: TokenSyntax) {
+    let isStartOfFile = token.previousToken == nil
+    let trivia = token.leadingTrivia
+
+    for (index, piece) in trivia.enumerated() {
       switch piece {
       case .lineComment(let text):
-        if offset > 0, case .newlines? = trivia[safe: offset - 1] {
-          return
-        } else {
-          appendToken(.break(size: 2, offset: 2))
+        if index > 0 || isStartOfFile {
+          if token.withoutTrivia().text == "}" {
+            if let previousToken = token.previousToken,
+               previousToken.withoutTrivia().text == "{" {
+              // do nothing
+            } else {
+              appendToken(.newline)
+            }
+          }
+
           appendToken(.comment(Comment(kind: .line, text: text)))
-          if let parent = nextToken.parent?.parent, !(parent is CodeBlockItemSyntax) {
+
+          if token.withoutTrivia().text == "}" {
+            appendToken(.newline(offset: -2))
+          } else {
             appendToken(.newline)
           }
-          return
         }
+
       case .blockComment(let text):
-        if offset > 0, case .newlines? = trivia[safe: offset - 1] {
-          return
-        } else {
-          appendToken(.break(size: 2, offset: 2))
+        if index > 0 || isStartOfFile {
+          if token.withoutTrivia().text == "}" {
+            if let previousToken = token.previousToken,
+              previousToken.withoutTrivia().text == "{" {
+              // do nothing
+            } else {
+              appendToken(.newline)
+            }
+          }
+
           appendToken(.comment(Comment(kind: .block, text: text)))
-          if let parent = nextToken.parent?.parent, !(parent is CodeBlockItemSyntax) {
+
+          if token.withoutTrivia().text == "}" {
+            appendToken(.newline(offset: -2))
+          } else {
             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 .blockComment(let text):
-        if fileStart {
-          appendToken(.comment(Comment(kind: .block, text: text)))
-          appendToken(.newline)
-        } else if offset > 0, case .newlines? = trivia[safe: offset - 1] {
-          appendToken(.comment(Comment(kind: .block, text: text)))
-          appendToken(.newline)
-        }
       case .docLineComment(let text):
         appendToken(.comment(Comment(kind: .docLine, text: text)))
-        if case .newlines? = trivia[safe: offset + 1],
-           case .docLineComment? = trivia[safe: offset + 2] {
-          /* do nothing */
+        if case .newlines? = trivia[safe: index + 1],
+           case .docLineComment? = trivia[safe: index + 2] {
+          // do nothing
         } else {
           appendToken(.newline)
         }
+
       case .docBlockComment(let text):
         appendToken(.comment(Comment(kind: .docBlock, text: text)))
-        if case .newlines? = trivia[safe: offset + 1],
-          case .docLineComment? = trivia[safe: offset + 2] {
-          /* do nothing */
-        } else {
-          appendToken(.newline)
-        }
+        appendToken(.newline)
+
       case .newlines(let n), .carriageReturns(let n), .carriageReturnLineFeeds(let n):
         if n > 1 {
           appendToken(.newlines(min(n - 1, config.maximumBlankLines), offset: 0))
         }
+
+      default:
+        break
+      }
+    }
+  }
+
+  func appendToken(_ token: Token) {
+    if let last = tokens.last {
+      switch (last, token) {
+      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)
+        return
       default:
         break
       }
     }
+    tokens.append(token)
+  }
+
+  private func shouldAddNewlineBefore(_ token: TokenSyntax?) -> Bool {
+    guard let token = token, let before = beforeMap[token] else { return false }
+    for item in before {
+      if case .newlines = item { return false }
+    }
+    return true
   }
 }
 

Tests/SwiftFormatPrettyPrintTests/CommentTests.swift

@@ -92,6 +92,23 @@ public class CommentTests: PrettyPrintTestCase {
 
       let reallyLongVariableName = 123 // This comment should wrap
 
+      func MyFun() {
+        // just a comment
+      }
+
+      func MyFun() {
+        // Comment 1
+        // Comment 2
+        let a = 123
+
+        let b = 456  // Comment 3
+      }
+
+      func MyFun() {
+        let c = 789 // Comment 4
+        // Comment 5
+      }
+
       let d = 123
       // Trailing Comment
       """
@@ -110,6 +127,23 @@ public class CommentTests: PrettyPrintTestCase {
       let reallyLongVariableName = 123
         // This comment should wrap
 
+      func MyFun() {
+        // just a comment
+      }
+
+      func MyFun() {
+        // Comment 1
+        // Comment 2
+        let a = 123
+
+        let b = 456  // Comment 3
+      }
+
+      func MyFun() {
+        let c = 789  // Comment 4
+        // Comment 5
+      }
+
       let d = 123
       // Trailing Comment