[WIP] Introduce per-module cache files to reduce FS overhead

Author: bjorn3

src/driver/aot.rs

@@ -328,7 +328,8 @@ fn make_module(sess: &Session, name: String) -> UnwindModule<ObjectModule> {
     let isa = crate::build_isa(sess, false);
 
     let mut builder =
-        ObjectBuilder::new(isa, name + ".o", cranelift_module::default_libcall_names()).unwrap();
+        ObjectBuilder::new(isa, name.clone() + ".o", cranelift_module::default_libcall_names())
+            .unwrap();
 
     // Disable function sections by default on MSVC as it causes significant slowdowns with link.exe.
     // Maybe link.exe has exponential behavior when there are many sections with the same name? Also
@@ -340,7 +341,7 @@ fn make_module(sess: &Session, name: String) -> UnwindModule<ObjectModule> {
         sess.opts.unstable_opts.function_sections.unwrap_or(default_function_sections),
     );
 
-    UnwindModule::new(ObjectModule::new(builder), true)
+    UnwindModule::new(ObjectModule::new(builder), &name, true)
 }
 
 fn emit_cgu(

src/driver/jit.rs

@@ -24,7 +24,7 @@ fn create_jit_module(tcx: TyCtxt<'_>) -> (UnwindModule<JITModule>, CodegenCx) {
     let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
     crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
     jit_builder.symbol_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info));
-    let mut jit_module = UnwindModule::new(JITModule::new(jit_builder), false);
+    let mut jit_module = UnwindModule::new(JITModule::new(jit_builder), "jit", false);
 
     let cx = crate::CodegenCx::new(tcx, jit_module.isa(), false, sym::dummy_cgu_name);
 

src/unwind_module.rs

@@ -1,3 +1,9 @@
+use std::borrow::Cow;
+use std::cell::Cell;
+use std::fs::File;
+use std::io::{BufWriter, Write};
+use std::path::PathBuf;
+
 use cranelift_codegen::Context;
 use cranelift_codegen::control::ControlPlane;
 use cranelift_codegen::incremental_cache::CacheKvStore;
@@ -8,19 +14,30 @@ use cranelift_module::{
     ModuleReloc, ModuleResult,
 };
 use cranelift_object::{ObjectModule, ObjectProduct};
+use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::owned_slice::{OwnedSlice, slice_owned};
 
 use crate::UnwindContext;
 
 /// A wrapper around a [Module] which adds any defined function to the [UnwindContext].
 pub(crate) struct UnwindModule<T> {
     pub(crate) module: T,
     unwind_context: UnwindContext,
+    cache: Option<Box<dyn CacheKvStore + Send>>,
 }
 
 impl<T: Module> UnwindModule<T> {
-    pub(crate) fn new(mut module: T, pic_eh_frame: bool) -> Self {
+    pub(crate) fn new(mut module: T, cgu_name: &str, pic_eh_frame: bool) -> Self {
         let unwind_context = UnwindContext::new(&mut module, pic_eh_frame);
-        UnwindModule { module, unwind_context }
+        let cache = match std::env::var("CG_CLIF_FUNCTION_CACHE").as_deref() {
+            Err(_) => None,
+            Ok("naive") => Some(Box::new(FileCache) as Box<dyn CacheKvStore + Send>),
+            Ok("module") => {
+                Some(Box::new(ModuleCache::new(cgu_name)) as Box<dyn CacheKvStore + Send>)
+            }
+            _ => panic!(),
+        };
+        UnwindModule { module, unwind_context, cache }
     }
 }
 
@@ -91,13 +108,13 @@ impl<T: Module> Module for UnwindModule<T> {
         ctx: &mut Context,
         ctrl_plane: &mut ControlPlane,
     ) -> ModuleResult<()> {
-        if std::env::var("CG_CLIF_FUNCTION_CACHE").as_deref() == Ok("naive") {
+        if let Some(cache) = &mut self.cache {
             if ctx.func.layout.blocks().nth(1).is_none()
                 || ctx.func.layout.blocks().nth(2).is_none()
             {
                 ctx.compile(self.module.isa(), ctrl_plane)?;
             } else {
-                ctx.compile_with_cache(self.module.isa(), &mut Cache, ctrl_plane)?;
+                ctx.compile_with_cache(self.module.isa(), &mut **cache, ctrl_plane)?;
             }
         } else {
             ctx.compile(self.module.isa(), ctrl_plane)?;
@@ -132,18 +149,18 @@ impl<T: Module> Module for UnwindModule<T> {
     }
 }
 
-struct Cache;
+struct FileCache;
 
-impl Cache {
+impl FileCache {
     fn file_for_key(&self, key: &[u8]) -> String {
         let mut path = key.iter().map(|b| format!("{:02x}", b)).collect::<String>();
         path.push_str(".clif_cache");
         "/home/bjorn/Projects/cg_clif/cache/".to_owned() + &path
     }
 }
 
-impl CacheKvStore for Cache {
-    fn get(&self, key: &[u8]) -> Option<std::borrow::Cow<'_, [u8]>> {
+impl CacheKvStore for FileCache {
+    fn get(&self, key: &[u8]) -> Option<Cow<'_, [u8]>> {
         let path = self.file_for_key(key);
         if std::fs::exists(&path).unwrap() {
             Some(std::fs::read(path).unwrap().into())
@@ -157,3 +174,64 @@ impl CacheKvStore for Cache {
         std::fs::write(path, val).unwrap();
     }
 }
+
+struct ModuleCache {
+    file: PathBuf,
+    entries: FxHashMap<[u8; 32], (OwnedSlice, Cell<bool>)>,
+}
+
+impl ModuleCache {
+    fn new(name: &str) -> Self {
+        let file = PathBuf::from("/home/bjorn/Projects/cg_clif/cache/".to_owned() + &name);
+
+        if !file.exists() {
+            return ModuleCache { file, entries: FxHashMap::default() };
+        }
+
+        let data = std::fs::read(&file).unwrap();
+        let mut data = slice_owned(data, |data| &data);
+        let mut entries = FxHashMap::default();
+        while !data.is_empty() {
+            let key: [u8; 32] = data[..32].try_into().unwrap();
+            let size = u32::from_le_bytes(data[32..36].try_into().unwrap());
+            entries.insert(
+                key,
+                (data.clone().slice(|data| &data[20..size as usize + 36]), Cell::new(false)),
+            );
+            data = data.slice(|data| &data[size as usize + 36..]);
+        }
+
+        ModuleCache { file, entries }
+    }
+}
+
+impl Drop for ModuleCache {
+    fn drop(&mut self) {
+        let mut buf_writer = BufWriter::new(File::create(&self.file).unwrap());
+        for (key, (data, accessed)) in &self.entries {
+            if !accessed.get() {
+                continue;
+            }
+            buf_writer.write_all(key).unwrap();
+            buf_writer.write_all(&u32::to_le_bytes(data.len() as u32)).unwrap();
+            buf_writer.write_all(data).unwrap();
+        }
+        buf_writer.into_inner().unwrap();
+    }
+}
+
+impl CacheKvStore for ModuleCache {
+    fn get(&self, key: &[u8]) -> Option<Cow<'_, [u8]>> {
+        if let Some((data, accessed)) = self.entries.get::<[u8; 32]>(key.try_into().unwrap()) {
+            accessed.set(true);
+            Some(Cow::Borrowed(&*data))
+        } else {
+            None
+        }
+    }
+
+    fn insert(&mut self, key: &[u8], val: Vec<u8>) {
+        self.entries
+            .insert(key.try_into().unwrap(), (slice_owned(val, |val| &val), Cell::new(true)));
+    }
+}