569 lines · 22.6 KB · Rust
RawDownload
use std::ops::Range;use std::sync::Arc; use anyhow::{Context as _, Result};use collections::{HashMap, HashSet};use futures::{    FutureExt as _,    future::{Shared, join_all},};use gpui::{AppContext as _, AsyncApp, Context, Entity, Task};use language::{Anchor, Buffer};use lsp::LanguageServerId;use rpc::{TypedEnvelope, proto};use settings::Settings as _;use std::time::Duration;use text::OffsetRangeExt as _;use util::ResultExt as _; use crate::{    CodeAction, LspAction, LspStore, LspStoreEvent, Project,    lsp_command::{GetCodeLens, LspCommand as _},    lsp_store::{        RunningFetch, missing_servers_to_query, next_lsp_fetch_id, upstream_lsp_query_server_filter,    },    project_settings::ProjectSettings,}; /// Opaque per-action identifier issued by [`LspStore`] at fetch time.////// LSP `CodeLens.data` is the server's private payload for resolve/// round-trips, so we can't use it (or anything derived from it) to/// disambiguate sibling lenses that share the same buffer `range`/// (TypeScript's references + implementations is the canonical case)./// We tag every cached action with this id and require it back on resolve/// so each lens routes to its own request and slot.////// Ids are issued in fetch order; sorting by id reproduces server-emit/// order, which is how callers recover a stable render order without/// paying for an ordered map.#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)]pub struct CodeLensActionId(u64); pub type CodeLensActions = HashMap<CodeLensActionId, CodeAction>; pub(super) type CodeLensTask =    Shared<Task<std::result::Result<Option<CodeLensActions>, Arc<anyhow::Error>>>>; pub type CodeLensResolveTask = Shared<Task<Option<(CodeLensActionId, CodeAction)>>>; #[derive(Debug, Default)]pub(super) struct CodeLensData {    pub(super) lens: HashMap<LanguageServerId, CodeLensActions>,    fetched_servers: HashSet<LanguageServerId>,    pub(super) next_id: u64,    pub(super) update: Option<RunningFetch<CodeLensTask>>,    pub(super) resolving: HashMap<(LanguageServerId, CodeLensActionId), CodeLensResolveTask>,} impl CodeLensData {    pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {        self.lens.remove(&server_id);        self.fetched_servers.remove(&server_id);        self.resolving.retain(|(s, _), _| *s != server_id);        RunningFetch::discard_if_queried(&mut self.update, server_id);    }     fn evict(&mut self, for_server: Option<LanguageServerId>) {        match for_server {            Some(server_id) => self.remove_server_data(server_id),            None => {                self.lens.clear();                self.fetched_servers.clear();                self.resolving.clear();            }        }        self.update = None;    }} fn flatten_cache(lens: &HashMap<LanguageServerId, CodeLensActions>) -> CodeLensActions {    let mut out = CodeLensActions::default();    out.reserve(lens.values().map(|per_server| per_server.len()).sum());    for per_server in lens.values() {        for (id, action) in per_server {            out.insert(*id, action.clone());        }    }    out} impl LspStore {    pub(super) fn refresh_code_lens(        &mut self,        for_server: Option<LanguageServerId>,        cx: &mut Context<Self>,    ) {        for lsp_data in self.lsp_data.values_mut() {            if let Some(code_lens) = &mut lsp_data.code_lens {                code_lens.evict(for_server);            }        }         cx.emit(LspStoreEvent::RefreshCodeLens {            server_id: for_server,        });        if let Some((downstream_client, project_id)) = self.downstream_client.as_ref() {            downstream_client                .send(proto::RefreshCodeLens {                    project_id: *project_id,                    server_id: for_server.map(|server_id| server_id.to_proto()),                })                .context("sending refresh code lens downstream")                .log_err();        }    }     /// Fetches all code lenses for the buffer, each tagged with the    /// [`CodeLensActionId`] that callers must pass back to    /// [`Self::resolve_code_lens`]. Resolution is the caller's job.    pub fn code_lens_actions(        &mut self,        buffer: &Entity<Buffer>,        cx: &mut Context<Self>,    ) -> Task<Result<Option<CodeLensActions>>> {        let buffer_id = buffer.read(cx).remote_id();        let fetch_task = self.fetch_code_lenses(buffer, cx);         cx.spawn(async move |lsp_store, cx| {            fetch_task                .await                .map_err(|e| anyhow::anyhow!("code lens fetch failed: {e:#}"))?;             let actions = lsp_store.read_with(cx, |lsp_store, _| {                lsp_store                    .lsp_data                    .get(&buffer_id)                    .and_then(|data| data.code_lens.as_ref())                    .map(|code_lens| flatten_cache(&code_lens.lens))            })?;            Ok(actions)        })    }     fn fetch_code_lenses(        &mut self,        buffer: &Entity<Buffer>,        cx: &mut Context<Self>,    ) -> CodeLensTask {        let version_queried_for = buffer.read(cx).version();        let buffer_id = buffer.read(cx).remote_id();        let current_servers = self.relevant_server_ids_for_capability_check(buffer, cx);         let mut servers_to_query = None;        if let Some(lsp_data) = self.current_lsp_data(buffer_id) {            if let Some(cached_lens) = &mut lsp_data.code_lens {                if !version_queried_for.changed_since(&lsp_data.buffer_version) {                    match missing_servers_to_query(                        &mut cached_lens.lens,                        &mut cached_lens.fetched_servers,                        &current_servers,                    ) {                        Some(missing_servers) => servers_to_query = Some(missing_servers),                        None => {                            return Task::ready(Ok(Some(flatten_cache(&cached_lens.lens))))                                .shared();                        }                    }                }                if let Some(running) = cached_lens.update.as_ref()                    && !version_queried_for.changed_since(&running.version)                    && servers_to_query                        .as_ref()                        .is_none_or(|missing| missing.is_subset(&running.servers))                {                    return running.task.clone();                }            }        }         let lens_lsp_data = self            .latest_lsp_data(buffer, cx)            .code_lens            .get_or_insert_default();        let fetch_id = next_lsp_fetch_id();        let queried_servers = servers_to_query            .clone()            .unwrap_or_else(|| current_servers.clone());        let buffer = buffer.clone();        let query_version_queried_for = version_queried_for.clone();        let new_task = cx            .spawn({                let queried_servers = queried_servers.clone();                async move |lsp_store, cx| {                    cx.background_executor()                        .timer(Duration::from_millis(30))                        .await;                    let fetched_lens = lsp_store                        .update(cx, |lsp_store, cx| {                            lsp_store.fetch_code_lens_for_buffer(&buffer, servers_to_query, cx)                        })                        .map_err(Arc::new)?                        .await                        .context("fetching code lens")                        .map_err(Arc::new);                    let fetched_lens = match fetched_lens {                        Ok(fetched_lens) => fetched_lens,                        Err(e) => {                            lsp_store                                .update(cx, |lsp_store, _| {                                    if let Some(lens_lsp_data) = lsp_store                                        .lsp_data                                        .get_mut(&buffer_id)                                        .and_then(|lsp_data| lsp_data.code_lens.as_mut())                                    {                                        RunningFetch::take_finished(                                            &mut lens_lsp_data.update,                                            fetch_id,                                        );                                    }                                })                                .ok();                            return Err(e);                        }                    };                     lsp_store                        .update(cx, |lsp_store, _| {                            let lsp_data = lsp_store.current_lsp_data(buffer_id)?;                            let code_lens = lsp_data.code_lens.as_mut()?;                            if !RunningFetch::take_finished(&mut code_lens.update, fetch_id) {                                return Some(flatten_cache(&code_lens.lens));                            }                            if let Some(fetched_lens) = fetched_lens {                                let mut tagged: HashMap<LanguageServerId, CodeLensActions> =                                    HashMap::default();                                for (server_id, actions) in fetched_lens {                                    let mut cache = CodeLensActions::default();                                    cache.reserve(actions.len());                                    for action in actions {                                        let id = CodeLensActionId(code_lens.next_id);                                        code_lens.next_id += 1;                                        cache.insert(id, action);                                    }                                    tagged.insert(server_id, cache);                                }                                if lsp_data.buffer_version == query_version_queried_for {                                    code_lens.lens.extend(tagged);                                    code_lens.fetched_servers.extend(queried_servers);                                } else if !lsp_data                                    .buffer_version                                    .changed_since(&query_version_queried_for)                                {                                    lsp_data.buffer_version = query_version_queried_for;                                    code_lens.lens = tagged;                                    code_lens.fetched_servers = queried_servers;                                }                            }                            Some(flatten_cache(&code_lens.lens))                        })                        .map_err(Arc::new)                }            })            .shared();        lens_lsp_data.update = Some(RunningFetch {            id: fetch_id,            version: version_queried_for,            servers: queried_servers,            task: new_task.clone(),        });        new_task    }     fn fetch_code_lens_for_buffer(        &mut self,        buffer: &Entity<Buffer>,        for_servers: Option<HashSet<LanguageServerId>>,        cx: &mut Context<Self>,    ) -> Task<Result<Option<HashMap<LanguageServerId, Vec<CodeAction>>>>> {        if let Some((upstream_client, project_id)) = self.upstream_client() {            let request = GetCodeLens;            if !self.is_capable_for_proto_request(buffer, &request, cx) {                return Task::ready(Ok(None));            }            let request_timeout = ProjectSettings::get_global(cx)                .global_lsp_settings                .get_request_timeout();            let request_task = upstream_client.request_lsp(                project_id,                upstream_lsp_query_server_filter(for_servers.as_ref()),                request_timeout,                cx.background_executor().clone(),                request.to_proto(project_id, buffer.read(cx)),            );            let buffer = buffer.clone();            cx.spawn(async move |weak_lsp_store, cx| {                let Some(lsp_store) = weak_lsp_store.upgrade() else {                    return Ok(None);                };                let Some(responses) = request_task.await? else {                    return Ok(None);                };                 let code_lens_actions = join_all(responses.payload.into_iter().map(|response| {                    let lsp_store = lsp_store.clone();                    let buffer = buffer.clone();                    let cx = cx.clone();                    async move {                        (                            LanguageServerId::from_proto(response.server_id),                            GetCodeLens                                .response_from_proto(response.response, lsp_store, buffer, cx)                                .await,                        )                    }                }))                .await;                 let mut has_errors = false;                let code_lens_actions = code_lens_actions                    .into_iter()                    .filter_map(|(server_id, code_lens)| match code_lens {                        Ok(code_lens) => Some((server_id, code_lens)),                        Err(e) => {                            has_errors = true;                            log::error!("{e:#}");                            None                        }                    })                    .collect::<HashMap<_, _>>();                anyhow::ensure!(                    !has_errors || !code_lens_actions.is_empty(),                    "Failed to fetch code lens"                );                Ok(Some(code_lens_actions))            })        } else {            let code_lens_actions_task = self.request_filtered_lsp_locally(                buffer,                None::<usize>,                GetCodeLens,                for_servers.as_ref(),                cx,            );            cx.background_spawn(async move {                Ok(Some(code_lens_actions_task.await.into_iter().collect()))            })        }    }     /// Resolves a single code lens via `codeLens/resolve`, identified by    /// the [`CodeLensActionId`] returned from [`Self::code_lens_actions`].    /// The returned task is shared and cached on [`CodeLensData::resolving`]    /// keyed by `(server, lens_id)`, so concurrent callers awaiting the    /// same lens only drive a single LSP request.    ///    /// `None` is yielded when the lens cannot be resolved (id no longer    /// cached, server gone, no `resolveProvider`, request failure, etc.).    /// On success, the cached entry is updated in place before the    /// `(id, resolved_action)` pair is returned.    ///    /// All visibility / batching policy lives in the caller. Remote (proto)    /// resolves are forwarded to the host via [`Self::resolve_code_action`].    pub fn resolve_code_lens(        &mut self,        buffer: &Entity<Buffer>,        server_id: LanguageServerId,        lens_id: CodeLensActionId,        cx: &mut Context<Self>,    ) -> CodeLensResolveTask {        let buffer_id = buffer.read(cx).remote_id();         let Some(code_lens) = self            .lsp_data            .get_mut(&buffer_id)            .and_then(|data| data.code_lens.as_mut())        else {            return Task::ready(None).shared();        };        let key = (server_id, lens_id);        if let Some(existing) = code_lens.resolving.get(&key) {            return existing.clone();        }        let Some(cached) = code_lens            .lens            .get(&server_id)            .and_then(|cache| cache.get(&lens_id))        else {            return Task::ready(None).shared();        };        if cached.resolved {            return Task::ready(Some((lens_id, cached.clone()))).shared();        }        let LspAction::CodeLens(lens) = &cached.lsp_action else {            return Task::ready(None).shared();        };        let lens = lens.clone();        let action = cached.clone();         if self.upstream_client().is_some() {            if !self.check_if_capable_for_proto_request(buffer, GetCodeLens::can_resolve_lens, cx) {                return Task::ready(None).shared();            }            let resolve = self.resolve_code_action(buffer, action, cx);            let task = cx                .spawn(async move |lsp_store, cx| {                    let resolved = resolve                        .await                        .context("resolving remote code lens")                        .log_err()?;                    lsp_store                        .update(cx, |lsp_store, _| {                            let code_lens = lsp_store                                .lsp_data                                .get_mut(&buffer_id)                                .and_then(|data| data.code_lens.as_mut())?;                            code_lens.resolving.remove(&key);                            let action = code_lens                                .lens                                .get_mut(&server_id)                                .and_then(|cache| cache.get_mut(&lens_id))?;                            action.resolved = true;                            action.lsp_action = resolved.lsp_action;                            Some((lens_id, action.clone()))                        })                        .ok()                        .flatten()                })                .shared();            if let Some(code_lens) = self                .lsp_data                .get_mut(&buffer_id)                .and_then(|data| data.code_lens.as_mut())            {                code_lens.resolving.insert(key, task.clone());            }            return task;        }         let Some(server) = self.language_server_for_id(server_id) else {            return Task::ready(None).shared();        };        if !GetCodeLens::can_resolve_lens(&server.capabilities()) {            return Task::ready(None).shared();        }        let request_timeout = ProjectSettings::get_global(cx)            .global_lsp_settings            .get_request_timeout();         let task = cx            .spawn({                async move |lsp_store, cx| {                    let response = server                        .request::<lsp::request::CodeLensResolve>(lens, request_timeout)                        .await                        .into_response();                    lsp_store                        .update(cx, |lsp_store, _| {                            let code_lens = lsp_store                                .lsp_data                                .get_mut(&buffer_id)                                .and_then(|data| data.code_lens.as_mut())?;                            code_lens.resolving.remove(&key);                            let resolved_lens = match response {                                Ok(resolved_lens) => resolved_lens,                                Err(e) => {                                    log::warn!("Failed to resolve code lens: {e:#}");                                    return None;                                }                            };                            let action = code_lens                                .lens                                .get_mut(&server_id)                                .and_then(|cache| cache.get_mut(&lens_id))?;                            action.resolved = true;                            action.lsp_action = LspAction::CodeLens(resolved_lens);                            Some((lens_id, action.clone()))                        })                        .ok()                        .flatten()                }            })            .shared();         if let Some(code_lens) = self            .lsp_data            .get_mut(&buffer_id)            .and_then(|data| data.code_lens.as_mut())        {            code_lens.resolving.insert(key, task.clone());        }        task    }     #[cfg(any(test, feature = "test-support"))]    pub fn forget_code_lens_task(&mut self, buffer_id: text::BufferId) -> Option<CodeLensTask> {        Some(            self.lsp_data                .get_mut(&buffer_id)?                .code_lens                .take()?

Lines 1–500 of 569.