From 1187e1367c5b520023b1e7c1cd02523977e4af74 Mon Sep 17 00:00:00 2001
From: Felipe Pena <felipensp@gmail.com>
Date: Tue, 17 Sep 2024 14:13:07 -0300
Subject: [PATCH] all: add comptime support for traversing the method
 parameters with `$for arg in method.params {` (#22229)

---
 vlib/builtin/builtin.v                        |  6 ++--
 vlib/v/ast/ast.v                              |  1 +
 vlib/v/ast/str.v                              |  1 +
 vlib/v/checker/comptime.v                     |  8 ++++-
 vlib/v/checker/if.v                           |  4 +--
 vlib/v/comptime/comptimeinfo.v                | 16 ++++++---
 vlib/v/gen/c/comptime.v                       | 36 ++++++++++++++-----
 vlib/v/gen/c/fn.v                             | 10 +++---
 vlib/v/gen/c/str_intp.v                       |  2 +-
 vlib/v/parser/comptime.v                      | 10 +++++-
 .../tests/comptime_unknown_meta_kind_err.out  |  2 +-
 .../comptime/comptime_for_method_param_test.v | 30 ++++++++++++++++
 12 files changed, 100 insertions(+), 26 deletions(-)
 create mode 100644 vlib/v/tests/comptime/comptime_for_method_param_test.v

diff --git a/vlib/builtin/builtin.v b/vlib/builtin/builtin.v
index 264cb5b2fa9ebb..9f7a2d30799510 100644
--- a/vlib/builtin/builtin.v
+++ b/vlib/builtin/builtin.v
@@ -90,8 +90,8 @@ fn __print_assert_failure(i &VAssertMetaInfo) {
 	}
 }
 
