module Unison.LSP.UCMWorker where

import Control.Monad.Reader
import Unison.Codebase qualified as Codebase
import Unison.Codebase.Branch qualified as Branch
import Unison.Codebase.Branch.Names qualified as Branch
import Unison.Codebase.ProjectPath (ProjectPath)
import Unison.Codebase.ProjectPath qualified as PP
import Unison.LSP.Completion
import Unison.LSP.Types
import Unison.LSP.Util.Signal (Signal)
import Unison.LSP.Util.Signal qualified as Signal
import Unison.LSP.VFS qualified as VFS
import Unison.Names (Names)
import Unison.Prelude
import Unison.PrettyPrintEnv.Names qualified as PPE
import Unison.PrettyPrintEnvDecl
import Unison.PrettyPrintEnvDecl.Names qualified as PPED
import Unison.Server.NameSearch (NameSearch)
import Unison.Server.NameSearch.FromNames qualified as NameSearch
import Unison.Sqlite qualified as Sqlite
import UnliftIO.STM

-- | Watches for state changes in UCM and updates cached LSP state accordingly
ucmWorker ::
  TMVar PrettyPrintEnvDecl ->
  TMVar Names ->
  TMVar (NameSearch Sqlite.Transaction) ->
  TMVar ProjectPath ->
  Signal PP.ProjectPathIds ->
  Lsp ()
ucmWorker :: TMVar PrettyPrintEnvDecl
-> TMVar Names
-> TMVar (NameSearch Transaction)
-> TMVar ProjectPath
-> Signal ProjectPathIds
-> Lsp ()
ucmWorker TMVar PrettyPrintEnvDecl
ppedVar TMVar Names
currentNamesVar TMVar (NameSearch Transaction)
nameSearchCacheVar TMVar ProjectPath
currentPathVar Signal ProjectPathIds
changeSignal = do
  STM ProjectPathIds
signalChanges <- Signal ProjectPathIds -> Lsp (STM ProjectPathIds)
forall (m :: * -> *) a. MonadIO m => Signal a -> m (STM a)
Signal.subscribe Signal ProjectPathIds
changeSignal
  STM ProjectPathIds -> Maybe (Branch IO) -> Lsp ()
forall a. STM ProjectPathIds -> Maybe (Branch IO) -> Lsp a
loop STM ProjectPathIds
signalChanges Maybe (Branch IO)
forall a. Maybe a
Nothing
  where
    loop :: STM PP.ProjectPathIds -> Maybe (Branch.Branch IO) -> Lsp a
    loop :: forall a. STM ProjectPathIds -> Maybe (Branch IO) -> Lsp a
loop STM ProjectPathIds
signalChanges Maybe (Branch IO)
currentBranch = do
      Env {Codebase IO Symbol Ann
codebase :: Codebase IO Symbol Ann
$sel:codebase:Env :: Env -> Codebase IO Symbol Ann
codebase, TMVar CompletionTree
completionsVar :: TMVar CompletionTree
$sel:completionsVar:Env :: Env -> TMVar CompletionTree
completionsVar} <- Lsp Env
forall r (m :: * -> *). MonadReader r m => m r
ask
      STM ProjectPathIds
-> Maybe (Branch IO) -> Lsp (ProjectPath, Maybe (Branch IO))
getChanges STM ProjectPathIds
signalChanges Maybe (Branch IO)
currentBranch Lsp (ProjectPath, Maybe (Branch IO))
-> ((ProjectPath, Maybe (Branch IO)) -> Lsp a) -> Lsp a
forall a b. Lsp a -> (a -> Lsp b) -> Lsp b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        (ProjectPath
_newPP, Maybe (Branch IO)
Nothing) -> STM ProjectPathIds -> Maybe (Branch IO) -> Lsp a
forall a. STM ProjectPathIds -> Maybe (Branch IO) -> Lsp a
loop STM ProjectPathIds
signalChanges Maybe (Branch IO)
currentBranch
        (ProjectPath
newPP, Just Branch IO
newBranch) -> do
          let newBranch0 :: Branch0 IO
