Merge pull request #9 from os-checker/filter-functions

Feat: Filter in functions annotated by kani

Author: zjp-CN

src/functions/callees.rs

@@ -3,17 +3,18 @@ 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,
+    CrateDef, DefId,
+    crate_def::Attribute,
     mir::mono::{Instance, MonoItem},
     ty::{FnDef, RigidTy, Ty, TyKind},
 };
 
-mod callees;
 mod kani;
+mod serialization;
+pub use serialization::SerFunction;
 
-pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<Function> {
+pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<SerFunction> {
     let local_items = stable_mir::all_local_items();
     let cap = local_items.len();
 
@@ -28,22 +29,29 @@ pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<Function> {
 
     let (mono_items, callgraph) = collect_reachable_items(tcx, &entries);
 
-    mono_items.iter().filter_map(|item| Function::new(item, &callgraph, tcx, src_map)).collect()
+    // Filter out non kanitool functions.
+    mono_items
+        .iter()
+        .filter_map(|item| Function::new(item, &callgraph, tcx, src_map, |x| !x.attrs.is_empty()))
+        .map(SerFunction::new)
+        .collect()
 }
 
 /// A Rust funtion with its file source, attributes, and raw function content.
-#[derive(Debug, Default, Serialize)]
+#[derive(Debug)]
 pub struct Function {
+    /// DefId in stable_mir.
+    def_id: DefId,
     /// Every funtion must be declared in a specific file, even for those
     /// generated by macros.
     file: String,
     /// Attributes are attached the function, but it seems that attributes
     /// and function must be separated to query.
-    attrs: Vec<String>,
+    attrs: Vec<Attribute>,
     /// Raw function string, including name, signature, and body.
     func: String,
     /// Recursive fnction calls inside the body.
-    callees: Vec<String>,
+    callees: IndexSet<Instance>,
 }
 
 impl Function {
@@ -52,34 +60,33 @@ impl Function {
         callgraph: &CallGraph,
         tcx: TyCtxt,
         src_map: &SourceMap,
+        filter: impl FnOnce(&Self) -> bool,
     ) -> Option<Self> {
         // skip non fn items
         let &MonoItem::Fn(inst) = item else {
             return None;
         };
 
-        let fn_def = ty_to_fndef(inst.ty())?;
         let inst_def = inst.def;
-        let span = inst_def.span();
 
+        // Only need kanitool attrs: proof, proof_for_contract, contract, ...
+        let attrs = KANI_TOOL_ATTRS.iter().flat_map(|v| inst_def.tool_attrs(v)).collect();
+
+        let span = inst_def.span();
         let file = span.get_filename();
+
+        let fn_def = ty_to_fndef(inst.ty())?;
         let body = fn_def.body()?;
 
         let mut callees = IndexSet::new();
         callgraph.recursive_callees(item, &mut callees);
-        let callees = callees.into_iter().map(|x| format!("{:?}", x.def.def_id())).collect();
 
         let func = source_code_with(body.span, tcx, src_map);
         info!(" - {:?} ({span:?}): {func}", inst_def.name());
 
-        // 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 })
+        let def_id = inst_def.def_id();
+        let this = Function { def_id, file, attrs, func, callees };
+        filter(&this).then_some(this)
     }
 }
 

src/functions/mod.rs

@@ -0,0 +1,35 @@
+use serde::Serialize;
+use stable_mir::{CrateDef, DefId};
+
+/// A Rust funtion with its file source, attributes, and raw function content.
+#[derive(Debug, Serialize)]
+pub struct SerFunction {
+    /// DefId in stable_mir.
+    pub def_id: String,
+    /// Every funtion must be declared in a specific file, even for those
+    /// generated by macros.
+    pub file: String,
+    /// Attributes are attached the function, but it seems that attributes
+    /// and function must be separated to query.
+    pub attrs: Vec<String>,
+    /// Raw function string, including name, signature, and body.
+    pub func: String,
+    /// Recursive fnction calls inside the body.
+    pub callees: Vec<String>,
+}
+
+impl SerFunction {
+    pub fn new(fun: super::Function) -> Self {
+        SerFunction {
+            def_id: format_def_id(&fun.def_id),
+            file: fun.file,
+            attrs: fun.attrs.iter().map(|a| a.as_str().to_owned()).collect(),
+            func: fun.func,
+            callees: fun.callees.into_iter().map(|x| format_def_id(&x.def.def_id())).collect(),
+        }
+    }
+}
+
+fn format_def_id(def_id: &DefId) -> String {
+    format!("{def_id:?}")
+}