From 5043e14e355659f0232f770241694d8ff3a242c9 Mon Sep 17 00:00:00 2001
From: Felipe Pena <felipensp@gmail.com>
Date: Fri, 22 Dec 2023 05:08:40 -0300
Subject: [PATCH] v: allow comptime-for to iterate over comptime variables, add
 `$string` comptime type, cleanup (#20233)

---
 vlib/v/ast/ast.v                              |  3 +
 vlib/v/ast/table.v                            | 55 -------------------
 vlib/v/checker/comptime.v                     | 29 ++++++++--
 vlib/v/checker/tests/comptime_for.out         |  6 +-
 .../checker/tests/wrong_comptime_for_err.out  | 28 ++++++++++
 .../v/checker/tests/wrong_comptime_for_err.vv | 33 +++++++++++
 vlib/v/comptime/comptimeinfo.v                |  3 +
 vlib/v/fmt/fmt.v                              | 11 +++-
 vlib/v/gen/c/assign.v                         |  2 +-
 vlib/v/gen/c/cgen.v                           |  2 +-
 vlib/v/gen/c/comptime.v                       | 10 +++-
 vlib/v/gen/c/fn.v                             |  2 +-
 vlib/v/gen/c/for.v                            |  4 +-
 vlib/v/gen/golang/golang.v                    |  1 +
 vlib/v/parser/comptime.v                      | 16 +++++-
 vlib/v/slow_tests/inout/comptime_iterate.out  |  4 ++
 vlib/v/slow_tests/inout/comptime_iterate.vv   | 35 ++++++++++++
 17 files changed, 170 insertions(+), 74 deletions(-)
 create mode 100644 vlib/v/checker/tests/wrong_comptime_for_err.out
 create mode 100644 vlib/v/checker/tests/wrong_comptime_for_err.vv
 create mode 100644 vlib/v/slow_tests/inout/comptime_iterate.out
 create mode 100644 vlib/v/slow_tests/inout/comptime_iterate.vv

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 73c96200e..d5b37894e 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -144,6 +144,7 @@ pub enum ComptimeTypeKind {
 	alias
 	function
 	option
+	string
 }
 
 pub struct ComptimeType {
@@ -168,6 +169,7 @@ pub fn (cty ComptimeType) str() string {
 		.alias { '\$alias' }
 		.function { '\$function' }
 		.option { '\$option' }
+		.string { '\$string' }
 	}
 }
 
@@ -1213,6 +1215,7 @@ pub:
 pub mut:
 	stmts []Stmt
 	typ   Type
+	expr  Expr
 }
 
 pub struct ForStmt {
diff --git a/vlib/v/ast/table.v b/vlib/v/ast/table.v
index 5ebf540e4..684b9271f 100644
--- a/vlib/v/ast/table.v
+++ b/vlib/v/ast/table.v
@@ -2353,61 +2353,6 @@ pub fn (mut t Table) check_if_elements_need_unwrap(root_typ Type, typ Type) bool
 	return false
 }
 
