LSP startup time speedup and CWD fix

unison-cli/src/Unison/CommandLine/Main.hs

@@ -16,6 +16,7 @@ import Ki qualified
 import System.Console.Haskeline qualified as Line
 import System.IO (hGetEcho, hPutStrLn, hSetEcho, stderr, stdin)
 import System.IO.Error (isDoesNotExistError)
+import U.Codebase.HashTags (CausalHash)
 import U.Codebase.Sqlite.Operations qualified as Operations
 import U.Codebase.Sqlite.Queries qualified as Queries
 import Unison.Auth.CredentialManager (newCredentialManager)
@@ -27,7 +28,7 @@ import Unison.Cli.Pretty (prettyProjectAndBranchName)
 import Unison.Cli.ProjectUtils (projectBranchPathPrism)
 import Unison.Codebase (Codebase)
 import Unison.Codebase qualified as Codebase
-import Unison.Codebase.Branch (Branch)
+import Unison.Codebase.Branch qualified as Branch
 import Unison.Codebase.Editor.HandleInput qualified as HandleInput
 import Unison.Codebase.Editor.Input (Event, Input (..))
 import Unison.Codebase.Editor.Output (Output)
@@ -126,7 +127,7 @@ main ::
   Codebase IO Symbol Ann ->
   Maybe Server.BaseUrl ->
   UCMVersion ->
-  (Branch IO -> STM ()) ->
+  (CausalHash -> STM ()) ->
   (Path.Absolute -> STM ()) ->
   ShouldWatchFiles ->
   IO ()
@@ -154,7 +155,7 @@ main dir welcome initialPath config initialInputs runtime sbRuntime nRuntime cod
           currentRoot <- atomically do
             currentRoot <- readTMVar rootVar
             guard $ Just currentRoot /= lastRoot
-            notifyBranchChange currentRoot
+            notifyBranchChange (Branch.headHash currentRoot)
             pure (Just currentRoot)
           loop currentRoot
     loop Nothing

unison-cli/src/Unison/LSP.hs

@@ -23,8 +23,8 @@ import Language.LSP.VFS
 import Network.Simple.TCP qualified as TCP
 import System.Environment (lookupEnv)
 import System.IO (hPutStrLn)
+import U.Codebase.HashTags
 import Unison.Codebase
-import Unison.Codebase.Branch (Branch)
 import Unison.Codebase.Path qualified as Path
 import Unison.Codebase.Runtime (Runtime)
 import Unison.Debug qualified as Debug
@@ -46,8 +46,6 @@ import Unison.LSP.UCMWorker (ucmWorker)
 import Unison.LSP.VFS qualified as VFS
 import Unison.Parser.Ann
 import Unison.Prelude
-import Unison.PrettyPrintEnvDecl qualified as PPED
-import Unison.Server.NameSearch.FromNames qualified as NameSearch
 import Unison.Symbol
 import UnliftIO
 import UnliftIO.Foreign (Errno (..), eADDRINUSE)
@@ -56,8 +54,8 @@ getLspPort :: IO String
 getLspPort = fromMaybe "5757" <$> lookupEnv "UNISON_LSP_PORT"
 
 -- | Spawn an LSP server on the configured port.
-spawnLsp :: Codebase IO Symbol Ann -> Runtime Symbol -> STM (Branch IO) -> STM (Path.Absolute) -> IO ()
-spawnLsp codebase runtime latestBranch latestPath =
+spawnLsp :: Codebase IO Symbol Ann -> Runtime Symbol -> STM CausalHash -> STM (Path.Absolute) -> IO ()
+spawnLsp codebase runtime latestRootHash latestPath =
   ifEnabled . TCP.withSocketsDo $ do
     lspPort <- getLspPort
     UnliftIO.handleIO (handleFailure lspPort) $ do
@@ -77,7 +75,7 @@ spawnLsp codebase runtime latestBranch latestPath =
           -- different un-saved state for the same file.
           initVFS $ \vfs -> do
             vfsVar <- newMVar vfs
-            void $ runServerWith lspServerLogger lspClientLogger clientInput clientOutput (serverDefinition vfsVar codebase runtime scope latestBranch latestPath)
+            void $ runServerWith lspServerLogger lspClientLogger clientInput clientOutput (serverDefinition vfsVar codebase runtime scope latestRootHash latestPath)
   where
     handleFailure :: String -> IOException -> IO ()
     handleFailure lspPort ioerr =
