From dec44e7f6c6867c06eff192147071fbb12909bee Mon Sep 17 00:00:00 2001
From: Jay Bosamiya <jayb@microsoft.com>
Date: Fri, 20 Sep 2024 15:07:25 -0700
Subject: [PATCH] Improve handling of separated clauses

Specifically, in the presence of comments around `expr_with_block`
cases (such as `assert ... by { ... }` blocks) that don't necessarily
have a semicolon, the previous version would end up reorganizing the
affinity of the comments across the clauses, preferring the previous
clause. With this commit, each comments stays loyal to its own clause.
---
 CHANGELOG.md               |  3 ++
 src/verus.pest             |  6 +++-
 tests/verus-consistency.rs | 60 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8dae773..a5f9fe9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,8 @@
 # Unreleased
 
+* Improved handling of comments around clauses/stanzas
+  - Each comment now maintains loyalty to the clause the user picked it to stay with, rather than automatically migrating to the previous clause in the presence of `assert ... by { ... }`-style constructs
+
 # v0.4.1
 
 * Minor fix to prevent panic on formatting files containing unbalanced parentheses in strings/chars/...
diff --git a/src/verus.pest b/src/verus.pest
index 04e3284..6a223b8 100644
--- a/src/verus.pest
+++ b/src/verus.pest
@@ -988,7 +988,11 @@ stmt = !{
   | proof_block
   | let_stmt
   | assignment_stmt
-  | expr_with_block ~ semi_str?
+  // We can't use `expr_with_block ~ semi_str?` here, and instead need to do it
+  // as the two separate ordered alternatives (`ewb ~ semi_str | ewb`) to
+  // prevent munching of multi-newlines that happens otherwise
+  | expr_with_block ~ semi_str
+  | expr_with_block
   | expr ~ semi_str
   | item_no_macro_call
   | macro_call_stmt
diff --git a/tests/verus-consistency.rs b/tests/verus-consistency.rs
index a82ad24..8cc1585 100644
--- a/tests/verus-consistency.rs
+++ b/tests/verus-consistency.rs
@@ -2453,6 +2453,36 @@ fn fff() {
     baz;
 }
 
+pub fn foo() {
+    // this should stay stuck to the `a`
+    assert(a) by {
+        // whatever
+    }
+    // this should also stay stuck to the `a`
+    // and so should this line
+
+    // but this line
+    // and this line should stick to the `b`
+    b;
+    // and this comment should also stick to the `b`
+
+    // similarly `c`
+    c;
+    // `c` again
+
+    // and an empty comment stands alone
+
+    // similarly `d`
+    d;
+    // `d` again
+
+    // and finally, `d`
+    assert(d) by {
+        // whatever
+    }
+    // well, now done with `d`
+}
+
 } // verus!
 "#;
 
@@ -2470,6 +2500,36 @@ fn fff() {
         baz;
     }
 
+    pub fn foo() {
+        // this should stay stuck to the `a`
+        assert(a) by {
+            // whatever
+        }
+        // this should also stay stuck to the `a`
+        // and so should this line
+
+        // but this line
+        // and this line should stick to the `b`
+        b;
+        // and this comment should also stick to the `b`
+
+        // similarly `c`
+        c;
+        // `c` again
+
+        // and an empty comment stands alone
+
+        // similarly `d`
+        d;
+        // `d` again
+
+        // and finally, `d`
+        assert(d) by {
+            // whatever
+        }
+        // well, now done with `d`
+    }
+
     } // verus!
     "###);
 }