Faster env snapshotting in proc-macro-srv

Author: Veykril

src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs

@@ -156,23 +156,26 @@ impl ProcMacro {
         let call_site = span_data_table.insert_full(call_site).0;
         let mixed_site = span_data_table.insert_full(mixed_site).0;
         let task = ExpandMacro {
-            macro_body: FlatTree::new(subtree, version, &mut span_data_table),
-            macro_name: self.name.to_string(),
-            attributes: attr.map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
+            data: msg::ExpandMacroData {
+                macro_body: FlatTree::new(subtree, version, &mut span_data_table),
+                macro_name: self.name.to_string(),
+                attributes: attr
+                    .map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
+                has_global_spans: ExpnGlobals {
+                    serialize: version >= HAS_GLOBAL_SPANS,
+                    def_site,
+                    call_site,
+                    mixed_site,
+                },
+                span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
+                    serialize_span_data_index_map(&span_data_table)
+                } else {
+                    Vec::new()
+                },
+            },
             lib: self.dylib_path.to_path_buf().into(),
             env: env.into(),
             current_dir,
-            has_global_spans: ExpnGlobals {
-                serialize: version >= HAS_GLOBAL_SPANS,
-                def_site,
-                call_site,
-                mixed_site,
-            },
-            span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
-                serialize_span_data_index_map(&span_data_table)
-            } else {
-                Vec::new()
-            },
         };
 
         let response = self.process.send_task(msg::Request::ExpandMacro(Box::new(task)))?;

src/tools/rust-analyzer/crates/proc-macro-api/src/msg.rs

@@ -72,6 +72,16 @@ pub struct PanicMessage(pub String);
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ExpandMacro {
+    pub lib: Utf8PathBuf,
+    /// Environment variables to set during macro expansion.
+    pub env: Vec<(String, String)>,
+    pub current_dir: Option<String>,
+    #[serde(flatten)]
+    pub data: ExpandMacroData,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ExpandMacroData {
     /// Argument of macro call.
     ///
     /// In custom derive this will be a struct or enum; in attribute-like macro - underlying
@@ -86,13 +96,6 @@ pub struct ExpandMacro {
 
     /// Possible attributes for the attribute-like macros.
     pub attributes: Option<FlatTree>,
-
-    pub lib: Utf8PathBuf,
-
-    /// Environment variables to set during macro expansion.
-    pub env: Vec<(String, String)>,
-
-    pub current_dir: Option<String>,
     /// marker for serde skip stuff
     #[serde(skip_serializing_if = "ExpnGlobals::skip_serializing_if")]
     #[serde(default)]
@@ -268,25 +271,30 @@ mod tests {
         let tt = fixture_token_tree();
         let mut span_data_table = Default::default();
         let task = ExpandMacro {
-            macro_body: FlatTree::new(&tt, CURRENT_API_VERSION, &mut span_data_table),
-            macro_name: Default::default(),
-            attributes: None,
+            data: ExpandMacroData {
+                macro_body: FlatTree::new(&tt, CURRENT_API_VERSION, &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(),
+            },
             lib: Utf8PathBuf::from_path_buf(std::env::current_dir().unwrap()).unwrap(),
             env: Default::default(),
             current_dir: Default::default(),
-            has_global_spans: ExpnGlobals {
-                serialize: true,
-                def_site: 0,
-                call_site: 0,
-                mixed_site: 0,
-            },
-            span_data_table: Vec::new(),
         };
 
         let json = serde_json::to_string(&task).unwrap();
         // println!("{}", json);
         let back: ExpandMacro = serde_json::from_str(&json).unwrap();
 
-        assert_eq!(tt, back.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
+        assert_eq!(
+            tt,
+            back.data.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table)
+        );
     }
 }

src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main.rs

@@ -33,12 +33,14 @@ fn run() -> io::Result<()> {
 #[cfg(any(feature = "sysroot-abi", rust_analyzer))]
 fn run() -> io::Result<()> {
     use proc_macro_api::msg::{self, Message};
+    use proc_macro_srv::EnvSnapshot;
 
     let read_request = |buf: &mut String| msg::Request::read(&mut io::stdin().lock(), buf);
 
     let write_response = |msg: msg::Response| msg.write(&mut io::stdout().lock());
 
-    let mut srv = proc_macro_srv::ProcMacroSrv::default();
+    let env = EnvSnapshot::new();
+    let mut srv = proc_macro_srv::ProcMacroSrv::new(&env);
     let mut buf = String::new();
 
     while let Some(req) = read_request(&mut buf)? {