Fix retain comments on matches (lambda rhs etc.)

source/library/Language/Haskell/Brittany/Internal/ToBriDoc/Decl.hs

@@ -250,64 +250,60 @@ layoutPatternBind
   -> BriDocNumbered
   -> LMatch GhcPs (LHsExpr GhcPs)
   -> ToBriDocM BriDocNumbered
-layoutPatternBind funId binderDoc lmatch@(L _ match) =
-  docHandleComms lmatch $ do
-    let pats = m_pats match
-    let (GRHSs _ grhss whereBinds) = m_grhss match
-    patDocs <- pats `forM` \p -> fmap return $ colsWrapPat =<< layoutPat p
-    let isInfix = isInfixMatch match
-    let matchEndLoc = case GHC.locA $ GHC.getLoc lmatch of
-          GHC.RealSrcSpan s _ -> Just $ GHC.realSrcSpanEnd s
-          GHC.UnhelpfulSpan{} -> Nothing
-    mIdDoc <- case match of
-      Match epAnn (FunRhs matchId _ _) _ _ ->
-        fmap Just
-          $ docHandleComms epAnn $ do
-            t <- lrdrNameToTextAnn matchId
-            let t' = fixPatternBindIdentifier match t
-            docLit t'
-      _ -> pure Nothing
-    patDoc <- docHandleComms lmatch $ case (mIdDoc, patDocs) of
-      (Just idDoc, p1 : p2 : pr) | isInfix -> if null pr
-        then docCols
-          ColPatternsFuncInfix
-          [ appSep $ docForceSingleline p1
-          , appSep $ pure idDoc
-          , docForceSingleline p2
-          ]
-        else docCols
-          ColPatternsFuncInfix
-          ([ docCols
-               ColPatterns
-               [ docParenL
-               , appSep $ docForceSingleline p1
-               , appSep $ pure idDoc
-               , docForceSingleline p2
-               , appSep $ docParenR
-               ]
-           ]
-          ++ (spacifyDocs $ docForceSingleline <$> pr)
-          )
-      (Just idDoc, []) -> pure idDoc
-      (Just idDoc, ps) ->
-        docCols ColPatternsFuncPrefix
-          $ appSep (pure idDoc)
-          : (spacifyDocs $ docForceSingleline <$> ps)
-      (Nothing, ps) ->
-        docCols ColPatterns
-          $ (List.intersperse docSeparator $ docForceSingleline <$> ps)
-    mWhereDocs <- layoutLocalBinds whereBinds
-    -- let mWhereArg = mWhereDocs <&> (,) (mkAnnKey lmatch)
-    let alignmentToken = if null pats then Nothing else funId
-    let hasComments = hasAnyCommentsBelow lmatch
-    docFlushCommsPost True matchEndLoc
-      $ layoutPatternBindFinal
-        alignmentToken
-        binderDoc
-        (Just patDoc)
-        (Right grhss)
-        mWhereDocs
-        hasComments
+layoutPatternBind funId binderDoc lmatch@(L _ match) = do
+  let Match epAnn matchCtx pats (GRHSs _ grhss whereBinds) = match
+  patDocs <- pats `forM` \p -> fmap return $ colsWrapPat =<< layoutPat p
+  let isInfix = isInfixMatch match
+  let matchEndLoc = case GHC.locA $ GHC.getLoc lmatch of
+        GHC.RealSrcSpan s _ -> Just $ GHC.realSrcSpanEnd s
+        GHC.UnhelpfulSpan{} -> Nothing
+  mIdDoc <- case matchCtx of
+    FunRhs matchId _ _ -> fmap Just $ do
+      t <- lrdrNameToTextAnn matchId
+      let t' = fixPatternBindIdentifier match t
+      docLit t'
+    _ -> pure Nothing
+  patDoc <- docHandleComms lmatch $ case (mIdDoc, patDocs) of
+    (Just idDoc, p1 : p2 : pr) | isInfix -> if null pr
+      then docCols
+        ColPatternsFuncInfix
+        [ appSep $ docForceSingleline p1
+        , appSep $ pure idDoc
+        , docForceSingleline p2
+        ]
+      else docCols
+        ColPatternsFuncInfix
+        ([ docCols
+             ColPatterns
+             [ docParenL
+             , appSep $ docForceSingleline p1
+             , appSep $ pure idDoc
+             , docForceSingleline p2
+             , appSep $ docParenR
+             ]
+         ]
+        ++ (spacifyDocs $ docForceSingleline <$> pr)
+        )
+    (Just idDoc, []) -> pure idDoc
+    (Just idDoc, ps) ->
+      docCols ColPatternsFuncPrefix
+        $ appSep (pure idDoc)
+        : (spacifyDocs $ docForceSingleline <$> ps)
+    (Nothing, ps) ->
+      docCols ColPatterns
+        $ (List.intersperse docSeparator $ docForceSingleline <$> ps)
+  mWhereDocs <- layoutLocalBinds whereBinds
+  -- let mWhereArg = mWhereDocs <&> (,) (mkAnnKey lmatch)
+  let alignmentToken = if null pats then Nothing else funId
+  let hasComments = hasAnyCommentsBelow lmatch
+  docHandleComms lmatch $ docHandleComms epAnn $ docFlushCommsPost True matchEndLoc
+    $ layoutPatternBindFinal
+      alignmentToken
+      binderDoc
+      (Just patDoc)
+      (Right grhss)
+      mWhereDocs
+      hasComments
 
 fixPatternBindIdentifier :: Match GhcPs (LHsExpr GhcPs) -> Text -> Text
 fixPatternBindIdentifier match idStr = go $ m_ctxt match

source/library/Language/Haskell/Brittany/Internal/ToBriDoc/Expr.hs

@@ -55,10 +55,9 @@ layoutExpr lexpr@(L _ expr) = do
     HsLit epAnn lit -> docHandleComms epAnn $ do
       allocateNode $ litBriDoc lit
     HsLam _ (MG _ (L _ [(L _ match)]) _)
-      | pats <- m_pats match
-      , GRHSs _ [lgrhs] llocals <- m_grhss match
+      | Match epAnn _matchCtx pats (GRHSs _ [lgrhs] llocals) <- match
       , EmptyLocalBinds{} <- llocals
-      , L _ (GRHS epAnn [] body) <- lgrhs
+      , L _ (GRHS rhsEpAnn [] body) <- lgrhs
       -> do
         patDocs <- zip (True : repeat False) pats `forM` \(isFirst, p) ->
           fmap return $ do
@@ -85,6 +84,7 @@ layoutExpr lexpr@(L _ expr) = do
           shareDoc
           $ docAddBaseY BrIndentRegular
           $ docHandleComms epAnn
+          $ docHandleComms rhsEpAnn
           $ layoutExpr body
         let funcPatternPartLine = docCols
               ColCasePattern