CommentRelocation.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | Comment relocation for pretty-printing comments correctly.
--
-- HIndent gathers all comments above a function, an import, a module
-- declaration, etc. For example, HIndent formats the following code
--
-- > f :: Int
-- > f = 1
-- >
-- > -- A comment between f and g
-- >
-- > -- Another comment between f and g
-- >
-- > g :: Int
-- > g = 2
--
-- to
--
-- > f :: Int
-- > f = 1
-- >
-- > -- A comment between f and g
-- > -- Another comment between f and g
-- > g :: Int
-- > g = 2
--
-- AST nodes must have the information of which comments are above, on the
-- same line, and below. However, AST nodes generated by a parser of
-- 'ghc-lib-parser' only contain comments after them. 'relocateComments' is
-- defined to solve the problem.
module HIndent.ModulePreprocessing.CommentRelocation
  ( relocateComments
  ) where

import Control.Exception
import Control.Monad.State
import Data.Foldable
import Data.Function
import Data.List
import Data.Traversable
import GHC.Data.Bag
import GHC.Types.SrcLoc
import Generics.SYB hiding (GT, typeOf, typeRep)
import HIndent.GhcLibParserWrapper.GHC.Hs
import HIndent.Pragma
import HIndent.Pretty.SigBindFamily
import Type.Reflection
#if MIN_VERSION_GLASGOW_HASKELL(9, 6, 0, 0)
import Control.Monad
#endif
-- | A wrapper type used in everywhereMEpAnnsBackwards' to collect all
-- 'EpAnn's to apply a function with them in order their positions.
data Wrapper =
  forall a. Typeable (EpAnn a) =>
            Wrapper (EpAnn a)

-- | 'State' with comments.
type WithComments = State [LEpaComment]

-- | This function collects all comments from the passed 'HsModule', and
-- modifies all 'EpAnn's so that all 'EpAnn's have 'EpaCommentsBalanced's.
relocateComments :: HsModule' -> [LEpaComment] -> HsModule'
relocateComments = evalState . relocate
  where
    relocate =
      relocatePragmas
        >=> relocateCommentsBeforePragmas
        >=> relocateCommentsInExportList
        >=> relocateCommentsInClass
        >=> relocateCommentsBeforeTopLevelDecls
        >=> relocateCommentsSameLine
        >=> relocateCommentsInDoExpr
        >=> relocateCommentsInCase
        >=> relocateCommentsTopLevelWhereClause
        >=> relocateCommentsAfter
        >=> assertAllCommentsAreConsumed
        >=> moveCommentsFromFunIdToMcFun
    assertAllCommentsAreConsumed x = do
      cs <- get
      assert (null cs) (pure x)
-- | This function locates pragmas to the module's EPA.
#if MIN_VERSION_ghc_lib_parser(9,6,1)
relocatePragmas :: HsModule GhcPs -> WithComments (HsModule GhcPs)
relocatePragmas m@HsModule {hsmodExt = xmod@XModulePs {hsmodAnn = epa@EpAnn {}}} = do
  newAnn <- insertComments (isPragma . ac_tok . unLoc) insertPriorComments epa
  return m {hsmodExt = xmod {hsmodAnn = newAnn}}
#else
relocatePragmas :: HsModule -> WithComments HsModule
relocatePragmas m@HsModule {hsmodAnn = epa@EpAnn {}} = do
  newAnn <- insertComments (isPragma . ac_tok . unLoc) insertPriorComments epa
  return m {hsmodAnn = newAnn}
#endif
relocatePragmas m = pure m
-- | This function locates comments that are located before pragmas to the
-- module's EPA.
#if MIN_VERSION_ghc_lib_parser(9, 10, 1)
relocateCommentsBeforePragmas :: HsModule GhcPs -> WithComments (HsModule GhcPs)
relocateCommentsBeforePragmas m@HsModule {hsmodExt = xmod@XModulePs {hsmodAnn = ann}}
  | pragmaExists m = do
    newAnn <- insertCommentsByPos (< startPosOfPragmas) insertPriorComments ann
    pure m {hsmodExt = xmod {hsmodAnn = newAnn}}
  | otherwise = pure m
  where
    startPosOfPragmas =
      let loc = getLoc $ head $ priorComments $ comments ann
       in case loc of
            EpaSpan (RealSrcSpan sp _) -> sp
            _ -> undefined
