builtin, checker, cgen: implement `x = a[k] or { ... }` for maps and arrays (#8193)

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 ? {

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)

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
 }
 

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 {
+					is_gen_or_and_assign_rhs := gen_or && g.is_assign_rhs
-						g.write('/*2*/string_clone(')
+					cur_line := if is_gen_or_and_assign_rhs {
-					}
+						line := g.go_before_stmt(0)
-					if g.is_fn_index_call {
+						g.out.write(tabs[g.indent])
-						if elem_typ.info is table.FnType {
+						line
-							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('*')
+					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_type := info.value_type
-				elem_typ := g.table.get_type_symbol(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 {
+					is_gen_or_and_assign_rhs := gen_or && g.is_assign_rhs
-						if elem_typ.info is table.FnType {
+					cur_line := if is_gen_or_and_assign_rhs {
-							g.write('((')
+						line := g.go_before_stmt(0)
-							g.write_fn_ptr_decl(&elem_typ.info, '')
+						g.out.write(tabs[g.indent])
-							g.write(')(*(voidptr*)map_get_1(')
+						line
-						}
-					} else if elem_typ.kind == .function {
-						g.write('(*(voidptr*)map_get_1(')
 					} 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(')

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
+}