feat: sync docs with incremental change

crates/llm-ls/src/document.rs

@@ -1,9 +1,11 @@
 use ropey::Rope;
 use tower_lsp::jsonrpc::Result;
-use tree_sitter::{Parser, Tree};
+use tower_lsp::lsp_types::Range;
+use tracing::info;
+use tree_sitter::{InputEdit, Parser, Point, Tree};
 
-use crate::internal_error;
 use crate::language_id::LanguageId;
+use crate::{get_position_idx, internal_error};
 
 fn get_parser(language_id: LanguageId) -> Result<Parser> {
     match language_id {
@@ -186,10 +188,85 @@ impl Document {
         })
     }
 
-    pub(crate) async fn change(&mut self, text: &str) -> Result<()> {
-        let rope = Rope::from_str(text);
-        self.tree = self.parser.parse(text, None);
-        self.text = rope;
+    pub(crate) async fn change(&mut self, range: Range, text: &str) -> Result<()> {
+        let start_idx = get_position_idx(
+            &self.text,
+            range.start.line as usize,
+            range.start.character as usize,
+        )?;
+        let start_byte = self
+            .text
+            .try_char_to_byte(start_idx)
+            .map_err(internal_error)?;
+        let old_end_idx = get_position_idx(
+            &self.text,
+            range.end.line as usize,
+            range.end.character as usize,
+        )?;
+        let old_end_byte = self
+            .text
+            .try_char_to_byte(old_end_idx)
+            .map_err(internal_error)?;
+        let start_position = Point {
+            row: range.start.line as usize,
+            column: range.start.character as usize,
+        };
+        let old_end_position = Point {
+            row: range.end.line as usize,
+            column: range.end.character as usize,
+        };
+        let (new_end_idx, new_end_position) = if range.start == range.end {
+            let row = range.start.line as usize;
+            let column = range.start.character as usize;
+            let idx = self.text.try_line_to_char(row).map_err(internal_error)? + column;
+            let rope = Rope::from_str(text);
+            let text_len = rope.len_chars();
+            let end_idx = idx + text_len;
+            self.text.insert(idx, text);
+            (
+                end_idx,
+                Point {
+                    row,
+                    column: column + text_len,
+                },
+            )
+        } else {
+            let slice_size = old_end_idx - start_idx;
+            self.text
+                .try_remove(start_idx..old_end_idx)
+                .map_err(internal_error)?;
+            self.text.insert(start_idx, text);
+            let rope = Rope::from_str(text);
+            let text_len = rope.len_chars();
+            let character_difference = text_len as isize - slice_size as isize;
+            let new_end_idx = if character_difference.is_negative() {
+                old_end_idx - character_difference.wrapping_abs() as usize
+            } else {
+                old_end_idx + character_difference as usize
+            };
+            let row = self
+                .text
+                .try_char_to_line(new_end_idx)
+                .map_err(internal_error)?;
+            let line_start = self.text.try_line_to_char(row).map_err(internal_error)?;
+            let column = new_end_idx - line_start;
+            (new_end_idx, Point { row, column })
+        };
+        if let Some(tree) = self.tree.as_mut() {
+            let edit = InputEdit {
+                start_byte,
+                old_end_byte,
+                new_end_byte: self
+                    .text
+                    .try_char_to_byte(new_end_idx)
+                    .map_err(internal_error)?,
+                start_position,
+                old_end_position,
+                new_end_position,
+            };
+            tree.edit(&edit);
+        }
+        self.tree = self.parser.parse(self.text.to_string(), self.tree.as_ref());
         Ok(())
     }
 }

crates/llm-ls/src/main.rs

@@ -25,6 +25,16 @@ const MAX_WARNING_REPEAT: Duration = Duration::from_secs(3_600);
 const NAME: &str = "llm-ls";
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
+fn get_position_idx(rope: &Rope, row: usize, col: usize) -> Result<usize> {
+    Ok(rope.try_line_to_char(row).map_err(internal_error)?
+        + col.min(
+            rope.get_line(row.min(rope.len_lines() - 1))
+                .ok_or_else(|| internal_error(format!("failed to find line at {row}")))?
+                .len_chars()
+                - 1,
+        ))
+}
+
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
 enum CompletionType {
     Empty,
@@ -42,45 +52,71 @@ impl Display for CompletionType {
     }
 }
 
