-- | Building blocks for "recompiling" (actually just type-checking)
-- the Futhark program managed by the language server.  The challenge
-- here is that if the program becomes type-invalid, we want to keep
-- the old state around.
module Futhark.LSP.Compile (tryTakeStateFromIORef, tryReCompile) where

import Colog.Core (logStringStderr, (<&))
import Control.Lens.Getter (view)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.IORef (IORef, readIORef, writeIORef)
import Data.Map qualified as M
import Data.Maybe (fromMaybe)
import Data.Text qualified as T
import Futhark.Compiler.Program (LoadedProg, lpFilePaths, lpWarnings, noLoadedProg, reloadProg)
import Futhark.LSP.Diagnostic (diagnosticSource, maxDiagnostic, publishErrorDiagnostics, publishWarningDiagnostics)
import Futhark.LSP.State (State (..), emptyState, updateStaleContent, updateStaleMapping)
import Futhark.LSP.Tool (computeMapping)
import Language.Futhark.Warnings (listWarnings)
import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFile, getVirtualFiles)
import Language.LSP.Types
  ( filePathToUri,
    fromNormalizedFilePath,
    toNormalizedUri,
    uriToNormalizedFilePath,
  )
import Language.LSP.VFS (VFS, vfsMap, virtualFileText)

-- | Try to take state from IORef, if it's empty, try to compile.
tryTakeStateFromIORef :: IORef State -> Maybe FilePath -> LspT () IO State
tryTakeStateFromIORef :: IORef State -> Maybe FilePath -> LspT () IO State
tryTakeStateFromIORef IORef State
state_mvar Maybe FilePath
file_path = do
  State
old_state <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef State
state_mvar
  case State -> Maybe LoadedProg
stateProgram State
old_state of
    Maybe LoadedProg
Nothing -> do
      State
new_state <- State -> Maybe FilePath -> LoadedProg -> LspT () IO State
tryCompile State
old_state Maybe FilePath
file_path LoadedProg
noLoadedProg
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar State
new_state
      forall (f :: * -> *) a. Applicative f => a -> f a
pure State
new_state
    Just LoadedProg
prog -> do
      -- If this is in the context of some file that is not part of
      -- the program, try to reload the program from that file.
      let files :: [FilePath]
files = LoadedProg -> [FilePath]
lpFilePaths LoadedProg
prog
      State
state <- case Maybe FilePath
file_path of
        Just FilePath
file_path'
          | FilePath
