fix Reference Counting #2194- #2907

crates/pyrefly_lsp_test/src/object_model.rs

@@ -71,6 +71,7 @@ use pretty_assertions::assert_eq;
 use pyrefly::commands::lsp::IndexingMode;
 use pyrefly::commands::lsp::LspArgs;
 use pyrefly::commands::lsp::run_lsp;
+pub use pyrefly::lsp::non_wasm::external_provider::ExternalProvider;
 use pyrefly::lsp::non_wasm::external_provider::NoExternalProvider;
 use pyrefly::lsp::non_wasm::module_helpers::ThriftRemapper;
 use pyrefly::lsp::non_wasm::protocol::JsonRpcMessage;
@@ -1438,6 +1439,7 @@ pub struct LspInteractionArgs {
     pub telemetry: Box<dyn Telemetry>,
     pub thread_count: ThreadCount,
     pub thrift_remapper: Option<ThriftRemapper>,
+    pub external_references: Arc<dyn ExternalProvider>,
 }
 
 impl Default for LspInteractionArgs {
@@ -1451,6 +1453,7 @@ impl Default for LspInteractionArgs {
             telemetry: Box::new(NoTelemetry),
             thread_count: TEST_THREAD_COUNT,
             thrift_remapper: None,
+            external_references: Arc::new(NoExternalProvider),
         }
     }
 }
@@ -1472,6 +1475,7 @@ impl LspInteraction {
             telemetry,
             thread_count,
             thrift_remapper,
+            external_references,
         } = args;
 
         let ((conn_client, _client_reader), (conn_server, server_reader)) = Connection::memory();
@@ -1489,7 +1493,7 @@ impl LspInteraction {
                 None,
                 thrift_remapper,
                 &*telemetry,
-                Arc::new(NoExternalProvider),
+                external_references,
                 None,
                 thread_count,
             );

crates/pyrefly_util/src/telemetry.rs

@@ -44,6 +44,7 @@ pub enum QueueName {
     LspQueue,
     RecheckQueue,
     FindReferenceQueue,
+    CodeLensQueue,
     SourceDbQueue,
 }
 
@@ -59,6 +60,7 @@ impl QueueName {
             Self::LspQueue => "lsp_queue",
             Self::RecheckQueue => "recheck_queue",
             Self::FindReferenceQueue => "find_reference_queue",
+            Self::CodeLensQueue => "code_lens_queue",
             Self::SourceDbQueue => "sourcedb_queue",
         }
     }

lsp/package.json

@@ -397,6 +397,12 @@
                     "description": "Enable Pyrefly's Run/Test CodeLens actions for Python files.",
                     "scope": "resource"
                 },
