From bd6005ed736926fcdfb3519a606a7030d81c3cd8 Mon Sep 17 00:00:00 2001
From: Felipe Pena <felipensp@gmail.com>
Date: Tue, 5 Nov 2024 10:45:03 -0300
Subject: [PATCH] checker: fix missing check for stack pointer return (fix
 #22726) (#22756)

* fix

* test

* fix
---
 vlib/v/checker/checker.v                      |  9 ++++--
 vlib/v/checker/return.v                       |  5 ++++
 vlib/v/checker/tests/return_stack_var_err.out |  7 +++++
 vlib/v/checker/tests/return_stack_var_err.vv  | 30 +++++++++++++++++++
 4 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 vlib/v/checker/tests/return_stack_var_err.out
 create mode 100644 vlib/v/checker/tests/return_stack_var_err.vv

diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 23c04857d43876..b37408d5b37d79 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -5310,11 +5310,14 @@ fn (mut c Checker) fail_if_stack_struct_action_outside_unsafe(mut ident ast.Iden
 		}
 		if obj.is_stack_obj && !c.inside_unsafe {
 			sym := c.table.sym(obj.typ.set_nr_muls(0))
-			if !sym.is_heap() && !c.pref.translated && !c.file.is_translated {
-				suggestion := if sym.kind == .struct {
+			is_heap := sym.is_heap()
+			if (!is_heap || !obj.typ.is_ptr()) && !c.pref.translated && !c.file.is_translated {
+				suggestion := if !is_heap && sym.kind == .struct {
 					'declaring `${sym.name}` as `@[heap]`'
-				} else {
+				} else if !is_heap {
 					'wrapping the `${sym.name}` object in a `struct` declared as `@[heap]`'
+				} else { // e.g. var from `for a in heap_object {`
+					'declaring `${ident.name}` mutable'
 				}
 				c.error('`${ident.name}` cannot be ${failed_action} outside `unsafe` blocks as it might refer to an object stored on stack. Consider ${suggestion}.',
 					ident.pos)
diff --git a/vlib/v/checker/return.v b/vlib/v/checker/return.v
index 6dd9a47c2277f6..d000dad20fd1ca 100644
--- a/vlib/v/checker/return.v
+++ b/vlib/v/checker/return.v
@@ -292,6 +292,11 @@ fn (mut c Checker) return_stmt(mut node ast.Return) {
 			mut r_expr := &node.exprs[expr_idxs[i]]
 			if mut r_expr is ast.Ident {
 				c.fail_if_stack_struct_action_outside_unsafe(mut r_expr, 'returned')
+			} else if mut r_expr is ast.PrefixExpr && r_expr.op == .amp {
+				// &var
+				if mut r_expr.right is ast.Ident {
+					c.fail_if_stack_struct_action_outside_unsafe(mut r_expr.right, 'returned')
+				}
 			}
 		}
 	}
diff --git a/vlib/v/checker/tests/return_stack_var_err.out b/vlib/v/checker/tests/return_stack_var_err.out
new file mode 100644
index 00000000000000..bc46b5b772047e
--- /dev/null
+++ b/vlib/v/checker/tests/return_stack_var_err.out
@@ -0,0 +1,7 @@
+vlib/v/checker/tests/return_stack_var_err.vv:26:12: error: `s` cannot be returned outside `unsafe` blocks as it might refer to an object stored on stack. Consider declaring `s` mutable.
+   24 |     for s in students {
+   25 |         if s.name == 'Amy' {
+   26 |             return &s
+      |                     ^
+   27 |         }
+   28 |     }
diff --git a/vlib/v/checker/tests/return_stack_var_err.vv b/vlib/v/checker/tests/return_stack_var_err.vv
new file mode 100644
index 00000000000000..514a270a8f5b58
--- /dev/null
+++ b/vlib/v/checker/tests/return_stack_var_err.vv
@@ -0,0 +1,30 @@
+module main
+
+@[heap]
+struct Student {
+	name string
+	age  int
+}
+
+fn main() {
+	mut students := []Student{}
+	students << Student{
+		name: 'Mike'
+		age:  16
+	}
+	students << Student{
+		name: 'Amy'
+		age:  15
+	}
+
+	_ := get_amy(students) or { unsafe { nil } }
+}
+
+fn get_amy(students []Student) ?&Student {
+	for s in students {
+		if s.name == 'Amy' {
+			return &s
+		}
+	}
+	return none
+}