fix Reference Counting #2194

pyrefly/lib/lsp/non_wasm.rs

@@ -7,7 +7,6 @@
 
 mod build_system;
 pub mod call_hierarchy;
-pub mod code_lens;
 pub mod connection;
 pub mod convert_module_package;
 pub mod document_symbols;

pyrefly/lib/lsp/non_wasm/server.rs

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+mod code_lens;
 mod event_loop;
 
 use std::cmp::min;
@@ -52,10 +53,9 @@ use lsp_types::CodeActionProvider;
 use lsp_types::CodeActionRequest;
 use lsp_types::CodeActionResponse;
 use lsp_types::CodeActionTriggerKind;
-use lsp_types::CodeLens;
 use lsp_types::CodeLensOptions;
-use lsp_types::CodeLensParams;
 use lsp_types::CodeLensRequest;
+use lsp_types::CodeLensResolveRequest;
 use lsp_types::CompletionItem;
 use lsp_types::CompletionList;
 use lsp_types::CompletionOptions;
@@ -299,7 +299,6 @@ use crate::lsp::non_wasm::call_hierarchy::find_function_at_position_in_ast;
 use crate::lsp::non_wasm::call_hierarchy::prepare_call_hierarchy_item;
 use crate::lsp::non_wasm::call_hierarchy::transform_incoming_calls;
 use crate::lsp::non_wasm::call_hierarchy::transform_outgoing_calls;
-use crate::lsp::non_wasm::code_lens::runnable_lsp_code_lens;
 use crate::lsp::non_wasm::convert_module_package::convert_module_package_code_actions;
 use crate::lsp::non_wasm::document_symbols::flatten_to_symbol_information;
 use crate::lsp::non_wasm::external_provider::ExternalProvider;
