vfmt: handle comments after type declarations (#6874)

vlib/v/ast/ast.v

@@ -738,6 +738,7 @@ pub:
 	is_pub      bool
 	parent_type table.Type
 	pos         token.Position
+	comments    []Comment
 }
 
 pub struct SumTypeDecl {
@@ -746,6 +747,7 @@ pub:
 	is_pub    bool
 	sub_types []table.Type
 	pos       token.Position
+	comments  []Comment
 }
 
 // New implementation of sum types
@@ -754,16 +756,18 @@ pub:
 	name      string
 	is_pub    bool
 	pos       token.Position
+	comments  []Comment
 pub mut:
 	sub_types []table.Type
 }
 
 pub struct FnTypeDecl {
 pub:
-	name   string
-	is_pub bool
-	typ    table.Type
-	pos    token.Position
+	name     string
+	is_pub   bool
+	typ      table.Type
+	pos      token.Position
+	comments []Comment
 }
 
 // TODO: handle this differently

vlib/v/fmt/fmt.v

@@ -495,6 +495,7 @@ pub fn (mut f Fmt) stmt(node ast.Stmt) {
 }
 
 pub fn (mut f Fmt) type_decl(node ast.TypeDecl) {
+	mut comments := []ast.Comment{}
 	match node {
 		ast.AliasTypeDecl {
 			if node.is_pub {
@@ -502,6 +503,7 @@ pub fn (mut f Fmt) type_decl(node ast.TypeDecl) {
 			}
 			ptype := f.table.type_to_str(node.parent_type)
 			f.write('type $node.name = $ptype')
+			comments << node.comments
 		}
 		ast.FnTypeDecl {
 			if node.is_pub {
@@ -544,6 +546,7 @@ pub fn (mut f Fmt) type_decl(node ast.TypeDecl) {
 			} else if fn_info.return_type.has_flag(.optional) {
 				f.write(' ?')
 			}
+			comments << node.comments
 		}
 		ast.SumTypeDecl {
 			if node.is_pub {
@@ -563,6 +566,7 @@ pub fn (mut f Fmt) type_decl(node ast.TypeDecl) {
 				f.wrap_long_line(2, true)
 			}
 			// f.write(sum_type_names.join(' | '))
+			comments << node.comments
 		}
 		ast.UnionSumTypeDecl {
 			if node.is_pub {
@@ -582,8 +586,15 @@ pub fn (mut f Fmt) type_decl(node ast.TypeDecl) {
 				f.wrap_long_line(2, true)
 			}
 			// f.write(sum_type_names.join(' | '))
+			comments << node.comments
 		}
 	}
+	if comments.len > 0 {
+		f.write(' ')
+		f.comments(comments, CommentsOptions{
+			has_nl: false
+		})
+	}
 	f.writeln('\n')
 }
 

vlib/v/fmt/tests/types_expected.vv

@@ -3,7 +3,7 @@ type FooBar = Bar | Foo
 
 pub type PublicBar = Bar | Foo | FooBar
 
-type Uint = byte | u16 | u32 | u64
+type Uint = byte | u16 | u32 | u64 // This should stay on the same line
 
 type Float = f32 | f64
 
@@ -19,7 +19,7 @@ type OneArgFn = fn (i int)
 
 type TwoDiffArgs = fn (i int, s string) bool
 
-type TwoSameArgs = fn (i int, j int) string
+type TwoSameArgs = fn (i int, j int) string // And a comment
 
 type VarArgs = fn (s ...string) int
 

vlib/v/fmt/tests/types_input.vv

@@ -6,7 +6,7 @@
 
 type   Uint =   u16 | u64
 	| u32
-	    | byte
+	    | byte // This should stay on the same line
 type
 Float =
 	f32 |
@@ -27,7 +27,7 @@ type TwoDiffArgs
 = fn (i    int, s string) bool
 
 
-	type TwoSameArgs = fn(i int, j int) string
+	type TwoSameArgs = fn(i int, j int) string // And a comment
 
 type VarArgs = fn
 (s ...string) int

vlib/v/parser/parser.v

@@ -557,6 +557,17 @@ pub fn (mut p Parser) eat_comments() []ast.Comment {
 	return comments
 }
 
+pub fn (mut p Parser) eat_lineend_comments() []ast.Comment {
+	mut comments := []ast.Comment{}
+	for {
+		if p.tok.kind != .comment || p.tok.line_nr != p.prev_tok.line_nr {
+			break
+		}
+		comments << p.comment()
+	}
+	return comments
+}
+
 pub fn (mut p Parser) stmt(is_top_level bool) ast.Stmt {
 	$if trace_parser ? {
 		tok_pos := p.tok.position()
@@ -1929,15 +1940,18 @@ fn (mut p Parser) type_decl() ast.TypeDecl {
 	}
 	mut sum_variants := []table.Type{}
 	p.check(.assign)
+	mut comments := []ast.Comment{}
 	if p.tok.kind == .key_fn {
 		// function type: `type mycallback fn(string, int)`
 		fn_name := p.prepend_mod(name)
 		fn_type := p.parse_fn_type(fn_name)
+		comments = p.eat_lineend_comments()
 		return ast.FnTypeDecl{
 			name: fn_name
 			is_pub: is_pub
 			typ: fn_type
 			pos: decl_pos
+			comments: comments
 		}
 	}
 	first_type := p.parse_type() // need to parse the first type before we can check if it's `type A = X | Y`
@@ -1964,11 +1978,13 @@ fn (mut p Parser) type_decl() ast.TypeDecl {
 			}
 			is_public: is_pub
 		})
+		comments = p.eat_lineend_comments()
 		return ast.SumTypeDecl{
 			name: name
 			is_pub: is_pub
 			sub_types: sum_variants
 			pos: decl_pos
+			comments: comments
 		}
 	}
 	// type MyType int
@@ -1995,11 +2011,13 @@ fn (mut p Parser) type_decl() ast.TypeDecl {
 		}
 		is_public: is_pub
 	})
+	comments = p.eat_lineend_comments()
 	return ast.AliasTypeDecl{
 		name: name
 		is_pub: is_pub
 		parent_type: parent_type
 		pos: decl_pos
+		comments: comments
 	}
 }