fix LSP: Go to Definition on nn.Module calls should navigate to forward() #5229

pyrefly/lib/state/lsp.rs

@@ -32,6 +32,7 @@ use pyrefly_python::module_path::ModuleStyle;
 use pyrefly_python::short_identifier::ShortIdentifier;
 use pyrefly_python::symbol_kind::SymbolKind;
 use pyrefly_python::sys_info::SysInfo;
+use pyrefly_types::class::ClassType;
 use pyrefly_types::function::FuncMetadata;
 use pyrefly_types::function::FunctionKind;
 use pyrefly_types::type_alias::TypeAliasData;
@@ -2188,11 +2189,47 @@ impl<'a> Transaction<'a> {
         }))
     }
 
+    /// Find `forward` definitions for a PyTorch module instance.
+    fn find_nn_module_forward_definitions(
+        &self,
+        handle: &Handle,
+        preference: FindPreference,
+        cls: &ClassType,
+        ty: &Type,
+    ) -> Option<Vec1<FindDefinitionItemWithDocstring>> {
+        let is_nn_module = self
+            .ad_hoc_solve(handle, "check_nn_module", |solver| {
+                let class = cls.class_object();
+                let mro = solver.get_mro_for_class(class);
+                iter::once(class)
+                    .chain(
+                        mro.ancestors_no_object()
+                            .iter()
+                            .map(|base| base.class_object()),
+                    )
+                    .any(|base| {
+                        base.has_toplevel_qname("torch.nn", "Module")
+                            || base.has_toplevel_qname("torch.nn.modules.module", "Module")
+                    })
+            })
+            .unwrap_or(false);
+        if !is_nn_module {
+            return None;
+        }
+        self.find_attribute_definition_for_base_type(
+            handle,
+            preference,
+            ty.clone(),
+            &Name::new_static("forward"),
+        )
+        .ok()
+    }
+
     /// When a name or attribute in a call position resolves to a class, find
     /// `__init__` and `__new__` definitions. When it resolves to a class
-    /// instance, find `__call__`. Returns all found definitions, or empty if
-    /// neither case applies. Does not match functions/callables — those should
-    /// use the normal go-to-definition path.
+    /// instance, find `forward` for PyTorch modules and `__call__` otherwise.
+    /// Returns all found definitions, or empty if neither case applies.
+    /// Functions and callables use the normal go-to-definition path.
     ///
     /// When a class synthesizes its own `__init__` or `__new__` (e.g.
     /// dataclass, pydantic, TypedDict, NamedTuple), we skip the corresponding
@@ -2239,10 +2276,21 @@ impl<'a> Transaction<'a> {
                 );
                 defs
             }
-            Type::ClassType(_) => self
-                .find_attribute_definition_for_base_type(handle, preference, ty, &dunder::CALL)
-                .map(Vec1::into_vec)
-                .unwrap_or_default(),
+            Type::ClassType(cls) | Type::SelfType(cls) => {
+                if let Some(defs) =
+                    self.find_nn_module_forward_definitions(handle, preference, cls, &ty)
+                {
+                    return defs.into_vec();
+                }
+                self.find_attribute_definition_for_base_type(handle, preference, ty, &dunder::CALL)
+                    .map(Vec1::into_vec)
+                    .unwrap_or_default()
+            }
+            Type::NNModule(module) => self.find_call_target_definitions(
+                handle,
+                preference,
+                module.class.clone().to_type(),
+            ),
             _ => vec![],
         }
     }

pyrefly/lib/test/lsp/definition.rs

@@ -3139,6 +3139,9 @@ class Adder:
     def __call__(self, x: int) -> int:
         return x + 1
 
+    def forward(self, x: int) -> int:
+        return x
+
 adder = Adder()
 adder(5)
 # ^
@@ -3147,8 +3150,8 @@ adder(5)
     assert_eq!(
         r#"
 # main.py
-7 | adder(5)
-      ^
+10 | adder(5)
+       ^
 Definition Result:
 3 |     def __call__(self, x: int) -> int:
             ^^^^^^^^
@@ -3158,6 +3161,92 @@ Definition Result:
     );
 }
 
+// Regression test for https://github.com/facebook/pyrefly/issues/5229.
+#[test]
+fn goto_def_nn_module_call_goes_to_forward() {
+    let module_code = r#"
+class Module:
+    def __call__(self, x: int) -> int: ...
+"#;
+    for module_name in ["torch.nn", "torch.nn.modules.module"] {
+        let code = format!(
+            r#"
+from {module_name} import Module
+
+class Model(Module):
+    def forward(self, x: int) -> int:
+        return x
+
+    def run(self) -> int:
+        return self(1)
+#               ^
+
+class Inherited(Model):
+    pass
+
+class Overridden(Model):
+    def forward(self, x: int) -> int:
+        return x + 1
+
+class Holder:
+    model: Model = Model()
+
+model = Model()
+model(1)
+# ^
+inherited = Inherited()
+inherited(1)
+# ^
+overridden = Overridden()
+overridden(1)
+# ^
+holder = Holder()
+holder.model(1)
+#        ^
+model.__call__(1)
+#       ^
+model
+# ^
+module = Module()
+module(1)
+# ^
+"#
+        );
+        let mut env = TestEnv::one(module_name, module_code);
+        env.add("main", &code);
+        let (state, handles) = env.to_state();
+        let handle = handles("main");
+        let forward = code.find("def forward").unwrap() + 4;
+        let overridden_forward = code.rfind("def forward").unwrap() + 4;
+        let call = module_code.find("__call__").unwrap();
+        let model = code.find("model = Model()").unwrap();
+        let expected = [
+            ("main", forward, "forward"),
+            ("main", forward, "forward"),
+            ("main", forward, "forward"),
+            ("main", overridden_forward, "forward"),
+            ("main", forward, "forward"),
+            (module_name, call, "__call__"),
+            ("main", model, "model"),
+            (module_name, call, "__call__"),
+        ];
+        let positions = extract_cursors_for_test(&code);
+        assert_eq!(positions.len(), expected.len());
+        for (position, (module, start, name)) in positions.into_iter().zip(expected) {
+            let definitions = state
+                .transaction()
+                .goto_definition(&handle, position)
+                .unwrap();
+            assert_eq!(definitions.len(), 1);
+            assert_eq!(definitions[0].module.name(), ModuleName::from_str(module));
+            assert_eq!(
+                definitions[0].range,
+                TextRange::at(TextSize::new(start as u32), TextSize::of(name)),
+            );
+        }
+    }
+}
+
 #[test]
 fn goto_def_non_constructor_call_goes_to_function() {
     let code = r#"