file_path' forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [FilePath]
files -> do
              forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& (FilePath
"File not part of program: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> FilePath
show FilePath
file_path')
              forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& (FilePath
"Program contains: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> FilePath
show [FilePath]
files)
              State -> Maybe FilePath -> LoadedProg -> LspT () IO State
tryCompile State
old_state Maybe FilePath
file_path LoadedProg
noLoadedProg
        Maybe FilePath
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure State
old_state
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar State
state
      forall (f :: * -> *) a. Applicative f => a -> f a
pure State
state

-- | Try to (re)-compile, replace old state if successful.
tryReCompile :: IORef State -> Maybe FilePath -> LspT () IO ()
tryReCompile :: IORef State -> Maybe FilePath -> LspT () IO ()
tryReCompile IORef State
state_mvar Maybe FilePath
file_path = do
  forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"(Re)-compiling ..."
  State
old_state <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef State
state_mvar
  let loaded_prog :: LoadedProg
loaded_prog = State -> LoadedProg
getLoadedProg State
old_state
  State
new_state <- State -> Maybe FilePath -> LoadedProg -> LspT () IO State
tryCompile State
old_state Maybe FilePath
file_path LoadedProg
loaded_prog
  case State -> Maybe LoadedProg
stateProgram State
new_state of
    Maybe LoadedProg
Nothing -> do
      forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Failed to (re)-compile, using old state or Nothing"
      forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Computing PositionMapping for: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> FilePath
show Maybe FilePath
file_path
      Maybe PositionMapping
mapping <- State -> Maybe FilePath -> LspM () (Maybe PositionMapping)
computeMapping State
old_state Maybe FilePath
file_path
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar forall a b. (a -> b) -> a -> b
$ Maybe FilePath -> Maybe PositionMapping -> State -> State
updateStaleMapping Maybe FilePath
file_path Maybe PositionMapping
mapping State
old_state
    Just LoadedProg
_ -> do
      forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"(Re)-compile successful"
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar State
new_state

-- | Try to compile, publish diagnostics on warnings and errors, return newly compiled state.
--  Single point where the compilation is done, and shouldn't be exported.
tryCompile :: State -> Maybe FilePath -> LoadedProg -> LspT () IO State
tryCompile :: State -> Maybe FilePath -> LoadedProg -> LspT () IO State
tryCompile State
_ Maybe FilePath
Nothing LoadedProg
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure State
emptyState
tryCompile State
state (Just FilePath
path) LoadedProg
old_loaded_prog = do
  forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Reloading program from " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> FilePath
show FilePath
path
  VFS
vfs <- forall config (m :: * -> *). MonadLsp config m => m VFS
getVirtualFiles
  Either (NonEmpty ProgError) LoadedProg
res <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ LoadedProg
-> [FilePath] -> VFS -> IO (Either (NonEmpty ProgError) LoadedProg)
reloadProg LoadedProg
old_loaded_prog [FilePath
path] (VFS -> VFS
transformVFS VFS
vfs) -- NOTE: vfs only keeps track of current opened files
  forall config (m :: * -> *).
MonadLsp config m =>
Int -> Maybe DiagnosticSource -> m ()
flushDiagnosticsBySource Int
maxDiagnostic Maybe DiagnosticSource
diagnosticSource
  case Either (NonEmpty ProgError) LoadedProg
res of
    Right LoadedProg
new_loaded_prog -> do
      forall a. [(SrcLoc, Doc a)] -> LspT () IO ()
publishWarningDiagnostics forall a b. (a -> b) -> a -> b
$ Warnings -> [(SrcLoc, Doc ())]
listWarnings forall a b. (a -> b) -> a -> b
$ LoadedProg -> Warnings
lpWarnings LoadedProg
new_loaded_prog
      Maybe VirtualFile
maybe_virtual_file <- forall config (m :: * -> *).
MonadLsp config m =>
NormalizedUri -> m (Maybe VirtualFile)
getVirtualFile forall a b. (a -> b) -> a -> b
$ Uri -> NormalizedUri
toNormalizedUri forall a b. (a -> b) -> a -> b
$ FilePath -> Uri
filePathToUri FilePath
path
      case Maybe VirtualFile
maybe_virtual_file of
        Maybe VirtualFile
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Maybe LoadedProg -> Map FilePath StaleFile -> State
State (forall a. a -> Maybe a
Just LoadedProg
new_loaded_prog) (State -> Map FilePath StaleFile
staleData State
state) -- should never happen
        Just VirtualFile
virtual_file ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ FilePath -> VirtualFile -> LoadedProg -> State -> State
updateStaleContent FilePath
path VirtualFile
virtual_file LoadedProg
new_loaded_prog State
state
    -- Preserve files that have been opened should be enoguth.
    -- But still might need an update on re-compile logic, don't discard all state afterwards,
    -- try to compile from root file, if there is a depencency relatetion, improve performance and provide more dignostic.
    Left NonEmpty ProgError
prog_error -> do
      forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Compilation failed, publishing diagnostics"
      NonEmpty ProgError -> LspT () IO ()
publishErrorDiagnostics NonEmpty ProgError
prog_error
      forall (f :: * -> *) a. Applicative f => a -> f a
pure State
emptyState

-- | Transform VFS to a map of file paths to file contents.
-- This is used to pass the file contents to the compiler.
transformVFS :: VFS -> M.Map FilePath T.Text
transformVFS :: VFS -> VFS
transformVFS VFS
vfs =
  forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
M.foldrWithKey
    ( \NormalizedUri
uri VirtualFile
virtual_file VFS
acc ->
        case NormalizedUri -> Maybe NormalizedFilePath
uriToNormalizedFilePath NormalizedUri
uri of
          Maybe NormalizedFilePath
Nothing -> VFS
acc
          Just NormalizedFilePath
file_path ->
            forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (NormalizedFilePath -> FilePath
fromNormalizedFilePath NormalizedFilePath
file_path) (VirtualFile -> DiagnosticSource
virtualFileText VirtualFile
virtual_file) VFS
acc
    )
    forall k a. Map k a
M.empty
    (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall s a. HasVfsMap s a => Lens' s a
vfsMap VFS
vfs)

getLoadedProg :: State -> LoadedProg
getLoadedProg :: State -> LoadedProg
getLoadedProg State
state = forall a. a -> Maybe a -> a
fromMaybe LoadedProg
noLoadedProg (State -> Maybe LoadedProg
stateProgram State
state)