From a331abf675ace4cef6d9761adc45fc1184562c1f Mon Sep 17 00:00:00 2001
From: Joe Conigliaro <joe.conigliaro@gmail.com>
Date: Sat, 21 Mar 2020 00:39:56 +1100
Subject: [PATCH] ast/cgen/parser: combine IfExpr branches into single node

---
 vlib/v/ast/ast.v         |  23 +++++---
 vlib/v/checker/checker.v |  46 +++++++--------
 vlib/v/fmt/fmt.v         |  42 +++++++-------
 vlib/v/gen/cgen.v        |  90 +++++++++++++++--------------
 vlib/v/gen/jsgen.v       |  28 ++++++++--
 vlib/v/gen/tests/2.c     |   2 +-
 vlib/v/parser/parser.v   | 118 ++++++++++++++++++---------------------
 7 files changed, 182 insertions(+), 167 deletions(-)

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index d1df59a823..54a63770b8 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -336,16 +336,21 @@ mut:
 
 pub struct IfExpr {
 pub:
-	tok_kind   token.Kind
-	cond       Expr
-	stmts      []Stmt
-	else_stmts []Stmt
-	left       Expr // `a` in `a := if ...`
-	pos        token.Position
+	tok_kind token.Kind
+	branches []IfBranch
+	left     Expr // `a` in `a := if ...`
+	pos      token.Position
 mut:
-	is_expr    bool
-	typ        table.Type
-	has_else   bool
+	is_expr  bool
+	typ      table.Type
+	has_else bool
+}
+
+pub struct IfBranch {
+pub:
+	cond  Expr
+	stmts []Stmt
+	pos   token.Position
 }
 
 pub struct MatchExpr {
diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 16ee005287..642b6a7479 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -909,32 +909,34 @@ pub fn (c mut Checker) if_expr(node mut ast.IfExpr) table.Type {
 		// println('$c.file.path  $node.pos.line_nr IF: checker exp type = ' + sym.name)
 		node.is_expr = true
 	}
-	typ := c.expr(node.cond)
 	node.typ = table.void_type
-	// node.typ = typ
-	typ_sym := c.table.get_type_symbol(typ)
-	// if typ_sym.kind != .bool {
-	if table.type_idx(typ) != table.bool_type_idx {
-		c.error('non-bool (`$typ_sym.name`) used as if condition', node.pos)
-	}
-	c.stmts(node.stmts)
-	if node.else_stmts.len > 0 {
-		c.stmts(node.else_stmts)
-	}
-	if node.stmts.len > 0 {
-		match node.stmts[node.stmts.len - 1] {
-			ast.ExprStmt {
-				// type_sym := p.table.get_type_symbol(it.typ)
-				// p.warn('if expr ret $type_sym.name')
-				t := c.expr(it.expr)
-				node.typ = t
-				return t
+	for i, branch in node.branches {
+		typ := c.expr(branch.cond)
+		if i < node.branches.len-1 || !node.has_else {
+			typ_sym := c.table.get_type_symbol(typ)
+			// if typ_sym.kind != .bool {
+			if table.type_idx(typ) != table.bool_type_idx {
+				c.error('non-bool (`$typ_sym.name`) used as if condition', node.pos)
 			}
-			else {}
+		}
+		c.stmts(branch.stmts)
 	}
+	if node.has_else && node.is_expr {
+		last_branch := node.branches[node.branches.len-1]
+		if last_branch.stmts.len > 0 {
+			match last_branch.stmts[last_branch.stmts.len - 1] {
+				ast.ExprStmt {
+					// type_sym := p.table.get_type_symbol(it.typ)
+					// p.warn('if expr ret $type_sym.name')
+					t := c.expr(it.expr)
+					node.typ = t
+					return t
+				}
+				else {}
+			}
+		}
 	}
-	return typ
-	// return table.void_type
+	return table.bool_type
 }
 
 pub fn (c mut Checker) postfix_expr(node ast.PostfixExpr) table.Type {
diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v
index df8c40d4df..0f81ff4ae0 100644
--- a/vlib/v/fmt/fmt.v
+++ b/vlib/v/fmt/fmt.v
@@ -378,38 +378,38 @@ fn (f mut Fmt) expr(node ast.Expr) {
 			f.write(it.val)
 		}
 		ast.IfExpr {
-			single_line := it.stmts.len == 1 && it.else_stmts.len == 1 && it.typ != table.void_type
+			single_line := it.branches.len == 2 && it.has_else //
+				&& it.branches[0].stmts.len == 1 && it.branches[1].stmts.len == 1
 			f.single_line_if = single_line
-			f.write('if ')
-			f.expr(it.cond)
-			if single_line {
-				f.write(' { ')
-			}
-			else {
-				f.writeln(' {')
-			}
-			f.stmts(it.stmts)
-			if single_line {
-				f.write(' ')
-			}
-			f.write('}')
-			if it.has_else {
-				f.write(' else ')
-			}
-			else if it.else_stmts.len > 0 {
-				f.write(' else {')
+			for i, branch in it.branches {
+				if i == 0 {
+					f.write('if ')
+					f.expr(branch.cond)
+					f.write(' {')
+				}
+				else if i < it.branches.len-1 || !it.has_else {
+					f.write('} else if ')
+					f.expr(branch.cond)
+					f.write(' {')
+				}
+				else if i == it.branches.len-1 && it.has_else {
+					f.write('} else {')
+				}
 				if single_line {
 					f.write(' ')
 				}
 				else {
 					f.writeln('')
 				}
-				f.stmts(it.else_stmts)
+				f.stmts(branch.stmts)
 				if single_line {
 					f.write(' ')
 				}
-				f.write('}')
 			}
+			// if !single_line {
+			// 	f.writeln('')
+			// }
+			f.write('}')
 			f.single_line_if = false
 		}
 		ast.Ident {
diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v
index b0204377f5..ead5c1dd7f 100644
--- a/vlib/v/gen/cgen.v
+++ b/vlib/v/gen/cgen.v
@@ -1176,64 +1176,68 @@ fn (g mut Gen) if_expr(node ast.IfExpr) {
 	}
 	// one line ?:
 	// TODO clean this up once `is` is supported
-	if node.is_expr && node.stmts.len == 1 && node.else_stmts.len == 1 && type_sym.kind != .void {
-		cond := node.cond
-		stmt1 := node.stmts[0]
-		else_stmt1 := node.else_stmts[0]
-		match stmt1 {
-			ast.ExprStmt {
-				g.inside_ternary = true
-				g.expr(cond)
-				g.write(' ? ')
-				expr_stmt := stmt1 as ast.ExprStmt
-				g.expr(expr_stmt.expr)
+	// TODO: make sure only one stmt in eac branch
+	if node.is_expr && node.branches.len >= 2 && node.has_else && type_sym.kind != .void {
+		g.inside_ternary = true
+		for i, branch in node.branches {
+			if i > 0 {
 				g.write(' : ')
-				g.stmt(else_stmt1)
 			}
-			else {}
-	}
+			if i < node.branches.len-1 || !node.has_else {
+				g.expr(branch.cond)
+				g.write(' ? ')
+			}
+			g.stmts(branch.stmts)
+		}
+		g.inside_ternary = false
 	}
 	else {
-		mut is_guard := false
 		guard_ok := g.new_tmp_var()
-		match node.cond {
-			ast.IfGuardExpr {
-				is_guard = true
-				g.writeln('bool $guard_ok;')
-				g.write('{ /* if guard */ ${g.typ(it.expr_type)} $it.var_name = ')
-				g.expr(it.expr)
-				g.writeln(';')
-				g.writeln('if (($guard_ok = ${it.var_name}.ok)) {')
+		mut is_guard := false
+		for i, branch in node.branches {
+			if i == 0 {
+				match branch.cond {
+					ast.IfGuardExpr {
+						is_guard = true
+						g.writeln('bool $guard_ok;')
+						g.write('{ /* if guard */ ${g.typ(it.expr_type)} $it.var_name = ')
+						g.expr(it.expr)
+						g.writeln(';')
+						g.writeln('if (($guard_ok = ${it.var_name}.ok)) {')
+					}
+					else {
+						g.write('if (')
+						g.expr(branch.cond)
+						g.writeln(') {')
+					}
+				}
+	
 			}
-			else {
-				g.inside_ternary = false
-				g.write('if (')
-				g.expr(node.cond)
-				g.writeln(') {')
+			else if i < node.branches.len-1 || !node.has_else {
+				g.writeln('} else if (')
+				g.expr(branch.cond)
+				g.write(') {')
+			}
+			else if i == node.branches.len-1 && node.has_else {
+				if is_guard {
+					g.writeln('} if (!$guard_ok) { /* else */')
+				}
+				else {
+					g.writeln('} else {')
+				}
 			}
-	}
-		for i, stmt in node.stmts {
 			// Assign ret value
-			if i == node.stmts.len - 1 && type_sym.kind != .void {}
+			// if i == node.stmts.len - 1 && type_sym.kind != .void {}
 			// g.writeln('$tmp =')
-			g.stmt(stmt)
+			// 
+			g.stmts(branch.stmts)
+			// g.writeln('')
 		}
 		if is_guard {
 			g.write('}')
 		}
 		g.writeln('}')
-		if node.else_stmts.len > 0 {
-			if is_guard {
-				g.writeln('if (!$guard_ok) { /* else */')
-			}
-			else {
-				g.writeln('else { ')
-			}
-			g.stmts(node.else_stmts)
-			g.writeln('}')
-		}
 	}
-	g.inside_ternary = false
 }
 
 fn (g mut Gen) index_expr(node ast.IndexExpr) {
diff --git a/vlib/v/gen/jsgen.v b/vlib/v/gen/jsgen.v
index fc3ea73394..67e2f32892 100644
--- a/vlib/v/gen/jsgen.v
+++ b/vlib/v/gen/jsgen.v
@@ -34,6 +34,12 @@ pub fn (g mut JsGen) writeln(s string) {
 	g.out.writeln(s)
 }
 
+fn (g mut JsGen) stmts(stmts []ast.Stmt) {
+	for stmt in stmts {
+		g.stmt(stmt)
+	}
+}
+
 fn (g mut JsGen) stmt(node ast.Stmt) {
 	match node {
 		ast.FnDecl {
@@ -163,13 +169,23 @@ fn (g mut JsGen) expr(node ast.Expr) {
 			}
 		}
 		ast.IfExpr {
-			g.write('if (')
-			g.expr(it.cond)
-			g.writeln(') {')
-			for stmt in it.stmts {
-				g.stmt(stmt)
+			for i, branch in it.branches {
+				if i == 0 {
+					g.write('if (')
+					g.expr(branch.cond)
+					g.writeln(') {')
+				}
+				else if i < it.branches.len-1 || !it.has_else {
+					g.write('else if (')
+					g.expr(branch.cond)
+					g.writeln(') {')
+				}
+				else {
+					g.write('else {')
+				}
+				g.stmts(branch.stmts)
+				g.writeln('}')
 			}
-			g.writeln('}')
 		}
 		else {
 			println(term.red('jsgen.expr(): bad node'))
diff --git a/vlib/v/gen/tests/2.c b/vlib/v/gen/tests/2.c
index cd8c6e017b..cab9c0f71a 100644
--- a/vlib/v/gen/tests/2.c
+++ b/vlib/v/gen/tests/2.c
@@ -52,7 +52,7 @@ void function2() {
 	}
 	if (false) {
 		foo(1);
-	}	else {
+	} else {
 		puts(tos3("else"));
 		foo(100);
 	}
diff --git a/vlib/v/parser/parser.v b/vlib/v/parser/parser.v
index 2b4a0f3905..d70a0c2d0c 100644
--- a/vlib/v/parser/parser.v
+++ b/vlib/v/parser/parser.v
@@ -1175,82 +1175,70 @@ fn (p mut Parser) for_statement() ast.Stmt {
 	}
 }
 
-fn (p mut Parser) if_expr() ast.Expr {
+fn (p mut Parser) if_expr() ast.IfExpr {
 	p.inside_if = true
-	// defer {
-	// }
-	mut node := ast.Expr{}
-	p.check(.key_if)
 	pos := p.tok.position()
-	// `if x := opt() {`
-	mut cond := ast.Expr{}
-	mut is_or := false
-	if p.peek_tok.kind == .decl_assign {
-		is_or = true
-		p.open_scope()
-		var_name := p.check_name()
-		p.check(.decl_assign)
-		expr := p.expr(0)
-		p.scope.register_var(ast.Var{
-			name: var_name
-			expr: expr
-		})
-		cond = ast.IfGuardExpr{
-			var_name: var_name
-			expr: expr
-		}
-	}
-	else {
-		cond = p.expr(0)
-	}
-	p.inside_if = false
+	mut branches := []ast.IfBranch
 	mut has_else := false
-	stmts := p.parse_block()
-	mut else_stmts := []ast.Stmt
-	if p.tok.kind == .key_else {
-		p.check(.key_else)
+	for p.tok.kind in [.key_if, .key_else] {
+		branch_pos := p.tok.position()
 		if p.tok.kind == .key_if {
-			// The next if block is handled by next if_expr()
-			has_else = true
+			p.check(.key_if)
 		}
-		// p.if_expr()
 		else {
-			else_stmts = p.parse_block()
+			p.check(.key_else)
+			if p.tok.kind == .key_if {
+				p.check(.key_if)
+			}
+			else {
+				has_else = true
+				branches << ast.IfBranch{
+					stmts: p.parse_block()
+					pos: branch_pos
+				}
+				break
+			}
+		}
+		mut cond := ast.Expr{}
+		mut is_or := false
+		// `if x := opt() {`
+		if p.peek_tok.kind == .decl_assign {
+			is_or = true
+			p.open_scope()
+			var_name := p.check_name()
+			p.check(.decl_assign)
+			expr := p.expr(0)
+			p.scope.register_var(ast.Var{
+				name: var_name
+				expr: expr
+			})
+			cond = ast.IfGuardExpr{
+				var_name: var_name
+				expr: expr
+			}
+		}
+		else {
+			cond = p.expr(0)
+		}
+		p.inside_if = false
+		stmts := p.parse_block()
+		if is_or {
+			p.close_scope()
+		}
+		branches << ast.IfBranch{
+			cond: cond
+			stmts: stmts
+			pos: branch_pos
+		}
+		if p.tok.kind != .key_else {
+			break
 		}
 	}
-	if is_or {
-		p.close_scope()
-	}
-	// mut typ := table.void_type
-	// mut left := ast.Expr{}
-	// If the last statement is an expression, return its type
-	/*
-	if stmts.len > 0 {
-		match stmts[stmts.len - 1] {
-			ast.ExprStmt {
-				type_sym := p.table.get_type_symbol(it.typ)
-				p.warn('if expr ret $type_sym.name')
-				typ = it.typ
-				// return node,it.ti
-				// left =
-			}
-			else {}
-	}
-	}
-	*/
-
-	node = ast.IfExpr{
-		cond: cond
-		stmts: stmts
-		else_stmts: else_stmts
-		// typ: typ
-		
+	return ast.IfExpr{
+		branches: branches
 		pos: pos
 		has_else: has_else
-		// left: left
-		
 	}
-	return node
 }
 
 fn (p mut Parser) string_expr() ast.Expr {