Compress token trees for best memory usage

Author: ChayimFriedman2

Cargo.lock

@@ -866,7 +866,8 @@ impl TtTreeSink<'_> {
     /// Parses a float literal as if it was a one to two name ref nodes with a dot inbetween.
     /// This occurs when a float literal is used as a field access.
     fn float_split(&mut self, has_pseudo_dot: bool) {
-        let (text, span) = match self.cursor.token_tree() {
+        let token_tree = self.cursor.token_tree();
+        let (text, span) = match &token_tree {
             Some(tt::TokenTree::Leaf(tt::Leaf::Literal(
                 lit @ tt::Literal { span, kind: tt::LitKind::Float, .. },
             ))) => (lit.text(), *span),
@@ -928,9 +929,15 @@ impl TtTreeSink<'_> {
                                 self.buf.push_str("r#");
                                 self.text_pos += TextSize::of("r#");
                             }
-                            let r = (ident.sym.as_str(), ident.span);
+                            let text = ident.sym.as_str();
+                            self.buf += text;
+                            self.text_pos += TextSize::of(text);
+                            combined_span = match combined_span {
+                                None => Some(ident.span),
+                                Some(prev_span) => Some(Self::merge_spans(prev_span, ident.span)),
+                            };
                             self.cursor.bump();
-                            r
+                            continue 'tokens;
                         }
                         tt::Leaf::Punct(punct) => {
                             assert!(punct.char.is_ascii());

crates/syntax-bridge/src/lib.rs

@@ -36,9 +36,9 @@ fn check_punct_spacing(fixture: &str) {
             if let tt::TokenTree::Leaf(Leaf::Punct(Punct {
                 spacing, span: Span { range, .. }, ..
             })) = token_tree
-                && let Some(expected) = annotations.remove(range)
+                && let Some(expected) = annotations.remove(&range)
             {
-                assert_eq!(expected, *spacing);
+                assert_eq!(expected, spacing);
             }
             cursor.bump();
         }

crates/syntax-bridge/src/tests.rs

@@ -15,6 +15,8 @@ doctest = false
 [dependencies]
 arrayvec.workspace = true
 text-size.workspace = true
+rustc-hash.workspace = true
+indexmap.workspace = true
 
 span = { path = "../span", version = "0.0", default-features = false }
 stdx.workspace = true

crates/tt/Cargo.toml

@@ -1,17 +1,17 @@
 //! Stateful iteration over token trees.
 //!
 //! We use this as the source of tokens for parser.
-use crate::{Leaf, Subtree, TokenTree, TokenTreesView};
+use crate::{Leaf, Subtree, TokenTree, TokenTreesView, dispatch_ref};
 
 pub struct Cursor<'a> {
-    buffer: &'a [TokenTree],
+    buffer: TokenTreesView<'a>,
     index: usize,
     subtrees_stack: Vec<usize>,
 }
 
 impl<'a> Cursor<'a> {
     pub fn new(buffer: TokenTreesView<'a>) -> Self {
-        Self { buffer: buffer.0, index: 0, subtrees_stack: Vec::new() }
+        Self { buffer, index: 0, subtrees_stack: Vec::new() }
     }
 
     /// Check whether it is eof
@@ -23,16 +23,22 @@ impl<'a> Cursor<'a> {
         self.subtrees_stack.is_empty()
     }
 
-    fn last_subtree(&self) -> Option<(usize, &'a Subtree)> {
+    fn at(&self, idx: usize) -> Option<TokenTree> {
+        dispatch_ref! {
+            match self.buffer.repr => tt => Some(tt.get(idx)?.to_api(self.buffer.span_parts))
+        }
+    }
+
+    fn last_subtree(&self) -> Option<(usize, Subtree)> {
         self.subtrees_stack.last().map(|&subtree_idx| {
-            let TokenTree::Subtree(subtree) = &self.buffer[subtree_idx] else {
+            let Some(TokenTree::Subtree(subtree)) = self.at(subtree_idx) else {
                 panic!("subtree pointing to non-subtree");
             };
             (subtree_idx, subtree)
         })
     }
 
-    pub fn end(&mut self) -> &'a Subtree {
+    pub fn end(&mut self) -> Subtree {
         let (last_subtree_idx, last_subtree) =
             self.last_subtree().expect("called `Cursor::end()` without an open subtree");
         // +1 because `Subtree.len` excludes the subtree itself.
@@ -46,14 +52,14 @@ impl<'a> Cursor<'a> {
     }
 
     /// Returns the `TokenTree` at the cursor if it is not at the end of a subtree.
-    pub fn token_tree(&self) -> Option<&'a TokenTree> {
+    pub fn token_tree(&self) -> Option<TokenTree> {
         if let Some((last_subtree_idx, last_subtree)) = self.last_subtree() {
             // +1 because `Subtree.len` excludes the subtree itself.
             if last_subtree_idx + last_subtree.usize_len() + 1 == self.index {
                 return None;
             }
         }
-        self.buffer.get(self.index)
+        self.at(self.index)
     }
 
     /// Bump the cursor, and enters a subtree if it is on one.
@@ -66,7 +72,7 @@ impl<'a> Cursor<'a> {
                 "called `Cursor::bump()` when at the end of a subtree"
             );
         }
-        if let TokenTree::Subtree(_) = self.buffer[self.index] {
+        if let Some(TokenTree::Subtree(_)) = self.at(self.index) {
             self.subtrees_stack.push(self.index);
         }
         self.index += 1;
@@ -81,28 +87,31 @@ impl<'a> Cursor<'a> {
             }
         }
         // +1 because `Subtree.len` excludes the subtree itself.
-        if let TokenTree::Subtree(_) = self.buffer[self.index] {
+        if let Some(TokenTree::Subtree(_)) = self.at(self.index) {
             self.subtrees_stack.push(self.index);
         }
         self.index += 1;
     }
 
-    pub fn peek_two_leaves(&self) -> Option<[&'a Leaf; 2]> {
+    pub fn peek_two_leaves(&self) -> Option<[Leaf; 2]> {
         if let Some((last_subtree_idx, last_subtree)) = self.last_subtree() {
             // +1 because `Subtree.len` excludes the subtree itself.
             let last_end = last_subtree_idx + last_subtree.usize_len() + 1;
             if last_end == self.index || last_end == self.index + 1 {
                 return None;
             }
         }
-        self.buffer.get(self.index..self.index + 2).and_then(|it| match it {
-            [TokenTree::Leaf(a), TokenTree::Leaf(b)] => Some([a, b]),
+        self.at(self.index).zip(self.at(self.index + 1)).and_then(|it| match it {
+            (TokenTree::Leaf(a), TokenTree::Leaf(b)) => Some([a, b]),
             _ => None,
         })
     }
 
     pub fn crossed(&self) -> TokenTreesView<'a> {
         assert!(self.is_root());
-        TokenTreesView::new(&self.buffer[..self.index])
+        TokenTreesView {
+            repr: self.buffer.repr.get(..self.index).unwrap(),
+            span_parts: self.buffer.span_parts,
+        }
     }
 }

crates/tt/src/buffer.rs

@@ -7,11 +7,14 @@ use arrayvec::ArrayVec;
 use intern::sym;
 use span::Span;
 
-use crate::{Ident, Leaf, MAX_GLUED_PUNCT_LEN, Punct, Spacing, Subtree, TokenTree, TokenTreesView};
+use crate::{
+    Ident, Leaf, MAX_GLUED_PUNCT_LEN, Punct, Spacing, Subtree, TokenTree, TokenTreesReprRef,
+    TokenTreesView, dispatch_ref,
+};
 
 #[derive(Clone)]
 pub struct TtIter<'a> {
-    inner: std::slice::Iter<'a, TokenTree>,
+    inner: TokenTreesView<'a>,
 }
 
 impl fmt::Debug for TtIter<'_> {
@@ -21,17 +24,17 @@ impl fmt::Debug for TtIter<'_> {
 }
 
 #[derive(Clone, Copy)]
-pub struct TtIterSavepoint<'a>(&'a [TokenTree]);
+pub struct TtIterSavepoint<'a>(TokenTreesView<'a>);
 
 impl<'a> TtIterSavepoint<'a> {
     pub fn remaining(self) -> TokenTreesView<'a> {
-        TokenTreesView::new(self.0)
+        self.0
     }
 }
 
 impl<'a> TtIter<'a> {
-    pub(crate) fn new(tt: &'a [TokenTree]) -> TtIter<'a> {
-        TtIter { inner: tt.iter() }
+    pub(crate) fn new(tt: TokenTreesView<'a>) -> TtIter<'a> {
+        TtIter { inner: tt }
     }
 
     pub fn expect_char(&mut self, char: char) -> Result<(), ()> {
@@ -141,8 +144,8 @@ impl<'a> TtIter<'a> {
                 let _ = self.next().unwrap();
                 let _ = self.next().unwrap();
                 res.push(first);
-                res.push(*second);
-                res.push(*third.unwrap());
+                res.push(second);
+                res.push(third.unwrap());
             }
             ('-' | '!' | '*' | '/' | '&' | '%' | '^' | '+' | '<' | '=' | '>' | '|', '=', _)
             | ('-' | '=' | '>', '>', _)
@@ -154,25 +157,29 @@ impl<'a> TtIter<'a> {
             | ('|', '|', _) => {
                 let _ = self.next().unwrap();
                 res.push(first);
-                res.push(*second);
+                res.push(second);
             }
             _ => res.push(first),
         }
         Ok(res)
     }
 
     /// This method won't check for subtrees, so the nth token tree may not be the nth sibling of the current tree.
-    fn peek_n(&self, n: usize) -> Option<&'a TokenTree> {
-        self.inner.as_slice().get(n)
+    fn peek_n(&self, n: usize) -> Option<TokenTree> {
+        dispatch_ref! {
+            match self.inner.repr => tt => Some(tt.get(n)?.to_api(self.inner.span_parts))
+        }
     }
 
     pub fn peek(&self) -> Option<TtElement<'a>> {
-        match self.inner.as_slice().first()? {
-            TokenTree::Leaf(leaf) => Some(TtElement::Leaf(leaf.clone())),
+        match self.peek_n(0)? {
+            TokenTree::Leaf(leaf) => Some(TtElement::Leaf(leaf)),
             TokenTree::Subtree(subtree) => {
-                let nested_iter =
-                    TtIter { inner: self.inner.as_slice()[1..][..subtree.usize_len()].iter() };
-                Some(TtElement::Subtree(*subtree, nested_iter))
+                let nested_repr = self.inner.repr.get(1..subtree.usize_len() + 1).unwrap();
+                let nested_iter = TtIter {
+                    inner: TokenTreesView { repr: nested_repr, span_parts: self.inner.span_parts },
+                };
+                Some(TtElement::Subtree(subtree, nested_iter))
             }
         }
     }
@@ -183,26 +190,51 @@ impl<'a> TtIter<'a> {
     }
 
     pub fn next_span(&self) -> Option<Span> {
-        Some(self.inner.as_slice().first()?.first_span())
+        Some(self.peek()?.first_span())
     }
 
     pub fn remaining(&self) -> TokenTreesView<'a> {
-        TokenTreesView::new(self.inner.as_slice())
+        self.inner
     }
 
     /// **Warning**: This advances `skip` **flat** token trees, subtrees account for children+1!
     pub fn flat_advance(&mut self, skip: usize) {
-        self.inner = self.inner.as_slice()[skip..].iter();
+        self.inner.repr = self.inner.repr.get(skip..).unwrap();
     }
 
     pub fn savepoint(&self) -> TtIterSavepoint<'a> {
-        TtIterSavepoint(self.inner.as_slice())
+        TtIterSavepoint(self.inner)
     }
 
     pub fn from_savepoint(&self, savepoint: TtIterSavepoint<'a>) -> TokenTreesView<'a> {
-        let len = (self.inner.as_slice().as_ptr() as usize - savepoint.0.as_ptr() as usize)
-            / size_of::<TokenTree>();
-        TokenTreesView::new(&savepoint.0[..len])
+        let len = match (self.inner.repr, savepoint.0.repr) {
+            (
+                TokenTreesReprRef::SpanStorage32(this),
+                TokenTreesReprRef::SpanStorage32(savepoint),
+            ) => {
+                (this.as_ptr() as usize - savepoint.as_ptr() as usize)
+                    / size_of::<crate::storage::TokenTree<crate::storage::SpanStorage32>>()
+            }
+            (
+                TokenTreesReprRef::SpanStorage64(this),
+                TokenTreesReprRef::SpanStorage64(savepoint),
+            ) => {
+                (this.as_ptr() as usize - savepoint.as_ptr() as usize)
+                    / size_of::<crate::storage::TokenTree<crate::storage::SpanStorage64>>()
+            }
+            (
+                TokenTreesReprRef::SpanStorage96(this),
+                TokenTreesReprRef::SpanStorage96(savepoint),
+            ) => {
+                (this.as_ptr() as usize - savepoint.as_ptr() as usize)
+                    / size_of::<crate::storage::TokenTree<crate::storage::SpanStorage96>>()
+            }
+            _ => panic!("savepoint did not originate from this TtIter"),
+        };
+        TokenTreesView {
+            repr: savepoint.0.repr.get(..len).unwrap(),
+            span_parts: savepoint.0.span_parts,
+        }
     }
 
     pub fn next_as_view(&mut self) -> Option<TokenTreesView<'a>> {
@@ -242,14 +274,12 @@ impl TtElement<'_> {
 impl<'a> Iterator for TtIter<'a> {
     type Item = TtElement<'a>;
     fn next(&mut self) -> Option<Self::Item> {
-        match self.inner.next()? {
-            TokenTree::Leaf(leaf) => Some(TtElement::Leaf(leaf.clone())),
-            TokenTree::Subtree(subtree) => {
-                let nested_iter =
-                    TtIter { inner: self.inner.as_slice()[..subtree.usize_len()].iter() };
-                self.inner = self.inner.as_slice()[subtree.usize_len()..].iter();
-                Some(TtElement::Subtree(*subtree, nested_iter))
-            }
-        }
+        let result = self.peek()?;
+        let skip = match &result {
+            TtElement::Leaf(_) => 1,
+            TtElement::Subtree(subtree, _) => subtree.usize_len() + 1,
+        };
+        self.inner.repr = self.inner.repr.get(skip..).unwrap();
+        Some(result)
     }
 }