Unwind Attribute without Python Changes (#846)

* Unwind Attribute without Python Changes

Author: jaisnan

library/kani_macros/src/lib.rs

@@ -36,9 +36,10 @@ pub fn proof(_attr: TokenStream, item: TokenStream) -> TokenStream {
 
 #[cfg(kani)]
 #[proc_macro_attribute]
-pub fn proof(_attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn proof(attr: TokenStream, item: TokenStream) -> TokenStream {
     let mut result = TokenStream::new();
 
+    assert!(attr.to_string().len() == 0, "#[kani::proof] does not take any arguments");
     result.extend("#[kanitool::proof]".parse::<TokenStream>().unwrap());
     // no_mangle is a temporary hack to make the function "public" so it gets codegen'd
     result.extend("#[no_mangle]".parse::<TokenStream>().unwrap());
@@ -49,3 +50,26 @@ pub fn proof(_attr: TokenStream, item: TokenStream) -> TokenStream {
     //     $item
     // )
 }
+
+#[cfg(not(kani))]
+#[proc_macro_attribute]
+pub fn unwind(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    // When the config is not kani, we should leave the function alone
+    item
+}
+
+/// Set Loop unwind limit for proof harnesses
+/// The attribute '#[kani::unwind(arg)]' can only be called alongside '#[kani::proof]'.
+/// arg - Takes in a integer value (u32) that represents the unwind value for the harness.
+#[cfg(kani)]
+#[proc_macro_attribute]
+pub fn unwind(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let mut result = TokenStream::new();
+
+    // Translate #[kani::unwind(arg)] to #[kanitool::unwind(arg)]
+    let insert_string = "#[kanitool::unwind(".to_owned() + &attr.clone().to_string() + ")]";
+    result.extend(insert_string.parse::<TokenStream>().unwrap());
+
+    result.extend(item);
+    result
+}

src/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs

@@ -8,9 +8,12 @@ use crate::codegen_cprover_gotoc::GotocCtx;
 use cbmc::goto_program::{Expr, Stmt, Symbol};
 use cbmc::InternString;
 use rustc_ast::ast;
+use rustc_ast::{Attribute, LitKind};
+use rustc_middle::mir::coverage::Op;
 use rustc_middle::mir::{HasLocalDecls, Local};
 use rustc_middle::ty::{self, Instance};
 use std::collections::BTreeMap;
+use std::convert::TryInto;
 use std::iter::FromIterator;
 use tracing::{debug, warn};
 
@@ -238,21 +241,50 @@ impl<'tcx> GotocCtx<'tcx> {
     /// This updates the goto context with any information that should be accumulated from a function's
     /// attributes.
     ///
-    /// Currently, this is only proof harness annotations.
-    /// i.e. `#[kani::proof]` (which kani_macros translates to `#[kanitool::proof]` for us to handle here)
+    /// Handle all attributes i.e. `#[kani::x]` (which kani_macros translates to `#[kanitool::x]` for us to handle here)
     fn handle_kanitool_attributes(&mut self) {
-        let instance = self.current_fn().instance();
+        let all_attributes = self.tcx.get_attrs(self.current_fn().instance().def_id());
+        let (proof_attributes, other_attributes) = partition_kanitool_attributes(all_attributes);
 
-        for attr in self.tcx.get_attrs(instance.def_id()) {
-            match kanitool_attr_name(attr).as_deref() {
-                Some("proof") => self.handle_kanitool_proof(),
-                _ => {}
+        if proof_attributes.is_empty() && !other_attributes.is_empty() {
+            self.tcx.sess.span_err(
+                other_attributes[0].1.span,
+                format!(
+                    "The {} attribute also requires the '#[kani::proof]' attribute",
+                    other_attributes[0].0
+                )
+                .as_str(),
+            );
+            return;
+        }
+        if proof_attributes.len() > 1 {
+            // No return because this only requires a warning
+            self.tcx.sess.span_warn(proof_attributes[0].span, "Only one '#[kani::proof]' allowed");
+        }
+        if !proof_attributes.is_empty() {
+            self.create_proof_harness(other_attributes);
+        }
+    }
+
+    /// Create the proof harness struct using the handler methods for various attributes
+    fn create_proof_harness(&mut self, other_attributes: Vec<(String, &Attribute)>) {
+        let mut harness = self.default_kanitool_proof();
+        for attr in other_attributes.iter() {
+            match attr.0.as_str() {
+                "unwind" => self.handle_kanitool_unwind(attr.1, &mut harness),
+                _ => {
+                    self.tcx.sess.span_err(
+                        attr.1.span,
+                        format!("Unsupported Annotation -> {}", attr.0.as_str()).as_str(),
+                    );
+                }
             }
         }
+        self.proof_harnesses.push(harness);
     }
 
-    /// Update `self` (the goto context) to add the current function as a listed proof harness
-    fn handle_kanitool_proof(&mut self) {
+    /// Create the default proof harness for the current function
+    fn default_kanitool_proof(&mut self) -> HarnessMetadata {
         let current_fn = self.current_fn();
         let pretty_name = current_fn.readable_name().to_owned();
         let mangled_name = current_fn.name();
@@ -263,20 +295,91 @@ impl<'tcx> GotocCtx<'tcx> {
             mangled_name,
             original_file: loc.filename().unwrap(),
             original_line: loc.line().unwrap().to_string(),
+            unwind_value: None,
         };
 
-        self.proof_harnesses.push(harness);
+        harness
+    }
+
+    /// Updates the proof harness with new unwind value
+    fn handle_kanitool_unwind(&mut self, attr: &Attribute, harness: &mut HarnessMetadata) {
+        // If some unwind value already exists, then the current unwind being handled is a duplicate
+        if !harness.unwind_value.is_none() {
+            self.tcx.sess.span_err(attr.span, "Only one '#[kani::unwind]' allowed");
+            return;
+        }
+        // Get Attribute value and if it's not none, assign it to the metadata
+        match extract_integer_argument(attr) {
+            None => {
+                // There are no integers or too many arguments given to the attribute
+                self.tcx
+                    .sess
+                    .span_err(attr.span, "Exactly one Unwind Argument as Integer accepted");
+                return;
+            }
+            Some(unwind_integer_value) => {
+                let val: Result<u32, _> = unwind_integer_value.try_into();
+                if val.is_err() {
+                    self.tcx
+                        .sess
+                        .span_err(attr.span, "Value above maximum permitted value - u32::MAX");
+                    return;
+                }
+                harness.unwind_value = Some(val.unwrap());
+            }
+        }
     }
 }
 
 /// If the attribute is named `kanitool::name`, this extracts `name`
 fn kanitool_attr_name(attr: &ast::Attribute) -> Option<String> {
     match &attr.kind {
         ast::AttrKind::Normal(ast::AttrItem { path: ast::Path { segments, .. }, .. }, _)
-            if segments.len() == 2 && segments[0].ident.as_str() == "kanitool" =>
+            if (!segments.is_empty()) && segments[0].ident.as_str() == "kanitool" =>
         {
             Some(segments[1].ident.as_str().to_string())
         }
         _ => None,
     }
 }
+
+/// Partition all the attributes into two buckets, proof_attributes and other_attributes
+fn partition_kanitool_attributes(
+    all_attributes: &[Attribute],
+) -> (Vec<&Attribute>, Vec<(String, &Attribute)>) {
+    let mut proof_attributes = vec![];
+    let mut other_attributes = vec![];
+
+    for attr in all_attributes {
+        // Get the string the appears after "kanitool::" in each attribute string.
+        // Ex - "proof" | "unwind" etc.
+        if let Some(attribute_string) = kanitool_attr_name(attr).as_deref() {
+            if attribute_string == "proof" {
+                proof_attributes.push(attr);
+            } else {
+                other_attributes.push((attribute_string.to_string(), attr));
+            }
+        }
+    }
+
+    (proof_attributes, other_attributes)
+}
+
+/// Extracts the integer value argument from the attribute provided
+/// For example, `unwind(8)` return `Some(8)`
+fn extract_integer_argument(attr: &Attribute) -> Option<u128> {
+    // Vector of meta items , that contain the arguments given the attribute
+    let attr_args = attr.meta_item_list()?;
+    // Only extracts one integer value as argument
+    if attr_args.len() == 1 {
+        let x = attr_args[0].literal()?;
+        match x.kind {
+            LitKind::Int(y, ..) => Some(y),
+            _ => None,
+        }
+    }
+    // Return none if there are no attributes or if there's too many attributes
+    else {
+        None
+    }
+}

src/kani-compiler/src/codegen_cprover_gotoc/context/metadata.rs

@@ -17,6 +17,8 @@ pub struct HarnessMetadata {
     pub original_file: String,
     /// The line in that file where the proof harness begins
     pub original_line: String,
+    /// Optional data to store unwind value
+    pub unwind_value: Option<u32>,
 }
 
 /// The structure of `.kani-metadata.json` files, which are emitted for each crate

tests/cargo-kani/simple-proof-annotation/Cargo.toml

@@ -9,4 +9,4 @@ edition = "2021"
 [dependencies]
 
 [kani.flags]
-output-format = "old"
+output-format = "old"

tests/cargo-kani/simple-unwind-annotation/Cargo.toml

@@ -0,0 +1,13 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+
+[package]
+name = "simple-unwind-annotation"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+
+[kani.flags]
+output-format = "old"
+function = "harness"

tests/cargo-kani/simple-unwind-annotation/main.expected

@@ -0,0 +1 @@
+[harness.assertion.2] line 21 assertion failed: counter < 10: FAILURE

tests/cargo-kani/simple-unwind-annotation/src/main.rs

@@ -0,0 +1,32 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: --function harness
+
+// The expected file presently looks for "1 == 2" above.
+// But eventually this test may start to fail as we might stop regarding 'main'
+// as a valid proof harness, since it isn't annotated as such.
+// This test should be updated if we go that route.
+
+fn main() {
+    assert!(1 == 2);
+}
+
+#[kani::proof]
+#[kani::unwind(9)]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}
+
+#[kani::proof]
+fn harness_2() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}

tests/kani/Unwind-Attribute/fixme_lib.rs

@@ -0,0 +1,21 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: --no-unwinding-checks
+
+// Added to fix me because there are no tests for the annotations themselves.
+// TODO : The unwind value needs to be parsed and the unwind needs to be applied to each harness, we do not test this behavior yet.
+
+fn main() {
+    assert!(1 == 2);
+}
+
+#[kani::proof]
+#[kani::unwind(10)]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}

tests/ui/arguments-proof/expected

@@ -0,0 +1,6 @@
+  --> main.rs:13:1
+   |
+13 | #[kani::proof(some, argument2)]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = help: message: #[kani::proof] does not take any arguments

tests/ui/arguments-proof/main.rs

@@ -0,0 +1,20 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: --no-unwinding-checks
+
+// This test is to check Kani's error handling for harnesses that have proof attributes
+// with arguments when the expected declaration takes no arguments.
+
+fn main() {
+    assert!(1 == 2);
+}
+
+#[kani::proof(some, argument2)]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}

tests/ui/multiple-proof-attributes/expected

@@ -0,0 +1,9 @@
+warning: Only one '#[kani::proof]' allowed
+  --> main.rs:14:1
+   |
+14 | #[kani::proof]
+   | ^^^^^^^^^^^^^^
+   |
+   = note: this warning originates in the attribute macro `kani::proof` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+warning: 1 warning emitted

tests/ui/multiple-proof-attributes/main.rs

@@ -0,0 +1,21 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: --no-unwinding-checks --verbose
+
+// This test is to check Kani's error handling for harnesses that have multiple proof annotations
+// declared.
+
+fn main() {
+    assert!(1 == 2);
+}
+
+#[kani::proof]
+#[kani::proof]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}

tests/ui/unsupported-annotation/expected

@@ -0,0 +1,7 @@
+error[E0433]: failed to resolve: could not find `test_annotation` in `kani`
+  --> main.rs:14:9
+   |
+14 | #[kani::test_annotation]
+   |         ^^^^^^^^^^^^^^^ could not find `test_annotation` in `kani`
+
+error: aborting due to previous error

tests/ui/unsupported-annotation/main.rs

@@ -0,0 +1,21 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: --no-unwinding-checks --verbose
+
+// This test is to check Kani's error handling for harnesses that have multiple proof annotations
+// declared.
+
+fn main() {
+    assert!(1 == 2);
+}
+
+#[kani::proof]
+#[kani::test_annotation]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}

tests/ui/unwind-multiple-arguments/expected

@@ -0,0 +1,5 @@
+error: Exactly one Unwind Argument as Integer accepted
+  --> main.rs:14:1
+   |
+14 | #[kani::unwind(10,5)]
+   | ^^^^^^^^^^^^^^^^^^^^^