Encode closing delimiter span in FlatTrees

5 files changed, 68 insertions, 26 deletions

diff --git a/crates/proc-macro-api/src/lib.rs b/crates/proc-macro-api/src/lib.rs
index 4525d9dfdd7..14a7b913bf0 100644
--- a/crates/proc-macro-api/src/lib.rs
+++ b/crates/proc-macro-api/src/lib.rs
@@ -146,15 +146,16 @@ impl ProcMacro {
         attr: Option<&tt::Subtree>,
         env: Vec<(String, String)>,
     ) -> Result<Result<tt::Subtree, PanicMessage>, ServerError> {
+        let version = self.process.lock().unwrap_or_else(|e| e.into_inner()).version();
         let current_dir = env
             .iter()
             .find(|(name, _)| name == "CARGO_MANIFEST_DIR")
             .map(|(_, value)| value.clone());
 
         let task = ExpandMacro {
-            macro_body: FlatTree::new(subtree),
+            macro_body: FlatTree::new(subtree, version),
             macro_name: self.name.to_string(),
-            attributes: attr.map(FlatTree::new),
+            attributes: attr.map(|subtree| FlatTree::new(subtree, version)),
             lib: self.dylib_path.to_path_buf().into(),
             env,
             current_dir,
@@ -163,7 +164,9 @@ impl ProcMacro {
         let request = msg::Request::ExpandMacro(task);
         let response = self.process.lock().unwrap_or_else(|e| e.into_inner()).send_task(request)?;
         match response {
-            msg::Response::ExpandMacro(it) => Ok(it.map(FlatTree::to_subtree)),
+            msg::Response::ExpandMacro(it) => {
+                Ok(it.map(|tree| FlatTree::to_subtree(tree, version)))
+            }
             msg::Response::ListMacros(..) | msg::Response::ApiVersionCheck(..) => {
                 Err(ServerError { message: "unexpected response".to_string(), io: None })
             }
diff --git a/crates/proc-macro-api/src/msg.rs b/crates/proc-macro-api/src/msg.rs
index 4040efe93f0..4b01643c2a2 100644
--- a/crates/proc-macro-api/src/msg.rs
+++ b/crates/proc-macro-api/src/msg.rs
@@ -12,8 +12,12 @@ use crate::ProcMacroKind;
 
 pub use crate::msg::flat::FlatTree;
 
+// The versions of the server protocol
 pub const NO_VERSION_CHECK_VERSION: u32 = 0;
-pub const CURRENT_API_VERSION: u32 = 1;
+pub const VERSION_CHECK_VERSION: u32 = 1;
+pub const ENCODE_CLOSE_SPAN_VERSION: u32 = 2;
+
+pub const CURRENT_API_VERSION: u32 = ENCODE_CLOSE_SPAN_VERSION;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Request {
@@ -146,7 +150,7 @@ mod tests {
     fn test_proc_macro_rpc_works() {
         let tt = fixture_token_tree();
         let task = ExpandMacro {
-            macro_body: FlatTree::new(&tt),
+            macro_body: FlatTree::new(&tt, CURRENT_API_VERSION),
             macro_name: Default::default(),
             attributes: None,
             lib: std::env::current_dir().unwrap(),
@@ -158,6 +162,6 @@ mod tests {
         // println!("{}", json);
         let back: ExpandMacro = serde_json::from_str(&json).unwrap();
 
-        assert_eq!(tt, back.macro_body.to_subtree());
+        assert_eq!(tt, back.macro_body.to_subtree(CURRENT_API_VERSION));
     }
 }
diff --git a/crates/proc-macro-api/src/msg/flat.rs b/crates/proc-macro-api/src/msg/flat.rs
index fd3202e0b28..44245336f01 100644
--- a/crates/proc-macro-api/src/msg/flat.rs
+++ b/crates/proc-macro-api/src/msg/flat.rs
@@ -39,7 +39,10 @@ use std::collections::{HashMap, VecDeque};
 
 use serde::{Deserialize, Serialize};
 
-use crate::tt::{self, TokenId};
+use crate::{
+    msg::ENCODE_CLOSE_SPAN_VERSION,
+    tt::{self, TokenId},
+};
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct FlatTree {
@@ -52,7 +55,8 @@ pub struct FlatTree {
 }
 
 struct SubtreeRepr {
-    id: tt::TokenId,
+    open: tt::TokenId,
+    close: tt::TokenId,
     kind: tt::DelimiterKind,
     tt: [u32; 2],
 }
@@ -74,7 +78,7 @@ struct IdentRepr {
 }
 
 impl FlatTree {
-    pub fn new(subtree: &tt::Subtree) -> FlatTree {
+    pub fn new(subtree: &tt::Subtree, version: u32) -> FlatTree {
         let mut w = Writer {
             string_table: HashMap::new(),
             work: VecDeque::new(),
@@ -89,7 +93,11 @@ impl FlatTree {
         w.write(subtree);
 
         return FlatTree {
-            subtree: write_vec(w.subtree, SubtreeRepr::write),
+            subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {
+                write_vec(w.subtree, SubtreeRepr::write_with_close_span)
+            } else {
+                write_vec(w.subtree, SubtreeRepr::write)
+            },
             literal: write_vec(w.literal, LiteralRepr::write),
             punct: write_vec(w.punct, PunctRepr::write),
             ident: write_vec(w.ident, IdentRepr::write),
@@ -102,9 +110,13 @@ impl FlatTree {
         }
     }
 
-    pub fn to_subtree(self) -> tt::Subtree {
+    pub fn to_subtree(self, version: u32) -> tt::Subtree {
         return Reader {
-            subtree: read_vec(self.subtree, SubtreeRepr::read),
+            subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {
+                read_vec(self.subtree, SubtreeRepr::read_with_close_span)
+            } else {
+                read_vec(self.subtree, SubtreeRepr::read)
+            },
             literal: read_vec(self.literal, LiteralRepr::read),
             punct: read_vec(self.punct, PunctRepr::read),
             ident: read_vec(self.ident, IdentRepr::read),
@@ -130,9 +142,9 @@ impl SubtreeRepr {
             tt::DelimiterKind::Brace => 2,
             tt::DelimiterKind::Bracket => 3,
         };
-        [self.id.0, kind, self.tt[0], self.tt[1]]
+        [self.open.0, kind, self.tt[0], self.tt[1]]
     }
-    fn read([id, kind, lo, len]: [u32; 4]) -> SubtreeRepr {
+    fn read([open, kind, lo, len]: [u32; 4]) -> SubtreeRepr {
         let kind = match kind {
             0 => tt::DelimiterKind::Invisible,
             1 => tt::DelimiterKind::Parenthesis,
@@ -140,7 +152,26 @@ impl SubtreeRepr {
             3 => tt::DelimiterKind::Bracket,
             other => panic!("bad kind {other}"),
         };
-        SubtreeRepr { id: TokenId(id), kind, tt: [lo, len] }
+        SubtreeRepr { open: TokenId(open), close: TokenId::UNSPECIFIED, kind, tt: [lo, len] }
+    }
+    fn write_with_close_span(self) -> [u32; 5] {
+        let kind = match self.kind {
+            tt::DelimiterKind::Invisible => 0,
+            tt::DelimiterKind::Parenthesis => 1,
+            tt::DelimiterKind::Brace => 2,
+            tt::DelimiterKind::Bracket => 3,
+        };
+        [self.open.0, self.close.0, kind, self.tt[0], self.tt[1]]
+    }
+    fn read_with_close_span([open, close, kind, lo, len]: [u32; 5]) -> SubtreeRepr {
+        let kind = match kind {
+            0 => tt::DelimiterKind::Invisible,
+            1 => tt::DelimiterKind::Parenthesis,
+            2 => tt::DelimiterKind::Brace,
+            3 => tt::DelimiterKind::Bracket,
+            other => panic!("bad kind {other}"),
+        };
+        SubtreeRepr { open: TokenId(open), close: TokenId(close), kind, tt: [lo, len] }
     }
 }
 
@@ -244,9 +275,10 @@ impl<'a> Writer<'a> {
 
     fn enqueue(&mut self, subtree: &'a tt::Subtree) -> u32 {
         let idx = self.subtree.len();
-        let delimiter_id = subtree.delimiter.open;
+        let open = subtree.delimiter.open;
+        let close = subtree.delimiter.close;
         let delimiter_kind = subtree.delimiter.kind;
-        self.subtree.push(SubtreeRepr { id: delimiter_id, kind: delimiter_kind, tt: [!0, !0] });
+        self.subtree.push(SubtreeRepr { open, close, kind: delimiter_kind, tt: [!0, !0] });
         self.work.push_back((idx, subtree));
         idx as u32
     }
@@ -277,11 +309,7 @@ impl Reader {
             let repr = &self.subtree[i];
             let token_trees = &self.token_tree[repr.tt[0] as usize..repr.tt[1] as usize];
             let s = tt::Subtree {
-                delimiter: tt::Delimiter {
-                    open: repr.id,
-                    close: TokenId::UNSPECIFIED,
-                    kind: repr.kind,
-                },
+                delimiter: tt::Delimiter { open: repr.open, close: repr.close, kind: repr.kind },
                 token_trees: token_trees
                     .iter()
                     .copied()
diff --git a/crates/proc-macro-api/src/process.rs b/crates/proc-macro-api/src/process.rs
index 1ccbd780fdd..a526a21831b 100644
--- a/crates/proc-macro-api/src/process.rs
+++ b/crates/proc-macro-api/src/process.rs
@@ -56,6 +56,10 @@ impl ProcMacroProcessSrv {
         }
     }
 
+    pub(crate) fn version(&self) -> u32 {
+        self.version
+    }
+
     pub(crate) fn version_check(&mut self) -> Result<u32, ServerError> {
         let request = Request::ApiVersionCheck {};
         let response = self.send_task(request)?;
diff --git a/crates/proc-macro-srv/src/lib.rs b/crates/proc-macro-srv/src/lib.rs
index 2aac3790313..84bd15efb8b 100644
--- a/crates/proc-macro-srv/src/lib.rs
+++ b/crates/proc-macro-srv/src/lib.rs
@@ -31,7 +31,10 @@ use std::{
     time::SystemTime,
 };
 
-use proc_macro_api::{msg, ProcMacroKind};
+use proc_macro_api::{
+    msg::{self, CURRENT_API_VERSION},
+    ProcMacroKind,
+};
 
 use ::tt::token_id as tt;
 
@@ -67,8 +70,8 @@ impl ProcMacroSrv {
             None => None,
         };
 
-        let macro_body = task.macro_body.to_subtree();
-        let attributes = task.attributes.map(|it| it.to_subtree());
+        let macro_body = task.macro_body.to_subtree(CURRENT_API_VERSION);
+        let attributes = task.attributes.map(|it| it.to_subtree(CURRENT_API_VERSION));
         let result = thread::scope(|s| {
             let thread = thread::Builder::new()
                 .stack_size(EXPANDER_STACK_SIZE)
@@ -76,7 +79,7 @@ impl ProcMacroSrv {
                 .spawn_scoped(s, || {
                     expander
                         .expand(&task.macro_name, &macro_body, attributes.as_ref())
-                        .map(|it| msg::FlatTree::new(&it))
+                        .map(|it| msg::FlatTree::new(&it, CURRENT_API_VERSION))
                 });
             let res = match thread {
                 Ok(handle) => handle.join(),