Fix retain comments on matches (lambda rhs etc.)

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

@@ -250,19 +250,15 @@ layoutPatternBind
   -> BriDocNumbered
   -> LMatch GhcPs (LHsExpr GhcPs)
   -> ToBriDocM BriDocNumbered
-layoutPatternBind funId binderDoc lmatch@(L _ match) =
+layoutPatternBind funId binderDoc lmatch@(L _ match) = do
-  docHandleComms lmatch $ do
+  let Match epAnn matchCtx pats (GRHSs _ grhss whereBinds) = match
-    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
+  mIdDoc <- case matchCtx of
-      Match epAnn (FunRhs matchId _ _) _ _ ->
+    FunRhs matchId _ _ -> fmap Just $ do
-        fmap Just
-          $ docHandleComms epAnn $ do
       t <- lrdrNameToTextAnn matchId
       let t' = fixPatternBindIdentifier match t
       docLit t'
@@ -300,7 +296,7 @@ layoutPatternBind funId binderDoc lmatch@(L _ match) =
   -- let mWhereArg = mWhereDocs <&> (,) (mkAnnKey lmatch)
   let alignmentToken = if null pats then Nothing else funId
   let hasComments = hasAnyCommentsBelow lmatch
-    docFlushCommsPost True matchEndLoc
+  docHandleComms lmatch $ docHandleComms epAnn $ docFlushCommsPost True matchEndLoc
     $ layoutPatternBindFinal
       alignmentToken
       binderDoc

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
+      | Match epAnn _matchCtx pats (GRHSs _ [lgrhs] llocals) <- match
-      , GRHSs _ [lgrhs] llocals <- m_grhss 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