From 5b93b4f37dfccae9c33ee68032a70194208da6f9 Mon Sep 17 00:00:00 2001
From: Delyan Angelov <delian66@gmail.com>
Date: Sat, 4 Jul 2020 16:13:58 +0300
Subject: [PATCH] vfmt: allow multiple comments after if branches

---
 vlib/v/ast/ast.v   |  8 ++++----
 vlib/v/fmt/fmt.v   |  6 +++---
 vlib/v/parser/if.v | 13 ++++++-------
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 4e39ec8c22..daf6015832 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -436,10 +436,10 @@ pub mut:
 
 pub struct IfBranch {
 pub:
-	cond    Expr
-	stmts   []Stmt
-	pos     token.Position
-	comment Comment
+	cond     Expr
+	stmts    []Stmt
+	pos      token.Position
+	comments []Comment
 }
 
 pub struct LockExpr {
diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v
index 6efb4d39ec..9eafc0ff12 100644
--- a/vlib/v/fmt/fmt.v
+++ b/vlib/v/fmt/fmt.v
@@ -836,7 +836,7 @@ pub fn (mut f Fmt) expr(node ast.Expr) {
 			f.write(node.val)
 		}
 		ast.LockExpr {
-			f.lock_expr(node) 
+			f.lock_expr(node)
 		}
 		ast.MapInit {
 			if node.keys.len == 0 {
@@ -1191,8 +1191,8 @@ pub fn (mut f Fmt) if_expr(it ast.IfExpr) {
 		(it.is_expr || f.is_assign)
 	f.single_line_if = single_line
 	for i, branch in it.branches {
-		if branch.comment.text != '' {
-			f.comment(branch.comment)
+		if branch.comments.len > 0 {
+			f.comments(branch.comments, true, .keep)
 		}
 		if i == 0 {
 			f.write('if ')
diff --git a/vlib/v/parser/if.v b/vlib/v/parser/if.v
index 449211f938..88d09767b3 100644
--- a/vlib/v/parser/if.v
+++ b/vlib/v/parser/if.v
@@ -16,17 +16,14 @@ fn (mut p Parser) if_expr() ast.IfExpr {
 	pos := p.tok.position()
 	mut branches := []ast.IfBranch{}
 	mut has_else := false
+	mut comments := []ast.Comment{}
 	for p.tok.kind in [.key_if, .key_else] {
 		p.inside_if = true
 		start_pos := p.tok.position()
-		mut comment := ast.Comment{}
 		if p.tok.kind == .key_if {
 			p.next()
 		} else {
-			// if p.tok.kind == .comment {
-			// p.error('place comments inside {}')
-			// }
-			comment = p.check_comment()
+			comments = p.eat_comments()
 			p.check(.key_else)
 			if p.tok.kind == .key_if {
 				p.next()
@@ -37,8 +34,9 @@ fn (mut p Parser) if_expr() ast.IfExpr {
 				branches << ast.IfBranch{
 					stmts: p.parse_block()
 					pos: start_pos.extend(end_pos)
-					comment: comment
+					comments: comments
 				}
+				comments = []
 				break
 			}
 		}
@@ -74,8 +72,9 @@ fn (mut p Parser) if_expr() ast.IfExpr {
 			cond: cond
 			stmts: stmts
 			pos: start_pos.extend(end_pos)
-			comment: comment
+			comments: comments
 		}
+		comments = []
 		if p.tok.kind != .key_else {
 			break
 		}