From 66570a0b9c8d9187161c0fb3dc08209d3ea3eb70 Mon Sep 17 00:00:00 2001
From: Bryan Parno <parno@cmu.edu>
Date: Wed, 13 Dec 2023 10:43:23 -0500
Subject: [PATCH] Don't add a newline after a final single-line prefix comment

---
 examples/syntax.rs         |  1 -
 examples/wip.rs            |  9 +--------
 src/lib.rs                 | 11 +++++++++--
 tests/verus-consistency.rs |  8 ++++++++
 4 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/examples/syntax.rs b/examples/syntax.rs
index dd1cf79..1373524 100644
--- a/examples/syntax.rs
+++ b/examples/syntax.rs
@@ -14,7 +14,6 @@ verus! {
 ///   - proof code: erased before compilation, may have requires/ensures
 ///   - spec code: erased before compilation, no requires/ensures, but may have recommends
 /// exec and proof functions may name their return values inside parentheses, before the return type
-
 fn my_exec_fun(x: u32, y: u32) -> (sum: u32)
     requires
         x < 100,
diff --git a/examples/wip.rs b/examples/wip.rs
index d31bb6b..b65ef5b 100644
--- a/examples/wip.rs
+++ b/examples/wip.rs
@@ -1,14 +1,7 @@
 verus! {
 
 // The first set of &&& get out-dented, while the second one is absorbed into a single line
-fn clone_value() -> (out: u8)
-    ensures
-        match value {
-            None => {
-                &&& out.is_None()
-            },
-        },
-{
+fn clone_value() -> (out: u8) {
 }
 
 } // verus!
diff --git a/src/lib.rs b/src/lib.rs
index 6d6b5e2..70933ca 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1127,6 +1127,11 @@ pub fn rustfmt(value: &str) -> Option<String> {
     None
 }
 
+fn is_multiline_comment(pair: &Pair<Rule>) -> bool {
+    matches!(&pair.as_span().as_str()[..2], "/*")
+}
+
+
 pub fn parse_and_format(s: &str) -> Result<String, pest::error::Error<Rule>> {
     let ctx = Context {
         inline_comment_lines: find_inline_comment_lines(s),
@@ -1169,10 +1174,12 @@ pub fn parse_and_format(s: &str) -> Result<String, pest::error::Error<Rule>> {
                     (prefix_comments, body, suffix_comments)
                 };
                 formatted_output += VERUS_PREFIX;
+                let mut final_prefix_comment_is_multiline = false;
                 for comment in &prefix_comments {
                     formatted_output += comment.as_str();
+                    final_prefix_comment_is_multiline = is_multiline_comment(&comment);
                 }
-                if prefix_comments.len() > 0 {
+                if prefix_comments.len() > 0 && final_prefix_comment_is_multiline {
                     formatted_output += "\n";
                 }
                 let items = body.into_inner();
@@ -1187,7 +1194,7 @@ pub fn parse_and_format(s: &str) -> Result<String, pest::error::Error<Rule>> {
                         }
                         prev_is_use = false;
                         formatted_output += item.as_str();
-                        if matches!(&item.as_span().as_str()[..2], "/*") {
+                        if is_multiline_comment(&item) {
                             formatted_output += "\n";
                         }
                     } else {
diff --git a/tests/verus-consistency.rs b/tests/verus-consistency.rs
index ec4c490..181908e 100644
--- a/tests/verus-consistency.rs
+++ b/tests/verus-consistency.rs
@@ -599,6 +599,10 @@ fn verus_comments() {
 /// External comment 2
 verus! {
 
+// This comment should stay with the function definition, even though it's parsed as a "prefix" comment
+fn clone_value() -> (out: u8) {
+}
+
 spec fn my_spec_fun() -> int {
     2
 }
@@ -707,6 +711,10 @@ impl AbstractEndPoint {
     /// External comment 2
     verus! {
 
+    // This comment should stay with the function definition, even though it's parsed as a "prefix" comment
+    fn clone_value() -> (out: u8) {
+    }
+
     spec fn my_spec_fun() -> int {
         2
     }