#elif MIN_VERSION_ghc_lib_parser(9, 6, 1)
relocateCommentsBeforePragmas :: HsModule GhcPs -> WithComments (HsModule GhcPs)
relocateCommentsBeforePragmas m@HsModule {hsmodExt = xmod@XModulePs {hsmodAnn = ann}}
  | pragmaExists m = do
    newAnn <- insertCommentsByPos (< startPosOfPragmas) insertPriorComments ann
    pure m {hsmodExt = xmod {hsmodAnn = newAnn}}
  | otherwise = pure m
  where
    startPosOfPragmas = anchor $ getLoc $ head $ priorComments $ comments ann
#else
relocateCommentsBeforePragmas :: HsModule -> WithComments HsModule
relocateCommentsBeforePragmas m@HsModule {hsmodAnn = ann}
  | pragmaExists m = do
    newAnn <- insertCommentsByPos (< startPosOfPragmas) insertPriorComments ann
    pure m {hsmodAnn = newAnn}
  | otherwise = pure m
  where
    startPosOfPragmas = anchor $ getLoc $ head $ priorComments $ comments ann
#endif
#if MIN_VERSION_ghc_lib_parser(9, 10, 1)
-- | This function locates comments that are located before each element of
-- an export list.
relocateCommentsInExportList :: HsModule' -> WithComments HsModule'
relocateCommentsInExportList =
  relocateCommentsBeforeEachElement
    elemGetter
    elemSetter
    annGetter
    annSetter
    cond
  where
    elemGetter :: HsModule' -> [LIE GhcPs]
    elemGetter HsModule {hsmodExports = Just (L _ xs)} = xs
    elemGetter _ = []
    elemSetter xs HsModule {hsmodExports = Just (L sp _), ..} =
      HsModule {hsmodExports = Just (L sp xs), ..}
    elemSetter _ x = x
    annGetter (L ann _) = ann
    annSetter newAnn (L _ x) = L newAnn x
    cond HsModule {hsmodExports = Just (L EpAnn {entry = EpaSpan (RealSrcSpan listAnc _)} _)} (L EpAnn {entry = EpaSpan (RealSrcSpan elemAnc _)} _) comAnc =
      srcSpanStartLine comAnc < srcSpanStartLine elemAnc
        && realSrcSpanStart listAnc < realSrcSpanStart comAnc
    cond _ _ _ = False

