3 files changed, 46 insertions, 42 deletions

diff --git a/crates/ra_proc_macro_srv/src/main.rs b/crates/ra_proc_macro_srv/src/cli.rs
index 70743c1f4..c771f2b38 100644
--- a/crates/ra_proc_macro_srv/src/main.rs
+++ b/crates/ra_proc_macro_srv/src/cli.rs
@@ -1,7 +1,7 @@
 //! Driver for proc macro server
 
+use crate::{expand_task, list_macros};
 use ra_proc_macro::msg::{self, Message};
-use ra_proc_macro_srv::{expand_task, list_macros};
 
 use std::io;
 
@@ -24,7 +24,8 @@ fn write_response(res: Result<msg::Response, String>) -> Result<(), io::Error> {
     let mut stdout = stdout.lock();
     msg.write(&mut stdout)
 }
-fn main() {
+
+pub fn run() {
     loop {
         let req = match read_request() {
             Err(err) => {
diff --git a/crates/ra_proc_macro_srv/src/dylib.rs b/crates/ra_proc_macro_srv/src/dylib.rs
index ec63d587b..16bd7466e 100644
--- a/crates/ra_proc_macro_srv/src/dylib.rs
+++ b/crates/ra_proc_macro_srv/src/dylib.rs
@@ -1,10 +1,12 @@
 //! Handles dynamic library loading for proc macro
 
 use crate::{proc_macro::bridge, rustc_server::TokenStream};
+use std::fs::File;
 use std::path::Path;
 
 use goblin::{mach::Mach, Object};
 use libloading::Library;
+use memmap::Mmap;
 use ra_proc_macro::ProcMacroKind;
 
 use std::io::Error as IoError;
@@ -16,55 +18,54 @@ fn invalid_data_err(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> I
     IoError::new(IoErrorKind::InvalidData, e)
 }
 
-fn get_symbols_from_lib(file: &Path) -> Result<Vec<String>, IoError> {
-    let buffer = std::fs::read(file)?;
+fn is_derive_registrar_symbol(symbol: &str) -> bool {
+    symbol.contains(NEW_REGISTRAR_SYMBOL)
+}
+
+fn find_registrar_symbol(file: &Path) -> Result<Option<String>, IoError> {
+    let file = File::open(file)?;
+    let buffer = unsafe { Mmap::map(&file)? };
     let object = Object::parse(&buffer).map_err(invalid_data_err)?;
 
     match object {
         Object::Elf(elf) => {
             let symbols = elf.dynstrtab.to_vec().map_err(invalid_data_err)?;
-            let names = symbols.iter().map(|s| s.to_string()).collect();
-            Ok(names)
+            let name =
+                symbols.iter().find(|s| is_derive_registrar_symbol(s)).map(|s| s.to_string());
+            Ok(name)
         }
         Object::PE(pe) => {
-            let symbol_names =
-                pe.exports.iter().flat_map(|s| s.name).map(|n| n.to_string()).collect();
-            Ok(symbol_names)
+            let name = pe
+                .exports
+                .iter()
+                .flat_map(|s| s.name)
+                .find(|s| is_derive_registrar_symbol(s))
+                .map(|s| s.to_string());
+            Ok(name)
         }
-        Object::Mach(mach) => match mach {
-            Mach::Binary(binary) => {
-                let exports = binary.exports().map_err(invalid_data_err)?;
-                let names = exports
-                    .into_iter()
-                    .map(|s| {
-                        // In macos doc:
-                        // 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 s.name.starts_with("_") {
-                            s.name[1..].to_string()
-                        } else {
-                            s.name
-                        }
-                    })
-                    .collect();
-                Ok(names)
-            }
-            Mach::Fat(_) => Ok(vec![]),
-        },
-        Object::Archive(_) | Object::Unknown(_) => Ok(vec![]),
+        Object::Mach(Mach::Binary(binary)) => {
+            let exports = binary.exports().map_err(invalid_data_err)?;
+            let name = exports
+                .iter()
+                .map(|s| {
+                    // In macos doc:
+                    // 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 s.name.starts_with("_") {
+                        &s.name[1..]
+                    } else {
+                        &s.name
+                    }
+                })
+                .find(|s| is_derive_registrar_symbol(s))
+                .map(|s| s.to_string());
+            Ok(name)
+        }
+        _ => Ok(None),
     }
 }
 
-fn is_derive_registrar_symbol(symbol: &str) -> bool {
-    symbol.contains(NEW_REGISTRAR_SYMBOL)
-}
-
-fn find_registrar_symbol(file: &Path) -> Result<Option<String>, IoError> {
-    let symbols = get_symbols_from_lib(file)?;
-    Ok(symbols.into_iter().find(|s| is_derive_registrar_symbol(s)))
-}
-
 /// Loads dynamic library in platform dependent manner.
 ///
 /// For unix, you have to use RTLD_DEEPBIND flag to escape problems described
diff --git a/crates/ra_proc_macro_srv/src/lib.rs b/crates/ra_proc_macro_srv/src/lib.rs
index 59716cbb3..c62b0ed89 100644
--- a/crates/ra_proc_macro_srv/src/lib.rs
+++ b/crates/ra_proc_macro_srv/src/lib.rs
@@ -22,7 +22,7 @@ mod dylib;
 use proc_macro::bridge::client::TokenStream;
 use ra_proc_macro::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask};
 
-pub fn expand_task(task: &ExpansionTask) -> Result<ExpansionResult, String> {
+pub(crate) fn expand_task(task: &ExpansionTask) -> Result<ExpansionResult, String> {
     let expander = dylib::Expander::new(&task.lib)
         .expect(&format!("Cannot expand with provided libraries: ${:?}", &task.lib));
 
@@ -39,7 +39,7 @@ pub fn expand_task(task: &ExpansionTask) -> Result<ExpansionResult, String> {
     }
 }
 
-pub fn list_macros(task: &ListMacrosTask) -> Result<ListMacrosResult, String> {
+pub(crate) fn list_macros(task: &ListMacrosTask) -> Result<ListMacrosResult, String> {
     let expander = dylib::Expander::new(&task.lib)
         .expect(&format!("Cannot expand with provided libraries: ${:?}", &task.lib));
 
@@ -53,5 +53,7 @@ pub fn list_macros(task: &ListMacrosTask) -> Result<ListMacrosResult, String> {
     }
 }
 
+pub mod cli;
+
 #[cfg(test)]
 mod tests;