Merge pull request #7 from os-checker/remove-attrs

Fix: Only extract kanitool attributes

Author: zjp-CN

src/functions/kani/mod.rs

@@ -2,3 +2,16 @@ mod coercion;
 mod reachability;
 
 pub use reachability::{CallGraph, collect_reachable_items};
+
+use std::sync::LazyLock;
+
+/// `#[kanitool::xxx]` attributes.
+pub static KANI_TOOL_ATTRS: LazyLock<Vec<[String; 2]>> = LazyLock::new(|| {
+    vec![
+        ["kanitool".into(), "proof".into()],
+        ["kanitool".into(), "proof_for_contract".into()],
+        // attrs for contracts
+        ["kanitool".into(), "requires".into()],
+        ["kanitool".into(), "ensures".into()],
+    ]
+});

src/functions/mod.rs

@@ -1,11 +1,11 @@
 use indexmap::IndexSet;
-use kani::{CallGraph, collect_reachable_items};
+use kani::{CallGraph, KANI_TOOL_ATTRS, collect_reachable_items};
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal::internal;
 use rustc_span::{Span, source_map::SourceMap};
 use serde::Serialize;
 use stable_mir::{
-    CrateDef, DefId,
+    CrateDef,
     mir::mono::{Instance, MonoItem},
     ty::{FnDef, RigidTy, Ty, TyKind},
 };
@@ -20,7 +20,7 @@ pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<Function> {
     let mut entries = Vec::with_capacity(cap);
 
     for item in local_items {
-        let _span = debug_span!("all_local_items", ?item).entered();
+        let _span = error_span!("all_local_items", ?item).entered();
 
         let Ok(inst) = Instance::try_from(item).inspect_err(|err| error!(?err)) else { continue };
         entries.push(MonoItem::from(inst));
@@ -72,14 +72,12 @@ impl Function {
         let func = source_code_with(body.span, tcx, src_map);
         info!(" - {:?} ({span:?}): {func}", inst_def.name());
 
-        // FIXME: kanitool and some other proc-macors attributes are generated by parsing,
-        // and these generated attrs share with span of hand-written attrs in source code.
-        // As a result, get_all_attributes through source span will emit duplicated attrs.
-        // Need to fix this in the future, by analyzing Symbols of these attrs.
-        // let attrs = get_all_attributes(inst_def.def_id(), tcx, src_map);
-        let attrs = vec![];
-
-        // TODO: filter in kanitool kind: proof, proof_for_contract, contract, ...
+        // Only need kanitool attrs: proof, proof_for_contract, contract, ...
+        let attrs = KANI_TOOL_ATTRS
+            .iter()
+            .flat_map(|v| inst_def.tool_attrs(v))
+            .map(|attr| attr.as_str().to_owned())
+            .collect();
 
         Some(Function { file, attrs, func, callees })
     }
@@ -113,15 +111,3 @@ fn source_code_with(
     let span = internal(tcx, stable_mir_span);
     source_code(span, src_map)
 }
-
-/// Get all attributes for the item.
-///
-/// We don't call [`all_tool_attrs`], because it only gives tool attributes,
-/// we want as raw attributes as possible.
-///
-/// [`all_tool_attrs`]: https://doc.rust-lang.org/nightly/nightly-rustc/stable_mir/struct.CrateItem.html#impl-CrateDef-for-CrateItem
-fn get_all_attributes(stable_mir_def_id: DefId, tcx: TyCtxt, src_map: &SourceMap) -> Vec<String> {
-    let def_id = internal(tcx, stable_mir_def_id);
-    // dbg!(tcx.hir_body_owned_by(def_id.expect_local())); // HIR body
-    tcx.get_all_attrs(def_id).map(|attr| source_code(attr.span(), src_map)).collect()
-}

tests/snapshots.rs

@@ -1,19 +1,25 @@
 use assert_cmd::Command;
 use expect_test::expect_file;
 
-fn cmd() -> Command {
-    Command::cargo_bin(env!("CARGO_PKG_NAME")).unwrap()
-}
-
-#[test]
-fn standard_proof() {
-    let output = cmd().args(["tests/standard_proof.rs"]).output().unwrap();
+fn cmd(args: &[&str]) -> String {
+    let output = Command::cargo_bin(env!("CARGO_PKG_NAME")).unwrap().args(args).output().unwrap();
     assert!(
         output.status.success(),
         "Failed to test standard_proof.rs:\n{}",
         std::str::from_utf8(&output.stderr).unwrap()
     );
 
-    let stdout = String::from_utf8_lossy(&output.stdout);
-    expect_file!["./snapshots/standard_proof.json"].assert_eq(&stdout);
+    String::from_utf8(output.stdout).unwrap()
+}
+
+#[test]
+fn standard_proofs() {
+    let json = cmd(&["tests/standard_proofs.rs"]);
+    expect_file!["./snapshots/standard_proofs.json"].assert_eq(&json);
+}
+
+#[test]
+fn standard_proofs_with_contracts() {
+    let json = cmd(&["tests/standard_proofs_with_contracts.rs"]);
+    expect_file!["./snapshots/standard_proofs_with_contracts.json"].assert_eq(&json);
 }