Merge pull request #32 from os-checker/test-macro-expansion

Feat: support spans on macro expansion

Author: zjp-CN

kani

@@ -1,19 +1,20 @@
 use indexmap::IndexSet;
 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 rustc_span::source_map::SourceMap;
 use stable_mir::{
-    CrateDef, DefId,
+    CrateDef,
     crate_def::Attribute,
-    mir::mono::{Instance, MonoItem},
-    ty::{FnDef, RigidTy, Ty, TyKind},
+    mir::{
+        Body,
+        mono::{Instance, MonoItem},
+    },
 };
-use std::cmp::Ordering;
 
 mod kani;
 mod serialization;
 pub use serialization::SerFunction;
+mod utils;
 
 pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<SerFunction> {
     let local_items = stable_mir::all_local_items();
@@ -34,25 +35,25 @@ pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<SerFunction> {
     let mut proofs: Vec<_> = mono_items
         .iter()
         .filter_map(|f| Function::new(f, &callgraph, tcx, src_map, |x| !x.attrs.is_empty()))
+        .map(|fun| SerFunction::new(fun, tcx, src_map))
         .collect();
     // Sort proofs by file path and source code.
-    proofs.sort_by(|a, b| (&*a.file, &*a.func).cmp(&(&*b.file, &*b.func)));
-    proofs.into_iter().map(|fun| SerFunction::new(fun, tcx, src_map)).collect()
+    proofs.sort_by(|a, b| a.cmp_by_file_and_func(b));
+    proofs
 }
 
 /// A Rust funtion with its file source, attributes, and raw function content.
 #[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.
+    /// Instance of the function.
+    instance: Instance,
+
+    /// kanitool's attributs.
     attrs: Vec<Attribute>,
-    /// Raw function string, including name, signature, and body.
-    func: String,
+
+    /// Function body.
+    body: Body,
+
     /// Recursive fnction calls inside the body.
     /// The elements are sorted by file path and fn source code to keep hash value stable.
     callees: IndexSet<Instance>,
@@ -66,98 +67,22 @@ impl Function {
         src_map: &SourceMap,
         filter: impl FnOnce(&Self) -> bool,
     ) -> Option<Self> {
-        // skip non fn items
-        let &MonoItem::Fn(inst) = item else {
+        // Skip non fn items
+        let &MonoItem::Fn(instance) = item else {
             return None;
         };
 
-        let inst_def = inst.def;
+        // Skip if no body.
+        let body = instance.body()?;
 
         // 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 = file_path(&inst);
-
-        let fn_def = ty_to_fndef(inst.ty())?;
-        let body = fn_def.body()?;
+        let attrs = KANI_TOOL_ATTRS.iter().flat_map(|v| instance.def.tool_attrs(v)).collect();
 
         let mut callees = IndexSet::new();
         callgraph.recursive_callees(item, &mut callees);
-        callees.sort_by(|a, b| cmp_callees(a, b, tcx, src_map));
-
-        let func = source_code_with(body.span, tcx, src_map);
-        info!(" - {:?} ({span:?}): {func}", inst_def.name());
+        callees.sort_by(|a, b| utils::cmp_callees(a, b, tcx, src_map));
 
-        let def_id = inst_def.def_id();
-        let this = Function { def_id, file, attrs, func, callees };
+        let this = Function { instance, attrs, body, callees };
         filter(&this).then_some(this)
     }
 }
-
-/// 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)
-}
-
-fn source_code_of_body(inst: &Instance, tcx: TyCtxt, src_map: &SourceMap) -> Option<String> {
-    inst.body().map(|body| source_code_with(body.span, tcx, src_map))
-}
-
-fn cmp_callees(a: &Instance, b: &Instance, tcx: TyCtxt, src_map: &SourceMap) -> Ordering {
-    let filename_a = file_path(a);
-    let filename_b = file_path(b);
-    match filename_a.cmp(&filename_b) {
-        Ordering::Equal => (),
-        ord => return ord,
-    }
-
-    let body_a = source_code_of_body(a, tcx, src_map);
-    let body_b = source_code_of_body(b, tcx, src_map);
-    body_a.cmp(&body_b)
-}
-
-fn file_path(inst: &Instance) -> String {
-    use std::sync::LazyLock;
-    static PREFIXES: LazyLock<[String; 2]> = LazyLock::new(|| {
-        let mut pwd = std::env::current_dir().unwrap().into_os_string().into_string().unwrap();
-        pwd.push('/');
-
-        let out = std::process::Command::new("rustc").arg("--print=sysroot").output().unwrap();
-        let sysroot = std::str::from_utf8(&out.stdout).unwrap().trim();
-        let sysroot = format!("{sysroot}/lib/rustlib/src/rust/");
-        [pwd, sysroot]
-    });
-
-    let file = inst.def.span().get_filename();
-    for prefix in &*PREFIXES {
-        if let Some(file) = file.strip_prefix(prefix) {
-            return file.to_owned();
-        }
-    }
-    file
-}

