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