Refactor: Cleanly seperate CurrentFnCtx from GotocCtx.

compiler/rustc_codegen_llvm/src/gotoc/current_fn.rs

@@ -0,0 +1,118 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+use crate::gotoc::GotocCtx;
+use crate::gotoc::Stmt;
+use rustc_hir::def_id::DefId;
+use rustc_middle::mir::BasicBlock;
+use rustc_middle::mir::Body;
+use rustc_middle::ty::Instance;
+use rustc_middle::ty::PolyFnSig;
+use rustc_middle::ty::TyCtxt;
+
+/// This structure represents useful data about the function we are currently compiling.
+pub struct CurrentFnCtx<'tcx> {
+    /// The GOTO block we are compiling into
+    block: Vec<Stmt>,
+    /// The current MIR basic block
+    current_bb: Option<BasicBlock>,
+    /// The codegen instance for the current function
+    instance: Instance<'tcx>,
+    /// The goto labels for all blocks
+    labels: Vec<String>,
+    /// The mir for the current instance
+    mir: &'tcx Body<'tcx>,
+    /// The symbol name of the current function
+    name: String,
+    /// The signature of the current function
+    sig: PolyFnSig<'tcx>,
+    /// A counter to enable creating temporary variables
+    temp_var_counter: u64,
+}
+
+/// Constructor
+impl CurrentFnCtx<'tcx> {
+    pub fn new(instance: Instance<'tcx>, gcx: &GotocCtx<'tcx>) -> Self {
+        Self {
+            block: vec![],
+            current_bb: None,
+            instance,
+            labels: vec![],
+            mir: gcx.tcx.instance_mir(instance.def),
+            name: gcx.symbol_name(instance),
+            sig: gcx.fn_sig_of_instance(instance),
+            temp_var_counter: 0,
+        }
+    }
+}
+
+/// Setters
+impl CurrentFnCtx<'tcx> {
+    /// Returns the current block, replacing it with an empty vector.
+    pub fn extract_block(&mut self) -> Vec<Stmt> {
+        std::mem::replace(&mut self.block, vec![])
+    }
+
+    pub fn get_and_incr_counter(&mut self) -> u64 {
+        let rval = self.temp_var_counter;
+        self.temp_var_counter += 1;
+        rval
+    }
+
+    pub fn push_onto_block(&mut self, s: Stmt) {
+        self.block.push(s)
+    }
+
+    pub fn reset_current_bb(&mut self) {
+        self.current_bb = None;
+    }
+
+    pub fn set_current_bb(&mut self, bb: BasicBlock) {
+        self.current_bb = Some(bb);
+    }
+
+    pub fn set_labels(&mut self, labels: Vec<String>) {
+        assert!(self.labels.is_empty());
+        self.labels = labels;
+    }
+}
+
+/// Getters
+impl CurrentFnCtx<'tcx> {
+    /// The basic block we are currently compiling
+    pub fn current_bb(&self) -> BasicBlock {
+        self.current_bb.unwrap()
+    }
+
+    /// The function we are currently compiling
+    pub fn instance(&self) -> Instance<'tcx> {
+        self.instance
+    }
+
+    /// The labels in the function we are currently compiling
+    pub fn labels(&self) -> &Vec<String> {
+        &self.labels
+    }
+
+    /// The MIR for the function we are currently compiling
+    pub fn mir(&self) -> &'tcx Body<'tcx> {
+        self.mir
+    }
+
+    /// The name of the function we are currently compiling
+    pub fn name(&self) -> String {
+        self.name.clone()
+    }
+
+    /// The signature of the function we are currently compiling
+    pub fn sig(&self) -> PolyFnSig<'tcx> {
+        self.sig
+    }
+}
+
+/// Utility functions
+impl CurrentFnCtx<'_> {
+    pub fn find_label(&self, bb: &BasicBlock) -> String {
+        self.labels[bb.index()].clone()
+    }
+}

compiler/rustc_codegen_llvm/src/gotoc/hooks.rs

@@ -104,7 +104,10 @@ impl<'tcx> GotocHook<'tcx> for Assume {
         let loc = tcx.codegen_span_option2(span);
 
         Stmt::block(
-            vec![Stmt::assume(cond, loc.clone()), Stmt::goto(tcx.find_label(&target), loc.clone())],
+            vec![
+                Stmt::assume(cond, loc.clone()),
+                Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
+            ],
             loc,
         )
     }
