Don't combine LineCommentAfterSourceCode and LineCommentOnSingleLine …

…as one trivia. (#1612) * Don't combine LineCommentAfterSourceCode and LineCommentOnSingleLine as one trivia. Fixes #1604. * Don't pass nextTokens.
fsprojects · Apr 6, 2021 · 376c409 · 376c409
1 parent 31922c0
commit 376c409
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 19 deletions.
diff --git a/src/Fantomas.Tests/LetBindingTests.fs b/src/Fantomas.Tests/LetBindingTests.fs
@@ -1624,3 +1624,20 @@ let myFun (a: decimal) b c : decimal = a + b + c
 let expensiveToCompute : int = 0
 let myFun (a: decimal) b c : decimal = a + b + c
 """
+
+[<Test>]
+let ``comments after short value binding, 1604`` () =
+    formatSourceString
+        false
+        """
+let foo = bar // bar
+//// hi!
+"""
+        config
+    |> prepend newline
+    |> should
+        equal
+        """
+let foo = bar // bar
+//// hi!
+"""
diff --git a/src/Fantomas/TokenParser.fs b/src/Fantomas/TokenParser.fs
@@ -718,30 +718,80 @@ let rec private getTriviaFromTokensThemSelves
         let prevToken = List.tryItem (headIndex - 1) allTokens
         let prevButOneToken = List.tryItem (headIndex - 2) allTokens
 
-        let commentType =
+        let isAfterSourceCode =
             match prevButOneToken, prevToken with
             | Some ({ LineNumber = ncln }), Some (WhiteSpaceToken _) ->
-                if ncln = headLineNumber then
-                    LineCommentAfterSourceCode
-                else
-                    LineCommentOnSingleLine
-            | Some ({ LineNumber = l1 }), Some ({ LineNumber = l2 }) when (headLineNumber = l1 && headLineNumber = l2) ->
-                LineCommentAfterSourceCode
-            | _ -> LineCommentOnSingleLine
+                // IDENT WHITESPACE LINE_COMMENT
+                ncln = headLineNumber
+            | Some ({ LineNumber = l1 }), Some ({ LineNumber = l2 }) ->
+                // IDENT LINE_COMMENT
+                headLineNumber = l1 && headLineNumber = l2
+            | _ -> false
 
-        let comment =
-            collectComment commentTokens
-            |> commentType
-            |> Comment
+        let info =
+            if isAfterSourceCode then
+                // Only capture the first line of the comment as LineCommentAfterSourceCode
+                // The next line(s) will be a LineCommentOnSingleLine
+                let commentsByLine =
+                    commentTokens
+                    |> List.groupBy (fun t -> t.LineNumber)
+
+                let firstComment =
+                    List.tryHead commentsByLine |> Option.map snd
+
+                match firstComment with
+                | Some (headToken :: _ as afterSourceTokens) ->
+                    let afterSourceCodeTrivia =
+                        let tc =
+                            collectComment afterSourceTokens
+                            |> LineCommentAfterSourceCode
+                            |> Comment
+
+                        let lastToken = List.tryLast afterSourceTokens
+
+                        let r =
+                            getRangeBetween mkRange headToken (Option.defaultValue headToken lastToken)
+
+                        Trivia.Create tc r
+
+                    let lineCommentOnSingleLine =
+                        if commentTokens.Length > afterSourceTokens.Length then
+                            let commentTokens =
+                                commentTokens
+                                |> List.skip afterSourceTokens.Length
+
+                            let triviaContent =
+                                collectComment commentTokens
+                                |> LineCommentOnSingleLine
+                                |> Comment
+
+                            let range =
+                                let headToken = List.head commentTokens
+                                let lastToken = List.tryLast commentTokens
+                                getRangeBetween mkRange headToken (Option.defaultValue headToken lastToken)
+
+                            Trivia.Create triviaContent range |> Some
+                        else
+                            None
+
+                    match lineCommentOnSingleLine with
+                    | Some lcsl -> afterSourceCodeTrivia :: lcsl :: foundTrivia
+                    | None -> afterSourceCodeTrivia :: foundTrivia
+                | _ ->
+                    // We should not hit this branch
+                    foundTrivia
+            else
+                let triviaContent =
+                    collectComment commentTokens
+                    |> LineCommentOnSingleLine
+                    |> Comment
 
-        let range =
-            let headToken = List.head commentTokens
-            let lastToken = List.tryLast commentTokens
-            getRangeBetween mkRange headToken (Option.defaultValue headToken lastToken)
+                let range =
+                    let headToken = List.head commentTokens
+                    let lastToken = List.tryLast commentTokens
+                    getRangeBetween mkRange headToken (Option.defaultValue headToken lastToken)
 
-        let info =
-            Trivia.Create comment range
-            |> List.appendItem foundTrivia
+                (Trivia.Create triviaContent range :: foundTrivia)
 
         getTriviaFromTokensThemSelves mkRange allTokens nextTokens info