Unwind Attribute without Python Changes

library/kani_macros/src/lib.rs

@@ -9,6 +9,7 @@
 //   RUSTFLAGS="-Zcrate-attr=feature(register_tool) -Zcrate-attr=register_tool(kanitool)"
 
 // proc_macro::quote is nightly-only, so we'll cobble things together instead
+extern crate proc_macro;
 use proc_macro::TokenStream;
 
 #[cfg(all(not(kani), not(test)))]
@@ -49,3 +50,29 @@ pub fn proof(_attr: TokenStream, item: TokenStream) -> TokenStream {
     //     $item
     // )
 }
+
+#[cfg(not(kani))]
+#[proc_macro_attribute]
+pub fn unwind(_attr: TokenStream, _item: TokenStream) -> TokenStream {
+    // Not-RMC, Not-Test means this code shouldn't exist, return nothing.
+    TokenStream::new()
+}
+
+#[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_] for easier handling
+    let insert_string = "#[kanitool::unwind_".to_owned() + &attr.clone().to_string() + "]";
+
+    // Add the string that looks like - #[kanitool::unwind_value_]
+    result.extend(insert_string.parse::<TokenStream>().unwrap());
+    // No mangle seems to be necessary as removing it prevents all the attributes in a lib from being read
+    result.extend("#[no_mangle]".parse::<TokenStream>().unwrap());
+
+    result.extend(item);
+    result
+
+    // / _attr
+}

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

@@ -3,7 +3,7 @@
 
 //! This file contains functions related to codegenning MIR functions into gotoc
 
-use crate::context::metadata::HarnessMetadata;
+use crate::context::metadata::{HarnessMetadata, UnwindMetadata};
 use crate::GotocCtx;
 use cbmc::goto_program::{Expr, Stmt, Symbol};
 use cbmc::InternString;
@@ -13,6 +13,8 @@ use rustc_middle::ty::{self, Instance, TyS};
 use std::collections::BTreeMap;
 use std::iter::FromIterator;
 use tracing::{debug, warn};
+extern crate regex;
+use regex::Regex;
 
 /// Utility to skip functions that can't currently be successfully codgenned.
 impl<'tcx> GotocCtx<'tcx> {
@@ -243,11 +245,33 @@ impl<'tcx> GotocCtx<'tcx> {
     fn handle_kanitool_attributes(&mut self) {
         let instance = self.current_fn().instance();
 
+        // This regex can be expanded as more attributes become available
+        // Can be thought as a menu, in expanded mode
+        let input_re = Regex::new(
+            r#"(?x)
+            (proof) |
+            (unwind_)(\d+)
+            "#,
+        )
+        .unwrap();
+
         for attr in self.tcx.get_attrs(instance.def_id()) {
-            match kanitool_attr_name(attr).as_deref() {
-                Some("proof") => self.handle_kanitool_proof(),
-                _ => {}
-            }
+            if let Some(attribute_string) = kanitool_attr_name(attr).as_deref() {
+                // This regex capture operation returns a vector of strings that can be matched for our purpose
+                // In later Attributes, it will be easy to parse the regex for the attribute arguments, regardless
+                // of the type of the argument itself if we capture all the patterns as a vector of strings of the format
+                // [command, argument]
+                let captures = input_re.captures(attribute_string).map(|captures| {
+                    captures.iter().skip(1).flat_map(|c| c).map(|c| c.as_str()).collect::<Vec<_>>()
+                });
+
+                // Match against the captured values as a slice
+                match captures.as_ref().map(|c| c.as_slice()) {
+                    Some(["proof"]) => self.handle_kanitool_proof(),
+                    Some(["unwind_", x]) => self.handle_kanitool_unwind(x),
+                    _ => {}
+                }
+            };
         }
     }
 
@@ -267,15 +291,36 @@ impl<'tcx> GotocCtx<'tcx> {
 
         self.proof_harnesses.push(harness);
     }
+
+    /// Unwind strings of the format 'unwind_x' and the mangled names are to be parsed for the value 'x'
+    fn handle_kanitool_unwind(&mut self, attribute_string: &str) {
+        // Parse parameter string for value of the argument
+        let argument_value: u32 =
+            attribute_string.parse().expect("Unwind Argument Value is not a number");
+
+        // Make sure the value lies between 0 and u32 max
+        assert!(argument_value < u32::MAX);
+
+        let current_fn = self.current_fn();
+        let mangled_name = current_fn.name();
+
+        let unwind_instance =
+            UnwindMetadata { mangled_name: mangled_name, unwind_value: argument_value };
+        self.unwind_metadata.push(unwind_instance);
+    }
 }
 
 /// 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())
+            let mut new_string = String::new();
+            for index in 1..segments.len() {
+                new_string.push_str(segments[index].ident.as_str());
+            }
+            Some(new_string)
         }
         _ => None,
     }