@@ -132,7 +135,7 @@ impl<'tcx> GotocHook<'tcx> for Nondet {
         let target = target.unwrap();
         let pt = tcx.place_ty(&p);
         if pt.is_unit() {
-            Stmt::goto(tcx.find_label(&target), loc)
+            Stmt::goto(tcx.current_fn().find_label(&target), loc)
         } else {
             let pe = tcx.codegen_place(&p).goto_expr;
             Stmt::block(
@@ -143,7 +146,7 @@ impl<'tcx> GotocHook<'tcx> for Nondet {
                         None => Stmt::skip(loc.clone()),
                         Some(f) => Stmt::assume(f.call(vec![pe.address_of()]), loc.clone()),
                     },
-                    Stmt::goto(tcx.find_label(&target), loc.clone()),
+                    Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
                 ],
                 loc,
             )
@@ -229,7 +232,7 @@ impl<'tcx> GotocHook<'tcx> for Intrinsic {
             Stmt::block(
                 vec![
                     tcx.codegen_intrinsic(instance, fargs, &p, span),
-                    Stmt::goto(tcx.find_label(&target), loc.clone()),
+                    Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
                 ],
                 loc,
             )
@@ -263,7 +266,7 @@ impl<'tcx> GotocHook<'tcx> for MemReplace {
         let place_layout = tcx.layout_of(place_type);
         let place_is_zst = place_layout.is_zst();
         if place_is_zst {
-            Stmt::block(vec![Stmt::goto(tcx.find_label(&target), loc.clone())], loc)
+            Stmt::block(vec![Stmt::goto(tcx.current_fn().find_label(&target), loc.clone())], loc)
         } else {
             let dest = fargs.remove(0);
             let src = fargs.remove(0);
@@ -273,7 +276,7 @@ impl<'tcx> GotocHook<'tcx> for MemReplace {
                         .goto_expr
                         .assign(dest.clone().dereference().with_location(loc.clone()), loc.clone()),
                     dest.dereference().assign(src, loc.clone()),
-                    Stmt::goto(tcx.find_label(&target), loc.clone()),
+                    Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
                 ],
                 loc,
             )
@@ -336,7 +339,7 @@ impl<'tcx> GotocHook<'tcx> for MemSwap {
         Stmt::block(
             vec![
                 tcx.find_function(&func_name).unwrap().call(vec![x, y]).as_stmt(loc.clone()),
-                Stmt::goto(tcx.find_label(&target), loc.clone()),
+                Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
             ],
             loc,
         )
@@ -374,7 +377,7 @@ impl<'tcx> GotocHook<'tcx> for PtrRead {
                 tcx.codegen_place(&p)
                     .goto_expr
                     .assign(src.dereference().with_location(loc.clone()), loc.clone()),
-                Stmt::goto(tcx.find_label(&target), loc.clone()),
+                Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
             ],
             loc,
         )
@@ -410,7 +413,7 @@ impl<'tcx> GotocHook<'tcx> for PtrWrite {
         Stmt::block(
             vec![
                 dst.dereference().assign(src, loc.clone()).with_location(loc.clone()),
-                Stmt::goto(tcx.find_label(&target), loc.clone()),
+                Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
             ],
             loc,
         )
@@ -446,7 +449,7 @@ impl<'tcx> GotocHook<'tcx> for RustAlloc {
                                 .cast_to(Type::unsigned_int(8).to_pointer()),
                             loc,
                         ),
-                        Stmt::goto(tcx.find_label(&target), Location::none()),
+                        Stmt::goto(tcx.current_fn().find_label(&target), Location::none()),
                     ],
                     Location::none(),
                 )
@@ -482,7 +485,7 @@ impl<'tcx> GotocHook<'tcx> for RustDealloc {
                         BuiltinFn::Free
                             .call(vec![ptr.cast_to(Type::void_pointer())], loc.clone())
                             .as_stmt(loc.clone()),
-                        Stmt::goto(tcx.find_label(&target), Location::none()),
+                        Stmt::goto(tcx.current_fn().find_label(&target), Location::none()),
                     ],
                     loc,
                 )
@@ -524,7 +527,7 @@ impl<'tcx> GotocHook<'tcx> for RustRealloc {
                         .cast_to(Type::unsigned_int(8).to_pointer()),
                     loc.clone(),
                 ),
-                Stmt::goto(tcx.find_label(&target), loc.clone()),
+                Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
             ],
             loc,
         )
@@ -560,7 +563,7 @@ impl<'tcx> GotocHook<'tcx> for RustAllocZeroed {
                         .cast_to(Type::unsigned_int(8).to_pointer()),
                     loc.clone(),
                 ),
-                Stmt::goto(tcx.find_label(&target), loc.clone()),
+                Stmt::goto(tcx.current_fn().find_label(&target), loc.clone()),
             ],
             loc,
         )
@@ -601,7 +604,7 @@ impl<'tcx> GotocHook<'tcx> for SliceFromRawPart {
                 loc.clone(),
             )
             .with_location(loc.clone());
-        Stmt::block(vec![code, Stmt::goto(tcx.find_label(&target), loc.clone())], loc)
+        Stmt::block(vec![code, Stmt::goto(tcx.current_fn().find_label(&target), loc.clone())], loc)
     }
 }