ast, parser, fmt: fixed and optimized comments for 'for, for_c, for_in stmts' (fix: #15922) (#15950)

vlib/v/ast/ast.v

@@ -1028,8 +1028,9 @@ pub mut:
 
 pub struct ForStmt {
 pub:
-	is_inf bool // `for {}`
+	is_inf   bool // `for {}`
-	pos    token.Pos
+	pos      token.Pos
+	comments []Comment
 pub mut:
 	cond  Expr
 	stmts []Stmt
@@ -1046,6 +1047,7 @@ pub:
 	high       Expr // `10` in `for i in 0..10 {`
 	stmts      []Stmt
 	pos        token.Pos
+	comments   []Comment
 	val_is_mut bool // `for mut val in vals {` means that modifying `val` will modify the array
 	// and the array cannot be indexed inside the loop
 pub mut:
@@ -1067,6 +1069,7 @@ pub:
 	has_inc  bool
 	is_multi bool // for a,b := 0,1; a < 10; a,b = a+b, a {...}
 	pos      token.Pos
+	comments []Comment
 pub mut:
 	init  Stmt // i := 0;
 	cond  Expr // i < 10;

vlib/v/fmt/fmt.v

@@ -1008,6 +1008,9 @@ pub fn (mut f Fmt) for_c_stmt(node ast.ForCStmt) {
 	if node.label.len > 0 {
 		f.write('$node.label: ')
 	}
+	if node.comments.len > 0 {
+		f.comments(node.comments)
+	}
 	f.write('for ')
 	if node.has_init {
 		f.single_line_if = true // to keep all for ;; exprs on the same line
@@ -1031,6 +1034,9 @@ pub fn (mut f Fmt) for_in_stmt(node ast.ForInStmt) {
 	if node.label.len > 0 {
 		f.write('$node.label: ')
 	}
+	if node.comments.len > 0 {
+		f.comments(node.comments)
+	}
 	f.write('for ')
 	if node.key_var != '' {
 		f.write(node.key_var)
@@ -1062,15 +1068,13 @@ pub fn (mut f Fmt) for_stmt(node ast.ForStmt) {
 	if node.label.len > 0 {
 		f.write('$node.label: ')
 	}
-	if node.cond is ast.Comment {
+	if node.comments.len > 0 {
-		f.comments([node.cond])
+		f.comments(node.comments)
 	}
 	f.write('for ')
-	if node.cond !is ast.Comment {
+	f.expr(node.cond)
-		f.expr(node.cond)
+	if !node.is_inf {
-		if !node.is_inf {
+		f.write(' ')
-			f.write(' ')
-		}
 	}
 	f.write('{')
 	if node.stmts.len > 0 || node.pos.line_nr < node.pos.last_line {

vlib/v/fmt/tests/comments_expected.vv

@@ -108,3 +108,43 @@ fn between_if_branches() {
 for {
 	break
 }
+
+// comments
+for {
+	break
+}
+
+// comments
+for i := 0; i < 1; i++ {
+	break
+}
+
+// comments
+for i := 0; i < 1; i++ {
+	break
+}
+
+// comments1
+// comments2
+// comments3
+// comments4
+for i := 0; i < 1; i++ {
+	break
+}
+
+// comments
+for _ in [1, 2] {
+	break
+}
+
+// comments
+for _ in [1, 2] {
+	break
+}
+
+// comments1
+// comments2
+// comments3
+for _ in [1, 2] {
+	break
+}

vlib/v/fmt/tests/comments_input.vv

@@ -92,3 +92,33 @@ for // comments
 {
 	break
 }
+
+for /* comments */ {
+	break
+}
+
+for i := 0; i < 1; i++ // comments
+{
+	break
+}
+
+for i := 0; i < 1; i++ /* comments */ {
+	break
+}
+
+for /* comments1 */ i := 0; /* comments2 */ i < 1; /* comments3 */ i++ /* comments4 */ {
+	break
+}
+
+for _ in [1, 2] /* comments */ {
+	break
+}
+
+for _ in [1, 2] // comments
+{
+	break
+}
+
+for /* comments1 */ _ in /* comments2 */ [1, 2] /* comments3 */ {
+	break
+}

vlib/v/parser/for.v

@@ -15,6 +15,8 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 	}
 	// defer { p.close_scope() }
 	// Infinite loop
+	mut comments := []ast.Comment{}
+	comments << p.eat_comments()
 	if p.tok.kind == .lcbr {
 		p.inside_for = false
 		stmts := p.parse_block_no_scope(false)
@@ -22,6 +24,7 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 		for_stmt := ast.ForStmt{
 			stmts: stmts
 			pos: pos
+			comments: comments
 			is_inf: true
 			scope: p.scope
 		}
@@ -46,6 +49,7 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 			init = p.assign_stmt()
 			has_init = true
 		}
+		comments << p.eat_comments()
 		// Allow `for ;; i++ {`
 		// Allow `for i = 0; i < ...`
 		p.check(.semicolon)
@@ -54,17 +58,21 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 			if p.tok.kind == .name && p.peek_tok.kind in [.inc, .dec] {
 				return p.error('cannot use $p.tok.lit$p.peek_tok.kind as value')
 			}
+			comments << p.eat_comments()
 			cond = p.expr(0)
 			has_cond = true
 		}
+		comments << p.eat_comments()
 		p.check(.semicolon)
 		if !is_multi {
 			is_multi = p.peek_tok.kind == .comma
 		}
+		comments << p.eat_comments()
 		if p.tok.kind != .lcbr {
 			inc = p.stmt(false)
 			has_inc = true
 		}
+		comments << p.eat_comments()
 		p.inside_for = false
 		stmts := p.parse_block_no_scope(false)
 		pos.update_last_line(p.prev_tok.line_nr)
@@ -78,6 +86,7 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 			cond: cond
 			inc: inc
 			pos: pos
+			comments: comments
 			scope: p.scope
 		}
 		p.close_scope()
@@ -128,10 +137,12 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 			return p.error_with_pos('redefinition of value iteration variable `$val_var_name`, use `for ($val_var_name in array) {` if you want to check for a condition instead',
 				val_var_pos)
 		}
+		comments << p.eat_comments()
 		p.check(.key_in)
 		if p.tok.kind == .name && p.tok.lit in [key_var_name, val_var_name] {
 			return p.error('in a `for x in array` loop, the key or value iteration variable `$p.tok.lit` can not be the same as the array variable')
 		}
+		comments << p.eat_comments()
 		// arr_expr
 		cond := p.expr(0)
 		// 0 .. 10
@@ -168,6 +179,7 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 				is_stack_obj: true
 			})
 		}
+		comments << p.eat_comments()
 		p.inside_for = false
 		stmts := p.parse_block_no_scope(false)
 		pos.update_last_line(p.prev_tok.line_nr)
@@ -180,6 +192,7 @@ fn (mut p Parser) for_stmt() ast.Stmt {
 			high: high_expr
 			is_range: is_range
 			pos: pos
+			comments: comments
 			val_is_mut: val_is_mut
 			scope: p.scope
 		}