1 files changed, 163 insertions, 118 deletions

diff --git a/crates/server/src/main_loop/mod.rs b/crates/server/src/main_loop/mod.rs
index db7d5ae34..34d077805 100644
--- a/crates/server/src/main_loop/mod.rs
+++ b/crates/server/src/main_loop/mod.rs
@@ -6,59 +6,97 @@ use std::{
 };
 
 use threadpool::ThreadPool;
-use crossbeam_channel::{Sender, Receiver};
+use serde::{Serialize, de::DeserializeOwned};
+use crossbeam_channel::{bounded, Sender, Receiver};
 use languageserver_types::{NumberOrString};
 use libanalysis::{FileId, JobHandle, JobToken};
+use gen_lsp_server::{RawRequest, RawNotification, RawMessage, RawResponse, ErrorCode};
 
 use {
-    req, dispatch,
-    Task, Result,
-    io::{Io, RawMsg, RawRequest, RawNotification},
-    vfs::FileEvent,
+    req,
+    Result,
+    vfs::{self, FileEvent},
     server_world::{ServerWorldState, ServerWorld},
     main_loop::subscriptions::{Subscriptions},
 };
 
+enum Task {
+    Respond(RawResponse),
+    Notify(RawNotification),
+}
+
 pub(super) fn main_loop(
-    io: &mut Io,
-    pool: &mut ThreadPool,
-    task_sender: Sender<Task>,
-    task_receiver: Receiver<Task>,
-    fs_events_receiver: Receiver<Vec<FileEvent>>,
+    receriver: &mut Receiver<RawMessage>,
+    sender: &mut Sender<RawMessage>,
 ) -> Result<()> {
+    let pool = ThreadPool::new(4);
+    let (task_sender, task_receiver) = bounded::<Task>(16);
+    let (fs_events_receiver, watcher) = vfs::watch(vec![
+        ::std::env::current_dir()?,
+    ]);
+
     info!("server initialized, serving requests");
     let mut state = ServerWorldState::new();
 
-    let mut pending_requests: HashMap<u64, JobHandle> = HashMap::new();
-    let mut fs_events_receiver = Some(&fs_events_receiver);
+    let mut pending_requests = HashMap::new();
     let mut subs = Subscriptions::new();
+    main_loop_inner(
+        &pool,
+        receriver,
+        sender,
+        task_receiver.clone(),
+        task_sender,
+        fs_events_receiver,
+        &mut state,
+        &mut pending_requests,
+        &mut subs,
+    )?;
+
+    info!("waiting for background jobs to finish...");
+    task_receiver.for_each(drop);
+    pool.join();
+    info!("...background jobs have finished");
+
+    info!("waiting for file watcher to finish...");
+    watcher.stop()?;
+    info!("...file watcher has finished");
+    Ok(())
+}
+
+fn main_loop_inner(
+    pool: &ThreadPool,
+    msg_receiver: &mut Receiver<RawMessage>,
+    msg_sender: &mut Sender<RawMessage>,
+    task_receiver: Receiver<Task>,
+    task_sender: Sender<Task>,
+    fs_receiver: Receiver<Vec<FileEvent>>,
+    state: &mut ServerWorldState,
+    pending_requests: &mut HashMap<u64, JobHandle>,
+    subs: &mut Subscriptions,
+) -> Result<u64> {
+    let mut fs_receiver = Some(fs_receiver);
     loop {
         enum Event {
-            Msg(RawMsg),
+            Msg(RawMessage),
             Task(Task),
             Fs(Vec<FileEvent>),
-            ReceiverDead,
             FsWatcherDead,
         }
         let event = select! {
-            recv(io.receiver(), msg) => match msg {
+            recv(msg_receiver, msg) => match msg {
                 Some(msg) => Event::Msg(msg),
-                None => Event::ReceiverDead,
+                None => bail!("client exited without shutdown"),
             },
             recv(task_receiver, task) => Event::Task(task.unwrap()),
-            recv(fs_events_receiver, events) => match events {
+            recv(fs_receiver, events) => match events {
                 Some(events) => Event::Fs(events),
                 None => Event::FsWatcherDead,
             }
         };
         let mut state_changed = false;
         match event {
-            Event::ReceiverDead => {
-                io.cleanup_receiver()?;
-                unreachable!();
-            }
             Event::FsWatcherDead => {
-                fs_events_receiver = None;
+                fs_receiver = None;
             }
             Event::Task(task) => {
                 match task {
@@ -66,12 +104,10 @@ pub(super) fn main_loop(
                         if let Some(handle) = pending_requests.remove(&response.id) {
                             assert!(handle.has_completed());
                         }
-                        io.send(RawMsg::Response(response))
+                        msg_sender.send(RawMessage::Response(response))
                     }
                     Task::Notify(n) =>
-                        io.send(RawMsg::Notification(n)),
-                    Task::Die(error) =>
-                        return Err(error),
+                        msg_sender.send(RawMessage::Notification(n)),
                 }
                 continue;
             }
@@ -82,16 +118,29 @@ pub(super) fn main_loop(
             }
             Event::Msg(msg) => {
                 match msg {
-                    RawMsg::Request(req) => {
-                        if !on_request(io, &mut state, &mut pending_requests, pool, &task_sender, req)? {
-                            return Ok(());
+                    RawMessage::Request(req) => {
+                        let req = match req.cast::<req::Shutdown>() {
+                            Ok((id, _params)) => return Ok(id),
+                            Err(req) => req,
+                        };
+                        match on_request(state, pending_requests, pool, &task_sender, req)? {
+                            None => (),
+                            Some(req) => {
+                                error!("unknown request: {:?}", req);
+                                let resp = RawResponse::err(
+                                    req.id,
+                                    ErrorCode::MethodNotFound as i32,
+                                    "unknown request".to_string(),
+                                );
+                                msg_sender.send(RawMessage::Response(resp))
+                            }
                         }
                     }
-                    RawMsg::Notification(not) => {
-                        on_notification(io, &mut state, &mut pending_requests, &mut subs, not)?;
+                    RawMessage::Notification(not) => {
+                        on_notification(msg_sender, state, pending_requests, subs, not)?;
                         state_changed = true;
                     }
-                    RawMsg::Response(resp) => {
+                    RawMessage::Response(resp) => {
                         error!("unexpected response: {:?}", resp)
                     }
                 }
@@ -110,13 +159,12 @@ pub(super) fn main_loop(
 }
 
 fn on_request(
-    io: &mut Io,
     world: &mut ServerWorldState,
     pending_requests: &mut HashMap<u64, JobHandle>,
     pool: &ThreadPool,
     sender: &Sender<Task>,
     req: RawRequest,
-) -> Result<bool> {
+) -> Result<Option<RawRequest>> {
     let mut pool_dispatcher = PoolDispatcher {
         req: Some(req),
         res: None,
@@ -141,81 +189,73 @@ fn on_request(
         Ok((id, handle)) => {
             let inserted = pending_requests.insert(id, handle).is_none();
             assert!(inserted, "duplicate request: {}", id);
+            Ok(None)
         },
-        Err(req) => {
-            let req = dispatch::handle_request::<req::Shutdown, _>(req, |(), resp| {
-                let resp = resp.into_response(Ok(()))?;
-                io.send(RawMsg::Response(resp));
-                Ok(())
-            })?;
-            match req {
-                Ok(_id) => {
-                    info!("lifecycle: initiating shutdown");
-                    return Ok(false);
-                }
-                Err(req) => {
-                    error!("unknown method: {:?}", req);
-                    io.send(RawMsg::Response(dispatch::unknown_method(req.id)?));
-                }
-            }
-        }
+        Err(req) => Ok(Some(req)),
     }
-    Ok(true)
 }
 
 fn on_notification(
-    io: &mut Io,
+    msg_sender: &mut Sender<RawMessage>,
     state: &mut ServerWorldState,
     pending_requests: &mut HashMap<u64, JobHandle>,
     subs: &mut Subscriptions,
     not: RawNotification,
 ) -> Result<()> {
-    let mut not = Some(not);
-    dispatch::handle_notification::<req::Cancel, _>(&mut not, |params| {
-        let id = match params.id {
-            NumberOrString::Number(id) => id,
-            NumberOrString::String(id) => {
-                panic!("string id's not supported: {:?}", id);
+    let not = match not.cast::<req::Cancel>() {
+        Ok(params) => {
+            let id = match params.id {
+                NumberOrString::Number(id) => id,
+                NumberOrString::String(id) => {
+                    panic!("string id's not supported: {:?}", id);
+                }
+            };
+            if let Some(handle) = pending_requests.remove(&id) {
+                handle.cancel();
             }
-        };
-        if let Some(handle) = pending_requests.remove(&id) {
-            handle.cancel();
+            return Ok(())
         }
-        Ok(())
-    })?;
-    dispatch::handle_notification::<req::DidOpenTextDocument, _>(&mut not, |params| {
-        let uri = params.text_document.uri;
-        let path = uri.to_file_path()
-            .map_err(|()| format_err!("invalid uri: {}", uri))?;
-        let file_id = state.add_mem_file(path, params.text_document.text);
-        subs.add_sub(file_id);
-        Ok(())
-    })?;
-    dispatch::handle_notification::<req::DidChangeTextDocument, _>(&mut not, |mut params| {
-        let uri = params.text_document.uri;
-        let path = uri.to_file_path()
-            .map_err(|()| format_err!("invalid uri: {}", uri))?;
-        let text = params.content_changes.pop()
-            .ok_or_else(|| format_err!("empty changes"))?
-            .text;
-        state.change_mem_file(path.as_path(), text)?;
-        Ok(())
-    })?;
-    dispatch::handle_notification::<req::DidCloseTextDocument, _>(&mut not, |params| {
-        let uri = params.text_document.uri;
-        let path = uri.to_file_path()
-            .map_err(|()| format_err!("invalid uri: {}", uri))?;
-        let file_id = state.remove_mem_file(path.as_path())?;
-        subs.remove_sub(file_id);
-        let not = req::PublishDiagnosticsParams { uri, diagnostics: Vec::new() };
-        let not = dispatch::send_notification::<req::PublishDiagnostics>(not);
-        io.send(RawMsg::Notification(not));
-        Ok(())
-    })?;
-
-    if let Some(not) = not {
-        error!("unhandled notification: {:?}", not);
-    }
+        Err(not) => not,
+    };
+    let not = match not.cast::<req::DidOpenTextDocument>() {
+        Ok(params) => {
+            let uri = params.text_document.uri;
+            let path = uri.to_file_path()
+                .map_err(|()| format_err!("invalid uri: {}", uri))?;
+            let file_id = state.add_mem_file(path, params.text_document.text);
+            subs.add_sub(file_id);
+            return Ok(())
+        }
+        Err(not) => not,
+    };
+    let not = match not.cast::<req::DidChangeTextDocument>() {
+        Ok(mut params) => {
+            let uri = params.text_document.uri;
+            let path = uri.to_file_path()
+                .map_err(|()| format_err!("invalid uri: {}", uri))?;
+            let text = params.content_changes.pop()
+                .ok_or_else(|| format_err!("empty changes"))?
+                .text;
+            state.change_mem_file(path.as_path(), text)?;
+            return Ok(())
+        }
+        Err(not) => not,
+    };
+    let not = match not.cast::<req::DidCloseTextDocument>() {
+        Ok(params) => {
+            let uri = params.text_document.uri;
+            let path = uri.to_file_path()
+                .map_err(|()| format_err!("invalid uri: {}", uri))?;
+            let file_id = state.remove_mem_file(path.as_path())?;
+            subs.remove_sub(file_id);
+            let params = req::PublishDiagnosticsParams { uri, diagnostics: Vec::new() };
+            let not = RawNotification::new::<req::PublishDiagnostics>(params);
+            msg_sender.send(RawMessage::Notification(not));
+            return Ok(())
+        }
+        Err(not) => not,
+    };
+    error!("unhandled notification: {:?}", not);
     Ok(())
 }
 
@@ -228,10 +268,14 @@ struct PoolDispatcher<'a> {
 }
 
 impl<'a> PoolDispatcher<'a> {
-    fn on<'b, R: req::ClientRequest>(
+    fn on<'b, R>(
         &'b mut self,
         f: fn(ServerWorld, R::Params, JobToken) -> Result<R::Result>
-    ) -> Result<&'b mut Self> {
+    ) -> Result<&'b mut Self>
+    where R: req::Request,
+          R::Params: DeserializeOwned + Send + 'static,
+          R::Result: Serialize + 'static,
+    {
         let req = match self.req.take() {
             None => return Ok(self),
             Some(req) => req,
@@ -239,23 +283,24 @@ impl<'a> PoolDispatcher<'a> {
         let world = self.world;
         let sender = self.sender;
         let pool = self.pool;
-        let (handle, token) = JobHandle::new();
-        let req = dispatch::handle_request::<R, _>(req, |params, resp| {
-            let world = world.snapshot();
-            let sender = sender.clone();
-            pool.execute(move || {
-                let res = f(world, params, token);
-                let task = match resp.into_response(res) {
-                    Ok(resp) => Task::Respond(resp),
-                    Err(e) => Task::Die(e),
-                };
-                sender.send(task);
-            });
-            Ok(())
-        })?;
-        match req {
-            Ok(id) => self.res = Some((id, handle)),
-            Err(req) => self.req = Some(req),
+        match req.cast::<R>() {
+            Ok((id, params)) => {
+                let (handle, token) = JobHandle::new();
+                let world = world.snapshot();
+                let sender = sender.clone();
+                pool.execute(move || {
+                    let resp = match f(world, params, token) {
+                        Ok(resp) => RawResponse::ok(id, resp),
+                        Err(e) => RawResponse::err(id, ErrorCode::InternalError as i32, e.to_string()),
+                    };
+                    let task = Task::Respond(resp);
+                    sender.send(task);
+                });
+                self.res = Some((id, handle));
+            }
+            Err(req) => {
+                self.req = Some(req)
+            }
         }
         Ok(self)
     }
@@ -282,7 +327,7 @@ fn update_file_notifications_on_threadpool(
                     error!("failed to compute diagnostics: {:?}", e)
                 }
                 Ok(params) => {
-                    let not = dispatch::send_notification::<req::PublishDiagnostics>(params);
+                    let not = RawNotification::new::<req::PublishDiagnostics>(params);
                     sender.send(Task::Notify(not));
                 }
             }
@@ -291,7 +336,7 @@ fn update_file_notifications_on_threadpool(
                     error!("failed to compute decorations: {:?}", e)
                 }
                 Ok(params) => {
-                    let not = dispatch::send_notification::<req::PublishDecorations>(params);
+                    let not = RawNotification::new::<req::PublishDecorations>(params);
                     sender.send(Task::Notify(not))
                 }
             }