{-# LANGUAGE CPP #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE GADTs  #-}
{-# LANGUAGE RankNTypes  #-}

module Refact.Utils ( -- * Synonyms
                      Module
                    , Stmt
                    , Expr
                    , Decl
                    , Name
                    , Pat
                    , Type
                    , Import
                    , FunBind
                    , AnnKeyMap
                    -- * Monad
                    , M
                    -- * Utility

                    , mergeAnns
                    , modifyAnnKey
                    , replaceAnnKey
                    , toGhcSrcSpan
                    , toGhcSrcSpan'
                    , setSrcSpanFile
                    , findParent
                    ) where

import Language.Haskell.GHC.ExactPrint
import Language.Haskell.GHC.ExactPrint.Types

import Data.Bifunctor (bimap)
import Data.Data
import Data.Map.Strict (Map)

import FastString (FastString)
import SrcLoc
import qualified SrcLoc as GHC
import qualified RdrName as GHC
import qualified ApiAnnotation as GHC
import qualified FastString    as GHC
import qualified GHC hiding (parseModule)

#if __GLASGOW_HASKELL__ >= 810
import GHC.Hs.Expr as GHC hiding (Stmt)
import GHC.Hs.ImpExp
#else
import HsExpr as GHC hiding (Stmt)
import HsImpExp
#endif

import Control.Monad.Trans.State

import qualified Data.Map as Map
import Data.Maybe


import qualified Refact.Types as R

import Data.Generics.Schemes
import Unsafe.Coerce


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

type AnnKeyMap = Map AnnKey [AnnKey]

type Module = (GHC.Located (GHC.HsModule GHC.GhcPs))

type Expr = GHC.Located (GHC.HsExpr GHC.GhcPs)

type Type = GHC.Located (GHC.HsType GHC.GhcPs)

type Decl = GHC.Located (GHC.HsDecl GHC.GhcPs)

type Pat =  GHC.Located (GHC.Pat GHC.GhcPs)

type Name = GHC.Located GHC.RdrName

type Stmt = ExprLStmt GHC.GhcPs

type Import = LImportDecl GHC.GhcPs

type FunBind = HsMatchContext GHC.RdrName

-- | Replaces an old expression with a new expression
--
-- Note that usually, new, inp and parent are all the same.
replace :: AnnKey  -- The thing we are replacing
        -> AnnKey  -- The thing which has the annotations we need for the new thing
        -> AnnKey  -- The thing which is going to be inserted
        -> AnnKey  -- The "parent", the largest thing which has he same SrcSpan
                   -- Usually the same as inp and new
        -> Anns -> Maybe Anns
replace :: AnnKey -> AnnKey -> AnnKey -> AnnKey -> Anns -> Maybe Anns
replace AnnKey
old AnnKey
new AnnKey
inp AnnKey
parent Anns
anns = do
  Annotation
oldan <- AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup AnnKey
old Anns
anns
  Annotation
newan <- AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup AnnKey
new Anns
anns
  DeltaPos
oldDelta <- Annotation -> DeltaPos
annEntryDelta  (Annotation -> DeltaPos) -> Maybe Annotation -> Maybe DeltaPos
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup AnnKey
parent Anns
anns
  Anns -> Maybe Anns
forall (m :: * -> *) a. Monad m => a -> m a
return (Anns -> Maybe Anns) -> Anns -> Maybe Anns
forall a b. (a -> b) -> a -> b
$ AnnKey -> Annotation -> Anns -> Anns
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert AnnKey
inp (DeltaPos -> AnnKey -> Annotation -> Annotation -> Annotation
combine DeltaPos
oldDelta AnnKey
new Annotation
oldan Annotation
newan) Anns
anns

combine :: DeltaPos -> AnnKey -> Annotation -> Annotation -> Annotation
combine :: DeltaPos -> AnnKey -> Annotation -> Annotation -> Annotation
combine DeltaPos
oldDelta AnnKey
newkey Annotation
oldann Annotation
newann =
  Ann :: DeltaPos
-> [(Comment, DeltaPos)]
-> [(Comment, DeltaPos)]
-> [(KeywordId, DeltaPos)]
-> Maybe [SrcSpan]
-> Maybe AnnKey
-> Annotation
Ann { annEntryDelta :: DeltaPos
annEntryDelta = DeltaPos
newEntryDelta
      , annPriorComments :: [(Comment, DeltaPos)]
annPriorComments = Annotation -> [(Comment, DeltaPos)]
annPriorComments Annotation
oldann [(Comment, DeltaPos)]
-> [(Comment, DeltaPos)] -> [(Comment, DeltaPos)]
forall a. [a] -> [a] -> [a]
++ Annotation -> [(Comment, DeltaPos)]
annPriorComments Annotation
newann
      , annFollowingComments :: [(Comment, DeltaPos)]
annFollowingComments = Annotation -> [(Comment, DeltaPos)]
annFollowingComments Annotation
oldann [(Comment, DeltaPos)]
-> [(Comment, DeltaPos)] -> [(Comment, DeltaPos)]
forall a. [a] -> [a] -> [a]
++ Annotation -> [(Comment, DeltaPos)]
annFollowingComments Annotation
newann
      , annsDP :: [(KeywordId, DeltaPos)]
annsDP = [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
removeComma (Annotation -> [(KeywordId, DeltaPos)]
annsDP Annotation
newann) [(KeywordId, DeltaPos)]
-> [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
forall a. [a] -> [a] -> [a]
++ [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
forall b. [(KeywordId, b)] -> [(KeywordId, b)]
extraComma (Annotation -> [(KeywordId, DeltaPos)]
annsDP Annotation
oldann)
      , annSortKey :: Maybe [SrcSpan]
annSortKey = Annotation -> Maybe [SrcSpan]
annSortKey Annotation
newann
      , annCapturedSpan :: Maybe AnnKey
annCapturedSpan = Annotation -> Maybe AnnKey
annCapturedSpan Annotation
newann}
  where
    -- Get rid of structural information when replacing, we assume that the
    -- structural information is already there in the new expression.
    removeComma :: [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
removeComma = ((KeywordId, DeltaPos) -> Bool)
-> [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(KeywordId
kw, DeltaPos
_) -> case KeywordId
kw of
                                         G AnnKeywordId
GHC.AnnComma
                                           | AnnKey SrcSpan
_ (CN [Char]
"ArithSeq") <- AnnKey
newkey -> Bool
True
                                           | Bool
otherwise -> Bool
False
                                         KeywordId
AnnSemiSep -> Bool
False
                                         KeywordId
_ -> Bool
True)

    -- Make sure to keep structural information in the template.
    extraComma :: [(KeywordId, b)] -> [(KeywordId, b)]
extraComma [] = []
    extraComma ([(KeywordId, b)] -> (KeywordId, b)
forall a. [a] -> a
last -> (KeywordId, b)
x) = case (KeywordId, b) -> KeywordId
forall a b. (a, b) -> a
fst (KeywordId, b)
x of
                              G AnnKeywordId
GHC.AnnComma -> [(KeywordId, b)
x]
                              KeywordId
AnnSemiSep -> [(KeywordId, b)
x]
                              G AnnKeywordId
GHC.AnnSemi -> [(KeywordId, b)
x]
                              KeywordId
_ -> []

    -- Keep the same delta if moving onto a new row
    newEntryDelta :: DeltaPos
newEntryDelta | DeltaPos -> Int
deltaRow DeltaPos
oldDelta Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = DeltaPos
oldDelta
                  | Bool
otherwise = Annotation -> DeltaPos
annEntryDelta Annotation
oldann


-- | A parent in this case is an element which has the same SrcSpan
findParent :: Data a => GHC.SrcSpan -> Anns -> a -> Maybe AnnKey
findParent :: SrcSpan -> Anns -> a -> Maybe AnnKey
findParent SrcSpan
ss Anns
as = GenericQ (Maybe AnnKey) -> GenericQ (Maybe AnnKey)
forall u. GenericQ (Maybe u) -> GenericQ (Maybe u)
something (SrcSpan -> Anns -> a -> Maybe AnnKey
forall a. Data a => SrcSpan -> Anns -> a -> Maybe AnnKey
findParentWorker SrcSpan
ss Anns
as)

-- Note that a parent must also have an annotation.
findParentWorker :: forall a . (Data a)
           => GHC.SrcSpan -> Anns -> a -> Maybe AnnKey
findParentWorker :: SrcSpan -> Anns -> a -> Maybe AnnKey
findParentWorker SrcSpan
oldSS Anns
as a
a
  | TyCon
con TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TypeRep -> TyCon
typeRepTyCon (Proxy (Located RdrName) -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy (Located RdrName)
forall k (t :: k). Proxy t
Proxy :: Proxy (GHC.Located GHC.RdrName))) Bool -> Bool -> Bool
&& TypeRep
x TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
== Proxy SrcSpan -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy SrcSpan
forall k (t :: k). Proxy t
Proxy :: Proxy GHC.SrcSpan)
      = if SrcSpan
ss SrcSpan -> SrcSpan -> Bool
forall a. Eq a => a -> a -> Bool
== SrcSpan
oldSS
            Bool -> Bool -> Bool
&& Maybe Annotation -> Bool
forall a. Maybe a -> Bool
isJust (AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (SrcSpan -> AnnConName -> AnnKey
AnnKey SrcSpan
ss AnnConName
cn) Anns
as)
          then AnnKey -> Maybe AnnKey
forall a. a -> Maybe a
Just (AnnKey -> Maybe AnnKey) -> AnnKey -> Maybe AnnKey
forall a b. (a -> b) -> a -> b
$ SrcSpan -> AnnConName -> AnnKey
AnnKey SrcSpan
ss AnnConName
cn
          else Maybe AnnKey
forall a. Maybe a
Nothing
  | Bool
otherwise = Maybe AnnKey
forall a. Maybe a
Nothing
  where
    (TyCon
con, ~[TypeRep
x, TypeRep
_]) = TypeRep -> (TyCon, [TypeRep])
splitTyConApp (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a)
    ss :: GHC.SrcSpan
    ss :: SrcSpan
ss = Int -> (forall d. Data d => d -> SrcSpan) -> a -> SrcSpan
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
0 forall d. Data d => d -> SrcSpan
forall a b. a -> b
unsafeCoerce a
a
    cn :: AnnConName
cn = Int -> (forall d. Data d => d -> AnnConName) -> a -> AnnConName
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
1 ([Char] -> AnnConName
CN ([Char] -> AnnConName) -> (d -> [Char]) -> d -> AnnConName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constr -> [Char]
forall a. Show a => a -> [Char]
show (Constr -> [Char]) -> (d -> Constr) -> d -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. d -> Constr
forall a. Data a => a -> Constr
toConstr) a
a


-- | 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 -> Located old -> Located new -> M (Located new)
modifyAnnKey :: mod -> Located old -> Located new -> M (Located new)
modifyAnnKey mod
m Located old
e1 Located new
e2 = do
  Anns
as <- ((Anns, AnnKeyMap) -> Anns) -> StateT (Anns, AnnKeyMap) IO Anns
forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets (Anns, AnnKeyMap) -> Anns
forall a b. (a, b) -> a
fst
  let parentKey :: AnnKey
parentKey = AnnKey -> Maybe AnnKey -> AnnKey
forall a. a -> Maybe a -> a
fromMaybe (Located new -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located new
e2) (SrcSpan -> Anns -> mod -> Maybe AnnKey
forall a. Data a => SrcSpan -> Anns -> a -> Maybe AnnKey
findParent (Located new -> SrcSpan
forall a. HasSrcSpan a => a -> SrcSpan
getLoc Located new
e2) Anns
as mod
m)
  Located new
e2 Located new -> StateT (Anns, AnnKeyMap) IO () -> M (Located new)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ((Anns, AnnKeyMap) -> (Anns, AnnKeyMap))
-> StateT (Anns, AnnKeyMap) IO ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify
          ( (Anns -> Anns)
-> (AnnKeyMap -> AnnKeyMap)
-> (Anns, AnnKeyMap)
-> (Anns, AnnKeyMap)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap
              ( Located old -> Located new -> Anns -> Anns
forall old new. Located old -> Located new -> Anns -> Anns
recoverBackquotes Located old
e1 Located new
e2
              (Anns -> Anns) -> (Anns -> Anns) -> Anns -> Anns
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AnnKey -> AnnKey -> AnnKey -> AnnKey -> Anns -> Anns
replaceAnnKey (Located old -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located old
e1) (Located new -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located new
e2) (Located new -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located new
e2) AnnKey
parentKey
              )
              (([AnnKey] -> [AnnKey] -> [AnnKey])
-> AnnKey -> [AnnKey] -> AnnKeyMap -> AnnKeyMap
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [AnnKey] -> [AnnKey] -> [AnnKey]
forall a. [a] -> [a] -> [a]
(++) (Located old -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located old
e1) [Located new -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located new
e2])
          )

-- | When the template contains a backquoted substitution variable, but the substitute
-- is not backquoted, we must add the corresponding 'GHC.AnnBackQuote's.
--
-- See tests/examples/Backquotes.hs for an example.
recoverBackquotes :: Located old -> Located new -> Anns -> Anns
recoverBackquotes :: Located old -> Located new -> Anns -> Anns
recoverBackquotes (L SrcSpan
old old
_) (L SrcSpan
new new
_) Anns
anns
  | Just Annotation
annOld <- AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (SrcSpan -> AnnConName -> AnnKey
AnnKey SrcSpan
old ([Char] -> AnnConName
CN [Char]
"Unqual")) Anns
anns
  , ( (G AnnKeywordId
GHC.AnnBackquote, DP (Int
i, Int
j))
    : rest :: [(KeywordId, DeltaPos)]
rest@( (G AnnKeywordId
GHC.AnnVal, DeltaPos
_)
           : (G AnnKeywordId
GHC.AnnBackquote, DeltaPos
_)
           : [(KeywordId, DeltaPos)]
_)
    ) <- Annotation -> [(KeywordId, DeltaPos)]
annsDP Annotation
annOld
  = let f :: Annotation -> Annotation
f Annotation
annNew = case Annotation -> [(KeywordId, DeltaPos)]
annsDP Annotation
annNew of
          [(G AnnKeywordId
GHC.AnnVal, DP (Int
i', Int
j'))] ->
            Annotation
annNew {annsDP :: [(KeywordId, DeltaPos)]
annsDP = (AnnKeywordId -> KeywordId
G AnnKeywordId
GHC.AnnBackquote, (Int, Int) -> DeltaPos
DP (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i', Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j')) (KeywordId, DeltaPos)
-> [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
forall a. a -> [a] -> [a]
: [(KeywordId, DeltaPos)]
rest}
          [(KeywordId, DeltaPos)]
_ -> Annotation
annNew
     in (Annotation -> Annotation) -> AnnKey -> Anns -> Anns
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust Annotation -> Annotation
f (SrcSpan -> AnnConName -> AnnKey
AnnKey SrcSpan
new ([Char] -> AnnConName
CN [Char]
"Unqual")) Anns
anns
  | Bool
otherwise = Anns
anns

-- | Lower level version of @modifyAnnKey@
replaceAnnKey :: AnnKey -> AnnKey -> AnnKey -> AnnKey -> Anns -> Anns
replaceAnnKey :: AnnKey -> AnnKey -> AnnKey -> AnnKey -> Anns -> Anns
replaceAnnKey AnnKey
old AnnKey
new AnnKey
inp AnnKey
deltainfo Anns
a =
  Anns -> Maybe Anns -> Anns
forall a. a -> Maybe a -> a
fromMaybe Anns
a (AnnKey -> AnnKey -> AnnKey -> AnnKey -> Anns -> Maybe Anns
replace AnnKey
old AnnKey
new AnnKey
inp AnnKey
deltainfo Anns
a)

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

-- | Convert a @Refact.Types.SrcSpan@ to a @SrcLoc.SrcSpan@
toGhcSrcSpan' :: FastString -> R.SrcSpan -> SrcSpan
toGhcSrcSpan' :: FastString -> SrcSpan -> SrcSpan
toGhcSrcSpan' FastString
file R.SrcSpan{Int
startLine :: SrcSpan -> Int
startCol :: SrcSpan -> Int
endLine :: SrcSpan -> Int
endCol :: SrcSpan -> Int
endCol :: Int
endLine :: Int
startCol :: Int
startLine :: Int
..} = SrcLoc -> SrcLoc -> SrcSpan
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
mkSrcLoc FastString
file

setSrcSpanFile :: FastString -> SrcSpan -> SrcSpan
setSrcSpanFile :: FastString -> SrcSpan -> SrcSpan
setSrcSpanFile FastString
file SrcSpan
s
  | RealSrcLoc RealSrcLoc
start <- SrcSpan -> SrcLoc
srcSpanStart SrcSpan
s
  , RealSrcLoc RealSrcLoc
end <- SrcSpan -> SrcLoc
srcSpanEnd SrcSpan
s
  = let start' :: SrcLoc
start' = FastString -> Int -> Int -> SrcLoc
mkSrcLoc FastString
file (RealSrcLoc -> Int
srcLocLine RealSrcLoc
start) (RealSrcLoc -> Int
srcLocCol RealSrcLoc
start)
        end' :: SrcLoc
end' = FastString -> Int -> Int -> SrcLoc
mkSrcLoc FastString
file (RealSrcLoc -> Int
srcLocLine RealSrcLoc
end) (RealSrcLoc -> Int
srcLocCol RealSrcLoc
end)
     in SrcLoc -> SrcLoc -> SrcSpan
mkSrcSpan SrcLoc
start' SrcLoc
end'
setSrcSpanFile FastString
_ SrcSpan
s = SrcSpan
s