Feat: Cache StableMir results and basic compiler ctx

src/functions/cache.rs

@@ -0,0 +1,114 @@
+//! Access cached data from local thread set for the given Instance.
+//! If the data hasn't been available, generate one and insert it.
+//! The data is always behind a borrow through the `get_*` callbacks.
+
+use super::utils::{SourceCode, source_code_with};
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::source_map::{SourceMap, get_source_map};
+use stable_mir::{
+    CrateDef, DefId,
+    mir::{Body, mono::Instance},
+};
+use std::{cell::RefCell, cmp::Ordering, sync::Arc};
+
+thread_local! {
+    static CACHE: RefCell<Cache> = RefCell::new(Cache::new());
+}
+
+pub fn set_rustc_ctx(tcx: TyCtxt) {
+    // Safety: TyCtxt<'short> is extended to TyCtxt<'static>,
+    // and we only use TyCtxt<'static> in stable_mir's callback.
+    let tcx = unsafe { std::mem::transmute::<TyCtxt<'_>, TyCtxt<'static>>(tcx) };
+    let src_map = get_source_map().expect("No source map.");
+    let rustc = RustcCxt { tcx, src_map };
+    CACHE.with(|c| c.borrow_mut().rustc = Some(rustc));
+}
+
+pub fn clear_rustc_ctx() {
+    CACHE.with(|c| c.borrow_mut().rustc = None);
+}
+
+fn get_cache<T>(f: impl FnOnce(&mut Cache) -> T) -> T {
+    CACHE.with(|c| f(&mut c.borrow_mut()))
+}
+
+fn get_cache_func<T>(inst: &Instance, f: impl FnOnce(&CacheFunction) -> T) -> Option<T> {
+    get_cache(|cache| cache.get_or_insert(inst).map(f))
+}
+
+pub fn get_body<T>(inst: &Instance, f: impl FnOnce(&Body) -> T) -> Option<T> {
+    get_cache_func(inst, |cf| f(&cf.body))
+}
+
+pub fn get_source_code(inst: &Instance) -> Option<SourceCode> {
+    get_cache_func(inst, |cf| cf.src.clone())
+}
+
+pub fn cmp_callees(a: &Instance, b: &Instance) -> Ordering {
+    get_cache(|cache| {
+        cache.get_or_insert(a);
+        cache.get_or_insert(b);
+        let func_a = cache.set.get(&a.def.def_id()).unwrap().as_ref().map(|f| &f.src);
+        let func_b = cache.set.get(&b.def.def_id()).unwrap().as_ref().map(|f| &f.src);
+        func_a.cmp(&func_b)
+    })
+}
+
+struct Cache {
+    set: FxHashMap<DefId, Option<CacheFunction>>,
+    rustc: Option<RustcCxt>,
+    path_prefixes: PathPrefixes,
+}
+
+impl Cache {
+    fn new() -> Self {
+        let (set, rustc) = Default::default();
+        let path_prefixes = PathPrefixes::new();
+        Cache { set, rustc, path_prefixes }
+    }
+
+    fn get_or_insert(&mut self, inst: &Instance) -> Option<&CacheFunction> {
+        self.set
+            .entry(inst.def.def_id())
+            .or_insert_with(|| {
+                let body = inst.body()?;
+                let rustc = self.rustc.as_ref()?;
+                let prefix = self.path_prefixes.prefixes();
+                let src = source_code_with(body.span, rustc.tcx, &rustc.src_map, prefix);
+                Some(CacheFunction { body, src })
+            })
+            .as_ref()
+    }
+}
+
+struct RustcCxt {
+    tcx: TyCtxt<'static>,
+    src_map: Arc<SourceMap>,
+}
+
+struct CacheFunction {
+    body: Body,
+    src: SourceCode,
+}
+
+struct PathPrefixes {
+    pwd: String,
+    sysroot: String,
+}
+
+impl PathPrefixes {
+    fn new() -> Self {
+        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/");
+        PathPrefixes { pwd, sysroot }
+    }
+
+    fn prefixes(&self) -> [&str; 2] {
+        [&*self.pwd, &self.sysroot]
+    }
+}

src/functions/kani/reachability.rs

@@ -17,9 +17,6 @@
 //! Note that this is a copy of `reachability.rs` that uses StableMIR but the public APIs are still
 //! kept with internal APIs.
 
-extern crate rustc_data_structures;
-extern crate rustc_session;
-
 use super::coercion;
 use indexmap::IndexSet;
 use rustc_data_structures::fingerprint::Fingerprint;

src/functions/mod.rs

@@ -1,22 +1,22 @@
 use indexmap::IndexSet;
 use kani::{CallGraph, KANI_TOOL_ATTRS, collect_reachable_items};
 use rustc_middle::ty::TyCtxt;
-use rustc_span::source_map::SourceMap;
 use stable_mir::{
     CrateDef,
     crate_def::Attribute,
-    mir::{
-        Body,
-        mono::{Instance, MonoItem},
-    },
+    mir::mono::{Instance, MonoItem},
 };
 
+mod cache;
+pub use cache::{clear_rustc_ctx, set_rustc_ctx};
+
 mod kani;
+mod utils;
+
 mod serialization;
 pub use serialization::SerFunction;
-mod utils;
 
-pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<SerFunction> {
+pub fn analyze(tcx: TyCtxt) -> Vec<SerFunction> {
     let local_items = stable_mir::all_local_items();
     let cap = local_items.len();
 
@@ -34,8 +34,8 @@ pub fn analyze(tcx: TyCtxt, src_map: &SourceMap) -> Vec<SerFunction> {
     // Filter out non kanitool functions.
     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))
+        .filter_map(|f| Function::new(f, &callgraph, |x| !x.attrs.is_empty()))
+        .map(SerFunction::new)
         .collect();
     // Sort proofs by file path and source code.
     proofs.sort_by(|a, b| a.cmp_by_file_and_func(b));
@@ -51,9 +51,6 @@ pub struct Function {
     /// kanitool's attributs.
     attrs: Vec<Attribute>,
 
-    /// 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>,
@@ -63,8 +60,6 @@ impl Function {
     pub fn new(
         item: &MonoItem,
         callgraph: &CallGraph,
-        tcx: TyCtxt,
-        src_map: &SourceMap,
         filter: impl FnOnce(&Self) -> bool,
     ) -> Option<Self> {
         // Skip non fn items
@@ -73,16 +68,16 @@ impl Function {
         };
 
         // Skip if no body.
-        let body = instance.body()?;
+        cache::get_body(&instance, |_| ())?;
 
         // Only need kanitool attrs: proof, proof_for_contract, contract, ...
         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| utils::cmp_callees(a, b, tcx, src_map));
+        callees.sort_by(cache::cmp_callees);
 
-        let this = Function { instance, attrs, body, callees };
+        let this = Function { instance, attrs, callees };
         filter(&this).then_some(this)
     }
 }

src/functions/serialization.rs

@@ -1,6 +1,4 @@
-use super::utils::{self, SourceCode};
-use rustc_middle::ty::TyCtxt;
-use rustc_span::source_map::SourceMap;
+use super::{cache, utils::SourceCode};
 use rustc_stable_hash::{FromStableHash, SipHasher128Hash, StableHasher, hashers::SipHasher128};
 use serde::Serialize;
 use stable_mir::{CrateDef, mir::mono::Instance};
@@ -12,9 +10,6 @@ pub struct SerFunction {
     hash: String,
     /// DefId in stable_mir.
     def_id: String,
-    /// 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>,
@@ -25,34 +20,31 @@ pub struct SerFunction {
 }
 
 impl SerFunction {
-    pub fn new(fun: super::Function, tcx: TyCtxt, src_map: &SourceMap) -> Self {
+    pub fn new(fun: super::Function) -> Self {
         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();
         // 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();
+        let func = cache::get_source_code(&inst).unwrap_or_default();
+        let callees: Vec<_> = fun.callees.iter().map(Callee::new).collect();
 
         // Hash
         let mut hasher = StableHasher::<SipHasher128>::new();
-        hasher.write_str(&file);
         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);
             c.func.with_hasher(&mut hasher);
         });
         let Hash128(hash) = hasher.finish();
 
-        SerFunction { hash, def_id, file, attrs, func, callees }
+        SerFunction { hash, def_id, 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))
+        self.func.cmp(&other.func)
     }
 }
 
@@ -75,15 +67,13 @@ fn format_def_id(inst: &Instance) -> String {
 #[derive(Debug, Serialize)]
 pub struct Callee {
     def_id: String,
-    file: String,
     func: SourceCode,
 }
 
 impl Callee {
-    fn new(inst: &Instance, tcx: TyCtxt, src_map: &SourceMap) -> Self {
+    fn new(inst: &Instance) -> Self {
         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 }
+        let func = cache::get_source_code(inst).unwrap_or_default();
+        Callee { def_id, func }
     }
 }

src/functions/utils.rs

@@ -3,24 +3,29 @@ 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};
+use std::hash::Hasher;
 
 /// Source code and potential source code before expansion.
 ///
-/// Refer to
-#[derive(Debug, Serialize, PartialEq, Eq, PartialOrd, Ord, Default)]
+/// The field order matters, since this struct implements Ord.
+#[derive(Clone, Debug, Serialize, PartialEq, Eq, PartialOrd, Ord, Default)]
 pub struct SourceCode {
-    // TODO:
-    // file: String,
-    //
+    // A file path where src lies.
+    // The path is stripped with pwd or sysroot prefix.
+    pub 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.
+    ///
+    /// Refer to [#31] to know sepecific cases.
+    ///
+    /// [#31]: https://github.com/os-checker/distributed-verification/issues/31
     pub before_expansion: Option<String>,
 }
 
@@ -47,65 +52,27 @@ fn span_to_source(span: Span, src_map: &SourceMap) -> String {
         .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,
+    path_prefixes: [&str; 2],
 ) -> 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 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)
     });
 
-    let file = inst.def.span().get_filename();
-    for prefix in &*PREFIXES {
-        if let Some(file) = file.strip_prefix(prefix) {
-            return file.to_owned();
+    let mut file = stable_mir_span.get_filename();
+    for prefix in path_prefixes {
+        if let Some(file_stripped) = file.strip_prefix(prefix) {
+            file = file_stripped.to_owned();
+            break;
         }
     }
-    file
+
+    SourceCode { file, src, before_expansion }
 }