diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v index bc7ebe959e..9f8a2f21da 100644 --- a/vlib/v/checker/checker.v +++ b/vlib/v/checker/checker.v @@ -2824,32 +2824,53 @@ fn (mut c Checker) stmt(node ast.Stmt) { } } else { sym := c.table.get_type_symbol(typ) - if sym.kind == .map && !(node.key_var.len > 0 && node.val_var.len > 0) { - c.error('declare a key and a value variable when ranging a map: `for key, val in map {`\n' + - 'use `_` if you do not need the variable', node.pos) - } - if node.key_var.len > 0 { - key_type := match sym.kind { - .map { sym.map_info().key_type } - else { table.int_type } - } - node.key_type = key_type - node.scope.update_var_type(node.key_var, key_type) - } - mut value_type := c.table.value_type(typ) - if value_type == table.void_type || typ.has_flag(.optional) { - if typ != table.void_type { - c.error('for in: cannot index `${c.table.type_to_str(typ)}`', + if sym.kind == .struct_ { + // iterators + next_fn := sym.find_method('next') or { + c.error('a struct must have a `next()` method to be an iterator', node.cond.position()) + return } + if !next_fn.return_type.has_flag(.optional) { + c.error('iterator method `next()` must return an optional', node.cond.position()) + } + // the receiver + if next_fn.params.len != 1 { + c.error('iterator method `next()` must have 0 parameters', node.cond.position()) + } + val_type := next_fn.return_type.clear_flag(.optional) + node.cond_type = typ + node.kind = sym.kind + node.val_type = val_type + node.scope.update_var_type(node.val_var, val_type) + } else { + if sym.kind == .map && !(node.key_var.len > 0 && node.val_var.len > 0) { + c.error('declare a key and a value variable when ranging a map: `for key, val in map {`\n' + + 'use `_` if you do not need the variable', node.pos) + } + if node.key_var.len > 0 { + key_type := match sym.kind { + .map { sym.map_info().key_type } + else { table.int_type } + } + node.key_type = key_type + node.scope.update_var_type(node.key_var, key_type) + } + mut value_type := c.table.value_type(typ) + if value_type == table.void_type || typ.has_flag(.optional) { + if typ != table.void_type { + c.error('for in: cannot index `${c.table.type_to_str(typ)}`', + node.cond.position()) + } + } + if node.val_is_mut { + value_type = value_type.to_ptr() + } + node.cond_type = typ + node.kind = sym.kind + node.val_type = value_type + node.scope.update_var_type(node.val_var, value_type) } - if node.val_is_mut { - value_type = value_type.to_ptr() - } - node.cond_type = typ - node.kind = sym.kind - node.val_type = value_type - node.scope.update_var_type(node.val_var, value_type) } c.check_loop_label(node.label, node.pos) c.stmts(node.stmts) diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v index a159592360..67626eb2f7 100644 --- a/vlib/v/gen/cgen.v +++ b/vlib/v/gen/cgen.v @@ -1305,6 +1305,27 @@ fn (mut g Gen) for_in(it ast.ForInStmt) { g.expr(it.cond) g.writeln('.str[$i];') } + } else if it.kind == .struct_ { + cond_type_sym := g.table.get_type_symbol(it.cond_type) + next_fn := cond_type_sym.find_method('next') or { + verror('`next` method not found') + return + } + ret_typ := next_fn.return_type + g.writeln('while (1) {') + t := g.new_tmp_var() + receiver_styp := g.typ(next_fn.params[0].typ) + fn_name := receiver_styp.replace_each(['*', '', '.', '__']) + '_next' + g.write('\t${g.typ(ret_typ)} $t = ${fn_name}(') + if !it.cond_type.is_ptr() { + g.write('&') + } + g.expr(it.cond) + g.writeln(');') + g.writeln('\tif (!${t}.ok) { break; }') + val := if it.val_var in ['', '_'] { g.new_tmp_var() } else { it.val_var } + val_styp := g.typ(it.val_type) + g.writeln('\t$val_styp $val = *($val_styp*)${t}.data;') } else { s := g.table.type_to_str(it.cond_type) g.error('for in: unhandled symbol `$it.cond` of type `$s`', it.pos) diff --git a/vlib/v/tests/for-in-iterator_test.v b/vlib/v/tests/for-in-iterator_test.v new file mode 100644 index 0000000000..62e7c1d0a3 --- /dev/null +++ b/vlib/v/tests/for-in-iterator_test.v @@ -0,0 +1,36 @@ +struct Doubler { +mut: + val int + until int +} + +fn (mut it Doubler) next() ?int { + v := it.val + if v > it.until { + return none + } + it.val *= 2 + return v +} + +fn doubler(start int, until int) Doubler { + return Doubler{start, until} +} + +fn test_for_in_iterator() { + mut d := doubler(5, 30) + mut vals := []int{} + for val in d { + vals << val + } + assert vals == [5, 10, 20] +} + +fn test_for_in_empty_iterator() { + mut d := doubler(5, 2) + mut vals := []int{} + for val in d { + vals << val + } + assert vals == [] +}