Merge remote-tracking branch 'origin/master' into line_ending_detection

3 files changed, 214 insertions, 58 deletions

diff --git a/helix-lsp/src/client.rs b/helix-lsp/src/client.rs
index 9ca708a7..101d2f9b 100644
--- a/helix-lsp/src/client.rs
+++ b/helix-lsp/src/client.rs
@@ -18,6 +18,7 @@ use tokio::{
 
 #[derive(Debug)]
 pub struct Client {
+    id: usize,
     _process: Child,
     server_tx: UnboundedSender<Payload>,
     request_counter: AtomicU64,
@@ -26,7 +27,11 @@ pub struct Client {
 }
 
 impl Client {
-    pub fn start(cmd: &str, args: &[String]) -> Result<(Self, UnboundedReceiver<Call>)> {
+    pub fn start(
+        cmd: &str,
+        args: &[String],
+        id: usize,
+    ) -> Result<(Self, UnboundedReceiver<(usize, Call)>)> {
         let process = Command::new(cmd)
             .args(args)
             .stdin(Stdio::piped())
@@ -43,9 +48,10 @@ impl Client {
         let reader = BufReader::new(process.stdout.take().expect("Failed to open stdout"));
         let stderr = BufReader::new(process.stderr.take().expect("Failed to open stderr"));
 
-        let (server_rx, server_tx) = Transport::start(reader, writer, stderr);
+        let (server_rx, server_tx) = Transport::start(reader, writer, stderr, id);
 
         let client = Self {
+            id,
             _process: process,
             server_tx,
             request_counter: AtomicU64::new(0),
@@ -59,6 +65,10 @@ impl Client {
         Ok((client, server_rx))
     }
 
+    pub fn id(&self) -> usize {
+        self.id
+    }
+
     fn next_request_id(&self) -> jsonrpc::Id {
         let id = self.request_counter.fetch_add(1, Ordering::Relaxed);
         jsonrpc::Id::Num(id)
@@ -165,31 +175,35 @@ impl Client {
     }
 
     /// Reply to a language server RPC call.
-    pub async fn reply(
+    pub fn reply(
         &self,
         id: jsonrpc::Id,
         result: core::result::Result<Value, jsonrpc::Error>,
-    ) -> Result<()> {
+    ) -> impl Future<Output = Result<()>> {
         use jsonrpc::{Failure, Output, Success, Version};
 
-        let output = match result {
-            Ok(result) => Output::Success(Success {
-                jsonrpc: Some(Version::V2),
-                id,
-                result,
-            }),
-            Err(error) => Output::Failure(Failure {
-                jsonrpc: Some(Version::V2),
-                id,
-                error,
-            }),
-        };
+        let server_tx = self.server_tx.clone();
 
-        self.server_tx
-            .send(Payload::Response(output))
-            .map_err(|e| Error::Other(e.into()))?;
+        async move {
+            let output = match result {
+                Ok(result) => Output::Success(Success {
+                    jsonrpc: Some(Version::V2),
+                    id,
+                    result,
+                }),
+                Err(error) => Output::Failure(Failure {
+                    jsonrpc: Some(Version::V2),
+                    id,
+                    error,
+                }),
+            };
 
-        Ok(())
+            server_tx
+                .send(Payload::Response(output))
+                .map_err(|e| Error::Other(e.into()))?;
+
+            Ok(())
+        }
     }
 
     // -------------------------------------------------------------------------------------------
@@ -229,8 +243,7 @@ impl Client {
                     ..Default::default()
                 }),
                 window: Some(lsp::WindowClientCapabilities {
-                    // TODO: temporarily disabled until we implement handling for window/workDoneProgress/create
-                    // work_done_progress: Some(true),
+                    work_done_progress: Some(true),
                     ..Default::default()
                 }),
                 ..Default::default()
@@ -259,6 +272,21 @@ impl Client {
         self.notify::<lsp::notification::Exit>(())
     }
 
+    /// Tries to shut down the language server but returns
+    /// early if server responds with an error.
+    pub async fn shutdown_and_exit(&self) -> Result<()> {
+        self.shutdown().await?;
+        self.exit().await
+    }
+
+    /// Forcefully shuts down the language server ignoring any errors.
+    pub async fn force_shutdown(&self) -> Result<()> {
+        if let Err(e) = self.shutdown().await {
+            log::warn!("language server failed to terminate gracefully - {}", e);
+        }
+        self.exit().await
+    }
+
     // -------------------------------------------------------------------------------------------
     // Text document
     // -------------------------------------------------------------------------------------------
@@ -465,6 +493,7 @@ impl Client {
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
         // ) -> Result<Vec<lsp::CompletionItem>> {
         let params = lsp::CompletionParams {
@@ -473,9 +502,7 @@ impl Client {
                 position,
             },
             // TODO: support these tokens by async receiving and updating the choice list
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
             partial_result_params: lsp::PartialResultParams {
                 partial_result_token: None,
             },
@@ -490,15 +517,14 @@ impl Client {
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
         let params = lsp::SignatureHelpParams {
             text_document_position_params: lsp::TextDocumentPositionParams {
                 text_document,
                 position,
             },
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
             context: None,
             // lsp::SignatureHelpContext
         };
@@ -510,15 +536,14 @@ impl Client {
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
         let params = lsp::HoverParams {
             text_document_position_params: lsp::TextDocumentPositionParams {
                 text_document,
                 position,
             },
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
             // lsp::SignatureHelpContext
         };
 
@@ -531,6 +556,7 @@ impl Client {
         &self,
         text_document: lsp::TextDocumentIdentifier,
         options: lsp::FormattingOptions,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> anyhow::Result<Vec<lsp::TextEdit>> {
         let capabilities = self.capabilities.as_ref().unwrap();
 
@@ -545,9 +571,7 @@ impl Client {
         let params = lsp::DocumentFormattingParams {
             text_document,
             options,
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
         };
 
         let response = self.request::<lsp::request::Formatting>(params).await?;
@@ -560,6 +584,7 @@ impl Client {
         text_document: lsp::TextDocumentIdentifier,
         range: lsp::Range,
         options: lsp::FormattingOptions,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> anyhow::Result<Vec<lsp::TextEdit>> {
         let capabilities = self.capabilities.as_ref().unwrap();
 
@@ -575,9 +600,7 @@ impl Client {
             text_document,
             range,
             options,
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
         };
 
         let response = self
@@ -596,15 +619,14 @@ impl Client {
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
         let params = lsp::GotoDefinitionParams {
             text_document_position_params: lsp::TextDocumentPositionParams {
                 text_document,
                 position,
             },
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
             partial_result_params: lsp::PartialResultParams {
                 partial_result_token: None,
             },
@@ -617,30 +639,42 @@ impl Client {
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
-        self.goto_request::<lsp::request::GotoDefinition>(text_document, position)
+        self.goto_request::<lsp::request::GotoDefinition>(text_document, position, work_done_token)
     }
 
     pub fn goto_type_definition(
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
-        self.goto_request::<lsp::request::GotoTypeDefinition>(text_document, position)
+        self.goto_request::<lsp::request::GotoTypeDefinition>(
+            text_document,
+            position,
+            work_done_token,
+        )
     }
 
     pub fn goto_implementation(
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
-        self.goto_request::<lsp::request::GotoImplementation>(text_document, position)
+        self.goto_request::<lsp::request::GotoImplementation>(
+            text_document,
+            position,
+            work_done_token,
+        )
     }
 
     pub fn goto_reference(
         &self,
         text_document: lsp::TextDocumentIdentifier,
         position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
     ) -> impl Future<Output = Result<Value>> {
         let params = lsp::ReferenceParams {
             text_document_position: lsp::TextDocumentPositionParams {
@@ -650,9 +684,7 @@ impl Client {
             context: lsp::ReferenceContext {
                 include_declaration: true,
             },
-            work_done_progress_params: lsp::WorkDoneProgressParams {
-                work_done_token: None,
-            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
             partial_result_params: lsp::PartialResultParams {
                 partial_result_token: None,
             },
diff --git a/helix-lsp/src/lib.rs b/helix-lsp/src/lib.rs
index cff62492..49d5527f 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -13,7 +13,10 @@ use helix_core::syntax::LanguageConfiguration;
 
 use std::{
     collections::{hash_map::Entry, HashMap},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
 };
 
 use serde::{Deserialize, Serialize};
@@ -182,6 +185,30 @@ pub mod util {
 }
 
 #[derive(Debug, PartialEq, Clone)]
+pub enum MethodCall {
+    WorkDoneProgressCreate(lsp::WorkDoneProgressCreateParams),
+}
+
+impl MethodCall {
+    pub fn parse(method: &str, params: jsonrpc::Params) -> Option<MethodCall> {
+        use lsp::request::Request;
+        let request = match method {
+            lsp::request::WorkDoneProgressCreate::METHOD => {
+                let params: lsp::WorkDoneProgressCreateParams = params
+                    .parse()
+                    .expect("Failed to parse WorkDoneCreate params");
+                Self::WorkDoneProgressCreate(params)
+            }
+            _ => {
+                log::warn!("unhandled lsp request: {}", method);
+                return None;
+            }
+        };
+        Some(request)
+    }
+}
+
+#[derive(Debug, PartialEq, Clone)]
 pub enum Notification {
     PublishDiagnostics(lsp::PublishDiagnosticsParams),
     ShowMessage(lsp::ShowMessageParams),
@@ -230,9 +257,10 @@ impl Notification {
 
 #[derive(Debug)]
 pub struct Registry {
-    inner: HashMap<LanguageId, Arc<Client>>,
+    inner: HashMap<LanguageId, (usize, Arc<Client>)>,
 
-    pub incoming: SelectAll<UnboundedReceiverStream<Call>>,
+    counter: AtomicUsize,
+    pub incoming: SelectAll<UnboundedReceiverStream<(usize, Call)>>,
 }
 
 impl Default for Registry {
@@ -245,10 +273,18 @@ impl Registry {
     pub fn new() -> Self {
         Self {
             inner: HashMap::new(),
+            counter: AtomicUsize::new(0),
             incoming: SelectAll::new(),
         }
     }
 
+    pub fn get_by_id(&mut self, id: usize) -> Option<&Client> {
+        self.inner
+            .values()
+            .find(|(client_id, _)| client_id == &id)
+            .map(|(_, client)| client.as_ref())
+    }
+
     pub fn get(&mut self, language_config: &LanguageConfiguration) -> Result<Arc<Client>> {
         if let Some(config) = &language_config.language_server {
             // avoid borrow issues
@@ -256,16 +292,17 @@ impl Registry {
             let s_incoming = &mut self.incoming;
 
             match inner.entry(language_config.scope.clone()) {
-                Entry::Occupied(language_server) => Ok(language_server.get().clone()),
+                Entry::Occupied(entry) => Ok(entry.get().1.clone()),
                 Entry::Vacant(entry) => {
                     // initialize a new client
-                    let (mut client, incoming) = Client::start(&config.command, &config.args)?;
+                    let id = self.counter.fetch_add(1, Ordering::Relaxed);
+                    let (mut client, incoming) = Client::start(&config.command, &config.args, id)?;
                     // TODO: run this async without blocking
                     futures_executor::block_on(client.initialize())?;
                     s_incoming.push(UnboundedReceiverStream::new(incoming));
                     let client = Arc::new(client);
 
-                    entry.insert(client.clone());
+                    entry.insert((id, client.clone()));
                     Ok(client)
                 }
             }
@@ -273,6 +310,84 @@ impl Registry {
             Err(Error::LspNotDefined)
         }
     }
+
+    pub fn iter_clients(&self) -> impl Iterator<Item = &Arc<Client>> {
+        self.inner.values().map(|(_, client)| client)
+    }
+}
+
+#[derive(Debug)]
+pub enum ProgressStatus {
+    Created,
+    Started(lsp::WorkDoneProgress),
+}
+
+impl ProgressStatus {
+    pub fn progress(&self) -> Option<&lsp::WorkDoneProgress> {
+        match &self {
+            ProgressStatus::Created => None,
+            ProgressStatus::Started(progress) => Some(&progress),
+        }
+    }
+}
+
+#[derive(Default, Debug)]
+/// Acts as a container for progress reported by language servers. Each server
+/// has a unique id assigned at creation through [`Registry`]. This id is then used
+/// to store the progress in this map.
+pub struct LspProgressMap(HashMap<usize, HashMap<lsp::ProgressToken, ProgressStatus>>);
+
+impl LspProgressMap {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Returns a map of all tokens coresponding to the lanaguage server with `id`.
+    pub fn progress_map(&self, id: usize) -> Option<&HashMap<lsp::ProgressToken, ProgressStatus>> {
+        self.0.get(&id)
+    }
+
+    /// Returns last progress status for a given server with `id` and `token`.
+    pub fn progress(&self, id: usize, token: &lsp::ProgressToken) -> Option<&ProgressStatus> {
+        self.0.get(&id).and_then(|values| values.get(token))
+    }
+
+    /// Checks if progress `token` for server with `id` is created.
+    pub fn is_created(&mut self, id: usize, token: &lsp::ProgressToken) -> bool {
+        self.0
+            .get(&id)
+            .map(|values| values.get(token).is_some())
+            .unwrap_or_default()
+    }
+
+    pub fn create(&mut self, id: usize, token: lsp::ProgressToken) {
+        self.0
+            .entry(id)
+            .or_default()
+            .insert(token, ProgressStatus::Created);
+    }
+
+    /// Ends the progress by removing the `token` from server with `id`, if removed returns the value.
+    pub fn end_progress(
+        &mut self,
+        id: usize,
+        token: &lsp::ProgressToken,
+    ) -> Option<ProgressStatus> {
+        self.0.get_mut(&id).and_then(|vals| vals.remove(token))
+    }
+
+    /// Updates the progess of `token` for server with `id` to `status`, returns the value replaced or `None`.
+    pub fn update(
+        &mut self,
+        id: usize,
+        token: lsp::ProgressToken,
+        status: lsp::WorkDoneProgress,
+    ) -> Option<ProgressStatus> {
+        self.0
+            .entry(id)
+            .or_default()
+            .insert(token, ProgressStatus::Started(status))
+    }
 }
 
 // REGISTRY = HashMap<LanguageId, Lazy/OnceCell<Arc<RwLock<Client>>>
diff --git a/helix-lsp/src/transport.rs b/helix-lsp/src/transport.rs
index e8068323..df55bbf6 100644
--- a/helix-lsp/src/transport.rs
+++ b/helix-lsp/src/transport.rs
@@ -1,4 +1,5 @@
 use crate::Result;
+use anyhow::Context;
 use jsonrpc_core as jsonrpc;
 use log::{error, info};
 use serde::{Deserialize, Serialize};
@@ -33,7 +34,8 @@ enum ServerMessage {
 
 #[derive(Debug)]
 pub struct Transport {
-    client_tx: UnboundedSender<jsonrpc::Call>,
+    id: usize,
+    client_tx: UnboundedSender<(usize, jsonrpc::Call)>,
     client_rx: UnboundedReceiver<Payload>,
 
     pending_requests: HashMap<jsonrpc::Id, Sender<Result<Value>>>,
@@ -48,11 +50,16 @@ impl Transport {
         server_stdout: BufReader<ChildStdout>,
         server_stdin: BufWriter<ChildStdin>,
         server_stderr: BufReader<ChildStderr>,
-    ) -> (UnboundedReceiver<jsonrpc::Call>, UnboundedSender<Payload>) {
+        id: usize,
+    ) -> (
+        UnboundedReceiver<(usize, jsonrpc::Call)>,
+        UnboundedSender<Payload>,
+    ) {
         let (client_tx, rx) = unbounded_channel();
         let (tx, client_rx) = unbounded_channel();
 
         let transport = Self {
+            id,
             server_stdout,
             server_stdin,
             server_stderr,
@@ -84,7 +91,7 @@ impl Transport {
 
             match (parts.next(), parts.next(), parts.next()) {
                 (Some("Content-Length"), Some(value), None) => {
-                    content_length = Some(value.parse().unwrap());
+                    content_length = Some(value.parse().context("invalid content length")?);
                 }
                 (Some(_), Some(_), None) => {}
                 _ => {
@@ -97,12 +104,12 @@ impl Transport {
             }
         }
 
-        let content_length = content_length.unwrap();
+        let content_length = content_length.context("missing content length")?;
 
         //TODO: reuse vector
         let mut content = vec![0; content_length];
         reader.read_exact(&mut content).await?;
-        let msg = String::from_utf8(content).unwrap();
+        let msg = String::from_utf8(content).context("invalid utf8 from server")?;
 
         info!("<- {}", msg);
 
@@ -156,7 +163,9 @@ impl Transport {
         match msg {
             ServerMessage::Output(output) => self.process_request_response(output).await?,
             ServerMessage::Call(call) => {
-                self.client_tx.send(call).unwrap();
+                self.client_tx
+                    .send((self.id, call))
+                    .context("failed to send a message to server")?;
                 // let notification = Notification::parse(&method, params);
             }
         };