diff --git a/cprover_bindings/src/goto_program/location.rs b/cprover_bindings/src/goto_program/location.rs index 9a2a046f6387..79bf96df4bc6 100644 --- a/cprover_bindings/src/goto_program/location.rs +++ b/cprover_bindings/src/goto_program/location.rs @@ -90,6 +90,23 @@ impl Location { Location::PropertyUnknownLocation { .. } => "".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>, + ) -> Option<()> { + if let Location::Loc { function, .. } = self { + if let Some(new_function) = new_function { + *function = Some(new_function.into()); + } + + Some(()) + } else { + None + } + } } /// Constructors diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs index c438a6eace2a..49246d25ed38 100644 --- a/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs +++ b/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,8 +77,9 @@ impl GotocCtx<'_> { .collect::>() .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, @@ -74,7 +87,16 @@ impl GotocCtx<'_> { 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 { diff --git a/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs b/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs index e5807a774041..ab8f08e589f0 100644 --- a/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs +++ b/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, + /// a cache of [Span](rustc_public::Span)s and their codegened [Location]s + pub span_cache: RefCell>, /// map (trait, method) pairs to possible implementations pub vtable_ctx: VtableCtx, pub current_fn: Option>, @@ -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(),