4 files changed, 62 insertions, 108 deletions

diff --git a/crates/proc_macro_api/Cargo.toml b/crates/proc_macro_api/Cargo.toml
index 3863e5189..1af2bbca7 100644
--- a/crates/proc_macro_api/Cargo.toml
+++ b/crates/proc_macro_api/Cargo.toml
@@ -17,3 +17,4 @@ crossbeam-channel = "0.5.0"
 jod-thread = "0.1.1"
 
 tt = { path = "../tt", version = "0.0.0" }
+base_db = { path = "../base_db", version = "0.0.0" }
diff --git a/crates/proc_macro_api/src/lib.rs b/crates/proc_macro_api/src/lib.rs
index d5e87cf7d..0d061fd53 100644
--- a/crates/proc_macro_api/src/lib.rs
+++ b/crates/proc_macro_api/src/lib.rs
@@ -16,6 +16,7 @@ use std::{
     sync::Arc,
 };
 
+use base_db::ProcMacro;
 use tt::{SmolStr, Subtree};
 
 use crate::process::{ProcMacroProcessSrv, ProcMacroProcessThread};
@@ -23,7 +24,7 @@ use crate::process::{ProcMacroProcessSrv, ProcMacroProcessThread};
 pub use rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};
 
 #[derive(Debug, Clone)]
