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 (PrettyPrintEnvDecl)
import Unison.PrettyPrintEnvDecl 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
  signalChanges <- Signal ProjectPathIds -> Lsp (STM ProjectPathIds)
forall (m :: * -> *) a. MonadIO m => Signal a -> m (STM a)
Signal.subscribe Signal ProjectPathIds
changeSignal
  loop signalChanges 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, completionsVar} <- Lsp Env
forall r (m :: * -> *). MonadReader r m => m r
ask
      getChanges signalChanges currentBranch >>= \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
          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 = Namer -> Suffixifier -> PrettyPrintEnvDecl
PPED.makePPED (Int -> Names -> Namer
PPE.hqNamer Int
hl Names
newNames) (Names -> Suffixifier
PPE.suffixifyByHash Names
newNames)
          atomically $ do
            writeTMVar currentPathVar newPP
            writeTMVar currentNamesVar newNames
            writeTMVar ppedVar pped
            writeTMVar nameSearchCacheVar (NameSearch.makeNameSearch hl newNames)
          -- Re-check everything with the new names and ppe
          VFS.markAllFilesDirty
          atomically do
            writeTMVar completionsVar (namesToCompletionTree newNames)
          loop signalChanges (Just 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} <- Lsp Env
forall r (m :: * -> *). MonadReader r m => m r
ask
      ppIds <- atomically signalChanges
      pp <- liftIO . Codebase.runTransaction codebase $ Codebase.resolveProjectPathIds ppIds
      atomically $ writeTMVar currentPathVar pp
      newBranch <- fmap (fromMaybe Branch.empty) . liftIO $ Codebase.getBranchAtProjectPath codebase pp
      pure $ (pp, if Just newBranch == currentBranch then Nothing else Just 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