Rename.hs

{-# LANGUAGE CPP               #-}
{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE GADTs             #-}
{-# LANGUAGE OverloadedLabels  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# LANGUAGE LambdaCase        #-}

module Ide.Plugin.Rename (descriptor, Log) where

import           Control.Lens                          ((^.))
import           Control.Monad
import           Control.Monad.Except                  (ExceptT, throwError)
import           Control.Monad.IO.Class                (MonadIO, liftIO)
import           Control.Monad.Trans.Class             (lift)
import           Data.Either                           (rights)
import           Data.Foldable                         (fold)
import           Data.Generics
import           Data.Hashable
import           Data.HashSet                          (HashSet)
import qualified Data.HashSet                          as HS
import           Data.List.NonEmpty                    (NonEmpty ((:|)),
                                                        groupWith)
import qualified Data.Map                              as M
import           Data.Maybe
import           Data.Mod.Word
import qualified Data.Text                             as T
import           Development.IDE                       (Recorder, WithPriority,
                                                        usePropertyAction)
import           Development.IDE.Core.FileStore        (getVersionedTextDoc)
import           Development.IDE.Core.PluginUtils
import           Development.IDE.Core.RuleTypes
import           Development.IDE.Core.Service          hiding (Log)
import           Development.IDE.Core.Shake            hiding (Log)
import           Development.IDE.GHC.Compat
import           Development.IDE.GHC.Compat.ExactPrint
import           Development.IDE.GHC.Error
import           Development.IDE.GHC.ExactPrint        hiding (Log)
import qualified Development.IDE.GHC.ExactPrint        as E
import           Development.IDE.Plugin.CodeAction
import           Development.IDE.Spans.AtPoint
import           Development.IDE.Types.Location
import           GHC.Iface.Ext.Types                   (HieAST (..),
                                                        HieASTs (..),
                                                        NodeOrigin (..),
                                                        SourcedNodeInfo (..))
import           GHC.Iface.Ext.Utils                   (generateReferencesMap)
import           HieDb                                 ((:.) (..))
import           HieDb.Query
import           HieDb.Types                           (RefRow (refIsGenerated))
import           Ide.Logger                            (Pretty (..),
                                                        cmapWithPrio)
import           Ide.Plugin.Error
import           Ide.Plugin.Properties
import qualified Ide.Plugin.Rename.ImportAlias         as ImportAlias
import qualified Ide.Plugin.Rename.ModuleName          as ModuleName
import           Ide.PluginUtils
import           Ide.Types
import qualified Language.LSP.Protocol.Lens            as L
import           Language.LSP.Protocol.Message
import           Language.LSP.Protocol.Types
import qualified Language.LSP.VFS                      as VFS

instance Hashable (Mod a) where hash n = hash (unMod n)

data Log
    = LogExactPrint E.Log
    | LogModuleName ModuleName.Log

instance Pretty Log where
    pretty = \ case
        LogExactPrint msg -> pretty msg
        LogModuleName msg -> pretty msg

descriptor :: Recorder (WithPriority Log) -> PluginId -> PluginDescriptor IdeState
descriptor recorder pluginId = mkExactprintPluginDescriptor exactPrintRecorder $
    (defaultPluginDescriptor pluginId "Provides renaming of Haskell identifiers and module names")
        { pluginHandlers = mconcat
              [ mkPluginHandler SMethod_TextDocumentRename renameProvider
              , mkPluginHandler SMethod_TextDocumentPrepareRename prepareRenameProvider
              , mkPluginHandler SMethod_TextDocumentCodeLens (ModuleName.codeLens moduleNameRecorder)
              ]
        , pluginCommands = [PluginCommand ModuleName.updateModuleNameCommand "Set name of module to match with file path" (ModuleName.command moduleNameRecorder)]
        , pluginConfigDescriptor = defaultConfigDescriptor
            { configCustomConfig = mkCustomConfig properties }
        }
    where
        exactPrintRecorder = cmapWithPrio LogExactPrint recorder
        moduleNameRecorder = cmapWithPrio LogModuleName recorder

prepareRenameProvider :: PluginMethodHandler IdeState Method_TextDocumentPrepareRename
prepareRenameProvider state _pluginId (PrepareRenameParams (TextDocumentIdentifier uri) lspPos _progressToken) = do
    nfp <- getNormalizedFilePathE uri
    codePointPos <- getCodePointPosition state nfp lspPos
    maybeParsed <- ImportAlias.getParsedModuleStale state nfp
    case maybeParsed of
        Nothing -> throwError $ PluginInternalError
            "The module hasn’t yet been parsed. Please wait for indexing to complete and try again."
        Just parsed -> do
            let hsModule = unLoc $ pm_parsed_source parsed
                exports = hsmodExports hsModule
                imports = hsmodImports hsModule
                hsDecls = hsmodDecls hsModule
            maybeAlias <- ImportAlias.resolveAliasAtPos
                getNamesAtPos state nfp lspPos codePointPos exports imports hsDecls
            case maybeAlias of
                Just (lspRange, _importAlias) ->
                    pure $ InL $ PrepareRenameResult $ InL $ lspRange
                Nothing -> do
                    HAR{hieAst} <- handleGetHieAst state nfp
                    let spansWithNamesUnderCursor =
                            [ srcSpan
                            | (names, srcSpan) <- getNamesSpansAtPoint' hieAst lspPos
                            , not (null names)]
                    -- When this handler says that rename is invalid, VSCode shows "The element can't be renamed"
                    -- and doesn't even allow you to create full rename request.
                    -- This handler deliberately approximates "things that definitely can't be renamed"
                    -- to mean "there is no Name at given position" (in which case
                    -- `spansWithNamesUnderCursor` would be empty).
                    --
                    -- In particular it allows some cases through (e.g. cross-module renames),
                    -- so that the full rename handler can give more informative error about them.
                    pure $ case spansWithNamesUnderCursor of
                        [] -> InR Null
                        srcSpan : _ -> InL $ PrepareRenameResult $ InL (realSrcSpanToRange srcSpan)

renameProvider :: PluginMethodHandler IdeState Method_TextDocumentRename
renameProvider state pluginId (RenameParams _prog (TextDocumentIdentifier uri) lspPos newNameText) = do
    nfp <- getNormalizedFilePathE uri
    codePointPos <- getCodePointPosition state nfp lspPos
    maybeParsed <- ImportAlias.getParsedModuleStale state nfp
    case maybeParsed of
        Nothing -> throwError $ PluginInternalError
            "The module hasn’t yet been parsed. Please wait for indexing to complete and try again."
        Just parsed -> do
            let hsModule = unLoc $ pm_parsed_source parsed
                exports = hsmodExports hsModule
                imports = hsmodImports hsModule
                hsDecls = hsmodDecls hsModule
            maybeAlias <- ImportAlias.resolveAliasAtPos
                getNamesAtPos state nfp lspPos codePointPos exports imports hsDecls
            case maybeAlias of
                Just (_lspRange, importAlias) -> ImportAlias.aliasBasedRename
                    state nfp uri importAlias exports hsDecls newNameText
                Nothing ->
                    nameBasedRename state pluginId nfp lspPos newNameText

-- | Name-based rename: the original rename logic.
nameBasedRename ::
    IdeState ->
    PluginId ->
    NormalizedFilePath ->
    Position ->
    T.Text ->
    ExceptT PluginError (HandlerM config) (MessageResult Method_TextDocumentRename)
nameBasedRename state pluginId nfp pos newNameText = do
    directOldNames <- getNamesAtPos state nfp pos
    directRefs <- concat <$> mapM (refsAtName state nfp) directOldNames

    {- References in HieDB are not necessarily transitive. With `NamedFieldPuns`, we can have
        indirect references through punned names. To find the transitive closure, we do a pass of
        the direct references to find the references for any punned names.
        See the `IndirectPuns` test for an example. -}
    indirectOldNames <- concat . filter ((>1) . length) <$>
        mapM (uncurry (getNamesAtPos state) <=< locToFilePos) directRefs
    let oldNames = filter matchesDirect indirectOldNames ++ directOldNames
           where
             matchesDirect n = occNameFS (nameOccName n) `elem` directFS
             directFS = map (occNameFS . nameOccName) directOldNames

    case oldNames of
        -- There were no Names at given position (e.g. rename triggered within a comment or on a keyword)
        [] -> throwError $ PluginInvalidParams "No symbol to rename at given position"
        _  -> do
            refs <- HS.fromList . concat <$> mapM (refsAtName state nfp) oldNames

            -- Validate rename
            crossModuleEnabled <- liftIO $ runAction "rename: config" state $ usePropertyAction #crossModule pluginId properties
            unless crossModuleEnabled $ failWhenImportOrExport state nfp refs oldNames
            when (any isBuiltInSyntax oldNames) $ throwError $ PluginInternalError "Invalid rename of built-in syntax"

            -- Perform rename
            let newName = mkTcOcc $ T.unpack newNameText
                filesRefs = collectWith locToUri refs
                getFileEdit (uri, locations) = do
                    verTxtDocId <- liftIO $ runAction "rename: getVersionedTextDoc" state $ getVersionedTextDoc (TextDocumentIdentifier uri)
                    getSrcEdit state verTxtDocId (replaceRefs newName locations)
            fileEdits <- mapM getFileEdit filesRefs
            pure $ InL $ fold fileEdits

-- | Limit renaming across modules.
failWhenImportOrExport ::
    IdeState ->
    NormalizedFilePath ->
    HashSet Location ->
    [Name] ->
    ExceptT PluginError (HandlerM config) ()
failWhenImportOrExport state nfp refLocs names = do
    pm <- runActionE "Rename.GetParsedModule" state
         (useE GetParsedModule nfp)
    let hsMod = unLoc $ pm_parsed_source pm
    case (unLoc <$> hsmodName hsMod, hsmodExports hsMod) of
        (mbModName, _) | not $ any (\n -> nameIsLocalOrFrom (replaceModName n mbModName) n) names
            -> throwError $ PluginInternalError "Renaming of an imported name is unsupported"
        (_, Just (L _ exports)) | any ((`HS.member` refLocs) . unsafeSrcSpanToLoc . getLoc) exports
            -> throwError $ PluginInternalError "Renaming of an exported name is unsupported"
        (Just _, Nothing) -> throwError $ PluginInternalError "Explicit export list required for renaming"
        _ -> pure ()

---------------------------------------------------------------------------------------------------
-- Source renaming

-- | Apply a function to a `ParsedSource` for a given `Uri` to compute a `WorkspaceEdit`.
getSrcEdit ::
    IdeState ->
    VersionedTextDocumentIdentifier ->
    (ParsedSource -> ParsedSource) ->
    ExceptT PluginError (HandlerM config) WorkspaceEdit
getSrcEdit state verTxtDocId updatePs = do
    ccs <- lift pluginGetClientCapabilities
    nfp <- getNormalizedFilePathE (verTxtDocId ^. L.uri)
    annAst <- runActionE "Rename.GetAnnotatedParsedSource" state
        (useE GetAnnotatedParsedSource nfp)
    let ps = annAst
        src = T.pack $ exactPrint ps
        res = T.pack $ exactPrint (updatePs ps)
    pure $ diffText ccs (verTxtDocId, src) res IncludeDeletions

-- | Replace names at every given `Location` (in a given `ParsedSource`) with a given new name.
replaceRefs ::
    OccName ->
    HashSet Location ->
    ParsedSource ->
    ParsedSource
replaceRefs newName refs = everywhere (mkT replaceLoc)
    where
        -- See Note [XRec and SrcSpans in the AST] in Language.Haskell.Syntax.Extension
        -- See Note [XRec and Anno in the AST] in GHC.Parser.Annotation
        -- GHC recommends using 'XRec' (available since 9.4.8 or earlier) to
        -- get the right annotation type for a given target type.
        -- XRec (GhcPass 'Parsed) RdrName = GenLocated (Anno RdrName) RdrName
        replaceLoc :: XRec (GhcPass 'Parsed) RdrName -> XRec (GhcPass 'Parsed) RdrName
        replaceLoc (L srcSpan oldRdrName)
            | isRef (locA srcSpan) = L srcSpan $ replace oldRdrName
        replaceLoc lOldRdrName = lOldRdrName
        replace :: RdrName -> RdrName
        replace (Qual modName _) = Qual modName newName
        replace _                = Unqual newName

        isRef :: SrcSpan -> Bool
        isRef = (`HS.member` refs) . unsafeSrcSpanToLoc

---------------------------------------------------------------------------------------------------
-- Reference finding

-- | Note: We only find exact name occurrences (i.e. type reference "depth" is 0).
refsAtName ::
    MonadIO m =>
    IdeState ->
    NormalizedFilePath ->
    Name ->
    ExceptT PluginError m [Location]
refsAtName state nfp name = do
    ShakeExtras{withHieDb} <- liftIO $ runAction "Rename.HieDb" state getShakeExtras
    ast <- handleGetHieAst state nfp
    dbRefs <- case nameModule_maybe name of
        Nothing -> pure []
        Just mod -> liftIO $ mapMaybe rowToLoc <$> withHieDb (\hieDb ->
            -- See Note [Generated references]
            -- REVIEW: Is this filter supposed to keep or remove generated references?
            filter (\(refRow HieDb.:. _) -> refIsGenerated refRow) <$>
            findReferences
                hieDb
                True
                (nameOccName name)
                (Just $ moduleName mod)
                (Just $ moduleUnit mod)
                [fromNormalizedFilePath nfp]
            )
    pure $ nameLocs name ast ++ dbRefs

nameLocs :: Name -> HieAstResult -> [Location]
nameLocs name (HAR _ _ rm _ _) =
    concatMap (map (realSrcSpanToLocation . fst))
              (M.lookup (Right name) rm)

---------------------------------------------------------------------------------------------------
-- Util

-- | Convert an LSP position (based on UTF-16 code units) to a position based on
-- whole Unicode code points.
getCodePointPosition ::
    MonadIO m =>
    IdeState ->
    NormalizedFilePath ->
    Position ->
    ExceptT PluginError m VFS.CodePointPosition
getCodePointPosition state nfp pos = do
    virtualFile <- runActionE "rename.getVirtualFile" state
        $ handleMaybeM (PluginInternalError
            ("Virtual file not found: " <> T.pack (show nfp)))
        $ getVirtualFile nfp
    case VFS.positionToCodePointPosition virtualFile pos of
        Nothing -> throwError $ PluginInvalidParams
            "The cursor position is inside a Unicode surrogate pair."
        Just codePointPosition -> pure codePointPosition

-- FIXME: 'getNamesAtPos' passes the LSP 'Position' directly to 'pointCommand',
-- which treats '_character' as a code-point column. This is incorrect for
-- files with supplementary-plane Unicode characters before the cursor.
-- Fixing it requires changes to 'pointCommand' in ghcide, not here.
getNamesAtPos :: MonadIO m => IdeState -> NormalizedFilePath -> Position -> ExceptT PluginError m [Name]
getNamesAtPos state nfp pos = do
    HAR{hieAst} <- handleGetHieAst state nfp
    pure $ getNamesAtPoint' hieAst pos

handleGetHieAst ::
    MonadIO m =>
    IdeState ->
    NormalizedFilePath ->
    ExceptT PluginError m HieAstResult
handleGetHieAst state nfp =
    -- We explicitly do not want to allow a stale version here - we only want to rename if
    -- the module compiles, otherwise we can't guarantee that we'll rename everything,
    -- which is bad (see https://github.com/haskell/haskell-language-server/issues/3799)
    fmap removeGenerated $ runActionE "Rename.GetHieAst" state $ useE GetHieAst nfp

{- Note [Generated references]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
GHC inserts `Use`s of record constructor everywhere where its record selectors are used,
which leads to record fields being renamed whenever corresponding constructor is renamed.
see https://github.com/haskell/haskell-language-server/issues/2915
To work around this, we filter out compiler-generated references.
-}
removeGenerated :: HieAstResult -> HieAstResult
removeGenerated HAR{..} =
    HAR{hieAst = sourceOnlyAsts, refMap = sourceOnlyRefMap, ..}
  where
    goAsts :: HieASTs a -> HieASTs a
    goAsts (HieASTs asts) = HieASTs (fmap goAst asts)

    goAst :: HieAST a -> HieAST a
    goAst (Node (SourcedNodeInfo sniMap) sp children) =
        let sourceOnlyNodeInfos = SourcedNodeInfo $ M.delete GeneratedInfo sniMap
        in Node sourceOnlyNodeInfos sp $ map goAst children

    sourceOnlyAsts = goAsts hieAst
    -- Also need to regenerate the RefMap, because the one in HAR
    -- is generated from HieASTs containing GeneratedInfo
    sourceOnlyRefMap = generateReferencesMap $ getAsts sourceOnlyAsts

collectWith :: (Hashable a, Eq b) => (a -> b) -> HashSet a -> [(b, HashSet a)]
collectWith f = map (\(a :| as) -> (f a, HS.fromList (a:as))) . groupWith f . HS.toList

-- | A variant 'getNamesAtPoint' that does not expect a 'PositionMapping'
-- FIXME: The use of 'pointCommand' is problematic. See 'getNamesAtPos' above.
getNamesAtPoint' :: HieASTs a -> Position -> [Name]
getNamesAtPoint' hf pos =
  concat $ pointCommand hf pos (rights . M.keys . getNodeIds)

-- | A variant of `getNamesAtPoint'` that also returns source spans.
-- FIXME: The use of 'pointCommand' is problematic. See 'getNamesAtPos' above.
getNamesSpansAtPoint' :: HieASTs a -> Position -> [([Name], RealSrcSpan)]
getNamesSpansAtPoint' hf pos =
  pointCommand hf pos $
    \astNode -> (rights . M.keys . getNodeIds $ astNode, nodeSpan astNode)

locToUri :: Location -> Uri
locToUri (Location uri _) = uri

unsafeSrcSpanToLoc :: SrcSpan -> Location
unsafeSrcSpanToLoc srcSpan =
    case srcSpanToLocation srcSpan of
        Nothing       -> error "Invalid conversion from UnhelpfulSpan to Location"
        Just location -> location

locToFilePos :: Monad m => Location -> ExceptT PluginError m (NormalizedFilePath, Position)
locToFilePos (Location uri (Range pos _)) = (,pos) <$> getNormalizedFilePathE uri

replaceModName :: Name -> Maybe ModuleName -> Module
replaceModName name mbModName =
    mkModule (moduleUnit $ nameModule name) (fromMaybe (mkModuleName "Main") mbModName)

---------------------------------------------------------------------------------------------------
-- Config

properties :: Properties '[ 'PropertyKey "crossModule" 'TBoolean]
properties = emptyProperties
  & defineBooleanProperty #crossModule
    "Enable experimental cross-module renaming" True