From 24cd3f184f6c5634772dca600d2ba509c5416013 Mon Sep 17 00:00:00 2001
From: Florian Verdonck <florian.verdonck@outlook.com>
Date: Sun, 19 Apr 2020 12:27:24 +0200
Subject: [PATCH] Fix line comment after then keyword. Fixes #730 (#768)

---
 src/Fantomas.Tests/IfThenElseTests.fs | 11 +++++++++++
 src/Fantomas/CodePrinter.fs           | 25 ++++++++++++++++++-------
 src/Fantomas/SourceParser.fs          |  4 ++--
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/Fantomas.Tests/IfThenElseTests.fs b/src/Fantomas.Tests/IfThenElseTests.fs
index 0aa38f8303..4a18570cbb 100644
--- a/src/Fantomas.Tests/IfThenElseTests.fs
+++ b/src/Fantomas.Tests/IfThenElseTests.fs
@@ -904,4 +904,15 @@ module String =
         else if String.length a' < String.length b' then a'
         else if String.length a' > String.length b' then b'
         else b'
+"""
+
+[<Test>]
+let ``comment after then in if/then, 730`` () =
+    formatSourceString false """if true then // comment
+    ()
+"""  config
+    |> prepend newline
+    |> should equal """
+if true then // comment
+    ()
 """
\ No newline at end of file
diff --git a/src/Fantomas/CodePrinter.fs b/src/Fantomas/CodePrinter.fs
index ae0f901e4e..fb66e71d0c 100644
--- a/src/Fantomas/CodePrinter.fs
+++ b/src/Fantomas/CodePrinter.fs
@@ -1181,26 +1181,37 @@ and genExpr astContext synExpr =
         // * https://github.com/fsprojects/fantomas/issues/513
         firstNewlineOrComment es +> atCurrentColumn (colEx (fun (e:SynExpr) -> sepConsideringTriviaContentBefore sepSemiNln e.Range) es (genExpr astContext))
 
-    | IfThenElse(e1, e2, None) ->
+    | IfThenElse(e1, e2, None, mIfToThen) ->
         fun (ctx:Context) ->
             let maxWidth = ctx.Config.MaxIfThenElseShortWidth
 
+            let thenKeywordHasLineComment =
+                TriviaHelpers.``has content after after that matches``
+                    (fun tn ->
+                        match tn.Type with
+                        | Token({ TokenInfo = ti }) when (ti.TokenName = "THEN") -> true
+                        | _ -> false)
+                    (function | Comment(LineCommentAfterSourceCode(_)) -> true | _ -> false)
+                    ctx.Trivia
+
+            let thenExpr = tokN mIfToThen "THEN"
+
             (leadingExpressionResult
-                (!- "if "+> genExpr astContext e1)
+                (!- "if " +> genExpr astContext e1)
                 (fun ((lb,cb),(la,ca)) ->
-                    let thenExpressionIsMultiline = futureNlnCheck (genExpr astContext e2) ctx
+                    let thenExpressionIsMultiline = thenKeywordHasLineComment || futureNlnCheck (genExpr astContext e2) ctx
 
                     if lb < la || thenExpressionIsMultiline then // if or then expression was multiline
-                        !- " then" +> indent +> sepNln +> genExpr astContext e2 +> unindent
+                        thenExpr (!- " then") +> indent +> sepNln +> genExpr astContext e2 +> unindent
                     elif (lb = la && (ca - cb) > maxWidth)
                          && not thenExpressionIsMultiline then // if expression is longer than maxWidth but not multiline
-                        sepNln +> !- "then " +> genExpr astContext e2
+                        sepNln +> thenExpr (!- "then ") +> genExpr astContext e2
                     elif (exceedsWidth maxWidth (genExpr astContext e2) ctx)
                          && not thenExpressionIsMultiline then // then is longer than maxWidth but not multiline
-                        sepNln +> !- "then " +> genExpr astContext e2
+                        sepNln +> thenExpr (!- "then ") +> genExpr astContext e2
                     else
                         // write out as short expression
-                        !- " then " +> genExpr astContext e2)
+                        thenExpr (!- " then ") +> genExpr astContext e2)
             |> atCurrentColumn) ctx
 
     // A generalization of IfThenElse
diff --git a/src/Fantomas/SourceParser.fs b/src/Fantomas/SourceParser.fs
index d25f9a585c..ff923f996e 100644
--- a/src/Fantomas/SourceParser.fs
+++ b/src/Fantomas/SourceParser.fs
@@ -834,8 +834,8 @@ let (|DotSet|_|) = function
     | _ -> None
 
 let (|IfThenElse|_|) = function
-    | SynExpr.IfThenElse(e1, e2, e3, _, _, _, _) ->
-        Some(e1, e2, e3)
+    | SynExpr.IfThenElse(e1, e2, e3, _, _, mIfToThen, _) ->
+        Some(e1, e2, e3, mIfToThen)
     | _ -> None
 
 let rec (|ElIf|_|) = function