-fn should_complete(document: &Document, position: Position) -> CompletionType {
+fn should_complete(document: &Document, position: Position) -> Result<CompletionType> {
     let row = position.line as usize;
     let column = position.character as usize;
     if let Some(tree) = &document.tree {
         let current_node = tree.root_node().descendant_for_point_range(
             tree_sitter::Point { row, column },
-            tree_sitter::Point { row, column },
+            tree_sitter::Point {
+                row,
+                column: column + 1,
+            },
         );
         if let Some(node) = current_node {
             if node == tree.root_node() {
-                return CompletionType::MultiLine;
+                return Ok(CompletionType::MultiLine);
             }
             let start = node.start_position();
             let end = node.end_position();
-            let mut start_offset = document.text.line_to_char(start.row) + start.column;
-            let mut end_offset = document.text.line_to_char(end.row) + end.column - 1;
-            let start_char = document.text.char(start_offset);
+            let mut start_offset = get_position_idx(&document.text, start.row, start.column)?;
+            let mut end_offset = get_position_idx(&document.text, end.row, end.column)? - 1;
+            let start_char = document
+                .text
+                .get_char(start_offset.min(document.text.len_chars() - 1))
+                .ok_or_else(|| {
+                    internal_error(format!("failed to find start char at {start_offset}"))
+                })?;
+            let end_char = document
+                .text
+                .get_char(end_offset.min(document.text.len_chars() - 1))
+                .ok_or_else(|| {
+                    internal_error(format!("failed to find end char at {end_offset}"))
+                })?;
             if !start_char.is_whitespace() {
                 start_offset += 1;
             }
-            let end_char = document.text.char(end_offset);
             if !end_char.is_whitespace() {
                 end_offset -= 1;
             }
             if start_offset >= end_offset {
-                return CompletionType::SingleLine;
+                return Ok(CompletionType::SingleLine);
             }
-            let slice = document.text.slice(start_offset..end_offset);
+            let slice = document
+                .text
+                .get_slice(start_offset..end_offset)
+                .ok_or_else(|| {
+                    internal_error(format!(
+                        "failed to find slice at {start_offset}..{end_offset}"
+                    ))
+                })?;
             if slice.to_string().trim().is_empty() {
-                return CompletionType::MultiLine;
+                return Ok(CompletionType::MultiLine);
             }
         }
     }
-    let start_idx = document.text.line_to_char(row);
-    let next_char = document.text.char(start_idx + column);
+    let start_idx = document
+        .text
+        .try_line_to_char(row)
+        .map_err(internal_error)?;
+    let next_char = document
+        .text
+        .get_char(start_idx + column)
+        .ok_or_else(|| internal_error(format!("failed to find char at {}", start_idx + column)))?;
     if next_char.is_whitespace() {
-        CompletionType::SingleLine
+        Ok(CompletionType::SingleLine)
     } else {
-        CompletionType::Empty
+        Ok(CompletionType::Empty)
     }
 }
 
@@ -271,12 +307,12 @@ fn build_prompt(
         let mut after_iter = text.lines_at(pos.line as usize);
         let mut before_line = before_iter.next();
         if let Some(line) = before_line {
-            let col = (pos.character as usize).clamp(0, line.len_chars());
+            let col = (pos.character as usize).clamp(0, line.len_chars() - 1);
             before_line = Some(line.slice(0..col));
         }
         let mut after_line = after_iter.next();
         if let Some(line) = after_line {
-            let col = (pos.character as usize).clamp(0, line.len_chars());
+            let col = (pos.character as usize).clamp(0, line.len_chars() - 1);
             after_line = Some(line.slice(col..));
         }
         let mut before = vec![];
@@ -334,7 +370,7 @@ fn build_prompt(
         let mut first = true;
         for mut line in text.lines_at(pos.line as usize + 1).reversed() {
             if first {
-                let col = (pos.character as usize).clamp(0, line.len_chars());
+                let col = (pos.character as usize).clamp(0, line.len_chars() - 1);
                 line = line.slice(0..col);
                 first = false;
             }
@@ -582,7 +618,7 @@ impl Backend {
                     *unauthenticated_warn_at = Instant::now();
                 }
             }
-            let completion_type = should_complete(document, params.text_document_position.position);
+            let completion_type = should_complete(document, params.text_document_position.position)?;
             info!(%completion_type, "completion type: {completion_type:?}");
             if completion_type == CompletionType::Empty {
                 return Ok(CompletionResult { request_id, completions: vec![]});
@@ -658,7 +694,7 @@ impl LanguageServer for Backend {
             }),
             capabilities: ServerCapabilities {
                 text_document_sync: Some(TextDocumentSyncCapability::Kind(
-                    TextDocumentSyncKind::FULL,
+                    TextDocumentSyncKind::INCREMENTAL,
                 )),
                 ..Default::default()
             },
@@ -702,9 +738,15 @@ impl LanguageServer for Backend {
         let mut document_map = self.document_map.write().await;
         let doc = document_map.get_mut(&uri);
         if let Some(doc) = doc {
-            match doc.change(&params.content_changes[0].text).await {
-                Ok(()) => info!("{uri} changed"),
-                Err(err) => error!("error when changing {uri}: {err}"),
+            for change in &params.content_changes {
+                if let Some(range) = change.range {
+                    match doc.change(range, &change.text).await {
+                        Ok(()) => info!("{uri} changed"),
+                        Err(err) => error!("error when changing {uri}: {err}"),
+                    }
+                } else {
+                    warn!("Could not update document, got change request with missing range");
+                }
             }
         } else {
             warn!("textDocument/didChange {uri}: document not found");

crates/testbed/src/holes_generator.rs

@@ -91,7 +91,7 @@ pub(crate) async fn generate_holes(
             if trimmed.starts_with(repo.language.comment_token()) || trimmed.is_empty() {
                 continue;
             }
-            let column_nb = rng.gen_range(0..15.min(line.len_chars()));
+            let column_nb = rng.gen_range(0..15.min(line.len_chars() - 1));
             holes.push(Hole::new(
                 line_nb as u32,
                 column_nb as u32,

crates/testbed/src/main.rs

@@ -454,7 +454,7 @@ async fn complete_holes(
                     .line(hole.cursor.line as usize)
                     .slice(hole.cursor.character as usize..)
                     .len_chars()
-                - 1;
+                - 1; // NOTE: -1 to preserve the trailing `\n`
             file_content.remove(hole_start..hole_end);
 
             let uri = Url::parse(&format!("file:/{file_path_str}"))?;