content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ImpredicativeTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  
  -- | Module preprocessing before pretty-printing.
  module HIndent.ModulePreprocessing
    ( modifyASTForPrettyPrinting
    ) where
  
  import Data.Function
  import Data.List (sortBy)
  import GHC.Hs
  import GHC.Types.SrcLoc
  import Generics.SYB hiding (GT, typeOf, typeRep)
  import HIndent.Fixity
  import HIndent.GhcLibParserWrapper.GHC.Hs
  import HIndent.ModulePreprocessing.CommentRelocation
  import Language.Haskell.GhclibParserEx.Fixity
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  import qualified GHC.Data.Strict as Strict
  import HIndent.GhcLibParserWrapper.GHC.Parser.Annotation
  #else
  import Control.Applicative
  import Data.Maybe
  import Type.Reflection
  #endif
  -- | This function modifies the given module AST for pretty-printing.
  --
  -- Pretty-printing a module without calling this function for it before may
  -- raise an error or not print it correctly.
  modifyASTForPrettyPrinting :: HsModule' -> HsModule'
  modifyASTForPrettyPrinting m = relocateComments (beforeRelocation m) allComments
    where
      beforeRelocation =
        resetListCompRange
          . resetLGRHSEndPositionInModule
          . removeAllDocDs
          . closeEpAnnOfHsFunTy
          . closeEpAnnOfMatchMExt
          . closePlaceHolderEpAnns
          . closeEpAnnOfFunBindFunId
          . resetModuleNameColumn
          . replaceAllNotUsedAnns
          . removeComments
          . sortExprLStmt
          . fixFixities
      allComments = listify (not . isEofComment . ac_tok . unLoc) m
  
  -- | This function modifies the given module AST to apply fixities of infix
  -- operators defined in the 'base' package.
  fixFixities :: HsModule' -> HsModule'
  fixFixities = applyFixities fixities
  
  -- | This function modifies the range of `HsDo` with `ListComp` so that it
  -- includes the whole list comprehension.
  --
  -- This function is necessary for `ghc-lib-parser>=9.10.1<9.12.1` because `HsDo`
  -- no longer includes brackets of list comprehensions in its range.
  resetListCompRange :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 12, 1)
  resetListCompRange = everywhere (mkT resetListCompRange')
    where
      resetListCompRange' :: HsExpr GhcPs -> HsExpr GhcPs
      resetListCompRange' (HsDo al@AnnList {al_brackets = ListSquare (EpTok (EpaSpan (RealSrcSpan open _))) (EpTok (EpaSpan (RealSrcSpan close _)))} ListComp (L EpAnn {..} xs)) =
        HsDo
          al
          ListComp
          (L EpAnn
               { entry =
                   EpaSpan
                     $ RealSrcSpan
                         (mkRealSrcSpan
                            (realSrcSpanStart open)
                            (realSrcSpanEnd close))
                         Strict.Nothing
               , ..
               }
             xs)
      resetListCompRange' x = x
  #elif MIN_VERSION_ghc_lib_parser(9, 10, 1)
  resetListCompRange = everywhere (mkT resetListCompRange')
    where
      resetListCompRange' :: HsExpr GhcPs -> HsExpr GhcPs
      resetListCompRange' (HsDo al@AnnList { al_open = Just (AddEpAnn _ (EpaSpan (RealSrcSpan open _)))
                                           , al_close = Just (AddEpAnn _ (EpaSpan (RealSrcSpan close _)))
                                           } ListComp (L EpAnn {..} xs)) =
        HsDo
          al
          ListComp
          (L EpAnn
               { entry =
                   EpaSpan
                     $ RealSrcSpan
                         (mkRealSrcSpan
                            (realSrcSpanStart open)
                            (realSrcSpanEnd close))
                         Strict.Nothing
               , ..
               }
             xs)
      resetListCompRange' x = x
  #else
  resetListCompRange = id
  #endif
  -- | This function sets an 'LGRHS's end position to the end position of the
  -- last RHS in the 'grhssGRHSs'.
  --
  -- The source span of an 'L?GRHS' contains the 'where' keyword, which
  -- locates comments in the wrong position in the process of comment
  -- relocation. This function prevents it by fixing the 'L?GRHS''s source
  -- span.
  resetLGRHSEndPositionInModule :: HsModule' -> HsModule'
  resetLGRHSEndPositionInModule = everywhere (mkT resetLGRHSEndPosition)
  
  -- | This function sorts lists of statements in order their positions.
  --
  -- For example, the last element of 'HsDo' of 'HsExpr' is the element
  -- before a bar, and the elements are not sorted by their locations. This
  -- function fixes the orderings.
  sortExprLStmt :: HsModule' -> HsModule'
  sortExprLStmt m@HsModule {hsmodDecls = xs} = m {hsmodDecls = sorted}
    where
      sorted = everywhere (mkT sortByLoc) xs
      sortByLoc :: [ExprLStmt GhcPs] -> [ExprLStmt GhcPs]
      sortByLoc = sortBy (compare `on` srcSpanToRealSrcSpan . locA . getLoc)
  
  -- | This function removes all comments from the given module not to
  -- duplicate them on comment relocation.
  removeComments :: HsModule' -> HsModule'
  removeComments = everywhere (mkT $ const emptyComments)
  
  -- | This function replaces all 'EpAnnNotUsed's in 'SrcSpanAnn''s with
  -- 'EpAnn's to make it possible to locate comments on them.
  replaceAllNotUsedAnns :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  -- 'EpAnnNotUsed' is not used since 9.10.1.
  replaceAllNotUsedAnns = id
  #else
  replaceAllNotUsedAnns = everywhere app
    where
      app ::
           forall a. Data a
        => (a -> a)
      app sp
        | App g (App y z) <- typeRep @a
        , Just HRefl <- eqTypeRep g (typeRep @SrcSpanAnn')
        , Just HRefl <- eqTypeRep y (typeRep @EpAnn) =
          fromMaybe sp $ do
            let try :: Typeable b => b -> Maybe a
                try ann = do
                  HRefl <- eqTypeRep (typeOf ann) z
                  pure sp {ann = EpAnn (spanAsAnchor $ locA sp) ann emptyComments}
            try emptyListItem
              <|> try emptyList
              <|> try emptyPragma
              <|> try emptyContext
              <|> try emptyNameAnn
              <|> try NoEpAnns
      app x = x
      emptyListItem = AnnListItem []
      emptyList = AnnList Nothing Nothing Nothing [] []
      emptyPragma = AnnPragma emptyAddEpAnn emptyAddEpAnn []
      emptyContext = AnnContext Nothing [] []
      emptyNameAnn = NameAnnTrailing []
      emptyAddEpAnn = AddEpAnn AnnAnyclass emptyEpaLocation
      emptyEpaLocation = EpaDelta (SameLine 0) []
  #endif
  -- | This function sets the start column of 'hsmodName' of the given
  -- 'HsModule' to 1 to correctly locate comments above the module name.
  resetModuleNameColumn :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  resetModuleNameColumn m@HsModule {hsmodName = Just (L epa@EpAnn {..} name)} =
    m {hsmodName = Just (L newAnn name)}
    where
      newAnn = epa {entry = realSpanAsAnchor newSpan}
      newSpan =
        mkRealSrcSpan
          (mkRealSrcLoc (srcSpanFile anc) (srcSpanStartLine anc) 1)
          (realSrcSpanEnd anc)
      anc =
        case entry of
          EpaSpan (RealSrcSpan a _) -> a
          _ -> error "resetModuleNameColumn: not a RealSrcSpan"
  #else
  resetModuleNameColumn m@HsModule {hsmodName = Just (L (SrcSpanAnn epa@EpAnn {..} sp) name)} =
    m {hsmodName = Just (L (SrcSpanAnn newAnn sp) name)}
    where
      newAnn = epa {entry = realSpanAsAnchor newSpan}
      newSpan =
        mkRealSrcSpan
          (mkRealSrcLoc (srcSpanFile anc) (srcSpanStartLine anc) 1)
          (realSrcSpanEnd anc)
      anc = anchor entry
  #endif
  resetModuleNameColumn m = m
  
  -- | This function replaces the 'EpAnn' of 'fun_id' in 'FunBind' with
  -- 'EpAnnNotUsed'.
  --
  -- The 'fun_id' contains the function's name. However, 'FunRhs' of 'Match'
  -- also contains the name, and we use the latter one. This function
  -- prevents comments from being located in 'fun_id'.
  closeEpAnnOfFunBindFunId :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  -- TODO: 'EpAnnNotUsed' is not used since 9.10.1. We need to find another
  -- way to close 'EpAnn's.
  closeEpAnnOfFunBindFunId = id
  #else
  closeEpAnnOfFunBindFunId = everywhere (mkT closeEpAnn)
    where
      closeEpAnn :: HsBind GhcPs -> HsBind GhcPs
      closeEpAnn bind@FunBind {fun_id = (L (SrcSpanAnn _ l) name)} =
        bind {fun_id = L (SrcSpanAnn EpAnnNotUsed l) name}
      closeEpAnn x = x
  #endif
  -- | This function replaces the 'EpAnn' of 'm_ext' in 'Match' with
  -- 'EpAnnNotUsed.
  --
  -- The field contains the annotation of the match LHS. However, the same
  -- information is also stored inside the 'Match'. This function removes the
  -- duplication not to locate comments on a wrong point.
  closeEpAnnOfMatchMExt :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  -- TODO: 'EpAnnNotUsed' is not used since 9.10.1. We need to find another
  -- way to close 'EpAnn's.
  closeEpAnnOfMatchMExt = id
  #else
  closeEpAnnOfMatchMExt = everywhere closeEpAnn
    where
      closeEpAnn ::
           forall a. Typeable a
        => a
        -> a
      closeEpAnn x
        | App (App g h) _ <- typeRep @a
        , Just HRefl <- eqTypeRep g (typeRep @Match)
        , Just HRefl <- eqTypeRep h (typeRep @GhcPs) = x {m_ext = EpAnnNotUsed}
        | otherwise = x
  #endif
  -- | This function replaces the 'EpAnn' of the first argument of 'HsFunTy'
  -- of 'HsType'.
  --
  -- 'HsFunTy' should not have any comments. Instead, its LHS and RHS should
  -- have them.
  closeEpAnnOfHsFunTy :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  -- TODO: 'EpAnnNotUsed' is not used since 9.10.1. We need to find another
  -- way to close 'EpAnn's.
  closeEpAnnOfHsFunTy = id
  #else
  closeEpAnnOfHsFunTy = everywhere (mkT closeEpAnn)
    where
      closeEpAnn :: HsType GhcPs -> HsType GhcPs
      closeEpAnn (HsFunTy _ p l r) = HsFunTy EpAnnNotUsed p l r
      closeEpAnn x = x
  #endif
  -- | This function replaces all 'EpAnn's that contain placeholder anchors
  -- to locate comments correctly. A placeholder anchor is an anchor pointing
  -- on (-1, -1).
  closePlaceHolderEpAnns :: HsModule' -> HsModule'
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  -- TODO: 'EpAnnNotUsed' is not used since 9.10.1. We need to find another
  -- way to close 'EpAnn's.
  closePlaceHolderEpAnns = id
  #else
  closePlaceHolderEpAnns = everywhere closeEpAnn
    where
      closeEpAnn ::
           forall a. Typeable a
        => a
        -> a
      closeEpAnn x
        | App g _ <- typeRep @a
        , Just HRefl <- eqTypeRep g (typeRep @EpAnn)
        , (EpAnn (Anchor sp _) _ _) <- x
        , srcSpanEndLine sp == -1 && srcSpanEndCol sp == -1 = EpAnnNotUsed
        | otherwise = x
  #endif
  -- | This function removes all 'DocD's from the given module. They have
  -- haddocks, but the same information is stored in 'EpaCommentTok's. Thus,
  -- we need to remove the duplication.
  removeAllDocDs :: HsModule' -> HsModule'
  removeAllDocDs x@HsModule {hsmodDecls = decls} =
    x {hsmodDecls = filter (not . isDocD . unLoc) decls}
    where
      isDocD DocD {} = True
      isDocD _ = False
  
  -- | This function sets the position of the given 'LGRHS' to the end
  -- position of the last RHS in it.
  --
  -- See the documentation of 'resetLGRHSEndPositionInModule' for the reason.
  resetLGRHSEndPosition ::
       LGRHS GhcPs (LHsExpr GhcPs) -> LGRHS GhcPs (LHsExpr GhcPs)
  #if MIN_VERSION_ghc_lib_parser(9, 10, 1)
  resetLGRHSEndPosition (L locAnn (GRHS ext@EpAnn {..} stmt body)) =
    let lastPosition =
          maximum
            $ realSrcSpanEnd . epaLocationToRealSrcSpan
                <$> listify collectEpaLocation' body
        newSpan =
          mkRealSrcSpan
            (realSrcSpanStart $ epaLocationToRealSrcSpan entry)
            lastPosition
        newLocAnn = locAnn {entry = realSpanAsAnchor newSpan}
        newAnn = ext {entry = realSpanAsAnchor newSpan}
     in L newLocAnn (GRHS newAnn stmt body)
    where
      collectEpaLocation' :: EpaLocation -> Bool
      collectEpaLocation' (EpaSpan RealSrcSpan {}) = True
      collectEpaLocation' _ = False
  #elif MIN_VERSION_ghc_lib_parser(9, 4, 1)
  resetLGRHSEndPosition (L (SrcSpanAnn locAnn@EpAnn {} sp) (GRHS ext@EpAnn {..} stmt body)) =
    let lastPosition =
          maximum $ realSrcSpanEnd . anchor <$> listify collectAnchor body
        newSpan = mkRealSrcSpan (realSrcSpanStart $ anchor entry) lastPosition
        newLocAnn = locAnn {entry = realSpanAsAnchor newSpan}
        newAnn = ext {entry = realSpanAsAnchor newSpan}
     in L (SrcSpanAnn newLocAnn sp) (GRHS newAnn stmt body)
    where
      collectAnchor :: Anchor -> Bool
      collectAnchor _ = True
  resetLGRHSEndPosition x = x
  #else
  resetLGRHSEndPosition (L _ (GRHS ext@EpAnn {..} stmt body)) =
    let lastPosition =
          maximum $ realSrcSpanEnd . anchor <$> listify collectAnchor body
        newSpan = mkRealSrcSpan (realSrcSpanStart $ anchor entry) lastPosition
        newLoc = RealSrcSpan newSpan Nothing
        newAnn = ext {entry = realSpanAsAnchor newSpan}
     in L newLoc (GRHS newAnn stmt body)
    where
      collectAnchor :: Anchor -> Bool
      collectAnchor _ = True
  resetLGRHSEndPosition x = x
  #endif
  isEofComment :: EpaCommentTok -> Bool
  #if !MIN_VERSION_ghc_lib_parser(9, 10, 1)
  isEofComment EpaEofComment = True
  #endif
  isEofComment _ = False
  