{-# LANGUAGE ScopedTypeVariables #-}
module Language.Rzk.VSCode.Env where

import           Control.Concurrent.Async   (Async, async, cancel)
import           Control.Concurrent.STM
import           Control.Exception          (catch)
import           Control.Monad.Reader
import qualified Data.Map.Strict            as Map
import qualified Data.Text                  as T
import           Language.LSP.Server
import           Language.Rzk.Syntax        (Module)
import qualified Language.Rzk.VSCode.Config as RzkConfig
import           Language.Rzk.VSCode.Logging
import qualified Language.Rzk.VSCode.ReferenceIndex as RefInd
import           Rzk.TypeCheck              (Checked, DeclView,
                                             TypeErrorInScopedContext)

-- | What checking one module produced.
--
-- 'cachedModuleChecked' is the state of the whole run /after/ this module: the
-- top-level scope and every declaration elaborated so far. That is what a resume
-- starts from — a cached elaborated term names the definitions it uses by their
-- foil name, which only means anything in the scope that produced it, so the scope
-- has to be cached with them.
--
-- 'cachedModuleDecls' is the rendered view of this module's own declarations,
-- which is all that completion, symbols and hover need.
data RzkCachedModule = RzkCachedModule
  { RzkCachedModule -> Checked
cachedModuleChecked :: Checked
  , RzkCachedModule -> [DeclView]
cachedModuleDecls   :: [DeclView]
  , RzkCachedModule -> [TypeErrorInScopedContext]
cachedModuleErrors  :: [TypeErrorInScopedContext]
  }

type RzkTypecheckCache = [(FilePath, RzkCachedModule)]

-- | Where a parse came from, deciding whether it may be reused.
data ParseSource
  = ParsedFromBuffer T.Text
    -- ^ Parsed from the editor buffer with this text; reusable for the
    -- current file while the buffer text is unchanged.
  | ParsedFromDisk
    -- ^ Parsed from the file on disk; trusted until invalidated by a
    -- file-change notification.
  | ParseInvalidated
    -- ^ The file changed; must be re-parsed. The module is kept as the last
    -- good parse, which hover falls back to while the file fails to parse.

-- | A parse result for the reference index. A failed parse is a 'Nothing'
-- module and is not retried until the source changes.
data ParsedModule = ParsedModule
  { ParsedModule -> ParseSource
parsedSource :: ParseSource
  , ParsedModule -> Maybe Module
parsedModule :: Maybe Module
  }

data ReferenceIndexCache = ReferenceIndexCache
  { ReferenceIndexCache -> Map FilePath ParsedModule
indexCacheModules :: Map.Map FilePath ParsedModule
    -- ^ Per-file parse cache. Entries parsed from disk are trusted until a
    -- file-change notification evicts them ('resetCacheForFiles'); the entry
    -- for the file being edited is revalidated against the editor buffer.
  , ReferenceIndexCache -> Maybe ([FilePath], ReferenceIndex)
indexCacheResult  :: Maybe ([FilePath], RefInd.ReferenceIndex)
    -- ^ The index last built, with the file set it was built from.
  }

emptyReferenceIndexCache :: ReferenceIndexCache
emptyReferenceIndexCache :: ReferenceIndexCache
emptyReferenceIndexCache = Map FilePath ParsedModule
-> Maybe ([FilePath], ReferenceIndex) -> ReferenceIndexCache
ReferenceIndexCache Map FilePath ParsedModule
forall k a. Map k a
Map.empty Maybe ([FilePath], ReferenceIndex)
forall a. Maybe a
Nothing

data RzkEnv = RzkEnv
  { RzkEnv -> TVar RzkTypecheckCache
rzkEnvTypecheckCache      :: TVar RzkTypecheckCache
  , RzkEnv -> TVar ReferenceIndexCache
rzkEnvReferenceIndexCache :: TVar ReferenceIndexCache
  , RzkEnv -> TVar (Maybe (Async ()))
rzkEnvTypecheckWorker     :: TVar (Maybe (Async ()))
    -- ^ The thread running the current project typecheck, if any.
  }

defaultRzkEnv :: IO RzkEnv
defaultRzkEnv :: IO RzkEnv
defaultRzkEnv = do
  typecheckCache <- RzkTypecheckCache -> IO (TVar RzkTypecheckCache)
forall a. a -> IO (TVar a)
newTVarIO []
  referenceIndexCache <- newTVarIO emptyReferenceIndexCache
  typecheckWorker <- newTVarIO Nothing
  return RzkEnv
    { rzkEnvTypecheckCache = typecheckCache
    , rzkEnvReferenceIndexCache = referenceIndexCache
    , rzkEnvTypecheckWorker = typecheckWorker
    }

type LSP = LspT RzkConfig.ServerConfig (ReaderT RzkEnv IO)

-- | Run the given action (a project typecheck) on a fresh worker thread,
-- cancelling the previous worker first. Cancellation waits for the old
-- worker to stop, so it can no longer write to the caches once the new
-- one starts. Typechecking runs on a worker so that the handler dispatch
-- thread stays responsive: 'lsp' dispatches messages sequentially, so a
-- multi-second re-check run directly in a notification handler would
-- block every later request (e.g. formatting on save).
spawnTypecheckWorker :: LSP () -> LSP ()
spawnTypecheckWorker :: LSP () -> LSP ()
spawnTypecheckWorker LSP ()
action = do
  lspEnv <- LspT
  ServerConfig (ReaderT RzkEnv IO) (LanguageContextEnv ServerConfig)
forall config (m :: * -> *).
MonadLsp config m =>
m (LanguageContextEnv config)
getLspEnv
  rzkEnv <- lift ask
  let run LspT ServerConfig (ReaderT RzkEnv m) a
act = ReaderT RzkEnv m a -> RzkEnv -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (LanguageContextEnv ServerConfig
-> LspT ServerConfig (ReaderT RzkEnv m) a -> ReaderT RzkEnv m a
forall config (m :: * -> *) a.
LanguageContextEnv config -> LspT config m a -> m a
runLspT LanguageContextEnv ServerConfig
lspEnv LspT ServerConfig (ReaderT RzkEnv m) a
act) RzkEnv
rzkEnv
  liftIO $ do
    oldWorker <- atomically $ swapTVar (rzkEnvTypecheckWorker rzkEnv) Nothing
    mapM_ cancel oldWorker
    worker <- async $
      run action `catch` \(ProgressCancelledException
_ :: ProgressCancelledException) ->
        LSP () -> IO ()
forall {m :: * -> *} {a}.
LspT ServerConfig (ReaderT RzkEnv m) a -> m a
run (Text -> LSP ()
forall c (m :: * -> *). MonadLsp c m => Text -> m ()
logInfo (FilePath -> Text
T.pack FilePath
"Typechecking was cancelled by the client"))
    atomically $ writeTVar (rzkEnvTypecheckWorker rzkEnv) (Just worker)

cacheTypecheckedModules :: RzkTypecheckCache -> LSP ()
cacheTypecheckedModules :: RzkTypecheckCache -> LSP ()
cacheTypecheckedModules RzkTypecheckCache
cache = ReaderT RzkEnv IO () -> LSP ()
forall (m :: * -> *) a. Monad m => m a -> LspT ServerConfig m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT RzkEnv IO () -> LSP ()) -> ReaderT RzkEnv IO () -> LSP ()
forall a b. (a -> b) -> a -> b
$ do
  typecheckCache <- (RzkEnv -> TVar RzkTypecheckCache)