src/kani-compiler/rustc_codegen_kani/src/compiler_interface.rs

@@ -134,7 +134,8 @@ impl CodegenBackend for GotocCodegenBackend {
             None
         };
 
-        let metadata = KaniMetadata { proof_harnesses: c.proof_harnesses };
+        let metadata =
+            KaniMetadata { proof_harnesses: c.proof_harnesses, unwind_metadata: c.unwind_metadata };
 
         // No output should be generated if user selected no_codegen.
         if !tcx.sess.opts.debugging_opts.no_codegen && tcx.sess.opts.output_types.should_codegen() {

src/kani-compiler/rustc_codegen_kani/src/context/goto_ctx.rs

@@ -14,10 +14,10 @@
 //! Any MIR specific functionality (e.g. codegen etc) should live in specialized files that use
 //! this structure as input.
 use super::current_fn::CurrentFnCtx;
-use super::metadata::HarnessMetadata;
-use super::vtable_ctx::VtableCtx;
+use super::metadata::{HarnessMetadata, UnwindMetadata};
 use crate::overrides::{fn_hooks, GotocHooks};
 use crate::utils::full_crate_name;
+use crate::VtableCtx;
 use cbmc::goto_program::{DatatypeComponent, Expr, Location, Stmt, Symbol, SymbolTable, Type};
 use cbmc::utils::aggr_tag;
 use cbmc::{InternStringOption, InternedString, NO_PRETTY_NAME};
@@ -59,6 +59,7 @@ pub struct GotocCtx<'tcx> {
     pub current_fn: Option<CurrentFnCtx<'tcx>>,
     pub type_map: FxHashMap<InternedString, Ty<'tcx>>,
     pub proof_harnesses: Vec<HarnessMetadata>,
+    pub unwind_metadata: Vec<UnwindMetadata>,
     /// a global counter for generating unique IDs for checks
     pub global_checks_count: u64,
 }
@@ -82,6 +83,7 @@ impl<'tcx> GotocCtx<'tcx> {
             current_fn: None,
             type_map: FxHashMap::default(),
             proof_harnesses: vec![],
+            unwind_metadata: vec![],
             global_checks_count: 0,
         }
     }

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

@@ -19,8 +19,19 @@ pub struct HarnessMetadata {
     pub original_line: String,
 }
 
+/// We emit this struct for every unwind we find
+
+#[derive(Serialize)]
+pub struct UnwindMetadata {
+    /// The name of the function in the CBMC symbol table, being used as a unique identifer.
+    pub mangled_name: String,
+    /// The value of the unwind attribute that the user wants to set
+    pub unwind_value: u32,
+}
+
 /// The structure of `.kani-metadata.json` files, which are emitted for each crate
 #[derive(Serialize)]
 pub struct KaniMetadata {
     pub proof_harnesses: Vec<HarnessMetadata>,
+    pub unwind_metadata: Vec<UnwindMetadata>,
 }

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

@@ -0,0 +1,12 @@
+# 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"

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

@@ -0,0 +1 @@
+line 5 assertion failed: 1 == 2: FAILURE

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

@@ -0,0 +1,41 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+fn main() {
+    assert!(1 == 2);
+}
+
+// rmc-flags: --no-unwinding-checks
+
+// 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.
+
+#[kani::proof]
+#[kani::unwind(10)]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}
+
+#[kani::unwind(8)]
+fn harness_2() {
+    let mut counter = 0;
+    for i in 0..7 {
+        counter += 1;
+        assert!(counter < 5);
+    }
+}
+
+#[kani::unwind(9)]
+fn harness_3() {
+    let mut counter = 0;
+    for i in 0..10 {
+        counter += 1;
+        assert!(counter < 8);
+    }
+}

tests/kani/Unwind-Attribute/fixme_lib.rs

@@ -0,0 +1,37 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+fn main() {
+    assert!(1 == 2);
+}
+
+// rmc-flags: --no-unwinding-checks
+
+// Fix me
+#[kani::proof]
+#[kani::unwind(10)]
+fn harness() {
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        assert!(counter < 10);
+    }
+}
+
+#[kani::unwind(8)]
+fn harness_2() {
+    let mut counter = 0;
+    for i in 0..7 {
+        counter += 1;
+        assert!(counter < 5);
+    }
+}
+
+#[kani::unwind(9)]
+fn harness_3() {
+    let mut counter = 0;
+    for i in 0..10 {
+        counter += 1;
+        assert!(counter < 8);
+    }
+}