cache implementation from branch

AlexanderPortland · AlexanderPortland · commit 7db6a0cef17d · 2025-08-19T17:38:36.000-07:00
diff --git a/cprover_bindings/src/goto_program/location.rs b/cprover_bindings/src/goto_program/location.rs
@@ -124,6 +124,23 @@ impl Location {
         }
     }
 
+    /// Tries to set the `function` field of a [Location::Loc], but returns `None` if the location
+    /// is of a different variant and the field couldn't be set.
+    pub fn try_set_function(
+        &mut self,
+        new_function: Option<impl Into<InternedString>>,
+    ) -> Option<()> {
+        if let Location::Loc { function, .. } = self {
+            if let Some(new_function) = new_function {
+                *function = Some(new_function.into());
+            }
+
+            Some(())
+        } else {
+            None
+        }
+    }
+
     /// Create a Property type Location
     pub fn property_location<T, U>(
         file: Option<U>,
diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/cache.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/cache.rs
@@ -1,18 +1,14 @@
-use std::hash::Hash;
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
 
-use std::collections::hash_map::Entry as HashMapEntry;
+use std::{cell::RefCell, hash::Hash};
 
-use cbmc::goto_program::{Location, Type};
 use fxhash::FxHashMap;
+use std::any::type_name;
 
-use crate::codegen_cprover_gotoc::context::SpanWrapper;
-
-// #[derive(Default)]
-// pub struct CodegenCache {
-
-//     types: FxHashMap<rustc_public::ty::Ty, Type>,
-//     spans: FxHashMap<SpanWrapper, Location>,
-// }
+thread_local! {
+    pub static CACHE: RefCell<CodegenCache> = RefCell::new(Default::default());
+}
 
 type HashImpl<K, V> = FxHashMap<K, V>;
 
@@ -21,109 +17,104 @@ where
     Self: Sized,
 {
     type Key: Hash + Eq;
-    fn get_individual_cache(cache: &mut CodegenCache) -> &mut HashImpl<Self::Key, Self>;
+    fn get_individual_cache(cache: &CodegenCache) -> &HashImpl<Self::Key, Self>;
+    fn get_individual_cache_mut(cache: &mut CodegenCache) -> &mut HashImpl<Self::Key, Self>;
 }
 
-// impl CodegenCacheEl for Type {
-//     type Key = rustc_public::ty::Ty;
-//     fn get_individual_cache(cache: &mut CodegenCache) -> &mut HashImpl<Self::Key, Self> {
-//         &mut cache.types
-//     }
-// }
-
-// impl CodegenCacheEl for Location {
-//     type Key = SpanWrapper;
-//     fn get_individual_cache(cache: &mut CodegenCache) -> &mut HashImpl<Self::Key, Self> {
-//         &mut cache.spans
+// impl Drop for CodegenCache {
+//     fn drop(&mut self) {
+//         self.calc_size();
 //     }
 // }
 
 macro_rules! generate_cache {
-    ($name:tt -- $([$field_name:tt] $key:path => $el:path),+) => {
+    ($name:tt -- $($(@$global:tt)? [$field_name:tt] $key:path => $el:path),+) => {
         #[derive(Default)]
         pub struct $name {
             $($field_name: HashImpl<$key, $el>,)*
         }
 
+        impl $name {
+            #[allow(dead_code)]
+            fn calc_size(&self) {
+                let mut total = 0;
+                $(
+                    let key_size = size_of::<$key>();
+                    let el_size = size_of::<$el>();
+                    let cap = self.$field_name.capacity();
+                    let field_total = (key_size + el_size) * cap;
+                    println!("for {:?} map -- key: {key_size}, el: {el_size}, cap: {cap} => {field_total}", type_name::<$el>());
+                    total += field_total;
+                )*
+                println!("{total} in total!");
+            }
+        }
+
+        pub fn clear_codegen_cache() {
+            CACHE.with_borrow_mut(|cache|{
+                $(
+                    clear_cache!($($global)? cache, $field_name);
+                )*
+            })
+        }
+
         $(
             impl CodegenCacheEl for $el {
                 type Key = $key;
-                fn get_individual_cache(cache: &mut CodegenCache) -> &mut HashImpl<Self::Key, Self> {
+                fn get_individual_cache(cache: &CodegenCache) -> &HashImpl<Self::Key, Self> {
+                    &cache.$field_name
+                }
+                fn get_individual_cache_mut(cache: &mut CodegenCache) -> &mut HashImpl<Self::Key, Self> {
                     &mut cache.$field_name
                 }
             }
         )*
     };
 }
 
+// TODO: add more granularity here...
+macro_rules! clear_cache {
+    ( $cache:tt, $field_name:tt) => {
+        $cache.$field_name.clear();
+    };
+    (global $cache:tt, $field_name:tt) => {
+        /* global field, don't clear cache */
+    };
+}
+
 generate_cache!(CodegenCache --
-    [types] rustc_public::ty::Ty => cbmc::goto_program::Type,
-    [spans] SpanWrapper => cbmc::goto_program::Location
+            [types] rustc_public::ty::Ty   => cbmc::goto_program::Type,
+    @global [spans] rustc_public::ty::Span => cbmc::goto_program::Location
 );
 
-// you give me a key i tell you if it's there
-
-// if it's there, you may want to modify before taking it
+pub struct FinalEntry<T: CodegenCacheEl>(Option<T>, T::Key);
 
-// if it's not, you want to insert it by running some code
-
-// they you want a ref to the newly inserted
-
-impl CodegenCache {
-    pub fn entry<E: CodegenCacheEl + Clone>(&mut self, key: E::Key) -> Entry<'_, E::Key, E> {
-        Entry(E::get_individual_cache(self).entry(key))
-    }
-
-    pub fn entry_ref<E: CodegenCacheEl>(&mut self, key: E::Key) -> EntryRef<'_, E::Key, E> {
-        EntryRef(E::get_individual_cache(self).entry(key))
-    }
+pub fn cache_entry<E: CodegenCacheEl + Clone>(key: E::Key) -> FinalEntry<E> {
+    let found_value = CACHE.with_borrow(|cache| E::get_individual_cache(cache).get(&key).cloned());
+    FinalEntry(found_value, key)
 }
 
-// an entry that we inevitably want to clone
-struct Entry<'a, K, V: Clone>(HashMapEntry<'a, K, V>);
-struct EntryRef<'a, K, V>(HashMapEntry<'a, K, V>);
-
-// an [Entry] that will have a function applied
-struct TweakedEntry<'a, K, V: Clone>(Option<V>, HashMapEntry<'a, K, V>);
-
-impl<'a, K, V: Clone> EntryRef<'a, K, V> {
-    pub fn or_insert_with<F: FnOnce() -> V>(self, f: F) -> &'a V {
-        match self.0 {
-            HashMapEntry::Occupied(existing) => existing.into_mut(),
-            HashMapEntry::Vacant(vacant) => vacant.insert(f()),
+impl<E: CodegenCacheEl> FinalEntry<E> {
+    pub fn tweak<F: FnOnce(&mut E)>(mut self, f: F) -> FinalEntry<E> {
+        if let Some(found_val) = &mut self.0 {
+            f(found_val)
         }
-    }
-}
-
-impl<'a, K, V: Clone> Entry<'a, K, V> {
-    pub fn tweak_from_cache<F: FnOnce(&mut V)>(self, f: F) -> TweakedEntry<'a, K, V> {
-        TweakedEntry(
-            match &self.0 {
-                HashMapEntry::Occupied(o) => {
-                    let mut new_val = o.get().clone();
-                    f(&mut new_val);
-                    Some(new_val)
-                }
-                HashMapEntry::Vacant(_) => None,
-            },
-            self.0,
-        )
-    }
 
-    pub fn or_insert_with<F: FnOnce() -> V>(self, f: F) -> V {
-        match self.0 {
-            HashMapEntry::Occupied(existing) => existing.get().clone(),
-            HashMapEntry::Vacant(vacant) => vacant.insert(f()).clone(),
-        }
+        self
     }
 }
 
-impl<'a, K, V: Clone> TweakedEntry<'a, K, V> {
-    pub fn or_insert_with<F: FnOnce() -> V>(self, f: F) -> V {
-        match (self.0, self.1) {
-            (Some(tweaked), HashMapEntry::Occupied(_)) => tweaked,
-            (None, HashMapEntry::Vacant(vacant)) => vacant.insert(f()).clone(),
-            _ => unreachable!(),
+impl<E: CodegenCacheEl + Clone> FinalEntry<E> {
+    pub fn or_insert_with<F: FnOnce() -> E>(self, f: F) -> E {
+        match self.0 {
+            Some(cached) => cached,
+            None => {
+                let calculated = f();
+                CACHE.with_borrow_mut(|cache| {
+                    E::get_individual_cache_mut(cache).insert(self.1, calculated.clone())
+                });
+                calculated
+            }
         }
     }
 }
diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/mod.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/mod.rs
@@ -6,6 +6,7 @@
 
 mod assert;
 mod block;
+pub(crate) mod cache;
 mod foreign_function;
 mod function;
 mod intrinsic;
diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs
@@ -4,6 +4,7 @@
 //! MIR Span related functions
 
 use crate::codegen_cprover_gotoc::GotocCtx;
+use crate::codegen_cprover_gotoc::codegen::cache::cache_entry;
 use cbmc::goto_program::Location;
 use lazy_static::lazy_static;
 use rustc_hir::Attribute;
@@ -37,6 +38,17 @@ impl GotocCtx<'_> {
     }
 
     pub fn codegen_span_stable(&self, sp: SpanStable) -> Location {
+        cache_entry::<Location>(sp)
+            .tweak(|res| {
+                res.try_set_function(
+                    self.current_fn.as_ref().map(|x| x.readable_name().to_string()),
+                )
+                .unwrap();
+            })
+            .or_insert_with(|| self.codegen_span_stable_inner(sp))
+    }
+
+    pub fn codegen_span_stable_inner(&self, sp: SpanStable) -> Location {
         // Attribute to mark functions as where automatic pointer checks should not be generated.
         let should_skip_ptr_checks_attr = vec![
             rustc_span::symbol::Symbol::intern("kanitool"),
@@ -66,6 +78,7 @@ impl GotocCtx<'_> {
                 .leak() // This is to preserve `Location` being Copy, but could blow up the memory utilization of compiler. 
         };
         let loc = sp.get_lines();
+
         Location::new(
             sp.get_filename().to_string(),
             self.current_fn.as_ref().map(|x| x.readable_name().to_string()),
diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/ty_stable.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/ty_stable.rs
@@ -7,6 +7,7 @@
 //! `typ.rs`.
 
 use crate::codegen_cprover_gotoc::GotocCtx;
+use crate::codegen_cprover_gotoc::codegen::cache::cache_entry;
 use crate::kani_middle::abi::LayoutOf;
 use cbmc::goto_program::Type;
 use rustc_middle::ty::layout::{LayoutOf as _, TyAndLayout};
@@ -20,8 +21,15 @@ impl<'tcx> GotocCtx<'tcx> {
         place.ty(self.current_fn().locals()).unwrap()
     }
 
+    // pub fn codegen_ty_stable(&mut self, ty: Ty) -> Type {
+    //     cache_entry::<Type>(ty).or_insert_with(||{
+    //         self.codegen_ty_stable_inner(ty)
+    //     })
+    // }
+
     pub fn codegen_ty_stable(&mut self, ty: Ty) -> Type {
-        self.codegen_ty(rustc_internal::internal(self.tcx, ty))
+        cache_entry::<Type>(ty)
+            .or_insert_with(|| self.codegen_ty(rustc_internal::internal(self.tcx, ty)))
     }
 
     pub fn codegen_ty_ref_stable(&mut self, ty: Ty) -> Type {
diff --git a/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs b/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs
@@ -16,9 +16,9 @@
 //! this structure as input.
 use super::current_fn::CurrentFnCtx;
 use super::vtable_ctx::VtableCtx;
-use crate::codegen_cprover_gotoc::UnsupportedConstructs;
 use crate::codegen_cprover_gotoc::overrides::{GotocHooks, fn_hooks};
 use crate::codegen_cprover_gotoc::utils::full_crate_name;
+use crate::codegen_cprover_gotoc::{UnsupportedConstructs, clear_codegen_cache};
 use crate::kani_middle::transform::BodyTransformation;
 use crate::kani_queries::QueryDb;
 use cbmc::goto_program::{
@@ -108,6 +108,7 @@ impl<'tcx> GotocCtx<'tcx> {
         machine_model: &MachineModel,
         transformer: BodyTransformation,
     ) -> GotocCtx<'tcx> {
+        clear_codegen_cache();
         let fhks = fn_hooks();
         let symbol_table = SymbolTable::new(machine_model.clone());
         let emit_vtable_restrictions = queries.args().emit_vtable_restrictions;
diff --git a/kani-compiler/src/codegen_cprover_gotoc/mod.rs b/kani-compiler/src/codegen_cprover_gotoc/mod.rs
@@ -6,6 +6,7 @@ mod context;
 mod overrides;
 mod utils;
 
+pub use codegen::cache::clear_codegen_cache;
 pub use compiler_interface::{GotocCodegenBackend, UnsupportedConstructs};
 pub use context::GotocCtx;
 pub use context::VtableCtx;
