Cache codegen of Spans

cprover_bindings/src/goto_program/location.rs

@@ -90,6 +90,23 @@ impl Location {
             Location::PropertyUnknownLocation { .. } => "<none>".to_string(),
         }
     }
+
+    /// Tries to set the `function` field of a [Location::Loc], returning `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
+        }
+    }
 }
 
 /// Constructors

kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs

@@ -37,6 +37,18 @@ impl GotocCtx<'_> {
     }
 
     pub fn codegen_span_stable(&self, sp: SpanStable) -> Location {
+        // First query the cache to see if we've already done codegen for this span.
+        if let Some(cached_loc) = self.span_cache.borrow_mut().get(&sp) {
+            let mut new_loc = *cached_loc;
+
+            // Recalculate the `current_fn` since it could be different than when we cached.
+            new_loc
+                .try_set_function(self.current_fn.as_ref().map(|x| x.readable_name().to_string()))
+                .unwrap();
+
+            return new_loc;
+        }
+
         // 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"),
@@ -65,16 +77,26 @@ impl GotocCtx<'_> {
                 .collect::<Vec<_>>()
                 .leak() // This is to preserve `Location` being Copy, but could blow up the memory utilization of compiler. 
         };
+
         let loc = sp.get_lines();
-        Location::new(
+        let new_loc = Location::new(
             sp.get_filename().to_string(),
             self.current_fn.as_ref().map(|x| x.readable_name().to_string()),
             loc.start_line,
             Some(loc.start_col),
             loc.end_line,
             Some(loc.end_col),
             pragmas,
-        )
+        );
+
+        // Insert codegened Location into the cache and sanity check it doesn't already exist.
+        let existing = self.span_cache.borrow_mut().insert(sp, new_loc);
+        debug_assert!(
+            existing.is_none(),
+            "if there was already an entry for this in the cache, we should've used that!"
+        );
+
+        new_loc
     }
 
     pub fn codegen_caller_span_stable(&self, sp: SpanStable) -> Location {

kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs

@@ -41,6 +41,7 @@ use rustc_public::ty::Allocation;
 use rustc_span::Span;
 use rustc_span::source_map::respan;
 use rustc_target::callconv::FnAbi;
+use std::cell::RefCell;
 use std::collections::{BTreeMap, HashSet};
 use std::fmt::Debug;
 
@@ -77,6 +78,8 @@ pub struct GotocCtx<'tcx> {
     pub global_var_count: u64,
     /// map a global allocation to a name in the symbol table
     pub alloc_map: FxHashMap<Allocation, String>,
+    /// a cache of [Span](rustc_public::Span)s and their codegened [Location]s
+    pub span_cache: RefCell<FxHashMap<rustc_public::ty::Span, Location>>,
     /// map (trait, method) pairs to possible implementations
     pub vtable_ctx: VtableCtx,
     pub current_fn: Option<CurrentFnCtx<'tcx>>,
@@ -119,6 +122,7 @@ impl<'tcx> GotocCtx<'tcx> {
             full_crate_name: full_crate_name(tcx),
             global_var_count: 0,
             alloc_map: FxHashMap::default(),
+            span_cache: RefCell::new(FxHashMap::default()),
             vtable_ctx: VtableCtx::new(emit_vtable_restrictions),
             current_fn: None,
             type_map: FxHashMap::default(),