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

   1 //! RA Proc Macro Server
   2 //!
   3 //! This library is able to call compiled Rust custom derive dynamic libraries on arbitrary code.
   4 //! The general idea here is based on <https://github.com/fedochet/rust-proc-macro-expander>.
   5 //!
   6 //! But we adapt it to better fit RA needs:
   7 //!
   8 //! * We use `tt` for proc-macro `TokenStream` server, it is easier to manipulate and interact with
   9 //!   RA than `proc-macro2` token stream.
  10 //! * By **copying** the whole rustc `lib_proc_macro` code, we are able to build this with `stable`
  11 //!   rustc rather than `unstable`. (Although in general ABI compatibility is still an issue)…
  12
  13 #![cfg(any(feature = "sysroot-abi", rust_analyzer))]
  14 #![feature(proc_macro_internals, proc_macro_diagnostic, proc_macro_span)]
  15 #![warn(rust_2018_idioms, unused_lifetimes)]
  16 #![allow(unreachable_pub, internal_features)]
  17
  18 extern crate proc_macro;
  19
  20 mod dylib;
  21 mod server;
  22 mod proc_macros;
  23
  24 use std::{
  25     collections::{hash_map::Entry, HashMap},
  26     env,
  27     ffi::OsString,
  28     fs,
  29     path::{Path, PathBuf},
  30     thread,
  31     time::SystemTime,
  32 };
  33
  34 use proc_macro_api::{
  35     msg::{self, ExpnGlobals, TokenId, CURRENT_API_VERSION},
  36     ProcMacroKind,
  37 };
  38
  39 mod tt {
  40     pub use proc_macro_api::msg::TokenId;
  41
  42     pub use ::tt::*;
  43
  44     pub type Subtree = ::tt::Subtree<TokenId>;
  45     pub type TokenTree = ::tt::TokenTree<TokenId>;
  46     pub type Delimiter = ::tt::Delimiter<TokenId>;
  47     pub type Leaf = ::tt::Leaf<TokenId>;
  48     pub type Literal = ::tt::Literal<TokenId>;
  49     pub type Punct = ::tt::Punct<TokenId>;
  50     pub type Ident = ::tt::Ident<TokenId>;
  51 }
  52
  53 // see `build.rs`
  54 include!(concat!(env!("OUT_DIR"), "/rustc_version.rs"));
  55
  56 #[derive(Default)]
  57 pub struct ProcMacroSrv {
  58     expanders: HashMap<(PathBuf, SystemTime), dylib::Expander>,
  59 }
  60
  61 const EXPANDER_STACK_SIZE: usize = 8 * 1024 * 1024;
  62
  63 impl ProcMacroSrv {
  64     pub fn expand(&mut self, task: msg::ExpandMacro) -> Result<msg::FlatTree, msg::PanicMessage> {
  65         let expander = self.expander(task.lib.as_ref()).map_err(|err| {
  66             debug_assert!(false, "should list macros before asking to expand");
  67             msg::PanicMessage(format!("failed to load macro: {err}"))
  68         })?;
  69
  70         let prev_env = EnvSnapshot::new();
  71         for (k, v) in &task.env {
  72             env::set_var(k, v);
  73         }
  74         let prev_working_dir = match task.current_dir {
  75             Some(dir) => {
  76                 let prev_working_dir = std::env::current_dir().ok();
  77                 if let Err(err) = std::env::set_current_dir(&dir) {
  78                     eprintln!("Failed to set the current working dir to {dir}. Error: {err:?}")
  79                 }
  80                 prev_working_dir
  81             }
  82             None => None,
  83         };
  84
  85         let ExpnGlobals { def_site, call_site, mixed_site, .. } = task.has_global_spans;
  86         let def_site = TokenId(def_site as u32);
  87         let call_site = TokenId(call_site as u32);
  88         let mixed_site = TokenId(mixed_site as u32);
  89
  90         let macro_body = task.macro_body.to_subtree_unresolved(CURRENT_API_VERSION);
  91         let attributes = task.attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION));
  92         let result = thread::scope(|s| {
  93             let thread = thread::Builder::new()
  94                 .stack_size(EXPANDER_STACK_SIZE)
  95                 .name(task.macro_name.clone())
  96                 .spawn_scoped(s, || {
  97                     expander
  98                         .expand(
  99                             &task.macro_name,
 100                             &macro_body,
 101                             attributes.as_ref(),
 102                             def_site,
 103                             call_site,
 104                             mixed_site,
 105                         )
 106                         .map(|it| msg::FlatTree::new_raw(&it, CURRENT_API_VERSION))
 107                 });
 108             let res = match thread {
 109                 Ok(handle) => handle.join(),
 110                 Err(e) => std::panic::resume_unwind(Box::new(e)),
 111             };
 112
 113             match res {
 114                 Ok(res) => res,
 115                 Err(e) => std::panic::resume_unwind(e),
 116             }
 117         });
 118
 119         prev_env.rollback();
 120
 121         if let Some(dir) = prev_working_dir {
 122             if let Err(err) = std::env::set_current_dir(&dir) {
 123                 eprintln!(
 124                     "Failed to set the current working dir to {}. Error: {:?}",
 125                     dir.display(),
 126                     err
 127                 )
 128             }
 129         }
 130
 131         result.map_err(msg::PanicMessage)
 132     }
 133
 134     pub fn list_macros(
 135         &mut self,
 136         dylib_path: &Path,
 137     ) -> Result<Vec<(String, ProcMacroKind)>, String> {
 138         let expander = self.expander(dylib_path)?;
 139         Ok(expander.list_macros())
 140     }
 141
 142     fn expander(&mut self, path: &Path) -> Result<&dylib::Expander, String> {
 143         let time = fs::metadata(path)
 144             .and_then(|it| it.modified())
 145             .map_err(|err| format!("Failed to get file metadata for {}: {err}", path.display()))?;
 146
 147         Ok(match self.expanders.entry((path.to_path_buf(), time)) {
 148             Entry::Vacant(v) => {
 149                 v.insert(dylib::Expander::new(path).map_err(|err| {
 150                     format!("Cannot create expander for {}: {err}", path.display())
 151                 })?)
 152             }
 153             Entry::Occupied(e) => e.into_mut(),
 154         })
 155     }
 156 }
 157
 158 pub struct PanicMessage {
 159     message: Option<String>,
 160 }
 161
 162 impl PanicMessage {
 163     pub fn into_string(self) -> Option<String> {
 164         self.message
 165     }
 166 }
 167
 168 struct EnvSnapshot {
 169     vars: HashMap<OsString, OsString>,
 170 }
 171
 172 impl EnvSnapshot {
 173     fn new() -> EnvSnapshot {
 174         EnvSnapshot { vars: env::vars_os().collect() }
 175     }
 176
 177     fn rollback(self) {}
 178 }
 179
 180 impl Drop for EnvSnapshot {
 181     fn drop(&mut self) {
 182         for (name, value) in env::vars_os() {
 183             let old_value = self.vars.remove(&name);
 184             if old_value != Some(value) {
 185                 match old_value {
 186                     None => env::remove_var(name),
 187                     Some(old_value) => env::set_var(name, old_value),
 188                 }
 189             }
 190         }
 191         for (name, old_value) in self.vars.drain() {
 192             env::set_var(name, old_value)
 193         }
 194     }
 195 }
 196
 197 #[cfg(test)]
 198 mod tests;
 199
 200 #[cfg(test)]
 201 pub fn proc_macro_test_dylib_path() -> std::path::PathBuf {
 202     proc_macro_test::PROC_MACRO_TEST_LOCATION.into()
 203 }