checker,cgen: fix comptime value recognition on generic arguments (#16768)

vlib/v/ast/ast.v

@@ -590,6 +590,7 @@ pub mut:
 	is_ctor_new        bool // if JS ctor calls requires `new` before call, marked as `[use_new]` in V
 	args               []CallArg
 	expected_arg_types []Type
+	comptime_ret_val   bool
 	language           Language
 	or_block           OrExpr
 	left               Expr // `user` in `user.register()`

vlib/v/checker/check_types.v

@@ -871,7 +871,25 @@ fn (mut c Checker) infer_fn_generic_types(func ast.Fn, mut node ast.CallExpr) {
 					func_.name = ''
 					idx := c.table.find_or_register_fn_type(func_, true, false)
 					typ = ast.new_type(idx).derive(arg.typ)
+				} else if c.inside_comptime_for_field && sym.info is ast.Struct
+					&& arg.expr is ast.ComptimeSelector {
+					compselector := arg.expr as ast.ComptimeSelector
+					if compselector.field_expr is ast.SelectorExpr {
+						selectorexpr := compselector.field_expr as ast.SelectorExpr
+						if selectorexpr.expr is ast.Ident {
+							ident := selectorexpr.expr as ast.Ident
+							if ident.name == c.comptime_for_field_var {
+								typ = c.comptime_fields_default_type
+
+								if func.return_type.has_flag(.generic)
+									&& gt_name == c.table.type_to_str(func.return_type) {
+									node.comptime_ret_val = true
+								}
+							}
+						}
+					}
 				}
+
 				if arg.expr.is_auto_deref_var() {
 					typ = typ.deref()
 				}

vlib/v/checker/checker.v

@@ -95,6 +95,7 @@ mut:
 	vweb_gen_types                   []ast.Type // vweb route checks
 	timers                           &util.Timers = util.get_timers()
 	for_in_any_val_type              ast.Type
+	comptime_for_field_var           string
 	comptime_fields_default_type     ast.Type
 	comptime_fields_type             map[string]ast.Type
 	fn_scope                         &ast.Scope = unsafe { nil }

vlib/v/checker/comptime.v

@@ -180,6 +180,7 @@ fn (mut c Checker) comptime_for(node ast.ComptimeFor) {
 			sym_info := sym.info as ast.Struct
 			c.inside_comptime_for_field = true
 			for field in sym_info.fields {
+				c.comptime_for_field_var = node.val_var
 				c.comptime_fields_type[node.val_var] = node.typ
 				c.comptime_fields_default_type = field.typ
 				c.stmts(node.stmts)
@@ -188,6 +189,7 @@ fn (mut c Checker) comptime_for(node ast.ComptimeFor) {
 				tsym := c.table.sym(unwrapped_expr_type)
 				c.table.dumps[int(unwrapped_expr_type.clear_flag(.optional).clear_flag(.result))] = tsym.cname
 			}
+			c.comptime_for_field_var = ''
 			c.inside_comptime_for_field = false
 		}
 	} else {

vlib/v/gen/c/assign.v

@@ -136,7 +136,7 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 				is_auto_heap = left.obj.is_auto_heap
 			}
 		}
-		styp := g.typ(var_type)
+		mut styp := g.typ(var_type)
 		mut is_fixed_array_init := false
 		mut has_val := false
 		match val {
@@ -146,7 +146,12 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 			}
 			ast.CallExpr {
 				is_call = true
-				return_type = val.return_type
+				if val.comptime_ret_val {
+					return_type = g.comptime_for_field_type
+					styp = g.typ(return_type)
+				} else {
+					return_type = val.return_type
+				}
 			}
 			// TODO: no buffer fiddling
 			ast.AnonFn {

vlib/v/gen/c/fn.v

@@ -1234,6 +1234,7 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 	mut is_interface_call := false
 	mut is_selector_call := false
 	mut has_comptime_field := false
+	mut comptime_args := []int{}
 	if node.left_type != 0 {
 		left_sym := g.table.sym(node.left_type)
 		if left_sym.kind == .interface_ {
@@ -1266,10 +1267,12 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 					node_.args[i].typ = call_arg.expr.obj.typ
 					if call_arg.expr.obj.is_comptime_field {
 						has_comptime_field = true
+						comptime_args << i
 					}
 				}
 			} else if mut call_arg.expr is ast.ComptimeSelector {
 				has_comptime_field = true
+				comptime_args << i
 			}
 		}
 	}
@@ -1362,7 +1365,11 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 			if func.generic_names.len > 0 {
 				if g.comptime_for_field_type != 0 && g.inside_comptime_for_field
 					&& has_comptime_field {
-					name = g.generic_fn_name([g.comptime_for_field_type], name)
+					mut concrete_types := node.concrete_types.map(g.unwrap_generic(it))
+					for k in comptime_args {
+						concrete_types[k] = g.comptime_for_field_type
+					}
+					name = g.generic_fn_name(concrete_types, name)
 				} else {
 					concrete_types := node.concrete_types.map(g.unwrap_generic(it))
 					name = g.generic_fn_name(concrete_types, name)

vlib/v/tests/comptime_on_generics_func_test.v

@@ -0,0 +1,24 @@
+module main
+
+struct MyStruct {
+	b string
+	c bool
+}
+
+fn get_type[T](t T) T {
+	return t
+}
+
+fn test_main() {
+	my := MyStruct{'cool', false}
+	$for field2 in MyStruct.fields {
+		$if field2.typ is string {
+			var := get_type(my.$(field2.name))
+			assert dump(var) == 'cool'
+		}
+		$if field2.typ is bool {
+			var := get_type(my.$(field2.name))
+			assert dump(var).str() == 'false'
+		}
+	}
+}

vlib/v/tests/comptime_var_on_multiple_args_test.v

@@ -0,0 +1,31 @@
+module main
+
+fn ok3[T](o T) T {
+	return o
+}
+
+fn ok2[T](o T, s string) string {
+	return s
+}
+
+fn ok[T](o T, s string) string {
+	return '${o}${s}'
+}
+
+struct Test {
+	s string
+}
+
+fn test_main() {
+	s := Test{'foo'}
+
+	$for f in Test.fields {
+		assert ok3(s.$(f.name)).str() == 'foo'
+	}
+	$for f in Test.fields {
+		assert ok(s.$(f.name), 'bar') == 'foobar'
+	}
+	$for f in Test.fields {
+		assert ok2(s.$(f.name), 'ok') == 'ok'
+	}
+}