Avoid materializing strings

4 files changed, 127 insertions, 19 deletions

diff --git a/crates/libsyntax2/src/yellow/green.rs b/crates/libsyntax2/src/yellow/green.rs
index 700f2704f..59aefb0de 100644
--- a/crates/libsyntax2/src/yellow/green.rs
+++ b/crates/libsyntax2/src/yellow/green.rs
@@ -43,21 +43,13 @@ impl GreenNode {
         }
     }
 
-    pub fn text(&self) -> String {
-        let mut buff = String::new();
-        go(self, &mut buff);
-        return buff;
-        fn go(node: &GreenNode, buff: &mut String) {
-            match node {
-                GreenNode::Leaf { text, .. } => buff.push_str(text.as_str()),
-                GreenNode::Branch(b) => b.children().iter().for_each(|child| go(child, buff)),
-            }
-        }
+    pub fn leaf_text(&self) -> Option<SmolStr> {
+        self.leaf_text_ref().map(Clone::clone)
     }
 
-    pub fn leaf_text(&self) -> Option<SmolStr> {
+    pub fn leaf_text_ref(&self) -> Option<&SmolStr> {
         match self {
-            GreenNode::Leaf { text, .. } => Some(text.clone()),
+            GreenNode::Leaf { text, .. } => Some(text),
             GreenNode::Branch(_) => None,
         }
     }
diff --git a/crates/libsyntax2/src/yellow/mod.rs b/crates/libsyntax2/src/yellow/mod.rs
index b94c794fe..82eda79d6 100644
--- a/crates/libsyntax2/src/yellow/mod.rs
+++ b/crates/libsyntax2/src/yellow/mod.rs
@@ -2,6 +2,7 @@ mod builder;
 mod green;
 mod red;
 mod syntax;
+mod syntax_text;
 
 use std::{
     sync::Arc,
@@ -12,6 +13,7 @@ pub(crate) use self::{
     builder::GreenBuilder,
     green::GreenNode,
     red::RedNode,
+    syntax_text::SyntaxText,
 };
 
 #[derive(Debug)]
@@ -32,6 +34,12 @@ pub struct OwnedRoot(Arc<SyntaxRoot>);
 #[derive(Clone, Copy, Debug)]
 pub struct RefRoot<'a>(&'a OwnedRoot); // TODO: shared_from_this instead of double indirection
 
+impl<'a> RefRoot<'a> {
+    fn syntax_root(&self) -> &'a SyntaxRoot {
+        self.0.syntax_root()
+    }
+}
+
 impl TreeRoot for OwnedRoot {
     fn borrowed(&self) -> RefRoot {
         RefRoot(&self)
@@ -78,7 +86,7 @@ impl RedPtr {
         RedPtr(red.into())
     }
 
-    unsafe fn get<'a>(self, _root: &'a impl TreeRoot) -> &'a RedNode {
+    unsafe fn get<'a>(self, _root: &'a SyntaxRoot) -> &'a RedNode {
         &*self.0.as_ptr()
     }
 }
diff --git a/crates/libsyntax2/src/yellow/syntax.rs b/crates/libsyntax2/src/yellow/syntax.rs
index 75b6cb7dc..ede9ad3c8 100644
--- a/crates/libsyntax2/src/yellow/syntax.rs
+++ b/crates/libsyntax2/src/yellow/syntax.rs
@@ -6,7 +6,7 @@ use std::{
 use smol_str::SmolStr;
 
 use {
-    yellow::{GreenNode, RedNode, TreeRoot, SyntaxRoot, RedPtr, RefRoot, OwnedRoot},
+    yellow::{GreenNode, RedNode, TreeRoot, SyntaxRoot, RedPtr, RefRoot, OwnedRoot, SyntaxText},
     SyntaxKind::{self, *},
     TextRange, TextUnit,
 };
@@ -58,6 +58,13 @@ impl SyntaxNode<OwnedRoot> {
     }
 }
 
+impl<'a> SyntaxNode<RefRoot<'a>> {
+    pub(crate) fn leaf_text_ref(self) -> Option<&'a SmolStr> {
+        let red = unsafe { self.red.get(self.root.syntax_root()) };
+        red.green().leaf_text_ref()
+    }
+}
+
 impl<R: TreeRoot> SyntaxNode<R> {
     pub fn borrowed<'a>(&'a self) -> SyntaxNodeRef<'a> {
         SyntaxNode {
@@ -66,7 +73,7 @@ impl<R: TreeRoot> SyntaxNode<R> {
         }
     }
 
