content: -- Copyright (c) Facebook, Inc. and its affiliates.
  --
  -- This source code is licensed under the MIT license found in the
  -- LICENSE file in the root directory of this source tree.
  --
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  module Retrie.Replace
    ( replace
    , Replacement(..)
    , Change(..)
    ) where
  
  import Control.Monad.Trans.Class
  import Control.Monad.Writer.Strict
  import Data.Char (isSpace)
  import Data.Generics
  
  import Retrie.ExactPrint
  import Retrie.Expr
  import Retrie.FreeVars
  import Retrie.GHC
  import Retrie.Subst
  import Retrie.Types
  import Retrie.Universe
  import Retrie.Util
  
  ------------------------------------------------------------------------
  
  -- | Specializes 'replaceImpl' to each of the AST types that retrie supports.
  replace
    :: (Data a, MonadIO m) => Context -> a -> TransformT (WriterT Change m) a
  replace c =
    mkM (replaceImpl @(HsExpr GhcPs) c)
      `extM` (replaceImpl @(Stmt GhcPs (LHsExpr GhcPs)) c)
      `extM` (replaceImpl @(HsType GhcPs) c)
      `extM` replacePat c
  
  replacePat :: MonadIO m => Context -> LPat GhcPs -> TransformT (WriterT Change m) (LPat GhcPs)
  -- We need to ensure we have a location available at the top level so we can
  -- transfer annotations. This ensures we don't try to rewrite a naked Pat.
  replacePat c p
    | Just lp <- dLPat p = cLPat <$> replaceImpl c lp
    | otherwise = return p
  
  -- | Generic replacement function. This is the thing that actually runs the
  -- 'Rewriter' carried by the context, instantiates templates, handles parens
  -- and other whitespace bookkeeping, and emits resulting 'Replacement's.
  replaceImpl
    :: forall ast m. (Data ast, ExactPrint ast, Matchable (LocatedA ast), MonadIO m)
    => Context -> LocatedA ast -> TransformT (WriterT Change m) (LocatedA ast)
  replaceImpl c e = do
    let
      -- Prevent rewriting source of the rewrite itself by refusing to
      -- match under a binding of something that appears in the template.
      f result@RewriterResult{..} = result
        { rrTransformer =
            fmap (fmap (check rrOrigin rrQuantifiers)) <$> rrTransformer
        }
      check origin quantifiers match
        | getLocA e `overlaps` origin = NoMatch
        | MatchResult _ Template{..} <- match
        , fvs <- freeVars quantifiers (astA tTemplate)
        , any (`elemFVs` fvs) (ctxtBinders c) = NoMatch
        | otherwise = match
  
    -- We want to match through HsPar so we can make a decision
    -- about whether to keep the parens or not based on the
    -- resulting expression, but we need to know the entry location
    -- of the parens, not the inner expression, so we have to
    -- keep both expressions around.
    match <- runRewriter f c (ctxtRewriter c) (getUnparened e)
  
    case match of
      NoMatch -> return e
      MatchResult sub Template{..} -> do
        -- graft template into target module
        t' <- graftA tTemplate
        -- substitute for quantifiers in grafted template
        r <- subst sub c t'
        -- copy appropriate annotations from old expression to template
        r0 <- addAllAnnsT e r
        -- add parens to template if needed
        res' <- (mkM (parenify c) `extM` parenifyT c `extM` parenifyP c) r0
        -- Make sure the replacement has the same anchor as the thing
        -- being replaced
        let res = transferAnchor e res'
  
        -- prune the resulting expression and log it with location
        orig <- printNoLeadingSpaces <$> pruneA e
        -- orig <- printA' <$> pruneA e
  
        repl <- printNoLeadingSpaces <$> pruneA res
        -- repl <- printA' <$> pruneA r
        -- repl <- printA' <$> pruneA res
        -- repl <- return $ showAst t'
  
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:orig="  [orig]
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:repl="  [repl]
  
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:e="  [showAst e]
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:r="  [showAst r]
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:r0="  [showAst r0]
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:t'=" [showAst t']
        -- lift $ liftIO $ debugPrint Loud "replaceImpl:res=" [showAst res]
  
        let replacement = Replacement (getLocA e) orig repl
        TransformT $ lift $ tell $ Change [replacement] [tImports]
        -- make the actual replacement
        return res'
  
  
  -- | Records a replacement made. In cases where we cannot use ghc-exactprint
  -- to print the resulting AST (e.g. CPP modules), we fall back on splicing
  -- strings. Can also be used by external tools (search, linters, etc).
  data Replacement = Replacement
    { replLocation :: SrcSpan
    , replOriginal :: String
    , replReplacement :: String
    } deriving Show
  
  -- | Used as the writer type during matching to indicate whether any change
  -- to the module should be made.
  data Change = NoChange | Change [Replacement] [AnnotatedImports]
  
  instance Semigroup Change where
    (<>) = mappend
  
  instance Monoid Change where
    mempty = NoChange
    mappend NoChange     other        = other
    mappend other        NoChange     = other
    mappend (Change rs1 is1) (Change rs2 is2) =
      Change (rs1 <> rs2) (is1 <> is2)
  
  -- The location of 'e' accurately points to the first non-space character
  -- of 'e', but when we exactprint 'e', we might get some leading spaces (if
  -- annEntryDelta of the first token is non-zero). This means we can't just
  -- splice in the printed expression at the desired location and call it a day.
  -- Unfortunately, its hard to find the right annEntryDelta (it may not be the
  -- top of the redex) and zero it out. As janky as it seems, its easier to just
  -- drop leading spaces like this.
  printNoLeadingSpaces :: (Data k, ExactPrint k) => Annotated k -> String
  printNoLeadingSpaces = dropWhile isSpace . printA
  