@@ -1672,7 +1671,7 @@ pub fn capabilities(
             ..Default::default()
         })),
         code_lens_provider: Some(CodeLensOptions {
-            resolve_provider: Some(false),
+            resolve_provider: Some(true),
             work_done_progress_options: Default::default(),
         }),
         completion_provider: Some(CompletionOptions {
@@ -2376,6 +2375,7 @@ impl Server {
                     DefinitionRequest::METHOD.as_str(),
                     ProvideType::METHOD.as_str(),
                     TypeErrorDisplayStatusRequest::METHOD.as_str(),
+                    CodeLensResolveRequest::METHOD.as_str(),
                 ];
 
                 let in_cancelled_requests = canceled_requests.remove(&x.id);
@@ -2700,18 +2700,6 @@ impl Server {
                         };
                         self.send_response(new_response(x.id, Ok(response)));
                     }
-                } else if let Some(params) = as_request::<CodeLensRequest>(&x) {
-                    if let Some(params) = self
-                        .extract_request_params_or_send_err_response::<CodeLensRequest>(
-                            params, &x.id,
-                        )
-                    {
-                        self.set_file_stats(params.text_document.uri.clone(), telemetry_event);
-                        self.send_response(new_response(
-                            x.id,
-                            Ok(self.code_lens(&transaction, params).unwrap_or_default()),
-                        ));
-                    }
                 } else if let Some(params) = as_request::<SemanticTokensRequest>(&x) {
                     if let Some(params) = self
                         .extract_request_params_or_send_err_response::<SemanticTokensRequest>(
@@ -2757,6 +2745,37 @@ impl Server {
                         };
                         self.send_response(new_response(x.id, Ok(response)));
                     }
+                } else if let Some(params) = as_request::<CodeLensRequest>(&x) {
+                    if let Some(params) = self
+                        .extract_request_params_or_send_err_response::<CodeLensRequest>(
+                            params, &x.id,
+                        )
+                    {
+                        self.set_file_stats(params.text_document.uri.clone(), telemetry_event);
+                        let response = match self.code_lens(&transaction, params) {
+                            Ok(response) => response,
+                            Err(reason) => {
+                                telemetry_event.set_empty_response_reason(reason);
+                                None
+                            }
+                        };
+                        self.send_response(new_response(x.id, Ok(response)));
+                    }
+                } else if let Some(params) = as_request::<CodeLensResolveRequest>(&x) {
+                    if let Some(params) = self
+                        .extract_request_params_or_send_err_response::<CodeLensResolveRequest>(
+                            params, &x.id,
+                        )
+                        && let Err(reason) = self.resolve_code_lens(
+                            x.id.clone(),
+                            &transaction,
+                            &params,
+                            telemetry_event.activity_key.clone(),
+                        )
+                    {
+                        telemetry_event.set_empty_response_reason(reason);
+                        self.send_response(new_response(x.id, Ok(params)));
+                    }
                 } else if let Some(params) = as_request::<WorkspaceSymbolRequest>(&x) {
                     if let Some(params) = self
                         .extract_request_params_or_send_err_response::<WorkspaceSymbolRequest>(
@@ -3190,24 +3209,6 @@ impl Server {
         telemetry.set_file_stats(TelemetryFileStats { uri, config_root });
     }
 
-    fn runnable_code_lens_cwd(&self, path: &std::path::Path) -> Option<String> {
-        let config = self.state.config_finder().python_file(
-            ModuleNameWithKind::guaranteed(ModuleName::unknown()),
-            &ModulePath::filesystem(path.to_path_buf()),
-        );
-        let cwd = config
-            .source
-            .root_from_file()
-            .map(std::path::Path::to_path_buf)
-            .or_else(|| {
-                self.workspaces
-                    .get_with(path.to_path_buf(), |(workspace_root, _)| {
-                        workspace_root.cloned()
-                    })
-            })?;
-        Some(cwd.to_string_lossy().into_owned())
-    }
-
     fn send_response(&self, x: Response) {
         self.connection.send(Message::Response(x))
     }
@@ -6019,36 +6020,6 @@ impl Server {
         Ok(Some(res))
     }
 
-    fn code_lens(
-        &self,
-        transaction: &Transaction<'_>,
-        params: CodeLensParams,
-    ) -> Option<Vec<CodeLens>> {
-        let uri = &params.text_document.uri;
-        let path = self.path_for_uri(uri)?;
-        let runnable_code_lens = self
-            .workspaces
-            .get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens);
-        let maybe_cell_idx = self.maybe_get_code_cell_index(uri);
-        let handle = self
-            .make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD.as_str()))
-            .ok()?;
-        let info = transaction.get_module_info(&handle)?;
-        let entries = transaction.runnable_code_lens_entries(&handle, uri, runnable_code_lens)?;
-        let cwd = self.runnable_code_lens_cwd(&path);
-
-        let mut lenses = Vec::new();
-        for entry in entries {
-            if info.to_cell_for_lsp(entry.range.start()) != maybe_cell_idx {
-                continue;
-            }
-            let range = info.to_lsp_range(entry.range);
-            lenses.push(runnable_lsp_code_lens(uri, range, entry, cwd.as_deref()));
-        }
-
-        Some(lenses)
-    }
-
     fn semantic_tokens_full(
         &self,
         transaction: &Transaction<'_>,

pyrefly/lib/lsp/non_wasm/server/code_lens.rs

@@ -0,0 +1,121 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+//! Server orchestration for runnable and reference-count CodeLens providers.
+
+mod reference_count;
+mod runnable;
+
+use lsp_server::ErrorCode;
+use lsp_server::RequestId;
+use lsp_types::CodeLens;
+use lsp_types::CodeLensParams;
+use lsp_types::CodeLensRequest;
+use lsp_types::Request as _;
+use pyrefly_python::module::TextRangeWithModule;
+use pyrefly_util::telemetry::ActivityKey;
+use pyrefly_util::telemetry::EmptyResponseReason;
+
+use super::Server;
+use crate::commands::lsp::IndexingMode;
+use crate::lsp::non_wasm::protocol::Message;
+use crate::lsp::non_wasm::protocol::Response;
+use crate::state::lsp::FindDefinitionItemWithDocstring;
+use crate::state::lsp::FindPreference;
+use crate::state::lsp::ImportBehavior;
+use crate::state::lsp::ReferenceOptions;
+use crate::state::state::Transaction;
+
+impl Server {
+    /// Return eager runnable lenses and unresolved reference-count lenses for a Python file.
+    pub(super) fn code_lens(
+        &self,
+        transaction: &Transaction<'_>,
+        params: CodeLensParams,
+    ) -> Result<Option<Vec<CodeLens>>, EmptyResponseReason> {
+        let uri = &params.text_document.uri;
+        if self.open_notebook_cells.read().contains_key(uri) {
+            return Ok(Some(Vec::new()));
+        }
+
+        let handle = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD.as_str()))?;
+        let mut lenses = runnable::code_lenses(self, transaction, &handle, uri)?;
+        if self.indexing_mode != IndexingMode::None {
+            lenses.extend(reference_count::unresolved_code_lenses(
+                transaction,
+                &handle,
+                uri,
+            )?);
+        }
+        Ok(Some(lenses))
+    }
+
+    /// Resolve one visible reference-count lens on the references queue.
+    pub(super) fn resolve_code_lens<'a>(
+        &'a self,
+        request_id: RequestId,
+        transaction: &Transaction<'a>,
+        lens: &CodeLens,
+        activity_key: Option<ActivityKey>,
+    ) -> Result<(), EmptyResponseReason> {
+        let data = match reference_count::resolve_data(lens) {
+            Ok(data) => data,
+            Err(message) => {
+                self.connection.send(Message::Response(Response::new_err(
+                    request_id,
+                    ErrorCode::InvalidParams as i32,
+                    message,
+                )));
+                return Ok(());
+            }
+        };
+        if self.open_notebook_cells.read().contains_key(&data.uri) {
+            return Err(EmptyResponseReason::NotebookNotSupported);
+        }
+
+        let handle =
+            self.make_handle_if_enabled(&data.uri, Some(CodeLensRequest::METHOD.as_str()))?;
+        let path_remapper = self.path_remapper.clone();
+        let source_uri = data.uri;
+        let source_uri_for_response = source_uri.clone();
+        let lens = lens.clone();
+        self.async_find_from_definition_helper(
+            request_id,
+            transaction,
+            handle,
+            &source_uri,
+            lens.range.start,
+            FindPreference {
+                import_behavior: ImportBehavior::StopAtRenamedImports,
+                ..Default::default()
+            },
+            activity_key,
+            move |transaction, handle, definition, _telemetry, _telemetry_event| {
+                let FindDefinitionItemWithDocstring {
+                    metadata,
+                    definition_range,
+                    module,
+                    ..
+                } = definition;
+                Ok(transaction.find_global_references_from_definition(
+                    *handle.sys_info(),
+                    metadata,
+                    TextRangeWithModule::new(module, definition_range),
+                    ReferenceOptions::all(false),
+                )?)
+            },
+            move |results| {
+                reference_count::resolve_code_lens(
+                    lens,
+                    &source_uri_for_response,
+                    results,
+                    path_remapper.as_ref(),
+                )
+            },
+        )
+    }
+}