content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE LambdaCase #-}
  
  -- | 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 (partition, sortBy)
  import GHC.Types.SrcLoc
  import Generics.SYB hiding (GT, typeOf, typeRep)
  import HIndent.GhcLibParserWrapper.GHC.Hs
  import HIndent.GhcLibParserWrapper.GHC.Parser.Annotation
  import HIndent.Pragma
  import HIndent.Pretty.SigBindFamily
  import Type.Reflection
  #if MIN_VERSION_GLASGOW_HASKELL(9, 6, 0, 0)
  import Control.Monad
  #endif
  #if !MIN_VERSION_ghc_lib_parser(9, 12, 1)
  import GHC.Data.Bag
  #endif
  #if MIN_VERSION_ghc_lib_parser(9, 6, 1)
  import Data.Maybe
  #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
  #if !MIN_VERSION_ghc_lib_parser(9, 10, 1)
  relocatePragmas m = pure m
  #endif
  -- | 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 =
              maybe (error "No prior comments") getLoc
                $ listToMaybe
                $ 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 =
        maybe (error "No prior comments.") (anchor . getLoc)
          $ listToMaybe
          $ 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, 12, 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 tcdMeths tcdATs tcdATDefs []
      elemGetter _ = []
      elemSetter xs (L sp (TyClD ext ClassDecl {..})) = L sp (TyClD ext newDecl)
        where
          newDecl =
            ClassDecl
              { tcdSigs = sigs
              , tcdMeths = 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 (filterLBind bindsSigs', filterLSig bindsSigs')
        where
          bindsSigs = mkSortedLSigBindFamilyList sigs binds [] [] []
      addCommentsBeforeEpAnn (L epa@EpAnn { entry = EpaSpan (RealSrcSpan anc _)
                                          , ..
                                          } x) = do
        cs <- get
        let (notAbove, above) =
              partitionAboveNotAbove (sortCommentsByLocation cs) anc
            epa' = epa {comments = insertPriorComments comments above}
        put notAbove
        pure $ L epa' x
      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
  #elif 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
      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 $ epaLocationToRealSrcSpan entry)
          insertFollowingComments
          epa
      isAfter anc comAnc = srcSpanEndLine anc <= srcSpanStartLine comAnc
  #else
  relocateCommentsAfter = everywhereMEpAnnsBackwards f
    where
      f epa@EpAnn {..} =
        insertCommentsByPos
          (isAfter $ epaLocationToRealSrcSpan 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'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  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 = do
            newEpa <-
              insertCommentsByPos
                (cond x element)
                insertPriorComments
                (annGetter element)
            pure $ annSetter newEpa element
  #else
  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
  #endif
  -- | 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 . epaLocationToRealSrcSpan . 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
              get >>= \case
                [] -> error "No comments."
                (i:is) -> do
                  put is
                  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 funann@EpAnn {} fun, ..}} ->
                     L
                       l'
                       x
                         { m_ctxt =
                             FunRhs {mc_fun = L funann {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` epaLocationToRealSrcSpan . 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
  