2 files changed, 26 insertions, 20 deletions

diff --git a/crates/mbe/src/syntax_bridge.rs b/crates/mbe/src/syntax_bridge.rs
index e648519f9..0cdc175be 100644
--- a/crates/mbe/src/syntax_bridge.rs
+++ b/crates/mbe/src/syntax_bridge.rs
@@ -16,16 +16,18 @@ use crate::ExpandError;
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub enum TokenTextRange {
     Token(TextRange),
-    Delimiter(TextRange, TextRange),
+    Delimiter(TextRange),
 }
 
 impl TokenTextRange {
     pub fn by_kind(self, kind: SyntaxKind) -> Option<TextRange> {
         match self {
             TokenTextRange::Token(it) => Some(it),
-            TokenTextRange::Delimiter(open, close) => match kind {
-                T!['{'] | T!['('] | T!['['] => Some(open),
-                T!['}'] | T![')'] | T![']'] => Some(close),
+            TokenTextRange::Delimiter(it) => match kind {
+                T!['{'] | T!['('] | T!['['] => Some(TextRange::at(it.start(), 1.into())),
+                T!['}'] | T![')'] | T![']'] => {
+                    Some(TextRange::at(it.end() - TextSize::of('}'), 1.into()))
+                }
                 _ => None,
             },
         }
@@ -51,6 +53,7 @@ pub fn syntax_node_to_token_tree(node: &SyntaxNode) -> Option<(tt::Subtree, Toke
     let global_offset = node.text_range().start();
     let mut c = Convertor::new(node, global_offset);
     let subtree = c.go()?;
+    c.id_alloc.map.entries.shrink_to_fit();
     Some((subtree, c.id_alloc.map))
 }
 
@@ -113,8 +116,10 @@ impl TokenMap {
     pub fn token_by_range(&self, relative_range: TextRange) -> Option<tt::TokenId> {
         let &(token_id, _) = self.entries.iter().find(|(_, range)| match range {
             TokenTextRange::Token(it) => *it == relative_range,
-            TokenTextRange::Delimiter(open, close) => {
-                *open == relative_range || *close == relative_range
+            TokenTextRange::Delimiter(it) => {
+                let open = TextRange::at(it.start(), 1.into());
+                let close = TextRange::at(it.end() - TextSize::of('}'), 1.into());
+                open == relative_range || close == relative_range
             }
         })?;
         Some(token_id)
@@ -136,15 +141,17 @@ impl TokenMap {
         close_relative_range: TextRange,
     ) -> usize {
         let res = self.entries.len();
-        self.entries
-            .push((token_id, TokenTextRange::Delimiter(open_relative_range, close_relative_range)));
+        let cover = open_relative_range.cover(close_relative_range);
+
+        self.entries.push((token_id, TokenTextRange::Delimiter(cover)));
         res
     }
 
     fn update_close_delim(&mut self, idx: usize, close_relative_range: TextRange) {
         let (_, token_text_range) = &mut self.entries[idx];
-        if let TokenTextRange::Delimiter(dim, _) = token_text_range {
-            *token_text_range = TokenTextRange::Delimiter(*dim, close_relative_range);
+        if let TokenTextRange::Delimiter(dim) = token_text_range {
+            let cover = dim.cover(close_relative_range);
+            *token_text_range = TokenTextRange::Delimiter(cover);
         }
     }
 
@@ -500,7 +507,7 @@ impl SrcToken for SynToken {
         }
     }
     fn to_text(&self) -> SmolStr {
-        self.token().text().clone()
+        self.token().text().into()
     }
 }
 
@@ -593,7 +600,8 @@ impl<'a> TtTreeSink<'a> {
         }
     }
 
-    fn finish(self) -> (Parse<SyntaxNode>, TokenMap) {
+    fn finish(mut self) -> (Parse<SyntaxNode>, TokenMap) {
+        self.token_map.entries.shrink_to_fit();
         (self.inner.finish(), self.token_map)
     }
 }
@@ -674,10 +682,8 @@ impl<'a> TreeSink for TtTreeSink<'a> {
             self.text_pos += TextSize::of(text);
         }
 
-        let text = SmolStr::new(self.buf.as_str());
+        self.inner.token(kind, self.buf.as_str());
         self.buf.clear();
-        self.inner.token(kind, text);
-
         // Add whitespace between adjoint puncts
         let next = last.bump();
         if let (
diff --git a/crates/mbe/src/tests.rs b/crates/mbe/src/tests.rs
index 9ff901e97..bd2977ebd 100644
--- a/crates/mbe/src/tests.rs
+++ b/crates/mbe/src/tests.rs
@@ -261,7 +261,6 @@ fn test_expr_order() {
 
     let dump = format!("{:#?}", expanded);
     assert_eq_text!(
-        dump.trim(),
         r#"[email protected]
   [email protected]
     [email protected] "fn"
@@ -285,6 +284,7 @@ fn test_expr_order() {
             [email protected] "2"
         [email protected] ";"
       [email protected] "}""#,
+        dump.trim()
     );
 }
 
@@ -989,7 +989,6 @@ fn test_tt_composite2() {
 
     let res = format!("{:#?}", &node);
     assert_eq_text!(
-        res.trim(),
         r###"[email protected]
   [email protected]
     [email protected]
@@ -1003,7 +1002,8 @@ fn test_tt_composite2() {
       [email protected] ">"
       [email protected] " "
       [email protected] "#"
-      [email protected] ")""###
+      [email protected] ")""###,
+        res.trim()
     );
 }
 
@@ -1742,7 +1742,7 @@ impl MacroFixture {
     fn assert_expand(&self, invocation: &str, expected: &str) {
         let expansion = self.expand_tt(invocation);
         let actual = format!("{:?}", expansion);
-        test_utils::assert_eq_text!(&actual.trim(), &expected.trim());
+        test_utils::assert_eq_text!(&expected.trim(), &actual.trim());
     }
 
     fn assert_expand_items(&self, invocation: &str, expected: &str) -> &MacroFixture {
@@ -1941,7 +1941,6 @@ fn test_no_space_after_semi_colon() {
 
     let dump = format!("{:#?}", expanded);
     assert_eq_text!(
-        dump.trim(),
         r###"[email protected]
   [email protected]
     [email protected]
@@ -1981,6 +1980,7 @@ fn test_no_space_after_semi_colon() {
     [email protected]
       [email protected] "f"
     [email protected] ";""###,
+        dump.trim()
     );
 }