src/functions/mod.rs

@@ -1,9 +1,10 @@
+use super::utils::{self, SourceCode};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::source_map::SourceMap;
 use rustc_stable_hash::{FromStableHash, SipHasher128Hash, StableHasher, hashers::SipHasher128};
 use serde::Serialize;
-use stable_mir::{CrateDef, DefId, mir::mono::Instance};
-use std::hash::Hasher;
+use stable_mir::{CrateDef, mir::mono::Instance};
+use std::{cmp::Ordering, hash::Hasher};
 
 /// A Rust funtion with its file source, attributes, and raw function content.
 #[derive(Debug, Serialize)]
@@ -18,34 +19,41 @@ pub struct SerFunction {
     /// and function must be separated to query.
     attrs: Vec<String>,
     /// Raw function string, including name, signature, and body.
-    func: String,
+    func: SourceCode,
     /// Recursive function calls inside the proof.
     callees: Vec<Callee>,
 }
 
 impl SerFunction {
     pub fn new(fun: super::Function, tcx: TyCtxt, src_map: &SourceMap) -> Self {
-        let def_id = format_def_id(&fun.def_id);
-        let file = fun.file;
+        let inst = fun.instance;
+        let def_id = format_def_id(&inst);
+        let file = utils::file_path(&inst);
         let attrs: Vec<_> = fun.attrs.iter().map(|a| a.as_str().to_owned()).collect();
-        let func = fun.func;
+        // Though this is from body span, fn name and signature are included.
+        let func = utils::source_code_with(fun.body.span, tcx, src_map);
         let callees: Vec<_> = fun.callees.iter().map(|x| Callee::new(x, tcx, src_map)).collect();
 
         // Hash
         let mut hasher = StableHasher::<SipHasher128>::new();
         hasher.write_str(&file);
-        hasher.write_str(&func);
+        func.with_hasher(&mut hasher);
         hasher.write_length_prefix(attrs.len());
         attrs.iter().for_each(|a| hasher.write_str(a));
         hasher.write_length_prefix(callees.len());
         callees.iter().for_each(|c| {
             hasher.write_str(&c.file);
-            hasher.write_str(&c.func);
+            c.func.with_hasher(&mut hasher);
         });
         let Hash128(hash) = hasher.finish();
 
         SerFunction { hash, def_id, file, attrs, func, callees }
     }
+
+    /// Compare by file and func string.
+    pub fn cmp_by_file_and_func(&self, other: &Self) -> Ordering {
+        (&*self.file, &self.func).cmp(&(&*other.file, &other.func))
+    }
 }
 
 // ************* hash *************
@@ -60,23 +68,22 @@ impl FromStableHash for Hash128 {
 }
 // ************* hash *************
 
