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);
+}