-pub fn (t &Table) is_comptime_type(x Type, y ComptimeType) bool {
-	x_kind := t.type_kind(x)
-	match y.kind {
-		.unknown {
-			return false
-		}
-		.map_ {
-			return x_kind == .map
-		}
-		.int {
-			return x_kind in [.i8, .i16, .int, .i64, .u8, .u16, .u32, .u64, .usize, .isize,
-				.int_literal]
-		}
-		.float {
-			return x_kind in [.f32, .f64, .float_literal]
-		}
-		.struct_ {
-			return x_kind == .struct_
-		}
-		.iface {
-			return x_kind == .interface_
-		}
-		.array {
-			return x_kind in [.array, .array_fixed]
-		}
-		.array_dynamic {
-			return x_kind == .array
-		}
-		.array_fixed {
-			return x_kind == .array_fixed
-		}
-		.sum_type {
-			return x_kind == .sum_type
-		}
-		.enum_ {
-			return x_kind == .enum_
-		}
-		.alias {
-			return x_kind == .alias
-		}
-		.function {
-			return x_kind == .function
-		}
-		.option {
-			return x.has_flag(.option)
-		}
-	}
-}
-
-@[inline]
-pub fn (t &Table) is_comptime_var(node Expr) bool {
-	return node is Ident && node.info is IdentVar && node.obj is Var
-		&& (node.obj as Var).ct_type_var != .no_comptime
-}
-
 pub fn (t &Table) dependent_names_in_expr(expr Expr) []string {
 	mut names := []string{}
 	match expr {
diff --git a/vlib/v/checker/comptime.v b/vlib/v/checker/comptime.v
index b48980216..5d207f3b3 100644
--- a/vlib/v/checker/comptime.v
+++ b/vlib/v/checker/comptime.v
@@ -213,11 +213,20 @@ fn (mut c Checker) comptime_selector(mut node ast.ComptimeSelector) ast.Type {
 }
 
 fn (mut c Checker) comptime_for(mut node ast.ComptimeFor) {
-	typ := c.unwrap_generic(node.typ)
+	typ := if node.typ != ast.void_type {
+		c.unwrap_generic(node.typ)
+	} else {
+		node.typ = c.expr(mut node.expr)
+		node.typ
+	}
 	sym := c.table.final_sym(typ)
 	if sym.kind == .placeholder || typ.has_flag(.generic) {
-		c.error('\$for expects a type name to be used here, but ${sym.name} is not a type name',
+		c.error('\$for expects a type name or variable name to be used here, but ${sym.name} is not a type or variable name',
 			node.typ_pos)
+		return
+	} else if sym.kind == .void {
+		c.error('only known compile-time variables can be used', node.typ_pos)
+		return
 	}
 	if node.kind == .fields {
 		if sym.kind in [.struct_, .interface_] {
@@ -261,7 +270,8 @@ fn (mut c Checker) comptime_for(mut node ast.ComptimeFor) {
 				}
 				c.pop_comptime_info()
 			}
-		} else if c.table.generic_type_names(node.typ).len == 0 && sym.kind != .placeholder {
+		} else if node.typ != ast.void_type && c.table.generic_type_names(node.typ).len == 0
+			&& sym.kind != .placeholder {
 			c.error('iterating over .fields is supported only for structs and interfaces, and ${sym.name} is neither',
 				node.typ_pos)
 			return
@@ -302,7 +312,18 @@ fn (mut c Checker) comptime_for(mut node ast.ComptimeFor) {
 		if c.variant_data_type == 0 {
 			c.variant_data_type = ast.Type(c.table.find_type_idx('VariantData'))
 		}
-		for variant in (sym.info as ast.SumType).variants {
+		mut variants := []ast.Type{}
+		if c.comptime.comptime_for_field_var != '' && typ == c.field_data_type {
+			sumtype_sym := c.table.sym(c.comptime.comptime_for_field_type)
+			if sumtype_sym.kind == .sum_type {
+				variants = (sumtype_sym.info as ast.SumType).variants.clone()
+			}
+		} else if sym.kind != .sum_type {
+			c.error('${sym.name} is not Sum type to use with .variants', node.typ_pos)
+		} else {
+			variants = (sym.info as ast.SumType).variants.clone()
+		}
+		for variant in variants {
 			c.push_new_comptime_info()
 			c.comptime.inside_comptime_for = true
 			c.comptime.comptime_for_variant_var = node.val_var
diff --git a/vlib/v/checker/tests/comptime_for.out b/vlib/v/checker/tests/comptime_for.out
index 577fe4f84..2e3d181b4 100644
--- a/vlib/v/checker/tests/comptime_for.out
+++ b/vlib/v/checker/tests/comptime_for.out
@@ -11,20 +11,20 @@ vlib/v/checker/tests/comptime_for.vv:3:7: warning: unused variable: `f`
       |          ^
     4 |     $for f in T.fields {
     5 |         $if f.typ is Huh {}
-vlib/v/checker/tests/comptime_for.vv:2:12: error: $for expects a type name to be used here, but Huh is not a type name
+vlib/v/checker/tests/comptime_for.vv:2:12: error: $for expects a type name or variable name to be used here, but Huh is not a type or variable name
     1 | fn unknown() {
     2 |     $for m in Huh.methods {}
       |               ~~~
     3 |     $for f in Huh.fields {}
     4 |     $for f in T.fields {
-vlib/v/checker/tests/comptime_for.vv:3:12: error: $for expects a type name to be used here, but Huh is not a type name
+vlib/v/checker/tests/comptime_for.vv:3:12: error: $for expects a type name or variable name to be used here, but Huh is not a type or variable name
     1 | fn unknown() {
     2 |     $for m in Huh.methods {}
     3 |     $for f in Huh.fields {}
       |               ~~~
     4 |     $for f in T.fields {
     5 |         $if f.typ is Huh {}
-vlib/v/checker/tests/comptime_for.vv:4:12: error: $for expects a type name to be used here, but T is not a type name
+vlib/v/checker/tests/comptime_for.vv:4:12: error: $for expects a type name or variable name to be used here, but T is not a type or variable name
     2 |     $for m in Huh.methods {}
     3 |     $for f in Huh.fields {}
     4 |     $for f in T.fields {
diff --git a/vlib/v/checker/tests/wrong_comptime_for_err.out b/vlib/v/checker/tests/wrong_comptime_for_err.out
new file mode 100644
index 000000000..7e7df2a80
--- /dev/null
+++ b/vlib/v/checker/tests/wrong_comptime_for_err.out
@@ -0,0 +1,28 @@
+vlib/v/checker/tests/wrong_comptime_for_err.vv:28:7: warning: unused variable: `f`
+   26 | 
+   27 | 
+   28 |     $for f in a.variants {
+      |          ^
+   29 |     }
+   30 |
+vlib/v/checker/tests/wrong_comptime_for_err.vv:31:7: warning: unused variable: `f`
+   29 |     }
+   30 | 
+   31 |     $for f in T.variants {
+      |          ^
+   32 |     }
+   33 | }
+vlib/v/checker/tests/wrong_comptime_for_err.vv:28:12: error: Foo is not Sum type to use with .variants
+   26 | 
+   27 | 
+   28 |     $for f in a.variants {
+      |               ^
+   29 |     }
+   30 |
+vlib/v/checker/tests/wrong_comptime_for_err.vv:31:12: error: $for expects a type name or variable name to be used here, but T is not a type or variable name
+   29 |     }
+   30 | 
+   31 |     $for f in T.variants {
+      |               ^
+   32 |     }
+   33 | }
diff --git a/vlib/v/checker/tests/wrong_comptime_for_err.vv b/vlib/v/checker/tests/wrong_comptime_for_err.vv
new file mode 100644
index 000000000..2e1367207
--- /dev/null
+++ b/vlib/v/checker/tests/wrong_comptime_for_err.vv
@@ -0,0 +1,33 @@
+type TestSum = int | string | []int
+
+struct Bar {
+    y int
+}
+
+struct Foo {
+mut:
+	a TestSum
+    b int   
+    c string
+    d Bar
+}
+
+fn main() {
+	a := Foo{}
+    $for f in a.fields {
+        $if f.typ is $sumtype {
+            $for f2 in f.variants {
+                dump(f2)
+            }
+        } $else $if f.typ is Bar {
+            assert true
+        }
+    }
+
+
+	$for f in a.variants {
+	}
+
+	$for f in T.variants {
+	}
+}
\ No newline at end of file
diff --git a/vlib/v/comptime/comptimeinfo.v b/vlib/v/comptime/comptimeinfo.v
index 39c7d37ec..e421b1beb 100644
--- a/vlib/v/comptime/comptimeinfo.v
+++ b/vlib/v/comptime/comptimeinfo.v
@@ -183,6 +183,9 @@ pub fn (mut ct ComptimeInfo) is_comptime_type(x ast.Type, y ast.ComptimeType) bo
 		.map_ {
 			return x_kind == .map
 		}
+		.string {
+			return x_kind == .string
+		}
 		.int {
 			return x_kind in [.i8, .i16, .int, .i64, .u8, .u16, .u32, .u64, .usize, .isize,
 				.int_literal]
diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v
index 163d234c6..a20a5fdd4 100644
--- a/vlib/v/fmt/fmt.v
+++ b/vlib/v/fmt/fmt.v
@@ -775,6 +775,7 @@ pub fn (mut f Fmt) expr(node_ ast.Expr) {
 				.alias { f.write('\$alias') }
 				.function { f.write('\$function') }
 				.option { f.write('\$option') }
+				.string { f.write('\$string') }
 			}
 		}
 	}
@@ -878,9 +879,15 @@ pub fn (mut f Fmt) branch_stmt(node ast.BranchStmt) {
 }
 
 pub fn (mut f Fmt) comptime_for(node ast.ComptimeFor) {
-	typ := f.no_cur_mod(f.table.type_to_str_using_aliases(node.typ, f.mod2alias))
+	typ := if node.typ != ast.void_type {
+		f.no_cur_mod(f.table.type_to_str_using_aliases(node.typ, f.mod2alias))
+	} else {
+		(node.expr as ast.Ident).name
+	}
 	f.write('\$for ${node.val_var} in ${typ}.${node.kind.str()} {')
-	f.mark_types_import_as_used(node.typ)
+	if node.typ != ast.void_type {
+		f.mark_types_import_as_used(node.typ)
+	}
 	if node.stmts.len > 0 || node.pos.line_nr < node.pos.last_line {
 		f.writeln('')
 		f.stmts(node.stmts)
diff --git a/vlib/v/gen/c/assign.v b/vlib/v/gen/c/assign.v
index 1ebf10db7..b26ff3815 100644
--- a/vlib/v/gen/c/assign.v
+++ b/vlib/v/gen/c/assign.v
@@ -235,7 +235,7 @@ fn (mut g Gen) assign_stmt(node_ ast.AssignStmt) {
 				}
 			}
 			if mut left.obj is ast.Var {
-				if val is ast.Ident && g.table.is_comptime_var(val) {
+				if val is ast.Ident && g.comptime.is_comptime_var(val) {
 					ctyp := g.unwrap_generic(g.comptime.get_comptime_var_type(val))
 					if ctyp != ast.void_type {
 						var_type = ctyp
diff --git a/vlib/v/gen/c/cgen.v b/vlib/v/gen/c/cgen.v
index 158abe562..1d839adbd 100644
--- a/vlib/v/gen/c/cgen.v
+++ b/vlib/v/gen/c/cgen.v
@@ -4591,7 +4591,7 @@ fn (mut g Gen) cast_expr(node ast.CastExpr) {
 	node_typ := g.unwrap_generic(node.typ)
 	mut expr_type := node.expr_type
 	sym := g.table.sym(node_typ)
-	if (node.expr is ast.Ident && g.table.is_comptime_var(node.expr))
+	if (node.expr is ast.Ident && g.comptime.is_comptime_var(node.expr))
 		|| node.expr is ast.ComptimeSelector {
 		expr_type = g.unwrap_generic(g.comptime.get_comptime_var_type(node.expr))
 	}
diff --git a/vlib/v/gen/c/comptime.v b/vlib/v/gen/c/comptime.v
index ae530e160..5b38ec5d6 100644
--- a/vlib/v/gen/c/comptime.v
+++ b/vlib/v/gen/c/comptime.v
@@ -467,7 +467,7 @@ fn (mut g Gen) comptime_if_cond(cond ast.Expr, pkg_exist bool) (bool, bool) {
 						&& cond.right in [ast.ComptimeType, ast.TypeNode] {
 						exp_type := g.get_expr_type(cond.left)
 						if cond.right is ast.ComptimeType {
-							is_true := g.table.is_comptime_type(exp_type, cond.right)
+							is_true := g.comptime.is_comptime_type(exp_type, cond.right)
 							if cond.op == .key_is {
 								if is_true {
 									g.write('1')
@@ -583,7 +583,7 @@ fn (mut g Gen) comptime_if_cond(cond ast.Expr, pkg_exist bool) (bool, bool) {
 
 						for expr in cond.right.exprs {
 							if expr is ast.ComptimeType {
-								if g.table.is_comptime_type(checked_type, expr as ast.ComptimeType) {
+								if g.comptime.is_comptime_type(checked_type, expr as ast.ComptimeType) {
 									if cond.op == .key_in {
 										g.write('1')
 									} else {
@@ -714,7 +714,11 @@ fn (mut g Gen) resolve_comptime_type(node ast.Expr, default_type ast.Type) ast.T
 }
 
 fn (mut g Gen) comptime_for(node ast.ComptimeFor) {
-	sym := g.table.final_sym(g.unwrap_generic(node.typ))
+	sym := if node.typ != g.field_data_type {
+		g.table.final_sym(g.unwrap_generic(node.typ))
+	} else {
+		g.table.final_sym(g.comptime.comptime_for_field_type)
+	}
 	g.writeln('/* \$for ${node.val_var} in ${sym.name}.${node.kind.str()} */ {')
 	g.indent++
 	mut i := 0
diff --git a/vlib/v/gen/c/fn.v b/vlib/v/gen/c/fn.v
index a75e8b66c..1d24d633c 100644
--- a/vlib/v/gen/c/fn.v
+++ b/vlib/v/gen/c/fn.v
@@ -1744,7 +1744,7 @@ fn (mut g Gen) fn_call(node ast.CallExpr) {
 	mut print_auto_str := false
 	if is_print && (node.args[0].typ != ast.string_type
 		|| g.comptime.comptime_for_method.len > 0
-		|| g.table.is_comptime_var(node.args[0].expr)) {
+		|| g.comptime.is_comptime_var(node.args[0].expr)) {
 		g.inside_interface_deref = true
 		defer {
 			g.inside_interface_deref = false
diff --git a/vlib/v/gen/c/for.v b/vlib/v/gen/c/for.v
index 0c1dbbd2c..57135c8cc 100644
--- a/vlib/v/gen/c/for.v
+++ b/vlib/v/gen/c/for.v
@@ -141,7 +141,7 @@ fn (mut g Gen) for_in_stmt(node_ ast.ForInStmt) {
 	mut node := unsafe { node_ }
 	mut is_comptime := false
 
-	if (node.cond is ast.Ident && g.table.is_comptime_var(node.cond))
+	if (node.cond is ast.Ident && g.comptime.is_comptime_var(node.cond))
 		|| node.cond is ast.ComptimeSelector {
 		mut unwrapped_typ := g.unwrap_generic(node.cond_type)
 		ctyp := g.comptime.get_comptime_var_type(node.cond)
@@ -221,7 +221,7 @@ fn (mut g Gen) for_in_stmt(node_ ast.ForInStmt) {
 		mut val_sym := g.table.sym(node.val_type)
 		op_field := g.dot_or_ptr(node.cond_type)
 
-		if is_comptime && g.table.is_comptime_var(node.cond) {
+		if is_comptime && g.comptime.is_comptime_var(node.cond) {
 			mut unwrapped_typ := g.unwrap_generic(node.cond_type)
 			ctyp := g.unwrap_generic(g.comptime.get_comptime_var_type(node.cond))
 			if ctyp != ast.void_type {
diff --git a/vlib/v/gen/golang/golang.v b/vlib/v/gen/golang/golang.v
index 1688f97ab..91125171a 100644
--- a/vlib/v/gen/golang/golang.v
+++ b/vlib/v/gen/golang/golang.v
@@ -674,6 +674,7 @@ pub fn (mut f Gen) expr(node_ ast.Expr) {
 				.alias { f.write('\$alias') }
 				.function { f.write('\$function') }
 				.option { f.write('\$option') }
+				.string { f.write('\$string') }
 			}
 		}
 	}
diff --git a/vlib/v/parser/comptime.v b/vlib/v/parser/comptime.v
index e19f3883b..177cdc107 100644
--- a/vlib/v/parser/comptime.v
+++ b/vlib/v/parser/comptime.v
@@ -11,7 +11,7 @@ import v.token
 const supported_comptime_calls = ['html', 'tmpl', 'env', 'embed_file', 'pkgconfig', 'compile_error',
 	'compile_warn', 'res']
 const comptime_types = ['map', 'array', 'array_dynamic', 'array_fixed', 'int', 'float', 'struct',
-	'interface', 'enum', 'sumtype', 'alias', 'function', 'option']
+	'interface', 'enum', 'sumtype', 'alias', 'function', 'option', 'string']
 
 fn (mut p Parser) parse_comptime_type() ast.ComptimeType {
 	pos := p.tok.pos()
@@ -61,6 +61,9 @@ fn (mut p Parser) parse_comptime_type() ast.ComptimeType {
 		'option' {
 			.option
 		}
+		'string' {
+			.string
+		}
 		else {
 			.unknown
 		}
@@ -298,14 +301,22 @@ fn (mut p Parser) comptime_for() ast.ComptimeFor {
 	// `$for val in App.values {`
 	// `$for field in App.fields {`
 	// `$for attr in App.attributes {`
+	// `$for variant in App.variants {`
 	p.next()
 	p.check(.key_for)
 	var_pos := p.tok.pos()
 	val_var := p.check_name()
 	p.check(.key_in)
+	mut expr := ast.empty_expr
 	mut typ_pos := p.tok.pos()
 	lang := p.parse_language()
-	typ := p.parse_any_type(lang, false, false, false)
+	mut typ := ast.void_type
+	if p.tok.lit[0].is_capital() {
+		typ = p.parse_any_type(lang, false, false, false)
+	} else {
+		expr = p.ident(lang)
+		p.mark_var_as_used((expr as ast.Ident).name)
+	}
 	typ_pos = typ_pos.extend(p.prev_tok.pos())
 	p.check(.dot)
 	for_val := p.check_name()
@@ -365,6 +376,7 @@ fn (mut p Parser) comptime_for() ast.ComptimeFor {
 		stmts: stmts
 		kind: kind
 		typ: typ
+		expr: expr
 		typ_pos: typ_pos
 		pos: spos.extend(p.tok.pos())
 	}
diff --git a/vlib/v/slow_tests/inout/comptime_iterate.out b/vlib/v/slow_tests/inout/comptime_iterate.out
new file mode 100644
index 000000000..65bcf3f05
--- /dev/null
+++ b/vlib/v/slow_tests/inout/comptime_iterate.out
@@ -0,0 +1,4 @@
+[vlib/v/slow_tests/inout/comptime_iterate.vv:18] f.name: a
+[vlib/v/slow_tests/inout/comptime_iterate.vv:22] 'int': int
+[vlib/v/slow_tests/inout/comptime_iterate.vv:24] 'string': string
+[vlib/v/slow_tests/inout/comptime_iterate.vv:31] f3.name: y
diff --git a/vlib/v/slow_tests/inout/comptime_iterate.vv b/vlib/v/slow_tests/inout/comptime_iterate.vv
new file mode 100644
index 000000000..6abd2b104
--- /dev/null
+++ b/vlib/v/slow_tests/inout/comptime_iterate.vv
@@ -0,0 +1,35 @@
+type TestSum = int | string
+
+struct Bar {
+	y int
+}
+
+struct Foo {
+mut:
+	a TestSum
+	b int
+	c string
+	d Bar
+}
+
+fn main() {
+	$for f in Foo.fields {
+		$if f.typ is $sumtype {
+			dump(f.name)
+			$for f2 in f.variants {
+				$if f2.typ is $int {
+					assert f2.typ == typeof[int]().idx
+					dump('int')
+				} $else $if f2.typ is $string {
+					dump('string')
+					assert f2.typ == typeof[string]().idx
+				}
+			}
+		} $else $if f.typ is Bar {
+			$for f3 in f.fields {
+				assert f3.typ == typeof[int]().idx
+				dump(f3.name)
+			}
+		}
+	}
+}
-- 
2.39.5