-- | 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 (Severity (Debug, Warning), (<&))
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.IORef (IORef, readIORef, writeIORef)
import Data.Int (Int32)
import Data.Maybe (fromMaybe, isJust)
import Futhark.Compiler.Program (LoadedProg, lpFilePaths, lpWarnings, noLoadedProg, reloadProg)
import Futhark.LSP.Diagnostic (diagnosticSource, maxDiagnostic, publishErrorDiagnostics, publishWarningDiagnostics)
import Futhark.LSP.State (State (..), updateStaleContent, updateStaleMapping)
import Futhark.LSP.Tool (computeMapping, logWithSeverity, transformVFS)
import Futhark.Util (showText)
import Language.Futhark.Warnings (listWarnings)
import Language.LSP.Protocol.Types
  ( filePathToUri,
    toNormalizedUri,
  )
import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFile, getVirtualFiles)
import Language.LSP.VFS (virtualFileVersion)

-- | Get the current state.  If the given file is not part of the
-- current program, first try to compile the program rooted at that
-- file, unless that already failed for this version of the file.
tryTakeStateFromIORef :: IORef State -> Maybe FilePath -> LspT () IO State
tryTakeStateFromIORef :: IORef State -> Maybe FilePath -> LspT () IO State
tryTakeStateFromIORef IORef State
state_mvar Maybe FilePath
Nothing = 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
tryTakeStateFromIORef IORef State
state_mvar (Just FilePath
file_path) = do
  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 stateProgram old_state of
    Just LoadedProg
prog
      | FilePath
file_path FilePath -> [FilePath] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` LoadedProg -> [FilePath]
lpFilePaths LoadedProg
prog -> State -> LspT () IO State
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure State
old_state
      | Bool
otherwise -> do
          Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Warning LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& (Text
"File not part of program: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> FilePath -> Text
forall a. Show a => a -> Text
showText FilePath
file_path)
          Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Warning LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& (Text
"Program contains: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [FilePath] -> Text
forall a. Show a => a -> Text
showText (LoadedProg -> [FilePath]
lpFilePaths LoadedProg
prog))
          State -> LspT () IO State
compileUnlessFailed State
old_state
    Maybe LoadedProg
Nothing -> State -> LspT () IO State
compileUnlessFailed State
old_state
  where
    compileUnlessFailed :: State -> LspT () IO State
compileUnlessFailed State
old_state = do
      attempt <- FilePath -> LspT () IO (Maybe (FilePath, Int32))
compileAttempt FilePath
file_path
      if isJust attempt && attempt == stateFailed old_state
        then pure old_state
        else do
          new_state <- either id id <$> tryCompile old_state file_path noLoadedProg
          liftIO $ writeIORef state_mvar new_state
          pure new_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
_ Maybe FilePath
Nothing = () -> LspT () IO ()
forall a. a -> LspT () IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
tryReCompile IORef State
state_mvar (Just FilePath
file_path) = do
  Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"(Re)-compiling ..."
  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
  result <- tryCompile old_state file_path $ getLoadedProg old_state
  case result of
    Left State
failed_state -> do
      Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"Failed to (re)-compile, using old state or Nothing"
      Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"Computing PositionMapping for: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> FilePath -> Text
forall a. Show a => a -> Text
showText FilePath
file_path
      mapping <- State -> Maybe FilePath -> LspM () (Maybe PositionMapping)
computeMapping State
failed_state (FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
file_path)
      liftIO $ writeIORef state_mvar $ updateStaleMapping (Just file_path) mapping failed_state
    Right State
new_state -> do
      Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"(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

-- | The file and its current version, as recorded in 'stateFailed'.
compileAttempt :: FilePath -> LspT () IO (Maybe (FilePath, Int32))
compileAttempt :: FilePath -> LspT () IO (Maybe (FilePath, Int32))
compileAttempt FilePath
file_path =
  (VirtualFile -> (FilePath, Int32))
-> Maybe VirtualFile -> Maybe (FilePath, Int32)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((FilePath
file_path,) (Int32 -> (FilePath, Int32))
-> (VirtualFile -> Int32) -> VirtualFile -> (FilePath, Int32)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VirtualFile -> Int32
virtualFileVersion)
    (Maybe VirtualFile -> Maybe (FilePath, Int32))
-> LspT () IO (Maybe VirtualFile)
-> LspT () IO (Maybe (FilePath, Int32))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NormalizedUri -> LspT () IO (Maybe VirtualFile)
forall config (m :: * -> *).
MonadLsp config m =>
NormalizedUri -> m (Maybe VirtualFile)
getVirtualFile (Uri -> NormalizedUri
toNormalizedUri (Uri -> NormalizedUri) -> Uri -> NormalizedUri
forall a b. (a -> b) -> a -> b
$ FilePath -> Uri
filePathToUri FilePath
file_path)

-- | Try to compile and publish diagnostics on warnings and errors.
-- Returns the new state on success, and on failure the old state
-- with the failed attempt recorded.  Single point where the
-- compilation is done, and shouldn't be exported.
tryCompile :: State -> FilePath -> LoadedProg -> LspT () IO (Either State State)
tryCompile :: State -> FilePath -> LoadedProg -> LspT () IO (Either State State)
tryCompile State
state FilePath
path LoadedProg
old_loaded_prog = do
  Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"Reloading program from " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> FilePath -> Text
forall a. Show a => a -> Text
showText FilePath
path
  vfs <- LspT () IO VFS
forall config (m :: * -> *). MonadLsp config m => m VFS
getVirtualFiles
  res <- liftIO $ reloadProg old_loaded_prog [path] (transformVFS vfs) -- NOTE: vfs only keeps track of current opened files
  flushDiagnosticsBySource maxDiagnostic diagnosticSource
  case res of
    Right LoadedProg
new_loaded_prog -> do
      Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"Successfully compiled, publishing warnings"
      [(Loc, Doc ())] -> LspT () IO ()
forall a. [(Loc, Doc a)] -> LspT () IO ()
publishWarningDiagnostics ([(Loc, Doc ())] -> LspT () IO ())
-> [(Loc, Doc ())] -> LspT () IO ()
forall a b. (a -> b) -> a -> b
$ Warnings -> [(Loc, Doc ())]
listWarnings (Warnings -> [(Loc, Doc ())]) -> Warnings -> [(Loc, Doc ())]
forall a b. (a -> b) -> a -> b
$ LoadedProg -> Warnings
lpWarnings LoadedProg
new_loaded_prog
      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
      let new_state = case Maybe VirtualFile
maybe_virtual_file of
            Maybe VirtualFile
Nothing -> State
state {stateProgram = Just new_loaded_prog} -- should never happen
            Just VirtualFile
virtual_file -> FilePath -> VirtualFile -> LoadedProg -> State -> State
updateStaleContent FilePath
path VirtualFile
virtual_file LoadedProg
new_loaded_prog State
state
      pure $ Right new_state {stateFailed = Nothing}
    Left NonEmpty ProgError
prog_error -> do
      Severity -> LogAction (LspT () IO) Text
forall c (m :: * -> *).
MonadLsp c m =>
Severity -> LogAction m Text
logWithSeverity Severity
Debug LogAction (LspT () IO) Text -> Text -> LspT () IO ()
forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text
"Compilation failed, publishing diagnostics"
      NonEmpty ProgError -> LspT () IO ()
publishErrorDiagnostics NonEmpty ProgError
prog_error
      attempt <- FilePath -> LspT () IO (Maybe (FilePath, Int32))
compileAttempt FilePath
path
      pure $ Left state {stateFailed = attempt}

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)