checker, cgen: allow iterating over optional array fields (#16858)

vlib/v/checker/for.v

@@ -154,7 +154,7 @@ fn (mut c Checker) for_in_stmt(mut node ast.ForInStmt) {
 			if sym.kind == .string {
 				value_type = ast.u8_type
 			}
-			if value_type == ast.void_type || typ.has_flag(.optional) || typ.has_flag(.result) {
+			if value_type == ast.void_type || typ.has_flag(.result) {
 				if typ != ast.void_type {
 					c.error('for in: cannot index `${c.table.type_to_str(typ)}`', node.cond.pos())
 				}

vlib/v/gen/c/cgen.v

@@ -4218,6 +4218,11 @@ fn (mut g Gen) select_expr(node ast.SelectExpr) {
 	}
 }
 
+pub fn (mut g Gen) is_comptime_var(node ast.Expr) bool {
+	return g.inside_comptime_for_field && node is ast.Ident
+		&& (node as ast.Ident).info is ast.IdentVar && ((node as ast.Ident).obj as ast.Var).is_comptime_field
+}
+
 fn (mut g Gen) ident(node ast.Ident) {
 	prevent_sum_type_unwrapping_once := g.prevent_sum_type_unwrapping_once
 	g.prevent_sum_type_unwrapping_once = false

vlib/v/gen/c/for.v

@@ -132,11 +132,15 @@ fn (mut g Gen) for_in_stmt(node_ ast.ForInStmt) {
 	mut node := unsafe { node_ }
 	if node.kind == .any {
 		g.inside_for_in_any_cond = true
-		unwrapped_typ := g.unwrap_generic(node.cond_type)
-		unwrapped_sym := g.table.sym(unwrapped_typ)
+		mut unwrapped_typ := g.unwrap_generic(node.cond_type)
+		mut unwrapped_sym := g.table.sym(unwrapped_typ)
 		node.kind = unwrapped_sym.kind
 		node.cond_type = unwrapped_typ
 		if node.key_var.len > 0 {
+			if g.is_comptime_var(node.cond) {
+				unwrapped_typ = g.unwrap_generic(g.comptime_for_field_type)
+				unwrapped_sym = g.table.sym(unwrapped_typ)
+			}
 			key_type := match unwrapped_sym.kind {
 				.map { unwrapped_sym.map_info().key_type }
 				else { ast.int_type }
@@ -163,10 +167,20 @@ fn (mut g Gen) for_in_stmt(node_ ast.ForInStmt) {
 	} else if node.kind == .array {
 		// `for num in nums {`
 		// g.writeln('// FOR IN array')
-		styp := g.typ(node.val_type)
-		val_sym := g.table.sym(node.val_type)
+		mut styp := g.typ(node.val_type)
+		mut val_sym := g.table.sym(node.val_type)
+
+		if g.is_comptime_var(node.cond) {
+			unwrapped_typ := g.unwrap_generic(g.comptime_for_field_type)
+			val_sym = g.table.sym(unwrapped_typ)
+			node.val_type = g.table.value_type(unwrapped_typ)
+			styp = g.typ(node.val_type)
+			node.scope.update_var_type(node.val_var, node.val_type)
+			node.cond_type = node.val_type
+		}
 		mut cond_var := ''
-		if node.cond is ast.Ident || node.cond is ast.SelectorExpr {
+		if (node.cond is ast.Ident && !node.cond_type.has_flag(.optional))
+			|| node.cond is ast.SelectorExpr {
 			cond_var = g.expr_string(node.cond)
 		} else {
 			cond_var = g.new_tmp_var()
@@ -180,7 +194,13 @@ fn (mut g Gen) for_in_stmt(node_ ast.ForInStmt) {
 		share_accessor := if node.cond_type.share() == .shared_t { 'val.' } else { '' }
 		op_field := field_accessor + share_accessor
 		g.empty_line = true
-		g.writeln('for (int ${i} = 0; ${i} < ${cond_var}${op_field}len; ++${i}) {')
+		opt_expr := '(*(${g.typ(node.cond_type.clear_flag(.optional))}*)${cond_var}${op_field}data)'
+		cond_expr := if node.cond_type.has_flag(.optional) {
+			'/*opt*/ ${opt_expr}${op_field}len'
+		} else {
+			'${cond_var}${op_field}len'
+		}
+		g.writeln('for (int ${i} = 0; ${i} < ${cond_expr}; ++${i}) {')
 		if node.val_var != '_' {
 			if val_sym.kind == .function {
 				g.write('\t')
@@ -196,7 +216,9 @@ fn (mut g Gen) for_in_stmt(node_ ast.ForInStmt) {
 				// instead of
 				// `int* val = ((int**)arr.data)[i];`
 				// right := if node.val_is_mut { styp } else { styp + '*' }
-				right := if node.val_is_mut || node.val_is_ref {
+				right := if node.cond_type.has_flag(.optional) {
+					'/*opt*/ ((${styp}*)${opt_expr}${op_field}data)[${i}]'
+				} else if node.val_is_mut || node.val_is_ref {
 					'((${styp})${cond_var}${op_field}data) + ${i}'
 				} else {
 					'((${styp}*)${cond_var}${op_field}data)[${i}]'

vlib/v/tests/for_loop_with_optional2_test.v

@@ -0,0 +1,44 @@
+struct Test {
+	a ?[]int
+	b ?[]string
+	c ?[]f64
+	d ?[][]string
+}
+
+fn test_for_in_optional_fields() {
+	mut out := []string{}
+	test := Test{
+		a: [1, 2, 3]
+		b: ['foo', 'bar']
+		c: [1.2, 2.3]
+		d: [['foo'], ['bar']]
+	}
+	//
+	for element in test.a {
+		out << '${element}'
+	}
+	dump(out)
+	assert out == ['1', '2', '3']
+	out.clear()
+	//
+	for element in test.b {
+		out << '${element}'
+	}
+	dump(out)
+	assert out == ['foo', 'bar']
+	out.clear()
+	//
+	for element in test.c {
+		out << '${element}'
+	}
+	dump(out)
+	assert out == ['1.2', '2.3']
+	out.clear()
+	//
+	for element in test.d {
+		out << '${element}'
+	}
+	dump(out)
+	assert out == ["['foo']", "['bar']"]
+	out.clear()
+}

vlib/v/tests/for_loop_with_optional_test.v

@@ -0,0 +1,35 @@
+struct Test {
+	a ?[]int
+	b ?[]string
+	c ?[]f64
+	d ?[][]string
+}
+
+fn run_loop[U](val U, field_name string) []string {
+	mut out := []string{}
+	$for field in U.fields {
+		variable := val.$(field.name)
+		if field_name == field.name {
+			for element in variable {
+				println(element)
+				out << element.str()
+			}
+		}
+	}
+	return out
+}
+
+fn test_main() {
+	test := Test{
+		a: [1, 2, 3]
+		b: ['foo', 'bar']
+		c: [1.2, 2.3]
+		d: [['foo'], ['bar']]
+	}
+
+	// println(run_loop(test, 'a'))
+	assert run_loop(test, 'a').str() == "['1', '2', '3']"
+	assert run_loop(test, 'b').str() == "['foo', 'bar']"
+	assert run_loop(test, 'c').str() == "['1.2', '2.3']"
+	assert run_loop(test, 'd').str() == "['['foo']', '['bar']']"
+}