state: share document updates and transaction management with WASI hosts

pyrefly/lib/lsp/non_wasm/lsp.rs

@@ -9,15 +9,12 @@
 
 use lsp_server::RequestId;
 use lsp_server::ResponseError;
-use lsp_types::TextDocumentContentChangeEvent;
-use ruff_source_file::LineIndex;
-use ruff_source_file::OneIndexed;
-use ruff_source_file::SourceLocation;
 use serde::de::DeserializeOwned;
 
 use crate::lsp::non_wasm::protocol::Notification;
 use crate::lsp::non_wasm::protocol::Request;
 use crate::lsp::non_wasm::protocol::Response;
+pub use crate::state::document::apply_change_events;
 
 pub fn as_notification<T>(x: &Notification) -> Option<Result<T::Params, serde_json::Error>>
 where
@@ -99,46 +96,3 @@ where
         },
     }
 }
-
-pub fn apply_change_events(original: &str, changes: Vec<TextDocumentContentChangeEvent>) -> String {
-    /// Convert lsp_types::Position to usize index for a given text.
-    fn position_to_usize(
-        position: lsp_types::Position,
-        index: &LineIndex,
-        source_text: &str,
-    ) -> usize {
-        let source_location = SourceLocation {
-            line: OneIndexed::from_zero_indexed(position.line as usize),
-            character_offset: OneIndexed::from_zero_indexed(position.character as usize),
-        };
-        let text_size = index.offset(
-            source_location,
-            source_text,
-            ruff_source_file::PositionEncoding::Utf16,
-        );
-        text_size.to_usize()
-    }
-
-    let mut result = original.to_owned();
-    for change in changes {
-        let (range, text) = match change {
-            TextDocumentContentChangeEvent::TextDocumentContentChangePartial(partial) => {
-                (Some(partial.range), partial.text)
-            }
-            TextDocumentContentChangeEvent::TextDocumentContentChangeWholeDocument(whole) => {
-                (None, whole.text)
-            }
-        };
-        // If no range is given, we can full text replace.
-        match range {
-            None => result = text,
-            Some(range) => {
-                let index = LineIndex::from_source_text(&result);
-                let start = position_to_usize(range.start, &index, &result);
-                let end = position_to_usize(range.end, &index, &result);
-                result.replace_range(start..end, &text);
-            }
-        }
-    }
-    result
-}

pyrefly/lib/lsp/non_wasm/server.rs

@@ -363,6 +363,10 @@ use crate::lsp::wasm::provide_type::ProvideTypeParams;
 use crate::lsp::wasm::provide_type::ProvideTypeResponse;
 use crate::lsp::wasm::provide_type::provide_type;
 use crate::module::bundled::BundledStub;
+use crate::state::document::DocumentVersions;
+use crate::state::document::OpenFiles;
+use crate::state::document::change_document;
+use crate::state::document::open_document;
 use crate::state::errors::BaselineApplyResult;
 use crate::state::load::Load;
 use crate::state::load::LspFile;
