From 4aea890736921d9a90cc1f8deb2c60dcbe34368e Mon Sep 17 00:00:00 2001
From: Travis Hance <tjhance7@gmail.com>
Date: Tue, 14 Jan 2025 19:07:44 -0500
Subject: [PATCH] Support `returns` clause (#98)

Co-authored-by: Jay Bosamiya <jayb@microsoft.com>
---
 CHANGELOG.md               |  2 +
 src/lib.rs                 |  5 +++
 src/verus.pest             | 12 +++++-
 tests/verus-consistency.rs | 86 +++++++++++++++++++++++++++++++++++++-
 4 files changed, 103 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f87ba16..78d6a58 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,7 @@
 # Unreleased
 
+* Add support for `returns` clauses (see [verus#1283](https://github.com/verus-lang/verus/pull/1283))
+
 * Add support for `assume_specification` (see [verus#1368](https://github.com/verus-lang/verus/pull/1368))
 
 # v0.5.0
diff --git a/src/lib.rs b/src/lib.rs
index 80294b7..e96265a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -768,6 +768,7 @@ fn to_doc<'a>(
 
         Rule::decreases_str
         | Rule::ensures_str
+        | Rule::returns_str
         | Rule::invariant_except_break_str
         | Rule::invariant_str
         | Rule::invariant_ensures_str
@@ -1345,9 +1346,13 @@ fn to_doc<'a>(
         Rule::groupable_comma_delimited_exprs_for_verus_clauses => {
             map_to_doc(ctx, arena, pair).nest(INDENT_SPACES).group()
         }
+        Rule::single_expr_for_verus_clause => {
+            comma_delimited_full(ctx, arena, pair).nest(INDENT_SPACES)
+        }
         Rule::verus_clause_non_expr => map_to_doc(ctx, arena, pair),
         Rule::requires_clause => map_to_doc(ctx, arena, pair),
         Rule::ensures_clause => map_to_doc(ctx, arena, pair),
+        Rule::returns_clause => map_to_doc(ctx, arena, pair),
         Rule::invariant_except_break_clause => map_to_doc(ctx, arena, pair),
         Rule::invariant_clause => map_to_doc(ctx, arena, pair),
         Rule::invariant_ensures_clause => map_to_doc(ctx, arena, pair),
diff --git a/src/verus.pest b/src/verus.pest
index b406788..bf174a8 100644
--- a/src/verus.pest
+++ b/src/verus.pest
@@ -267,6 +267,7 @@ do_str = ${ "do" ~ !("_" | ASCII_ALPHANUMERIC) }
 dyn_str = ${ "dyn" ~ !("_" | ASCII_ALPHANUMERIC) }
 else_str = ${ "else" ~ !("_" | ASCII_ALPHANUMERIC) }
 ensures_str = ${ "ensures" ~ !("_" | ASCII_ALPHANUMERIC) }
+returns_str = ${ "returns" ~ !("_" | ASCII_ALPHANUMERIC) }
 enum_str = ${ "enum" ~ !("_" | ASCII_ALPHANUMERIC) }
 exec_str = ${ "exec" ~ !("_" | ASCII_ALPHANUMERIC) }
 exists_str = ${ "exists" ~ !("_" | ASCII_ALPHANUMERIC) }
@@ -713,7 +714,7 @@ use_tree_list = {
 }
 
 fn_qualifier = {
-    (requires_clause | recommends_clause | ensures_clause | decreases_clause | opens_invariants_clause | unwind_clause)*
+    (requires_clause | recommends_clause | ensures_clause | returns_clause | decreases_clause | opens_invariants_clause | unwind_clause)*
 }
 
 fn_terminator = {
@@ -1744,10 +1745,15 @@ groupable_comma_delimited_exprs_for_verus_clauses = {
     comma_delimited_exprs_for_verus_clauses
 }
 
+single_expr_for_verus_clause = {
+    !verus_clause_non_expr ~ expr_no_struct ~ ","?
+}
+
 verus_clause_non_expr = _{
     "{"             // TODO: Why is this permitted here?
   | requires_str
   | ensures_str
+  | returns_str
   | invariant_except_break_str
   | invariant_str
   | invariant_ensures_str
@@ -1766,6 +1772,10 @@ ensures_clause = {
     ensures_str ~ comma_delimited_exprs_for_verus_clauses?
 }
 
+returns_clause = {
+    returns_str ~ single_expr_for_verus_clause
+}
+
 invariant_except_break_clause = {
     invariant_except_break_str ~ comma_delimited_exprs_for_verus_clauses?
 }
diff --git a/tests/verus-consistency.rs b/tests/verus-consistency.rs
index a1b1d33..e13d439 100644
--- a/tests/verus-consistency.rs
+++ b/tests/verus-consistency.rs
@@ -511,7 +511,7 @@ verus!{
             cong;
             done;
         ");
-    
+
     fn foo(x: usize) {
         match x {
             inj_ord_choice_pat!((_,x), *, *) => (),
@@ -2667,3 +2667,87 @@ pub assume_specification<T, const N: usize> [ <[T; N]>::as_slice ](ar: &[T; N])
     } // verus!
     ")
 }
+
+#[test]
+fn verus_support_returns_clause() {
+    let file = r#"
+verus!{
+    fn test(b: bool) -> (c: bool)
+        requires x,
+        ensures c == !b,
+        returns !b,
+    {
+    }
+
+    fn test2(b: bool) -> (c: bool)
+        requires x,
+        ensures c == !b,
+        returns !b
+    {
+    }
+
+    fn test3(b: bool) -> (c: bool)
+        requires x,
+        ensures c == !b,
+        returns !b
+        opens_invariants any
+    {
+    }
+
+    fn test4(b: bool) -> (c: bool)
+        requires x,
+        ensures c == !b,
+        returns !b,
+        opens_invariants any
+    {
+    }
+}
+"#;
+    assert_snapshot!(parse_and_format(file).unwrap(), @r"
+    verus! {
+
+    fn test(b: bool) -> (c: bool)
+        requires
+            x,
+        ensures
+            c == !b,
+        returns
+            !b,
+    {
+    }
+
+    fn test2(b: bool) -> (c: bool)
+        requires
+            x,
+        ensures
+            c == !b,
+        returns
+            !b,
+    {
+    }
+
+    fn test3(b: bool) -> (c: bool)
+        requires
+            x,
+        ensures
+            c == !b,
+        returns
+            !b,
+        opens_invariants any
+    {
+    }
+
+    fn test4(b: bool) -> (c: bool)
+        requires
+            x,
+        ensures
+            c == !b,
+        returns
+            !b,
+        opens_invariants any
+    {
+    }
+
+    } // verus!
+    ")
+}