Move cached_context out of CodegenCx

4 files changed, 42 insertions, 16 deletions

diff --git a/src/base.rs b/src/base.rs
index 8e3f905166b..c68d33465bc 100644
--- a/src/base.rs
+++ b/src/base.rs
@@ -25,6 +25,7 @@ struct CodegenedFunction<'tcx> {
 
 pub(crate) fn codegen_and_compile_fn<'tcx>(
     cx: &mut crate::CodegenCx<'tcx>,
+    cached_context: &mut Context,
     module: &mut dyn Module,
     instance: Instance<'tcx>,
 ) {
@@ -32,13 +33,15 @@ pub(crate) fn codegen_and_compile_fn<'tcx>(
     let _inst_guard =
         crate::PrintOnPanic(|| format!("{:?} {}", instance, tcx.symbol_name(instance).name));
 
-    let codegened_func = codegen_fn(cx, module, instance);
+    let cached_func = std::mem::replace(&mut cached_context.func, Function::new());
+    let codegened_func = codegen_fn(cx, cached_func, module, instance);
 
-    compile_fn(cx, module, codegened_func);
+    compile_fn(cx, cached_context, module, codegened_func);
 }
 
 fn codegen_fn<'tcx>(
     cx: &mut crate::CodegenCx<'tcx>,
+    cached_func: Function,
     module: &mut dyn Module,
     instance: Instance<'tcx>,
 ) -> CodegenedFunction<'tcx> {
@@ -61,11 +64,10 @@ fn codegen_fn<'tcx>(
     let sig = get_function_sig(tcx, module.isa().triple(), instance);
     let func_id = module.declare_function(symbol_name.name, Linkage::Local, &sig).unwrap();
 
-    cx.cached_context.clear();
-
     // Make the FunctionBuilder
     let mut func_ctx = FunctionBuilderContext::new();
-    let mut func = std::mem::replace(&mut cx.cached_context.func, Function::new());
+    let mut func = cached_func;
+    func.clear();
     func.name = ExternalName::user(0, func_id.as_u32());
     func.signature = sig;
     func.collect_debug_info();
@@ -140,6 +142,7 @@ fn codegen_fn<'tcx>(
 
 fn compile_fn<'tcx>(
     cx: &mut crate::CodegenCx<'tcx>,
+    cached_context: &mut Context,
     module: &mut dyn Module,
     codegened_func: CodegenedFunction<'tcx>,
 ) {
@@ -148,7 +151,7 @@ fn compile_fn<'tcx>(
     let mut clif_comments = codegened_func.clif_comments;
 
     // Store function in context
-    let context = &mut cx.cached_context;
+    let context = cached_context;
     context.clear();
     context.func = codegened_func.func;
 
diff --git a/src/driver/aot.rs b/src/driver/aot.rs
index 802e8ebd6f6..6aa28637943 100644
--- a/src/driver/aot.rs
+++ b/src/driver/aot.rs
@@ -128,11 +128,17 @@ fn module_codegen(
         cgu_name,
     );
     super::predefine_mono_items(tcx, &mut module, &mono_items);
+    let mut cached_context = Context::new();
     for (mono_item, _) in mono_items {
         match mono_item {
             MonoItem::Fn(inst) => {
                 cx.tcx.sess.time("codegen fn", || {
-                    crate::base::codegen_and_compile_fn(&mut cx, &mut module, inst)
+                    crate::base::codegen_and_compile_fn(
+                        &mut cx,
+                        &mut cached_context,
+                        &mut module,
+                        inst,
+                    )
                 });
             }
             MonoItem::Static(def_id) => crate::constant::codegen_static(tcx, &mut module, def_id),
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index a7ea2b182ab..1b046d7ec6e 100644
--- a/src/driver/jit.rs
+++ b/src/driver/jit.rs
@@ -111,6 +111,7 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         &backend_config,
         matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
     );
+    let mut cached_context = Context::new();
 
     let (_, cgus) = tcx.collect_and_partition_mono_items(());
     let mono_items = cgus
@@ -129,10 +130,17 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
                     CodegenMode::Aot => unreachable!(),
                     CodegenMode::Jit => {
                         cx.tcx.sess.time("codegen fn", || {
-                            crate::base::codegen_and_compile_fn(&mut cx, &mut jit_module, inst)
+                            crate::base::codegen_and_compile_fn(
+                                &mut cx,
+                                &mut cached_context,
+                                &mut jit_module,
+                                inst,
+                            )
                         });
                     }
-                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
+                    CodegenMode::JitLazy => {
+                        codegen_shim(&mut cx, &mut cached_context, &mut jit_module, inst)
+                    }
                 },
                 MonoItem::Static(def_id) => {
                     crate::constant::codegen_static(tcx, &mut jit_module, def_id);
@@ -260,7 +268,12 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
                 Symbol::intern("dummy_cgu_name"),
             );
             tcx.sess.time("codegen fn", || {
-                crate::base::codegen_and_compile_fn(&mut cx, jit_module, instance)
+                crate::base::codegen_and_compile_fn(
+                    &mut cx,
+                    &mut Context::new(),
+                    jit_module,
+                    instance,
+                )
             });
 
             assert!(cx.global_asm.is_empty());
@@ -336,7 +349,12 @@ fn load_imported_symbols_for_jit(
     imported_symbols
 }
 
-fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
+fn codegen_shim<'tcx>(
+    cx: &mut CodegenCx<'tcx>,
+    cached_context: &mut Context,
+    module: &mut JITModule,
+    inst: Instance<'tcx>,
+) {
     let tcx = cx.tcx;
 
     let pointer_type = module.target_config().pointer_type();
@@ -359,8 +377,9 @@ fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: In
         )
         .unwrap();
 
-    cx.cached_context.clear();
-    let trampoline = &mut cx.cached_context.func;
+    let context = cached_context;
+    context.clear();
+    let trampoline = &mut context.func;
     trampoline.signature = sig.clone();
 
     let mut builder_ctx = FunctionBuilderContext::new();
@@ -383,5 +402,5 @@ fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: In
     let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
     trampoline_builder.ins().return_(&ret_vals);
 
-    module.define_function(func_id, &mut cx.cached_context).unwrap();
+    module.define_function(func_id, context).unwrap();
 }
diff --git a/src/lib.rs b/src/lib.rs
index bb0793b1deb..a3f8cc4dfa3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -123,7 +123,6 @@ struct CodegenCx<'tcx> {
     tcx: TyCtxt<'tcx>,
     global_asm: String,
     inline_asm_index: Cell<usize>,
-    cached_context: Context,
     debug_context: Option<DebugContext<'tcx>>,
     unwind_context: UnwindContext,
     cgu_name: Symbol,
@@ -150,7 +149,6 @@ impl<'tcx> CodegenCx<'tcx> {
             tcx,
             global_asm: String::new(),
             inline_asm_index: Cell::new(0),
-            cached_context: Context::new(),
             debug_context,
             unwind_context,
             cgu_name,