From 80a4914989ea1362e7ee54732e8b3ce82869d0cf Mon Sep 17 00:00:00 2001
From: Carolyn Zech <cmzech@amazon.com>
Date: Tue, 17 Sep 2024 15:56:27 -0400
Subject: [PATCH] call check_proof_attribute for contract harnesses

---
 kani-compiler/src/kani_middle/attributes.rs  | 27 +++++++++++++++-----
 tests/ui/invalid-contract-harness/expected   | 17 ++++++++++++
 tests/ui/invalid-contract-harness/invalid.rs | 25 ++++++++++++++++++
 tests/ui/invalid-harnesses/expected          |  2 +-
 tests/ui/mir-linker/generic-harness/expected |  2 +-
 5 files changed, 65 insertions(+), 8 deletions(-)
 create mode 100644 tests/ui/invalid-contract-harness/expected
 create mode 100644 tests/ui/invalid-contract-harness/invalid.rs

diff --git a/kani-compiler/src/kani_middle/attributes.rs b/kani-compiler/src/kani_middle/attributes.rs
index 9a3ff7c1d6a6..86e97f633dfb 100644
--- a/kani-compiler/src/kani_middle/attributes.rs
+++ b/kani-compiler/src/kani_middle/attributes.rs
@@ -358,7 +358,7 @@ impl<'tcx> KaniAttributes<'tcx> {
                         );
                     }
                     expect_single(self.tcx, kind, &attrs);
-                    attrs.iter().for_each(|attr| self.check_proof_attribute(attr))
+                    attrs.iter().for_each(|attr| self.check_proof_attribute(kind, attr))
                 }
                 KaniAttributeKind::Unstable => attrs.iter().for_each(|attr| {
                     let _ = UnstableAttribute::try_from(*attr).map_err(|err| err.report(self.tcx));
@@ -370,6 +370,7 @@ impl<'tcx> KaniAttributes<'tcx> {
                         );
                     }
                     expect_single(self.tcx, kind, &attrs);
+                    attrs.iter().for_each(|attr| self.check_proof_attribute(kind, attr))
                 }
                 KaniAttributeKind::StubVerified => {
                     expect_single(self.tcx, kind, &attrs);
@@ -583,15 +584,29 @@ impl<'tcx> KaniAttributes<'tcx> {
     }
 
     /// Check that if this item is tagged with a proof_attribute, it is a valid harness.
-    fn check_proof_attribute(&self, proof_attribute: &Attribute) {
+    fn check_proof_attribute(&self, kind: KaniAttributeKind, proof_attribute: &Attribute) {
         let span = proof_attribute.span;
         let tcx = self.tcx;
-        expect_no_args(tcx, KaniAttributeKind::Proof, proof_attribute);
+        if let KaniAttributeKind::Proof = kind {
+            expect_no_args(tcx, kind, proof_attribute);
+        }
+
         if tcx.def_kind(self.item) != DefKind::Fn {
-            tcx.dcx().span_err(span, "the `proof` attribute can only be applied to functions");
+            tcx.dcx().span_err(
+                span,
+                format!(
+                    "the '#[kani::{}]' attribute can only be applied to functions",
+                    kind.as_ref()
+                ),
+            );
         } else if tcx.generics_of(self.item).requires_monomorphization(tcx) {
-            tcx.dcx()
-                .span_err(span, "the `proof` attribute cannot be applied to generic functions");
+            tcx.dcx().span_err(
+                span,
+                format!(
+                    "the '#[kani::{}]' attribute cannot be applied to generic functions",
+                    kind.as_ref()
+                ),
+            );
         } else {
             let instance = Instance::mono(tcx, self.item);
             if !super::fn_abi(tcx, instance).args.is_empty() {
diff --git a/tests/ui/invalid-contract-harness/expected b/tests/ui/invalid-contract-harness/expected
new file mode 100644
index 000000000000..1e07c4ca4b16
--- /dev/null
+++ b/tests/ui/invalid-contract-harness/expected
@@ -0,0 +1,17 @@
+error: only one '#[kani::proof_for_contract]' attribute is allowed per harness\
+invalid.rs:\
+|\
+| #[kani::proof_for_contract(foo)]\
+| ^^^^^^^^^^^^^^
+
+error: functions used as harnesses cannot have any arguments\
+invalid.rs:\
+|\
+| #[kani::proof_for_contract(foo)]
+| ^^^^^^^^^^^^^^
+
+error: the '#[kani::proof_for_contract]' attribute cannot be applied to generic functions\
+invalid.rs:\
+|\
+| #[kani::proof_for_contract(foo)]\
+| ^^^^^^^^^^^^^^
diff --git a/tests/ui/invalid-contract-harness/invalid.rs b/tests/ui/invalid-contract-harness/invalid.rs
new file mode 100644
index 000000000000..e30382e3c602
--- /dev/null
+++ b/tests/ui/invalid-contract-harness/invalid.rs
@@ -0,0 +1,25 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Zfunction-contracts
+
+// This test is to check Kani's error handling of invalid usages of the `proof_for_contract` harness.
+// We also ensure that all errors and warnings are printed in one compilation.
+
+#[kani::requires(true)]
+fn foo() {}
+
+#[kani::proof_for_contract(foo)]
+#[kani::proof_for_contract(foo)]
+fn multiple_proof_annotations() {
+    foo();
+}
+
+#[kani::proof_for_contract(foo)]
+fn proof_with_arg(arg: bool) {
+    foo();
+}
+
+#[kani::proof_for_contract(foo)]
+fn generic_harness<T: Default>() {
+    foo();
+}
diff --git a/tests/ui/invalid-harnesses/expected b/tests/ui/invalid-harnesses/expected
index 7def51a0d85a..265af44a1685 100644
--- a/tests/ui/invalid-harnesses/expected
+++ b/tests/ui/invalid-harnesses/expected
@@ -10,7 +10,7 @@ invalid.rs:\
 | #[kani::proof]
 | ^^^^^^^^^^^^^^
 
-error: the `proof` attribute cannot be applied to generic functions\
+error: the '#[kani::proof]' attribute cannot be applied to generic functions\
 invalid.rs:\
 |\
 | #[kani::proof]\
diff --git a/tests/ui/mir-linker/generic-harness/expected b/tests/ui/mir-linker/generic-harness/expected
index 0798bb9e99a3..176ca1a5b7db 100644
--- a/tests/ui/mir-linker/generic-harness/expected
+++ b/tests/ui/mir-linker/generic-harness/expected
@@ -1 +1 @@
-error: the `proof` attribute cannot be applied to generic functions
+error: the '#[kani::proof]' attribute cannot be applied to generic functions