diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v index 61381bfb65..f6510953d0 100644 --- a/vlib/v/ast/ast.v +++ b/vlib/v/ast/ast.v @@ -1048,14 +1048,15 @@ pub: val_is_mut bool // `for mut val in vals {` means that modifying `val` will modify the array // and the array cannot be indexed inside the loop pub mut: - cond Expr - key_type Type - val_type Type - cond_type Type - high_type Type - kind Kind // array/map/string - label string // `label: for {` - scope &Scope + val_is_ref bool // `for val in &arr {` means that value of `val` will be the reference of the value in `arr` + cond Expr + key_type Type + val_type Type + cond_type Type + high_type Type + kind Kind // array/map/string + label string // `label: for {` + scope &Scope } pub struct ForCStmt { diff --git a/vlib/v/checker/for.v b/vlib/v/checker/for.v index d1ac711ac9..1292a758ce 100644 --- a/vlib/v/checker/for.v +++ b/vlib/v/checker/for.v @@ -70,6 +70,22 @@ fn (mut c Checker) for_in_stmt(mut node ast.ForInStmt) { node.scope.update_var_type(node.val_var, node.val_type) } else { sym := c.table.final_sym(typ) + if sym.kind != .string { + match mut node.cond { + ast.PrefixExpr { + node.val_is_ref = node.cond.op == .amp + } + ast.Ident { + match mut node.cond.info { + ast.IdentVar { + node.val_is_ref = !node.cond.is_mut() && node.cond.info.typ.is_ptr() + } + else {} + } + } + else {} + } + } if sym.kind == .struct_ { // iterators next_fn := sym.find_method_with_generic_parent('next') or { @@ -148,6 +164,8 @@ fn (mut c Checker) for_in_stmt(mut node ast.ForInStmt) { } else {} } + } else if node.val_is_ref { + value_type = value_type.ref() } node.cond_type = typ node.kind = sym.kind diff --git a/vlib/v/gen/c/for.v b/vlib/v/gen/c/for.v index b91c2ccd34..e895ac0137 100644 --- a/vlib/v/gen/c/for.v +++ b/vlib/v/gen/c/for.v @@ -178,7 +178,7 @@ 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 { + right := if node.val_is_mut || node.val_is_ref { '(($styp)$cond_var${op_field}data) + $i' } else { '(($styp*)$cond_var${op_field}data)[$i]' @@ -295,7 +295,7 @@ fn (mut g Gen) for_in_stmt(node ast.ForInStmt) { } else { val_styp := g.typ(node.val_type) if node.val_type.is_ptr() { - if node.val_is_mut { + if node.val_is_mut || node.val_is_ref { g.write('$val_styp ${c_name(node.val_var)} = &(*($val_styp)') } else { g.write('$val_styp ${c_name(node.val_var)} = (*($val_styp*)') diff --git a/vlib/v/tests/for_in_ref_arr_test.v b/vlib/v/tests/for_in_ref_arr_test.v new file mode 100644 index 0000000000..8dab119fc0 --- /dev/null +++ b/vlib/v/tests/for_in_ref_arr_test.v @@ -0,0 +1,31 @@ +fn test_for_in_ref_val_ref_arr() { + arr := [1, 2, 3, 4, 5] + mut rets := []&int{} + mut expects := unsafe { []&int{len: 5, init: &arr[it]} } + + for val in &arr { + println(val) + rets << val + } + + for i, val in &arr { + assert voidptr(val) == voidptr(rets[i]) + } + assert rets == expects +} + +fn test_for_in_ref_val_ref_arr_ident() { + arr_ := [1, 2, 3, 4, 5] + arr := &arr_ + mut rets := []&int{} + mut expects := unsafe { []&int{len: 5, init: &arr_[it]} } + + for val in arr { + rets << val + } + + for i, val in arr { + assert voidptr(val) == voidptr(rets[i]) + } + assert rets == expects +} diff --git a/vlib/v/tests/for_in_ref_map_test.v b/vlib/v/tests/for_in_ref_map_test.v new file mode 100644 index 0000000000..4f20db6e88 --- /dev/null +++ b/vlib/v/tests/for_in_ref_map_test.v @@ -0,0 +1,48 @@ +fn test_for_in_ref_val_ref_map() { + mut mp := { + 'a': 1 + 'b': 2 + 'c': 3 + 'd': 4 + 'e': 5 + } + mut rets := map[string]&int{} + mut expects := map[string]&int{} + + for k, mut val in mp { + expects[k] = unsafe { val } + } + for k, val in &mp { + rets[k] = unsafe { val } + } + + for k, val in &mp { + assert voidptr(val) == voidptr(expects[k]) + } + assert rets == expects +} + +fn test_for_in_ref_val_ref_map_ident() { + mut mp_ := { + 'a': 1 + 'b': 2 + 'c': 3 + 'd': 4 + 'e': 5 + } + mp := &mp_ + mut rets := map[string]&int{} + mut expects := map[string]&int{} + + for k, mut val in mp_ { + expects[k] = unsafe { val } + } + for k, val in mp { + rets[k] = unsafe { val } + } + + for k, val in mp { + assert voidptr(val) == voidptr(expects[k]) + } + assert rets == expects +} diff --git a/vlib/v/tests/inout/printing_for_v_in_a.out b/vlib/v/tests/inout/printing_for_v_in_a.out index 4f19f936f5..a7182a7c21 100644 --- a/vlib/v/tests/inout/printing_for_v_in_a.out +++ b/vlib/v/tests/inout/printing_for_v_in_a.out @@ -2,10 +2,10 @@ 98 99 abc -23 -77 +&23 +&77 abc -> k: abc | v: xyz -> k: def | v: jkl +> k: abc | v: &xyz +> k: def | v: &jkl abc abc