-// MethodArgs holds type information for function and/or method arguments.
-pub struct MethodArgs {
+// MethodParam holds type information for function and/or method arguments.
+pub struct MethodParam {
 pub:
 	typ  int
 	name string
@@ -102,7 +102,7 @@ pub struct FunctionData {
 pub:
 	name        string
 	attrs       []string
-	args        []MethodArgs
+	args        []MethodParam
 	return_type int
 	typ         int
 }
diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index b4cd5ce81c2efe..bb3f8826190a99 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -1250,6 +1250,7 @@ pub enum ComptimeForKind {
 	attributes
 	values
 	variants
+	params
 }
 
 pub struct ComptimeFor {
diff --git a/vlib/v/ast/str.v b/vlib/v/ast/str.v
index 89dbaf45b95e65..cf7aa132fa9e1c 100644
--- a/vlib/v/ast/str.v
+++ b/vlib/v/ast/str.v
@@ -896,5 +896,6 @@ pub fn (e ComptimeForKind) str() string {
 		.attributes { return 'attributes' }
 		.values { return 'values' }
 		.variants { return 'variants' }
+		.params { return 'params' }
 	}
 }
diff --git a/vlib/v/checker/comptime.v b/vlib/v/checker/comptime.v
index 6b0fb8d5693d89..286c35d34abfaa 100644
--- a/vlib/v/checker/comptime.v
+++ b/vlib/v/checker/comptime.v
@@ -318,13 +318,19 @@ fn (mut c Checker) comptime_for(mut node ast.ComptimeFor) {
 		for method in methods {
 			c.push_new_comptime_info()
 			c.comptime.inside_comptime_for = true
-			c.comptime.comptime_for_method = method.name
+			c.comptime.comptime_for_method = unsafe { &method }
 			c.comptime.comptime_for_method_var = node.val_var
 			c.comptime.comptime_for_method_ret_type = method.return_type
 			c.comptime.type_map['${node.val_var}.return_type'] = method.return_type
 			c.stmts(mut node.stmts)
 			c.pop_comptime_info()
 		}
+	} else if node.kind == .params {
+		c.push_new_comptime_info()
+		c.comptime.inside_comptime_for = true
+		c.comptime.comptime_for_method_param_var = node.val_var
+		c.stmts(mut node.stmts)
+		c.pop_comptime_info()
 	} else if node.kind == .variants {
 		if c.variant_data_type == 0 {
 			c.variant_data_type = ast.idx_to_type(c.table.find_type_idx('VariantData'))
diff --git a/vlib/v/checker/if.v b/vlib/v/checker/if.v
index ab487f41efc666..c45755ac8ccb63 100644
--- a/vlib/v/checker/if.v
+++ b/vlib/v/checker/if.v
@@ -284,14 +284,14 @@ fn (mut c Checker) if_expr(mut node ast.IfExpr) ast.Type {
 									is_comptime_type_is_expr = true
 									match branch.cond.op {
 										.eq {
-											skip_state = if c.comptime.comptime_for_method == right.val.str() {
+											skip_state = if c.comptime.comptime_for_method.name == right.val.str() {
 												ComptimeBranchSkipState.eval
 											} else {
 												ComptimeBranchSkipState.skip
 											}
 										}
 										.ne {
-											skip_state = if c.comptime.comptime_for_method == right.val.str() {
+											skip_state = if c.comptime.comptime_for_method.name == right.val.str() {
 												ComptimeBranchSkipState.skip
 											} else {
 												ComptimeBranchSkipState.eval
diff --git a/vlib/v/comptime/comptimeinfo.v b/vlib/v/comptime/comptimeinfo.v
index efc7e9d249e01e..3d340a8277bded 100644
--- a/vlib/v/comptime/comptimeinfo.v
+++ b/vlib/v/comptime/comptimeinfo.v
@@ -95,7 +95,7 @@ pub fn (mut ct ComptimeInfo) get_comptime_var_type(node ast.Expr) ast.Type {
 	} else if node is ast.SelectorExpr && ct.is_comptime_selector_type(node) {
 		return ct.get_type_from_comptime_var(node.expr as ast.Ident)
 	} else if node is ast.ComptimeCall {
-		method_name := ct.comptime_for_method
+		method_name := ct.comptime_for_method.name
 		left_sym := ct.table.sym(ct.resolver.unwrap_generic(node.left_type))
 		f := left_sym.find_method(method_name) or {
 			ct.error('could not find method `${method_name}` on compile-time resolution',
@@ -114,6 +114,9 @@ pub fn (mut ct ComptimeInfo) get_type_from_comptime_var(var ast.Ident) ast.Type
 		ct.comptime_for_variant_var {
 			ct.type_map['${ct.comptime_for_variant_var}.typ']
 		}
+		ct.comptime_for_method_param_var {
+			ct.type_map['${ct.comptime_for_method_param_var}.typ']
+		}
 		else {
 			// field_var.typ from $for field
 			ct.comptime_for_field_type
@@ -152,7 +155,7 @@ pub fn (mut ct ComptimeInfo) is_comptime_selector_field_name(node ast.SelectorEx
 pub fn (mut ct ComptimeInfo) is_comptime_selector_type(node ast.SelectorExpr) bool {
 	if ct.inside_comptime_for && node.expr is ast.Ident {
 		return
-			node.expr.name in [ct.comptime_for_enum_var, ct.comptime_for_variant_var, ct.comptime_for_field_var]
+			node.expr.name in [ct.comptime_for_enum_var, ct.comptime_for_variant_var, ct.comptime_for_field_var, ct.comptime_for_method_param_var]
 			&& node.field_name == 'typ'
 	}
 	return false
@@ -251,7 +254,7 @@ pub fn (mut ct ComptimeInfo) is_comptime_type(x ast.Type, y ast.ComptimeType) bo
 	}
 }
 
-// comptime_get_kind_var identifies the comptime variable kind (i.e. if it is about .values, .fields, .methods  etc)
+// comptime_get_kind_var identifies the comptime variable kind (i.e. if it is about .values, .fields, .methods, .args etc)
 fn (mut ct ComptimeInfo) comptime_get_kind_var(var ast.Ident) ?ast.ComptimeForKind {
 	if ct.inside_comptime_for {
 		return none
@@ -273,6 +276,9 @@ fn (mut ct ComptimeInfo) comptime_get_kind_var(var ast.Ident) ?ast.ComptimeForKi
 		ct.comptime_for_attr_var {
 			return .attributes
 		}
+		ct.comptime_for_method_param_var {
+			return .params
+		}
 		else {
 			return none
 		}
@@ -339,6 +345,8 @@ pub mut:
 	comptime_for_attr_var string
 	// .methods
 	comptime_for_method_var      string
-	comptime_for_method          string
+	comptime_for_method          &ast.Fn = unsafe { nil }
 	comptime_for_method_ret_type ast.Type
+	// .args
+	comptime_for_method_param_var string
 }
diff --git a/vlib/v/gen/c/comptime.v b/vlib/v/gen/c/comptime.v
index f29dc7825ca147..3539aa6e7ecac5 100644
--- a/vlib/v/gen/c/comptime.v
+++ b/vlib/v/gen/c/comptime.v
@@ -139,7 +139,7 @@ fn (mut g Gen) comptime_call(mut node ast.ComptimeCall) {
 	g.trace_autofree('// \$method call. sym="${sym.name}"')
 	if node.method_name == 'method' {
 		// `app.$method()`
-		m := sym.find_method(g.comptime.comptime_for_method) or { return }
+		m := sym.find_method(g.comptime.comptime_for_method.name) or { return }
 		/*
 		vals := m.attrs[0].split('/')
 		args := vals.filter(it.starts_with(':')).map(it[1..])
@@ -199,7 +199,7 @@ fn (mut g Gen) comptime_call(mut node ast.ComptimeCall) {
 			g.write('(*(${g.base_type(m.return_type)}*)')
 		}
 		// TODO: check argument types
-		g.write('${util.no_dots(sym.name)}_${g.comptime.comptime_for_method}(')
+		g.write('${util.no_dots(sym.name)}_${g.comptime.comptime_for_method.name}(')
 
 		// try to see if we need to pass a pointer
 		if mut node.left is ast.Ident {
@@ -583,15 +583,15 @@ fn (mut g Gen) comptime_if_cond(cond ast.Expr, pkg_exist bool) (bool, bool) {
 				.eq, .ne {
 					// TODO: Implement `$if method.args.len == 1`
 					if cond.left is ast.SelectorExpr && (g.comptime.comptime_for_field_var.len > 0
-						|| g.comptime.comptime_for_method.len > 0) {
+						|| g.comptime.comptime_for_method != unsafe { nil }) {
 						if cond.right is ast.StringLiteral {
 							if cond.left.expr is ast.Ident && cond.left.field_name == 'name' {
 								if g.comptime.comptime_for_method_var.len > 0
 									&& cond.left.expr.name == g.comptime.comptime_for_method_var {
 									is_true := if cond.op == .eq {
-										g.comptime.comptime_for_method == cond.right.val
+										g.comptime.comptime_for_method.name == cond.right.val
 									} else {
-										g.comptime.comptime_for_method != cond.right.val
+										g.comptime.comptime_for_method.name != cond.right.val
 									}
 									if is_true {
 										g.write('1')
@@ -828,7 +828,7 @@ fn (mut g Gen) comptime_for(node ast.ComptimeFor) {
 					}
 				}
 			}
-			g.comptime.comptime_for_method = method.name
+			g.comptime.comptime_for_method = unsafe { &method }
 			g.comptime.comptime_for_method_var = node.val_var
 			g.writeln('/* method ${i} */ {')
 			g.writeln('\t${node.val_var}.name = _SLIT("${method.name}");')
@@ -842,10 +842,10 @@ fn (mut g Gen) comptime_for(node ast.ComptimeFor) {
 			}
 			if method.params.len < 2 {
 				// 0 or 1 (the receiver) args
-				g.writeln('\t${node.val_var}.args = __new_array_with_default(0, 0, sizeof(MethodArgs), 0);')
+				g.writeln('\t${node.val_var}.args = __new_array_with_default(0, 0, sizeof(MethodParam), 0);')
 			} else {
 				len := method.params.len - 1
-				g.write('\t${node.val_var}.args = new_array_from_c_array(${len}, ${len}, sizeof(MethodArgs), _MOV((MethodArgs[${len}]){')
+				g.write('\t${node.val_var}.args = new_array_from_c_array(${len}, ${len}, sizeof(MethodParam), _MOV((MethodParam[${len}]){')
 				// Skip receiver arg
 				for j, arg in method.params[1..] {
 					typ := arg.typ.idx()
@@ -1020,6 +1020,26 @@ fn (mut g Gen) comptime_for(node ast.ComptimeFor) {
 			}
 			g.pop_comptime_info()
 		}
+	} else if node.kind == .params {
+		method := g.comptime.comptime_for_method
+
+		if method.params.len > 0 {
+			g.writeln('\tMethodParam ${node.val_var} = {0};')
+		}
+		g.push_new_comptime_info()
+		g.comptime.inside_comptime_for = true
+		g.comptime.comptime_for_method_param_var = node.val_var
+		for param in method.params[1..] {
+			g.comptime.type_map['${node.val_var}.typ'] = param.typ
+
+			g.writeln('/* method param ${i} */ {')
+			g.writeln('\t${node.val_var}.typ = ${int(param.typ)};')
+			g.writeln('\t${node.val_var}.name = _SLIT("${param.name}");')
+			g.stmts(node.stmts)
+			g.writeln('}')
+			i++
+		}
+		g.pop_comptime_info()
 	}
 	g.indent--
 	g.writeln('}// \$for')
diff --git a/vlib/v/gen/c/fn.v b/vlib/v/gen/c/fn.v
index a4594c959bf33d..4554d739bf251f 100644
--- a/vlib/v/gen/c/fn.v
+++ b/vlib/v/gen/c/fn.v
@@ -1122,7 +1122,7 @@ fn (mut g Gen) gen_to_str_method_call(node ast.CallExpr) bool {
 	} else if left_node is ast.ComptimeCall {
 		if left_node.method_name == 'method' {
 			sym := g.table.sym(g.unwrap_generic(left_node.left_type))
-			if m := sym.find_method(g.comptime.comptime_for_method) {
+			if m := sym.find_method(g.comptime.comptime_for_method.name) {
 				rec_type = m.return_type
 				g.gen_expr_to_string(left_node, rec_type)
 				return true
@@ -1436,7 +1436,7 @@ fn (mut g Gen) resolve_comptime_args(func ast.Fn, mut node_ ast.CallExpr, concre
 				if call_arg.expr.method_name == 'method' {
 					sym := g.table.sym(g.unwrap_generic(call_arg.expr.left_type))
 					// `app.$method()`
-					if m := sym.find_method(g.comptime.comptime_for_method) {
+					if m := sym.find_method(g.comptime.comptime_for_method.name) {
 						comptime_args[k] = m.return_type
 					}
 				}
@@ -2043,7 +2043,7 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 	// Handle `print(x)`
 	mut print_auto_str := false
 	if is_print && (node.args[0].typ != ast.string_type
-		|| g.comptime.comptime_for_method.len > 0
+		|| g.comptime.comptime_for_method != unsafe { nil }
 		|| g.comptime.is_comptime_var(node.args[0].expr)) {
 		g.inside_interface_deref = true
 		defer {
@@ -2059,7 +2059,7 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 		if typ == 0 {
 			g.checker_bug('print arg.typ is 0', node.pos)
 		}
-		if typ != ast.string_type || g.comptime.comptime_for_method.len > 0 {
+		if typ != ast.string_type || g.comptime.comptime_for_method != unsafe { nil } {
 			expr := node.args[0].expr
 			typ_sym := g.table.sym(typ)
 			if typ_sym.kind == .interface_ && (typ_sym.info as ast.Interface).defines_method('str') {
@@ -2090,7 +2090,7 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 				} else if expr is ast.ComptimeCall {
 					if expr.method_name == 'method' {
 						sym := g.table.sym(g.unwrap_generic(expr.left_type))
-						if m := sym.find_method(g.comptime.comptime_for_method) {
+						if m := sym.find_method(g.comptime.comptime_for_method.name) {
 							typ = m.return_type
 						}
 					}
diff --git a/vlib/v/gen/c/str_intp.v b/vlib/v/gen/c/str_intp.v
index cfe5ea970be21d..4f1e894cde33de 100644
--- a/vlib/v/gen/c/str_intp.v
+++ b/vlib/v/gen/c/str_intp.v
@@ -165,7 +165,7 @@ fn (mut g Gen) str_val(node ast.StringInterLiteral, i int, fmts []u8) {
 	fmt := fmts[i]
 	typ := g.unwrap_generic(node.expr_types[i])
 	typ_sym := g.table.sym(typ)
-	if typ == ast.string_type && g.comptime.comptime_for_method.len == 0 {
+	if typ == ast.string_type && g.comptime.comptime_for_method == unsafe { nil } {
 		if g.inside_vweb_tmpl {
 			g.write('${g.vweb_filter_fn_name}(')
 			if expr.is_auto_deref_var() && fmt != `p` {
diff --git a/vlib/v/parser/comptime.v b/vlib/v/parser/comptime.v
index a6f7b68f77a18d..15e74b3508ae43 100644
--- a/vlib/v/parser/comptime.v
+++ b/vlib/v/parser/comptime.v
@@ -363,6 +363,14 @@ fn (mut p Parser) comptime_for() ast.ComptimeFor {
 		p.close_scope()
 	}
 	match for_val {
+		'params' {
+			p.scope.register(ast.Var{
+				name: val_var
+				typ:  p.table.find_type_idx('MethodParam')
+				pos:  var_pos
+			})
+			kind = .params
+		}
 		'methods' {
 			p.scope.register(ast.Var{
 				name: val_var
@@ -403,7 +411,7 @@ fn (mut p Parser) comptime_for() ast.ComptimeFor {
 			kind = .attributes
 		}
 		else {
-			p.error_with_pos('unknown kind `${for_val}`, available are: `methods`, `fields`, `values`, `variants` or `attributes`',
+			p.error_with_pos('unknown kind `${for_val}`, available are: `methods`, `fields`, `values`, `variants`, `attributes` or `params`',
 				p.prev_tok.pos())
 			return ast.ComptimeFor{}
 		}
diff --git a/vlib/v/parser/tests/comptime_unknown_meta_kind_err.out b/vlib/v/parser/tests/comptime_unknown_meta_kind_err.out
index d6d77199239b7f..b50cdb7c9bf9d8 100644
--- a/vlib/v/parser/tests/comptime_unknown_meta_kind_err.out
+++ b/vlib/v/parser/tests/comptime_unknown_meta_kind_err.out
@@ -1,4 +1,4 @@
-vlib/v/parser/tests/comptime_unknown_meta_kind_err.vv:7:20: error: unknown kind `abcde`, available are: `methods`, `fields`, `values`, `variants` or `attributes`
+vlib/v/parser/tests/comptime_unknown_meta_kind_err.vv:7:20: error: unknown kind `abcde`, available are: `methods`, `fields`, `values`, `variants`, `attributes` or `params`
     5 | 
     6 | fn test_main() {
     7 |     $for item in Test.abcde {
diff --git a/vlib/v/tests/comptime/comptime_for_method_param_test.v b/vlib/v/tests/comptime/comptime_for_method_param_test.v
new file mode 100644
index 00000000000000..d01c43b625868f
--- /dev/null
+++ b/vlib/v/tests/comptime/comptime_for_method_param_test.v
@@ -0,0 +1,30 @@
+module main
+
+struct Struct1 {
+	num f64
+}
+
+fn (mut s Struct1) do_thing(a f32, b voidptr) bool {
+	return false
+}
+
+fn register[T]() []string {
+	mut args := []string{}
+	$for method in T.methods {
+		$for arg in method.params {
+			$if arg.typ is f32 {
+				args << 'f32: ${arg.name} ${typeof(arg.typ).name}'
+			} $else $if arg.typ is voidptr {
+				args << '&void: ${arg.name} ${typeof(arg.typ).name}'
+			} $else {
+			}
+		}
+	}
+	return args
+}
+
+pub fn test_main() {
+	args := register[Struct1]()
+	assert args[0] == 'f32: a f32'
+	assert args[1] == '&void: b voidptr'
+}