From 33e226b1ad32404207c39c5bcbabe1cf969c64de Mon Sep 17 00:00:00 2001
From: Aaron Bembenek <113437016+aaronbembenek-aws@users.noreply.github.com>
Date: Mon, 31 Oct 2022 17:18:19 -0400
Subject: [PATCH] Add `kani::stub` attribute (#1814)

---
 .../codegen_cprover_gotoc/codegen/function.rs |  4 +++
 library/kani_macros/src/lib.rs                | 27 +++++++++++++++++++
 tests/ui/stub-attribute/attribute.rs          | 12 +++++++++
 tests/ui/stub-attribute/expected              |  1 +
 4 files changed, 44 insertions(+)
 create mode 100644 tests/ui/stub-attribute/attribute.rs
 create mode 100644 tests/ui/stub-attribute/expected

diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs
index a770775cdd46..6abba55ac955 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs
@@ -375,6 +375,10 @@ impl<'tcx> GotocCtx<'tcx> {
         let mut harness = self.default_kanitool_proof();
         for attr in other_attributes.iter() {
             match attr.0.as_str() {
+                "stub" => self
+                    .tcx
+                    .sess
+                    .span_warn(attr.1.span, "Attribute `kani::stub` is currently ignored by Kani"),
                 "unwind" => self.handle_kanitool_unwind(attr.1, &mut harness),
                 _ => {
                     self.tcx.sess.span_err(
diff --git a/library/kani_macros/src/lib.rs b/library/kani_macros/src/lib.rs
index d9ebe55a2e38..aa9ee4799a8b 100644
--- a/library/kani_macros/src/lib.rs
+++ b/library/kani_macros/src/lib.rs
@@ -118,3 +118,30 @@ pub fn unwind(attr: TokenStream, item: TokenStream) -> TokenStream {
     result.extend(item);
     result
 }
+
+#[cfg(not(kani))]
+#[proc_macro_attribute]
+pub fn stub(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    // When the config is not kani, we should leave the function alone
+    item
+}
+
+/// Specify a function/method stub pair to use for proof harness
+///
+/// The attribute `#[kani::stub(original, replacement)]` can only be used alongside `#[kani::proof]`.
+///
+/// # Arguments
+/// * `original` - The function or method to replace, specified as a path.
+/// * `replacement` - The function or method to use as a replacement, specified as a path.
+#[cfg(kani)]
+#[proc_macro_attribute]
+pub fn stub(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let mut result = TokenStream::new();
+
+    // Translate #[kani::stub(original, replacement)] to #[kanitool::stub(original, replacement)]
+    let insert_string = "#[kanitool::stub(".to_owned() + &attr.to_string() + ")]";
+    result.extend(insert_string.parse::<TokenStream>().unwrap());
+
+    result.extend(item);
+    result
+}
diff --git a/tests/ui/stub-attribute/attribute.rs b/tests/ui/stub-attribute/attribute.rs
new file mode 100644
index 000000000000..2e3b81fb09ea
--- /dev/null
+++ b/tests/ui/stub-attribute/attribute.rs
@@ -0,0 +1,12 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+//! Checks that the `kani::stub` attribute is accepted
+
+fn foo() {}
+
+fn bar() {}
+
+#[kani::proof]
+#[kani::stub(foo, bar)]
+fn main() {}
diff --git a/tests/ui/stub-attribute/expected b/tests/ui/stub-attribute/expected
new file mode 100644
index 000000000000..7610fce8e249
--- /dev/null
+++ b/tests/ui/stub-attribute/expected
@@ -0,0 +1 @@
+warning: Attribute `kani::stub` is currently ignored by Kani