fix: More proc-macro-srv proto fixes

crates/proc-macro-api/src/legacy_protocol/msg.rs

@@ -185,7 +185,7 @@ mod tests {
 
     use super::*;
 
-    fn fixture_token_tree() -> TopSubtree<Span> {
+    fn fixture_token_tree_top_many_none() -> TopSubtree<Span> {
         let anchor = SpanAnchor {
             file_id: span::EditionedFileId::new(
                 span::FileId::from_raw(0xe4e4e),
@@ -201,7 +201,7 @@ mod tests {
                 ctx: SyntaxContext::root(Edition::CURRENT),
             },
             close: Span {
-                range: TextRange::empty(TextSize::new(19)),
+                range: TextRange::empty(TextSize::new(0)),
                 anchor,
                 ctx: SyntaxContext::root(Edition::CURRENT),
             },
@@ -259,56 +259,159 @@ mod tests {
                 ctx: SyntaxContext::root(Edition::CURRENT),
             },
         );
+        builder.open(
+            DelimiterKind::Bracket,
+            Span {
+                range: TextRange::at(TextSize::new(15), TextSize::of('[')),
+                anchor,
+                ctx: SyntaxContext::root(Edition::CURRENT),
+            },
+        );
         builder.push(Leaf::Literal(Literal {
             symbol: sym::INTEGER_0,
             span: Span {
-                range: TextRange::at(TextSize::new(15), TextSize::of("0u32")),
+                range: TextRange::at(TextSize::new(16), TextSize::of("0u32")),
                 anchor,
                 ctx: SyntaxContext::root(Edition::CURRENT),
             },
             kind: tt::LitKind::Integer,
             suffix: Some(sym::u32),
         }));
         builder.close(Span {
-            range: TextRange::at(TextSize::new(19), TextSize::of('}')),
+            range: TextRange::at(TextSize::new(20), TextSize::of(']')),
             anchor,
             ctx: SyntaxContext::root(Edition::CURRENT),
         });
 
+        builder.close(Span {
+            range: TextRange::at(TextSize::new(21), TextSize::of('}')),
+            anchor,
+            ctx: SyntaxContext::root(Edition::CURRENT),
+        });
+
+        builder.build()
+    }
+
+    fn fixture_token_tree_top_empty_none() -> TopSubtree<Span> {
+        let anchor = SpanAnchor {
+            file_id: span::EditionedFileId::new(
+                span::FileId::from_raw(0xe4e4e),
+                span::Edition::CURRENT,
+            ),
+            ast_id: ROOT_ERASED_FILE_AST_ID,
+        };
+
+        let builder = TopSubtreeBuilder::new(Delimiter {
+            open: Span {
+                range: TextRange::empty(TextSize::new(0)),
+                anchor,
+                ctx: SyntaxContext::root(Edition::CURRENT),
+            },
+            close: Span {
+                range: TextRange::empty(TextSize::new(0)),
+                anchor,
+                ctx: SyntaxContext::root(Edition::CURRENT),
+            },
+            kind: DelimiterKind::Invisible,
+        });
+
+        builder.build()
+    }
+
+    fn fixture_token_tree_top_empty_brace() -> TopSubtree<Span> {
+        let anchor = SpanAnchor {
+            file_id: span::EditionedFileId::new(
+                span::FileId::from_raw(0xe4e4e),
+                span::Edition::CURRENT,
+            ),
+            ast_id: ROOT_ERASED_FILE_AST_ID,
+        };
+
+        let builder = TopSubtreeBuilder::new(Delimiter {
+            open: Span {
+                range: TextRange::empty(TextSize::new(0)),
+                anchor,
+                ctx: SyntaxContext::root(Edition::CURRENT),
+            },
+            close: Span {
+                range: TextRange::empty(TextSize::new(0)),
+                anchor,
+                ctx: SyntaxContext::root(Edition::CURRENT),
+            },
+            kind: DelimiterKind::Brace,
+        });
+
         builder.build()
     }
 
     #[test]
     fn test_proc_macro_rpc_works() {
-        let tt = fixture_token_tree();
-        for v in version::RUST_ANALYZER_SPAN_SUPPORT..=version::CURRENT_API_VERSION {
-            let mut span_data_table = Default::default();
-            let task = ExpandMacro {
-                data: ExpandMacroData {
-                    macro_body: FlatTree::from_subtree(tt.view(), v, &mut span_data_table),
-                    macro_name: Default::default(),
-                    attributes: None,
-                    has_global_spans: ExpnGlobals {
-                        serialize: true,
-                        def_site: 0,
-                        call_site: 0,
-                        mixed_site: 0,
+        for tt in [
+            fixture_token_tree_top_many_none,
+            fixture_token_tree_top_empty_none,
+            fixture_token_tree_top_empty_brace,
+        ] {
+            for v in version::RUST_ANALYZER_SPAN_SUPPORT..=version::CURRENT_API_VERSION {
+                let tt = tt();
+                let mut span_data_table = Default::default();
+                let task = ExpandMacro {
+                    data: ExpandMacroData {
+                        macro_body: FlatTree::from_subtree(tt.view(), v, &mut span_data_table),
+                        macro_name: Default::default(),
+                        attributes: None,
+                        has_global_spans: ExpnGlobals {
+                            serialize: true,
+                            def_site: 0,
+                            call_site: 0,
+                            mixed_site: 0,
+                        },
+                        span_data_table: Vec::new(),
                     },
-                    span_data_table: Vec::new(),
-                },
-                lib: Utf8PathBuf::from_path_buf(std::env::current_dir().unwrap()).unwrap(),
-                env: Default::default(),
-                current_dir: Default::default(),
-            };
-
-            let json = serde_json::to_string(&task).unwrap();
-            // println!("{}", json);
-            let back: ExpandMacro = serde_json::from_str(&json).unwrap();
-
-            assert!(
-                tt == back.data.macro_body.to_subtree_resolved(v, &span_data_table),
-                "version: {v}"
-            );
+                    lib: Utf8PathBuf::from_path_buf(std::env::current_dir().unwrap()).unwrap(),
+                    env: Default::default(),
+                    current_dir: Default::default(),
+                };
+
+                let json = serde_json::to_string(&task).unwrap();
+                // println!("{}", json);
+                let back: ExpandMacro = serde_json::from_str(&json).unwrap();
+
+                assert_eq!(
+                    tt,
+                    back.data.macro_body.to_subtree_resolved(v, &span_data_table),
+                    "version: {v}"
+                );
+            }
+        }
+    }
+
+    #[test]
+    #[cfg(feature = "sysroot-abi")]
+    fn test_proc_macro_rpc_works_ts() {
+        for tt in [
+            fixture_token_tree_top_many_none,
+            fixture_token_tree_top_empty_none,
+            fixture_token_tree_top_empty_brace,
+        ] {
+            let tt = tt();
+            for v in version::RUST_ANALYZER_SPAN_SUPPORT..=version::CURRENT_API_VERSION {
+                let mut span_data_table = Default::default();
+                let flat_tree = FlatTree::from_subtree(tt.view(), v, &mut span_data_table);
+                assert_eq!(
+                    tt,
+                    flat_tree.clone().to_subtree_resolved(v, &span_data_table),
+                    "version: {v}"
+                );
+                let ts = flat_tree.to_tokenstream_resolved(v, &span_data_table, |a, b| a.cover(b));
+                let call_site = *span_data_table.first().unwrap();
+                let mut span_data_table = Default::default();
+                assert_eq!(
+                    tt,
+                    FlatTree::from_tokenstream(ts.clone(), v, call_site, &mut span_data_table)
+                        .to_subtree_resolved(v, &span_data_table),
+                    "version: {v}, ts:\n{ts:#?}"
+                );
+            }
         }
     }
 }

crates/proc-macro-api/src/legacy_protocol/msg/flat.rs

@@ -85,7 +85,7 @@ pub fn deserialize_span_data_index_map(map: &[u32]) -> SpanDataIndexMap {
         .collect()
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct FlatTree {
     subtree: Vec<u32>,
     literal: Vec<u32>,
@@ -615,14 +615,17 @@ impl<'a, T: SpanTransformer>
         root: &'a proc_macro_srv::TokenStream<T::Span>,
     ) {
         let call_site = self.token_id_of(call_site);
-        self.subtree.push(SubtreeRepr {
-            open: call_site,
-            close: call_site,
-            kind: tt::DelimiterKind::Invisible,
-            tt: [!0, !0],
-        });
-        self.work.push_back((0, root.len(), Some(root.iter())));
-
+        if let Some(group) = root.as_single_group() {
+            self.enqueue(group);
+        } else {
+            self.subtree.push(SubtreeRepr {
+                open: call_site,
+                close: call_site,
+                kind: tt::DelimiterKind::Invisible,
+                tt: [!0, !0],
+            });
+            self.work.push_back((0, root.len(), Some(root.iter())));
+        }
         while let Some((idx, len, group)) = self.work.pop_front() {
             self.group(idx, len, group);
         }
@@ -962,6 +965,11 @@ impl<T: SpanTransformer> Reader<'_, T> {
             };
             res[i] = Some(g);
         }
-        res[0].take().unwrap().stream.unwrap_or_default()
+        let group = res[0].take().unwrap();
+        if group.delimiter == proc_macro_srv::Delimiter::None {
+            group.stream.unwrap_or_default()
+        } else {
+            TokenStream::new(vec![proc_macro_srv::TokenTree::Group(group)])
+        }
     }
 }

crates/proc-macro-srv/src/lib.rs

@@ -83,31 +83,10 @@ impl<'env> ProcMacroSrv<'env> {
     }
 
     pub fn join_spans(&self, first: Span, second: Span) -> Option<Span> {
-        // We can't modify the span range for fixup spans, those are meaningful to fixup, so just
-        // prefer the non-fixup span.
-        if first.anchor.ast_id == span::FIXUP_ERASED_FILE_AST_ID_MARKER {
-            return Some(second);
-        }
-        if second.anchor.ast_id == span::FIXUP_ERASED_FILE_AST_ID_MARKER {
-            return Some(first);
-        }
-        // FIXME: Once we can talk back to the client, implement a "long join" request for anchors
-        // that differ in [AstId]s as joining those spans requires resolving the AstIds.
-        if first.anchor != second.anchor {
-            return None;
-        }
-        // Differing context, we can't merge these so prefer the one that's root
-        if first.ctx != second.ctx {
-            if first.ctx.is_root() {
-                return Some(second);
-            } else if second.ctx.is_root() {
-                return Some(first);
-            }
-        }
-        Some(Span {
-            range: first.range.cover(second.range),
-            anchor: second.anchor,
-            ctx: second.ctx,
+        first.join(second, |_, _| {
+            // FIXME: Once we can talk back to the client, implement a "long join" request for anchors
+            // that differ in [AstId]s as joining those spans requires resolving the AstIds.
+            None
         })
     }
 }

crates/proc-macro-srv/src/token_stream.rs

@@ -40,6 +40,13 @@ impl<S> TokenStream<S> {
         TokenStreamIter::new(self)
     }
 
+    pub fn as_single_group(&self) -> Option<&Group<S>> {
+        match &**self.0 {
+            [TokenTree::Group(group)] => Some(group),
+            _ => None,
+        }
+    }
+
     pub(crate) fn from_str(s: &str, span: S) -> Result<Self, String>
     where
         S: SpanLike + Copy,

crates/span/src/ast_id.rs

@@ -44,6 +44,7 @@ pub const ROOT_ERASED_FILE_AST_ID: ErasedFileAstId =
 
 /// ErasedFileAstId used as the span for syntax node fixups. Any Span containing this file id is to be
 /// considered fake.
+/// Do not modify this, it is used by the proc-macro server.
 pub const FIXUP_ERASED_FILE_AST_ID_MARKER: ErasedFileAstId =
     ErasedFileAstId(pack_hash_index_and_kind(0, 0, ErasedFileAstIdKind::Fixup as u32));
 

crates/span/src/lib.rs

@@ -28,6 +28,33 @@ impl Span {
         let range = self.range.cover(other.range);
         Span { range, ..self }
     }
+
+    pub fn join(
+        self,
+        other: Span,
+        differing_anchor: impl FnOnce(Span, Span) -> Option<Span>,
+    ) -> Option<Span> {
+        // We can't modify the span range for fixup spans, those are meaningful to fixup, so just
+        // prefer the non-fixup span.
+        if self.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return Some(other);
+        }
+        if other.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return Some(self);
+        }
+        if self.anchor != other.anchor {
+            return differing_anchor(self, other);
+        }
+        // Differing context, we can't merge these so prefer the one that's root
+        if self.ctx != other.ctx {
+            if self.ctx.is_root() {
+                return Some(other);
+            } else if other.ctx.is_root() {
+                return Some(self);
+            }
+        }
+        Some(Span { range: self.range.cover(other.range), anchor: other.anchor, ctx: other.ctx })
+    }
 }
 
 /// Spans represent a region of code, used by the IDE to be able link macro inputs and outputs