diff --git a/vlib/builtin/array.v b/vlib/builtin/array.v index 241fb787b4..85f6213c58 100644 --- a/vlib/builtin/array.v +++ b/vlib/builtin/array.v @@ -229,6 +229,16 @@ fn (a array) get(i int) voidptr { } } +// Private function. Used to implement x = a[i] or { ... } +fn (a array) get_with_check(i int) voidptr { + if i < 0 || i >= a.len { + return 0 + } + unsafe { + return byteptr(a.data) + i * a.element_size + } +} + // first returns the first element of the array. pub fn (a array) first() voidptr { $if !no_bounds_checking ? { diff --git a/vlib/builtin/map.v b/vlib/builtin/map.v index 7e749c3594..0f68737595 100644 --- a/vlib/builtin/map.v +++ b/vlib/builtin/map.v @@ -552,6 +552,29 @@ fn (m &map) get_1(key voidptr, zero voidptr) voidptr { return zero } +// If `key` matches the key of an element in the container, +// the method returns a reference to its mapped value. +// If not, a zero pointer is returned. +// This is used in `x := m['key'] or { ... }` +fn (m &map) get_1_check(key voidptr) voidptr { + mut index, mut meta := m.key_to_index(key) + for { + if meta == unsafe { m.metas[index] } { + kv_index := int(unsafe { m.metas[index + 1] }) + pkey := unsafe { m.key_values.key(kv_index) } + if m.key_eq_fn(key, pkey) { + return unsafe { byteptr(pkey) + m.key_values.key_bytes } + } + } + index += 2 + meta += probe_inc + if meta > unsafe { m.metas[index] } { + break + } + } + return 0 +} + // Checks whether a particular key exists in the map. fn (m &map) exists_1(key voidptr) bool { mut index, mut meta := m.key_to_index(key) diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v index a2346fa769..e05ea5fdc8 100644 --- a/vlib/v/checker/checker.v +++ b/vlib/v/checker/checker.v @@ -4804,7 +4804,7 @@ fn (mut c Checker) check_index_type(typ_sym &table.TypeSymbol, index_type table. } pub fn (mut c Checker) index_expr(mut node ast.IndexExpr) table.Type { - typ := c.expr(node.left) + mut typ := c.expr(node.left) node.left_type = typ typ_sym := c.table.get_final_type_symbol(typ) if typ_sym.kind !in [.array, .array_fixed, .string, .map] && !typ.is_ptr() && typ !in [table.byteptr_type, table.charptr_type] && @@ -4844,9 +4844,10 @@ pub fn (mut c Checker) index_expr(mut node ast.IndexExpr) table.Type { if typ_sym.kind == .array_fixed { elem_type := c.table.value_type(typ) idx := c.table.find_or_register_array(elem_type) - return table.new_type(idx) + typ = table.new_type(idx) + } else { + typ = typ.set_nr_muls(0) } - return typ.set_nr_muls(0) } else { // [1] index_type := c.expr(node.index) if typ_sym.kind == .map { @@ -4860,9 +4861,10 @@ pub fn (mut c Checker) index_expr(mut node ast.IndexExpr) table.Type { } value_type := c.table.value_type(typ) if value_type != table.void_type { - return value_type + typ = value_type } } + c.stmts(node.or_expr.stmts) return typ } diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v index 8245304523..e1a686c855 100644 --- a/vlib/v/gen/cgen.v +++ b/vlib/v/gen/cgen.v @@ -3915,6 +3915,7 @@ fn (mut g Gen) if_expr(node ast.IfExpr) { } fn (mut g Gen) index_expr(node ast.IndexExpr) { + gen_or := node.or_expr.kind != .absent match node.index { ast.RangeExpr { sym := g.table.get_type_symbol(node.left_type) @@ -3972,7 +3973,8 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) { if sym.kind == .array { info := sym.info as table.Array elem_type_str := g.typ(info.elem_type) - elem_typ := g.table.get_type_symbol(info.elem_type) + elem_type := info.elem_type + elem_typ := g.table.get_type_symbol(elem_type) // `vals[i].field = x` is an exception and requires `array_get`: // `(*(Val*)array_get(vals, i)).field = x;` is_selector := node.left is ast.SelectorExpr @@ -4051,21 +4053,35 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) { } needs_clone := info.elem_type == table.string_type_idx && g.is_autofree && !g.is_assign_lhs - if needs_clone { - g.write('/*2*/string_clone(') - } - if g.is_fn_index_call { - if elem_typ.info is table.FnType { - g.write('((') - g.write_fn_ptr_decl(&elem_typ.info, '') - g.write(')(*($array_ptr_type_str)/*ee elem_typ */array_get(') - } - } else if is_direct_array_access { - g.write('(($array_ptr_type_str)') + is_gen_or_and_assign_rhs := gen_or && g.is_assign_rhs + cur_line := if is_gen_or_and_assign_rhs { + line := g.go_before_stmt(0) + g.out.write(tabs[g.indent]) + line } else { - g.write('(*($array_ptr_type_str)/*ee elem_typ */array_get(') - if left_is_ptr && !node.left_type.has_flag(.shared_f) { - g.write('*') + '' + } + tmp_opt := if gen_or { g.new_tmp_var() } else { '' } + tmp_opt_ptr := if gen_or { g.new_tmp_var() } else { '' } + if gen_or { + g.write('$array_ptr_type_str $tmp_opt_ptr = ($array_ptr_type_str)/*ee elem_ptr_typ */(array_get_with_check(') + } else { + if needs_clone { + g.write('/*2*/string_clone(') + } + if g.is_fn_index_call { + if elem_typ.info is table.FnType { + g.write('((') + g.write_fn_ptr_decl(&elem_typ.info, '') + g.write(')(*($array_ptr_type_str)/*ee elem_typ */array_get(') + } + } else if is_direct_array_access { + g.write('(($array_ptr_type_str)') + } else { + g.write('(*($array_ptr_type_str)/*ee elem_typ */array_get(') + if left_is_ptr && !node.left_type.has_flag(.shared_f) { + g.write('*') + } } } g.expr(node.left) @@ -4098,12 +4114,26 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) { if needs_clone { g.write(')') } + if gen_or { + g.writeln(';') + opt_elem_type := g.typ(elem_type.set_flag(.optional)) + g.writeln('$opt_elem_type $tmp_opt = {0};') + g.writeln('if ($tmp_opt_ptr) {') + g.writeln('\t${tmp_opt}.ok = true; ${tmp_opt}.is_none = false; ${tmp_opt}.v_error = (string){.str=(byteptr)""}; ${tmp_opt}.ecode = 0;') + g.writeln('\t*(($elem_type_str*)&${tmp_opt}.data) = *(($elem_type_str*)$tmp_opt_ptr);') + g.writeln('} else {') + g.writeln('\t${tmp_opt}.ok = false; ${tmp_opt}.is_none = false; ${tmp_opt}.v_error = (string){.str=(byteptr)"array index out of range"}; ${tmp_opt}.ecode = 0;') + g.writeln('}') + g.or_block(tmp_opt, node.or_expr, elem_type) + g.write('\n$cur_line*($elem_type_str*)${tmp_opt}.data') + } } } else if sym.kind == .map { info := sym.info as table.Map key_type_str := g.typ(info.key_type) - elem_type_str := g.typ(info.value_type) - elem_typ := g.table.get_type_symbol(info.value_type) + elem_type := info.value_type + elem_type_str := g.typ(elem_type) + elem_typ := g.table.get_type_symbol(elem_type) get_and_set_types := elem_typ.kind in [.struct_, .map] if g.is_assign_lhs && !g.is_array_set && !get_and_set_types { if g.assign_op == .assign || info.value_type == table.string_type { @@ -4149,16 +4179,30 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) { g.write('}, &($elem_type_str[]){ $zero }))') } else { zero := g.type_default(info.value_type) - if g.is_fn_index_call { - if elem_typ.info is table.FnType { - g.write('((') - g.write_fn_ptr_decl(&elem_typ.info, '') - g.write(')(*(voidptr*)map_get_1(') - } - } else if elem_typ.kind == .function { - g.write('(*(voidptr*)map_get_1(') + is_gen_or_and_assign_rhs := gen_or && g.is_assign_rhs + cur_line := if is_gen_or_and_assign_rhs { + line := g.go_before_stmt(0) + g.out.write(tabs[g.indent]) + line } else { - g.write('(*($elem_type_str*)map_get_1(') + '' + } + tmp_opt := if gen_or { g.new_tmp_var() } else { '' } + tmp_opt_ptr := if gen_or { g.new_tmp_var() } else { '' } + if gen_or { + g.write('$elem_type_str* $tmp_opt_ptr = ($elem_type_str*)/*ee elem_ptr_typ */(map_get_1_check(') + } else { + if g.is_fn_index_call { + if elem_typ.info is table.FnType { + g.write('((') + g.write_fn_ptr_decl(&elem_typ.info, '') + g.write(')(*(voidptr*)map_get_1(') + } + } else if elem_typ.kind == .function { + g.write('(*(voidptr*)map_get_1(') + } else { + g.write('(*($elem_type_str*)map_get_1(') + } } if !left_is_ptr || node.left_type.has_flag(.shared_f) { g.write('ADDR(map, ') @@ -4177,13 +4221,28 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) { g.write('), &($key_type_str[]){') g.expr(node.index) g.write('}') - if g.is_fn_index_call { + if gen_or { + g.write('))') + } else if g.is_fn_index_call { g.write(', &(voidptr[]){ $zero })))') } else if elem_typ.kind == .function { g.write(', &(voidptr[]){ $zero }))') } else { g.write(', &($elem_type_str[]){ $zero }))') } + if gen_or { + g.writeln(';') + opt_elem_type := g.typ(elem_type.set_flag(.optional)) + g.writeln('$opt_elem_type $tmp_opt = {0};') + g.writeln('if ($tmp_opt_ptr) {') + g.writeln('\t${tmp_opt}.ok = true; ${tmp_opt}.is_none = false; ${tmp_opt}.v_error = (string){.str=(byteptr)""}; ${tmp_opt}.ecode = 0;') + g.writeln('\t*(($elem_type_str*)&${tmp_opt}.data) = *(($elem_type_str*)$tmp_opt_ptr);') + g.writeln('} else {') + g.writeln('\t${tmp_opt}.ok = false; ${tmp_opt}.is_none = false; ${tmp_opt}.v_error = (string){.str=(byteptr)"array index out of range"}; ${tmp_opt}.ecode = 0;') + g.writeln('}') + g.or_block(tmp_opt, node.or_expr, elem_type) + g.write('\n$cur_line*($elem_type_str*)${tmp_opt}.data') + } } } else if sym.kind == .string && !node.left_type.is_ptr() { g.write('string_at(') diff --git a/vlib/v/tests/array_map_or_test.v b/vlib/v/tests/array_map_or_test.v new file mode 100644 index 0000000000..ba76ffeeca --- /dev/null +++ b/vlib/v/tests/array_map_or_test.v @@ -0,0 +1,27 @@ +fn test_array_or() { + m := [3, 4, 5] + mut testvar := 17 + el := m[4] or { + testvar = -43 + } + good := m[1] or { + testvar = 11 + } + assert testvar == -43 + assert el == 0 + assert good == 4 +} + +fn test_map_or() { + m := {'as': 3, 'qw': 4, 'kl': 5} + mut testvar := -21 + el := m['pp'] or { + testvar = 7931 + } + good := m['kl'] or { + testvar = -45 + } + assert testvar == 7931 + assert el == 0 + assert good == 5 +}