Skip to content

Commit

Permalink
Merge pull request rust-lang#18669 from Veykril/push-qqkuxtvsmsut
Browse files Browse the repository at this point in the history
internal: Only parse the object file once in proc-macro-srv
  • Loading branch information
Veykril authored Dec 12, 2024
2 parents 096e3e5 + 8a86fa0 commit 3bd4597
Show file tree
Hide file tree
Showing 8 changed files with 31 additions and 32 deletions.
2 changes: 1 addition & 1 deletion crates/proc-macro-srv-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ fn run() -> io::Result<()> {

let write_response = |msg: msg::Response| msg.write(write_json, &mut io::stdout().lock());

let env = EnvSnapshot::new();
let env = EnvSnapshot::default();
let mut srv = proc_macro_srv::ProcMacroSrv::new(&env);
let mut buf = String::new();

Expand Down
16 changes: 8 additions & 8 deletions crates/proc-macro-srv/src/dylib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use proc_macro::bridge;
use std::{fmt, fs, io, time::SystemTime};

use libloading::Library;
use memmap2::Mmap;
use object::Object;
use paths::{Utf8Path, Utf8PathBuf};
use proc_macro_api::ProcMacroKind;
Expand All @@ -23,8 +22,8 @@ fn is_derive_registrar_symbol(symbol: &str) -> bool {
symbol.contains(NEW_REGISTRAR_SYMBOL)
}

fn find_registrar_symbol(buffer: &[u8]) -> object::Result<Option<String>> {
Ok(object::File::parse(buffer)?
fn find_registrar_symbol(obj: &object::File<'_>) -> object::Result<Option<String>> {
Ok(obj
.exports()?
.into_iter()
.map(|export| export.name())
Expand Down Expand Up @@ -109,15 +108,16 @@ struct ProcMacroLibraryLibloading {

impl ProcMacroLibraryLibloading {
fn open(path: &Utf8Path) -> Result<Self, LoadProcMacroDylibError> {
let buffer = unsafe { Mmap::map(&fs::File::open(path)?)? };
let file = fs::File::open(path)?;
let file = unsafe { memmap2::Mmap::map(&file) }?;
let obj = object::File::parse(&*file)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
let version_info = version::read_dylib_info(&obj)?;
let symbol_name =
find_registrar_symbol(&buffer).map_err(invalid_data_err)?.ok_or_else(|| {
find_registrar_symbol(&obj).map_err(invalid_data_err)?.ok_or_else(|| {
invalid_data_err(format!("Cannot find registrar symbol in file {path}"))
})?;

let version_info = version::read_dylib_info(&buffer)?;
drop(buffer);

let lib = load_library(path).map_err(invalid_data_err)?;
let proc_macros = crate::proc_macros::ProcMacros::from_lib(
&lib,
Expand Down
24 changes: 11 additions & 13 deletions crates/proc-macro-srv/src/dylib/version.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
use std::io::{self, Read};

use object::read::{File as BinaryFile, Object, ObjectSection};
use object::read::{Object, ObjectSection};

#[derive(Debug)]
#[allow(dead_code)]
Expand All @@ -16,14 +16,14 @@ pub struct RustCInfo {
}

/// Read rustc dylib information
pub fn read_dylib_info(buffer: &[u8]) -> io::Result<RustCInfo> {
pub fn read_dylib_info(obj: &object::File<'_>) -> io::Result<RustCInfo> {
macro_rules! err {
($e:literal) => {
io::Error::new(io::ErrorKind::InvalidData, $e)
};
}

let ver_str = read_version(buffer)?;
let ver_str = read_version(obj)?;
let mut items = ver_str.split_whitespace();
let tag = items.next().ok_or_else(|| err!("version format error"))?;
if tag != "rustc" {
Expand Down Expand Up @@ -70,10 +70,8 @@ pub fn read_dylib_info(buffer: &[u8]) -> io::Result<RustCInfo> {

/// This is used inside read_version() to locate the ".rustc" section
/// from a proc macro crate's binary file.
fn read_section<'a>(dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'a [u8]> {
BinaryFile::parse(dylib_binary)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
.section_by_name(section_name)
fn read_section<'a>(obj: &object::File<'a>, section_name: &str) -> io::Result<&'a [u8]> {
obj.section_by_name(section_name)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "section read error"))?
.data()
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
Expand Down Expand Up @@ -101,8 +99,8 @@ fn read_section<'a>(dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'
///
/// Check this issue for more about the bytes layout:
/// <https://github.com/rust-lang/rust-analyzer/issues/6174>
pub fn read_version(buffer: &[u8]) -> io::Result<String> {
let dot_rustc = read_section(buffer, ".rustc")?;
pub fn read_version(obj: &object::File<'_>) -> io::Result<String> {
let dot_rustc = read_section(obj, ".rustc")?;

// check if magic is valid
if &dot_rustc[0..4] != b"rust" {
Expand Down Expand Up @@ -151,10 +149,10 @@ pub fn read_version(buffer: &[u8]) -> io::Result<String> {

#[test]
fn test_version_check() {
let info = read_dylib_info(&unsafe {
memmap2::Mmap::map(&std::fs::File::open(crate::proc_macro_test_dylib_path()).unwrap())
.unwrap()
})
let info = read_dylib_info(
&object::File::parse(&*std::fs::read(crate::proc_macro_test_dylib_path()).unwrap())
.unwrap(),
)
.unwrap();

assert_eq!(
Expand Down
10 changes: 5 additions & 5 deletions crates/proc-macro-srv/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
#![cfg(any(feature = "sysroot-abi", rust_analyzer))]
#![cfg_attr(feature = "in-rust-tree", feature(rustc_private))]
#![feature(proc_macro_internals, proc_macro_diagnostic, proc_macro_span)]
#![allow(unreachable_pub, internal_features)]
#![allow(unreachable_pub, internal_features, clippy::disallowed_types, clippy::print_stderr)]

extern crate proc_macro;
#[cfg(feature = "in-rust-tree")]
Expand Down Expand Up @@ -65,7 +65,7 @@ impl<'env> ProcMacroSrv<'env> {

const EXPANDER_STACK_SIZE: usize = 8 * 1024 * 1024;

impl<'env> ProcMacroSrv<'env> {
impl ProcMacroSrv<'_> {
pub fn set_span_mode(&mut self, span_mode: SpanMode) {
self.span_mode = span_mode;
}
Expand Down Expand Up @@ -248,8 +248,8 @@ pub struct EnvSnapshot {
vars: HashMap<OsString, OsString>,
}

impl EnvSnapshot {
pub fn new() -> EnvSnapshot {
impl Default for EnvSnapshot {
fn default() -> EnvSnapshot {
EnvSnapshot { vars: env::vars_os().collect() }
}
}
Expand Down Expand Up @@ -305,7 +305,7 @@ impl Drop for EnvChange<'_> {
}

if let Some(dir) = &self.prev_working_dir {
if let Err(err) = std::env::set_current_dir(&dir) {
if let Err(err) = std::env::set_current_dir(dir) {
eprintln!(
"Failed to set the current working dir to {}. Error: {:?}",
dir.display(),
Expand Down
2 changes: 1 addition & 1 deletion crates/proc-macro-srv/src/proc_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ pub(crate) struct ProcMacros {

impl From<bridge::PanicMessage> for crate::PanicMessage {
fn from(p: bridge::PanicMessage) -> Self {
Self { message: p.as_str().map(|s| s.to_string()) }
Self { message: p.as_str().map(|s| s.to_owned()) }
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -498,7 +498,7 @@ mod tests {
})),
tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident {
sym: Symbol::intern("T"),
span: span,
span,
is_raw: tt::IdentIsRaw::No,
})),
tt::TokenTree::Subtree(tt::Subtree {
Expand Down
5 changes: 3 additions & 2 deletions crates/proc-macro-srv/src/server_impl/token_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ pub(super) struct TokenStreamBuilder<S> {
}

/// pub(super)lic implementation details for the `TokenStream` type, such as iterators.
pub(super) mod token_stream {
pub(super) mod token_stream_impls {

use core::fmt;

Expand Down Expand Up @@ -137,6 +137,7 @@ pub(super) mod token_stream {
}
}

#[allow(clippy::to_string_trait_impl)]
impl<S> ToString for TokenStream<S> {
fn to_string(&self) -> String {
::tt::pretty(&self.token_trees)
Expand All @@ -150,7 +151,7 @@ impl<S> TokenStreamBuilder<S> {
}

pub(super) fn push(&mut self, stream: TokenStream<S>) {
self.acc.extend(stream.into_iter())
self.acc.extend(stream)
}

pub(super) fn build(self) -> TokenStream<S> {
Expand Down
2 changes: 1 addition & 1 deletion crates/proc-macro-srv/src/tests/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ fn assert_expand_impl(

pub(crate) fn list() -> Vec<String> {
let dylib_path = proc_macro_test_dylib_path();
let env = EnvSnapshot::new();
let env = EnvSnapshot::default();
let mut srv = ProcMacroSrv::new(&env);
let res = srv.list_macros(&dylib_path).unwrap();
res.into_iter().map(|(name, kind)| format!("{name} [{kind:?}]")).collect()
Expand Down

0 comments on commit 3bd4597

Please sign in to comment.