newBranch0 = Branch IO -> Branch0 IO
forall (m :: * -> *). Branch m -> Branch0 m
Branch.head Branch IO
newBranch
          let newNames :: Names
newNames = Branch0 IO -> Names
forall (m :: * -> *). Branch0 m -> Names
Branch.toNames Branch0 IO
newBranch0
          Int
hl <- IO Int -> Lsp Int
forall a. IO a -> Lsp a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> Lsp Int) -> IO Int -> Lsp Int
forall a b. (a -> b) -> a -> b
$ Codebase IO Symbol Ann -> Transaction Int -> IO Int
forall (m :: * -> *) v a b.
MonadIO m =>
Codebase m v a -> Transaction b -> m b
Codebase.runTransaction Codebase IO Symbol Ann
codebase Transaction Int
Codebase.hashLength
          let pped :: PrettyPrintEnvDecl
pped = Namer -> Suffixifier -> PrettyPrintEnvDecl
PPED.makePPED (Int -> Names -> Namer
PPE.hqNamer Int
hl Names
newNames) (Names -> Suffixifier
PPE.suffixifyByHash Names
newNames)
          STM () -> Lsp ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Lsp ()) -> STM () -> Lsp ()
forall a b. (a -> b) -> a -> b
$ do
            TMVar ProjectPath -> ProjectPath -> STM ()
forall a. TMVar a -> a -> STM ()
writeTMVar TMVar ProjectPath
currentPathVar ProjectPath
newPP
            TMVar Names -> Names -> STM ()
forall a. TMVar a -> a -> STM ()
writeTMVar TMVar Names
currentNamesVar Names
newNames
            TMVar PrettyPrintEnvDecl -> PrettyPrintEnvDecl -> STM ()
forall a. TMVar a -> a -> STM ()
writeTMVar TMVar PrettyPrintEnvDecl
ppedVar PrettyPrintEnvDecl
pped
            TMVar (NameSearch Transaction) -> NameSearch Transaction -> STM ()
forall a. TMVar a -> a -> STM ()
writeTMVar TMVar (NameSearch Transaction)
nameSearchCacheVar (Int -> Names -> NameSearch Transaction
forall (m :: * -> *). Applicative m => Int -> Names -> NameSearch m
NameSearch.makeNameSearch Int
hl Names
newNames)
          -- Re-check everything with the new names and ppe
          Lsp ()
VFS.markAllFilesDirty
          STM () -> Lsp ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically do
            TMVar CompletionTree -> CompletionTree -> STM ()
forall a. TMVar a -> a -> STM ()
writeTMVar TMVar CompletionTree
completionsVar (Names -> CompletionTree
namesToCompletionTree Names
newNames)
          STM ProjectPathIds -> Maybe (Branch IO) -> Lsp a
forall a. STM ProjectPathIds -> Maybe (Branch IO) -> Lsp a
loop STM ProjectPathIds
signalChanges (Branch IO -> Maybe (Branch IO)
forall a. a -> Maybe a
Just Branch IO
newBranch)
    -- Waits for a possible change, then checks if there's actually any difference to the branches we care about.
    -- If so, returns the new branch, otherwise Nothing.
    getChanges :: STM PP.ProjectPathIds -> Maybe (Branch.Branch IO) -> Lsp (ProjectPath, Maybe (Branch.Branch IO))
    getChanges :: STM ProjectPathIds
-> Maybe (Branch IO) -> Lsp (ProjectPath, Maybe (Branch IO))
getChanges STM ProjectPathIds
signalChanges Maybe (Branch IO)
currentBranch = do
      Env {Codebase IO Symbol Ann
$sel:codebase:Env :: Env -> Codebase IO Symbol Ann
codebase :: Codebase IO Symbol Ann
codebase} <- Lsp Env
forall r (m :: * -> *). MonadReader r m => m r
ask
      ProjectPathIds