-> ReaderT RzkEnv IO (TVar RzkTypecheckCache)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks RzkEnv -> TVar RzkTypecheckCache
rzkEnvTypecheckCache
  referenceIndexCache <- asks rzkEnvReferenceIndexCache
  liftIO $ atomically $ do
    writeTVar typecheckCache cache
    -- Drop the built index (the declaration set may have changed), but keep
    -- the parse entries: files changed on disk are already evicted per path
    -- by 'resetCacheForFiles' on file-change notifications, and a kept entry
    -- may hold the last good parse of a file that currently has a syntax
    -- error, which hover falls back to.
    modifyTVar referenceIndexCache $ \ReferenceIndexCache
c -> ReferenceIndexCache
c { indexCacheResult = Nothing }

resetCacheForAllFiles :: LSP ()
resetCacheForAllFiles :: LSP ()
resetCacheForAllFiles = RzkTypecheckCache -> LSP ()
cacheTypecheckedModules []

resetCacheForFiles :: [FilePath] -> LSP ()
resetCacheForFiles :: [FilePath] -> LSP ()
resetCacheForFiles [FilePath]
paths = ReaderT RzkEnv IO () -> LSP ()
forall (m :: * -> *) a. Monad m => m a -> LspT ServerConfig m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT RzkEnv IO () -> LSP ()) -> ReaderT RzkEnv IO () -> LSP ()
forall a b. (a -> b) -> a -> b
$ do
  typecheckCache <- (RzkEnv -> TVar RzkTypecheckCache)
