From 89b64d4daac2907ad1939534b8fd3898e9a2f3ce Mon Sep 17 00:00:00 2001
From: Nikita Revenco <154856872+NikitaRevenco@users.noreply.github.com>
Date: Fri, 6 Dec 2024 22:38:56 +0000
Subject: [PATCH] feat: replace existing tree sitter document instead of
 creating new one

---
 helix-term/src/commands/typed.rs | 18 ++++++++++++---
 helix-view/src/editor.rs         | 39 ++++++++++++++++++++++++++------
 helix-view/src/lib.rs            |  6 +++++
 3 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index cd422a6956e4..958c0885c1df 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -11,7 +11,7 @@ use helix_core::indent::MAX_INDENT;
 use helix_core::{line_ending, shellwords::Shellwords};
 use helix_stdx::path::home_dir;
 use helix_view::document::{read_to_string, DEFAULT_LANGUAGE_NAME};
-use helix_view::editor::{CloseError, ConfigEvent};
+use helix_view::editor::{CloseError, ConfigEvent, TREE_SITTER_TREE_DOCUMENT_ID};
 use serde_json::Value;
 use ui::completers::{self, Completer};
 
@@ -2154,6 +2154,10 @@ fn tree_sitter_tree(
         return Ok(());
     }
 
+    // Safety: 100_000_000 != 0
+    let doc_id =
+        DocumentId::new(unsafe { NonZeroUsize::new_unchecked(TREE_SITTER_TREE_DOCUMENT_ID) });
+
     let (_view, doc) = current_ref!(cx.editor);
 
     if let Some(syntax) = doc.syntax() {
@@ -2165,11 +2169,19 @@ fn tree_sitter_tree(
             let mut contents = String::new();
             helix_core::syntax::pretty_print_tree(&mut contents, selected_node)?;
 
-            cx.editor.new_file_from_document(
+            // If this document doesn't exist, create it. If it exists, override the contents with our new string
+
+            if let Some(_) = cx.editor.documents.get_mut(&doc_id) {
+                if let Ok(_) = cx.editor.close_document(doc_id, true) {
+                } else {
+                    bail!("Couldn't close the previous Tree Sitter document")
+                };
+            }
+            cx.editor.new_file_from_document_with_id(
                 Action::VerticalSplit,
                 Document::from(Rope::from(contents), None, cx.editor.config.clone()),
+                doc_id,
             );
-
             let (_view, doc) = current!(cx.editor);
 
             doc.set_language_by_language_id("tsq", cx.editor.syn_loader.clone())?
diff --git a/helix-view/src/editor.rs b/helix-view/src/editor.rs
index f492a5f8db30..3e549f7b236a 100644
--- a/helix-view/src/editor.rs
+++ b/helix-view/src/editor.rs
@@ -59,6 +59,11 @@ use arc_swap::{
 
 pub const DEFAULT_AUTO_SAVE_DELAY: u64 = 3000;
 
+// No document is likely to ever reach this size
+// This is needed to be able to replace the previous Tree Sitter Tree Document with a new one
+// WARNING: do not set this to 0, as it will result in undefined behaviour.
+pub const TREE_SITTER_TREE_DOCUMENT_ID: usize = 100_000_000;
+
 fn deserialize_duration_millis<'de, D>(deserializer: D) -> Result<Duration, D::Error>
 where
     D: serde::Deserializer<'de>,
@@ -1671,13 +1676,9 @@ impl Editor {
         self._refresh();
     }
 
-    /// Generate an id for a new document and register it.
-    fn new_document(&mut self, mut doc: Document) -> DocumentId {
-        let id = self.next_document_id;
-        // Safety: adding 1 from 1 is fine, probably impossible to reach usize max
-        self.next_document_id =
-            DocumentId(unsafe { NonZeroUsize::new_unchecked(self.next_document_id.0.get() + 1) });
+    fn new_document_impl(&mut self, mut doc: Document, id: DocumentId) -> DocumentId {
         doc.id = id;
+
         self.documents.insert(id, doc);
 
         let (save_sender, save_receiver) = tokio::sync::mpsc::unbounded_channel();
@@ -1685,16 +1686,40 @@ impl Editor {
 
         let stream = UnboundedReceiverStream::new(save_receiver).flatten();
         self.save_queue.push(stream);
-
         id
     }
 
+    fn new_document_with_id(&mut self, doc: Document, id: DocumentId) -> DocumentId {
+        self.new_document_impl(doc, id)
+    }
+
+    /// Generate an id for a new document and register it.
+    fn new_document(&mut self, doc: Document) -> DocumentId {
+        let id = self.next_document_id;
+
+        // Safety: adding 1 from 1 is fine, probably impossible to reach usize max
+        self.next_document_id =
+            DocumentId(unsafe { NonZeroUsize::new_unchecked(self.next_document_id.0.get() + 1) });
+
+        self.new_document_impl(doc, id)
+    }
+
     pub fn new_file_from_document(&mut self, action: Action, doc: Document) -> DocumentId {
         let id = self.new_document(doc);
         self.switch(id, action);
         id
     }
 
+    pub fn new_file_from_document_with_id(
+        &mut self,
+        action: Action,
+        doc: Document,
+        id: DocumentId,
+    ) {
+        self.new_document_with_id(doc, id);
+        self.switch(id, action);
+    }
+
     pub fn new_file(&mut self, action: Action) -> DocumentId {
         self.new_file_from_document(action, Document::default(self.config.clone()))
     }
diff --git a/helix-view/src/lib.rs b/helix-view/src/lib.rs
index d54b49ef5400..f672de870b10 100644
--- a/helix-view/src/lib.rs
+++ b/helix-view/src/lib.rs
@@ -24,6 +24,12 @@ use std::num::NonZeroUsize;
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
 pub struct DocumentId(NonZeroUsize);
 
+impl DocumentId {
+    pub fn new(id: NonZeroUsize) -> DocumentId {
+        DocumentId(id)
+    }
+}
+
 impl Default for DocumentId {
     fn default() -> DocumentId {
         // Safety: 1 is non-zero