From 38a155ac81884c2d51f6fc68aec09e90397aac9a Mon Sep 17 00:00:00 2001
From: Felipe Pena <felipensp@gmail.com>
Date: Sun, 21 May 2023 10:22:40 -0300
Subject: [PATCH] cgen: make comptime call works with or-block (#18215)

---
 cmd/tools/vast/vast.v                      |  1 +
 vlib/v/ast/ast.v                           |  1 +
 vlib/v/checker/comptime.v                  |  1 +
 vlib/v/fmt/fmt.v                           |  1 +
 vlib/v/gen/c/comptime.v                    | 16 ++++++++++++++--
 vlib/v/parser/comptime.v                   | 13 +++++++++++--
 vlib/v/tests/comptime_call_or_block_test.v | 21 +++++++++++++++++++++
 7 files changed, 50 insertions(+), 4 deletions(-)
 create mode 100644 vlib/v/tests/comptime_call_or_block_test.v

diff --git a/cmd/tools/vast/vast.v b/cmd/tools/vast/vast.v
index 481bf1ad5a..faa7e138d8 100644
--- a/cmd/tools/vast/vast.v
+++ b/cmd/tools/vast/vast.v
@@ -1001,6 +1001,7 @@ fn (t Tree) comptime_call(node ast.ComptimeCall) &Node {
 	obj.add_terse('env_value', t.string_node(node.env_value))
 	obj.add('pos', t.pos(node.pos))
 	obj.add_terse('args', t.array_node_call_arg(node.args))
+	obj.add_terse('or_block', t.or_expr(node.or_block))
 	return obj
 }
 
diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 1cc962ba9a..f4491cdd4a 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -1845,6 +1845,7 @@ pub:
 	is_env       bool
 	env_pos      token.Pos
 	is_pkgconfig bool
+	or_block     OrExpr
 pub mut:
 	left_type   Type
 	result_type Type
diff --git a/vlib/v/checker/comptime.v b/vlib/v/checker/comptime.v
index 6e29d93f5d..e8f95aa410 100644
--- a/vlib/v/checker/comptime.v
+++ b/vlib/v/checker/comptime.v
@@ -141,6 +141,7 @@ fn (mut c Checker) comptime_call(mut node ast.ComptimeCall) ast.Type {
 			// check each arg expression
 			node.args[i].typ = c.expr(arg.expr)
 		}
+		c.stmts_ending_with_expression(node.or_block.stmts)
 		// assume string for now
 		return ast.string_type
 	}
diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v
index 6cdfe2ad09..78a412cff7 100644
--- a/vlib/v/fmt/fmt.v
+++ b/vlib/v/fmt/fmt.v
@@ -1968,6 +1968,7 @@ pub fn (mut f Fmt) comptime_call(node ast.ComptimeCall) {
 			}
 			f.expr(node.left)
 			f.write('.$${method_expr}')
+			f.or_expr(node.or_block)
 		}
 	}
 }
diff --git a/vlib/v/gen/c/comptime.v b/vlib/v/gen/c/comptime.v
index c49fc929b2..0a96569c77 100644
--- a/vlib/v/gen/c/comptime.v
+++ b/vlib/v/gen/c/comptime.v
@@ -161,7 +161,8 @@ fn (mut g Gen) comptime_call(mut node ast.ComptimeCall) {
 			}
 		}
 
-		if !g.inside_call && (m.return_type.has_flag(.option) || m.return_type.has_flag(.result)) {
+		if !g.inside_call && node.or_block.kind != .block
+			&& (m.return_type.has_flag(.option) || m.return_type.has_flag(.result)) {
 			g.write('(*(${g.base_type(m.return_type)}*)')
 		}
 		// TODO: check argument types
@@ -222,9 +223,20 @@ fn (mut g Gen) comptime_call(mut node ast.ComptimeCall) {
 			}
 		}
 		g.write(')')
-		if !g.inside_call && (m.return_type.has_flag(.option) || m.return_type.has_flag(.result)) {
+		if !g.inside_call && node.or_block.kind != .block
+			&& (m.return_type.has_flag(.option) || m.return_type.has_flag(.result)) {
 			g.write('.data)')
 		}
+		if node.or_block.kind != .absent
+			&& (m.return_type.has_flag(.option) || m.return_type.has_flag(.result)) {
+			if !g.inside_assign {
+				cur_line := g.go_before_stmt(0)
+				tmp_var := g.new_tmp_var()
+				g.write('${g.typ(m.return_type)} ${tmp_var} = ')
+				g.write(cur_line)
+				g.or_block(tmp_var, node.or_block, m.return_type)
+			}
+		}
 		return
 	}
 	mut j := 0
diff --git a/vlib/v/parser/comptime.v b/vlib/v/parser/comptime.v
index 6eeff4cca2..3d61e5c335 100644
--- a/vlib/v/parser/comptime.v
+++ b/vlib/v/parser/comptime.v
@@ -357,9 +357,13 @@ fn (mut p Parser) comptime_selector(left ast.Expr) ast.Expr {
 		p.check(.lpar)
 		args := p.call_args()
 		p.check(.rpar)
+		mut or_kind := ast.OrKind.absent
+		mut or_pos := p.tok.pos()
+		mut or_stmts := []ast.Stmt{}
 		if p.tok.kind == .key_orelse {
-			p.next()
-			p.check(.lcbr)
+			// `$method() or {}``
+			or_kind = .block
+			or_stmts, or_pos = p.or_block(.with_err_var)
 		}
 		return ast.ComptimeCall{
 			left: left
@@ -369,6 +373,11 @@ fn (mut p Parser) comptime_selector(left ast.Expr) ast.Expr {
 			args_var: ''
 			args: args
 			pos: start_pos.extend(p.prev_tok.pos())
+			or_block: ast.OrExpr{
+				stmts: or_stmts
+				kind: or_kind
+				pos: or_pos
+			}
 		}
 	}
 	mut has_parens := false
diff --git a/vlib/v/tests/comptime_call_or_block_test.v b/vlib/v/tests/comptime_call_or_block_test.v
new file mode 100644
index 0000000000..ea78b9f3c4
--- /dev/null
+++ b/vlib/v/tests/comptime_call_or_block_test.v
@@ -0,0 +1,21 @@
+pub struct App {
+}
+
+pub fn (mut app App) app_index() ! {
+	return error('hhh')
+}
+
+pub fn (mut app App) no_error() {
+}
+
+fn test_main() {
+	mut app := App{}
+	mut ret2 := ''
+	$for method in App.methods {
+		$if method.is_pub {
+			app.$method() or { ret2 = err.msg() }
+			dump(ret2)
+		}
+	}
+	assert ret2 == 'hhh'
+}