feat: add source

Author: HerringtonDarkholme

crates/core/src/lib.rs

@@ -6,6 +6,7 @@ mod node;
 pub mod ops;
 mod pattern;
 mod replacer;
+mod source;
 pub mod traversal;
 mod ts_parser;
 
@@ -21,7 +22,6 @@ use language::Language;
 use node::Root;
 use ts_parser::{Edit, TSParseError};
 
-#[derive(Clone)]
 pub struct AstGrep<L: Language> {
   inner: Root<L>,
 }
@@ -64,7 +64,7 @@ impl<L: Language> AstGrep<L> {
   }
 
   pub fn generate(self) -> String {
-    self.inner.source
+    self.inner.source.to_string()
   }
 }
 

crates/core/src/match_tree.rs

@@ -268,7 +268,7 @@ mod test {
       inner: goal.root_node().child(0).unwrap(),
       root: &Root {
         inner: goal.clone(),
-        source: s1.to_string(),
+        source: s1.into(),
         lang: Tsx,
       },
     };
@@ -277,7 +277,7 @@ mod test {
       inner: cand.root_node(),
       root: &Root {
         inner: cand.clone(),
-        source: s2.to_string(),
+        source: s2.into(),
         lang: Tsx,
       },
     };
@@ -298,7 +298,7 @@ mod test {
       inner: goal.root_node().child(0).unwrap(),
       root: &Root {
         inner: goal.clone(),
-        source: s1.to_string(),
+        source: s1.into(),
         lang: Tsx,
       },
     };
@@ -307,7 +307,7 @@ mod test {
       inner: cand.root_node(),
       root: &Root {
         inner: cand.clone(),
-        source: s2.to_string(),
+        source: s2.into(),
         lang: Tsx,
       },
     };

crates/core/src/node.rs

@@ -1,17 +1,17 @@
 use crate::language::Language;
 use crate::matcher::{FindAllNodes, Matcher, NodeMatch};
 use crate::replacer::Replacer;
+use crate::source::{Content, Source};
 use crate::traversal::{Pre, Visitor};
 use crate::ts_parser::{parse, perform_edit, Edit, TSParseError};
 
 use std::borrow::Cow;
 
 /// Represents [`tree_sitter::Tree`] and owns source string
 /// Note: Root is generic against [`Language`](crate::language::Language)
-#[derive(Clone)]
 pub struct Root<L: Language> {
   pub(crate) inner: tree_sitter::Tree,
-  pub(crate) source: String,
+  pub(crate) source: Source,
   pub(crate) lang: L,
 }
 
@@ -30,7 +30,7 @@ impl<L: Language> Root<L> {
   }
   // extract non generic implementation to reduce code size
   pub fn do_edit(&mut self, edit: Edit) -> Result<(), TSParseError> {
-    let input = unsafe { self.source.as_mut_vec() };
+    let input = self.source.as_mut_vec();
     let input_edit = perform_edit(&mut self.inner, input, &edit);
     self.inner.edit(&input_edit);
     self.inner = parse(&self.source, Some(&self.inner), self.lang.get_ts_language())?;

crates/core/src/pattern.rs

@@ -2,12 +2,13 @@ use crate::language::Language;
 use crate::match_tree::{extract_var_from_node, match_end_non_recursive, match_node_non_recursive};
 use crate::matcher::{KindMatcher, Matcher};
 use crate::{meta_var::MetaVarEnv, Node, Root};
+use std::sync::Arc;
 
 use bit_set::BitSet;
 
 #[derive(Clone)]
 pub struct Pattern<L: Language> {
-  pub root: Root<L>,
+  pub root: Arc<Root<L>>,
   /// used in contextual pattern, specify which AST subpart is considered as pattern
   /// e.g. in js`class { $F }` we set selector to public_field_definition
   selector: Option<KindMatcher<L>>,
@@ -22,7 +23,7 @@ impl<L: Language> Pattern<L> {
       todo!("multi-children pattern is not supported yet.")
     }
     Self {
-      root,
+      root: Arc::new(root),
       selector: None,
     }
   }
@@ -39,7 +40,7 @@ impl<L: Language> Pattern<L> {
       todo!("use result to indicate failure");
     }
     Self {
-      root,
+      root: Arc::new(root),
       selector: Some(kind_matcher),
     }
   }

crates/core/src/source.rs

@@ -0,0 +1,46 @@
+use std::ops::Deref;
+
+pub trait Content: ToString + Deref<Target = str> {
+  fn as_mut_vec(&mut self) -> &mut Vec<u8>;
+}
+
+pub enum Source {
+  Plain(String),
+  Customized(Box<dyn Content + Sync + Send>),
+}
+
+use Source::*;
+
+impl From<&str> for Source {
+  fn from(s: &str) -> Self {
+    Plain(s.into())
+  }
+}
+
+impl Deref for Source {
+  type Target = str;
+  fn deref(&self) -> &Self::Target {
+    match self {
+      Plain(s) => s.deref(),
+      Customized(c) => c.deref(),
+    }
+  }
+}
+
+impl ToString for Source {
+  fn to_string(&self) -> String {
+    match self {
+      Self::Plain(s) => s.to_owned(),
+      Self::Customized(c) => c.to_string(),
+    }
+  }
+}
+
+impl Content for Source {
+  fn as_mut_vec(&mut self) -> &mut Vec<u8> {
+    match self {
+      Plain(s) => unsafe { s.as_mut_vec() },
+      Customized(c) => c.as_mut_vec(),
+    }
+  }
+}

crates/lsp/src/lib.rs

@@ -14,7 +14,6 @@ pub use tower_lsp::{LspService, Server};
 pub trait LSPLang: Language + Eq + Send + Sync + 'static {}
 impl<T> LSPLang for T where T: Language + Eq + Send + Sync + 'static {}
 
-#[derive(Clone)]
 struct VersionedAst<L: Language> {
   version: i32,
   root: AstGrep<L>,
@@ -223,17 +222,16 @@ impl<L: LSPLang> Backend<L> {
   }
   async fn on_open(&self, params: DidOpenTextDocumentParams) -> Option<()> {
     let text_doc = params.text_document;
-    let uri = text_doc.uri.as_str();
+    let uri = text_doc.uri.as_str().to_owned();
     let text = text_doc.text;
     let lang = Self::infer_lang_from_uri(&text_doc.uri)?;
     let root = AstGrep::new(text, lang);
     let versioned = VersionedAst {
       version: text_doc.version,
       root,
     };
-    let copied = versioned.clone();
+    self.publish_diagnostics(text_doc.uri, &versioned).await;
     self.map.insert(uri.to_owned(), versioned); // don't lock dashmap
-    self.publish_diagnostics(text_doc.uri, &copied).await;
     Some(())
   }
   async fn on_change(&self, params: DidChangeTextDocumentParams) -> Option<()> {
@@ -251,9 +249,7 @@ impl<L: LSPLang> Backend<L> {
       version: text_doc.version,
       root,
     };
-    let copied = versioned.clone();
-    drop(versioned); // don't lock dashmap
-    self.publish_diagnostics(text_doc.uri, &copied).await;
+    self.publish_diagnostics(text_doc.uri, &versioned).await;
     Some(())
   }
   async fn on_close(&self, params: DidCloseTextDocumentParams) {