Improve comment placement in if-then-else blocks

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 * Don't error when the `JavaScriptFFI` language pragma is present. [Issue
   1087](https://github.com/tweag/ormolu/issues/1087).
+* Improve comment placement in if-then-else blocks. [Issue
+  998](https://github.com/tweag/ormolu/issues/998).
 
 ## Ormolu 0.7.3.0
 

data/examples/declaration/value/function/if-single-line-functions-do-out.hs

@@ -0,0 +1,5 @@
+foo x = do
+  y <- quux
+  if x > 2
+    then bar x
+    else baz y

data/examples/declaration/value/function/if-single-line-functions-do.hs

@@ -0,0 +1,5 @@
+foo x = do
+  y <- quux
+  if x > 2
+    then bar x
+    else baz y

data/examples/declaration/value/function/if-single-line-functions-out.hs

@@ -0,0 +1,4 @@
+foo x =
+  if x > 2
+    then bar x
+    else baz x

data/examples/declaration/value/function/if-single-line-functions.hs

@@ -0,0 +1,4 @@
+foo x =
+  if x > 2
+    then bar x
+    else baz x

data/examples/declaration/value/function/if-with-comment-above-branches-out.hs

@@ -0,0 +1,7 @@
+foo =
+  if undefined
+    -- then comment
+    then undefined
+    -- else comment
+    else do
+      undefined

data/examples/declaration/value/function/if-with-comment-before-do-blocks-out.hs

@@ -0,0 +1,10 @@
+foo =
+  if something
+    -- then comment
+    then do
+      stuff
+      stuff
+    -- else comment
+    else do
+      stuff
+      stuff

data/examples/declaration/value/function/if-with-comment-before-do-blocks.hs

@@ -0,0 +1,10 @@
+foo =
+  if something
+    -- then comment
+    then do
+      stuff
+      stuff
+    -- else comment
+    else do
+      stuff
+      stuff

data/examples/declaration/value/function/if-with-comment-in-branches-out.hs

@@ -0,0 +1,14 @@
+foo =
+  if x
+    then
+      -- comment 1
+      -- comment 2
+      case a of
+        Just b -> _
+        Nothing -> _
+    else
+      -- comment 3
+      -- comment 4
+      case a of
+        Just b -> _
+        Nothing -> _

data/examples/declaration/value/function/if-with-comment-in-branches-with-functions-out.hs

@@ -0,0 +1,12 @@
+foo =
+  if x
+    then
+      -- comment 1
+      -- comment 2
+      foo 1 2 3
+    else
+      -- comment 3
+      -- comment 4
+      foo
+        "here"
+        "there"

data/examples/declaration/value/function/if-with-comment-in-branches-with-functions.hs

@@ -0,0 +1,12 @@
+foo =
+  if x
+    then
+      -- comment 1
+      -- comment 2
+      foo 1 2 3
+    else
+      -- comment 3
+      -- comment 4
+      foo
+        "here"
+        "there"

data/examples/declaration/value/function/if-with-comment-in-branches.hs

@@ -0,0 +1,14 @@
+foo =
+  if x
+    then
+      -- comment 1
+      -- comment 2
+      case a of
+        Just b -> _
+        Nothing -> _
+    else
+      -- comment 3
+      -- comment 4
+      case a of
+        Just b -> _
+        Nothing -> _

data/examples/declaration/value/function/if-with-comment-next-to-keyword-out.hs

@@ -0,0 +1,10 @@
+foo =
+  if x
+    then -- comment 1
+    -- comment 2
+      foo 1 2 3
+    else -- comment 3
+    -- comment 4
+      foo
+        "here"
+        "there"

data/examples/declaration/value/function/if-with-comment-next-to-keyword.hs

@@ -0,0 +1,10 @@
+foo =
+  if x
+    then -- comment 1
+      -- comment 2
+      foo 1 2 3
+    else -- comment 3
+      -- comment 4
+      foo
+        "here"
+        "there"

src/Ormolu/Printer/Meat/Declaration/Value.hs

@@ -363,8 +363,8 @@ p_hsCmd' isApp s = \case
     p_case isApp cmdPlacement p_hsCmd e mgroup
   HsCmdLamCase _ variant mgroup ->
     p_lamcase isApp variant cmdPlacement p_hsCmd mgroup
-  HsCmdIf _ _ if' then' else' ->
-    p_if cmdPlacement p_hsCmd if' then' else'
+  HsCmdIf anns _ if' then' else' ->
+    p_if cmdPlacement p_hsCmd anns if' then' else'
   HsCmdLet _ _ localBinds _ c ->
     p_let p_hsCmd localBinds c
   HsCmdDo _ es -> do
@@ -731,8 +731,8 @@ p_hsExpr' isApp s = \case
     p_unboxedSum N tag arity (located e p_hsExpr)
   HsCase _ e mgroup ->
     p_case isApp exprPlacement p_hsExpr e mgroup
-  HsIf _ if' then' else' ->
-    p_if exprPlacement p_hsExpr if' then' else'
+  HsIf anns if' then' else' ->
+    p_if exprPlacement p_hsExpr anns if' then' else'
   HsMultiIf _ guards -> do
     txt "if"
     breakpoint
@@ -971,28 +971,56 @@ p_if ::
   (body -> Placement) ->
   -- | Render
   (body -> R ()) ->
+  -- | Annotations
+  EpAnn AnnsIf ->
   -- | If
   LHsExpr GhcPs ->
   -- | Then
   LocatedA body ->
   -- | Else
   LocatedA body ->
   R ()
-p_if placer render if' then' else' = do
+p_if placer render epAnn if' then' else' = do
   txt "if"
   space
   located if' p_hsExpr
   breakpoint
   inci $ do
-    txt "then"
+    locatedToken thenSpan "then"
     space
-    located then' $ \x ->
-      placeHanging (placer x) (render x)
+    placeHangingLocated thenSpan then'
     breakpoint
-    txt "else"
+    locatedToken elseSpan "else"
     space
-    located else' $ \x ->
-      placeHanging (placer x) (render x)
+    placeHangingLocated elseSpan else'
+  where
+    (thenSpan, elseSpan, commentSpans) =
+      case epAnn of
+        EpAnn {anns = AnnsIf {aiThen, aiElse}, comments} ->
+          ( loc' $ epaLocationRealSrcSpan aiThen,
+            loc' $ epaLocationRealSrcSpan aiElse,
+            map (anchor . getLoc) $
+              case comments of
+                EpaComments cs -> cs
+                EpaCommentsBalanced pre post -> pre <> post
+          )
+        EpAnnNotUsed ->
+          (noSrcSpan, noSrcSpan, [])
+
+    locatedToken tokenSpan token =
+      located (L tokenSpan ()) $ \_ -> txt token
+
+    betweenSpans spanA spanB s = spanA < s && s < spanB
+
+    placeHangingLocated tokenSpan bodyLoc@(L _ body) = do
+      let bodySpan = getLoc' bodyLoc
+          hasComments = fromMaybe False $ do
+            tokenRealSpan <- srcSpanToRealSrcSpan tokenSpan
+            bodyRealSpan <- srcSpanToRealSrcSpan bodySpan
+            pure $ any (betweenSpans tokenRealSpan bodyRealSpan) commentSpans
+          placement = if hasComments then Normal else placer body
+      switchLayout [tokenSpan, bodySpan] $
+        placeHanging placement (located bodyLoc render)
 
 p_let ::
   -- | Render

src/Ormolu/Utils.hs

@@ -147,6 +147,9 @@ class HasSrcSpan l where
 instance HasSrcSpan SrcSpan where
   loc' = id
 
+instance HasSrcSpan RealSrcSpan where
+  loc' l = RealSrcSpan l Strict.Nothing
+
 instance HasSrcSpan (SrcSpanAnn' ann) where
   loc' = locA