{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}

module Refact.Utils
  ( -- * Synonyms
    Module,
    Stmt,
    Expr,
    Decl,
    Name,
    Pat,
    Type,
    Import,
    FunBind,
    pattern RealSrcLoc',
    pattern RealSrcSpan',

    -- * Monad
    M,

    -- * Utility
    modifyAnnKey,
    getAnnSpan,
    getAnnSpanA,
    toGhcSrcSpan,
    toGhcSrcSpan',
    annSpanToSrcSpan,
    srcSpanToAnnSpan,
    setAnnSpanFile,
    setSrcSpanFile,
    setRealSrcSpanFile,
  )
where

import Control.Monad.Trans.State.Strict (StateT)
import Data.Data
  ( Data (),
  )
import Data.Generics (everywhere, mkT)
import Data.Typeable
import qualified GHC
import Language.Haskell.GHC.ExactPrint hiding (transferEntryDP)
import Refact.Compat
  ( AnnSpan,
    FastString,
    FunBind,
    Module,
    annSpanToSrcSpan,
    mkFastString,
    setAnnSpanFile,
    setRealSrcSpanFile,
    setSrcSpanFile,
    srcSpanToAnnSpan,
    pattern RealSrcLoc',
    pattern RealSrcSpan', AnnConstraint,
    transferEntryDP
  )
import qualified Refact.Types as R


-- Types
-- type M a = StateT (Anns, AnnKeyMap) IO a
type M a = StateT () IO a

type Expr = GHC.LHsExpr GHC.GhcPs

type Type = GHC.LHsType GHC.GhcPs

type Decl = GHC.LHsDecl GHC.GhcPs

type Pat = GHC.LPat GHC.GhcPs

type Name = GHC.LocatedN GHC.RdrName

type Stmt = GHC.ExprLStmt GHC.GhcPs

type Import = GHC.LImportDecl GHC.GhcPs

