Add method return type inference for user-defined methods

rust/src/analyzer/definitions.rs

@@ -8,6 +8,7 @@
 
 use crate::env::{GlobalEnv, LocalEnv};
 use crate::graph::{ChangeSet, VertexId};
+use crate::types::Type;
 use ruby_prism::Node;
 
 use super::install::install_statements;
@@ -64,16 +65,29 @@ pub(crate) fn process_def_node(
     def_node: &ruby_prism::DefNode,
 ) -> Option<VertexId> {
     let method_name = String::from_utf8_lossy(def_node.name().as_slice()).to_string();
-    install_method(genv, method_name);
+    install_method(genv, method_name.clone());
 
     // Process parameters BEFORE processing body
     if let Some(params_node) = def_node.parameters() {
         install_parameters(genv, lenv, changes, source, &params_node);
     }
 
+    let mut last_vtx = None;
     if let Some(body) = def_node.body() {
         if let Some(statements) = body.as_statements_node() {
-            install_statements(genv, lenv, changes, source, &statements);
+            last_vtx = install_statements(genv, lenv, changes, source, &statements);
+        }
+    }
+
+    // Register user-defined method with return vertex (before exiting scope)
+    if let Some(return_vtx) = last_vtx {
+        let recv_type_name = genv
+            .scope_manager
+            .current_class_name()
+            .or_else(|| genv.scope_manager.current_module_name());
+
+        if let Some(name) = recv_type_name {
+            genv.register_user_method(Type::instance(&name), &method_name, return_vtx);
         }
     }
 
@@ -151,7 +165,9 @@ fn extract_constant_path(node: &Node) -> Option<String> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::graph::ChangeSet;
     use crate::parser::ParseSession;
+    use crate::types::Type;
 
     #[test]
     fn test_enter_exit_class_scope() {
@@ -288,4 +304,27 @@ mod tests {
         let name = extract_module_name(&module_node);
         assert_eq!(name, "Api::V1");
     }
+
+    #[test]
+    fn test_process_def_node_registers_user_method() {
+        let source = "class User; def name; \"Alice\"; end; end";
+        let session = ParseSession::new();
+        let parse_result = session.parse_source(source, "test.rb").unwrap();
+        let root = parse_result.node();
+        let program = root.as_program_node().unwrap();
+
+        let mut genv = GlobalEnv::new();
+        let mut lenv = LocalEnv::new();
+        let mut changes = ChangeSet::new();
+
+        let stmt = program.statements().body().first().unwrap();
+        let class_node = stmt.as_class_node().unwrap();
+        process_class_node(&mut genv, &mut lenv, &mut changes, source, &class_node);
+
+        // User#name should be registered as a user-defined method
+        let info = genv
+            .resolve_method(&Type::instance("User"), "name")
+            .expect("User#name should be registered");
+        assert!(info.return_vertex.is_some());
+    }
 }

rust/src/analyzer/install.rs

@@ -81,15 +81,17 @@ pub(crate) fn install_node(
     None
 }
 
-/// Process multiple statements
+/// Process multiple statements (returns last expression's VertexId)
 pub(crate) fn install_statements(
     genv: &mut GlobalEnv,
     lenv: &mut LocalEnv,
     changes: &mut ChangeSet,
     source: &str,
     statements: &ruby_prism::StatementsNode,
-) {
+) -> Option<VertexId> {
+    let mut last_vtx = None;
     for stmt in &statements.body() {
-        install_node(genv, lenv, changes, source, &stmt);
+        last_vtx = install_node(genv, lenv, changes, source, &stmt);
     }
+    last_vtx
 }

rust/src/env/global_env.rs

@@ -162,6 +162,17 @@ impl GlobalEnv {
             .register_with_block(recv_ty, method_name, ret_ty, block_param_types);
     }
 
+    /// Register a user-defined method (return type resolved via graph)
+    pub fn register_user_method(
+        &mut self,
+        recv_ty: Type,
+        method_name: &str,
+        return_vertex: VertexId,
+    ) {
+        self.method_registry
+            .register_user_method(recv_ty, method_name, return_vertex);
+    }
+
     // ===== Type Errors =====
 
     /// Record a type error (undefined method)

rust/src/env/method_registry.rs

@@ -1,5 +1,6 @@
 //! Method registration and resolution
 
+use crate::graph::VertexId;
 use crate::types::Type;
 use std::collections::HashMap;
 
@@ -8,6 +9,7 @@ use std::collections::HashMap;
 pub struct MethodInfo {
     pub return_type: Type,
     pub block_param_types: Option<Vec<Type>>,
+    pub return_vertex: Option<VertexId>,
 }
 
 /// Registry for method definitions
@@ -42,6 +44,24 @@ impl MethodRegistry {
             MethodInfo {
                 return_type: ret_ty,
                 block_param_types,
+                return_vertex: None,
+            },
+        );
+    }
+
+    /// Register a user-defined method (return type resolved via graph)
+    pub fn register_user_method(
+        &mut self,
+        recv_ty: Type,
+        method_name: &str,
+        return_vertex: VertexId,
+    ) {
+        self.methods.insert(
+            (recv_ty, method_name.to_string()),
+            MethodInfo {
+                return_type: Type::Bot,
+                block_param_types: None,
+                return_vertex: Some(return_vertex),
             },
         );
     }