+                "python.pyrefly.referencesCodeLens": {
+                    "type": "boolean",
+                    "default": false,
+                    "description": "Show reference counts above Python definitions. Requires indexing. Run/Test CodeLens uses a separate setting.",
+                    "scope": "resource"
+                },
                 "python.pyrefly.streamDiagnostics": {
                     "type": "boolean",
                     "default": true,

lsp/src/code-lens.ts

@@ -11,6 +11,7 @@ import {execFile} from 'child_process';
 import * as path from 'path';
 import * as vscode from 'vscode';
 import {ExtensionContext} from 'vscode';
+import type {Location as ProtocolLocation} from 'vscode-languageclient/node';
 import {PythonEnvironment} from './python-environment';
 
 type CodeLensPosition = {
@@ -323,6 +324,34 @@ export function registerCodeLensCommands(
   context: ExtensionContext,
   pythonEnv: PythonEnvironment,
 ): void {
+  context.subscriptions.push(
+    vscode.commands.registerCommand(
+      'pyrefly.showReferences',
+      async (
+        uri: string,
+        position: CodeLensPosition,
+        locations: ProtocolLocation[],
+      ) => {
+        await vscode.commands.executeCommand(
+          'editor.action.showReferences',
+          vscode.Uri.parse(uri),
+          new vscode.Position(position.line, position.character),
+          locations.map(location =>
+            new vscode.Location(
+              vscode.Uri.parse(location.uri),
+              new vscode.Range(
+                location.range.start.line,
+                location.range.start.character,
+                location.range.end.line,
+                location.range.end.character,
+              ),
+            ),
+          ),
+        );
+      },
+    ),
+  );
+
   context.subscriptions.push(
     vscode.commands.registerCommand('pyrefly.runTest', async args => {
       const parsedArgs = parseRunTestArgs(args);

lsp/src/extension.ts

@@ -297,6 +297,7 @@ export async function activate(
       customHoverProvider: supportsHoverVerbosity,
       findFileUsages: true,
       pushTypeErrorDisplayStatus: true,
+      supportsReferencesCodeLens: true,
     },
   };
 

lsp/src/test/extension.test.ts

@@ -47,4 +47,38 @@ suite('Extension Test Suite', () => {
       await fs.rm(directory, {recursive: true, force: true});
     }
   });
+
+  test('Show reference CodeLens arguments through the VS Code command', async function () {
+    this.timeout(30000);
+    assert.ok(extension);
+    await extension.activate();
+
+    const directory = await fs.mkdtemp(join(tmpdir(), 'pyrefly-references-'));
+    const uri = vscode.Uri.file(join(directory, 'references.py'));
+    try {
+      await vscode.workspace.fs.writeFile(
+        uri,
+        Buffer.from('def foo(): pass\nfoo()\n'),
+      );
+      const document = await vscode.workspace.openTextDocument(uri);
+      await vscode.window.showTextDocument(document);
+      // The built-in command rejects the serialized URI and position without conversion.
+      await vscode.commands.executeCommand(
+        'pyrefly.showReferences',
+        uri.toString(),
+        {line: 0, character: 4},
+        [{
+          uri: uri.toString(),
+          range: {
+            start: {line: 1, character: 0},
+            end: {line: 1, character: 3},
+          },
+        }],
+      );
+    } finally {
+      await vscode.commands.executeCommand('closeReferenceSearch');
+      await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+      await fs.rm(directory, {recursive: true, force: true});
+    }
+  });
 });

pyrefly/Cargo.toml

@@ -20,6 +20,10 @@ path = "bin/main.rs"
 name = "pyrefly_laziness_tests"
 path = "test_laziness/mod.rs"
 
+[[bench]]
+name = "code_lens"
+harness = false
+
 [[bench]]
 name = "commit"
 harness = false

pyrefly/benches/README.md

@@ -1,3 +1,17 @@
+# Reference CodeLens benchmark
+
+`code_lens` measures warm LSP requests against a generated 10,000-line Python
+file with 2,500 function definitions and ten importing modules. It measures
+listing unresolved lenses separately from resolving twenty visible lenses.
+Fixture creation, initial analysis, and server shutdown are not timed.
+The fixture does not use an external reference provider.
+
+```bash
+PYREFLY_LOG=off cargo bench -p pyrefly --bench code_lens -- --noplot
+```
+
+These synthetic results do not establish performance for every workspace.
+
 # PyTorch benchmarks
 
 Real-world benchmarks over a large, pinned PyTorch checkout (15k+ Python files)

pyrefly/benches/code_lens.rs

@@ -0,0 +1,138 @@
+/*
+ * 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.
+ */
+
+//! Measure listing and resolving visible reference lenses in a 10,000-line file.
+
+use std::cell::RefCell;
+use std::fs;
+use std::time::Duration;
+
+use criterion::Criterion;
+use criterion::criterion_group;
+use criterion::criterion_main;
+use lsp_types::CodeLens;
+use lsp_types::CodeLensRequest;
+use lsp_types::CodeLensResolveRequest;
+use lsp_types::Uri;
+use pyrefly_lsp_test::IndexingMode;
+use pyrefly_lsp_test::LspArgs;
+use pyrefly_lsp_test::object_model::InitializeSettings;
+use pyrefly_lsp_test::object_model::LspInteraction;
+use pyrefly_lsp_test::object_model::LspInteractionArgs;
+use pyrefly_util::telemetry::NoTelemetry;
+use serde_json::json;
+use tempfile::TempDir;
+
+const DEFINITIONS: usize = 2500;
+const VISIBLE_LENSES: usize = 20;
+const IMPORTERS: usize = 10;
+
+/// Measure warm LSP requests with fixture creation and initial analysis excluded.
+fn code_lens(c: &mut Criterion) {
+    let root = TempDir::new().unwrap();
+    fs::write(
+        root.path().join("pyrefly.toml"),
+        "project-includes = [\"**/*.py\"]\n",
+    )
+    .unwrap();
+    let mut source = String::new();
+    for index in 0..DEFINITIONS {
+        source.push_str(&format!(
+            "def function_{index}() -> int:\n    return {index}\n\nfunction_{index}()\n"
+        ));
+    }
+    assert_eq!(source.lines().count(), 10_000);
+    fs::write(root.path().join("large.py"), source).unwrap();
+    let mut usage = "import large\n".to_owned();
+    for index in 0..VISIBLE_LENSES {
+        usage.push_str(&format!("large.function_{index}()\n"));
+    }
+    for index in 0..IMPORTERS {
+        fs::write(root.path().join(format!("usage_{index}.py")), &usage).unwrap();
+    }
+
+    let uri = Uri::from_file_path(root.path().join("large.py")).unwrap();
+    let mut interaction = LspInteraction::new_with_args(LspInteractionArgs {
+        args: LspArgs {
+            indexing_mode: IndexingMode::LazyBlocking,
+            ..LspInteractionArgs::default().args
+        },
+        telemetry: Box::new(NoTelemetry),
+        ..Default::default()
+    });
+    interaction.client.set_message_logging(false);
+    interaction.set_root(root.path().to_path_buf());
+    interaction
+        .initialize(InitializeSettings {
+            workspace_folders: Some(vec![(
+                "benchmark".to_owned(),
+                Uri::from_file_path(root.path()).unwrap(),
+            )]),
+            initialization_options: Some(json!({"pyrefly": {
+                "supportsReferencesCodeLens": true,
+                "referencesCodeLens": true
+            }})),
+            ..Default::default()
+        })
+        .unwrap();
+    interaction.client.did_open("large.py");
+    let params = json!({"textDocument": {"uri": uri}});
+    let lenses = RefCell::new(Vec::<CodeLens>::new());
+    interaction
+        .client
+        .send_request::<CodeLensRequest>(params.clone())
+        .expect_response_with(|response| {
+            *lenses.borrow_mut() = response.unwrap();
+            assert_eq!(lenses.borrow().len(), DEFINITIONS);
+            true
+        })
+        .unwrap();
+    let visible: Vec<_> = lenses
+        .into_inner()
+        .into_iter()
+        .take(VISIBLE_LENSES)
+        .map(|lens| json!(lens))
+        .collect();
+    let resolve_visible = || {
+        for lens in &visible {
+            interaction
+                .client
+                .send_request::<CodeLensResolveRequest>(lens.clone())
+                .expect_response_with(|resolved| {
+                    assert_eq!(resolved.command.unwrap().title, "11 references");
+                    true
+                })
+                .unwrap();
+        }
+    };
+    resolve_visible();
+
+    let mut group = c.benchmark_group("code_lens/10000_lines");
+    group.sample_size(10);
+    group.warm_up_time(Duration::from_secs(1));
+    group.measurement_time(Duration::from_secs(2));
+    group.bench_function("list", |b| {
+        b.iter(|| {
+            interaction
+                .client
+                .send_request::<CodeLensRequest>(params.clone())
+                .expect_response_with(|response| {
+                    let lenses = response.unwrap();
+                    assert_eq!(lenses.len(), DEFINITIONS);
+                    assert!(lenses.iter().all(|lens| lens.command.is_none()));
+                    true
+                })
+                .unwrap();
+        });
+    });
+    group.bench_function("resolve_20_visible", |b| b.iter(resolve_visible));
+    group.finish();
+    interaction.shutdown().unwrap();
+}
+
+criterion_group!(benches, code_lens);
+criterion_main!(benches);

pyrefly/lib/lsp/non_wasm.rs

@@ -20,6 +20,7 @@ pub mod move_symbol_new_file;
 mod mru;
 pub mod protocol;
 pub mod queue;
+mod runnable;
 pub mod safe_delete_file;
 pub mod server;
 pub mod stdlib;