-fn format_def_id(def_id: &DefId) -> String {
-    format!("{def_id:?}")
+fn format_def_id(inst: &Instance) -> String {
+    format!("{:?}", inst.def.def_id())
 }
 
 #[derive(Debug, Serialize)]
 pub struct Callee {
     def_id: String,
     file: String,
-    func: String,
+    func: SourceCode,
 }
 
 impl Callee {
     fn new(inst: &Instance, tcx: TyCtxt, src_map: &SourceMap) -> Self {
-        let inst_def = &inst.def;
-        let def_id = format_def_id(&inst_def.def_id());
-        let file = super::file_path(inst);
-        let func = super::source_code_of_body(inst, tcx, src_map).unwrap_or_default();
+        let def_id = format_def_id(inst);
+        let file = utils::file_path(inst);
+        let func = utils::source_code_of_body(inst, tcx, src_map).unwrap_or_default();
         Callee { def_id, file, func }
     }
 }

src/functions/serialization.rs

@@ -0,0 +1,111 @@
+use rustc_middle::ty::TyCtxt;
+use rustc_smir::rustc_internal::internal;
+use rustc_span::{Span, source_map::SourceMap};
+use rustc_stable_hash::{StableHasher, hashers::SipHasher128};
+use serde::Serialize;
+use stable_mir::{CrateDef, mir::mono::Instance};
+use std::{cmp::Ordering, hash::Hasher};
+
+/// Source code and potential source code before expansion.
+///
+/// Refer to
+#[derive(Debug, Serialize, PartialEq, Eq, PartialOrd, Ord, Default)]
+pub struct SourceCode {
+    // TODO:
+    // file: String,
+    //
+    /// Source that a stable_mir span points to.
+    pub src: String,
+    /// Is the stable_mir span from a macro expansion?
+    /// If it is from an expansion, what's the source code before expansion?
+    /// * Some(_) happens when the src (stable_mir) span comes from expansion, and tells
+    ///   the source before the expansion.
+    /// * None if the src is not from a macro expansion.
+    pub before_expansion: Option<String>,
+}
+
+impl SourceCode {
+    pub fn with_hasher(&self, hasher: &mut StableHasher<SipHasher128>) {
+        hasher.write_str(&self.src);
+        match &self.before_expansion {
+            Some(text) => {
+                hasher.write_u8(1);
+                hasher.write_str(text);
+            }
+            None => hasher.write_u8(0),
+        }
+    }
+}
+
+fn span_to_source(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 span.
+fn source_code(span: Span, src_map: &SourceMap) -> SourceCode {
+    let src = span_to_source(span, src_map);
+    let before_expansion = span.from_expansion().then(|| {
+        let ancestor_span = span.find_oldest_ancestor_in_same_ctxt();
+        span_to_source(ancestor_span, src_map)
+    });
+    SourceCode { src, before_expansion }
+}
+
+/// Source code for a stable_mir span.
+pub fn source_code_with(
+    stable_mir_span: stable_mir::ty::Span,
+    tcx: TyCtxt,
+    src_map: &SourceMap,
+) -> SourceCode {
+    let span = internal(tcx, stable_mir_span);
+    source_code(span, src_map)
+}
+
+// FIXME: takes body and returns SourceCode
+pub fn source_code_of_body(
+    inst: &Instance,
+    tcx: TyCtxt,
+    src_map: &SourceMap,
+) -> Option<SourceCode> {
+    inst.body().map(|body| source_code_with(body.span, tcx, src_map))
+}
+
+pub fn cmp_callees(a: &Instance, b: &Instance, tcx: TyCtxt, src_map: &SourceMap) -> Ordering {
+    let filename_a = file_path(a);
+    let filename_b = file_path(b);
+    match filename_a.cmp(&filename_b) {
+        Ordering::Equal => (),
+        ord => return ord,
+    }
+
+    let body_a = source_code_of_body(a, tcx, src_map);
+    let body_b = source_code_of_body(b, tcx, src_map);
+    body_a.cmp(&body_b)
+}
+
+pub fn file_path(inst: &Instance) -> String {
+    use std::sync::LazyLock;
+    static PREFIXES: LazyLock<[String; 2]> = LazyLock::new(|| {
+        let mut pwd = std::env::current_dir().unwrap().into_os_string().into_string().unwrap();
+        pwd.push('/');
+
+        let out = std::process::Command::new("rustc").arg("--print=sysroot").output().unwrap();
+        let sysroot = std::str::from_utf8(&out.stdout).unwrap().trim();
+        let sysroot = format!("{sysroot}/lib/rustlib/src/rust/");
+        [pwd, sysroot]
+    });
+
+    let file = inst.def.span().get_filename();
+    for prefix in &*PREFIXES {
+        if let Some(file) = file.strip_prefix(prefix) {
+            return file.to_owned();
+        }
+    }
+    file
+}