From a1f123bd4242aae97ba34da15ffd2719eaeb08ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Uwe=20Kr=C3=BCger?=
 <45282134+UweKrueger@users.noreply.github.com>
Date: Mon, 5 Jul 2021 20:00:30 +0200
Subject: [PATCH] cgen: enable string index error handling `s[i] or {...}`
 (#10670)

---
 vlib/builtin/string.v               | 11 ++++++
 vlib/v/gen/c/index.v                | 27 ++++++++++++---
 vlib/v/tests/string_index_or_test.v | 52 +++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 5 deletions(-)
 create mode 100644 vlib/v/tests/string_index_or_test.v

diff --git a/vlib/builtin/string.v b/vlib/builtin/string.v
index e12cca7ddd..78b06bd25f 100644
--- a/vlib/builtin/string.v
+++ b/vlib/builtin/string.v
@@ -1263,6 +1263,17 @@ fn (s string) at(idx int) byte {
 	}
 }
 
+// version of `at()` that is used in `a[i] or {`
+// return an error when the index is out of range
+fn (s string) at_with_check(idx int) ?byte {
+	if idx < 0 || idx >= s.len {
+		return error('string index out of range')
+	}
+	unsafe {
+		return s.str[idx]
+	}
+}
+
 // is_space returns `true` if the byte is a white space character.
 // The following list is considered white space characters: ` `, `\t`, `\n`, `\v`, `\f`, `\r`, 0x85, 0xa0
 // Example: assert byte(` `).is_space() == true
diff --git a/vlib/v/gen/c/index.v b/vlib/v/gen/c/index.v
index e676428463..d7cb79c32f 100644
--- a/vlib/v/gen/c/index.v
+++ b/vlib/v/gen/c/index.v
@@ -25,11 +25,28 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) {
 				g.expr(node.index)
 				g.write(']')
 			} else {
-				g.write('string_at(')
-				g.expr(node.left)
-				g.write(', ')
-				g.expr(node.index)
-				g.write(')')
+				gen_or := node.or_expr.kind != .absent || node.is_option
+				if gen_or {
+					tmp_opt := g.new_tmp_var()
+					cur_line := g.go_before_stmt(0)
+					g.out.write_string(util.tabs(g.indent))
+					opt_elem_type := g.typ(ast.byte_type.set_flag(.optional))
+					g.write('$opt_elem_type $tmp_opt = string_at_with_check(')
+					g.expr(node.left)
+					g.write(', ')
+					g.expr(node.index)
+					g.writeln(');')
+					if !node.is_option {
+						g.or_block(tmp_opt, node.or_expr, ast.byte_type)
+					}
+					g.write('\n$cur_line*(byte*)&${tmp_opt}.data')
+				} else {
+					g.write('string_at(')
+					g.expr(node.left)
+					g.write(', ')
+					g.expr(node.index)
+					g.write(')')
+				}
 			}
 		} else {
 			g.expr(node.left)
diff --git a/vlib/v/tests/string_index_or_test.v b/vlib/v/tests/string_index_or_test.v
new file mode 100644
index 0000000000..be48207b94
--- /dev/null
+++ b/vlib/v/tests/string_index_or_test.v
@@ -0,0 +1,52 @@
+fn test_empty_string_access() {
+	mut res := ''
+	a := ''
+	if a[0] or { `0` } == `1` {
+		res = 'good'
+	} else {
+		res = 'bad'
+	}
+	assert res == 'bad'
+}
+
+fn test_good_string_access() {
+	mut res := ''
+	a := 'abcd'
+	if a[2] or { `0` } == `c` {
+		res = 'good'
+	} else {
+		res = 'bad'
+	}
+	assert res == 'good'
+}
+
+fn test_if_guard_bad() {
+	mut res := 'bad'
+	a := 'xy'
+	if z := a[2] {
+		res = '${z:c}'
+	}
+	assert res == 'bad'
+}
+
+fn test_if_guard_good() {
+	mut res := 'bad'
+	a := 'xy123'
+	if z := a[2] {
+		res = '${z:c}'
+	}
+	assert res == '1'
+}
+
+fn get_propagate(s string, i int) ?string {
+	c := s[i] ?
+	return 'got `${c:c}`'
+}
+
+fn test_propagation() {
+	s := 'abcd'
+	x := get_propagate(s, 2) or { '$err' }
+	y := get_propagate(s, 5) or { '$err' }
+	assert x == 'got `c`'
+	assert y == 'string index out of range'
+}