Better explanations and extra functionality in MutableBody (rust-lang#3382)

`MutableBody` is the core data structure for MIR manipulation in
instrumentation passes. However, its public methods have limited
documentation and unclear semantics. This slowed down the development of
instrumentation passes.

This PR aims to fix that by:
- Clarifying how source instruction shifts when the methods are called;
- Adding functionality for inserting basic blocks;
- Expanding the support for different terminator types.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.

Author: artemagvanian

kani-compiler/src/kani_middle/transform/body.rs

@@ -88,16 +88,20 @@ impl MutableBody {
 
     pub fn new_str_operand(&mut self, msg: &str, span: Span) -> Operand {
         let literal = MirConst::from_str(msg);
-        Operand::Constant(ConstOperand { span, user_ty: None, const_: literal })
+        self.new_const_operand(literal, span)
     }
 
-    pub fn new_const_operand(&mut self, val: u128, uint_ty: UintTy, span: Span) -> Operand {
+    pub fn new_uint_operand(&mut self, val: u128, uint_ty: UintTy, span: Span) -> Operand {
         let literal = MirConst::try_from_uint(val, uint_ty).unwrap();
+        self.new_const_operand(literal, span)
+    }
+
+    fn new_const_operand(&mut self, literal: MirConst, span: Span) -> Operand {
         Operand::Constant(ConstOperand { span, user_ty: None, const_: literal })
     }
 
     /// Create a raw pointer of `*mut type` and return a new local where that value is stored.
-    pub fn new_cast_ptr(
+    pub fn insert_ptr_cast(
         &mut self,
         from: Operand,
         pointee_ty: Ty,
@@ -108,13 +112,13 @@ impl MutableBody {
         assert!(from.ty(self.locals()).unwrap().kind().is_raw_ptr());
         let target_ty = Ty::new_ptr(pointee_ty, mutability);
         let rvalue = Rvalue::Cast(CastKind::PtrToPtr, from, target_ty);
-        self.new_assignment(rvalue, source, position)
+        self.insert_assignment(rvalue, source, position)
     }
 
     /// Add a new assignment for the given binary operation.
     ///
     /// Return the local where the result is saved.
-    pub fn new_binary_op(
+    pub fn insert_binary_op(
         &mut self,
         bin_op: BinOp,
         lhs: Operand,
@@ -123,13 +127,13 @@ impl MutableBody {
         position: InsertPosition,
     ) -> Local {
         let rvalue = Rvalue::BinaryOp(bin_op, lhs, rhs);
-        self.new_assignment(rvalue, source, position)
+        self.insert_assignment(rvalue, source, position)
     }
 
     /// Add a new assignment.
     ///
-    /// Return  local where the result is saved.
-    pub fn new_assignment(
+    /// Return the local where the result is saved.
+    pub fn insert_assignment(
         &mut self,
         rvalue: Rvalue,
         source: &mut SourceInstruction,
@@ -146,9 +150,10 @@ impl MutableBody {
     /// Add a new assert to the basic block indicated by the given index.
     ///
     /// The new assertion will have the same span as the source instruction, and the basic block
-    /// will be split. The source instruction will be adjusted to point to the first instruction in
-    /// the new basic block.
-    pub fn add_check(
+    /// will be split. If `InsertPosition` is `InsertPosition::Before`, `source` will point to the
+    /// same instruction as before. If `InsertPosition` is `InsertPosition::After`, `source` will
+    /// point to the new terminator.
+    pub fn insert_check(
         &mut self,
         tcx: TyCtxt,
         check_type: &CheckType,
@@ -183,7 +188,7 @@ impl MutableBody {
                     unwind: UnwindAction::Terminate,
                 };
                 let terminator = Terminator { kind, span };
-                self.split_bb(source, position, terminator);
+                self.insert_terminator(source, position, terminator);
             }
             CheckType::Panic | CheckType::NoCore => {
                 tcx.sess
@@ -199,10 +204,11 @@ impl MutableBody {
 
     /// Add a new call to the basic block indicated by the given index.
     ///
-    /// The new call will have the same span as the source instruction, and the basic block
-    /// will be split. The source instruction will be adjusted to point to the first instruction in
-    /// the new basic block.
-    pub fn add_call(
+    /// The new call will have the same span as the source instruction, and the basic block will be
+    /// split. If `InsertPosition` is `InsertPosition::Before`, `source` will point to the same
+    /// instruction as before. If `InsertPosition` is `InsertPosition::After`, `source` will point
+    /// to the new terminator.
+    pub fn insert_call(
         &mut self,
         callee: &Instance,
         source: &mut SourceInstruction,
@@ -222,13 +228,14 @@ impl MutableBody {
             unwind: UnwindAction::Terminate,
         };
         let terminator = Terminator { kind, span };
-        self.split_bb(source, position, terminator);
+        self.insert_terminator(source, position, terminator);
     }
 
-    /// Split a basic block and use the new terminator in the basic block that was split.
-    ///
-    /// The source is updated to point to the same instruction which is now in the new basic block.
-    pub fn split_bb(
+    /// Split a basic block and use the new terminator in the basic block that was split. If
+    /// `InsertPosition` is `InsertPosition::Before`, `source` will point to the same instruction as
+    /// before. If `InsertPosition` is `InsertPosition::After`, `source` will point to the new
+    /// terminator.
+    fn split_bb(
         &mut self,
         source: &mut SourceInstruction,
         position: InsertPosition,
@@ -245,6 +252,7 @@ impl MutableBody {
     }
 
     /// Split a basic block right before the source location.
+    /// `source` will point to the same instruction as before after the function is done.
     fn split_bb_before(&mut self, source: &mut SourceInstruction, new_term: Terminator) {
         let new_bb_idx = self.blocks.len();
         let (idx, bb) = match source {
@@ -268,46 +276,77 @@ impl MutableBody {
     }
 
     /// Split a basic block right after the source location.
+    /// `source` will point to the new terminator after the function is done.
     fn split_bb_after(&mut self, source: &mut SourceInstruction, mut new_term: Terminator) {
         let new_bb_idx = self.blocks.len();
         match source {
             // Split the current block after the statement located at `source`
             // and move the remaining statements into the new one.
             SourceInstruction::Statement { idx, bb } => {
                 let (orig_idx, orig_bb) = (*idx, *bb);
-                *idx = 0;
-                *bb = new_bb_idx;
                 let old_term = mem::replace(&mut self.blocks[orig_bb].terminator, new_term);
                 let bb_stmts = &mut self.blocks[orig_bb].statements;
                 let remaining = bb_stmts.split_off(orig_idx + 1);
                 let new_bb = BasicBlock { statements: remaining, terminator: old_term };
                 self.blocks.push(new_bb);
+                // Update the source to point at the terminator.
+                *source = SourceInstruction::Terminator { bb: orig_bb };
             }
             // Make the terminator at `source` point at the new block,
             // the terminator of which is a simple Goto instruction.
             SourceInstruction::Terminator { bb } => {
-                let current_terminator = &mut self.blocks.get_mut(*bb).unwrap().terminator;
-                // Kani can only instrument function calls like this.
-                match (&mut current_terminator.kind, &mut new_term.kind) {
-                    (
-                        TerminatorKind::Call { target: Some(target_bb), .. },
-                        TerminatorKind::Call { target: Some(new_target_bb), .. },
-                    ) => {
-                        // Set the new terminator to point where previous terminator pointed.
-                        *new_target_bb = *target_bb;
-                        // Point the current terminator to the new terminator's basic block.
-                        *target_bb = new_bb_idx;
-                        // Update the current poisition.
-                        *bb = new_bb_idx;
-                        self.blocks.push(BasicBlock { statements: vec![], terminator: new_term });
-                    }
-                    _ => unimplemented!("Kani can only split blocks after calls."),
-                };
+                let current_term = &mut self.blocks.get_mut(*bb).unwrap().terminator;
+                let target_bb = get_mut_target_ref(current_term);
+                let new_target_bb = get_mut_target_ref(&mut new_term);
+                // Set the new terminator to point where previous terminator pointed.
+                *new_target_bb = *target_bb;
+                // Point the current terminator to the new terminator's basic block.
+                *target_bb = new_bb_idx;
+                // Update the source to point at the terminator.
+                *bb = new_bb_idx;
+                self.blocks.push(BasicBlock { statements: vec![], terminator: new_term });
             }
         };
     }
 
-    /// Insert statement before or after the source instruction and update the source as needed.
+    /// Insert basic block before or after the source instruction and update `source` accordingly. If
+    /// `InsertPosition` is `InsertPosition::Before`, `source` will point to the same instruction as
+    /// before. If `InsertPosition` is `InsertPosition::After`, `source` will point to the
+    /// terminator of the newly inserted basic block.
+    pub fn insert_bb(
+        &mut self,
+        mut bb: BasicBlock,
+        source: &mut SourceInstruction,
+        position: InsertPosition,
+    ) {
+        // Splitting adds 1 block, so the added block index is len + 1;
+        let split_bb_idx = self.blocks().len();
+        let inserted_bb_idx = self.blocks().len() + 1;
+        // Update the terminator of the basic block to point at the remaining part of the split
+        // basic block.
+        let target = get_mut_target_ref(&mut bb.terminator);
+        *target = split_bb_idx;
+        let new_term = Terminator {
+            kind: TerminatorKind::Goto { target: inserted_bb_idx },
+            span: source.span(&self.blocks),
+        };
+        self.split_bb(source, position, new_term);
+        self.blocks.push(bb);
+    }
+
+    pub fn insert_terminator(
+        &mut self,
+        source: &mut SourceInstruction,
+        position: InsertPosition,
+        terminator: Terminator,
+    ) {
+        self.split_bb(source, position, terminator);
+    }
+
+    /// Insert statement before or after the source instruction and update the source as needed. If
+    /// `InsertPosition` is `InsertPosition::Before`, `source` will point to the same instruction as
+    /// before. If `InsertPosition` is `InsertPosition::After`, `source` will point to the
+    /// newly inserted statement.
     pub fn insert_stmt(
         &mut self,
         new_stmt: Statement,
@@ -338,22 +377,18 @@ impl MutableBody {
                     SourceInstruction::Terminator { bb } => {
                         // Create a new basic block, as we need to append a statement after the terminator.
                         let current_terminator = &mut self.blocks.get_mut(*bb).unwrap().terminator;
-                        // Kani can only instrument function calls in this way.
-                        match &mut current_terminator.kind {
-                            TerminatorKind::Call { target: Some(target_bb), .. } => {
-                                *source = SourceInstruction::Statement { idx: 0, bb: new_bb_idx };
-                                let new_bb = BasicBlock {
-                                    statements: vec![new_stmt],
-                                    terminator: Terminator {
-                                        kind: TerminatorKind::Goto { target: *target_bb },
-                                        span,
-                                    },
-                                };
-                                *target_bb = new_bb_idx;
-                                self.blocks.push(new_bb);
-                            }
-                            _ => unimplemented!("Kani can only insert statements after calls."),
+                        // Update target of the terminator.
+                        let target_bb = get_mut_target_ref(current_terminator);
+                        *source = SourceInstruction::Statement { idx: 0, bb: new_bb_idx };
+                        let new_bb = BasicBlock {
+                            statements: vec![new_stmt],
+                            terminator: Terminator {
+                                kind: TerminatorKind::Goto { target: *target_bb },
+                                span,
+                            },
                         };
+                        *target_bb = new_bb_idx;
+                        self.blocks.push(new_bb);
                     }
                 }
             }
@@ -574,3 +609,15 @@ pub trait MutMirVisitor {
         }
     }
 }
+
+fn get_mut_target_ref(terminator: &mut Terminator) -> &mut BasicBlockIdx {
+    match &mut terminator.kind {
+        TerminatorKind::Assert { target, .. }
+        | TerminatorKind::Drop { target, .. }
+        | TerminatorKind::Goto { target }
+        | TerminatorKind::Call { target: Some(target), .. } => target,
+        _ => unimplemented!(
+            "Kani can only insert instructions after terminators that have a `target` field."
+        ),
+    }
+}

kani-compiler/src/kani_middle/transform/check_uninit/mod.rs

@@ -230,7 +230,7 @@ impl UninitPass {
                     *pointee_info.ty(),
                 );
                 collect_skipped(&operation, body, skip_first);
-                body.add_call(
+                body.insert_call(
                     &is_ptr_initialized_instance,
                     source,
                     operation.position(),
@@ -257,7 +257,7 @@ impl UninitPass {
                 let layout_operand =
                     mk_layout_operand(body, source, operation.position(), &element_layout);
                 collect_skipped(&operation, body, skip_first);
-                body.add_call(
+                body.insert_call(
                     &is_ptr_initialized_instance,
                     source,
                     operation.position(),
@@ -276,7 +276,7 @@ impl UninitPass {
         // Make sure all non-padding bytes are initialized.
         collect_skipped(&operation, body, skip_first);
         let ptr_operand_ty = ptr_operand.ty(body.locals()).unwrap();
-        body.add_check(
+        body.insert_check(
             tcx,
             &self.check_type,
             source,
@@ -345,7 +345,7 @@ impl UninitPass {
                     *pointee_info.ty(),
                 );
                 collect_skipped(&operation, body, skip_first);
-                body.add_call(
+                body.insert_call(
                     &set_ptr_initialized_instance,
                     source,
                     operation.position(),
@@ -372,7 +372,7 @@ impl UninitPass {
                 let layout_operand =
                     mk_layout_operand(body, source, operation.position(), &element_layout);
                 collect_skipped(&operation, body, skip_first);
-                body.add_call(
+                body.insert_call(
                     &set_ptr_initialized_instance,
                     source,
                     operation.position(),
@@ -408,8 +408,8 @@ impl UninitPass {
             span,
             user_ty: None,
         }));
-        let result = body.new_assignment(rvalue, source, position);
-        body.add_check(tcx, &self.check_type, source, position, result, reason);
+        let result = body.insert_assignment(rvalue, source, position);
+        body.insert_check(tcx, &self.check_type, source, position, result, reason);
     }
 }
 
@@ -432,7 +432,7 @@ pub fn mk_layout_operand(
     layout_byte_mask: &[bool],
 ) -> Operand {
     Operand::Move(Place {
-        local: body.new_assignment(
+        local: body.insert_assignment(
             Rvalue::Aggregate(
                 AggregateKind::Array(Ty::bool_ty()),
                 layout_byte_mask
@@ -531,7 +531,7 @@ fn inject_memory_init_setup(
     )
     .unwrap();
 
-    new_body.add_call(
+    new_body.insert_call(
         &memory_initialization_init,
         &mut source,
         InsertPosition::Before,

kani-compiler/src/kani_middle/transform/check_uninit/uninit_visitor.rs

@@ -57,7 +57,7 @@ impl MemoryInitOp {
                         Operand::Copy(place) | Operand::Move(place) => place,
                         Operand::Constant(_) => unreachable!(),
                     };
-                    body.new_assignment(
+                    body.insert_assignment(
                         Rvalue::AddressOf(Mutability::Not, place.clone()),
                         source,
                         self.position(),