From 6f0e4a0a74e23308d89f9d8ea836c96544bd4b0b Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sun, 12 Apr 2026 10:52:06 +0000
Subject: [PATCH] Analyzer: validate return statement types

Tracks the current function's declared return type and checks that
`return value` statements match it. Returning from a void function
with a value produces E0203; returning a value with the wrong type
produces E0201. State handler `return value` is still silently
accepted (the value is discarded).

https://claude.ai/code/session_01W6eQFStA66EuMKHUFo2rx3
---
 src/analyzer/mod.rs   | 33 +++++++++++++++++++++++++++++---
 src/analyzer/tests.rs | 44 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/src/analyzer/mod.rs b/src/analyzer/mod.rs
index 903ad5f..f8a6d92 100644
--- a/src/analyzer/mod.rs
+++ b/src/analyzer/mod.rs
@@ -50,6 +50,8 @@ pub fn analyze(program: &Program) -> AnalysisResult {
         in_loop: false,
         used_vars: HashSet::new(),
         function_signatures: HashMap::new(),
+        current_return_type: None,
+        in_function_body: false,
     };
     analyzer.analyze_program(program);
 
@@ -78,6 +80,15 @@ struct Analyzer {
     /// Function name to parameter types (in order). Used to validate
     /// call arity and argument types.
     function_signatures: HashMap<String, Vec<NesType>>,
+    /// Return type of the function currently being analyzed, or None
+    /// when the function has no declared return type. Only meaningful
+    /// when `in_function_body` is true.
+    current_return_type: Option<NesType>,
+    /// True while analyzing a function body (as opposed to a state
+    /// handler's `on_enter` / `on_exit` / `on_frame` block). Used to
+    /// distinguish "void function" from "state handler" when checking
+    /// `return value` statements.
+    in_function_body: bool,
 }
 
 impl Analyzer {
@@ -153,7 +164,11 @@ impl Analyzer {
                 }
                 self.mark_var_used(&param.name);
             }
+            self.current_return_type.clone_from(&fun.return_type);
+            self.in_function_body = true;
             self.check_block(&fun.body, &state_names);
+            self.current_return_type = None;
+            self.in_function_body = false;
             for name in &added_params {
                 self.symbols.remove(name);
             }
@@ -571,10 +586,22 @@ impl Analyzer {
                     self.check_expr_type(frame, &NesType::U8);
                 }
             }
-            Statement::Return(Some(expr), _) => {
-                // For M1, just validate the expression without checking return type
+            Statement::Return(Some(expr), span) => {
                 self.walk_expr_reads(expr);
-                let _ = self.infer_type(expr);
+                if let Some(ret_ty) = self.current_return_type.clone() {
+                    // Function with declared return type — check the value.
+                    self.check_expr_type(expr, &ret_ty);
+                } else if self.in_function_body {
+                    // Function with no declared return type ("void"),
+                    // but the return statement has a value.
+                    self.diagnostics.push(Diagnostic::error(
+                        ErrorCode::E0203,
+                        "return value in function with no declared return type",
+                        *span,
+                    ));
+                }
+                // State handlers (`in_function_body == false`) accept
+                // `return value` silently — the value is simply discarded.
             }
             Statement::Call(name, args, span) => {
                 if self.symbols.contains_key(name) {
diff --git a/src/analyzer/tests.rs b/src/analyzer/tests.rs
index 70a3f82..08c0e79 100644
--- a/src/analyzer/tests.rs
+++ b/src/analyzer/tests.rs
@@ -270,6 +270,50 @@ fn analyze_call_arity_in_expr_context() {
     );
 }
 
+#[test]
+fn analyze_return_type_ok() {
+    analyze_ok(
+        r#"
+        game "Test" { mapper: NROM }
+        fun get_five() -> u8 { return 5 }
+        on frame { wait_frame }
+        start Main
+    "#,
+    );
+}
+
+#[test]
+fn analyze_return_wrong_type() {
+    let errors = analyze_errors(
+        r#"
+        game "Test" { mapper: NROM }
+        fun is_ok() -> bool { return 5 }
+        on frame { wait_frame }
+        start Main
+    "#,
+    );
+    assert!(
+        errors.contains(&ErrorCode::E0201),
+        "returning wrong type should produce E0201, got: {errors:?}"
+    );
+}
+
+#[test]
+fn analyze_return_value_from_void_fn() {
+    let errors = analyze_errors(
+        r#"
+        game "Test" { mapper: NROM }
+        fun do_nothing() { return 5 }
+        on frame { wait_frame }
+        start Main
+    "#,
+    );
+    assert!(
+        errors.contains(&ErrorCode::E0203),
+        "returning value from void function should produce E0203, got: {errors:?}"
+    );
+}
+
 #[test]
 fn analyze_const_assignment_error() {
     let errors = analyze_errors(