-- | 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.Protocol.Types
  ( filePathToUri,
    fromNormalizedFilePath,
    toNormalizedUri,
    uriToNormalizedFilePath,
  )
import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFile, getVirtualFiles)
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 <- IO State -> LspT () IO State
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO State -> LspT () IO State) -> IO State -> LspT () IO State
forall a b. (a -> b) -> a -> b
$ IORef State -> IO State
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
      IO () -> LspT () IO ()
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> LspT () IO ()) -> IO () -> LspT () IO ()
forall a b. (a -> b) -> a -> b
$ IORef State -> State -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar State
new_state
      State -> LspT () IO State
forall a. a -> LspT () IO a
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' FilePath -> [FilePath] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [FilePath]
files -> do
              LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& (FilePath
"File not part of program: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath -> FilePath
forall a. Show a => a -> FilePath
show FilePath
file_path')
              LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& (FilePath
"Program contains: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> [FilePath] -> FilePath
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
_ -> State -> LspT () IO State
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure State
old_state
      IO () -> LspT () IO ()
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> LspT () IO ()) -> IO () -> LspT () IO ()
forall a b. (a -> b) -> a -> b
$ IORef State -> State -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar State
state
      State -> LspT () IO State
forall a. a -> LspT () IO a
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
  LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"(Re)-compiling ..."
  State
old_state <- IO State -> LspT () IO State
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO State -> LspT () IO State) -> IO State -> LspT () IO State
forall a b. (a -> b) -> a -> b
$ IORef State -> IO State
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
      LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Failed to (re)-compile, using old state or Nothing"
      LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Computing PositionMapping for: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> Maybe FilePath -> FilePath
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
      IO () -> LspT () IO ()
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> LspT () IO ()) -> IO () -> LspT () IO ()
forall a b. (a -> b) -> a -> b
$ IORef State -> State -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef State
state_mvar (State -> IO ()) -> State -> IO ()
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
      LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"(Re)-compile successful"
      IO () -> LspT () IO ()
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> LspT () IO ()) -> IO () -> LspT () IO ()
forall a b. (a -> b) -> a -> b
$ IORef State -> State -> IO ()
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
_ = State -> LspT () IO State
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure State
emptyState
tryCompile State
state (Just FilePath
path) LoadedProg
old_loaded_prog = do
  LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Reloading program from " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath -> FilePath
forall a. Show a => a -> FilePath
show FilePath
path
  VFS
vfs <- LspT () IO VFS
forall config (m :: * -> *). MonadLsp config m => m VFS
getVirtualFiles
  Either (NonEmpty ProgError) LoadedProg
res <- IO (Either (NonEmpty ProgError) LoadedProg)
-> LspT () IO (Either (NonEmpty ProgError) LoadedProg)
forall a. IO a -> LspT () IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either (NonEmpty ProgError) LoadedProg)
 -> LspT () IO (Either (NonEmpty ProgError) LoadedProg))
-> IO (Either (NonEmpty ProgError) LoadedProg)
-> LspT () IO (Either (NonEmpty ProgError) LoadedProg)
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
  Int -> Maybe Text -> LspT () IO ()
forall config (m :: * -> *).
MonadLsp config m =>
Int -> Maybe Text -> m ()
flushDiagnosticsBySource Int
maxDiagnostic Maybe Text
diagnosticSource
  case Either (NonEmpty ProgError) LoadedProg
res of
    Right LoadedProg
new_loaded_prog -> do
      [(SrcLoc, Doc ())] -> LspT () IO ()
forall a. [(SrcLoc, Doc a)] -> LspT () IO ()
publishWarningDiagnostics ([(SrcLoc, Doc ())] -> LspT () IO ())
-> [(SrcLoc, Doc ())] -> LspT () IO ()
forall a b. (a -> b) -> a -> b
$ Warnings -> [(SrcLoc, Doc ())]
listWarnings (Warnings -> [(SrcLoc, Doc ())]) -> Warnings -> [(SrcLoc, Doc ())]
forall a b. (a -> b) -> a -> b
$ LoadedProg -> Warnings
lpWarnings LoadedProg
new_loaded_prog
      Maybe VirtualFile
maybe_virtual_file <- NormalizedUri -> LspT () IO (Maybe VirtualFile)
forall config (m :: * -> *).
MonadLsp config m =>
NormalizedUri -> m (Maybe VirtualFile)
getVirtualFile (NormalizedUri -> LspT () IO (Maybe VirtualFile))
-> NormalizedUri -> LspT () IO (Maybe VirtualFile)
forall a b. (a -> b) -> a -> b
$ Uri -> NormalizedUri
toNormalizedUri (Uri -> NormalizedUri) -> Uri -> NormalizedUri
forall a b. (a -> b) -> a -> b
$ FilePath -> Uri
filePathToUri FilePath
path
      case Maybe VirtualFile
maybe_virtual_file of
        Maybe VirtualFile
Nothing -> State -> LspT () IO State
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (State -> LspT () IO State) -> State -> LspT () IO State
forall a b. (a -> b) -> a -> b
$ Maybe LoadedProg -> Map FilePath StaleFile -> State
State (LoadedProg -> Maybe LoadedProg
forall a. a -> Maybe a
Just LoadedProg
new_loaded_prog) (State -> Map FilePath StaleFile
staleData State
state) -- should never happen
        Just VirtualFile
virtual_file ->
          State -> LspT () IO State
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (State -> LspT () IO State) -> State -> LspT () IO State
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
      LogAction (LspT () IO) FilePath
forall (m :: * -> *). MonadIO m => LogAction m FilePath
logStringStderr LogAction (LspT () IO) FilePath -> FilePath -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& FilePath
"Compilation failed, publishing diagnostics"
      NonEmpty ProgError -> LspT () IO ()
publishErrorDiagnostics NonEmpty ProgError
prog_error
      State -> LspT () IO State
forall a. a -> LspT () IO a
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 =
  (NormalizedUri -> VirtualFile -> VFS -> VFS)
-> VFS -> Map NormalizedUri VirtualFile -> 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 ->
            FilePath -> Text -> VFS -> VFS
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (NormalizedFilePath -> FilePath
fromNormalizedFilePath NormalizedFilePath
file_path) (VirtualFile -> Text
virtualFileText VirtualFile
virtual_file) VFS
acc
    )
    VFS
forall k a. Map k a
M.empty
    (Getting
  (Map NormalizedUri VirtualFile) VFS (Map NormalizedUri VirtualFile)
-> VFS -> Map NormalizedUri VirtualFile
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting
  (Map NormalizedUri VirtualFile) VFS (Map NormalizedUri VirtualFile)
forall s a. HasVfsMap s a => Lens' s a
Lens' VFS (Map NormalizedUri VirtualFile)
vfsMap VFS
vfs)

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