Merge branch 'arg-forwarding' into 'master'

Pass simple args to the closure and preserve their types

See merge request karroffel/contracts!5

Author: karroffel

Cargo.toml

@@ -35,7 +35,7 @@ override_log = []
 mirai_assertions = []
 
 [dependencies]
-syn = { version = "1.0", features = ["extra-traits", "full", "visit-mut"] }
+syn = { version = "1.0", features = ["extra-traits", "full", "visit", "visit-mut"] }
 quote = "1.0"
 proc-macro2 = "1.0"
 

src/implementation/codegen.rs

@@ -316,29 +316,14 @@ pub(crate) fn generate(
     };
 
     //
-    // wrap the function body in a closure
+    // wrap the function body in a closure if we have any postconditions
     //
 
-    let block = func.function.block.clone();
-
-    let ret_ty = match &func.function.sig.output {
-        ReturnType::Type(_, ty) => {
-            let span = ty.span();
-            match ty.as_ref() {
-                syn::Type::ImplTrait(_) => quote::quote! {},
-                ty => quote::quote_spanned! { span=>
-                    -> #ty
-                },
-            }
-        }
-        ReturnType::Default => quote::quote! {},
-    };
-
-    let body = quote::quote! {
-        #[allow(unused_mut)]
-        let mut run = || #ret_ty #block;
-
-        let ret = run();
+    let body = if post.is_empty() {
+        let block = &func.function.block;
+        quote::quote! { let ret = #block; }
+    } else {
+        create_closure_body_and_adjust_signature(&func.function)
     };
 
     //
@@ -371,3 +356,157 @@ pub(crate) fn generate(
 
     func.function.into_token_stream()
 }
+
+struct SelfReplacer<'a> {
+    replace_with: &'a syn::Ident,
+}
+
+impl syn::visit_mut::VisitMut for SelfReplacer<'_> {
+    fn visit_ident_mut(&mut self, i: &mut Ident) {
+        if i == "self" {
+            *i = self.replace_with.clone();
+        }
+    }
+}
+
+fn ty_contains_impl_trait(ty: &syn::Type) -> bool {
+    use syn::visit::Visit;
+
+    struct TyContainsImplTrait {
+        seen_impl_trait: bool,
+    }
+
+    impl syn::visit::Visit<'_> for TyContainsImplTrait {
+        fn visit_type_impl_trait(&mut self, _: &syn::TypeImplTrait) {
+            self.seen_impl_trait = true;
+        }
+    }
+
+    let mut vis = TyContainsImplTrait {
+        seen_impl_trait: false,
+    };
+    vis.visit_type(ty);
+    vis.seen_impl_trait
+}
+
+fn create_closure_body_and_adjust_signature(func: &syn::ItemFn) -> TokenStream {
+    let is_method = func.sig.receiver().is_some();
+
+    // If the function has a receiver (e.g. `self`, `&mut self`, or `self: Box<Self>`) rename it
+    // to `this__` within the closure
+
+    let mut block = func.block.clone();
+    let mut closure_args = vec![];
+    let mut arg_names = vec![];
+
+    if is_method {
+        let this_ident = syn::Ident::new("this__", Span::call_site());
+
+        let mut receiver = func.sig.inputs[0].clone();
+        match receiver {
+            // self, &self, &mut self
+            syn::FnArg::Receiver(rcv) => {
+                // The `Self` type.
+                let self_ty = Box::new(syn::Type::Path(syn::TypePath {
+                    qself: None,
+                    path: syn::Path::from(syn::Ident::new("Self", rcv.span())),
+                }));
+
+                let ty = if let Some((and_token, ref lifetime)) = rcv.reference
+                {
+                    Box::new(syn::Type::Reference(syn::TypeReference {
+                        and_token,
+                        lifetime: lifetime.clone(),
+                        mutability: rcv.mutability,
+                        elem: self_ty,
+                    }))
+                } else {
+                    self_ty
+                };
+
+                let pat_mut = if rcv.reference.is_none() {
+                    rcv.mutability
+                } else {
+                    None
+                };
+
+                // this__: [& [mut]] Self
+                let new_rcv = syn::PatType {
+                    attrs: rcv.attrs.clone(),
+                    pat: Box::new(syn::Pat::Ident(syn::PatIdent {
+                        attrs: vec![],
+                        by_ref: None,
+                        mutability: pat_mut,
+                        ident: this_ident.clone(),
+                        subpat: None,
+                    })),
+                    colon_token: syn::Token![:](rcv.span()),
+                    ty,
+                };
+
+                receiver = syn::FnArg::Typed(new_rcv);
+            }
+
+            // self: Box<Self>
+            syn::FnArg::Typed(ref mut pat) => {
+                if let syn::Pat::Ident(ref mut ident) = *pat.pat {
+                    if ident.ident == "self" {
+                        ident.ident = this_ident.clone();
+                    }
+                }
+            }
+        }
+
+        closure_args.push(receiver);
+        arg_names.push(syn::Ident::new("self", Span::call_site()));
+
+        // Replace any references to `self` in the function body with references to `this__`.
+        syn::visit_mut::visit_block_mut(
+            &mut SelfReplacer {
+                replace_with: &this_ident,
+            },
+            &mut block,
+        );
+    }
+
+    // Replace any pattern matching in the function signature with placeholder identifiers.
+    // Pattern matching gets done in the closure instead.
+    let args = func.sig.inputs.iter().skip(if is_method { 1 } else { 0 });
+    for arg in args {
+        match arg {
+            syn::FnArg::Receiver(_) => unreachable!("Multiple receivers?"),
+
+            syn::FnArg::Typed(syn::PatType { pat, ty, .. }) => {
+                if !ty_contains_impl_trait(ty) {
+                    if let syn::Pat::Ident(ident) = &**pat {
+                        arg_names.push(ident.ident.clone());
+                        closure_args.push(arg.clone());
+                    }
+                }
+            }
+        }
+    }
+
+    let ret_ty = match &func.sig.output {
+        ReturnType::Type(_, ty) => {
+            let span = ty.span();
+            match ty.as_ref() {
+                syn::Type::ImplTrait(_) => quote::quote! {},
+                ty => quote::quote_spanned! { span=>
+                    -> #ty
+                },
+            }
+        }
+        ReturnType::Default => quote::quote! {},
+    };
+
+    let closure_args = closure_args.iter();
+    let arg_names = arg_names.iter();
+
+    quote::quote! {
+        #[allow(unused_mut)]
+        let mut run = |#(#closure_args),*| #ret_ty #block;
+
+        let ret = run(#(#arg_names),*);
+    }
+}

tests/functions.rs

@@ -89,13 +89,52 @@ fn test_early_return() {
     abs(-4);
 }
 
+#[test]
+fn test_mut_ref_and_lifetimes() {
+    #[requires(i < s.len())]
+    #[ensures(*ret == 0)]
+    fn insert_zero<'a>(s: &'a mut [u8], i: usize) -> &'a mut u8 {
+        s[i] = 0;
+        &mut s[i]
+    }
+
+    insert_zero(&mut [4, 2], 1);
+}
+
+#[test]
+fn test_pattern_match() {
+    #[ensures(ret > a && ret > b)]
+    fn add((a, b): (u8, u8)) -> u8 {
+        a.saturating_add(b)
+    }
+
+    assert_eq!(add((4, 2)), 6);
+}
+
 #[test]
 fn test_impl_trait_return() {
     // make sure that compiling functions that return existentially
     // qualified types works properly.
 
     #[requires(x >= 10)]
-    fn impl_test(x: isize) -> impl Clone + std::fmt::Debug {
+    #[ensures(ret.clone() == ret)]
+    fn impl_test(x: isize) -> impl Clone + PartialEq + std::fmt::Debug {
+        "it worked"
+    }
+
+    let x = impl_test(200);
+    let y = x.clone();
+    assert_eq!(
+        format!("{:?} and {:?}", x, y),
+        r#""it worked" and "it worked""#
+    );
+}
+
+#[test]
+fn test_impl_trait_arg() {
+    #[requires(x.clone() == x)]
+    #[ensures(ret.clone() == ret)]
+    fn impl_test(x: impl Clone + PartialEq + std::fmt::Debug) -> &'static str {
         "it worked"
     }
 

tests/methods.rs

@@ -26,6 +26,15 @@ fn methods() {
             self.count += 2;
         }
 
+        // Manually express the invariant in terms of `ret` since `self.count` is mutably borrowed.
+        #[requires(is_even(self.count))]
+        #[ensures(is_even(*ret))]
+        #[ensures(*ret == old(self.count) + 2)]
+        fn next_even_and_get<'a>(&'a mut self) -> &'a mut usize {
+            self.count += 2;
+            &mut self.count
+        }
+
         #[invariant(is_even(self.count))]
         #[requires(self.count >= 2)]
         #[ensures(self.count == old(self.count) - 2)]
@@ -41,6 +50,8 @@ fn methods() {
 
     adder.prev_even();
     adder.prev_even();
+
+    assert_eq!(*adder.next_even_and_get(), 2);
 }
 
 #[test]