fix Go-to definition only finds definition from one element in a union #1470 (#1569)

Summary:
fix #1470

Playground has this issue in the sandbox because its go-to-definition call only returned the first Range, even though the core analyzer already reports every branch of a union.

Now collects every TextRangeWithModule from the transaction.goto_definition instead of truncating to the first match, serializes all of those ranges so the wasm bridge preserves the multiplicity.

Pull Request resolved: #1569

Reviewed By: stroxler

Differential Revision: D86973443

Pulled By: yangdanny97

fbshipit-source-id: 98332098260745edac0bbaef3461923ee718729c

Author: asukaminato0721

pyrefly/lib/playground.rs

@@ -468,16 +468,21 @@ impl Playground {
         SemanticTokensLegends::lsp_semantic_token_legends()
     }
 
-    pub fn goto_definition(&mut self, pos: Position) -> Option<Range> {
-        let handle = self.handles.get(&self.active_filename)?;
+    pub fn goto_definition(&mut self, pos: Position) -> Vec<Range> {
+        let handle = match self.handles.get(&self.active_filename) {
+            Some(handle) => handle,
+            None => return Vec::new(),
+        };
         let transaction = self.state.transaction();
-        let position = self.to_text_size(&transaction, pos)?;
-        // TODO: Support goto multiple definitions
+        let position = match self.to_text_size(&transaction, pos) {
+            Some(position) => position,
+            None => return Vec::new(),
+        };
         transaction
             .goto_definition(handle, position)
             .into_iter()
-            .next()
             .map(|r| Range::new(r.module.display_range(r.range)))
+            .collect()
     }
 
     pub fn autocomplete(&self, pos: Position) -> Vec<AutoCompletionItem> {

pyrefly/lib/test/lsp/definition.rs

@@ -65,6 +65,7 @@ x = 1 # go-to-definition is unsupported for literals
 #   ^
 "#;
     let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    println!("REPORT=>{}<=", report);
     assert_eq!(
         r#"
 # main.py
@@ -1287,6 +1288,55 @@ Definition Result:
     );
 }
 
+#[test]
+fn union_method_access_test() {
+    let code = r#"
+class State:
+    def get_location(self) -> str: ...
+
+class NewYork:
+    def get_location(self) -> str:
+        return "10016"
+
+class Massachusetts:
+    def get_location(self) -> str:
+        return "02108"
+
+def find_location(x: Massachusetts | NewYork):
+    x.get_location()
+      # ^
+
+def find_location2(x: NewYork | Massachusetts):
+    x.get_location()
+      # ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+14 |     x.get_location()
+             ^
+Definition Result:
+10 |     def get_location(self) -> str:
+             ^^^^^^^^^^^^
+Definition Result:
+6 |     def get_location(self) -> str:
+            ^^^^^^^^^^^^
+
+18 |     x.get_location()
+             ^
+Definition Result:
+10 |     def get_location(self) -> str:
+             ^^^^^^^^^^^^
+Definition Result:
+6 |     def get_location(self) -> str:
+            ^^^^^^^^^^^^
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
 // todo(kylei) go-to definition on x should go to the definition
 #[test]
 fn global_keyword() {

pyrefly_wasm/lib.rs

@@ -76,10 +76,12 @@ impl State {
 
     #[wasm_bindgen(js_name=gotoDefinition)]
     pub fn goto_definition(&mut self, line: i32, column: i32) -> JsValue {
-        self.0
-            .goto_definition(Position::new(line, column))
-            .map(|result| serde_wasm_bindgen::to_value(&result).unwrap())
-            .unwrap_or(JsValue::NULL)
+        let results = self.0.goto_definition(Position::new(line, column));
+        if results.is_empty() {
+            JsValue::NULL
+        } else {
+            serde_wasm_bindgen::to_value(&results).unwrap_or(JsValue::NULL)
+        }
     }
 
     #[wasm_bindgen(js_name=autoComplete)]

website/src/__tests__/__mocks__/configuredMonacoMock.ts

@@ -72,7 +72,7 @@ function setAutoCompleteFunction(
  */
 function setGetDefFunction(
     _model: MockEditorModel,
-    _getDefFunction: (line: number, column: number) => any
+    _getDefFunction: (line: number, column: number) => any[] | null
 ): void {}
 
 /**

website/src/__tests__/__mocks__/pyreflyWasmMock.ts

@@ -61,8 +61,8 @@ class MockState {
      * @param column Column number
      * @returns Empty definition result
      */
-    gotoDefinition(line?: number, column?: number): Record<string, any> {
-        return {};
+    gotoDefinition(_line?: number, _column?: number): Record<string, any>[] {
+        return [];
     }
 
     /**

website/src/__tests__/pyrefly_wasm.test.ts

@@ -118,12 +118,13 @@ movie: Movie = {'name': 'Blade Runner',
         it('should return definition location for function call', () => {
             pyreService.setActiveFile('main.py');
             // Position of "test" in "test(42)"
-            const definition = pyreService.gotoDefinition(18, 13);
+            const definitions = pyreService.gotoDefinition(18, 13);
 
-            // Should return a location object
-            expect(definition).toBeDefined();
+            expect(definitions).toBeDefined();
+            expect(definitions).not.toBeNull();
+            expect(definitions!.length).toBeGreaterThan(0);
 
-            // expect that location is correct
+            const definition = definitions![0];
             expect(definition.startLineNumber).toBe(13);
             expect(definition.startColumn).toBe(5);
             expect(definition.endLineNumber).toBe(13);

website/src/sandbox/Sandbox.tsx

@@ -45,7 +45,7 @@ export interface PyreflyState {
     setActiveFile: (filename: string) => void;
     getErrors: () => ReadonlyArray<PyreflyErrorMessage>;
     autoComplete: (line: number, column: number) => any;
-    gotoDefinition: (line: number, column: number) => any;
+    gotoDefinition: (line: number, column: number) => monaco.IRange[] | null;
     hover: (line: number, column: number) => any;
     inlayHint: () => any;
     semanticTokens: (range: any) => any;

website/src/sandbox/configured-monaco.ts

@@ -18,7 +18,10 @@ type SemanticTokens = monaco.languages.SemanticTokens;
 type SemanticTokensLegend = monaco.languages.SemanticTokensLegend;
 
 type AutoCompleteFunction = (line: number, column: number) => CompletionItem[];
-type GetDefFunction = (line: number, column: number) => Range | null;
+type GetDefFunction = (
+    line: number,
+    column: number
+) => ReadonlyArray<Range> | null;
 type HoverFunction = (line: number, column: number) => Hover | null;
 type InlayHintFunction = () => InlayHint[];
 type SemanticTokensFunction = (range: Range | null) => SemanticTokens | null;
@@ -49,7 +52,7 @@ function setAutoCompleteFunction(
 const defaultGetDefFunctionForMonaco: GetDefFunction = (
     _l: number,
     _c: number
-): Range | null => null;
+): ReadonlyArray<Range> | null => null;
 const getDefFunctionsForMonaco = new Map<
     monaco.editor.ITextModel,
     GetDefFunction
@@ -217,8 +220,14 @@ monaco.languages.registerDefinitionProvider('python', {
             const f =
                 getDefFunctionsForMonaco.get(model) ??
                 defaultGetDefFunctionForMonaco;
-            const range = f(position.lineNumber, position.column);
-            return range != null ? { uri: model.uri, range } : null;
+            const ranges = f(position.lineNumber, position.column);
+            if (!ranges || ranges.length === 0) {
+                return null;
+            }
+            if (ranges.length === 1) {
+                return { uri: model.uri, range: ranges[0] };
+            }
+            return ranges.map((range) => ({ uri: model.uri, range }));
         } catch (e) {
             console.error(e);
             return null;