Switch to thread local

zhassan-aws · zhassan-aws · commit 677422433d5b · 2026-01-14T14:53:13.000-08:00
diff --git a/kani-compiler/src/codegen_aeneas_llbc/compiler_interface.rs b/kani-compiler/src/codegen_aeneas_llbc/compiler_interface.rs
@@ -11,7 +11,7 @@ use crate::kani_middle::codegen_units::{CodegenUnit, CodegenUnits};
 use crate::kani_middle::provide;
 use crate::kani_middle::reachability::{collect_reachable_items, filter_crate_items};
 use crate::kani_middle::transform::{BodyTransformation, GlobalPasses};
-use crate::kani_queries::QueryDb;
+use crate::kani_queries::QUERY_DB;
 use charon_lib::ast::{AnyTransId, TranslatedCrate, meta::ItemOpacity::*, meta::Span};
 use charon_lib::errors::ErrorCtx;
 use charon_lib::name_matcher::NamePattern;
@@ -43,22 +43,15 @@ use rustc_target::spec::Arch;
 use std::any::Any;
 use std::fs::File;
 use std::path::Path;
-use std::sync::{Arc, Mutex};
 use std::time::Instant;
 use tracing::{debug, info, trace};
 
 #[derive(Clone)]
-pub struct LlbcCodegenBackend {
-    /// The query is shared with `KaniCompiler` and it is initialized as part of `rustc`
-    /// initialization, which may happen after this object is created.
-    /// Since we don't have any guarantees on when the compiler creates the Backend object, neither
-    /// in which thread it will be used, we prefer to explicitly synchronize any query access.
-    queries: Arc<Mutex<QueryDb>>,
-}
+pub struct LlbcCodegenBackend {}
 
 impl LlbcCodegenBackend {
-    pub fn new(queries: Arc<Mutex<QueryDb>>) -> Self {
-        LlbcCodegenBackend { queries }
+    pub fn new() -> Self {
+        LlbcCodegenBackend {}
     }
 
     /// Generate code that is reachable from the given starting points.
@@ -91,16 +84,18 @@ impl LlbcCodegenBackend {
             .collect();
 
         // Apply all transformation passes, including global passes.
-        let mut global_passes = GlobalPasses::new(&self.queries.lock().unwrap(), tcx);
-        global_passes.run_global_passes(
-            &mut transformer,
-            tcx,
-            starting_items,
-            instances,
-            call_graph,
-        );
+        QUERY_DB.with(|db| {
+            let mut global_passes = GlobalPasses::new(&db.borrow(), tcx);
+            global_passes.run_global_passes(
+                &mut transformer,
+                tcx,
+                starting_items,
+                instances,
+                call_graph,
+            );
+        });
 
-        let queries = self.queries.lock().unwrap().clone();
+        let queries = QUERY_DB.with(|db| db.borrow().clone());
         check_reachable_items(tcx, &queries, &items);
 
         // Follow rustc naming convention (cx is abbrev for context).
@@ -192,7 +187,7 @@ impl LlbcCodegenBackend {
 
 impl CodegenBackend for LlbcCodegenBackend {
     fn provide(&self, providers: &mut Providers) {
-        provide::provide(providers, &self.queries.lock().unwrap());
+        QUERY_DB.with(|db| provide::provide(providers, &db.borrow()));
     }
 
     fn print_version(&self) {
@@ -211,7 +206,7 @@ impl CodegenBackend for LlbcCodegenBackend {
     fn codegen_crate(&self, tcx: TyCtxt) -> Box<dyn Any> {
         let ret_val = rustc_internal::run(tcx, || {
             // Queries shouldn't change today once codegen starts.
-            let queries = self.queries.lock().unwrap().clone();
+            let queries = QUERY_DB.with(|db| db.borrow().clone());
 
             // Codegen all items that need to be processed according to the selected reachability mode:
             //
diff --git a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
@@ -14,7 +14,7 @@ use crate::kani_middle::codegen_units::{CodegenUnit, CodegenUnits};
 use crate::kani_middle::provide;
 use crate::kani_middle::reachability::{collect_reachable_items, filter_crate_items};
 use crate::kani_middle::transform::{BodyTransformation, GlobalPasses};
-use crate::kani_queries::QueryDb;
+use crate::kani_queries::QUERY_DB;
 use cbmc::goto_program::Location;
 use cbmc::{InternedString, MachineModel};
 use cbmc::{RoundingMode, WithInterner};
@@ -52,7 +52,6 @@ use std::fs::File;
 use std::io::BufWriter;
 use std::num::NonZero;
 use std::path::Path;
-use std::sync::{Arc, Mutex};
 use std::thread::available_parallelism;
 use std::time::Instant;
 use tracing::{debug, info};
@@ -66,17 +65,11 @@ const MAX_SENSIBLE_FILE_EXPORT_THREADS: usize = 4;
 
 pub type UnsupportedConstructs = FxHashMap<InternedString, Vec<Location>>;
 
-pub struct GotocCodegenBackend {
-    /// The query is shared with `KaniCompiler` and it is initialized as part of `rustc`
-    /// initialization, which may happen after this object is created.
-    /// Since we don't have any guarantees on when the compiler creates the Backend object, neither
-    /// in which thread it will be used, we prefer to explicitly synchronize any query access.
-    queries: Arc<Mutex<QueryDb>>,
-}
+pub struct GotocCodegenBackend {}
 
 impl GotocCodegenBackend {
-    pub fn new(queries: Arc<Mutex<QueryDb>>) -> Self {
-        GotocCodegenBackend { queries }
+    pub fn new() -> Self {
+        GotocCodegenBackend {}
     }
 
     /// Generate code that is reachable from the given starting points.
@@ -140,7 +133,7 @@ impl GotocCodegenBackend {
         // Follow rustc naming convention (cx is abbrev for context).
         // https://rustc-dev-guide.rust-lang.org/conventions.html#naming-conventions
         let mut gcx =
-            GotocCtx::new(tcx, (*self.queries.lock().unwrap()).clone(), machine_model, transformer);
+            QUERY_DB.with(|db| GotocCtx::new(tcx, db.borrow().clone(), machine_model, transformer));
         check_reachable_items(gcx.tcx, &gcx.queries, &items);
 
         let contract_info = with_timer(
@@ -221,7 +214,7 @@ impl GotocCodegenBackend {
 
         // No output should be generated if user selected no_codegen.
         if !tcx.sess.opts.unstable_opts.no_codegen && tcx.sess.opts.output_types.should_codegen() {
-            let pretty = self.queries.lock().unwrap().args().output_pretty_json;
+            let pretty = QUERY_DB.with(|db| db.borrow().args().output_pretty_json);
 
             // Save all the data needed to write this goto file
             // so another thread can handle it in parallel.
@@ -283,7 +276,7 @@ impl GotocCodegenBackend {
 
 impl CodegenBackend for GotocCodegenBackend {
     fn provide(&self, providers: &mut Providers) {
-        provide::provide(providers, &self.queries.lock().unwrap());
+        QUERY_DB.with(|db| provide::provide(providers, &db.borrow()));
     }
 
     fn print_version(&self) {
@@ -341,7 +334,7 @@ impl CodegenBackend for GotocCodegenBackend {
             // needed for generating code to the given crate.
             // The cached information must not outlive the stable-mir `run` scope.
             // See [QueryDb::kani_functions] for more information.
-            let queries = self.queries.lock().unwrap().clone();
+            let queries = QUERY_DB.with(|db| db.borrow().clone());
 
             check_target(tcx.sess);
             check_options(tcx.sess);
diff --git a/kani-compiler/src/kani_compiler.rs b/kani-compiler/src/kani_compiler.rs
@@ -21,7 +21,7 @@ use crate::codegen_aeneas_llbc::LlbcCodegenBackend;
 #[cfg(feature = "cprover")]
 use crate::codegen_cprover_gotoc::GotocCodegenBackend;
 use crate::kani_middle::check_crate_items;
-use crate::kani_queries::QueryDb;
+use crate::kani_queries::QUERY_DB;
 use crate::session::init_session;
 use clap::Parser;
 use rustc_codegen_ssa::traits::CodegenBackend;
@@ -30,7 +30,6 @@ use rustc_interface::Config;
 use rustc_middle::ty::TyCtxt;
 use rustc_public::rustc_internal;
 use rustc_session::config::ErrorOutputType;
-use std::sync::{Arc, Mutex};
 use tracing::debug;
 
 /// Run the Kani flavour of the compiler.
@@ -42,54 +41,55 @@ pub fn run(args: Vec<String>) {
 
 /// Configure the LLBC backend (Aeneas's IR).
 #[allow(unused)]
-fn llbc_backend(_queries: Arc<Mutex<QueryDb>>) -> Box<dyn CodegenBackend> {
+fn llbc_backend(args: Arguments) -> Box<dyn CodegenBackend> {
     #[cfg(feature = "llbc")]
-    return Box::new(LlbcCodegenBackend::new(_queries));
+    {
+        QUERY_DB.with(|db| db.borrow_mut().set_args(args));
+        return Box::new(LlbcCodegenBackend::new());
+    }
     #[cfg(not(feature = "llbc"))]
     unreachable!()
 }
 
 /// Configure the cprover backend that generates goto-programs.
 #[allow(unused)]
-fn cprover_backend(_queries: Arc<Mutex<QueryDb>>) -> Box<dyn CodegenBackend> {
+fn cprover_backend(args: Arguments) -> Box<dyn CodegenBackend> {
     #[cfg(feature = "cprover")]
-    return Box::new(GotocCodegenBackend::new(_queries));
+    {
+        QUERY_DB.with(|db| db.borrow_mut().set_args(args));
+        return Box::new(GotocCodegenBackend::new());
+    }
     #[cfg(not(feature = "cprover"))]
     unreachable!()
 }
 
 #[cfg(any(feature = "cprover", feature = "llbc"))]
-fn backend(queries: Arc<Mutex<QueryDb>>) -> Box<dyn CodegenBackend> {
-    let backend = queries.lock().unwrap().args().backend;
+fn backend(args: Arguments) -> Box<dyn CodegenBackend> {
+    let backend = args.backend;
     match backend {
         #[cfg(feature = "cprover")]
-        BackendOption::CProver => cprover_backend(queries),
+        BackendOption::CProver => cprover_backend(args),
         #[cfg(feature = "llbc")]
-        BackendOption::Llbc => llbc_backend(queries),
+        BackendOption::Llbc => llbc_backend(args),
     }
 }
 
 /// Fallback backend. It will trigger an error if no backend has been enabled.
 #[cfg(not(any(feature = "cprover", feature = "llbc")))]
-fn backend(queries: Arc<Mutex<QueryDb>>) -> Box<CodegenBackend> {
+fn backend(args: Arguments) -> Box<CodegenBackend> {
     compile_error!("No backend is available. Use `cprover` or `llbc`.");
 }
 
 /// This object controls the compiler behavior.
 ///
 /// It is responsible for initializing the query database, as well as controlling the compiler
 /// state machine.
-struct KaniCompiler {
-    /// Store the query database. The queries should be initialized as part of `config` when the
-    /// compiler state is Init.
-    /// Note that we need to share the queries with the backend before `config` is called.
-    pub queries: Arc<Mutex<QueryDb>>,
-}
+struct KaniCompiler {}
 
 impl KaniCompiler {
     /// Create a new [KaniCompiler] instance.
     pub fn new() -> KaniCompiler {
-        KaniCompiler { queries: QueryDb::new() }
+        KaniCompiler {}
     }
 
     /// Compile the current crate with the given arguments.
@@ -107,19 +107,17 @@ impl Callbacks for KaniCompiler {
     /// Configure the [KaniCompiler] `self` object during the [CompilationStage::Init].
     fn config(&mut self, config: &mut Config) {
         let mut args = vec!["kani-compiler".to_string()];
-        config.make_codegen_backend = Some(Box::new({
-            let queries = self.queries.clone();
-            move |_cfg| backend(queries)
-        }));
         // `kani-driver` passes the `kani-compiler` specific arguments through llvm-args, so extract them here.
         args.extend(config.opts.cg.llvm_args.iter().cloned());
         let args = Arguments::parse_from(args);
         init_session(&args, matches!(config.opts.error_format, ErrorOutputType::Json { .. }));
 
-        // Configure queries.
-        let queries = &mut (*self.queries.lock().unwrap());
-        queries.set_args(args);
-        debug!(?queries, "config end");
+        // Capture args in the closure so they're available when the backend is created
+        // (potentially on a different thread).
+        config.make_codegen_backend = Some(Box::new({
+            let args = args.clone();
+            move |_cfg| backend(args)
+        }));
     }
 
     /// After analysis, we check the crate items for Kani API misuse or configuration issues.
@@ -129,7 +127,8 @@ impl Callbacks for KaniCompiler {
         tcx: TyCtxt<'_>,
     ) -> Compilation {
         rustc_internal::run(tcx, || {
-            check_crate_items(tcx, self.queries.lock().unwrap().args().ignore_global_asm);
+            let ignore_global_asm = QUERY_DB.with(|db| db.borrow().args().ignore_global_asm);
+            check_crate_items(tcx, ignore_global_asm);
         })
         .unwrap();
         Compilation::Continue
diff --git a/kani-compiler/src/kani_queries/mod.rs b/kani-compiler/src/kani_queries/mod.rs
@@ -7,11 +7,14 @@ use crate::kani_middle::kani_functions::{
     KaniFunction, find_kani_functions, validate_kani_functions,
 };
 use rustc_public::ty::FnDef;
-use std::cell::OnceCell;
+use std::cell::{OnceCell, RefCell};
 use std::collections::HashMap;
-use std::sync::{Arc, Mutex};
 
-/// This structure should only be used behind a synchronized reference or a snapshot.
+thread_local! {
+    pub static QUERY_DB: RefCell<QueryDb> = RefCell::new(QueryDb::default());
+}
+
+/// This structure is only accessed via thread_local storage to ensure thread safety.
 ///
 /// TODO: Merge this with arguments
 #[derive(Debug, Default, Clone)]
@@ -20,17 +23,7 @@ pub struct QueryDb {
     kani_functions: OnceCell<HashMap<KaniFunction, FnDef>>,
 }
 
-// SAFETY: QueryDb is always used behind Arc<Mutex<>> which provides proper synchronization.
-// The FnDef type from rustc_public is not Send due to ThreadLocalIndex, but we ensure
-// it's only accessed within the appropriate rustc_public context and never actually sent
-// between threads.
-unsafe impl Send for QueryDb {}
-
 impl QueryDb {
-    pub fn new() -> Arc<Mutex<QueryDb>> {
-        Arc::new(Mutex::new(QueryDb::default()))
-    }
-
     pub fn set_args(&mut self, args: Arguments) {
         self.args = Some(args);
     }
