v: add compile-time enum evaluation with $for item in MyEnum.fields { dump(item.value) dump(item.name) } (#17517)

vlib/builtin/builtin.v

@@ -107,6 +107,12 @@ pub:
 	typ         int
 }
 
+pub struct EnumData {
+pub:
+	name  string
+	value i64
+}
+
 // FieldData holds information about a field. Fields reside on structs.
 pub struct FieldData {
 pub:

vlib/v/ast/ast.v

@@ -1033,6 +1033,7 @@ pub enum ComptimeForKind {
 	methods
 	fields
 	attributes
+	values
 }
 
 pub struct ComptimeFor {

vlib/v/ast/str.v

@@ -691,5 +691,6 @@ pub fn (e ComptimeForKind) str() string {
 		.methods { return 'methods' }
 		.fields { return 'fields' }
 		.attributes { return 'attributes' }
+		.values { return 'values' }
 	}
 }

vlib/v/checker/checker.v

@@ -100,6 +100,7 @@ mut:
 	comptime_fields_default_type     ast.Type
 	comptime_fields_type             map[string]ast.Type
 	comptime_for_field_value         ast.StructField // value of the field variable
+	comptime_enum_field_value        string // current enum value name
 	fn_scope                         &ast.Scope = unsafe { nil }
 	main_fn_decl_node                ast.FnDecl
 	match_exhaustive_cutoff_limit    int = 10
@@ -115,6 +116,7 @@ mut:
 	is_index_assign                  bool
 	comptime_call_pos                int // needed for correctly checking use before decl for templates
 	goto_labels                      map[string]ast.GotoLabel // to check for unused goto labels
+	enum_data_type                   ast.Type
 }
 
 pub fn new_checker(table &ast.Table, pref_ &pref.Preferences) &Checker {
@@ -1301,6 +1303,10 @@ fn (mut c Checker) selector_expr(mut node ast.SelectorExpr) ast.Type {
 		c.error('`${node.expr}` does not return a value', node.pos)
 		node.expr_type = ast.void_type
 		return ast.void_type
+	} else if c.inside_comptime_for_field && typ == c.enum_data_type && node.field_name == 'value' {
+		// for comp-time enum.values
+		node.expr_type = c.comptime_fields_type[c.comptime_for_field_var]
+		return node.expr_type
 	}
 	node.expr_type = typ
 	if !(node.expr is ast.Ident && (node.expr as ast.Ident).kind == .constant) {

vlib/v/checker/comptime.v

@@ -210,6 +210,20 @@ fn (mut c Checker) comptime_for(node ast.ComptimeFor) {
 			c.comptime_for_field_var = ''
 			c.inside_comptime_for_field = false
 		}
+	} else if node.kind == .values {
+		if sym.kind == .enum_ {
+			sym_info := sym.info as ast.Enum
+			c.inside_comptime_for_field = true
+			if c.enum_data_type == 0 {
+				c.enum_data_type = ast.Type(c.table.find_type_idx('EnumData'))
+			}
+			for field in sym_info.vals {
+				c.comptime_enum_field_value = field
+				c.comptime_for_field_var = node.val_var
+				c.comptime_fields_type[node.val_var] = node.typ
+				c.stmts(node.stmts)
+			}
+		}
 	} else {
 		c.stmts(node.stmts)
 	}

vlib/v/gen/c/cgen.v

@@ -46,6 +46,7 @@ fn string_array_to_map(a []string) map[string]bool {
 pub struct Gen {
 	pref                &pref.Preferences = unsafe { nil }
 	field_data_type     ast.Type // cache her to avoid map lookups
+	enum_data_type      ast.Type // cache her to avoid map lookups
 	module_built        string
 	timers_should_print bool
 	table               &ast.Table = unsafe { nil }
@@ -198,7 +199,8 @@ mut:
 	comptime_for_method_var          string // $for method in T.methods {}; the variable name
 	comptime_for_field_var           string // $for field in T.fields {}; the variable name
 	comptime_for_field_value         ast.StructField // value of the field variable
-	comptime_for_field_type          ast.Type        // type of the field variable inferred from `$if field.typ is T {}`
+	comptime_enum_field_value        string   // value of enum name
+	comptime_for_field_type          ast.Type // type of the field variable inferred from `$if field.typ is T {}`
 	comptime_var_type_map            map[string]ast.Type
 	comptime_values_stack            []CurrentComptimeValues // stores the values from the above on each $for loop, to make nesting them easier
 	prevent_sum_type_unwrapping_once bool // needed for assign new values to sum type
@@ -308,6 +310,7 @@ pub fn gen(files []&ast.File, table &ast.Table, pref_ &pref.Preferences) (string
 		timers: util.new_timers(should_print: timers_should_print, label: 'global_cgen')
 		inner_loop: &ast.empty_stmt
 		field_data_type: ast.Type(table.find_type_idx('FieldData'))
+		enum_data_type: ast.Type(table.find_type_idx('EnumData'))
 		is_cc_msvc: pref_.ccompiler == 'msvc'
 		use_segfault_handler: !('no_segfault_handler' in pref_.compile_defines
 			|| pref_.os in [.wasm32, .wasm32_emscripten])
@@ -663,6 +666,7 @@ fn cgen_process_one_file_cb(mut p pool.PoolProcessor, idx int, wid int) &Gen {
 		)
 		inner_loop: &ast.empty_stmt
 		field_data_type: ast.Type(global_g.table.find_type_idx('FieldData'))
+		enum_data_type: ast.Type(global_g.table.find_type_idx('EnumData'))
 		array_sort_fn: global_g.array_sort_fn
 		waiter_fns: global_g.waiter_fns
 		threaded_fns: global_g.threaded_fns
@@ -3444,6 +3448,13 @@ fn (mut g Gen) selector_expr(node ast.SelectorExpr) {
 				g.error('unknown generic field', node.pos)
 			}
 		}
+	} else {
+		// for comp-time enum value evaluation
+		if node.expr_type == g.enum_data_type && node.expr is ast.Ident
+			&& (node.expr as ast.Ident).name == 'value' {
+			g.write(node.str())
+			return
+		}
 	}
 	if node.expr_type == 0 {
 		g.checker_bug('unexpected SelectorExpr.expr_type = 0', node.pos)

vlib/v/gen/c/comptime.v

@@ -827,6 +827,30 @@ fn (mut g Gen) comptime_for(node ast.ComptimeFor) {
 				g.pop_existing_comptime_values()
 			}
 		}
+	} else if node.kind == .values {
+		if sym.kind == .enum_ {
+			if sym.info is ast.Enum {
+				if sym.info.vals.len > 0 {
+					g.writeln('\tEnumData ${node.val_var} = {0};')
+				}
+				for val in sym.info.vals {
+					g.comptime_enum_field_value = val
+					g.comptime_for_field_type = node.typ
+
+					g.writeln('/* enum vals ${i} */ {')
+					g.writeln('\t${node.val_var}.name = _SLIT("${val}");')
+					g.write('\t${node.val_var}.value = ')
+					if g.pref.translated && node.typ.is_number() {
+						g.writeln('_const_main__${g.comptime_enum_field_value};')
+					} else {
+						g.writeln('${g.typ(g.comptime_for_field_type)}__${g.comptime_enum_field_value};')
+					}
+					g.stmts(node.stmts)
+					g.writeln('}')
+					i++
+				}
+			}
+		}
 	} else if node.kind == .attributes {
 		if sym.info is ast.Struct {
 			if sym.info.attrs.len > 0 {
@@ -840,6 +864,7 @@ fn (mut g Gen) comptime_for(node ast.ComptimeFor) {
 				g.writeln('\t${node.val_var}.kind = AttributeKind__${attr.kind};')
 				g.stmts(node.stmts)
 				g.writeln('}')
+				i++
 			}
 		}
 	}

vlib/v/parser/comptime.v

@@ -276,6 +276,13 @@ fn (mut p Parser) comptime_for() ast.ComptimeFor {
 			typ: p.table.find_type_idx('FunctionData')
 			pos: var_pos
 		})
+	} else if for_val == 'values' {
+		p.scope.register(ast.Var{
+			name: val_var
+			typ: p.table.find_type_idx('EnumData')
+			pos: var_pos
+		})
+		kind = .values
 	} else if for_val == 'fields' {
 		p.scope.register(ast.Var{
 			name: val_var
@@ -291,7 +298,7 @@ fn (mut p Parser) comptime_for() ast.ComptimeFor {
 		})
 		kind = .attributes
 	} else {
-		p.error_with_pos('unknown kind `${for_val}`, available are: `methods`, `fields` or `attributes`',
+		p.error_with_pos('unknown kind `${for_val}`, available are: `methods`, `fields`, `values`, or `attributes`',
 			p.prev_tok.pos())
 		return ast.ComptimeFor{}
 	}

vlib/v/parser/tests/comptime_unknown_meta_kind_err.out

@@ -0,0 +1,7 @@
+vlib/v/parser/tests/comptime_unknown_meta_kind_err.vv:7:20: error: unknown kind `abcde`, available are: `methods`, `fields`, `values`, or `attributes`
+    5 | 
+    6 | fn test_main() {
+    7 |     $for item in Test.abcde {
+      |                       ~~~~~
+    8 |         dump(item)
+    9 |     }

vlib/v/parser/tests/comptime_unknown_meta_kind_err.vv

@@ -0,0 +1,10 @@
+enum Test {
+	foo
+	bar
+}
+
+fn test_main() {
+	$for item in Test.abcde {
+		dump(item)
+	}
+}

vlib/v/tests/comptime_enum_test.v

@@ -0,0 +1,27 @@
+enum Test {
+	foo
+	bar
+}
+
+fn test_main() {
+	$for item in Test.values {
+		assert item.name in ['foo', 'bar']
+		match item.value {
+			.foo {
+				println('foo>> item: ${item.name}')
+				assert item.value == .foo
+			}
+			.bar {
+				println('foo>> item: ${item.name}')
+				assert item.value == .bar
+			}
+		}
+		if item.value == .foo {
+			println('foo>> item: ${item.name}')
+			assert item.value == .foo
+		} else if item.value == .bar {
+			println('foo>> item: ${item.name}')
+			assert item.value == .bar
+		}
+	}
+}