Allow specifying the scheduling strategy in #[kani_proof] for async f…

…unctions (#1661)

Instead of having to manually wrap the code in `kani::spawnable_block_on` as in #1659, this PR allows `#[kani::proof]` to be applied to harnesses that use `spawn` as well. For this to happen, the user has to specify a scheduling strategy.

Co-authored-by: Celina G. Val <celinval@amazon.com>

library/kani/src/futures.rs

@@ -182,10 +182,12 @@ impl Future for JoinHandle {
 #[crate::unstable(feature = "async-lib", issue = 2559, reason = "experimental async support")]
 pub fn spawn<F: Future<Output = ()> + Sync + 'static>(fut: F) -> JoinHandle {
  unsafe {
- GLOBAL_EXECUTOR
- .as_mut()
- .expect("`spawn` should only be called within `block_on_with_spawn`")
- .spawn(fut)
+ if let Some(executor) = GLOBAL_EXECUTOR.as_mut() {
+ executor.spawn(fut)
+ } else {
+ // An explicit panic instead of `.expect(...)` has better location information in Kani's output
+ panic!("`spawn` should only be called within `block_on_with_spawn`")
+ }
  }
 }
 

library/kani_macros/src/lib.rs

@@ -22,7 +22,12 @@ use regular as attr_impl;
 
 /// Marks a Kani proof harness
 ///
-/// For async harnesses, this will call [`kani::block_on`] (see its documentation for more information).
+/// For async harnesses, this will call [`block_on`] to drive the future to completion (see its documentation for more information).
+///
+/// If you want to spawn tasks in an async harness, you have to pass a schedule to the `#[kani::proof]` attribute,
+/// e.g. `#[kani::proof(schedule = kani::RoundRobin::default())]`.
+/// This will wrap the async function in a call to [`block_on_with_spawn`] (see its documentation for more information).
+#[proc_macro_error]
 #[proc_macro_attribute]
 pub fn proof(attr: TokenStream, item: TokenStream) -> TokenStream {
  attr_impl::proof(attr, item)
@@ -93,10 +98,13 @@ pub fn derive_arbitrary(item: TokenStream) -> TokenStream {
 /// This code should only be activated when pre-building Kani's sysroot.
 #[cfg(kani_sysroot)]
 mod sysroot {
+ use proc_macro_error::{abort, abort_call_site};
+
  use super::*;
 
  use {
  quote::{format_ident, quote},
+ syn::parse::{Parse, ParseStream},
  syn::{parse_macro_input, ItemFn},
  };
 
@@ -126,7 +134,31 @@ mod sysroot {
  };
  }
 
+ struct ProofOptions {
+ schedule: Option<syn::Expr>,
+ }
+
+ impl Parse for ProofOptions {
+ fn parse(input: ParseStream) -> syn::Result<Self> {
+ if input.is_empty() {
+ Ok(ProofOptions { schedule: None })
+ } else {
+ let ident = input.parse::<syn::Ident>()?;
+ if ident != "schedule" {
+ abort_call_site!("`{}` is not a valid option for `#[kani::proof]`.", ident;
+ help = "did you mean `schedule`?";
+ note = "for now, `schedule` is the only option for `#[kani::proof]`.";
+ );
+ }
+ let _ = input.parse::<syn::Token![=]>()?;
+ let schedule = Some(input.parse::<syn::Expr>()?);
+ Ok(ProofOptions { schedule })
+ }
+ }
+ }
+
  pub fn proof(attr: TokenStream, item: TokenStream) -> TokenStream {
+ let proof_options = parse_macro_input!(attr as ProofOptions);
  let fn_item = parse_macro_input!(item as ItemFn);
  let attrs = fn_item.attrs;
  let vis = fn_item.vis;
@@ -138,9 +170,13 @@ mod sysroot {
  #[kanitool::proof]
  );
 
- assert!(attr.is_empty(), "#[kani::proof] does not take any arguments currently");
-
  if sig.asyncness.is_none() {
+ if proof_options.schedule.is_some() {
+ abort_call_site!(
+ "`#[kani::proof(schedule = ...)]` can only be used with `async` functions.";
+ help = "did you mean to make this function `async`?";
+ );
+ }
  // Adds `#[kanitool::proof]` and other attributes
  quote!(
  #kani_attributes
@@ -152,32 +188,44 @@ mod sysroot {
  // For async functions, it translates to a synchronous function that calls `kani::block_on`.
  // Specifically, it translates
  // ```ignore
- // #[kani::async_proof]
+ // #[kani::proof]
  // #[attribute]
  // pub async fn harness() { ... }
  // ```
  // to
  // ```ignore
- // #[kani::proof]
+ // #[kanitool::proof]
  // #[attribute]
  // pub fn harness() {
  // async fn harness() { ... }
  // kani::block_on(harness())
+ // // OR
+ // kani::spawnable_block_on(harness(), schedule)
+ // // where `schedule` was provided as an argument to `#[kani::proof]`.
  // }
  // ```
- assert!(
- sig.inputs.is_empty(),
- "#[kani::proof] cannot be applied to async functions that take inputs for now"
- );
+ if !sig.inputs.is_empty() {
+ abort!(
+ sig.inputs,
+ "`#[kani::proof]` cannot be applied to async functions that take arguments for now";
+ help = "try removing the arguments";
+ );
+ }
  let mut modified_sig = sig.clone();
  modified_sig.asyncness = None;
  let fn_name = &sig.ident;
+ let schedule = proof_options.schedule;
+ let block_on_call = if let Some(schedule) = schedule {
+ quote!(kani::block_on_with_spawn(#fn_name(), #schedule))
+ } else {
+ quote!(kani::block_on(#fn_name()))
+ };
  quote!(
  #kani_attributes
  #(#attrs)*
  #vis #modified_sig {
  #sig #body
- kani::block_on(#fn_name())
+ #block_on_call
  }
  )
  .into()

tests/expected/async_proof/expected

@@ -1,5 +1,3 @@
-error: custom attribute panicked
-#[kani::proof] does not take any arguments currently
+`foo` is not a valid option for `#[kani::proof]`.
 
-error: custom attribute panicked
-#[kani::proof] cannot be applied to async functions that take inputs for now
+`#[kani::proof]` cannot be applied to async functions that take arguments for now

tests/expected/async_proof/main.rs

@@ -8,7 +8,7 @@
 fn main() {}
 
 #[kani::proof(foo)]
-async fn test_async_proof_with_arguments() {}
+async fn test_async_proof_with_options() {}
 
 #[kani::proof]
-async fn test_async_proof_on_function_with_inputs(_: ()) {}
+async fn test_async_proof_on_function_with_arguments(_: ()) {}

tests/kani/AsyncAwait/spawn.rs

@@ -11,9 +11,24 @@ use std::sync::{
  Arc,
 };
 
+#[kani::proof(schedule = kani::RoundRobin::default())]
+#[kani::unwind(4)]
+async fn round_robin_schedule() {
+ let x = Arc::new(AtomicI64::new(0)); // Surprisingly, Arc verified faster than Rc
+ let x2 = x.clone();
+ let x3 = x2.clone();
+ let handle = kani::spawn(async move {
+ x3.fetch_add(1, Ordering::Relaxed);
+ });
+ kani::yield_now().await;
+ x2.fetch_add(1, Ordering::Relaxed);
+ handle.await;
+ assert_eq!(x.load(Ordering::Relaxed), 2);
+}
+
 #[kani::proof]
 #[kani::unwind(4)]
-fn round_robin_schedule() {
+fn round_robin_schedule_manual() {
  let x = Arc::new(AtomicI64::new(0)); // Surprisingly, Arc verified faster than Rc
  let x2 = x.clone();
  kani::block_on_with_spawn(

tests/ui/arguments-proof/expected

@@ -1 +1,3 @@
- = help: message: #[kani::proof] does not take any arguments
+`some` is not a valid option for `#[kani::proof]`.
+
+the trait bound `NotASchedule: kani::futures::SchedulingStrategy` is not satisfied

tests/ui/arguments-proof/main.rs

@@ -1,7 +1,7 @@
 // Copyright Kani Contributors
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
-// kani-flags: --no-unwinding-checks
+// kani-flags: --no-unwinding-checks -Z async-lib
 
 // This test is to check Kani's error handling for harnesses that have proof attributes
 // with arguments when the expected declaration takes no arguments.
@@ -19,3 +19,12 @@ fn harness() {
  assert!(counter < 10);
  }
 }
+
+// Test what happens if the schedule option is incorrect:
+
+struct NotASchedule;
+
+#[kani::proof(schedule = NotASchedule)]
+async fn test() {
+ assert!(true);
+}

tests/ui/arguments-proof/missing-unstable-flag/expected

@@ -0,0 +1 @@
+error: Use of unstable feature `async-lib`: experimental async support

tests/ui/arguments-proof/missing-unstable-flag/main.rs

@@ -0,0 +1,12 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// compile-flags: --edition 2018
+// kani-flags: --no-unwinding-checks
+
+// This test is to check that the `schedule` argument requires an unstable flag.
+
+#[kani::proof(schedule = kani::RoundRobin::default())]
+async fn test() {
+ assert!(true);
+}