Cleanup proc-macro dylib handling

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

@@ -9,37 +9,7 @@ use libloading::Library;
 use object::Object;
 use paths::{Utf8Path, Utf8PathBuf};
 
-use crate::{proc_macros::ProcMacroKind, ProcMacroSrvSpan};
-
-const NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";
-
-fn invalid_data_err(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
-    io::Error::new(io::ErrorKind::InvalidData, e)
-}
-
-fn is_derive_registrar_symbol(symbol: &str) -> bool {
-    symbol.contains(NEW_REGISTRAR_SYMBOL)
-}
-
-fn find_registrar_symbol(obj: &object::File<'_>) -> object::Result<Option<String>> {
-    Ok(obj
-        .exports()?
-        .into_iter()
-        .map(|export| export.name())
-        .filter_map(|sym| String::from_utf8(sym.into()).ok())
-        .find(|sym| is_derive_registrar_symbol(sym))
-        .map(|sym| {
-            // From MacOS docs:
-            // https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/dlsym.3.html
-            // Unlike other dyld API's, the symbol name passed to dlsym() must NOT be
-            // prepended with an underscore.
-            if cfg!(target_os = "macos") && sym.starts_with('_') {
-                sym[1..].to_owned()
-            } else {
-                sym
-            }
-        }))
-}
+use crate::{proc_macros::ProcMacros, ProcMacroKind, ProcMacroSrvSpan};
 
 /// Loads dynamic library in platform dependent manner.
 ///
@@ -99,13 +69,14 @@ impl From<libloading::Error> for LoadProcMacroDylibError {
     }
 }
 
-struct ProcMacroLibraryLibloading {
+struct ProcMacroLibrary {
+    // 'static is actually the lifetime of library, so make sure this drops before _lib
+    proc_macros: &'static ProcMacros,
     // Hold on to the library so it doesn't unload
     _lib: Library,
-    proc_macros: crate::proc_macros::ProcMacros,
 }
 
-impl ProcMacroLibraryLibloading {
+impl ProcMacroLibrary {
     fn open(path: &Utf8Path) -> Result<Self, LoadProcMacroDylibError> {
         let file = fs::File::open(path)?;
         let file = unsafe { memmap2::Mmap::map(&file) }?;
@@ -118,27 +89,22 @@ impl ProcMacroLibraryLibloading {
             })?;
 
         let lib = load_library(path).map_err(invalid_data_err)?;
-        let proc_macros = crate::proc_macros::ProcMacros::from_lib(
+        let proc_macros = unsafe {
+            // SAFETY: We extend the lifetime here to avoid referential borrow problems
+            // We never reveal proc_macros to the outside and drop it before _lib
+            std::mem::transmute::<&ProcMacros, &'static ProcMacros>(ProcMacros::from_lib(
                 &lib,
                 symbol_name,
                 &version_info.version_string,
-        )?;
-        Ok(ProcMacroLibraryLibloading { _lib: lib, proc_macros })
-    }
-}
-
-struct RemoveFileOnDrop(Utf8PathBuf);
-impl Drop for RemoveFileOnDrop {
-    fn drop(&mut self) {
-        #[cfg(windows)]
-        std::fs::remove_file(&self.0).unwrap();
-        _ = self.0;
+            )?)
+        };
+        Ok(ProcMacroLibrary { _lib: lib, proc_macros })
     }
 }
 
 // Drop order matters as we can't remove the dylib before the library is unloaded
 pub(crate) struct Expander {
-    inner: ProcMacroLibraryLibloading,
+    inner: ProcMacroLibrary,
     _remove_on_drop: RemoveFileOnDrop,
     modified_time: SystemTime,
 }
@@ -151,7 +117,7 @@ impl Expander {
         let modified_time = fs::metadata(&lib).and_then(|it| it.modified())?;
 
         let path = ensure_file_with_lock_free_access(&lib)?;
-        let library = ProcMacroLibraryLibloading::open(path.as_ref())?;
+        let library = ProcMacroLibrary::open(path.as_ref())?;
 
         Ok(Expander { inner: library, _remove_on_drop: RemoveFileOnDrop(path), modified_time })
     }
@@ -184,6 +150,44 @@ impl Expander {
     }
 }
 