-- | Locates comments before each case branch.
relocateCommentsInCase :: HsModule' -> WithComments HsModule'
relocateCommentsInCase =
  relocateCommentsBeforeEachElement
    elemGetter
    elemSetter
    annGetter
    annSetter
    cond
  where
    elemGetter :: LHsExpr GhcPs -> [LMatch GhcPs (LHsExpr GhcPs)]
    elemGetter (L _ (HsCase _ _ (MG {mg_alts = L _ xs}))) = xs
    elemGetter _ = []
    elemSetter xs (L sp (HsCase ext expr (MG {mg_alts = L sp' _, ..}))) =
      L sp (HsCase ext expr (MG {mg_alts = L sp' xs, ..}))
    elemSetter _ x = x
    annGetter (L ann _) = ann
    annSetter newAnn (L _ x) = L newAnn x
    cond (L EpAnn {entry = EpaSpan (RealSrcSpan caseAnchor _)} _) (L EpAnn {entry = EpaSpan (RealSrcSpan branchAnchor _)} _) comAnc =
      srcSpanStartLine comAnc < srcSpanStartLine branchAnchor
        && realSrcSpanStart caseAnchor < realSrcSpanStart comAnc
    cond _ _ _ = False

-- | Locates comments before each class element.
relocateCommentsInClass :: HsModule' -> WithComments HsModule'
relocateCommentsInClass =
  relocateCommentsBeforeEachElement
    elemGetter
    elemSetter
    annGetter
    annSetter
    cond
  where
    elemGetter :: LHsDecl GhcPs -> [LSigBindFamily]
    elemGetter (L _ (TyClD _ ClassDecl {..})) =
      mkSortedLSigBindFamilyList
        tcdSigs
        (bagToList tcdMeths)
        tcdATs
        tcdATDefs
        []
    elemGetter _ = []
    elemSetter xs (L sp (TyClD ext ClassDecl {..})) = L sp (TyClD ext newDecl)
      where
        newDecl =
          ClassDecl
            { tcdSigs = sigs
            , tcdMeths = listToBag binds
            , tcdATs = typeFamilies
            , tcdATDefs = tyFamInsts
            , ..
            }
        (sigs, binds, typeFamilies, tyFamInsts, _) =
          destructLSigBindFamilyList xs
    elemSetter _ x = x
    annGetter (L ann _) = ann
    annSetter newAnn (L _ x) = L newAnn x
    cond (L EpAnn {entry = EpaSpan (RealSrcSpan classAnchor _)} _) (L EpAnn {entry = EpaSpan (RealSrcSpan elemAnchor _)} _) comAnc =
      srcSpanStartLine comAnc < srcSpanStartLine elemAnchor
        && 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 ann _) = ann
    annSetter newAnn (L _ x) = L newAnn x
    cond (L EpAnn {entry = EpaSpan (RealSrcSpan doAnchor _)} _) (L EpAnn {entry = EpaSpan (RealSrcSpan 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)
  where
    f epa@EpAnn {..}
      | EpaSpan (RealSrcSpan anc _) <- entry =
        insertCommentsByPos (isBefore anc) insertPriorComments epa
      | otherwise = pure epa
    isBefore anc comAnc =
      srcSpanStartCol anc == 1
        && srcSpanStartCol comAnc == 1
        && srcSpanStartLine comAnc < srcSpanStartLine anc

-- | This function scans the given AST from bottom to top and locates
-- comments that are on the same line as the node.  Comments are stored in
-- the 'followingComments' of 'EpaCommentsBalanced'.
relocateCommentsSameLine :: HsModule' -> WithComments HsModule'
relocateCommentsSameLine = everywhereMEpAnnsBackwards f
  where
    f epa@EpAnn {..}
      | EpaSpan (RealSrcSpan anc _) <- entry =
        insertCommentsByPos (isOnSameLine anc) insertFollowingComments epa
      | otherwise = pure epa
    isOnSameLine anc comAnc =
      srcSpanStartLine comAnc == srcSpanStartLine anc
        && srcSpanStartLine comAnc == srcSpanEndLine anc

-- | This function locates comments above the top-level declarations in
-- a 'where' clause in the topmost declaration.
relocateCommentsTopLevelWhereClause :: HsModule' -> WithComments HsModule'
relocateCommentsTopLevelWhereClause m@HsModule {..} = do
  hsmodDecls' <- mapM relocateCommentsDeclWhereClause hsmodDecls
  pure m {hsmodDecls = hsmodDecls'}
  where
    relocateCommentsDeclWhereClause (L l (ValD ext fb@(FunBind {fun_matches = MG {..}}))) = do
      mg_alts' <- mapM (mapM relocateCommentsMatch) mg_alts
      pure $ L l (ValD ext fb {fun_matches = MG {mg_alts = mg_alts', ..}})
    relocateCommentsDeclWhereClause x = pure x
    relocateCommentsMatch (L l match@Match {m_grhss = gs@GRHSs {grhssLocalBinds = (HsValBinds ext (ValBinds ext' binds sigs))}}) = do
      (binds', sigs') <- relocateCommentsBindsSigs binds sigs
      let localBinds = HsValBinds ext (ValBinds ext' binds' sigs')
      pure $ L l match {m_grhss = gs {grhssLocalBinds = localBinds}}
    relocateCommentsMatch x = pure x
    relocateCommentsBindsSigs ::
         LHsBindsLR GhcPs GhcPs
      -> [LSig GhcPs]
      -> WithComments (LHsBindsLR GhcPs GhcPs, [LSig GhcPs])
    relocateCommentsBindsSigs binds sigs = do
      bindsSigs' <- mapM addCommentsBeforeEpAnn bindsSigs
      pure (listToBag $ filterLBind bindsSigs', filterLSig bindsSigs')
      where
        bindsSigs = mkSortedLSigBindFamilyList sigs (bagToList binds) [] [] []
    addCommentsBeforeEpAnn (L epa@EpAnn {..} x)
      | EpaSpan (RealSrcSpan anc _) <- entry = do
        cs <- get
        let (notAbove, above) =
              partitionAboveNotAbove (sortCommentsByLocation cs) anc
            epa' = epa {comments = insertPriorComments comments above}
        put notAbove
        pure $ L epa' x
      | otherwise = undefined
    addCommentsBeforeEpAnn x = pure x
    partitionAboveNotAbove cs sp =
      fst
        $ foldr'
            (\c@(L l _) ((ls, rs), lastSpan) ->
               case l of
                 EpaSpan (RealSrcSpan anc _) ->
                   if anc `isAbove` lastSpan
                     then ((ls, c : rs), anc)
                     else ((c : ls, rs), lastSpan)
                 _ -> undefined)
            (([], []), sp)
            cs
    isAbove comAnc anc =
      srcSpanStartCol comAnc == srcSpanStartCol anc
        && srcSpanEndLine comAnc + 1 == srcSpanStartLine anc
#else
-- | This function locates comments that are located before each element of
-- an export list.
relocateCommentsInExportList :: HsModule' -> WithComments HsModule'
relocateCommentsInExportList =
  relocateCommentsBeforeEachElement
    elemGetter
    elemSetter
    annGetter
    annSetter
    cond
  where
    elemGetter :: HsModule' -> [LIE GhcPs]
    elemGetter HsModule {hsmodExports = Just (L _ xs)} = xs
    elemGetter _ = []
    elemSetter xs HsModule {hsmodExports = Just (L sp _), ..} =
      HsModule {hsmodExports = Just (L sp xs), ..}
    elemSetter _ x = x
    annGetter (L SrcSpanAnn {..} _) = ann
    annSetter newAnn (L SrcSpanAnn {..} x) = L SrcSpanAnn {ann = newAnn, ..} x
    cond HsModule {hsmodExports = Just (L SrcSpanAnn {ann = EpAnn {entry = Anchor {anchor = listAnc}}} _)} (L SrcSpanAnn {ann = EpAnn {entry = Anchor {anchor = elemAnc}}} _) comAnc =
      srcSpanStartLine comAnc < srcSpanStartLine elemAnc
        && realSrcSpanStart listAnc < realSrcSpanStart comAnc
    cond _ _ _ = False

-- | Locates comments before each case branch.
relocateCommentsInCase :: HsModule' -> WithComments HsModule'
relocateCommentsInCase =
  relocateCommentsBeforeEachElement
    elemGetter
    elemSetter
    annGetter
    annSetter
    cond
  where
    elemGetter :: LHsExpr GhcPs -> [LMatch GhcPs (LHsExpr GhcPs)]
    elemGetter (L _ (HsCase _ _ (MG {mg_alts = L _ xs}))) = xs
    elemGetter _ = []
    elemSetter xs (L sp (HsCase ext expr (MG {mg_alts = L sp' _, ..}))) =
      L sp (HsCase ext expr (MG {mg_alts = 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 = caseAnchor}}} _) (L SrcSpanAnn {ann = EpAnn {entry = Anchor {anchor = branchAnchor}}} _) comAnc =
      srcSpanStartLine comAnc < srcSpanStartLine branchAnchor
        && realSrcSpanStart caseAnchor < realSrcSpanStart comAnc
    cond _ _ _ = False

-- | Locates comments before each class element.
relocateCommentsInClass :: HsModule' -> WithComments HsModule'
relocateCommentsInClass =
  relocateCommentsBeforeEachElement
    elemGetter
    elemSetter
    annGetter
    annSetter
    cond
  where
    elemGetter :: LHsDecl GhcPs -> [LSigBindFamily]
    elemGetter (L _ (TyClD _ ClassDecl {..})) =
      mkSortedLSigBindFamilyList
        tcdSigs
        (bagToList tcdMeths)
        tcdATs
        tcdATDefs
        []
    elemGetter _ = []
    elemSetter xs (L sp (TyClD ext ClassDecl {..})) = L sp (TyClD ext newDecl)
      where
        newDecl =
          ClassDecl
            { tcdSigs = sigs
            , tcdMeths = listToBag binds
            , tcdATs = typeFamilies
            , tcdATDefs = tyFamInsts
            , ..
            }
        (sigs, binds, typeFamilies, tyFamInsts, _) =
          destructLSigBindFamilyList 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 = classAnchor}}} _) (L SrcSpanAnn {ann = EpAnn {entry = Anchor {anchor = elemAnchor}}} _) comAnc =
      srcSpanStartLine comAnc < srcSpanStartLine elemAnchor
        && 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)
  where
    f epa@EpAnn {..} =
      insertCommentsByPos (isBefore $ anchor entry) insertPriorComments epa
    f EpAnnNotUsed = pure EpAnnNotUsed
    isBefore anc comAnc =
      srcSpanStartCol anc == 1
        && srcSpanStartCol comAnc == 1
        && srcSpanStartLine comAnc < srcSpanStartLine anc

-- | This function scans the given AST from bottom to top and locates
-- comments that are on the same line as the node.  Comments are stored in
-- the 'followingComments' of 'EpaCommentsBalanced'.
relocateCommentsSameLine :: HsModule' -> WithComments HsModule'
relocateCommentsSameLine = everywhereMEpAnnsBackwards f
  where
    f epa@EpAnn {..} =
      insertCommentsByPos
        (isOnSameLine $ anchor entry)
        insertFollowingComments
        epa
    f EpAnnNotUsed = pure EpAnnNotUsed
    isOnSameLine anc comAnc =
      srcSpanStartLine comAnc == srcSpanStartLine anc
        && srcSpanStartLine comAnc == srcSpanEndLine anc

-- | This function locates comments above the top-level declarations in
-- a 'where' clause in the topmost declaration.
relocateCommentsTopLevelWhereClause :: HsModule' -> WithComments HsModule'
relocateCommentsTopLevelWhereClause m@HsModule {..} = do
  hsmodDecls' <- mapM relocateCommentsDeclWhereClause hsmodDecls
  pure m {hsmodDecls = hsmodDecls'}
  where
    relocateCommentsDeclWhereClause (L l (ValD ext fb@(FunBind {fun_matches = MG {..}}))) = do
      mg_alts' <- mapM (mapM relocateCommentsMatch) mg_alts
      pure $ L l (ValD ext fb {fun_matches = MG {mg_alts = mg_alts', ..}})
    relocateCommentsDeclWhereClause x = pure x
    relocateCommentsMatch (L l match@Match {m_grhss = gs@GRHSs {grhssLocalBinds = (HsValBinds ext (ValBinds ext' binds sigs))}}) = do
      (binds', sigs') <- relocateCommentsBindsSigs binds sigs
      let localBinds = HsValBinds ext (ValBinds ext' binds' sigs')
      pure $ L l match {m_grhss = gs {grhssLocalBinds = localBinds}}
    relocateCommentsMatch x = pure x
    relocateCommentsBindsSigs ::
         LHsBindsLR GhcPs GhcPs
      -> [LSig GhcPs]
      -> WithComments (LHsBindsLR GhcPs GhcPs, [LSig GhcPs])
    relocateCommentsBindsSigs binds sigs = do
      bindsSigs' <- mapM addCommentsBeforeEpAnn bindsSigs
      pure (listToBag $ filterLBind bindsSigs', filterLSig bindsSigs')
      where
        bindsSigs = mkSortedLSigBindFamilyList sigs (bagToList binds) [] [] []
    addCommentsBeforeEpAnn (L (SrcSpanAnn epa@EpAnn {..} sp) x) = do
      cs <- get
      let (notAbove, above) =
            partitionAboveNotAbove (sortCommentsByLocation cs) entry
          epa' = epa {comments = insertPriorComments comments above}
      put notAbove
      pure $ L (SrcSpanAnn epa' sp) x
    addCommentsBeforeEpAnn x = pure x
    partitionAboveNotAbove cs sp =
      fst
        $ foldr'
            (\c@(L l _) ((ls, rs), lastSpan) ->
               if anchor l `isAbove` anchor lastSpan
                 then ((ls, c : rs), l)
                 else ((c : ls, rs), lastSpan))
            (([], []), sp)
            cs
    isAbove comAnc anc =
      srcSpanStartCol comAnc == srcSpanStartCol anc
        && srcSpanEndLine comAnc + 1 == srcSpanStartLine anc
#endif
-- | This function scans the given AST from bottom to top and locates
-- comments in the comment pool after each node on it.
relocateCommentsAfter :: HsModule' -> WithComments HsModule'
#if MIN_VERSION_ghc_lib_parser(9, 10, 1)
relocateCommentsAfter = everywhereMEpAnnsBackwards f
  where
    f epa@EpAnn {..} =
      insertCommentsByPos (isAfter $ anchor entry) insertFollowingComments epa
    isAfter anc comAnc = srcSpanEndLine anc <= srcSpanStartLine comAnc
#else
relocateCommentsAfter = everywhereMEpAnnsBackwards f
  where
    f epa@EpAnn {..} =
      insertCommentsByPos (isAfter $ anchor entry) insertFollowingComments epa
    f EpAnnNotUsed = pure EpAnnNotUsed
    isAfter anc comAnc = srcSpanEndLine anc <= srcSpanStartLine comAnc
#endif
-- | Locates comments before each element in a parent.
relocateCommentsBeforeEachElement ::
     forall a b c. Typeable a
  => (a -> [b]) -- ^ Element getter
  -> ([b] -> a -> a) -- ^ Element setter
  -> (b -> EpAnn c) -- ^ Annotation getter
  -> (EpAnn c -> b -> b) -- ^ Annotation setter
  -> (a -> b -> RealSrcSpan -> Bool) -- ^ The function to decide whether to locate comments
  -> HsModule'
  -> WithComments HsModule'
relocateCommentsBeforeEachElement elemGetter elemSetter annGetter annSetter cond =
  everywhereM (mkM f)
  where
    f :: a -> WithComments a
    f x = do
      newElems <- mapM insertCommentsBeforeElement (elemGetter x)
      pure $ elemSetter newElems x
      where
        insertCommentsBeforeElement element
          | elemAnn@EpAnn {} <- annGetter element = do
            newEpa <-
              insertCommentsByPos (cond x element) insertPriorComments elemAnn
            pure $ annSetter newEpa element
          | otherwise = pure element

-- | This function applies the given function to all 'EpAnn's.
applyM ::
     forall a. Typeable a
  => (forall b. EpAnn b -> WithComments (EpAnn b))
  -> (a -> WithComments a)
applyM f
  | App g _ <- typeRep @a
  , Just HRefl <- eqTypeRep g (typeRep @EpAnn) = f
  | otherwise = pure

-- | This function drains comments whose positions satisfy the given
-- predicate and inserts them to the given node using the given inserter.
insertCommentsByPos ::
     (RealSrcSpan -> Bool)
  -> (EpAnnComments -> [LEpaComment] -> EpAnnComments)
  -> EpAnn a
  -> WithComments (EpAnn a)
insertCommentsByPos cond = insertComments (cond . anchor . getLoc)

-- | This function drains comments that satisfy the given predicate and
-- inserts them to the given node using the given inserter.
insertComments ::
     (LEpaComment -> Bool)
  -> (EpAnnComments -> [LEpaComment] -> EpAnnComments)
  -> EpAnn a
  -> WithComments (EpAnn a)
insertComments cond inserter epa@EpAnn {..} = do
  coms <- drainComments cond
  pure $ epa {comments = inserter comments coms}
#if !MIN_VERSION_ghc_lib_parser(9, 10, 1)
insertComments _ _ EpAnnNotUsed = pure EpAnnNotUsed
#endif
-- | This function inserts comments to `priorComments`.
insertPriorComments :: EpAnnComments -> [LEpaComment] -> EpAnnComments
insertPriorComments (EpaComments prior) cs =
  EpaComments (sortCommentsByLocation $ prior ++ cs)
insertPriorComments (EpaCommentsBalanced prior following) cs =
  EpaCommentsBalanced (sortCommentsByLocation $ prior ++ cs) following

-- | This function inserts comments to `followingComments`.
insertFollowingComments :: EpAnnComments -> [LEpaComment] -> EpAnnComments
insertFollowingComments (EpaComments prior) cs = EpaCommentsBalanced prior cs
insertFollowingComments (EpaCommentsBalanced prior following) cs =
  EpaCommentsBalanced prior (sortCommentsByLocation $ following ++ cs)

-- | This function drains comments that satisfy the given predicate.
drainComments :: (LEpaComment -> Bool) -> WithComments [LEpaComment]
drainComments cond = do
  coms <- get
  let (xs, others) = partition cond coms
  put others
  return xs

-- | 'everywhereM' but applies the given function to EPAs in order their
-- positions from backwards.
everywhereMEpAnnsBackwards ::
     Data a
  => (forall b. EpAnn b -> WithComments (EpAnn b))
  -> a
  -> WithComments a
everywhereMEpAnnsBackwards =
  everywhereMEpAnnsInOrder (flip compareEpaByEndPosition)

-- | 'everywhereM' but applies the given function to EPAs in order
-- specified by the given ordering function.
everywhereMEpAnnsInOrder ::
     Data a
  => (forall b c. EpAnn b -> EpAnn c -> Ordering)
  -> (forall b. EpAnn b -> WithComments (EpAnn b))
  -> a
  -> WithComments a
everywhereMEpAnnsInOrder cmp f hm =
  collectEpAnnsInOrderEverywhereMTraverses
    >>= applyFunctionInOrderEpAnnEndPositions
    >>= putModifiedEpAnnsToModule
  where
    collectEpAnnsInOrderEverywhereMTraverses
      -- This function uses 'everywhereM' to collect 'EpAnn's because they
      -- should be collected in the same order as 'putModifiedEpAnnsToModule'
      -- puts them to the AST.
     = reverse <$> execStateT (everywhereM collectEpAnnsST hm) []
      where
        collectEpAnnsST x = do
          modify $ collectEpAnns x
          pure x
        collectEpAnns ::
             forall a. Typeable a
          => a
          -> ([Wrapper] -> [Wrapper])
        collectEpAnns x
          -- If 'a' is 'EpAnn b' ('b' can be any type), wrap 'x' with a 'Wrapper'.
          | App g _ <- typeRep @a
          , Just HRefl <- eqTypeRep g (typeRep @EpAnn) = (Wrapper x :)
          | otherwise = id
    applyFunctionInOrderEpAnnEndPositions ::
         [Wrapper]
      -> WithComments [(Int, Wrapper)] -- ^ The first element of the tuple
                                       -- indicates how many 'Wrapper's were there before 'everywhereM'
                                       -- accessed the second element.
    applyFunctionInOrderEpAnnEndPositions anns =
      forM sorted $ \(i, Wrapper x) -> do
        x' <- f x
        pure (i, Wrapper x')
      where
        indexed = zip [0 :: Int ..] anns
        sorted = sortBy (\(_, Wrapper a) (_, Wrapper b) -> cmp a b) indexed
    putModifiedEpAnnsToModule anns = evalStateT (everywhereM setEpAnn hm) [0 ..]
      where
        setEpAnn ::
             forall a. Typeable a
          => a
          -> StateT [Int] WithComments a
        setEpAnn x
          -- This guard arm checks if 'a' is 'EpAnn b' ('b' can be any type).
          | App g g' <- typeRep @a
          , Just HRefl <- eqTypeRep g (typeRep @EpAnn) = do
            i <- gets head
            modify tail
            case lookup i anns of
              Just (Wrapper y)
                | App _ h <- typeOf y
                , Just HRefl <- eqTypeRep g' h -> pure y
              _ -> error "Unmatches"
          | otherwise = pure x

-- | This function moves comments in `fun_id` of `FunBind` to
-- `mc_fun` of `HsMatchContext`.
--
-- This is a workaround for the issue that `EpAnn`s in `mc_fun` cannot be
-- closed since 9.10.1.
moveCommentsFromFunIdToMcFun :: HsModule' -> WithComments HsModule'
#if MIN_VERSION_ghc_lib_parser(9, 10, 1)
moveCommentsFromFunIdToMcFun = pure . everywhere (mkT f)
  where
    f :: HsBind GhcPs -> HsBind GhcPs
    f fb@FunBind { fun_id = L EpAnn {comments = from, ..} fid
                 , fun_matches = MG {mg_alts = L l alts, ..}
                 } =
      fb
        { fun_id = L EpAnn {comments = EpaCommentsBalanced [] [], ..} fid
        , fun_matches = MG {mg_alts = L l alts', ..}
        }
      where
        alts' =
          fmap
            (\(L l' x) ->
               case x of
                 Match {m_ctxt = FunRhs {mc_fun = L EpAnn {..} fun, ..}, ..} ->
                   L
                     l'
                     Match
                       { m_ctxt =
                           FunRhs
                             {mc_fun = L EpAnn {comments = from, ..} fun, ..}
                       , ..
                       }
                 x'' -> L l' x'')
            alts
    f x = x
#else
moveCommentsFromFunIdToMcFun = pure
#endif
-- | This function sorts comments by its location.
sortCommentsByLocation :: [LEpaComment] -> [LEpaComment]
sortCommentsByLocation = sortBy (compare `on` anchor . getLoc)

-- | This function compares given EPAs by their end positions.
compareEpaByEndPosition :: EpAnn a -> EpAnn b -> Ordering
#if MIN_VERSION_ghc_lib_parser(9, 10, 1)
compareEpaByEndPosition (EpAnn (EpaSpan a) _ _) (EpAnn (EpaSpan b) _ _) =
  case (a, b) of
    (RealSrcSpan a' _, RealSrcSpan b' _) ->
      compare (realSrcSpanEnd a') (realSrcSpanEnd b')
    (UnhelpfulSpan _, UnhelpfulSpan _) -> EQ
    (_, UnhelpfulSpan _) -> GT
    (UnhelpfulSpan _, _) -> LT
compareEpaByEndPosition (EpAnn a _ _) (EpAnn b _ _) =
  case (a, b) of
    (EpaDelta {}, EpaDelta {}) -> EQ
    (_, EpaDelta {}) -> GT
    (EpaDelta {}, _) -> LT
#else
compareEpaByEndPosition (EpAnn a _ _) (EpAnn b _ _) =
  on compare (realSrcSpanEnd . anchor) a b
compareEpaByEndPosition EpAnnNotUsed EpAnnNotUsed = EQ
compareEpaByEndPosition _ EpAnnNotUsed = GT
compareEpaByEndPosition EpAnnNotUsed _ = LT
#endif