From d7d0e6f8188636b679ceed67c848126cfceeec0c Mon Sep 17 00:00:00 2001 From: AlexanderPortland Date: Wed, 6 Aug 2025 12:50:31 -0700 Subject: [PATCH 1/2] cache codegen of Spans (using transmute hack to hash them) --- cprover_bindings/src/goto_program/location.rs | 17 ++++++++++++ .../src/codegen_cprover_gotoc/codegen/span.rs | 27 +++++++++++++++++-- .../codegen_cprover_gotoc/context/goto_ctx.rs | 17 ++++++++++++ .../src/codegen_cprover_gotoc/context/mod.rs | 2 +- 4 files changed, 60 insertions(+), 3 deletions(-) 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..e352477ab3cd 100644 --- 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::context::SpanWrapper; use cbmc::goto_program::Location; use lazy_static::lazy_static; use rustc_hir::Attribute; @@ -37,6 +38,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(&SpanWrapper(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 +78,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 +88,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(SpanWrapper(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..47ea1f916a93 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; @@ -62,6 +63,19 @@ pub struct MinimalGotocCtx { pub has_loop_contracts: bool, } +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub struct OurSpan(usize); + +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct SpanWrapper(pub rustc_public::ty::Span); + +impl std::hash::Hash for SpanWrapper { + fn hash(&self, state: &mut H) { + let our_span: OurSpan = unsafe { std::mem::transmute(self.0) }; + our_span.0.hash(state); + } +} + pub struct GotocCtx<'tcx> { /// the typing context pub tcx: TyCtxt<'tcx>, @@ -77,6 +91,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 +135,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(), diff --git a/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs b/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs index 75052039615d..0796480a291a 100644 --- a/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs +++ b/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs @@ -10,5 +10,5 @@ mod current_fn; mod goto_ctx; mod vtable_ctx; -pub use goto_ctx::{GotocCtx, MinimalGotocCtx}; +pub use goto_ctx::{GotocCtx, MinimalGotocCtx, SpanWrapper}; pub use vtable_ctx::VtableCtx; From ca82871afac585a8e05ed3eafeab433de2bb40b0 Mon Sep 17 00:00:00 2001 From: AlexanderPortland Date: Wed, 6 Aug 2025 13:03:09 -0700 Subject: [PATCH 2/2] fix transmute hack (for when `rustc_public::Span` hopefully becomes `Hash`) --- .../src/codegen_cprover_gotoc/codegen/span.rs | 5 ++--- .../src/codegen_cprover_gotoc/context/goto_ctx.rs | 15 +-------------- .../src/codegen_cprover_gotoc/context/mod.rs | 2 +- 3 files changed, 4 insertions(+), 18 deletions(-) diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs index e352477ab3cd..49246d25ed38 100644 --- a/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs +++ b/kani-compiler/src/codegen_cprover_gotoc/codegen/span.rs @@ -4,7 +4,6 @@ //! MIR Span related functions use crate::codegen_cprover_gotoc::GotocCtx; -use crate::codegen_cprover_gotoc::context::SpanWrapper; use cbmc::goto_program::Location; use lazy_static::lazy_static; use rustc_hir::Attribute; @@ -39,7 +38,7 @@ 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(&SpanWrapper(sp)) { + 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. @@ -91,7 +90,7 @@ impl GotocCtx<'_> { ); // Insert codegened Location into the cache and sanity check it doesn't already exist. - let existing = self.span_cache.borrow_mut().insert(SpanWrapper(sp), new_loc); + 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!" 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 47ea1f916a93..ab8f08e589f0 100644 --- a/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs +++ b/kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs @@ -63,19 +63,6 @@ pub struct MinimalGotocCtx { pub has_loop_contracts: bool, } -#[derive(Clone, Copy, PartialEq, Eq, Hash)] -pub struct OurSpan(usize); - -#[derive(Clone, Copy, PartialEq, Eq)] -pub struct SpanWrapper(pub rustc_public::ty::Span); - -impl std::hash::Hash for SpanWrapper { - fn hash(&self, state: &mut H) { - let our_span: OurSpan = unsafe { std::mem::transmute(self.0) }; - our_span.0.hash(state); - } -} - pub struct GotocCtx<'tcx> { /// the typing context pub tcx: TyCtxt<'tcx>, @@ -92,7 +79,7 @@ pub struct GotocCtx<'tcx> { /// 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>, + pub span_cache: RefCell>, /// map (trait, method) pairs to possible implementations pub vtable_ctx: VtableCtx, pub current_fn: Option>, diff --git a/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs b/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs index 0796480a291a..75052039615d 100644 --- a/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs +++ b/kani-compiler/src/codegen_cprover_gotoc/context/mod.rs @@ -10,5 +10,5 @@ mod current_fn; mod goto_ctx; mod vtable_ctx; -pub use goto_ctx::{GotocCtx, MinimalGotocCtx, SpanWrapper}; +pub use goto_ctx::{GotocCtx, MinimalGotocCtx}; pub use vtable_ctx::VtableCtx;