Extract out ParserState only used by exclusive ref handlers

crates/ark/src/lsp/documents.rs

@@ -5,9 +5,6 @@
 //
 //
 
-use std::sync::Arc;
-use std::sync::Mutex;
-
 use anyhow::*;
 use ropey::Rope;
 use tower_lsp::lsp_types::DidChangeTextDocumentParams;
@@ -42,31 +39,12 @@ pub struct Document {
     // The document's textual contents.
     pub contents: Rope,
 
-    // The version of the document we last synchronized with.
-    // None if the document hasn't been synchronized yet.
-    pub version: Option<i32>,
-
     // The document's AST.
     pub ast: Tree,
 
-    // The parser used to generate the AST. Currently only used in `did_change()`.
-    // In principle this should be out of the world state and owned by the main
-    // loop in a separate hashmap but it's much easier to store it here in the
-    // document.
-    //
-    // Under a mutex because `Parser` is not clonable, which is needed for
-    // snapshots. Write handlers (i.e. `did_change()` might need an exclusive
-    // reference while background threads have acquired snapshots of the world
-    // state, so this can't be an `RwLock`.
-    //
-    // SAFETY: Background threads holding clones of the world state or documents
-    // should never access `parser`.
-    //
-    // TODO: This could also be in a `SyncUnsafeCell` instead of a mutex as we
-    // assume only one writer at a time which must have an exclusive reference
-    // to the world state, and no concurrent reads from clones. However that's
-    // currently an unstable feature.
-    parser: Arc<Mutex<Parser>>,
+    // The version of the document we last synchronized with.
+    // None if the document hasn't been synchronized yet.
+    pub version: Option<i32>,
 }
 
 impl std::fmt::Debug for Document {
@@ -80,28 +58,27 @@ impl std::fmt::Debug for Document {
 
 impl Document {
     pub fn new(contents: &str, version: Option<i32>) -> Self {
-        // create initial document from rope
-        let document = Rope::from(contents);
-
+        // A one-shot parser, assumes the `Document` won't be incrementally reparsed.
+        // Useful for testing, `with_document()`, and `index_file()`.
         let language = tree_sitter_r::language();
-
-        // create a parser for this document
         let mut parser = Parser::new();
+        parser.set_language(&language).unwrap();
 
-        parser
-            .set_language(&language)
-            .expect("failed to create parser");
+        Self::new_with_parser(contents, &mut parser, version)
+    }
+
+    pub fn new_with_parser(contents: &str, parser: &mut Parser, version: Option<i32>) -> Self {
+        let document = Rope::from(contents);
         let ast = parser.parse(contents, None).unwrap();
 
         Self {
             contents: document,
             version,
-            parser: Arc::new(Mutex::new(parser)),
             ast,
         }
     }
 
-    pub fn on_did_change(&mut self, params: &DidChangeTextDocumentParams) {
+    pub fn on_did_change(&mut self, parser: &mut Parser, params: &DidChangeTextDocumentParams) {
         let new_version = params.text_document.version;
 
         // Check for out-of-order change notifications
@@ -118,7 +95,7 @@ impl Document {
         }
 
         for event in &params.content_changes {
-            if let Err(err) = self.update(event) {
+            if let Err(err) = self.update(parser, event) {
                 panic!("Failed to update document: {err:?}");
             }
         }
@@ -127,7 +104,11 @@ impl Document {
         self.version = Some(new_version);
     }
 
-    fn update(&mut self, change: &TextDocumentContentChangeEvent) -> Result<()> {
+    fn update(
+        &mut self,
+        parser: &mut Parser,
+        change: &TextDocumentContentChangeEvent,
+    ) -> Result<()> {
         // Extract edit range. Nothing to do if there wasn't an edit.
         let range = match change.range {
             Some(r) => r,
@@ -174,11 +155,7 @@ impl Document {
         let contents = &self.contents;
         let callback = &mut |byte, point| Self::parse_callback(contents, byte, point);
 
-        let ast = self
-            .parser
-            .lock()
-            .unwrap()
-            .parse_with(callback, Some(&self.ast));
+        let ast = parser.parse_with(callback, Some(&self.ast));
         self.ast = ast.unwrap();
 
         Ok(())

crates/ark/src/lsp/main_loop.rs

@@ -26,6 +26,7 @@ use crate::lsp::backend::LspResponse;
 use crate::lsp::diagnostics;
 use crate::lsp::documents::Document;
 use crate::lsp::handlers;
+use crate::lsp::state::ParserState;
 use crate::lsp::state::WorldState;
 use crate::lsp::state_handlers;
 use crate::lsp::state_handlers::ConsoleInputs;
@@ -85,6 +86,10 @@ pub(crate) struct GlobalState {
     /// (clones) to handlers.
     world: WorldState,
 
+    /// The state containing tree-sitter parsers for documents contained in the
+    /// `WorldState`. Only used in exclusive ref handlers, and is not cloneable.
+    parsers: ParserState,
+
     /// LSP client shared with tower-lsp and the log loop
     client: Client,
 
@@ -125,6 +130,7 @@ impl GlobalState {
 
         Self {
             world: WorldState::default(),
+            parsers: ParserState::new(),
             client,
             events_tx,
             events_rx,
@@ -206,16 +212,16 @@ impl GlobalState {
                             // TODO: Re-index the changed files.
                         },
                         LspNotification::DidOpenTextDocument(params) => {
-                            state_handlers::did_open(params, &mut self.world)?;
+                            state_handlers::did_open(params, &mut self.world, &mut self.parsers)?;
                         },
                         LspNotification::DidChangeTextDocument(params) => {
-                            state_handlers::did_change(params, &mut self.world)?;
+                            state_handlers::did_change(params, &mut self.world, &mut self.parsers)?;
                         },
                         LspNotification::DidSaveTextDocument(_params) => {
                             // Currently ignored
                         },
                         LspNotification::DidCloseTextDocument(params) => {
-                            state_handlers::did_close(params, &mut self.world)?;
+                            state_handlers::did_close(params, &mut self.world, &mut self.parsers)?;
                         },
                     }
                 },

crates/ark/src/lsp/state.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::path::Path;
 
 use anyhow::anyhow;
+use tree_sitter::Parser;
 use url::Url;
 
 use crate::lsp::documents::Document;
@@ -49,6 +50,13 @@ pub(crate) struct Workspace {
     pub folders: Vec<Url>,
 }
 
+/// The set of tree-sitter document parsers managed by the `GlobalState`. Unlike
+/// `WorldState`, `ParserState` cannot be cloned and is only accessed by exclusive
+/// handlers.
+pub(crate) struct ParserState {
+    inner: HashMap<Url, Parser>,
+}
+
 impl WorldState {
     pub(crate) fn get_document(&self, uri: &Url) -> anyhow::Result<&Document> {
         if let Some(doc) = self.documents.get(uri) {
@@ -67,6 +75,30 @@ impl WorldState {
     }
 }
 
+impl ParserState {
+    pub(crate) fn new() -> Self {
+        Self {
+            inner: HashMap::new(),
+        }
+    }
+
+    pub(crate) fn insert(&mut self, uri: Url, parser: Parser) -> Option<Parser> {
+        self.inner.insert(uri, parser)
+    }
+
+    pub(crate) fn remove(&mut self, uri: &Url) -> Option<Parser> {
+        self.inner.remove(uri)
+    }
+
+    pub(crate) fn get_mut(&mut self, uri: &Url) -> anyhow::Result<&mut Parser> {
+        if let Some(parser) = self.inner.get_mut(uri) {
+            Ok(parser)
+        } else {
+            Err(anyhow!("Can't find parser for URI {uri}"))
+        }
+    }
+}
+
 pub(crate) fn with_document<T, F>(
     path: &Path,
     state: &WorldState,

crates/ark/src/lsp/state_handlers.rs

@@ -27,11 +27,13 @@ use tower_lsp::lsp_types::TextDocumentSyncKind;
 use tower_lsp::lsp_types::WorkDoneProgressOptions;
 use tower_lsp::lsp_types::WorkspaceFoldersServerCapabilities;
 use tower_lsp::lsp_types::WorkspaceServerCapabilities;
+use tree_sitter::Parser;
 
 use crate::lsp;
 use crate::lsp::documents::Document;
 use crate::lsp::encoding::get_position_encoding_kind;
 use crate::lsp::indexer;
+use crate::lsp::state::ParserState;
 use crate::lsp::state::WorldState;
 
 // Handlers that mutate the world state
@@ -127,12 +129,19 @@ pub(crate) fn initialize(
 pub(crate) fn did_open(
     params: DidOpenTextDocumentParams,
     state: &mut WorldState,
+    parsers: &mut ParserState,
 ) -> anyhow::Result<()> {
     let contents = params.text_document.text.as_str();
     let uri = params.text_document.uri;
     let version = params.text_document.version;
 
-    let document = Document::new(contents, Some(version));
+    let language = tree_sitter_r::language();
+    let mut parser = Parser::new();
+    parser.set_language(&language).unwrap();
+
+    let document = Document::new_with_parser(contents, &mut parser, Some(version));
+
+    parsers.insert(uri.clone(), parser);
     state.documents.insert(uri.clone(), document.clone());
 
     lsp::spawn_diagnostics_refresh(uri, document, state.clone());
@@ -143,12 +152,14 @@ pub(crate) fn did_open(
 pub(crate) fn did_change(
     params: DidChangeTextDocumentParams,
     state: &mut WorldState,
+    parsers: &mut ParserState,
 ) -> anyhow::Result<()> {
     let uri = &params.text_document.uri;
     let doc = state.get_document_mut(uri)?;
+    let mut parser = parsers.get_mut(uri)?;
 
     // Respond to document updates
-    doc.on_did_change(&params);
+    doc.on_did_change(&mut parser, &params);
 
     // Update index
     if let Ok(path) = uri.to_file_path() {
@@ -167,6 +178,7 @@ pub(crate) fn did_change(
 pub(crate) fn did_close(
     params: DidCloseTextDocumentParams,
     state: &mut WorldState,
+    parsers: &mut ParserState,
 ) -> anyhow::Result<()> {
     let uri = params.text_document.uri;
 
@@ -178,6 +190,10 @@ pub(crate) fn did_close(
         .remove(&uri)
         .ok_or(anyhow!("Failed to remove document for URI: {uri}"))?;
 
+    parsers
+        .remove(&uri)
+        .ok_or(anyhow!("Failed to remove parser for URI: {uri}"))?;
+
     lsp::log_info!("did_close(): closed document with URI: '{uri}'.");
 
     Ok(())