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

pull/8197/head
Uwe Krüger 2021-01-19 06:06:57 +01:00 committed by GitHub
parent a65b73d3e4
commit b74690cbec
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 152 additions and 31 deletions

View File

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

View File

@ -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)

View File

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

View File

@ -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(')

View File

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