From f57bb8034c426ec99c13006a02d8ed18a2c568b9 Mon Sep 17 00:00:00 2001
From: Viet Dinh <54ckb0y789@gmail.com>
Date: Sat, 11 Nov 2023 23:12:23 -0500
Subject: [PATCH] feat: consult all LSPs for goto-definitions

---
 helix-term/src/commands/lsp.rs | 64 +++++++++++++++++++++++++++++-----
 1 file changed, 56 insertions(+), 8 deletions(-)

diff --git a/helix-term/src/commands/lsp.rs b/helix-term/src/commands/lsp.rs
index 47d1f8d1ab70c..9e7d05a7f8c74 100644
--- a/helix-term/src/commands/lsp.rs
+++ b/helix-term/src/commands/lsp.rs
@@ -1,7 +1,7 @@
 use futures_util::{
     future::{join_all, BoxFuture},
     stream::FuturesUnordered,
-    FutureExt,
+    FutureExt, TryFutureExt,
 };
 use helix_lsp::{
     block_on,
@@ -47,6 +47,7 @@ use std::{
     future::Future,
     path::PathBuf,
     sync::Arc,
+    task::Poll,
 };
 
 /// Gets the first language server that is attached to a document which supports a specific feature.
@@ -1064,16 +1065,63 @@ where
 {
     let (view, doc) = current!(cx.editor);
 
-    let language_server = language_server_with_feature!(cx.editor, doc, feature);
-    let offset_encoding = language_server.offset_encoding();
-    let pos = doc.position(view.id, offset_encoding);
-    let future = request_provider(language_server, pos, doc.identifier()).unwrap();
+    let language_servers = doc.language_servers_with_feature(feature);
+    let mut futures = vec![];
+    for server in language_servers {
+        let offset_encoding = server.offset_encoding();
+        let pos = doc.position(view.id, offset_encoding);
+        let fut = request_provider(server, pos, doc.identifier()).unwrap();
+        let fut = fut.map(move |res| {
+            res.map(|mut resp| {
+                if !resp.is_null() {
+                    resp = serde_json::json!([offset_encoding as usize, resp]);
+                }
+                resp
+            })
+        });
+        futures.push(Box::pin(fut));
+    }
+
+    let future =
+        futures_util::future::poll_fn(move |ctx| {
+            if futures.is_empty() {
+                return Poll::Ready(Ok(Value::Null));
+            }
+            loop {
+                // .. select the next ready future
+                let Some((idx, res)) = futures.iter_mut().enumerate().find_map(|(idx, fut)| {
+                    match fut.try_poll_unpin(ctx) {
+                        Poll::Pending => None,
+                        Poll::Ready(res) => Some((idx, res)),
+                    }
+                }) else {
+                    return Poll::Pending;
+                };
+
+                _ = futures.swap_remove(idx);
+                match res {
+                    // .. that succeeds AND is not null
+                    Ok(res) if !res.is_null() => return Poll::Ready(Ok(res)),
+                    // .. or go back to polling immediately
+                    _ => continue,
+                }
+            }
+        });
 
     cx.callback(
         future,
-        move |editor, compositor, response: Option<lsp::GotoDefinitionResponse>| {
-            let items = to_locations(response);
-            goto_impl(editor, compositor, items, offset_encoding);
+        move |editor, compositor, response: Option<(usize, lsp::GotoDefinitionResponse)>| {
+            let encoding = response
+                .as_ref()
+                .map(|(encoding, _)| match encoding {
+                    0 => OffsetEncoding::Utf8,
+                    1 => OffsetEncoding::Utf32,
+                    2 => OffsetEncoding::Utf16,
+                    _ => unreachable!(),
+                })
+                .unwrap_or_default();
+            let items = to_locations(response.map(|(_, resp)| resp));
+            goto_impl(editor, compositor, items, encoding);
         },
     );
 }