@@ -87,4 +107,15 @@ mod tests {
         let registry = MethodRegistry::new();
         assert!(registry.resolve(&Type::string(), "unknown").is_none());
     }
+
+    #[test]
+    fn test_register_user_method_and_resolve() {
+        let mut registry = MethodRegistry::new();
+        let return_vtx = VertexId(42);
+        registry.register_user_method(Type::instance("User"), "name", return_vtx);
+
+        let info = registry.resolve(&Type::instance("User"), "name").unwrap();
+        assert_eq!(info.return_vertex, Some(VertexId(42)));
+        assert_eq!(info.return_type, Type::Bot);
+    }
 }

rust/src/graph/box.rs

@@ -86,11 +86,14 @@ impl BoxTrait for MethodCallBox {
         for recv_ty in recv_types {
             // Resolve method
             if let Some(method_info) = genv.resolve_method(&recv_ty, &self.method_name) {
-                // Create return type as Source
-                let ret_src_id = genv.new_source(method_info.return_type.clone());
-
-                // Add edge to return value
-                changes.add_edge(ret_src_id, self.ret);
+                if let Some(return_vtx) = method_info.return_vertex {
+                    // User-defined: edge from body's last expr → call site return
+                    changes.add_edge(return_vtx, self.ret);
+                } else {
+                    // RBS/builtin: create Source with fixed return type
+                    let ret_src_id = genv.new_source(method_info.return_type.clone());
+                    changes.add_edge(ret_src_id, self.ret);
+                }
             } else {
                 // Record type error for diagnostic reporting
                 genv.record_type_error(
@@ -488,4 +491,36 @@ mod tests {
         assert_eq!(genv.get_vertex(key_vtx).unwrap().show(), "String");
         assert_eq!(genv.get_vertex(value_vtx).unwrap().show(), "Integer");
     }
+
+    #[test]
+    fn test_method_call_box_user_defined_method() {
+        let mut genv = GlobalEnv::new();
+
+        // Simulate: def name; "Alice"; end
+        let body_src = genv.new_source(Type::string());
+
+        // Register user-defined method User#name with return_vertex
+        genv.register_user_method(Type::instance("User"), "name", body_src);
+
+        // Simulate: user.name (receiver has type User)
+        let recv_vtx = genv.new_vertex();
+        let recv_src = genv.new_source(Type::instance("User"));
+        genv.add_edge(recv_src, recv_vtx);
+
+        let ret_vtx = genv.new_vertex();
+        let box_id = genv.alloc_box_id();
+        let call_box = MethodCallBox::new(
+            box_id,
+            recv_vtx,
+            "name".to_string(),
+            ret_vtx,
+            None,
+        );
+        genv.register_box(box_id, Box::new(call_box));
+
+        genv.run_all();
+
+        // Return type should be String (propagated from body's last expression)
+        assert_eq!(genv.get_vertex(ret_vtx).unwrap().show(), "String");
+    }
 }

test/method_return_type_test.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+class MethodReturnTypeTest < Minitest::Test
+  include CLITestHelper
+
+  # ============================================
+  # No Error
+  # ============================================
+
+  def test_user_method_returns_string
+    source = <<~RUBY
+      class User
+        def name
+          "Alice"
+        end
+
+        def greet
+          self.name.upcase
+        end
+      end
+    RUBY
+
+    assert_no_check_errors(source)
+  end
+
+  def test_user_method_returns_integer
+    source = <<~RUBY
+      class Calculator
+        def answer
+          42
+        end
+
+        def compute
+          self.answer.even?
+        end
+      end
+    RUBY
+
+    assert_no_check_errors(source)
+  end
+
+  def test_user_method_chain
+    source = <<~RUBY
+      class Formatter
+        def title
+          "hello"
+        end
+
+        def format
+          self.title.upcase.downcase
+        end
+      end
+    RUBY
+
+    assert_no_check_errors(source)
+  end
+
+  # ============================================
+  # Error Detection
+  # ============================================
+
+  def test_user_method_return_type_error
+    source = <<~RUBY
+      class User
+        def name
+          "Alice"
+        end
+
+        def greet
+          self.name.even?
+        end
+      end
+    RUBY
+
+    assert_check_error(source, method_name: 'even?', receiver_type: 'String')
+  end
+
+  def test_user_method_returns_integer_type_error
+    source = <<~RUBY
+      class Calculator
+        def answer
+          42
+        end
+
+        def compute
+          self.answer.upcase
+        end
+      end
+    RUBY
+
+    assert_check_error(source, method_name: 'upcase', receiver_type: 'Integer')
+  end
+end