From 2c3949ab2c6725b9d67fbc9c84c0f834d8f46d25 Mon Sep 17 00:00:00 2001
From: Travis Hance <tjhance7@gmail.com>
Date: Tue, 1 Oct 2024 09:12:18 -0400
Subject: [PATCH] returns clause

---
 src/lib.rs                 |  2 +
 src/verus.pest             |  8 +++-
 tests/verus-consistency.rs | 84 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 5ea07e1..90a2f08 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
@@ -1335,6 +1336,7 @@ fn to_doc<'a>(
         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 cb33fe8..a07a54e 100644
--- a/src/verus.pest
+++ b/src/verus.pest
@@ -260,6 +260,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) }
@@ -701,7 +702,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 = {
@@ -1715,6 +1716,7 @@ 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
@@ -1733,6 +1735,10 @@ ensures_clause = {
     ensures_str ~ comma_delimited_exprs_for_verus_clauses?
 }
 
+returns_clause = {
+    returns_str ~ comma_delimited_exprs_for_verus_clauses?
+}
+
 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 7c2c4dc..91e9295 100644
--- a/tests/verus-consistency.rs
+++ b/tests/verus-consistency.rs
@@ -2601,3 +2601,87 @@ verus! {
     } // 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!
+    "###)
+}