From d80f5165dd65f7fa1a0caa53975b16618144752e Mon Sep 17 00:00:00 2001
From: Larpon <Larpon@users.noreply.github.com>
Date: Fri, 25 Feb 2022 14:36:48 +0100
Subject: [PATCH] fmt: keep trailing comments after fn header decl (#13596)

---
 vlib/v/ast/ast.v                              |  3 ++-
 vlib/v/fmt/fmt.v                              | 20 +++++++++++++++++++
 .../fn_headers_with_comments_expected.vv      | 16 +++++++++++++++
 .../tests/fn_headers_with_comments_input.vv   |  9 +++++++++
 .../tests/fn_headers_with_comments_keep.vv    |  5 +++++
 vlib/v/parser/fn.v                            |  1 +
 6 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 vlib/v/fmt/tests/fn_headers_with_comments_expected.vv
 create mode 100644 vlib/v/fmt/tests/fn_headers_with_comments_input.vv
 create mode 100644 vlib/v/fmt/tests/fn_headers_with_comments_keep.vv

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 9daee0c62..407f18db3 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -530,7 +530,8 @@ pub mut:
 	has_await         bool // 'true' if this function uses JS.await
 	//
 	comments      []Comment // comments *after* the header, but *before* `{`; used for InterfaceDecl
-	next_comments []Comment // coments that are one line after the decl; used for InterfaceDecl
+	end_comments  []Comment // comments *after* header declarations. E.g.: `fn C.C_func(x int) int // Comment`
+	next_comments []Comment // comments that are one line after the decl; used for InterfaceDecl
 	//
 	source_file &File = 0
 	scope       &Scope
diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v
index 698c61335..82c5c3564 100644
--- a/vlib/v/fmt/fmt.v
+++ b/vlib/v/fmt/fmt.v
@@ -914,6 +914,26 @@ pub fn (mut f Fmt) enum_decl(node ast.EnumDecl) {
 pub fn (mut f Fmt) fn_decl(node ast.FnDecl) {
 	f.attrs(node.attrs)
 	f.write(node.stringify(f.table, f.cur_mod, f.mod2alias)) // `Expr` instead of `ast.Expr` in mod ast
+	// Handle trailing comments after fn header declarations
+	if node.end_comments.len > 0 {
+		first_comment := node.end_comments[0]
+		if first_comment.text.contains('\n') {
+			f.writeln('\n')
+		} else {
+			f.write(' ')
+		}
+		f.comment(first_comment)
+		if node.end_comments.len > 1 {
+			f.writeln('\n')
+			comments := node.end_comments[1..]
+			for i, comment in comments {
+				f.comment(comment)
+				if i != comments.len - 1 {
+					f.writeln('\n')
+				}
+			}
+		}
+	}
 	f.fn_body(node)
 }
 
diff --git a/vlib/v/fmt/tests/fn_headers_with_comments_expected.vv b/vlib/v/fmt/tests/fn_headers_with_comments_expected.vv
new file mode 100644
index 000000000..84cd47472
--- /dev/null
+++ b/vlib/v/fmt/tests/fn_headers_with_comments_expected.vv
@@ -0,0 +1,16 @@
+fn C.Mix_LoadMUS1(file byteptr) voidptr // *Mix_Music
+
+fn C.Mix_LoadMUS2(file byteptr) voidptr //*Mix_Music
+
+fn C.Mix_LoadMUS3(file byteptr) voidptr // 1
+
+// 2
+
+// 3
+
+// Loads music
+fn C.Mix_LoadMUS4(file byteptr) voidptr
+
+/*
+Test
+*/
diff --git a/vlib/v/fmt/tests/fn_headers_with_comments_input.vv b/vlib/v/fmt/tests/fn_headers_with_comments_input.vv
new file mode 100644
index 000000000..9474830d6
--- /dev/null
+++ b/vlib/v/fmt/tests/fn_headers_with_comments_input.vv
@@ -0,0 +1,9 @@
+fn C.Mix_LoadMUS1(file byteptr) voidptr // *Mix_Music
+
+fn C.Mix_LoadMUS2(file byteptr) voidptr /* *Mix_Music */
+
+fn C.Mix_LoadMUS3(file byteptr) voidptr /* 1 */ /* 2 */ /* 3 */
+
+// Loads music
+fn C.Mix_LoadMUS4(file byteptr) voidptr /* Test
+*/
diff --git a/vlib/v/fmt/tests/fn_headers_with_comments_keep.vv b/vlib/v/fmt/tests/fn_headers_with_comments_keep.vv
new file mode 100644
index 000000000..dd4d451c5
--- /dev/null
+++ b/vlib/v/fmt/tests/fn_headers_with_comments_keep.vv
@@ -0,0 +1,5 @@
+fn C.Mix_LoadMUS1(file byteptr) voidptr // Mix_Music
+
+fn C.Mix_LoadMUS2(file byteptr) voidptr // Mix_Music 2
+
+// Comment
diff --git a/vlib/v/parser/fn.v b/vlib/v/parser/fn.v
index 8a09e9e20..4d30f4d57 100644
--- a/vlib/v/parser/fn.v
+++ b/vlib/v/parser/fn.v
@@ -527,6 +527,7 @@ fn (mut p Parser) fn_decl() ast.FnDecl {
 		is_builtin: p.builtin_mod || p.mod in util.builtin_module_parts
 		scope: p.scope
 		label_names: p.label_names
+		end_comments: p.eat_comments(same_line: true)
 	}
 	if generic_names.len > 0 {
 		p.table.register_fn_generic_types(fn_decl.fkey())
-- 
2.39.5