ppIds <- STM ProjectPathIds -> Lsp ProjectPathIds
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically STM ProjectPathIds
signalChanges
      ProjectPath
pp <- IO ProjectPath -> Lsp ProjectPath
forall a. IO a -> Lsp a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ProjectPath -> Lsp ProjectPath)
-> (Transaction ProjectPath -> IO ProjectPath)
-> Transaction ProjectPath
-> Lsp ProjectPath
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Codebase IO Symbol Ann -> Transaction ProjectPath -> IO ProjectPath
forall (m :: * -> *) v a b.
MonadIO m =>
Codebase m v a -> Transaction b -> m b
Codebase.runTransaction Codebase IO Symbol Ann
codebase (Transaction ProjectPath -> Lsp ProjectPath)
-> Transaction ProjectPath -> Lsp ProjectPath
forall a b. (a -> b) -> a -> b
$ ProjectPathIds -> Transaction ProjectPath
Codebase.resolveProjectPathIds ProjectPathIds
ppIds
      STM () -> Lsp ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Lsp ()) -> STM () -> Lsp ()
forall a b. (a -> b) -> a -> b
$ TMVar ProjectPath -> ProjectPath -> STM ()
forall a. TMVar a -> a -> STM ()
writeTMVar TMVar ProjectPath
currentPathVar ProjectPath
pp
      Branch IO
newBranch <- (Maybe (Branch IO) -> Branch IO)
-> Lsp (Maybe (Branch IO)) -> Lsp (Branch IO)
forall a b. (a -> b) -> Lsp a -> Lsp b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Branch IO -> Maybe (Branch IO) -> Branch IO
forall a. a -> Maybe a -> a
fromMaybe Branch IO
forall (m :: * -> *). Branch m
Branch.empty) (Lsp (Maybe (Branch IO)) -> Lsp (Branch IO))
-> (IO (Maybe (Branch IO)) -> Lsp (Maybe (Branch IO)))
-> IO (Maybe (Branch IO))
-> Lsp (Branch IO)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Maybe (Branch IO)) -> Lsp (Maybe (Branch IO))
forall a. IO a -> Lsp a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (Branch IO)) -> Lsp (Branch IO))
-> IO (Maybe (Branch IO)) -> Lsp (Branch IO)
forall a b. (a -> b) -> a -> b
$ Codebase IO Symbol Ann -> ProjectPath -> IO (Maybe (Branch IO))
forall (m :: * -> *) v a.
MonadIO m =>
Codebase m v a -> ProjectPath -> m (Maybe (Branch m))
Codebase.getBranchAtProjectPath Codebase IO Symbol Ann
codebase ProjectPath
pp
      pure $ (ProjectPath
pp, if Branch IO -> Maybe (Branch IO)
forall a. a -> Maybe a
Just Branch IO
newBranch Maybe (Branch IO) -> Maybe (Branch IO) -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe (Branch IO)
currentBranch then Maybe (Branch IO)
forall a. Maybe a
Nothing else Branch IO -> Maybe (Branch IO)
forall a. a -> Maybe a
Just Branch IO
newBranch)
    -- This is added in stm-2.5.1, remove this if we upgrade.
    writeTMVar :: TMVar a -> a -> STM ()
    writeTMVar :: forall a. TMVar a -> a -> STM ()
writeTMVar TMVar a
var a
a =
      TMVar a -> STM (Maybe a)
forall a. TMVar a -> STM (Maybe a)
tryReadTMVar TMVar a
var STM (Maybe a) -> (Maybe a -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe a
Nothing -> TMVar a -> a -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar a
var a
a
        Just a
_ -> STM a -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM a -> STM ()) -> STM a -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar a -> a -> STM a
forall a. TMVar a -> a -> STM a
swapTMVar TMVar a
var a
a