@@ -1175,7 +1179,7 @@ pub struct Server {
     /// Notebook cell URIs are entirely arbitrary, and any URI received from the language client
     /// should be mapped through here in case they correspond to a cell.
     open_notebook_cells: RwLock<HashMap<Uri, PathBuf>>,
-    open_files: RwLock<HashMap<PathBuf, Arc<LspFile>>>,
+    open_files: OpenFiles,
     /// Files saved since the last recheck started, awaiting a baseline prune, each held
     /// with the buffer that was current at the save. A recheck claims the whole map
     /// before it runs, so the state it commits reflects everything it claimed; files
@@ -1214,7 +1218,7 @@ pub struct Server {
     /// Exact file paths that already have a dedicated, permanent watcher registration.
     /// These registrations are additive and never removed, so this set only grows.
     watched_exact_paths: Mutex<SmallSet<PathBuf>>,
-    version_info: Mutex<HashMap<PathBuf, i32>>,
+    version_info: DocumentVersions,
     id: Uuid,
     /// The surface/entrypoint for the language server (`--from` CLI arg)
     surface: Option<String>,
@@ -4301,8 +4305,13 @@ impl Server {
         } else {
             None
         };
-        self.version_info.lock().insert(path.clone(), version);
-        self.open_files.write().insert(path.clone(), contents);
+        open_document(
+            &self.open_files,
+            &self.version_info,
+            path.clone(),
+            version,
+            contents,
+        );
         self.queue_source_db_rebuild_and_recheck(telemetry, telemetry_event, false);
         if !subsequent_mutation {
             info!(
@@ -4338,8 +4347,9 @@ impl Server {
         let VersionedTextDocumentIdentifier {
             text_document_identifier: TextDocumentIdentifier { uri },
             version,
-        } = params.text_document;
-        let Some(file_path) = self.path_for_uri(&uri) else {
+        } = &params.text_document;
+        let version = *version;
+        let Some(file_path) = self.path_for_uri(uri) else {
             return Err(anyhow::anyhow!(
                 "Received textDocument/didChange for unknown uri: {uri}"
             ));
@@ -4356,39 +4366,20 @@ impl Server {
             );
         }
         drop(version_info);
-        let mut lock = self.open_files.write();
-        let Some(original) = lock.get_mut(&file_path) else {
-            return Err(anyhow::anyhow!(
-                "File not found in open_files: {}",
-                file_path.display()
-            ));
-        };
-        let language_id = match original.as_ref() {
-            LspFile::Source { language_id, .. } => Arc::clone(language_id),
-            LspFile::Notebook(_) => {
-                return Err(anyhow::anyhow!(
-                    "Expected text file for {}, but got notebook",
-                    uri
-                ));
-            }
-        };
-        *original = Arc::new(LspFile::Source {
-            contents: Arc::new(apply_change_events(
-                original.get_string(),
-                params.content_changes,
-            )),
-            language_id,
-        });
-        drop(lock);
-        // Update version_info only after the mutation has fully succeeded.
-        self.version_info.lock().insert(file_path.clone(), version);
+        change_document(
+            &self.open_files,
+            &self.version_info,
+            &file_path,
+            &params.text_document,
+            params.content_changes,
+        )?;
         if !subsequent_mutation {
             info!(
                 "File {} changed, prepare to validate open files.",
                 file_path.display()
             );
             if let Ok(handle) = self.make_handle_if_enabled(
-                &uri,
+                uri,
                 Some(DidChangeTextDocumentNotification::METHOD.as_str()),
             ) {
                 self.currently_streaming_diagnostics_for_handles

pyrefly/lib/lsp/non_wasm/transaction_manager.rs

@@ -5,166 +5,4 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-use pyrefly_util::telemetry::TelemetryEvent;
-
-use crate::state::require::Require;
-use crate::state::state::CommittingTransaction;
-use crate::state::state::State;
-use crate::state::state::Transaction;
-use crate::state::state::TransactionData;
-
-/// `TransactionManager` aims to always produce a transaction that contains the up-to-date
-/// in-memory contents.
-#[derive(Default)]
-pub struct TransactionManager<'a> {
-    /// Invariant:
-    /// If it's None, then the main `State` already contains up-to-date checked content
-    /// of all in-memory files.
-    /// Otherwise, it will contain up-to-date checked content of all in-memory files.
-    saved_state: Option<TransactionData<'a>>,
-}
-
-impl<'a> TransactionManager<'a> {
-    #[expect(clippy::result_large_err)] // Both results are basically the same size
-    /// Produce a possibly committable transaction in order to recheck in-memory files.
-    pub fn get_possibly_committable_transaction(
-        &mut self,
-        state: &'a State,
-    ) -> Result<CommittingTransaction<'a>, Transaction<'a>> {
-        // If there is no ongoing recheck due to on-disk changes, we should prefer to commit
-        // the in-memory changes into the main state.
-        if let Some(transaction) = state.try_new_committable_transaction(Require::Exports, None) {
-            // If we can commit in-memory changes, then there is no point of holding the
-            // non-committable transaction with a possibly outdated view of the `ReadableState`
-            // so we can destroy the saved state.
-            self.saved_state = None;
-            Ok(transaction)
-        } else {
-            // If there is an ongoing recheck, trying to get a committable transaction will block
-            // until the recheck is finished. This is bad for perceived perf. Therefore, we will
-            // temporarily use a non-committable transaction to hold the information that's necessary
-            // to power IDE services.
-            Err(self.non_committable_transaction(state))
-        }
-    }
-
-    /// Produce a `Transaction` to power readonly IDE services.
-    /// This transaction will never be able to be committed.
-    /// After using it, the state should be saved by calling the `save` method.
-    ///
-    /// The `Transaction` will always contain the handles of all open files with the latest content.
-    /// It might be created fresh from state, or reused from previously saved state.
-    ///
-    /// If we were unable to restore a transaction from saved state, we create a fresh transaction.
-    /// Callers may need to re-validate open files in this case.
-    pub fn non_committable_transaction(&mut self, state: &'a State) -> Transaction<'a> {
-        let previous_blocking = match self.saved_state.take() {
-            Some(saved_state) => match saved_state.restore() {
-                Ok(tx) => return tx,
-                Err(blocked) => Some(blocked),
-            },
-            None => None,
-        };
-        let mut tx = state.transaction();
-        tx.set_fresh();
-        if let Some(d) = previous_blocking {
-            tx.add_locked_blocking_duration(d);
-        }
-        tx
-    }
-
-    /// This function should be called once we finished using transaction for an LSP request.
-    pub fn save(&mut self, transaction: Transaction<'a>, telemetry: &mut TelemetryEvent) {
-        self.saved_state = Some(transaction.save(telemetry))
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::time::Instant;
-
-    use dupe::Dupe;
-    use pyrefly_build::handle::Handle;
-    use pyrefly_python::module_name::ModuleName;
-    use pyrefly_util::telemetry::QueueName;
-    use pyrefly_util::telemetry::TelemetryEventKind;
-    use pyrefly_util::telemetry::TelemetryServerState;
-    use pyrefly_util::thread_pool::TEST_THREAD_COUNT;
-    use uuid::Uuid;
-
-    use super::*;
-    use crate::module::finder::DirEntryCache;
-    use crate::module::finder::FindImportOptions;
-    use crate::module::finder::find_import;
-    use crate::test::util::TestEnv;
-
-    /// A recheck cancels the in-flight reads that block its commit. That cancellation belongs
-    /// to the request being aborted, so it should *not* survive into the restored transaction.
-    /// Clearing it must also keep the work the saved transaction already did, so the restored
-    /// transaction reuses the cached stdlib instead of recomputing it.
-    #[test]
-    fn test_restored_transaction_is_not_still_cancelled() {
-        let mut test_env = TestEnv::new();
-        test_env.add("first", "x: int = 1\n");
-        test_env.add("second", "y: int = \"not an int\"\n");
-        let config_file = test_env.config();
-        let sys_info = test_env.sys_info();
-        let state = State::new(test_env.config_finder(), TEST_THREAD_COUNT);
-        let handle = |name: &str| {
-            let name = ModuleName::from_str(name);
-            let path = find_import(
-                &config_file,
-                name,
-                FindImportOptions::new(&DirEntryCache::new()),
-            )
-            .finding()
-            .unwrap();
-            Handle::new(name, path, sys_info.dupe())
-        };
-
-        let mut manager = TransactionManager::default();
-        let mut transaction = manager.non_committable_transaction(&state);
-        transaction.set_memory(test_env.get_memory());
-        transaction.run(&[handle("first")], Require::Everything, None);
-        let old_cancellation = transaction.get_cancellation_handle();
-        old_cancellation.cancel();
-        let mut telemetry = TelemetryEvent::new_task(
-            TelemetryEventKind::InvalidateConfig,
-            TelemetryServerState {
-                has_sourcedb: false,
-                id: Uuid::new_v4(),
-                surface: None,
-                server_start_time: Instant::now(),
-                agent_session_id: None,
-                agent_invocation_id: None,
-                active_experiments: Vec::new(),
-            },
-            QueueName::RecheckQueue,
-            0,
-            Instant::now(),
-        );
-        manager.save(transaction, &mut telemetry);
-
-        let second = handle("second");
-        let mut transaction = manager.saved_state.take().unwrap().restore().unwrap();
-        assert!(old_cancellation.is_cancelled());
-        assert!(
-            !transaction.get_cancellation_handle().is_cancelled(),
-            "restore should replace the saved cancellation handle"
-        );
-        transaction.set_memory(test_env.get_memory());
-        transaction.run(&[second.dupe()], Require::Everything, None);
-        assert!(
-            transaction.compute_stdlib_cached(),
-            "restored transaction should reuse the stdlib computed before the save"
-        );
-        assert_eq!(
-            transaction
-                .get_errors([&second])
-                .collect_errors()
-                .ordinary
-                .len(),
-            1
-        );
-    }
-}
+pub use crate::state::transaction_manager::TransactionManager;

pyrefly/lib/state.rs

@@ -6,6 +6,7 @@
  */
 
 pub mod dirty;
+pub mod document;
 pub mod epoch;
 pub mod errors;
 pub mod ide;
@@ -24,3 +25,4 @@ pub mod state;
 mod step_slot;
 pub mod steps;
 pub mod subscriber;
+pub mod transaction_manager;