getAnnSpanA :: forall an a. GHC.LocatedAn an a -> AnnSpan
getAnnSpanA :: forall an a. LocatedAn an a -> AnnSpan
getAnnSpanA = SrcSpan -> AnnSpan
srcSpanToAnnSpan (SrcSpan -> AnnSpan)
-> (LocatedAn an a -> SrcSpan) -> LocatedAn an a -> AnnSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocatedAn an a -> SrcSpan
forall a e. GenLocated (SrcSpanAnn' a) e -> SrcSpan
GHC.getLocA

getAnnSpan :: forall a. GHC.Located a -> AnnSpan
getAnnSpan :: forall a. Located a -> AnnSpan
getAnnSpan = SrcSpan -> AnnSpan
srcSpanToAnnSpan (SrcSpan -> AnnSpan)
-> (Located a -> SrcSpan) -> Located a -> AnnSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Located a -> SrcSpan
forall l e. GenLocated l e -> l
GHC.getLoc

-- | Perform the necessary adjustments to annotations when replacing
-- one Located thing with another Located thing.
--
-- For example, this function will ensure the correct relative position and
-- make sure that any trailing semi colons or commas are transferred.
-- modifyAnnKey ::
--   (Data old, Data new, Data mod) =>
--   mod ->
--   GHC.Located old ->
--   GHC.Located new ->
--   M (GHC.Located new)
modifyAnnKey ::
  (Data mod, Data t, Data old, Data new, AnnConstraint t, Typeable t) =>
  mod ->
  GHC.LocatedAn t old ->
  GHC.LocatedAn t new ->
  M (GHC.LocatedAn t new)
modifyAnnKey :: forall mod t old new.
(Data mod, Data t, Data old, Data new, AnnConstraint t,
 Typeable t) =>
mod -> LocatedAn t old -> LocatedAn t new -> M (LocatedAn t new)
modifyAnnKey mod
_m LocatedAn t old
e1 LocatedAn t new
e2 = do
  -- liftIO $ putStrLn $ "modifyAnnKey:e1" ++ showAst e1
  -- liftIO $ putStrLn $ "modifyAnnKey:e2" ++ showAst e2
  let e2_0 :: LocatedAn t new
e2_0 = LocatedAn t old -> LocatedAn t new -> LocatedAn t new
forall t old new.
(Data t, Data old, Data new, AnnConstraint t, Typeable t) =>
LocatedAn t old -> LocatedAn t new -> LocatedAn t new
handleBackquotes LocatedAn t old
e1 LocatedAn t new
e2
  -- liftIO $ putStrLn $ "modifyAnnKey:e2_0" ++ showAst e2_0
  let (LocatedAn t new
e2', Int
_, [String]
_) = Transform (LocatedAn t new) -> (LocatedAn t new, Int, [String])
forall a. Transform a -> (a, Int, [String])
runTransform (Transform (LocatedAn t new) -> (LocatedAn t new, Int, [String]))
-> Transform (LocatedAn t new) -> (LocatedAn t new, Int, [String])
forall a b. (a -> b) -> a -> b
$ LocatedAn t old -> LocatedAn t new -> Transform (LocatedAn t new)
forall (m :: * -> *) t2 t1 a b.
(Monad m, Monoid t2, Typeable t1, Typeable t2) =>
LocatedAn t1 a -> LocatedAn t2 b -> TransformT m (LocatedAn t2 b)
transferEntryDP LocatedAn t old
e1 LocatedAn t new
e2_0
  -- liftIO $ putStrLn $ "modifyAnnKey:e2'" ++ showAst e2'
  LocatedAn t new -> M (LocatedAn t new)
forall a. a -> StateT () IO a
forall (m :: * -> *) a. Monad m => a -> m a
return LocatedAn t new
e2'

-- | This function handles backquotes in two scenarios:
--
--     1. When the template contains a backquoted substitution variable, but the substitute
--        is not backquoted, we must add the corresponding 'GHC.NameBackquotes'. See
--        tests/examples/Backquotes.hs for an example.
--     2. When the template contains a substitution variable without backquote, and the
--        substitute is backquoted, we remove the 'GHC.NameBackquotes' annotation. See
--        tests/examples/Uncurry.hs for an example.
--        N.B.: this is not always correct, since it is possible that the refactoring output
--        should keep the backquotes, but currently no test case fails because of it.
handleBackquotes ::
  forall t old new.
  (Data t, Data old, Data new, AnnConstraint t, Typeable t) =>
  GHC.LocatedAn t old ->
  GHC.LocatedAn t new ->
  GHC.LocatedAn t new
#if MIN_VERSION_ghc(9,12,0)
handleBackquotes old new@(GHC.L loc _) =
  everywhere (mkT update) new
  where
    update :: GHC.LHsExpr GHC.GhcPs -> GHC.LHsExpr GHC.GhcPs
    update (GHC.L l (GHC.HsVar x (GHC.L ln n))) = GHC.L l (GHC.HsVar x (GHC.L ln' n))
      where
        ln' =
          if GHC.locA l == GHC.locA loc
            then case cast old :: Maybe (GHC.LHsExpr GHC.GhcPs) of
              Just (GHC.L _ (GHC.HsVar _ (GHC.L (GHC.EpAnn _ ann _) _)))
                -- scenario 1
                | GHC.NameAnn (GHC.NameBackquotes _ _) _ _ <- ann ->
                  case ln of
                    (GHC.EpAnn a _ cs) -> (GHC.EpAnn a ann cs)
                -- scenario 2
                | (GHC.EpAnn a ann' cs) <- ln,
                  GHC.NameAnn (GHC.NameBackquotes _ _) _ _ <- ann' ->
                  (GHC.EpAnn a ann cs)
              Just _ -> ln
              Nothing -> ln
            else ln
    update x = x
#else
handleBackquotes :: forall t old new.
(Data t, Data old, Data new, AnnConstraint t, Typeable t) =>
LocatedAn t old -> LocatedAn t new -> LocatedAn t new
handleBackquotes LocatedAn t old
old new :: LocatedAn t new
new@(GHC.L SrcAnn t
loc new
_) =
  (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhere ((GenLocated SrcSpanAnnA (HsExpr GhcPs)
 -> GenLocated SrcSpanAnnA (HsExpr GhcPs))
-> a -> a
forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
mkT LHsExpr GhcPs -> LHsExpr GhcPs
GenLocated SrcSpanAnnA (HsExpr GhcPs)
-> GenLocated SrcSpanAnnA (HsExpr GhcPs)
update) LocatedAn t new
new
  where
    update :: GHC.LHsExpr GHC.GhcPs -> GHC.LHsExpr GHC.GhcPs
    update :: LHsExpr GhcPs -> LHsExpr GhcPs
update (GHC.L SrcSpanAnnA
l (GHC.HsVar XVar GhcPs
x (GHC.L SrcSpanAnn' (EpAnn NameAnn)
ln RdrName
n))) = SrcSpanAnnA
-> HsExpr GhcPs -> GenLocated SrcSpanAnnA (HsExpr GhcPs)
forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
l (XVar GhcPs -> LIdP GhcPs -> HsExpr GhcPs
forall p. XVar p -> LIdP p -> HsExpr p
GHC.HsVar XVar GhcPs
x (SrcSpanAnn' (EpAnn NameAnn)
-> RdrName -> GenLocated (SrcSpanAnn' (EpAnn NameAnn)) RdrName
forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnn' (EpAnn NameAnn)
ln' RdrName
n))
      where
        ln' :: SrcSpanAnn' (EpAnn NameAnn)
ln' =
          if SrcSpanAnnA -> SrcSpan
forall a. SrcSpanAnn' a -> SrcSpan
GHC.locA SrcSpanAnnA
l SrcSpan -> SrcSpan -> Bool
forall a. Eq a => a -> a -> Bool
== SrcAnn t -> SrcSpan
forall a. SrcSpanAnn' a -> SrcSpan
GHC.locA SrcAnn t
loc
            then case LocatedAn t old -> Maybe (GenLocated SrcSpanAnnA (HsExpr GhcPs))
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast LocatedAn t old
old :: Maybe (GHC.LHsExpr GHC.GhcPs) of
              Just (GHC.L SrcSpanAnnA
_ (GHC.HsVar XVar GhcPs
_ (GHC.L (GHC.SrcSpanAnn (GHC.EpAnn Anchor
_ NameAnn
ann EpAnnComments
_) SrcSpan
_) RdrName
_)))
                -- scenario 1
                | GHC.NameAnn NameAdornment
GHC.NameBackquotes EpaLocation
_ EpaLocation
_ EpaLocation
_ [TrailingAnn]
_ <- NameAnn
ann ->
                  case SrcSpanAnn' (EpAnn NameAnn)
ln of
                    (GHC.SrcSpanAnn (GHC.EpAnn Anchor
a NameAnn
_ EpAnnComments
cs) SrcSpan
ll) -> EpAnn NameAnn -> SrcSpan -> SrcSpanAnn' (EpAnn NameAnn)
forall a. a -> SrcSpan -> SrcSpanAnn' a
GHC.SrcSpanAnn (Anchor -> NameAnn -> EpAnnComments -> EpAnn NameAnn
forall ann. Anchor -> ann -> EpAnnComments -> EpAnn ann
GHC.EpAnn Anchor
a NameAnn
ann EpAnnComments
cs) SrcSpan
ll
                    (GHC.SrcSpanAnn EpAnn NameAnn
GHC.EpAnnNotUsed SrcSpan
ll) ->
                      EpAnn NameAnn -> SrcSpan -> SrcSpanAnn' (EpAnn NameAnn)
forall a. a -> SrcSpan -> SrcSpanAnn' a
GHC.SrcSpanAnn (Anchor -> NameAnn -> EpAnnComments -> EpAnn NameAnn
forall ann. Anchor -> ann -> EpAnnComments -> EpAnn ann
GHC.EpAnn (SrcSpan -> Anchor
GHC.spanAsAnchor SrcSpan
ll) NameAnn
ann EpAnnComments
GHC.emptyComments) SrcSpan
ll
                -- scenario 2
                | GHC.SrcSpanAnn (GHC.EpAnn Anchor
a NameAnn
ann' EpAnnComments
cs) SrcSpan
ll <- SrcSpanAnn' (EpAnn NameAnn)
ln,
                  GHC.NameAnn NameAdornment
GHC.NameBackquotes EpaLocation
_ EpaLocation
_ EpaLocation
_ [TrailingAnn]
_ <- NameAnn
ann' ->
                  EpAnn NameAnn -> SrcSpan -> SrcSpanAnn' (EpAnn NameAnn)
forall a. a -> SrcSpan -> SrcSpanAnn' a
GHC.SrcSpanAnn (Anchor -> NameAnn -> EpAnnComments -> EpAnn NameAnn
forall ann. Anchor -> ann -> EpAnnComments -> EpAnn ann
GHC.EpAnn Anchor
a NameAnn
ann EpAnnComments
cs) SrcSpan
ll
              Just LHsExpr GhcPs
_ -> SrcSpanAnn' (EpAnn NameAnn)
ln
              Maybe (LHsExpr GhcPs)
Nothing -> SrcSpanAnn' (EpAnn NameAnn)
ln
            else SrcSpanAnn' (EpAnn NameAnn)
ln
    update LHsExpr GhcPs
x = LHsExpr GhcPs
x
#endif

-- | Convert a @Refact.Types.SrcSpan@ to a @SrcLoc.SrcSpan@
toGhcSrcSpan :: FilePath -> R.SrcSpan -> GHC.SrcSpan
toGhcSrcSpan :: String -> SrcSpan -> SrcSpan
toGhcSrcSpan = FastString -> SrcSpan -> SrcSpan
toGhcSrcSpan' (FastString -> SrcSpan -> SrcSpan)
-> (String -> FastString) -> String -> SrcSpan -> SrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> FastString
mkFastString

-- | Convert a @Refact.Types.SrcSpan@ to a @SrcLoc.SrcSpan@
toGhcSrcSpan' :: FastString -> R.SrcSpan -> GHC.SrcSpan
toGhcSrcSpan' :: FastString -> SrcSpan -> SrcSpan
toGhcSrcSpan' FastString
file R.SrcSpan {Int
startLine :: Int
startCol :: Int
endLine :: Int
endCol :: Int
startLine :: SrcSpan -> Int
startCol :: SrcSpan -> Int
endLine :: SrcSpan -> Int
endCol :: SrcSpan -> Int
..} = SrcLoc -> SrcLoc -> SrcSpan
GHC.mkSrcSpan (Int -> Int -> SrcLoc
f Int
startLine Int
startCol) (Int -> Int -> SrcLoc
f Int
endLine Int
endCol)
  where
    f :: Int -> Int -> SrcLoc
f = FastString -> Int -> Int -> SrcLoc
GHC.mkSrcLoc FastString
file