+fn invalid_data_err(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
+    io::Error::new(io::ErrorKind::InvalidData, e)
+}
+
+fn is_derive_registrar_symbol(symbol: &str) -> bool {
+    const NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";
+    symbol.contains(NEW_REGISTRAR_SYMBOL)
+}
+
+fn find_registrar_symbol(obj: &object::File<'_>) -> object::Result<Option<String>> {
+    Ok(obj
+        .exports()?
+        .into_iter()
+        .map(|export| export.name())
+        .filter_map(|sym| String::from_utf8(sym.into()).ok())
+        .find(|sym| is_derive_registrar_symbol(sym))
+        .map(|sym| {
+            // From MacOS docs:
+            // https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/dlsym.3.html
+            // Unlike other dyld API's, the symbol name passed to dlsym() must NOT be
+            // prepended with an underscore.
+            if cfg!(target_os = "macos") && sym.starts_with('_') {
+                sym[1..].to_owned()
+            } else {
+                sym
+            }
+        }))
+}
+
+struct RemoveFileOnDrop(Utf8PathBuf);
+impl Drop for RemoveFileOnDrop {
+    fn drop(&mut self) {
+        #[cfg(windows)]
+        std::fs::remove_file(&self.0).unwrap();
+        _ = self.0;
+    }
+}
+
 /// Copy the dylib to temp directory to prevent locking in Windows
 #[cfg(windows)]
 fn ensure_file_with_lock_free_access(path: &Utf8Path) -> io::Result<Utf8PathBuf> {

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

@@ -43,7 +43,12 @@ use span::{Span, TokenId};
 
 use crate::server_impl::TokenStream;
 
-pub use crate::proc_macros::ProcMacroKind;
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+pub enum ProcMacroKind {
+    CustomDerive,
+    Attr,
+    Bang,
+}
 
 pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");
 

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

@@ -4,18 +4,10 @@ use proc_macro::bridge;
 
 use libloading::Library;
 
-use crate::{dylib::LoadProcMacroDylibError, ProcMacroSrvSpan};
+use crate::{dylib::LoadProcMacroDylibError, ProcMacroKind, ProcMacroSrvSpan};
 
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
-pub enum ProcMacroKind {
-    CustomDerive,
-    Attr,
-    Bang,
-}
-
-pub(crate) struct ProcMacros {
-    exported_macros: Vec<bridge::client::ProcMacro>,
-}
+#[repr(transparent)]
+pub(crate) struct ProcMacros([bridge::client::ProcMacro]);
 
 impl From<bridge::PanicMessage> for crate::PanicMessage {
     fn from(p: bridge::PanicMessage) -> Self {
@@ -33,18 +25,17 @@ impl ProcMacros {
     /// *`info` - RustCInfo about the compiler that was used to compile the
     ///           macro crate. This is the information we use to figure out
     ///           which ABI to return
-    pub(crate) fn from_lib(
-        lib: &Library,
+    pub(crate) fn from_lib<'l>(
+        lib: &'l Library,
         symbol_name: String,
         version_string: &str,
-    ) -> Result<ProcMacros, LoadProcMacroDylibError> {
-        if version_string == crate::RUSTC_VERSION_STRING {
-            let macros =
-                unsafe { lib.get::<&&[bridge::client::ProcMacro]>(symbol_name.as_bytes()) }?;
-
-            return Ok(Self { exported_macros: macros.to_vec() });
+    ) -> Result<&'l ProcMacros, LoadProcMacroDylibError> {
+        if version_string != crate::RUSTC_VERSION_STRING {
+            return Err(LoadProcMacroDylibError::AbiMismatch(version_string.to_owned()));
         }
-        Err(LoadProcMacroDylibError::AbiMismatch(version_string.to_owned()))
+        unsafe { lib.get::<&'l &'l ProcMacros>(symbol_name.as_bytes()) }
+            .map(|it| **it)
+            .map_err(Into::into)
     }
 
     pub(crate) fn expand<S: ProcMacroSrvSpan>(
@@ -63,7 +54,7 @@ impl ProcMacros {
                 crate::server_impl::TokenStream::with_subtree(attr)
             });
 
-        for proc_macro in &self.exported_macros {
+        for proc_macro in &self.0 {
             match proc_macro {
                 bridge::client::ProcMacro::CustomDerive { trait_name, client, .. }
                     if *trait_name == macro_name =>
@@ -109,7 +100,7 @@ impl ProcMacros {
     }
 
     pub(crate) fn list_macros(&self) -> Vec<(String, ProcMacroKind)> {
-        self.exported_macros
+        self.0
             .iter()
             .map(|proc_macro| match proc_macro {
                 bridge::client::ProcMacro::CustomDerive { trait_name, .. } => {