From 99d7031949bd35704639861edea56bef66914dfd Mon Sep 17 00:00:00 2001
From: Jay Bosamiya <jayb@microsoft.com>
Date: Sat, 11 Jan 2025 19:38:31 -0800
Subject: [PATCH] Add support for `assume_specification`

---
 CHANGELOG.md               |  2 ++
 src/lib.rs                 | 13 +++++++++++++
 src/verus.pest             | 27 +++++++++++++++++++++++++--
 tests/verus-consistency.rs | 23 +++++++++++++++++++++++
 4 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1662950..f87ba16 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,7 @@
 # Unreleased
 
+* Add support for `assume_specification` (see [verus#1368](https://github.com/verus-lang/verus/pull/1368))
+
 # v0.5.0
 
 * Improve handling of inner-docstring comments (i.e., `///`-prefixed comments)
diff --git a/src/lib.rs b/src/lib.rs
index 180d849..80294b7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -781,6 +781,7 @@ fn to_doc<'a>(
         Rule::any_str
         | Rule::assert_str
         | Rule::assume_str
+        | Rule::assume_specification_str
         | Rule::checked_str
         | Rule::choose_str
         | Rule::exec_str
@@ -802,6 +803,7 @@ fn to_doc<'a>(
         Rule::path_no_generics => map_to_doc(ctx, arena, pair),
         Rule::path_segment => map_to_doc(ctx, arena, pair),
         Rule::path_segment_no_generics => map_to_doc(ctx, arena, pair),
+        Rule::path_segment_type => map_to_doc(ctx, arena, pair),
         Rule::generic_arg_list => map_to_doc(ctx, arena, pair),
         Rule::generic_arg_list_with_colons => map_to_doc(ctx, arena, pair),
         Rule::generic_args => comma_delimited(ctx, arena, pair, false).group(),
@@ -948,6 +950,17 @@ fn to_doc<'a>(
                 }
             }))
         }
+        Rule::assume_specification => arena.concat(pair.into_inner().map(|p| {
+            let rule = p.as_rule();
+            let d = to_doc(ctx, p, arena);
+            match rule {
+                Rule::semi_str => arena.hardline().append(d),
+                _ => d,
+            }
+        })),
+        Rule::assume_specification_for => {
+            map_to_doc(ctx, arena, pair).enclose(arena.space(), arena.space())
+        }
         Rule::abi => map_to_doc(ctx, arena, pair).append(arena.space()),
         Rule::param_list => comma_delimited(ctx, arena, pair, false).parens().group(),
         Rule::closure_param_list => comma_delimited(ctx, arena, pair, false)
diff --git a/src/verus.pest b/src/verus.pest
index fda124f..b406788 100644
--- a/src/verus.pest
+++ b/src/verus.pest
@@ -244,6 +244,7 @@ any_str = ${ "any" ~ !("_" | ASCII_ALPHANUMERIC) }
 assert_space_str = ${ "assert" ~ !("_" | ASCII_ALPHANUMERIC) }
 assert_str = ${ "assert" ~ !("_" | ASCII_ALPHANUMERIC) }
 assume_str = ${ "assume" ~ !("_" | ASCII_ALPHANUMERIC) }
+assume_specification_str = ${ "assume_specification" ~ !("_" | ASCII_ALPHANUMERIC) }
 async_str = ${ "async" ~ !("_" | ASCII_ALPHANUMERIC) }
 as_str = ${ "as" ~ !("_" | ASCII_ALPHANUMERIC) }
 auto_str = ${ "auto" ~ !("_" | ASCII_ALPHANUMERIC) }
@@ -446,14 +447,18 @@ path_segment = {
     colons_str? ~ name_ref ~ generic_arg_list?
   | colons_str? ~ name_ref ~ param_list ~ ret_type?
   | colons_str? ~ name_ref
-  | langle_str ~ path_type ~ (as_str ~ path_type)? ~ rangle_str
+  | langle_str ~ path_segment_type ~ (as_str ~ path_segment_type)? ~ rangle_str
 }
 
 path_segment_no_generics = {
     colons_str? ~ name_ref ~ generic_arg_list_with_colons?
   | colons_str? ~ name_ref ~ param_list ~ ret_type?
   | colons_str? ~ name_ref
-  | langle_str ~ path_type ~ (as_str ~ path_type)? ~ rangle_str
+  | langle_str ~ path_segment_type ~ (as_str ~ path_segment_type)? ~ rangle_str
+}
+
+path_segment_type = {
+    type
 }
 
 generic_args = {
@@ -631,6 +636,7 @@ item_no_macro_call = _{
   | enum
   | extern_block
   | extern_crate
+  | assume_specification
   | fn
   | global
   | impl
@@ -724,6 +730,23 @@ fn = {
     fn_terminator
 }
 
+assume_specification = {
+    attr* ~
+    visibility? ~
+    assume_specification_str ~
+    generic_param_list? ~
+    lbracket_str ~ assume_specification_for ~ rbracket_str ~
+    param_list ~
+    where_clause? ~
+    ret_type? ~
+    fn_qualifier ~
+    semi_str
+}
+
+assume_specification_for = {
+    path
+}
+
 abi = {
     extern_str ~ string?
 }
diff --git a/tests/verus-consistency.rs b/tests/verus-consistency.rs
index 8cb8852..a1b1d33 100644
--- a/tests/verus-consistency.rs
+++ b/tests/verus-consistency.rs
@@ -2644,3 +2644,26 @@ fn foo() {
     } // verus!
     "###)
 }
+
+#[test]
+fn verus_assume_specification() {
+    let file = r#"
+verus! {
+
+pub assume_specification<T, const N: usize> [ <[T; N]>::as_slice ](ar: &[T; N]) -> (out: &[T])
+    ensures
+        ar@ == out@;
+
+} // verus!
+"#;
+    assert_snapshot!(parse_and_format(file).unwrap(), @r"
+    verus! {
+
+    pub assume_specification<T, const N: usize>[ <[T; N]>::as_slice ](ar: &[T; N]) -> (out: &[T])
+        ensures
+            ar@ == out@,
+    ;
+
+    } // verus!
+    ")
+}