-> ReaderT RzkEnv IO (TVar RzkTypecheckCache)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks RzkEnv -> TVar RzkTypecheckCache
rzkEnvTypecheckCache
  referenceIndexCache <- asks rzkEnvReferenceIndexCache
  liftIO $ atomically $ do
    modifyTVar typecheckCache (takeWhile ((`notElem` paths) . fst))
    modifyTVar referenceIndexCache $ \ReferenceIndexCache
c -> ReferenceIndexCache
      { indexCacheModules :: Map FilePath ParsedModule
indexCacheModules =
          (FilePath
 -> Map FilePath ParsedModule -> Map FilePath ParsedModule)
-> Map FilePath ParsedModule
-> [FilePath]
-> Map FilePath ParsedModule
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((ParsedModule -> ParsedModule)
-> FilePath
-> Map FilePath ParsedModule
-> Map FilePath ParsedModule
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (\ParsedModule
pm -> ParsedModule
pm { parsedSource = ParseInvalidated }))
                (ReferenceIndexCache -> Map FilePath ParsedModule
indexCacheModules ReferenceIndexCache
c) [FilePath]
paths
      , indexCacheResult :: Maybe ([FilePath], ReferenceIndex)
indexCacheResult  = Maybe ([FilePath], ReferenceIndex)
forall a. Maybe a
Nothing
      }

getCachedTypecheckedModules :: LSP RzkTypecheckCache
getCachedTypecheckedModules :: LSP RzkTypecheckCache
getCachedTypecheckedModules = ReaderT RzkEnv IO RzkTypecheckCache -> LSP RzkTypecheckCache
forall (m :: * -> *) a. Monad m => m a -> LspT ServerConfig m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT RzkEnv IO RzkTypecheckCache -> LSP RzkTypecheckCache)
-> ReaderT RzkEnv IO RzkTypecheckCache -> LSP RzkTypecheckCache
forall a b. (a -> b) -> a -> b
$ do
  typecheckCache <- (RzkEnv -> TVar RzkTypecheckCache)
-> ReaderT RzkEnv IO (TVar RzkTypecheckCache)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks RzkEnv -> TVar RzkTypecheckCache
rzkEnvTypecheckCache
  liftIO $ readTVarIO typecheckCache

cacheReferenceIndex :: ReferenceIndexCache -> LSP ()
cacheReferenceIndex :: ReferenceIndexCache -> LSP ()
cacheReferenceIndex ReferenceIndexCache
cache = ReaderT RzkEnv IO () -> LSP ()
forall (m :: * -> *) a. Monad m => m a -> LspT ServerConfig m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT RzkEnv IO () -> LSP ()) -> ReaderT RzkEnv IO () -> LSP ()
forall a b. (a -> b) -> a -> b
$ do
  referenceIndexCache <- (RzkEnv -> TVar ReferenceIndexCache)
-> ReaderT RzkEnv IO (TVar ReferenceIndexCache)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks RzkEnv -> TVar ReferenceIndexCache
rzkEnvReferenceIndexCache
  liftIO $ atomically $ writeTVar referenceIndexCache cache

getCachedReferenceIndex :: LSP ReferenceIndexCache
getCachedReferenceIndex :: LSP ReferenceIndexCache
getCachedReferenceIndex = ReaderT RzkEnv IO ReferenceIndexCache -> LSP ReferenceIndexCache
forall (m :: * -> *) a. Monad m => m a -> LspT ServerConfig m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT RzkEnv IO ReferenceIndexCache -> LSP ReferenceIndexCache)
-> ReaderT RzkEnv IO ReferenceIndexCache -> LSP ReferenceIndexCache
forall a b. (a -> b) -> a -> b
$ do
  referenceIndexCache <- (RzkEnv -> TVar ReferenceIndexCache)
-> ReaderT RzkEnv IO (TVar ReferenceIndexCache)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks RzkEnv -> TVar ReferenceIndexCache
rzkEnvReferenceIndexCache
  liftIO $ readTVarIO referenceIndexCache