Merge pull request #4 from os-checker/stable_mir

Switch from HIR to StableMIR

Author: zjp-CN

src/cli/mod.rs

@@ -1,5 +1,4 @@
 use clap::Parser;
-use std::path::PathBuf;
 
 /// Parse cli arguments.
 pub fn parse() -> Args {
@@ -10,9 +9,12 @@ pub fn parse() -> Args {
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 pub struct Args {
-    /// Json output file path. Print to stdout if not set.
+    /// Possible one of these values:
+    /// * `--json false`: skip serializing to json
+    /// * `--json path/to/file.json`
+    /// * print to stdout if not set
     #[arg(long)]
-    pub json: Option<PathBuf>,
+    pub json: Option<String>,
 
     /// Args for rustc.
     pub rustc_args: Vec<String>,

src/functions/call_graph.rs

@@ -0,0 +1,56 @@
+use super::ty_to_fndef;
+use indexmap::IndexSet;
+use stable_mir::{CrateDef, mir::*, ty::FnDef};
+
+pub fn calls_in_body(body: &Body) -> Vec<FnDef> {
+    let locals = body.locals();
+    let mut direct = Vec::new();
+    for basic_block in &body.blocks {
+        terminator_to_fn_def(&basic_block.terminator, locals, &mut direct);
+    }
+    // all direct calls in a body
+    direct
+}
+
+/// Extract fn call from the terminator.
+fn terminator_to_fn_def(terminator: &Terminator, locals: &[LocalDecl], direct: &mut Vec<FnDef>) {
+    let TerminatorKind::Call { func, args, .. } = &terminator.kind else { return };
+
+    push_fn_def(func, locals, direct);
+
+    // FIXME: Will there be a FnDef in an arg in MIR? Probably no?
+    // Anyway, need find a test for this. But now assume yes.
+    // Maybe a function pointer or raw pointer cast to fn pointer suit?
+    for arg in args {
+        push_fn_def(arg, locals, direct);
+    }
+
+    // Will there ever be a FnDef in destination? Skip destination for now.
+}
+
+/// Push a FnDef Operand if any.
+fn push_fn_def(func: &Operand, locals: &[LocalDecl], direct: &mut Vec<FnDef>) {
+    let ty = func.ty(locals).expect("malformed operand or incompatible locals list");
+    if let Some(fn_def) = ty_to_fndef(ty) {
+        direct.push(fn_def);
+    }
+}
+
+/// Recursively retrieve calls for a call.
+pub fn recursive_callees(fn_def: FnDef, visited: &mut IndexSet<FnDef>) {
+    // print!("[{}] ", fn_def.name());
+
+    let Some(body) = fn_def.body() else { return };
+
+    // let mut buf = Vec::with_capacity(1024);
+    // body.dump(&mut buf, &format!("{fn_def:?}")).unwrap();
+    // println!("{}", String::from_utf8(buf).unwrap());
+
+    let mut direct = calls_in_body(&body);
+    while let Some(call) = direct.pop() {
+        // the call hasn't been reached, traverse
+        if visited.insert(call) {
+            recursive_callees(call, visited);
+        }
+    }
+}

src/functions/callees.rs

@@ -1,9 +1,15 @@
-use rustc_hir::{ExprKind, Item, ItemKind};
+use indexmap::IndexSet;
 use rustc_middle::ty::TyCtxt;
-use rustc_span::source_map::SourceMap;
+use rustc_smir::rustc_internal::internal;
+use rustc_span::{Span, source_map::SourceMap};
 use serde::Serialize;
+use stable_mir::{
+    CrateDef, CrateItem, DefId, ItemKind,
+    mir::mono::Instance,
+    ty::{FnDef, RigidTy, Ty, TyKind},
+};
 
-mod call_graph;
+mod callees;
 
 /// A Rust funtion with its file source, attributes, and raw function content.
 #[derive(Debug, Default, Serialize)]
@@ -16,54 +22,84 @@ pub struct Function {
     attrs: Vec<String>,
     /// Raw function string, including name, signature, and body.
     func: String,
+    /// Recursive fnction calls inside the body.
+    callees: Vec<String>,
 }
 
 impl Function {
-    pub fn new(item: &Item, src_map: &SourceMap, tcx: TyCtxt) -> Option<Self> {
-        if let ItemKind::Fn { has_body, body, .. } = &item.kind
-            && *has_body
-        {
-            let mut func = Function {
-                file: src_map
-                    .span_to_filename(item.span)
-                    .prefer_remapped_unconditionaly()
-                    .to_string(),
-                ..Default::default()
-            };
+    pub fn new(item: CrateItem, tcx: TyCtxt, src_map: &SourceMap) -> Option<Self> {
+        if !matches!(item.kind(), ItemKind::Fn) {
+            // skip non fn items
+            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 body = fn_def.body()?;
 
-            // add attributes
-            for (i, attr) in tcx.get_all_attrs(item.owner_id).enumerate() {
-                // FIXME: kani rewrites attributes from source to `#[allow(dead_code)]`
-                // and `#[kanitool::...]`, but share the same span with the source span.
-                // As a result, repeated `#[kani::proof]` are obtained from the source text,
-                // but `attr: &AttrItem` contains real expanded attributes.
-                let src_attr = src_map
-                    .span_to_source(attr.span(), |text, x, y| {
-                        let src = &text[x..y];
-                        debug!("[attr {i}] [{x}:{y}]\n{src}");
-                        Ok(src.to_owned())
-                    })
-                    .unwrap();
-                func.attrs.push(src_attr);
-            }
+        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();
 
-            // add function
-            func.func = src_map
-                .span_to_source(item.span, |text, x, y| {
-                    let src = &text[x..y];
-                    debug!("[{x}:{y}]\n{src}");
-                    Ok(src.to_owned())
-                })
-                .unwrap();
+        let func = source_code_with(body.span, tcx, src_map);
+        info!(" - {:?} ({:?}): {func}", item.name(), item.span());
 
-            // dbg!(tcx.hir_body(*body));
-            let fn_body = tcx.hir_body(*body);
-            if let ExprKind::Block(block, _) = fn_body.value.kind {
-                dbg!(block);
-            }
+        // 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);
 
-            return Some(func);
-        }
-        None
+        // TODO: kanitool kind: proof, proof_for_contract, contract, ...
+
+        Some(Function { file, attrs, func, callees })
     }
 }
+
+/// Extract FnDef from Ty.
+fn ty_to_fndef(ty: Ty) -> Option<FnDef> {
+    let TyKind::RigidTy(RigidTy::FnDef(fn_def, _)) = ty.kind() else {
+        return None;
+    };
+    Some(fn_def)
+}
+
+/// Source code for a span.
+fn source_code(span: Span, src_map: &SourceMap) -> String {
+    src_map
+        .span_to_source(span, |text, x, y| {
+            let src = &text[x..y];
+            debug!("[{x}:{y}]\n{src}");
+            Ok(src.to_owned())
+        })
+        .unwrap()
+}
+
+/// Source code for a stable_mir span.
+fn source_code_with(
+    stable_mir_span: stable_mir::ty::Span,
+    tcx: TyCtxt,
+    src_map: &SourceMap,
+) -> String {
+    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()
+}