fix: place function bodies in labelled block expression instead of closures

robjtede · robjtede · commit dee2b82fc49f · 2025-05-02T23:20:53.000+01:00
fixes #5 #7 #8 taken from https://gitlab.com/karroffel/contracts/-/merge_requests/9/diffs
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## Unreleased
 
+- Improve support for async functions and mutable borrows in contracts.
+
 ## 0.6.4
 
 - Add effective contract type to invariant violation messages.
diff --git a/README.md b/README.md
@@ -1,10 +1,20 @@
-# _Design By Contract_ for Rust
+# `contracts`
 
-[![License][license]][LICENSE] ![Build status][build] ![Lines of Code][loc]
+> _Design By Contract_ for Rust
 
-[license]: https://img.shields.io/badge/license-MPL%202.0-blue.svg
-[build]: https://gitlab.com/karroffel/contracts/badges/master/pipeline.svg
-[loc]: https://tokei.rs/b1/gitlab/karroffel/contracts?category=code
+<!-- prettier-ignore-start -->
+
+[![crates.io](https://img.shields.io/crates/v/contracts?label=latest)](https://crates.io/crates/contracts)
+[![Documentation](https://docs.rs/contracts/badge.svg?version=0.6.4)](https://docs.rs/contracts/0.6.4)
+[![dependency status](https://deps.rs/crate/contracts/0.6.4/status.svg)](https://deps.rs/crate/contracts/0.6.4)
+![MIT or Apache 2.0 licensed](https://img.shields.io/crates/l/contracts.svg)
+<br />
+[![CI](https://github.com/x52dev/contracts-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/x52dev/contracts-rs/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/x52dev/contracts-rs/graph/badge.svg?token=OpYe6I7dj5)](https://codecov.io/gh/x52dev/contracts-rs)
+![Version](https://img.shields.io/crates/msrv/contracts.svg)
+[![Download](https://img.shields.io/crates/d/contracts.svg)](https://crates.io/crates/contracts)
+
+<!-- prettier-ignore-end -->
 
 Annotate functions and methods with "contracts", using _invariants_, _pre-conditions_ and _post-conditions_.
 
diff --git a/src/implementation/codegen.rs b/src/implementation/codegen.rs
@@ -4,7 +4,12 @@
 
 use proc_macro2::{Ident, Span, TokenStream};
 use quote::ToTokens;
-use syn::{spanned::Spanned, visit_mut as visitor, Attribute, Expr, ExprCall, ReturnType};
+use syn::{
+    spanned::Spanned,
+    visit::{visit_return_type, Visit},
+    visit_mut::{self as visitor, visit_block_mut, visit_expr_mut, VisitMut},
+    Attribute, Expr, ExprCall, ReturnType, TypeImplTrait,
+};
 
 use crate::implementation::{Contract, ContractMode, ContractType, FuncWithContracts};
 
@@ -98,7 +103,6 @@ pub(crate) fn extract_old_calls(contracts: &mut [Contract]) -> Vec<OldExpr> {
         }
 
         for assertion in &mut contract.assertions {
-            use visitor::VisitMut;
             extractor.visit_expr_mut(assertion);
         }
     }
@@ -169,11 +173,11 @@ pub(crate) fn generate(
     // creates an assertion appropriate for the current mode
     let make_assertion = |mode: ContractMode,
                           ctype: ContractType,
-                          display: proc_macro2::TokenStream,
+                          display: TokenStream,
                           exec_expr: &Expr,
                           desc: &str| {
         let span = display.span();
-        let mut result = proc_macro2::TokenStream::new();
+        let mut result = TokenStream::new();
 
         let format_args = quote::quote_spanned! { span=>
             concat!(concat!(#desc, ": "), stringify!(#display))
@@ -208,7 +212,7 @@ pub(crate) fn generate(
     // generate assertion code for pre-conditions
     //
 
-    let pre: proc_macro2::TokenStream = func
+    let pre = func
         .contracts
         .iter()
         .filter(|c| c.ty == ContractType::Requires || c.ty == ContractType::Invariant)
@@ -222,7 +226,7 @@ pub(crate) fn generate(
             let desc = if let Some(desc) = c.desc.as_ref() {
                 format!("{} of {} violated: {}", contract_type_name, func_name, desc)
             } else {
-                format!("{} of {} violated", c.ty.message_name(), func_name)
+                format!("{} of {} violated", contract_type_name, func_name)
             };
 
             c.assertions
@@ -240,13 +244,13 @@ pub(crate) fn generate(
                     )
                 })
         })
-        .collect();
+        .collect::<TokenStream>();
 
     //
     // generate assertion code for post-conditions
     //
 
-    let post: proc_macro2::TokenStream = func
+    let post = func
         .contracts
         .iter()
         .filter(|c| c.ty == ContractType::Ensures || c.ty == ContractType::Invariant)
@@ -260,7 +264,7 @@ pub(crate) fn generate(
             let desc = if let Some(desc) = c.desc.as_ref() {
                 format!("{} of {} violated: {}", contract_type_name, func_name, desc)
             } else {
-                format!("{} of {} violated", c.ty.message_name(), func_name)
+                format!("{} of {} violated", contract_type_name, func_name)
             };
 
             c.assertions
@@ -278,14 +282,14 @@ pub(crate) fn generate(
                     )
                 })
         })
-        .collect();
+        .collect::<TokenStream>();
 
     //
     // bind "old()" expressions
     //
 
     let olds = {
-        let mut toks = proc_macro2::TokenStream::new();
+        let mut toks = TokenStream::new();
 
         for old in olds {
             let span = old.expr.span();
@@ -305,14 +309,27 @@ pub(crate) fn generate(
     };
 
     //
-    // wrap the function body in a closure if we have any postconditions
+    // wrap the function body in a block so that we can use its return value
     //
 
-    let body = if post.is_empty() {
-        let block = &func.function.block;
-        quote::quote! { let ret = #block; }
-    } else {
-        create_body_closure(&func.function)
+    let body = 'blk: {
+        let mut block = func.function.block.clone();
+        visit_block_mut(&mut ReturnReplacer, &mut block);
+
+        let mut impl_detector = ImplDetector { found_impl: false };
+        visit_return_type(&mut impl_detector, &func.function.sig.output);
+
+        if !impl_detector.found_impl {
+            if let ReturnType::Type(.., ref return_type) = func.function.sig.output {
+                break 'blk quote::quote! {
+                    let ret: #return_type = 'run: #block;
+                };
+            }
+        }
+
+        quote::quote! {
+            let ret = 'run: #block;
+        }
     };
 
     //
@@ -346,177 +363,25 @@ pub(crate) fn generate(
     func.function.into_token_stream()
 }
 
-struct SelfReplacer<'a> {
-    replace_with: &'a syn::Ident,
-}
+struct ReturnReplacer;
 
-impl syn::visit_mut::VisitMut for SelfReplacer<'_> {
-    fn visit_ident_mut(&mut self, i: &mut Ident) {
-        if i == "self" {
-            *i = self.replace_with.clone();
+impl VisitMut for ReturnReplacer {
+    fn visit_expr_mut(&mut self, node: &mut Expr) {
+        if let Expr::Return(ret_expr) = node {
+            let ret_expr_expr = ret_expr.expr.clone();
+            *node = syn::parse_quote!(break 'run #ret_expr_expr);
         }
-    }
-}
-
-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;
-        }
+        visit_expr_mut(self, node);
     }
-
-    let mut vis = TyContainsImplTrait {
-        seen_impl_trait: false,
-    };
-    vis.visit_type(ty);
-    vis.seen_impl_trait
 }
 
-fn create_body_closure(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 {
-        // `mixed_site` makes the identifier hygienic so it won't collide with a local variable
-        // named `this`.
-        let this_ident = syn::Ident::new("this", Span::mixed_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);
-
-        match &func.sig.inputs[0] {
-            syn::FnArg::Receiver(receiver) => {
-                arg_names.push(syn::Ident::new("self", receiver.self_token.span()));
-            }
-            syn::FnArg::Typed(pat) => {
-                if let syn::Pat::Ident(ident) = &*pat.pat {
-                    arg_names.push(ident.ident.clone());
-                } else {
-                    // Non-trivial receiver pattern => do not capture
-                    closure_args.pop();
-                }
-            }
-        };
-
-        // 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,
-        );
-    }
-
-    // Any function arguments of the form `ident: ty` become closure arguments and get passed
-    // explicitly. More complex ones, e.g. pattern matching like `(a, b): (u32, u32)`, are
-    // captured 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 {
-                        let ident_str = ident.ident.to_string();
-
-                        // Any function argument identifier starting with '_' signals
-                        // that the binding will not be used.
-                        if !ident_str.starts_with('_') || ident_str.starts_with("__") {
-                            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;
+struct ImplDetector {
+    found_impl: bool,
+}
 
-        let ret = run(#(#arg_names),*);
+impl<'a> Visit<'a> for ImplDetector {
+    fn visit_type_impl_trait(&mut self, _node: &'a TypeImplTrait) {
+        self.found_impl = true;
     }
 }
diff --git a/tests/issues.rs b/tests/issues.rs
