From 8aef09451abbb31dd16ef66c89421d4573c65283 Mon Sep 17 00:00:00 2001 From: dak2 Date: Thu, 12 Feb 2026 08:46:58 +0900 Subject: [PATCH] Add method return type inference for user-defined methods Register the last expression's VertexId from def body as return_vertex in MethodInfo, and propagate types via edge connection at call sites. --- rust/src/analyzer/definitions.rs | 43 ++++++++++++++- rust/src/analyzer/install.rs | 8 ++- rust/src/env/global_env.rs | 11 ++++ rust/src/env/method_registry.rs | 31 +++++++++++ rust/src/graph/box.rs | 45 +++++++++++++-- test/method_return_type_test.rb | 95 ++++++++++++++++++++++++++++++++ 6 files changed, 223 insertions(+), 10 deletions(-) create mode 100644 test/method_return_type_test.rb diff --git a/rust/src/analyzer/definitions.rs b/rust/src/analyzer/definitions.rs index 2ce5408..58996a3 100644 --- a/rust/src/analyzer/definitions.rs +++ b/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 { 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, ¶ms_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 { #[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()); + } } diff --git a/rust/src/analyzer/install.rs b/rust/src/analyzer/install.rs index 57a771b..c9aa0e2 100644 --- a/rust/src/analyzer/install.rs +++ b/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 { + 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 } diff --git a/rust/src/env/global_env.rs b/rust/src/env/global_env.rs index d7e0908..cfba3f9 100644 --- a/rust/src/env/global_env.rs +++ b/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) diff --git a/rust/src/env/method_registry.rs b/rust/src/env/method_registry.rs index 65cf9f6..b01ad06 100644 --- a/rust/src/env/method_registry.rs +++ b/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>, + pub return_vertex: Option, } /// 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); + } } diff --git a/rust/src/graph/box.rs b/rust/src/graph/box.rs index d2d76eb..344185a 100644 --- a/rust/src/graph/box.rs +++ b/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"); + } } diff --git a/test/method_return_type_test.rb b/test/method_return_type_test.rb new file mode 100644 index 0000000..dbbf165 --- /dev/null +++ b/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