feat: define and apply CallGraph::recursive_callees

Author: zjp-CN

src/functions/kani/mod.rs

@@ -1,4 +1,4 @@
 mod coercion;
 mod reachability;
 
-pub use reachability::collect_reachable_items;
+pub use reachability::{CallGraph, collect_reachable_items};

src/functions/kani/reachability.rs

@@ -20,6 +20,8 @@
 extern crate rustc_data_structures;
 extern crate rustc_session;
 
+use super::coercion;
+use indexmap::IndexSet;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
@@ -42,8 +44,6 @@ use std::{
     io::{BufWriter, Write},
 };
 
-use super::coercion;
-
 /// Collect all reachable items starting from the given starting points.
 pub fn collect_reachable_items(
     tcx: TyCtxt,
@@ -68,15 +68,6 @@ pub fn collect_reachable_items(
     // order of the errors and warnings is stable.
     let mut sorted_items: Vec<_> = collector.collected.into_iter().collect();
     sorted_items.sort_by_cached_key(|item| to_fingerprint(tcx, item));
-    // print collected fn names
-    let item_names: Vec<_> = sorted_items
-        .iter()
-        .filter_map(|item| match item {
-            MonoItem::Fn(f) => Some(f.name()),
-            _ => None,
-        })
-        .collect();
-    dbg!(item_names);
     (sorted_items, collector.call_graph)
 }
 
@@ -585,6 +576,25 @@ struct Node(pub MonoItem);
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 struct CollectedNode(pub CollectedItem);
 
+impl CallGraph {
+    pub fn recursive_callees(&self, item: &MonoItem, callees: &mut IndexSet<Instance>) {
+        let key = &Node(item.clone());
+        let nodes = self.edges.get(key).unwrap_or_else(|| panic!("No {item:?} in the call graph."));
+
+        for node in nodes {
+            let item = &node.0.item;
+            match item {
+                MonoItem::Fn(inst) => _ = callees.insert(*inst),
+                // TODO: only consider functions items.
+                // Functions may be in statics, but need testing to comfirm.
+                MonoItem::Static(_) => (),
+                _ => continue,
+            }
+            self.recursive_callees(item, callees);
+        }
+    }
+}
+
 impl CallGraph {
     /// Add a new node into a graph.
     fn add_node(&mut self, item: MonoItem) {

src/functions/mod.rs

@@ -1,11 +1,11 @@
 use indexmap::IndexSet;
-use kani::collect_reachable_items;
+use kani::{CallGraph, 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, CrateItem, DefId, ItemKind,
+    CrateDef, DefId,
     mir::mono::{Instance, MonoItem},
     ty::{FnDef, RigidTy, Ty, TyKind},
 };
@@ -17,7 +17,6 @@ pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<Function> {
     let local_items = stable_mir::all_local_items();
     let cap = local_items.len();
 
-    let mut outputs = Vec::with_capacity(cap);
     let mut entries = Vec::with_capacity(cap);
 
     for item in local_items {
@@ -29,7 +28,7 @@ pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<Function> {
 
     let (mono_items, callgraph) = collect_reachable_items(tcx, &entries);
 
-    outputs
+    mono_items.iter().filter_map(|item| Function::new(item, &callgraph, tcx, src_map)).collect()
 }
 
 /// A Rust funtion with its file source, attributes, and raw function content.
@@ -48,37 +47,39 @@ pub struct Function {
 }
 
 impl Function {
-    pub fn new(item: CrateItem, tcx: TyCtxt, src_map: &SourceMap) -> Option<Self> {
-        if !matches!(item.kind(), ItemKind::Fn) {
-            // skip non fn items
+    pub fn new(
+        item: &MonoItem,
+        callgraph: &CallGraph,
+        tcx: TyCtxt,
+        src_map: &SourceMap,
+    ) -> Option<Self> {
+        // skip non fn items
+        let &MonoItem::Fn(inst) = item else {
             return None;
-        }
-        // item.emit_mir(&mut std::io::stdout()).unwrap(); // MIR body
-        let inst = Instance::try_from(item).inspect_err(|err| error!(?err)).ok()?;
+        };
+
         let fn_def = ty_to_fndef(inst.ty())?;
-        let file = item.span().get_filename();
+        let inst_def = inst.def;
+        let span = inst_def.span();
+
+        let file = span.get_filename();
         let body = fn_def.body()?;
 
         let mut callees = IndexSet::new();
-        // retrieve direct calls
-        callees.extend(callees::calls_in_body(&body));
-        // recursive calls
-        let direct_calls: Vec<_> = callees.iter().copied().collect();
-        for call in direct_calls {
-            callees::recursive_callees(call, &mut callees);
-        }
-        let callees = callees.into_iter().map(|x| format!("{x:?}")).collect();
+        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!(" - {:?} ({:?}): {func}", item.name(), item.span());
+        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(item.def_id(), tcx, src_map);
+        // let attrs = get_all_attributes(inst_def.def_id(), tcx, src_map);
+        let attrs = vec![];
 
-        // TODO: kanitool kind: proof, proof_for_contract, contract, ...
+        // TODO: filter in kanitool kind: proof, proof_for_contract, contract, ...
 
         Some(Function { file, attrs, func, callees })
     }