Add Document::url for verbatim wire output

Author: lionel-

crates/ark/src/lsp/document.rs

@@ -11,13 +11,26 @@ use aether_lsp_utils::proto::PositionEncoding;
 use anyhow::anyhow;
 use oak_semantic::semantic_index::SemanticIndex;
 use tower_lsp::lsp_types;
+use tower_lsp::lsp_types::Url;
 use tree_sitter::Parser;
 use tree_sitter::Tree;
 
 use crate::lsp::config::DocumentConfig;
 
+/// Placeholder URL used by `Document::new`. Production paths replace
+/// this with the editor's verbatim URL via `WorldState::insert_document`.
+/// Tests that don't go through `WorldState` leave the placeholder.
+const PLACEHOLDER_URL: &str = "ark://internal/document-url-unset";
+
 #[derive(Clone)]
 pub struct Document {
+    /// The editor's URL for this document, byte-for-byte as received in
+    /// `did_open`. Used for wire output (diagnostics, etc.) so the
+    /// frontend always sees the URI it sent us, never a normalised
+    /// round-trip. `WorldState::insert_document` populates this; the
+    /// default constructor uses a placeholder.
+    pub url: Url,
+
     /// The document's textual contents.
     pub contents: String,
 
@@ -76,6 +89,7 @@ impl Document {
         let line_index = biome_line_index::LineIndex::new(&contents);
 
         Self {
+            url: Url::parse(PLACEHOLDER_URL).expect("placeholder URL parses"),
             contents,
             version,
             ast,

crates/ark/src/lsp/goto_definition.rs

@@ -164,7 +164,7 @@ mod parked_salsa_tests {
 
     fn make_state(uri: &lsp_types::Url, doc: &Document) -> WorldState {
         let mut state = WorldState::default();
-        state.documents.insert(uri.clone(), doc.clone());
+        state.insert_document(uri.clone(), doc.clone());
         state
     }
 
@@ -229,8 +229,8 @@ mod parked_salsa_tests {
         let pkg = Package::from_parts(pkg_root, desc, ns);
 
         let mut state = WorldState::default();
-        state.documents.insert(uri_aaa.clone(), doc_aaa.clone());
-        state.documents.insert(uri_ccc.clone(), doc_ccc.clone());
+        state.insert_document(uri_aaa.clone(), doc_aaa.clone());
+        state.insert_document(uri_ccc.clone(), doc_ccc.clone());
         state.root = Some(SourceRoot::Package(pkg));
 
         // ccc.R sees bbb.R's definition (later in collation)
@@ -283,9 +283,9 @@ mod parked_salsa_tests {
         let pkg = Package::from_parts(pkg_root, desc, ns);
 
         let mut state = WorldState::default();
-        state.documents.insert(uri_aaa.clone(), doc_aaa.clone());
-        state.documents.insert(uri_bbb.clone(), doc_bbb);
-        state.documents.insert(uri_ccc.clone(), doc_ccc.clone());
+        state.insert_document(uri_aaa.clone(), doc_aaa.clone());
+        state.insert_document(uri_bbb.clone(), doc_bbb);
+        state.insert_document(uri_ccc.clone(), doc_ccc.clone());
         state.root = Some(SourceRoot::Package(pkg));
 
         // aaa.R is now last in collation, so it can see bbb.R's definition
@@ -332,7 +332,7 @@ mod parked_salsa_tests {
         let pkg = Package::from_parts(pkg_root, desc, ns);
 
         let mut state = WorldState::default();
-        state.documents.insert(uri_aaa.clone(), doc_aaa.clone());
+        state.insert_document(uri_aaa.clone(), doc_aaa.clone());
         state.root = Some(SourceRoot::Package(pkg));
 
         // Cursor on `helper` inside the function body (line 0, col 16)
@@ -464,7 +464,7 @@ mod parked_salsa_tests {
         // `is.null` is NOT in the INDEX
         let doc_null = Document::new("is.null(1)\n", None);
         let uri_null = lsp_types::Url::from_file_path(pkg_root.join("R/bar.R")).unwrap();
-        state.documents.insert(uri_null.clone(), doc_null.clone());
+        state.insert_document(uri_null.clone(), doc_null.clone());
 
         let params = make_params(uri_null, 0, 0);
         let result = goto_definition(&doc_null, params, &state).unwrap();
@@ -524,8 +524,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
         state.workspace.folders = vec![lsp_types::Url::from_directory_path(dir.path()).unwrap()];
 
         let params = make_params(script_uri, 1, 0);
@@ -551,10 +550,9 @@ mod parked_salsa_tests {
         let script_uri = lsp_types::Url::from_file_path(script_dir.join("script.R")).unwrap();
 
         let mut state = WorldState::default();
-        state.documents.insert(helpers_uri.clone(), helpers_doc);
+        state.insert_document(helpers_uri.clone(), helpers_doc);
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         // Cursor on `helper` (line 1, col 0)
         let params = make_params(script_uri, 1, 0);
@@ -594,7 +592,7 @@ mod parked_salsa_tests {
 
         let mut state = make_state(&script_uri, &script_doc);
         state.library = library;
-        state.documents.insert(helpers_uri.clone(), helpers_doc);
+        state.insert_document(helpers_uri.clone(), helpers_doc);
 
         // `mutate` (line 1) resolves via dplyr, attached by helpers.R's library() call.
         // Package symbol, no NavigationTarget.
@@ -635,8 +633,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
         // helpers.R is intentionally NOT inserted into state.documents
 
         let params = make_params(script_uri, 1, 0);
@@ -756,8 +753,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         // Should resolve without hanging. Both symbols are reachable
         // because a.R is visited first (gets its exports + b.R's exports),
@@ -802,8 +798,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         let helpers_uri = lsp_types::Url::from_file_path(dir.path().join("helpers.R")).unwrap();
 
@@ -849,8 +844,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         // `foo` on line 2 should resolve to the LOCAL definition on line 1,
         // not to the sourced one from helpers.R.
@@ -898,8 +892,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         let helpers_uri = lsp_types::Url::from_file_path(dir.path().join("helpers.R")).unwrap();
         let a_uri = lsp_types::Url::from_file_path(dir.path().join("a.R")).unwrap();
@@ -948,8 +941,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         // `my_var` (line 2) should resolve to its own definition on line 1,
         // not to a Sourced duplicate.
@@ -994,8 +986,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         let (index, file_scope) = state.file_analysis(&script_uri, &script_doc);
 
@@ -1034,8 +1025,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         let b_uri = lsp_types::Url::from_file_path(dir.path().join("b.R")).unwrap();
 
@@ -1064,8 +1054,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         let helpers_uri = lsp_types::Url::from_file_path(dir.path().join("helpers.R")).unwrap();
 
@@ -1104,8 +1093,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         let helpers_uri = lsp_types::Url::from_file_path(dir.path().join("helpers.R")).unwrap();
 
@@ -1139,8 +1127,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         // `helper` on line 2 should resolve to helpers.R (source() shadows local)
         let params = make_params(script_uri, 2, 0);
@@ -1178,8 +1165,7 @@ mod parked_salsa_tests {
 
         let mut state = WorldState::default();
         state
-            .documents
-            .insert(script_uri.clone(), script_doc.clone());
+            .insert_document(script_uri.clone(), script_doc.clone());
 
         // `fn` on line 1 should resolve to the SECOND def in helpers.R (line 1)
         let params = make_params(script_uri, 1, 0);

crates/ark/src/lsp/main_loop.rs

@@ -443,12 +443,7 @@ impl GlobalState {
                 // time: a buffer may have opened or closed since the scan
                 // kicked off. The buffer-drain inside `apply_scan_completed` uses
                 // this set as its watcher-event `skip` argument.
-                let editor_owned: HashSet<FilePath> = self.world
-                    .documents
-                    .keys()
-                    .map(FilePath::from_url)
-                    .collect();
-
+                let editor_owned: HashSet<FilePath> = self.world.documents.keys().cloned().collect();
                 let followups = self.lsp_state.oak_scheduler.apply_scan_completed(
                     &mut self.world.oak,
                     scan,
@@ -989,7 +984,7 @@ async fn process_diagnostics_batch(batch: Vec<RefreshDiagnosticsTask>) {
         futures.push(task::spawn_blocking(move || {
             let _span = tracing::info_span!("diagnostics_refresh", uri = %uri).entered();
 
-            if let Some(document) = state.documents.get(&uri) {
+            if let Some(document) = state.documents.get(&FilePath::from_url(&uri)) {
                 // Special case testthat-specific behaviour. This is a simple
                 // stopgap approach that has some false positives (e.g. when we
                 // work on testthat itself the flag will always be true), but
@@ -1125,8 +1120,8 @@ pub(crate) fn diagnostics_refresh_all(state: WorldState) {
         n = state.documents.len()
     );
 
-    for (uri, _document) in state.documents.iter() {
-        if !ExtUrl::should_diagnose(uri) {
+    for document in state.documents.values() {
+        if !ExtUrl::should_diagnose(&document.url) {
             continue;
         }
 
@@ -1136,7 +1131,7 @@ pub(crate) fn diagnostics_refresh_all(state: WorldState) {
         // non-oak state).
         INDEXER_QUEUE
             .send(IndexerQueueTask::Diagnostics(RefreshDiagnosticsTask {
-                uri: uri.clone(),
+                uri: document.url.clone(),
                 state: state.legacy_snapshot(),
             }))
             .unwrap_or_else(|err| lsp::log_error!("Failed to queue diagnostics refresh: {err}"));

crates/ark/src/lsp/state.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::path::Path;
 
+use aether_path::FilePath;
 use anyhow::anyhow;
 use oak_db::OakDatabase;
 use oak_semantic::library::Library;
@@ -15,8 +16,10 @@ use crate::lsp::inputs::source_root::SourceRoot;
 /// code. This is a pure value. There is no interior mutability in this data
 /// structure. It can be cloned and safely sent to other threads.
 pub(crate) struct WorldState {
-    /// Watched documents
-    pub(crate) documents: HashMap<Url, Document>,
+    /// Watched documents, keyed on the normalised [`FilePath`] form.
+    /// The verbatim editor URL is preserved on each [`Document::url`]
+    /// for wire output.
+    pub(crate) documents: HashMap<FilePath, Document>,
 
     /// Watched folders
     pub(crate) workspace: Workspace,
@@ -79,15 +82,17 @@ impl WorldState {
     }
 
     pub(crate) fn get_document(&self, uri: &Url) -> anyhow::Result<&Document> {
-        if let Some(doc) = self.documents.get(uri) {
+        let key = FilePath::from_url(uri);
+        if let Some(doc) = self.documents.get(&key) {
             Ok(doc)
         } else {
             Err(anyhow!("Can't find document for URI {uri}"))
         }
     }
 
     pub(crate) fn get_document_mut(&mut self, uri: &Url) -> anyhow::Result<&mut Document> {
-        if let Some(doc) = self.documents.get_mut(uri) {
+        let key = FilePath::from_url(uri);
+        if let Some(doc) = self.documents.get_mut(&key) {
             Ok(doc)
         } else {
             Err(anyhow!("Can't find document for URI {uri}"))
@@ -115,6 +120,15 @@ impl WorldState {
             ..self.clone()
         }
     }
+
+    /// Insert a document, keying on the normalised [`FilePath`] and
+    /// stashing the verbatim editor URL on [`Document::url`] for wire
+    /// output.
+    pub(crate) fn insert_document(&mut self, uri: Url, mut doc: Document) {
+        let key = FilePath::from_url(&uri);
+        doc.url = uri;
+        self.documents.insert(key, doc);
+    }
 }
 
 pub(crate) fn with_document<T, F>(
@@ -151,8 +165,11 @@ where
 }
 
 pub(crate) fn workspace_uris(state: &WorldState) -> Vec<Url> {
-    let uris: Vec<Url> = state.documents.iter().map(|elt| elt.0.clone()).collect();
-    uris
+    state
+        .documents
+        .values()
+        .map(|doc| doc.url.clone())
+        .collect()
 }
 
 #[cfg(test)]

crates/ark/src/lsp/state_handlers.rs

@@ -143,7 +143,7 @@ pub(crate) fn initialize(
     // Start first round of indexing. `state.documents` is empty at init since
     // no `didOpen` has fired yet, but build the set through the same shape we
     // use elsewhere so the call site reads consistently.
-    let editor_owned: HashSet<FilePath> = state.documents.keys().map(FilePath::from_url).collect();
+    let editor_owned: HashSet<FilePath> = state.documents.keys().cloned().collect();
     let pending = lsp_state.oak_scheduler.set_workspace_paths(
         &mut state.oak,
         &workspace_paths,
@@ -271,7 +271,7 @@ pub(crate) fn did_open(
     let document = Document::new_with_parser(contents, &mut parser, Some(version));
 
     lsp_state.parsers.insert(uri.clone(), parser);
-    state.documents.insert(uri.clone(), document.clone());
+    state.insert_document(uri.clone(), document.clone());
 
     let url_id = FilePath::from_url(&uri);
     state.oak.upsert_editor(url_id, contents.to_string());
@@ -330,7 +330,7 @@ pub(crate) fn did_close(
 
     state
         .documents
-        .remove(&uri)
+        .remove(&FilePath::from_url(&uri))
         .ok_or(anyhow!("Failed to remove document for URI: {uri}"))?;
 
     lsp_state
@@ -405,7 +405,7 @@ pub(crate) fn did_change_watched_files(
 ) -> anyhow::Result<Vec<ScanRequest>> {
     // Editor owns the contents of files it has open: Oak should ignore
     // disk-side events for those URLs.
-    let editor_owned: HashSet<FilePath> = state.documents.keys().map(FilePath::from_url).collect();
+    let editor_owned: HashSet<FilePath> = state.documents.keys().cloned().collect();
 
     let events: Vec<FileEvent> = params
         .changes
@@ -456,7 +456,7 @@ pub(crate) fn did_change_workspace_folders(
     // Editor-owned URLs survive eviction in `OrphanRoot` so the user's
     // open buffers keep getting analysed even when their workspace
     // folder goes away.
-    let editor_owned: HashSet<FilePath> = state.documents.keys().map(FilePath::from_url).collect();
+    let editor_owned: HashSet<FilePath> = state.documents.keys().cloned().collect();
 
     let pending = lsp_state.oak_scheduler.set_workspace_paths(
         &mut state.oak,