-    pub fn owned<'a>(&'a self) -> SyntaxNode {
+    pub fn owned(&self) -> SyntaxNode {
         SyntaxNode {
             root: self.root.owned(),
             red: self.red,
@@ -82,8 +89,8 @@ impl<R: TreeRoot> SyntaxNode<R> {
         TextRange::offset_len(red.start_offset(), red.green().text_len())
     }
 
-    pub fn text(&self) -> String {
-        self.red().green().text()
+    pub fn text(&self) -> SyntaxText {
+        SyntaxText::new(self.borrowed())
     }
 
     pub fn children(&self) -> impl Iterator<Item = SyntaxNode<R>> {
@@ -91,7 +98,7 @@ impl<R: TreeRoot> SyntaxNode<R> {
         let n_children = self.red().n_children();
         let root = self.root.clone();
         (0..n_children).map(move |i| {
-            let red = unsafe { red.get(&root) };
+            let red = unsafe { red.get(root.syntax_root()) };
             SyntaxNode {
                 root: root.clone(),
                 red: red.get_child(i).unwrap(),
@@ -171,7 +178,7 @@ impl<R: TreeRoot> SyntaxNode<R> {
     }
 
     fn red(&self) -> &RedNode {
-        unsafe { self.red.get(&self.root) }
+        unsafe { self.red.get(self.root.syntax_root()) }
     }
 }
 
diff --git a/crates/libsyntax2/src/yellow/syntax_text.rs b/crates/libsyntax2/src/yellow/syntax_text.rs
new file mode 100644
index 000000000..268547470
--- /dev/null
+++ b/crates/libsyntax2/src/yellow/syntax_text.rs
@@ -0,0 +1,101 @@
+use std::{
+    fmt, ops,
+};
+
+use {
+    SyntaxNodeRef, TextRange, TextUnit,
+    algo::walk::preorder,
+    text_utils::{intersect, contains_offset_nonstrict},
+};
+
+#[derive(Clone)]
+pub struct SyntaxText<'a> {
+    node: SyntaxNodeRef<'a>,
+    range: TextRange,
+}
+
+impl<'a> SyntaxText<'a> {
+    pub(crate) fn new(node: SyntaxNodeRef<'a>) -> SyntaxText<'a> {
+        SyntaxText {
+            node,
+            range: node.range()
+        }
+    }
+    pub fn chunks(&self) -> impl Iterator<Item=&'a str> {
+        let range = self.range;
+        preorder(self.node)
+            .filter_map(move |node| {
+                let text = node.leaf_text_ref()?;
+                let range = intersect(range, node.range())?;
+                let range = range - node.range().start();
+                Some(&text[range])
+            })
+    }
+    pub fn to_string(&self) -> String {
+        self.chunks().collect()
+    }
+    pub fn contains(&self, c: char) -> bool {
+        self.chunks().any(|it| it.contains(c))
+    }
+    pub fn find(&self, c: char) -> Option<TextUnit> {
+        let mut acc: TextUnit = 0.into();
+        for chunk in self.chunks() {
+            if let Some(pos) = chunk.find(c) {
+                let pos: TextUnit = (pos as u32).into();
+                return Some(acc + pos);
+            }
+            acc += TextUnit::of_str(chunk);
+        }
+        None
+    }
+    pub fn len(&self) -> TextUnit {
+        self.range.len()
+    }
+    pub fn slice(&self, range: impl SyntaxTextSlice) -> SyntaxText<'a> {
+        let range = range.restrict(self.range)
+            .unwrap_or_else(|| {
+                panic!("invalid slice, range: {:?}, slice: {:?}", self.range, range)
+            });
+        SyntaxText { node: self.node, range }
+    }
+}
+
+impl<'a> fmt::Debug for SyntaxText<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Debug::fmt(&self.to_string(), f)
+    }
+}
+
+impl<'a> fmt::Display for SyntaxText<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Display::fmt(&self.to_string(), f)
+    }
+}
+
+pub trait SyntaxTextSlice: fmt::Debug {
+    fn restrict(&self, range: TextRange) -> Option<TextRange>;
+}
+
+impl SyntaxTextSlice for TextRange {
+    fn restrict(&self, range: TextRange) -> Option<TextRange> {
+        intersect(*self, range)
+    }
+}
+
+impl SyntaxTextSlice for ops::RangeTo<TextUnit> {
+    fn restrict(&self, range: TextRange) -> Option<TextRange> {
+        if !contains_offset_nonstrict(range, self.end) {
+            return None;
+        }
+        Some(TextRange::from_to(range.start(), self.end))
+    }
+}
+
+impl SyntaxTextSlice for ops::RangeFrom<TextUnit> {
+    fn restrict(&self, range: TextRange) -> Option<TextRange> {
+        if !contains_offset_nonstrict(range, self.start) {
+            return None;
+        }
+        Some(TextRange::from_to(self.start, range.end()))
+    }
+}