Migrate hover, signature help, and completions to ArkFile (#1264)

Branched from #1263
Progress towards #1212

The handlers migrated use a slightly different approach because they use
`r_task()` heavily and the Oak DB shouldn't be touched from an
`r_task()` because of the risk of cancellation panics, which would be UB
over the R stack. We could catch panics in `r_task()` and unwind for
cancellation but I didn't want to think through the various
implications.

- `DocumentContext` drops `Document` and stores the pieces it needs from
the `ArkDb` and LSP state (`tree`, `contents`, `line_index`,
`encoding`). This way the DB is not sent over an `r_task()`.

- Add LSP <> TS position/range converters as methods on
`DocumentContext` (parallel to the ones on `ArkFile`, also replacements
for the `Document` methods).

- Migrate hover, signature help, and completion to the refactored
`DocumentContext`.

Author: lionel-

crates/ark/src/lsp/ark_file.rs

@@ -90,21 +90,15 @@ impl ArkFile {
         db: &dyn ArkDb,
         point: tree_sitter::Point,
     ) -> anyhow::Result<lsp_types::Position> {
-        let line_col = biome_line_index::LineCol {
-            line: point.row as u32,
-            col: point.column as u32,
-        };
-        to_proto::position_from_line_col(line_col, self.line_index(db), self.encoding)
+        lsp_position_from_tree_sitter_point(point, self.line_index(db), self.encoding)
     }
 
     pub(crate) fn lsp_range_from_tree_sitter_range(
         &self,
         db: &dyn ArkDb,
         range: tree_sitter::Range,
     ) -> anyhow::Result<lsp_types::Range> {
-        let start = self.lsp_position_from_tree_sitter_point(db, range.start_point)?;
-        let end = self.lsp_position_from_tree_sitter_point(db, range.end_point)?;
-        Ok(lsp_types::Range::new(start, end))
+        lsp_range_from_tree_sitter_range(range, self.line_index(db), self.encoding)
     }
 
     pub(crate) fn tree_sitter_range_from_lsp_range(
@@ -126,6 +120,31 @@ impl ArkFile {
     }
 }
 
+/// Free functions over `LineIndex` + `PositionEncoding`, so anything holding
+/// those two (an `ArkFile` plus its `db`, or a `DocumentContext`) can convert
+/// without each carrying its own copy of the logic.
+pub(crate) fn lsp_position_from_tree_sitter_point(
+    point: tree_sitter::Point,
+    line_index: &biome_line_index::LineIndex,
+    encoding: PositionEncoding,
+) -> anyhow::Result<lsp_types::Position> {
+    let line_col = biome_line_index::LineCol {
+        line: point.row as u32,
+        col: point.column as u32,
+    };
+    to_proto::position_from_line_col(line_col, line_index, encoding)
+}
+
+pub(crate) fn lsp_range_from_tree_sitter_range(
+    range: tree_sitter::Range,
+    line_index: &biome_line_index::LineIndex,
+    encoding: PositionEncoding,
+) -> anyhow::Result<lsp_types::Range> {
+    let start = lsp_position_from_tree_sitter_point(range.start_point, line_index, encoding)?;
+    let end = lsp_position_from_tree_sitter_point(range.end_point, line_index, encoding)?;
+    Ok(lsp_types::Range::new(start, end))
+}
+
 #[cfg(test)]
 pub(crate) fn test_ark_file(code: &str) -> (oak_db::OakDatabase, ArkFile) {
     use aether_path::FilePath;

crates/ark/src/lsp/completions/completion_item.rs

@@ -36,6 +36,7 @@ use tower_lsp::lsp_types::Range;
 use tower_lsp::lsp_types::TextEdit;
 use tree_sitter::Node;
 
+use crate::lsp::ark_file::lsp_position_from_tree_sitter_point;
 use crate::lsp::completions::function_context::ArgumentsStatus;
 use crate::lsp::completions::function_context::FunctionContext;
 use crate::lsp::completions::function_context::FunctionRefUsage;
@@ -100,7 +101,7 @@ pub(super) fn completion_item_from_assignment(
     let lhs = node.child_by_field_name("lhs").into_result()?;
     let rhs = node.child_by_field_name("rhs").into_result()?;
 
-    let label = lhs.node_as_str(&context.document.contents)?.to_string();
+    let label = lhs.node_as_str(context.contents)?.to_string();
 
     // TODO: Resolve functions that exist in-document here.
     let mut item = completion_item(label.clone(), CompletionData::ScopeVariable {
@@ -123,7 +124,7 @@ pub(super) fn completion_item_from_assignment(
     // benefit from the logic in completion_item_from_function() :(
     if rhs.node_type() == NodeType::FunctionDefinition {
         if let Some(parameters) = rhs.child_by_field_name("parameters") {
-            let parameters = parameters.node_as_str(&context.document.contents)?;
+            let parameters = parameters.node_as_str(context.contents)?;
             item.detail = Some(join!(label, parameters));
         }
 
@@ -644,9 +645,8 @@ fn completion_item_from_dot_dot_dot(
 
     item.kind = Some(CompletionItemKind::FIELD);
 
-    let position = context
-        .document
-        .lsp_position_from_tree_sitter_point(context.point)?;
+    let position =
+        lsp_position_from_tree_sitter_point(context.point, context.line_index, context.encoding)?;
 
     let range = Range {
         start: position,

crates/ark/src/lsp/completions/function_context.rs

@@ -10,6 +10,8 @@ use tower_lsp::lsp_types;
 use tower_lsp::lsp_types::Range;
 use tree_sitter::Node;
 
+use crate::lsp::ark_file::lsp_position_from_tree_sitter_point;
+use crate::lsp::ark_file::lsp_range_from_tree_sitter_range;
 use crate::lsp::document_context::DocumentContext;
 use crate::lsp::traits::node::NodeExt;
 use crate::treesitter::node_find_parent_call;
@@ -61,9 +63,11 @@ impl FunctionContext {
             // We shouldn't ever attempt to instantiate a FunctionContext or
             // function-flavored CompletionItem in this degenerate case, but we
             // return a dummy FunctionContext just to be safe.
-            let node_end = document_context
-                .document
-                .lsp_position_from_tree_sitter_point(completion_node.range().end_point)?;
+            let node_end = lsp_position_from_tree_sitter_point(
+                completion_node.range().end_point,
+                document_context.line_index,
+                document_context.encoding,
+            )?;
 
             return Ok(Self {
                 name: String::new(),
@@ -74,10 +78,7 @@ impl FunctionContext {
             });
         };
 
-        let usage = determine_function_usage(
-            &effective_function_node,
-            &document_context.document.contents,
-        );
+        let usage = determine_function_usage(&effective_function_node, document_context.contents);
 
         let function_name_node = if effective_function_node.is_namespace_operator() {
             // Note: this could be 'None', in the case of, e.g., `dplyr::@`
@@ -93,7 +94,7 @@ impl FunctionContext {
 
         let name = match function_name_node {
             Some(node) => node
-                .node_to_string(&document_context.document.contents)
+                .node_to_string(document_context.contents)
                 .unwrap_or_default(),
             None => String::new(),
         };
@@ -111,16 +112,18 @@ impl FunctionContext {
         Ok(Self {
             name,
             range: match function_name_node {
-                Some(node) => document_context
-                    .document
-                    .lsp_range_from_tree_sitter_range(node.range())?,
+                Some(node) => lsp_range_from_tree_sitter_range(
+                    node.range(),
+                    document_context.line_index,
+                    document_context.encoding,
+                )?,
                 None => {
                     // Create a zero-width range at the end of the effective_function_node
-                    let node_end = document_context
-                        .document
-                        .lsp_position_from_tree_sitter_point(
-                            effective_function_node.range().end_point,
-                        )?;
+                    let node_end = lsp_position_from_tree_sitter_point(
+                        effective_function_node.range().end_point,
+                        document_context.line_index,
+                        document_context.encoding,
+                    )?;
                     lsp_types::Range::new(node_end, node_end)
                 },
             },

crates/ark/src/lsp/completions/provide.rs

@@ -25,7 +25,7 @@ pub(crate) fn provide_completions(
         "provide_completions() - Completion node text: '{node_text}', Node type: '{node_type:?}'",
         node_text = document_context
             .node
-            .node_as_str(&document_context.document.contents)
+            .node_as_str(document_context.contents)
             .unwrap_or_default(),
         node_type = document_context.node.node_type()
     );

crates/ark/src/lsp/completions/sources/composite.rs

@@ -220,8 +220,7 @@ mod tests {
     use crate::lsp::completions::completion_context::CompletionContext;
     use crate::lsp::completions::sources::composite::get_completions;
     use crate::lsp::completions::sources::composite::is_identifier_like;
-    use crate::lsp::document::Document;
-    use crate::lsp::document_context::DocumentContext;
+    use crate::lsp::document_context::TestDocument;
     use crate::lsp::state::WorldState;
     use crate::r_task;
     use crate::treesitter::NodeType;
@@ -235,8 +234,8 @@ mod tests {
             // identifiers that we provide completions for
             for keyword in ["if", "for", "while"] {
                 let (text, point) = point_from_cursor(&format!("{keyword}@"));
-                let document = Document::new(text.as_str(), None);
-                let context = DocumentContext::new(&document, point, None);
+                let doc = TestDocument::new(&text);
+                let context = doc.context(point);
 
                 assert!(is_identifier_like(context.node));
                 assert_eq!(
@@ -251,8 +250,8 @@ mod tests {
     fn test_get_completions_on_empty_document() {
         r_task(|| {
             let (text, point) = point_from_cursor("@");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
 
@@ -269,8 +268,8 @@ mod tests {
         r_task(|| {
             let code = "x <- 1:3\n@\nrnorm(3)";
             let (text, point) = point_from_cursor(code);
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
 

crates/ark/src/lsp/completions/sources/composite/call.rs

@@ -76,7 +76,7 @@ fn completions_from_call(
         return Ok(None);
     };
 
-    let callee = callee.node_as_str(&document_context.document.contents)?;
+    let callee = callee.node_as_str(document_context.contents)?;
 
     // - Prefer `root` as the first argument if it exists
     // - Then fall back to looking it up, if possible
@@ -122,7 +122,7 @@ fn get_first_argument(context: &DocumentContext, node: &Node) -> anyhow::Result<
         return Ok(None);
     };
 
-    let text = value.node_as_str(&context.document.contents)?;
+    let text = value.node_as_str(context.contents)?;
 
     let options = RParseEvalOptions {
         forbid_function_calls: true,
@@ -280,8 +280,7 @@ mod tests {
     use crate::fixtures::point_from_cursor;
     use crate::lsp::completions::completion_context::CompletionContext;
     use crate::lsp::completions::sources::composite::call::completions_from_call;
-    use crate::lsp::document::Document;
-    use crate::lsp::document_context::DocumentContext;
+    use crate::lsp::document_context::TestDocument;
     use crate::lsp::state::WorldState;
     use crate::r_task;
 
@@ -290,8 +289,8 @@ mod tests {
         r_task(|| {
             // Right after `tab`
             let (text, point) = point_from_cursor("match(tab@)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -303,8 +302,8 @@ mod tests {
 
             // Right after `tab`
             let (text, point) = point_from_cursor("match(1, tab@)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -323,8 +322,8 @@ mod tests {
         r_task(|| {
             // Place cursor between `()`
             let (text, point) = point_from_cursor("not_a_known_function(@)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap();
@@ -343,8 +342,8 @@ mod tests {
 
             // Place cursor between `()`
             let (text, point) = point_from_cursor("my_fun(@)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -360,17 +359,17 @@ mod tests {
 
             // Place just before the `()`
             let (text, point) = point_from_cursor("my_fun@()");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap();
             assert!(completions.is_none());
 
             // Place just after the `()`
             let (text, point) = point_from_cursor("my_fun()@");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap();
@@ -392,8 +391,8 @@ mod tests {
 
             // Place cursor between `()`
             let (text, point) = point_from_cursor("my_fun(@)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -409,8 +408,8 @@ mod tests {
         r_task(|| {
             // No arguments typed yet
             let (text, point) = point_from_cursor("match(\n  @\n)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -421,8 +420,8 @@ mod tests {
 
             // Partially typed argument
             let (text, point) = point_from_cursor("match(\n  tab@\n)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -433,8 +432,8 @@ mod tests {
 
             // Partially typed second argument
             let (text, point) = point_from_cursor("match(\n  1,\n  tab@\n)");
-            let document = Document::new(text.as_str(), None);
-            let document_context = DocumentContext::new(&document, point, None);
+            let doc = TestDocument::new(&text);
+            let document_context = doc.context(point);
             let state = WorldState::default();
             let context = CompletionContext::new(&document_context, &state);
             let completions = completions_from_call(&context).unwrap().unwrap();
@@ -450,8 +449,8 @@ mod tests {
         r_task(|| {
             fn assert_no_call_completions(code_with_cursor: &str) {
                 let (text, point) = point_from_cursor(code_with_cursor);
-                let document = Document::new(text.as_str(), None);
-                let document_context = DocumentContext::new(&document, point, None);
+                let doc = TestDocument::new(&text);
+                let document_context = doc.context(point);
                 let state = WorldState::default();
                 let context = CompletionContext::new(&document_context, &state);
                 let completions = completions_from_call(&context).unwrap();

crates/ark/src/lsp/completions/sources/composite/document.rs

@@ -168,7 +168,7 @@ fn completions_from_document_function_arguments(
             continue;
         }
 
-        let parameter = node.node_as_str(&context.document.contents)?.to_string();
+        let parameter = node.node_as_str(context.contents)?.to_string();
         match completion_item_from_scope_parameter(parameter.as_str(), context) {
             Ok(item) => completions.push(item),
             Err(err) => log::error!("{err:?}"),
@@ -185,7 +185,7 @@ fn call_uses_nse(node: &Node, context: &DocumentContext) -> bool {
         lhs.is_identifier_or_string().into_result()?;
 
         let value = lhs
-            .node_as_str(&context.document.contents)?
+            .node_as_str(context.contents)?
             .to_string();
         matches!(value.as_str(), "expression" | "local" | "quote" | "enquote" | "substitute" | "with" | "within").into_result()?;