chore: move file field to SourceCode

Author: zjp-CN

src/functions/cache.rs

@@ -10,10 +10,10 @@ use stable_mir::{
     CrateDef, DefId,
     mir::{Body, mono::Instance},
 };
-use std::{cell::RefCell, sync::Arc};
+use std::{cell::RefCell, cmp::Ordering, sync::Arc};
 
 thread_local! {
-    static CACHE: RefCell<Cache> = RefCell::new(Cache::default());
+    static CACHE: RefCell<Cache> = RefCell::new(Cache::new());
 }
 
 pub fn set_rustc_ctx(tcx: TyCtxt) {
@@ -29,8 +29,12 @@ 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> {
-    CACHE.with(|c| c.borrow_mut().get_or_insert(inst).map(f))
+    get_cache(|cache| cache.get_or_insert(inst).map(f))
 }
 
 pub fn get_body<T>(inst: &Instance, f: impl FnOnce(&Body) -> T) -> Option<T> {
@@ -41,20 +45,37 @@ pub fn get_source_code(inst: &Instance) -> Option<SourceCode> {
     get_cache_func(inst, |cf| cf.src.clone())
 }
 
-#[derive(Default)]
+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 {
-    rustc: Option<RustcCxt>,
     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 src = source_code_with(body.span, rustc.tcx, &rustc.src_map);
+                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()
@@ -70,3 +91,24 @@ 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/mod.rs

@@ -75,7 +75,7 @@ impl Function {
 
         let mut callees = IndexSet::new();
         callgraph.recursive_callees(item, &mut callees);
-        callees.sort_by(utils::cmp_callees);
+        callees.sort_by(cache::cmp_callees);
 
         let this = Function { instance, attrs, callees };
         filter(&this).then_some(this)

src/functions/serialization.rs

@@ -1,7 +1,4 @@
-use super::{
-    cache,
-    utils::{self, SourceCode},
-};
+use super::{cache, utils::SourceCode};
 use rustc_stable_hash::{FromStableHash, SipHasher128Hash, StableHasher, hashers::SipHasher128};
 use serde::Serialize;
 use stable_mir::{CrateDef, mir::mono::Instance};
@@ -13,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>,
@@ -29,31 +23,28 @@ impl SerFunction {
     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 = cache::get_source_code(&inst).unwrap_or_default();
-        let callees: Vec<_> = fun.callees.iter().map(|x| Callee::new(x)).collect();
+        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)
     }
 }
 
@@ -76,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) -> Self {
         let def_id = format_def_id(inst);
-        let file = utils::file_path(inst);
         let func = cache::get_source_code(inst).unwrap_or_default();
-        Callee { def_id, file, func }
+        Callee { def_id, func }
     }
 }

src/functions/utils.rs

@@ -3,26 +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 super::cache;
+use std::hash::Hasher;
 
 /// Source code and potential source code before expansion.
 ///
-/// Refer to
+/// 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>,
 }
 
@@ -54,46 +57,22 @@ 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);
     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 }
-}
-
-pub fn cmp_callees(a: &Instance, b: &Instance) -> 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 = cache::get_source_code(a);
-    let body_b = cache::get_source_code(b);
-    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();
+    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 }
 }

src/lib.rs

@@ -6,9 +6,6 @@ pub struct SerFunction {
     pub hash: String,
     /// 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>,
@@ -21,15 +18,15 @@ pub struct SerFunction {
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Callee {
     pub def_id: String,
-    pub file: String,
     pub func: SourceCode,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct SourceCode {
-    // TODO:
-    // file: String,
-    //
+    // A file path where src lies.
+    // The path is stripped with pwd or sysroot prefix.
+    file: String,
+
     /// Source that a stable_mir span points to.
     pub src: String,
     /// Is the stable_mir span from a macro expansion?