@@ -107,16 +105,16 @@ serverDefinition ::
   Codebase IO Symbol Ann ->
   Runtime Symbol ->
   Ki.Scope ->
-  STM (Branch IO) ->
+  STM CausalHash ->
   STM (Path.Absolute) ->
   ServerDefinition Config
-serverDefinition vfsVar codebase runtime scope latestBranch latestPath =
+serverDefinition vfsVar codebase runtime scope latestRootHash latestPath =
   ServerDefinition
     { defaultConfig = defaultLSPConfig,
       configSection = "unison",
       parseConfig = Config.parseConfig,
       onConfigChange = Config.updateConfig,
-      doInitialize = lspDoInitialize vfsVar codebase runtime scope latestBranch latestPath,
+      doInitialize = lspDoInitialize vfsVar codebase runtime scope latestRootHash latestPath,
       staticHandlers = lspStaticHandlers,
       interpretHandler = lspInterpretHandler,
       options = lspOptions
@@ -128,26 +126,31 @@ lspDoInitialize ::
   Codebase IO Symbol Ann ->
   Runtime Symbol ->
   Ki.Scope ->
-  STM (Branch IO) ->
+  STM CausalHash ->
   STM (Path.Absolute) ->
   LanguageContextEnv Config ->
   Msg.TMessage 'Msg.Method_Initialize ->
   IO (Either Msg.ResponseError Env)
-lspDoInitialize vfsVar codebase runtime scope latestBranch latestPath lspContext _initMsg = do
-  -- TODO: some of these should probably be MVars so that we correctly wait for names and
-  -- things to be generated before serving requests.
+lspDoInitialize vfsVar codebase runtime scope latestRootHash latestPath lspContext _initMsg = do
   checkedFilesVar <- newTVarIO mempty
   dirtyFilesVar <- newTVarIO mempty
-  ppedCacheVar <- newTVarIO PPED.empty
-  parseNamesCacheVar <- newTVarIO mempty
-  currentPathCacheVar <- newTVarIO Path.absoluteEmpty
+  ppedCacheVar <- newEmptyTMVarIO
+  currentNamesCacheVar <- newEmptyTMVarIO
+  currentPathCacheVar <- newEmptyTMVarIO
   cancellationMapVar <- newTVarIO mempty
-  completionsVar <- newTVarIO mempty
-  nameSearchCacheVar <- newTVarIO $ NameSearch.makeNameSearch 0 mempty
-  let env = Env {ppedCache = readTVarIO ppedCacheVar, parseNamesCache = readTVarIO parseNamesCacheVar, currentPathCache = readTVarIO currentPathCacheVar, nameSearchCache = readTVarIO nameSearchCacheVar, ..}
+  completionsVar <- newEmptyTMVarIO
+  nameSearchCacheVar <- newEmptyTMVarIO
+  let env =
+        Env
+          { ppedCache = atomically $ readTMVar ppedCacheVar,
+            currentNamesCache = atomically $ readTMVar currentNamesCacheVar,
+            currentPathCache = atomically $ readTMVar currentPathCacheVar,
+            nameSearchCache = atomically $ readTMVar nameSearchCacheVar,
+            ..
+          }
   let lspToIO = flip runReaderT lspContext . unLspT . flip runReaderT env . runLspM
   Ki.fork scope (lspToIO Analysis.fileAnalysisWorker)
-  Ki.fork scope (lspToIO $ ucmWorker ppedCacheVar parseNamesCacheVar nameSearchCacheVar latestBranch latestPath)
+  Ki.fork scope (lspToIO $ ucmWorker ppedCacheVar currentNamesCacheVar nameSearchCacheVar currentPathCacheVar latestRootHash latestPath)
   pure $ Right $ env
 
 -- | LSP request handlers that don't register/unregister dynamically

unison-cli/src/Unison/LSP/Completion.hs

@@ -53,7 +53,7 @@ completionHandler m respond =
   respond . maybe (Right $ InL mempty) (Right . InR . InL) =<< runMaybeT do
     let fileUri = (m ^. params . textDocument . uri)
     (range, prefix) <- VFS.completionPrefix (m ^. params . textDocument . uri) (m ^. params . position)
-    ppe <- PPED.suffixifiedPPE <$> lift globalPPED
+    ppe <- PPED.suffixifiedPPE <$> lift currentPPED
     codebaseCompletions <- lift getCodebaseCompletions
     Config {maxCompletions} <- lift getConfig
     let defMatches = matchCompletions codebaseCompletions prefix

unison-cli/src/Unison/LSP/FileAnalysis.hs

@@ -82,7 +82,7 @@ checkFile doc = runMaybeT do
   currentPath <- lift getCurrentPath
   let fileUri = doc ^. uri
   (fileVersion, contents) <- VFS.getFileContents fileUri
-  parseNames <- lift getParseNames
+  parseNames <- lift getCurrentNames
   let sourceName = getUri $ doc ^. uri
   let lexedSource@(srcText, tokens) = (contents, L.lexer (Text.unpack sourceName) (Text.unpack contents))
   let ambientAbilities = []
@@ -158,7 +158,7 @@ fileAnalysisWorker = forever do
 
 analyseFile :: (Foldable f) => Uri -> Text -> f (Note Symbol Ann) -> Lsp ([Diagnostic], [RangedCodeAction])
 analyseFile fileUri srcText notes = do
-  pped <- PPED.suffixifiedPPE <$> LSP.globalPPED
+  pped <- PPED.suffixifiedPPE <$> LSP.currentPPED
   (noteDiags, noteActions) <- analyseNotes fileUri pped (Text.unpack srcText) notes
   pure (noteDiags, noteActions)
 
@@ -167,7 +167,7 @@ analyseFile fileUri srcText notes = do
 computeConflictWarningDiagnostics :: Uri -> FileSummary -> Lsp [Diagnostic]
 computeConflictWarningDiagnostics fileUri fileSummary@FileSummary {fileNames} = do
   let defLocations = fileDefLocations fileSummary
-  conflictedNames <- Names.conflicts <$> getParseNames
+  conflictedNames <- Names.conflicts <$> getCurrentNames
   let locationForName :: Name -> Set Ann
       locationForName name = fold $ Map.lookup (Name.toVar name) defLocations
   let conflictedTermLocations =
@@ -360,7 +360,7 @@ analyseNotes fileUri ppe src notes = do
       | not (isUserBlank v) = pure []
       | otherwise = do
           Env {codebase} <- ask
-          ppe <- PPED.suffixifiedPPE <$> globalPPED
+          ppe <- PPED.suffixifiedPPE <$> currentPPED
           let cleanedTyp = Context.generalizeAndUnTypeVar typ -- TODO: is this right?
           refs <- liftIO . Codebase.runTransaction codebase $ Codebase.termsOfType codebase cleanedTyp
           forMaybe (toList refs) $ \ref -> runMaybeT $ do
@@ -395,7 +395,10 @@ getFileAnalysis uri = do
         writeTVar checkedFilesV $ Map.insert uri mvar checkedFiles
         pure mvar
       Just mvar -> pure mvar
-  atomically (readTMVar tmvar)
+  Debug.debugM Debug.LSP "Waiting on file analysis" uri
+  r <- atomically (readTMVar tmvar)
+  Debug.debugM Debug.LSP "Got file analysis" uri
+  pure r
 
 -- | Build a Names from a file if it's parseable.
 --
@@ -427,7 +430,7 @@ ppedForFile fileUri = do
 
 ppedForFileHelper :: Maybe (UF.UnisonFile Symbol a) -> Maybe (UF.TypecheckedUnisonFile Symbol a) -> Lsp PPED.PrettyPrintEnvDecl
 ppedForFileHelper uf tf = do
-  codebasePPED <- globalPPED
+  codebasePPED <- currentPPED
   hashLen <- asks codebase >>= \codebase -> liftIO (Codebase.runTransaction codebase Codebase.hashLength)
   pure $ case (uf, tf) of
     (Nothing, Nothing) -> codebasePPED

unison-cli/src/Unison/LSP/Types.hs

@@ -69,7 +69,7 @@ data Env = Env
   { -- contains handlers for talking to the client.
     lspContext :: LanguageContextEnv Config,
     codebase :: Codebase IO Symbol Ann,
-    parseNamesCache :: IO Names,
+    currentNamesCache :: IO Names,
     ppedCache :: IO PrettyPrintEnvDecl,
     nameSearchCache :: IO (NameSearch Sqlite.Transaction),
     currentPathCache :: IO Path.Absolute,
@@ -83,7 +83,7 @@ data Env = Env
     -- A map  of request IDs to an action which kills that request.
     cancellationMapVar :: TVar (Map (Int32 |? Text) (IO ())),
     -- A lazily computed map of all valid completion suffixes from the current path.
-    completionsVar :: TVar CompletionTree,
+    completionsVar :: TMVar CompletionTree,
     scope :: Ki.Scope
   }
 
@@ -133,16 +133,16 @@ getCurrentPath :: Lsp Path.Absolute
 getCurrentPath = asks currentPathCache >>= liftIO
 
 getCodebaseCompletions :: Lsp CompletionTree
-getCodebaseCompletions = asks completionsVar >>= readTVarIO
+getCodebaseCompletions = asks completionsVar >>= atomically . readTMVar
 
-globalPPED :: Lsp PrettyPrintEnvDecl
-globalPPED = asks ppedCache >>= liftIO
+currentPPED :: Lsp PrettyPrintEnvDecl
+currentPPED = asks ppedCache >>= liftIO
 
 getNameSearch :: Lsp (NameSearch Sqlite.Transaction)
 getNameSearch = asks nameSearchCache >>= liftIO
 
-getParseNames :: Lsp Names
-getParseNames = asks parseNamesCache >>= liftIO
+getCurrentNames :: Lsp Names
+getCurrentNames = asks currentNamesCache >>= liftIO
 
 data Config = Config
   { formattingWidth :: Int,

unison-cli/src/Unison/LSP/UCMWorker.hs

@@ -1,8 +1,8 @@
 module Unison.LSP.UCMWorker where
 
 import Control.Monad.Reader
+import U.Codebase.HashTags
 import Unison.Codebase qualified as Codebase
-import Unison.Codebase.Branch (Branch)
 import Unison.Codebase.Branch qualified as Branch
 import Unison.Codebase.Branch.Names qualified as Branch
 import Unison.Codebase.Path qualified as Path
@@ -21,37 +21,48 @@ import UnliftIO.STM
 
 -- | Watches for state changes in UCM and updates cached LSP state accordingly
 ucmWorker ::
-  TVar PrettyPrintEnvDecl ->
-  TVar Names ->
-  TVar (NameSearch Sqlite.Transaction) ->
-  STM (Branch IO) ->
+  TMVar PrettyPrintEnvDecl ->
+  TMVar Names ->
+  TMVar (NameSearch Sqlite.Transaction) ->
+  TMVar Path.Absolute ->
+  STM CausalHash ->
   STM Path.Absolute ->
   Lsp ()
-ucmWorker ppedVar parseNamesVar nameSearchCacheVar getLatestRoot getLatestPath = do
+ucmWorker ppedVar parseNamesVar nameSearchCacheVar currentPathVar getLatestRoot getLatestPath = do
   Env {codebase, completionsVar} <- ask
-  let loop :: (Branch IO, Path.Absolute) -> Lsp a
+  let loop :: (CausalHash, Path.Absolute) -> Lsp a
       loop (currentRoot, currentPath) = do
         Debug.debugM Debug.LSP "LSP path: " currentPath
-        let currentBranch0 = Branch.getAt0 (Path.unabsolute currentPath) (Branch.head currentRoot)
+        currentBranch0 <- fmap Branch.head . liftIO $ (Codebase.getBranchAtPath codebase currentPath)
         let parseNames = Branch.toNames currentBranch0
         hl <- liftIO $ Codebase.runTransaction codebase Codebase.hashLength
         let pped = PPED.makePPED (PPE.hqNamer hl parseNames) (PPE.suffixifyByHash parseNames)
         atomically $ do
-          writeTVar parseNamesVar parseNames
-          writeTVar ppedVar pped
-          writeTVar nameSearchCacheVar (NameSearch.makeNameSearch hl parseNames)
+          writeTMVar currentPathVar currentPath
+          writeTMVar parseNamesVar parseNames
+          writeTMVar ppedVar pped
+          writeTMVar nameSearchCacheVar (NameSearch.makeNameSearch hl parseNames)
         -- Re-check everything with the new names and ppe
         VFS.markAllFilesDirty
         atomically do
-          writeTVar completionsVar (namesToCompletionTree parseNames)
+          writeTMVar completionsVar (namesToCompletionTree parseNames)
+        Debug.debugLogM Debug.LSP "LSP Initialized"
         latest <- atomically $ do
           latestRoot <- getLatestRoot
           latestPath <- getLatestPath
           guard $ (currentRoot /= latestRoot || currentPath /= latestPath)
           pure (latestRoot, latestPath)
+        Debug.debugLogM Debug.LSP "LSP Change detected"
         loop latest
-
-  -- Bootstrap manually from codebase just in case we're in headless mode and don't get any
-  -- updates from UCM
-  rootBranch <- liftIO $ Codebase.getRootBranch codebase
-  loop (rootBranch, Path.absoluteEmpty)
+  (rootBranch, currentPath) <- atomically $ do
+    rootBranch <- getLatestRoot
+    currentPath <- getLatestPath
+    pure (rootBranch, currentPath)
+  loop (rootBranch, currentPath)
+  where
+    -- This is added in stm-2.5.1, remove this if we upgrade.
+    writeTMVar :: TMVar a -> a -> STM ()
+    writeTMVar var a =
+      tryReadTMVar var >>= \case
+        Nothing -> putTMVar var a
+        Just _ -> void $ swapTMVar var a

unison-cli/unison/Main.hs

@@ -50,10 +50,10 @@ import System.IO.CodePage (withCP65001)
 import System.IO.Error (catchIOError)
 import System.IO.Temp qualified as Temp
 import System.Path qualified as Path
+import U.Codebase.HashTags (CausalHash)
 import U.Codebase.Sqlite.Queries qualified as Queries
 import Unison.Codebase (Codebase, CodebasePath)
 import Unison.Codebase qualified as Codebase
-import Unison.Codebase.Branch (Branch)
 import Unison.Codebase.Editor.Input qualified as Input
 import Unison.Codebase.Execute (execute)
 import Unison.Codebase.Init (CodebaseInitOptions (..), InitError (..), InitResult (..), SpecifiedCodebase (..))
@@ -280,22 +280,20 @@ main = do
                         segments <- Codebase.runTransaction theCodebase Queries.expectMostRecentNamespace
                         pure (Path.Absolute (Path.fromList (map NameSegment.NameSegment segments)))
                   Headless -> pure $ fromMaybe defaultInitialPath mayStartingPath
-                rootVar <- newEmptyTMVarIO
+                rootCausalHash <- Codebase.runTransaction theCodebase (Queries.expectNamespaceRoot >>= Queries.expectCausalHash)
+                rootCausalHashVar <- newTVarIO rootCausalHash
                 pathVar <- newTVarIO startingPath
-                let notifyOnRootChanges :: Branch IO -> STM ()
+                let notifyOnRootChanges :: CausalHash -> STM ()
                     notifyOnRootChanges b = do
-                      isEmpty <- isEmptyTMVar rootVar
-                      if isEmpty
-                        then putTMVar rootVar b
-                        else void $ swapTMVar rootVar b
+                      writeTVar rootCausalHashVar b
                 let notifyOnPathChanges :: Path.Absolute -> STM ()
                     notifyOnPathChanges = writeTVar pathVar
                 -- Unfortunately, the windows IO manager on GHC 8.* is prone to just hanging forever
                 -- when waiting for input on handles, so if we listen for LSP connections it will
                 -- prevent UCM from shutting down properly. Hopefully we can re-enable LSP on
                 -- Windows when we move to GHC 9.*
                 -- https://gitlab.haskell.org/ghc/ghc/-/merge_requests/1224
-                void . Ki.fork scope $ LSP.spawnLsp theCodebase runtime (readTMVar rootVar) (readTVar pathVar)
+                void . Ki.fork scope $ LSP.spawnLsp theCodebase runtime (readTVar rootCausalHashVar) (readTVar pathVar)
                 Server.startServer (Backend.BackendEnv {Backend.useNamesIndex = False}) codebaseServerOpts sbRuntime theCodebase $ \baseUrl -> do
                   case exitOption of
                     DoNotExit -> do
@@ -503,7 +501,7 @@ launch ::
   Maybe Server.BaseUrl ->
   Maybe Path.Absolute ->
   InitResult ->
-  (Branch IO -> STM ()) ->
+  (CausalHash -> STM ()) ->
   (Path.Absolute -> STM ()) ->
   CommandLine.ShouldWatchFiles ->
   IO ()