-pub struct ProcMacroProcessExpander {
+struct ProcMacroProcessExpander {
     process: Arc<ProcMacroProcessSrv>,
     dylib_path: PathBuf,
     name: SmolStr,
@@ -42,21 +43,24 @@ impl tt::TokenExpander for ProcMacroProcessExpander {
     fn expand(
         &self,
         subtree: &Subtree,
-        _attr: Option<&Subtree>,
+        attr: Option<&Subtree>,
     ) -> Result<Subtree, tt::ExpansionError> {
-        self.process.custom_derive(&self.dylib_path, subtree, &self.name)
-    }
-}
+        let task = ExpansionTask {
+            macro_body: subtree.clone(),
+            macro_name: self.name.to_string(),
+            attributes: attr.cloned(),
+            lib: self.dylib_path.to_path_buf(),
+        };
 
-#[derive(Debug)]
-enum ProcMacroClientKind {
-    Process { process: Arc<ProcMacroProcessSrv>, thread: ProcMacroProcessThread },
-    Dummy,
+        let result: ExpansionResult = self.process.send_task(msg::Request::ExpansionMacro(task))?;
+        Ok(result.expansion)
+    }
 }
 
 #[derive(Debug)]
 pub struct ProcMacroClient {
-    kind: ProcMacroClientKind,
+    process: Arc<ProcMacroProcessSrv>,
+    thread: ProcMacroProcessThread,
 }
 
 impl ProcMacroClient {
@@ -65,47 +69,35 @@ impl ProcMacroClient {
         args: impl IntoIterator<Item = impl AsRef<OsStr>>,
     ) -> io::Result<ProcMacroClient> {
         let (thread, process) = ProcMacroProcessSrv::run(process_path, args)?;
-        Ok(ProcMacroClient {
-            kind: ProcMacroClientKind::Process { process: Arc::new(process), thread },
-        })
+        Ok(ProcMacroClient { process: Arc::new(process), thread })
     }
 
-    pub fn dummy() -> ProcMacroClient {
-        ProcMacroClient { kind: ProcMacroClientKind::Dummy }
-    }
+    pub fn by_dylib_path(&self, dylib_path: &Path) -> Vec<ProcMacro> {
+        let macros = match self.process.find_proc_macros(dylib_path) {
+            Err(err) => {
+                eprintln!("Failed to find proc macros. Error: {:#?}", err);
+                return vec![];
+            }
+            Ok(macros) => macros,
+        };
 
-    pub fn by_dylib_path(&self, dylib_path: &Path) -> Vec<(SmolStr, Arc<dyn tt::TokenExpander>)> {
-        match &self.kind {
-            ProcMacroClientKind::Dummy => vec![],
-            ProcMacroClientKind::Process { process, .. } => {
-                let macros = match process.find_proc_macros(dylib_path) {
-                    Err(err) => {
-                        eprintln!("Failed to find proc macros. Error: {:#?}", err);
-                        return vec![];
-                    }
-                    Ok(macros) => macros,
+        macros
+            .into_iter()
+            .map(|(name, kind)| {
+                let name = SmolStr::new(&name);
+                let kind = match kind {
+                    ProcMacroKind::CustomDerive => base_db::ProcMacroKind::CustomDerive,
+                    ProcMacroKind::FuncLike => base_db::ProcMacroKind::FuncLike,
+                    ProcMacroKind::Attr => base_db::ProcMacroKind::Attr,
                 };
+                let expander: Arc<dyn tt::TokenExpander> = Arc::new(ProcMacroProcessExpander {
+                    process: self.process.clone(),
+                    name: name.clone(),
+                    dylib_path: dylib_path.into(),
+                });
 
-                macros
-                    .into_iter()
-                    .filter_map(|(name, kind)| {
-                        match kind {
-                            ProcMacroKind::CustomDerive | ProcMacroKind::FuncLike => {
-                                let name = SmolStr::new(&name);
-                                let expander: Arc<dyn tt::TokenExpander> =
-                                    Arc::new(ProcMacroProcessExpander {
-                                        process: process.clone(),
-                                        name: name.clone(),
-                                        dylib_path: dylib_path.into(),
-                                    });
-                                Some((name, expander))
-                            }
-                            // FIXME: Attribute macro are currently unsupported.
-                            ProcMacroKind::Attr => None,
-                        }
-                    })
-                    .collect()
-            }
-        }
+                ProcMacro { name, kind, expander }
+            })
+            .collect()
     }
 }
diff --git a/crates/proc_macro_api/src/process.rs b/crates/proc_macro_api/src/process.rs
index 907cb3db7..d68723ada 100644
--- a/crates/proc_macro_api/src/process.rs
+++ b/crates/proc_macro_api/src/process.rs
@@ -10,16 +10,15 @@ use std::{
 };
 
 use crossbeam_channel::{bounded, Receiver, Sender};
-use tt::Subtree;
 
 use crate::{
     msg::{ErrorCode, Message, Request, Response, ResponseError},
-    rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind},
+    rpc::{ListMacrosResult, ListMacrosTask, ProcMacroKind},
 };
 
 #[derive(Debug, Default)]
 pub(crate) struct ProcMacroProcessSrv {
-    inner: Option<Weak<Sender<Task>>>,
+    inner: Weak<Sender<Task>>,
 }
 
 #[derive(Debug)]
@@ -42,7 +41,7 @@ impl ProcMacroProcessSrv {
         });
 
         let task_tx = Arc::new(task_tx);
-        let srv = ProcMacroProcessSrv { inner: Some(Arc::downgrade(&task_tx)) };
+        let srv = ProcMacroProcessSrv { inner: Arc::downgrade(&task_tx) };
         let thread = ProcMacroProcessThread { handle, sender: task_tx };
 
         Ok((thread, srv))
@@ -58,43 +57,22 @@ impl ProcMacroProcessSrv {
         Ok(result.macros)
     }
 
-    pub(crate) fn custom_derive(
-        &self,
-        dylib_path: &Path,
-        subtree: &Subtree,
-        derive_name: &str,
-    ) -> Result<Subtree, tt::ExpansionError> {
-        let task = ExpansionTask {
-            macro_body: subtree.clone(),
-            macro_name: derive_name.to_string(),
-            attributes: None,
-            lib: dylib_path.to_path_buf(),
-        };
-
-        let result: ExpansionResult = self.send_task(Request::ExpansionMacro(task))?;
-        Ok(result.expansion)
-    }
-
     pub(crate) fn send_task<R>(&self, req: Request) -> Result<R, tt::ExpansionError>
     where
         R: TryFrom<Response, Error = &'static str>,
     {
-        let sender = match &self.inner {
-            None => return Err(tt::ExpansionError::Unknown("No sender is found.".to_string())),
-            Some(it) => it,
-        };
-
         let (result_tx, result_rx) = bounded(0);
-        let sender = match sender.upgrade() {
-            None => {
-                return Err(tt::ExpansionError::Unknown("Proc macro process is closed.".into()))
-            }
+        let sender = match self.inner.upgrade() {
+            None => return Err(tt::ExpansionError::Unknown("proc macro process is closed".into())),
             Some(it) => it,
         };
-        sender.send(Task { req, result_tx }).unwrap();
+        sender
+            .send(Task { req, result_tx })
+            .map_err(|_| tt::ExpansionError::Unknown("proc macro server crashed".into()))?;
+
         let res = result_rx
             .recv()
-            .map_err(|_| tt::ExpansionError::Unknown("Proc macro thread is closed.".into()))?;
+            .map_err(|_| tt::ExpansionError::Unknown("proc macro server crashed".into()))?;
 
         match res {
             Some(Response::Error(err)) => {
@@ -109,32 +87,23 @@ impl ProcMacroProcessSrv {
 }
 
 fn client_loop(task_rx: Receiver<Task>, mut process: Process) {
-    let (mut stdin, mut stdout) = match process.stdio() {
-        None => return,
-        Some(it) => it,
-    };
-
-    for task in task_rx {
-        let Task { req, result_tx } = task;
+    let (mut stdin, mut stdout) = process.stdio().expect("couldn't access child stdio");
 
+    for Task { req, result_tx } in task_rx {
         match send_request(&mut stdin, &mut stdout, req) {
             Ok(res) => result_tx.send(res).unwrap(),
             Err(_err) => {
+                log::error!(
+                    "proc macro server crashed, server process state: {:?}",
+                    process.child.try_wait()
+                );
                 let res = Response::Error(ResponseError {
                     code: ErrorCode::ServerErrorEnd,
-                    message: "Server closed".into(),
+                    message: "proc macro server crashed".into(),
                 });
                 result_tx.send(res.into()).unwrap();
-                // Restart the process
-                if process.restart().is_err() {
-                    break;
-                }
-                let stdio = match process.stdio() {
-                    None => break,
-                    Some(it) => it,
-                };
-                stdin = stdio.0;
-                stdout = stdio.1;
+                // Exit the thread.
+                break;
             }
         }
     }
@@ -146,8 +115,6 @@ struct Task {
 }
 
 struct Process {
-    path: PathBuf,
-    args: Vec<OsString>,
     child: Child,
 }
 
@@ -162,15 +129,9 @@ impl Process {
         path: PathBuf,
         args: impl IntoIterator<Item = impl AsRef<OsStr>>,
     ) -> io::Result<Process> {
-        let args = args.into_iter().map(|s| s.as_ref().into()).collect();
+        let args: Vec<OsString> = args.into_iter().map(|s| s.as_ref().into()).collect();
         let child = mk_child(&path, &args)?;
-        Ok(Process { path, args, child })
-    }
-
-    fn restart(&mut self) -> io::Result<()> {
-        let _ = self.child.kill();
-        self.child = mk_child(&self.path, &self.args)?;
-        Ok(())
+        Ok(Process { child })
     }
 
     fn stdio(&mut self) -> Option<(impl Write, impl BufRead)> {
diff --git a/crates/proc_macro_api/src/rpc.rs b/crates/proc_macro_api/src/rpc.rs
index 203109ca4..b85f92eea 100644
--- a/crates/proc_macro_api/src/rpc.rs
+++ b/crates/proc_macro_api/src/rpc.rs
@@ -19,7 +19,7 @@ pub struct ListMacrosTask {
     pub lib: PathBuf,
 }
 
-#[derive(Clone, Eq, PartialEq, Debug, Serialize, Deserialize)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, Serialize, Deserialize)]
 pub enum ProcMacroKind {
     CustomDerive,
     FuncLike,