Move macro to lib.rs, fix tests

library/kani_macros/src/derive.rs

@@ -157,7 +157,7 @@ fn inv_conds_inner(ident: &Ident, fields: &Fields) -> Option<TokenStream> {
 /// ```
 /// which allows us to refer to the struct fields without using `self`.
 /// Note that the actual stream is generated in the `field_refs_inner` function.
-fn field_refs(ident: &Ident, data: &Data) -> TokenStream {
+pub fn field_refs(ident: &Ident, data: &Data) -> TokenStream {
  match data {
  Data::Struct(struct_data) => field_refs_inner(ident, &struct_data.fields),
  Data::Enum(_) => unreachable!(),
@@ -344,7 +344,7 @@ pub fn expand_derive_invariant(item: proc_macro::TokenStream) -> proc_macro::Tok
 }
 
 /// Add a bound `T: Invariant` to every type parameter T.
-fn add_trait_bound_invariant(mut generics: Generics) -> Generics {
+pub fn add_trait_bound_invariant(mut generics: Generics) -> Generics {
  generics.params.iter_mut().for_each(|param| {
  if let GenericParam::Type(type_param) = param {
  type_param.bounds.push(parse_quote!(kani::Invariant));
@@ -422,43 +422,3 @@ fn struct_invariant_conjunction(ident: &Ident, fields: &Fields) -> TokenStream {
  }
  }
 }
-
-/// Generates an `Invariant` implementation where the `is_safe` function body is
-/// the `attr` expression passed to the attribute macro.
-/// Only available for structs.
-pub fn attr_custom_invariant(
- attr: TokenStream,
- item: proc_macro::TokenStream,
-) -> proc_macro::TokenStream {
- let derive_item = parse_macro_input!(item as DeriveInput);
- let item_name = &derive_item.ident;
-
- if !matches!(derive_item.data, Data::Struct(..)) {
- abort!(Span::call_site(), "Cannot define invariant for `{}`", item_name;
- note = item_name.span() =>
- "`#[kani::invariant(..)]` is only available for structs"
- )
- }
-
- // Keep a copy of the original item to re-emit it later.
- // Note that this isn't a derive macro
- let original_item = derive_item.clone();
-
- // Add a bound `T: Invariant` to every type parameter T.
- let generics = add_trait_bound_invariant(derive_item.generics);
- // Generate an expression to sum up the heap size of each field.
- let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
-
- let expanded = quote! {
- // Re-emit the original item
- #original_item
-
- // The generated implementation.
- impl #impl_generics kani::Invariant for #item_name #ty_generics #where_clause {
- fn is_safe(&self) -> bool {
- #attr
- }
- }
- };
- proc_macro::TokenStream::from(expanded)
-}

library/kani_macros/src/lib.rs

@@ -12,8 +12,11 @@
 mod derive;
 
 // proc_macro::quote is nightly-only, so we'll cobble things together instead
+use derive::{add_trait_bound_invariant, field_refs};
 use proc_macro::TokenStream;
-use proc_macro_error::proc_macro_error;
+use proc_macro2::{Span, TokenStream as TokenStream2};
+use proc_macro_error::{abort, proc_macro_error};
+use quote::quote;
 
 #[cfg(kani_sysroot)]
 use sysroot as attr_impl;
@@ -95,7 +98,7 @@ pub fn solver(attr: TokenStream, item: TokenStream) -> TokenStream {
 #[proc_macro_attribute]
 #[proc_macro_error]
 pub fn invariant(attr: TokenStream, item: TokenStream) -> TokenStream {
- derive::attr_custom_invariant(attr.into(), item)
+ attr_custom_invariant(attr.into(), item)
 }
 
 /// Mark an API as unstable. This should only be used inside the Kani sysroot.
@@ -398,3 +401,46 @@ mod regular {
  no_op!(proof_for_contract);
  no_op!(stub_verified);
 }
+
+/// Generates an `Invariant` implementation where the `is_safe` function body is
+/// the `attr` expression passed to the attribute macro.
+/// Only available for structs.
+fn attr_custom_invariant(
+ attr: TokenStream2,
+ item: proc_macro::TokenStream,
+) -> proc_macro::TokenStream {
+ let derive_item = syn::parse_macro_input!(item as syn::DeriveInput);
+ let item_name = &derive_item.ident;
+
+ if !matches!(derive_item.data, syn::Data::Struct(..)) {
+ abort!(Span::call_site(), "Cannot define invariant for `{}`", item_name;
+ note = item_name.span() =>
+ "`#[kani::invariant(..)]` is only available for structs"
+ )
+ }
+
+ // Keep a copy of the original item to re-emit it later.
+ // Note that this isn't a derive macro
+ let original_item = derive_item.clone();
+
+ // Add a bound `T: Invariant` to every type parameter T.
+ let generics = add_trait_bound_invariant(derive_item.generics);
+ // Generate an expression to sum up the heap size of each field.
+ let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
+ let field_refs = field_refs(&item_name, &derive_item.data);
+ let expanded = quote! {
+ // Re-emit the original item
+ #original_item
+
+ // The generated implementation.
+ impl #impl_generics kani::Invariant for #item_name #ty_generics #where_clause {
+ fn is_safe(&self) -> bool {
+ let obj = self;
+ #field_refs
+ #attr
+ }
+ }
+ };
+ TokenStream::from(expanded)
+}

tests/expected/attr-invariant/check_invariant/check_invariant.rs

@@ -8,7 +8,7 @@ extern crate kani;
 use kani::Invariant;
 
 #[derive(kani::Arbitrary)]
-#[kani::invariant(self.x.is_safe() && self.y.is_safe())]
+#[kani::invariant(x.is_safe() && y.is_safe())]
 struct Point {
  x: i32,
  y: i32,