Fix the misplacement of comments in a do expression (#707)

* Fix the misplacement of comments in a do expression * Add a changelog * Fix * Fix
mihaimaruseac · Apr 22, 2023 · c0cd349 · c0cd349
1 parent 8971072
commit c0cd349
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@
 
 - Fixed module names being removed from uses of qualified `do` ([#696]).
 - Misplaced haddocks for class declarations ([#706]).
+- Misplaced comments in do expressions ([#707]).
 
 ### Removed
 
@@ -339,6 +340,7 @@ This version is accidentally pushlished, and is the same as 5.3.3.
 [@uhbif19]: https://github.com/uhbif19
 [@toku-sa-n]: https://github.com/toku-sa-n
 
+[#707]: https://github.com/mihaimaruseac/hindent/pull/707
 [#706]: https://github.com/mihaimaruseac/hindent/pull/706
 [#699]: https://github.com/mihaimaruseac/hindent/pull/699
 [#696]: https://github.com/mihaimaruseac/hindent/pull/696

diff --git a/TESTS.md b/TESTS.md
@@ -3041,15 +3041,12 @@ gamma = do
   -- the very last block is detected differently
 ```
 
-Doesn't work yet (wrong comment position detection)
+Comments in a do expression
 
-```haskell pending
+```haskell
 gamma = do
   -- in the beginning of a do-block
   delta
-  where
-    -- before alpha
-    alpha = alpha
 ```
 
 Comments in a class instance

diff --git a/src/HIndent/ModulePreprocessing/CommentRelocation.hs b/src/HIndent/ModulePreprocessing/CommentRelocation.hs
@@ -75,6 +75,7 @@ relocateComments = evalState . relocate
       relocateCommentsInClass >=>
       relocateCommentsBeforeTopLevelDecls >=>
       relocateCommentsSameLine >=>
+      relocateCommentsInDoExpr >=>
       relocateCommentsTopLevelWhereClause >=>
       relocateCommentsAfter >=> assertAllCommentsAreConsumed
     assertAllCommentsAreConsumed x = do
@@ -177,6 +178,32 @@ relocateCommentsInClass =
       realSrcSpanStart classAnchor < realSrcSpanStart comAnc
     cond _ _ _ = False
 
+-- | Locates comments before each statement in a do expression.
+relocateCommentsInDoExpr :: HsModule' -> WithComments HsModule'
+relocateCommentsInDoExpr =
+  relocateCommentsBeforeEachElement
+    elemGetter
+    elemSetter
+    annGetter
+    annSetter
+    cond
+  where
+    elemGetter :: LHsExpr GhcPs -> [ExprLStmt GhcPs]
+    elemGetter (L _ (HsDo _ DoExpr {} (L _ xs))) = xs
+    elemGetter (L _ (HsDo _ MDoExpr {} (L _ xs))) = xs
+    elemGetter _ = []
+    elemSetter xs (L sp (HsDo ext flavor@DoExpr {} (L sp' _))) =
+      L sp (HsDo ext flavor (L sp' xs))
+    elemSetter xs (L sp (HsDo ext flavor@MDoExpr {} (L sp' _))) =
+      L sp (HsDo ext flavor (L sp' xs))
+    elemSetter _ x = x
+    annGetter (L SrcSpanAnn {..} _) = ann
+    annSetter newAnn (L SrcSpanAnn {..} x) = L SrcSpanAnn {ann = newAnn, ..} x
+    cond (L SrcSpanAnn {ann = EpAnn {entry = Anchor {anchor = doAnchor}}} _) (L SrcSpanAnn {ann = EpAnn {entry = Anchor {anchor = elemAnchor}}} _) comAnc =
+      srcSpanStartLine comAnc < srcSpanStartLine elemAnchor &&
+      realSrcSpanStart doAnchor < realSrcSpanStart comAnc
+    cond _ _ _ = False
+
 -- | This function locates comments located before top-level declarations.
 relocateCommentsBeforeTopLevelDecls :: HsModule' -> WithComments HsModule'
 relocateCommentsBeforeTopLevelDecls = everywhereM (applyM f)