Merge pull request #847 from hacspec/fix-844

Fix ensures for unit-returning functions with &mut inputs

Author: W95Psp

hax-lib-macros/src/utils.rs

@@ -201,6 +201,19 @@ fn create_future_ident(name: &str) -> syn::Ident {
     proc_macro2::Ident::new(&format!("{name}_future"), proc_macro2::Span::call_site())
 }
 
+/// The engine translates functions of arity zero to functions that
+/// takes exactly one unit argument. The zero-arity functions we
+/// generate are translated correctly as well. But in the case of a
+/// `ensures` clause, that's an issue: we produce a function of arity
+/// one, whose first argument is the result of the function. Instead,
+/// we need a function of arity two.
+/// `fix_signature_arity` adds a `unit` if needed.
+fn add_unit_to_sig_if_needed(signature: &mut Signature) {
+    if signature.inputs.is_empty() {
+        signature.inputs.push(parse_quote! {_: ()})
+    }
+}
+
 /// Common logic when generating a function decoration
 pub fn make_fn_decoration(
     mut phi: Expr,
@@ -223,12 +236,13 @@ pub fn make_fn_decoration(
     };
     let decoration = {
         let decoration_sig = {
-            let mut sig = signature;
+            let mut sig = signature.clone();
             sig.ident = format_ident!("{}", kind.to_string());
             if let FnDecorationKind::Ensures { ret_binder } = &kind {
-                let output = match sig.output {
-                    syn::ReturnType::Default => quote! {()},
-                    syn::ReturnType::Type(_, t) => quote! {#t},
+                add_unit_to_sig_if_needed(&mut sig);
+                let output_typ = match sig.output {
+                    syn::ReturnType::Default => parse_quote! {()},
+                    syn::ReturnType::Type(_, t) => t,
                 };
                 let mut_ref_inputs = mut_ref_inputs
                     .iter()
@@ -240,16 +254,26 @@ pub fn make_fn_decoration(
                 let mut rewrite_future =
                     RewriteFuture(mut_ref_inputs.clone().map(|x| x.0).collect());
                 rewrite_future.visit_expr_mut(&mut phi);
-                let (pats, tys): (Vec<_>, Vec<_>) = mut_ref_inputs
+                let (mut pats, mut tys): (Vec<_>, Vec<_>) = mut_ref_inputs
                     .map(|(name, ty)| {
                         (
                             create_future_ident(&name).to_token_stream(),
                             ty.to_token_stream(),
                         )
                     })
-                    .chain([(ret_binder.to_token_stream(), output)].into_iter())
                     .unzip();
 
+                let is_output_typ_unit = if let syn::Type::Tuple(tuple) = &*output_typ {
+                    tuple.elems.is_empty()
+                } else {
+                    false
+                };
+
+                if !is_output_typ_unit || pats.is_empty() {
+                    pats.push(ret_binder.to_token_stream());
+                    tys.push(quote! {#output_typ});
+                }
+
                 sig.inputs
                     .push(syn::parse_quote! {(#(#pats),*): (#(#tys),*)});
             }

test-harness/src/snapshots/toolchain__attributes into-fstar.snap

@@ -26,6 +26,28 @@ exit = 0
 diagnostics = []
 
 [stdout.files]
+"Attributes.Ensures_on_arity_zero_fns.fst" = '''
+module Attributes.Ensures_on_arity_zero_fns
+#set-options "--fuel 0 --ifuel 1 --z3rlimit 15"
+open Core
+open FStar.Mul
+
+let basically_a_constant (_: Prims.unit)
+    : Prims.Pure u8
+      (requires true)
+      (ensures
+        fun x ->
+          let x:u8 = x in
+          x >. 100uy) = 127uy
+
+let doing_nothing (_: Prims.unit)
+    : Prims.Pure Prims.unit
+      (requires true)
+      (ensures
+        fun v__x ->
+          let v__x:Prims.unit = v__x in
+          true) = ()
+'''
 "Attributes.Inlined_code_ensures_requires.fst" = '''
 module Attributes.Inlined_code_ensures_requires
 #set-options "--fuel 0 --ifuel 1 --z3rlimit 15"
@@ -36,8 +58,8 @@ let increment_array (v: t_Array u8 (sz 4))
     : Prims.Pure (t_Array u8 (sz 4))
       (requires forall i. FStar.Seq.index v i <. 254uy)
       (ensures
-        fun temp_0_ ->
-          let vv_future, ():(t_Array u8 (sz 4) & Prims.unit) = temp_0_ in
+        fun vv_future ->
+          let vv_future:t_Array u8 (sz 4) = vv_future in
           let future_v:t_Array u8 (sz 4) = vv_future in
           forall i. FStar.Seq.index future_v i >. 0uy) =
   let v:t_Array u8 (sz 4) =
@@ -368,6 +390,14 @@ let add3_lemma (x: u32)
 
 let inlined_code__V: u8 = 12uy
 
+let issue_844_ (v__x: u8)
+    : Prims.Pure u8
+      Prims.l_True
+      (ensures
+        fun v__x_future ->
+          let v__x_future:u8 = v__x_future in
+          true) = v__x
+
 let u32_max: u32 = 90000ul
 
 unfold let some_function _ = "hello from F*"

tests/attributes/src/lib.rs

@@ -20,6 +20,21 @@ fn swap_and_mut_req_ens(x: &mut u32, y: &mut u32) -> u32 {
     *x + *y
 }
 
+#[hax_lib::ensures(|_| true)]
+fn issue_844(_x: &mut u8) {}
+
+// From issue #845
+mod ensures_on_arity_zero_fns {
+    #[hax_lib::requires(true)]
+    #[hax_lib::ensures(|_x| true)]
+    fn doing_nothing() {}
+    #[hax_lib::requires(true)]
+    #[hax_lib::ensures(|x| x > 100)]
+    fn basically_a_constant() -> u8 {
+        127
+    }
+}
+
 #[hax::lemma]
 fn add3_lemma(x: u32) -> Proof<{ x <= 10 || x >= u32_max / 3 || add3(x, x, x) == x * 3 }> {}