Stop using visitor for instrumentation target finding

kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/instrumentation_visitor.rs

@@ -16,9 +16,8 @@ use crate::kani_middle::{
 };
 use rustc_middle::ty::TyCtxt;
 use stable_mir::mir::{
-    mono::Instance,
-    visit::{Location, PlaceContext},
-    BasicBlock, MirVisitor, Operand, Place, Rvalue,
+    mono::Instance, BasicBlock, CopyNonOverlapping, InlineAsmOperand, NonDivergingIntrinsic,
+    Operand, Place, Rvalue, Statement, StatementKind, Terminator, TerminatorKind,
 };
 use std::collections::HashSet;
 
@@ -33,6 +32,13 @@ pub struct InstrumentationVisitor<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
 }
 
+enum PlaceOperation {
+    Read,
+    Write,
+    Deinit,
+    Noop,
+}
+
 impl<'a, 'tcx> TargetFinder for InstrumentationVisitor<'a, 'tcx> {
     fn find_next(
         &mut self,
@@ -44,14 +50,11 @@ impl<'a, 'tcx> TargetFinder for InstrumentationVisitor<'a, 'tcx> {
             SourceInstruction::Statement { idx, bb } => {
                 let BasicBlock { statements, .. } = &body.blocks()[bb];
                 let stmt = &statements[idx];
-                self.visit_statement(stmt, self.__location_hack_remove_before_merging(stmt.span))
+                self.check_statement(stmt)
             }
             SourceInstruction::Terminator { bb } => {
                 let BasicBlock { terminator, .. } = &body.blocks()[bb];
-                self.visit_terminator(
-                    terminator,
-                    self.__location_hack_remove_before_merging(terminator.span),
-                )
+                self.check_terminator(terminator)
             }
         }
         self.target.clone()
@@ -76,17 +79,152 @@ impl<'a, 'tcx> InstrumentationVisitor<'a, 'tcx> {
     }
 }
 
-impl<'a, 'tcx> MirVisitor for InstrumentationVisitor<'a, 'tcx> {
-    fn visit_rvalue(&mut self, rvalue: &Rvalue, location: Location) {
+impl<'a, 'tcx> InstrumentationVisitor<'a, 'tcx> {
+    fn check_statement(&mut self, stmt: &Statement) {
+        let Statement { kind, .. } = stmt;
+        match kind {
+            StatementKind::Assign(place, rvalue) => {
+                self.check_place(place, PlaceOperation::Write);
+                self.check_rvalue(rvalue);
+            }
+            StatementKind::FakeRead(_, place) => {
+                // According to the compiler docs, "When executed at runtime this is a nop." For
+                // more info, see
+                // https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.StatementKind.html#variant.FakeRead,
+                self.check_place(place, PlaceOperation::Noop);
+            }
+            StatementKind::SetDiscriminant { place, .. } => {
+                self.check_place(place, PlaceOperation::Write);
+            }
+            StatementKind::Deinit(place) => {
+                self.check_place(place, PlaceOperation::Deinit);
+            }
+            StatementKind::Retag(_, place) => {
+                // According to the compiler docs, "These statements are currently only interpreted
+                // by miri and only generated when -Z mir-emit-retag is passed." For more info, see
+                // https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.StatementKind.html#variant.Retag,
+                self.check_place(place, PlaceOperation::Noop);
+            }
+            StatementKind::PlaceMention(place) => {
+                // According to the compiler docs, "When executed at runtime, this computes the
+                // given place, but then discards it without doing a load." For more info, see
+                // https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.StatementKind.html#variant.PlaceMention,
+                self.check_place(place, PlaceOperation::Noop);
+            }
+            StatementKind::AscribeUserType { place, .. } => {
+                // According to the compiler docs, "When executed at runtime this is a nop." For
+                // more info, see
+                // https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.StatementKind.html#variant.AscribeUserType,
+                self.check_place(place, PlaceOperation::Noop);
+            }
+
+            StatementKind::Intrinsic(intrisic) => match intrisic {
+                NonDivergingIntrinsic::Assume(operand) => {
+                    self.check_operand(operand);
+                }
+                NonDivergingIntrinsic::CopyNonOverlapping(CopyNonOverlapping {
+                    src,
+                    dst,
+                    count,
+                }) => {
+                    self.check_operand(src);
+                    self.check_operand(dst);
+                    self.check_operand(count);
+                }
+            },
+            StatementKind::StorageLive(_)
+            | StatementKind::StorageDead(_)
+            | StatementKind::Coverage(_)
+            | StatementKind::ConstEvalCounter
+            | StatementKind::Nop => {}
+        }
+    }
+
+    fn check_terminator(&mut self, term: &Terminator) {
+        let Terminator { kind, .. } = term;
+        match kind {
+            TerminatorKind::Goto { .. }
+            | TerminatorKind::Resume
+            | TerminatorKind::Abort
+            | TerminatorKind::Unreachable
+            | TerminatorKind::Return => {}
+            TerminatorKind::Assert { cond, .. } => {
+                self.check_operand(cond);
+            }
+            TerminatorKind::Drop { place, .. } => {
+                // According to the documentation, "After drop elaboration: Drop terminators are a
+                // complete nop for types that have no drop glue. For other types, Drop terminators
+                // behave exactly like a call to core::mem::drop_in_place with a pointer to the
+                // given place." Since we check arguments when instrumenting function calls, perhaps
+                // we need to check that one, too. For more info, see:
+                // https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/mir/enum.TerminatorKind.html#variant.Drop
+                self.check_place(place, PlaceOperation::Read);
+            }
+            TerminatorKind::Call { func, args, destination, target: _, unwind: _ } => {
+                self.check_operand(func);
+                for arg in args {
+                    self.check_operand(arg);
+                }
+                self.check_place(destination, PlaceOperation::Write);
+            }
+            TerminatorKind::InlineAsm { operands, .. } => {
+                for op in operands {
+                    let InlineAsmOperand { in_value, out_place, raw_rpr: _ } = op;
+                    if let Some(input) = in_value {
+                        self.check_operand(input);
+                    }
+                    if let Some(output) = out_place {
+                        self.check_place(output, PlaceOperation::Write);
+                    }
+                }
+            }
+            TerminatorKind::SwitchInt { discr, .. } => {
+                self.check_operand(discr);
+            }
+        }
+    }
+
+    fn check_rvalue(&mut self, rvalue: &Rvalue) {
         match rvalue {
-            Rvalue::AddressOf(..) | Rvalue::Ref(..) => {
-                // These operations are always legitimate for us.
+            Rvalue::AddressOf(_, place) | Rvalue::Ref(_, _, place) => {
+                self.check_place(place, PlaceOperation::Noop);
+            }
+            Rvalue::Aggregate(_, operands) => {
+                for op in operands {
+                    self.check_operand(op);
+                }
+            }
+            Rvalue::BinaryOp(_, lhs, rhs) | Rvalue::CheckedBinaryOp(_, lhs, rhs) => {
+                self.check_operand(lhs);
+                self.check_operand(rhs);
+            }
+            Rvalue::Cast(_, op, _)
+            | Rvalue::Repeat(op, _)
+            | Rvalue::ShallowInitBox(op, ..)
+            | Rvalue::UnaryOp(_, op)
+            | Rvalue::Use(op) => {
+                self.check_operand(op);
+            }
+            Rvalue::CopyForDeref(place) | Rvalue::Discriminant(place) | Rvalue::Len(place) => {
+                self.check_place(place, PlaceOperation::Read);
+            }
+            Rvalue::ThreadLocalRef(..) | Rvalue::NullaryOp(..) => {}
+        }
+    }
+
+    fn check_operand(&mut self, operand: &Operand) {
+        match operand {
+            Operand::Copy(place) | Operand::Move(place) => {
+                self.check_place(place, PlaceOperation::Read)
+            }
+            Operand::Constant(_) => {
+                // Those should be safe to skip, as they are either constants or statics. In the
+                // latter case, we handle them in regular uninit visior
             }
-            _ => self.super_rvalue(rvalue, location),
         }
     }
 
-    fn visit_place(&mut self, place: &Place, ptx: PlaceContext, location: Location) {
+    fn check_place(&mut self, place: &Place, place_operation: PlaceOperation) {
         // Match the place by whatever it is pointing to and find an intersection with the targets.
         if self
             .points_to
@@ -95,18 +233,31 @@ impl<'a, 'tcx> MirVisitor for InstrumentationVisitor<'a, 'tcx> {
             .next()
             .is_some()
         {
-            // If we are mutating the place, initialize it.
-            if ptx.is_mutating() {
-                self.push_target(MemoryInitOp::SetRef {
-                    operand: Operand::Copy(place.clone()),
-                    value: true,
-                    position: InsertPosition::After,
-                });
-            } else {
-                // Otherwise, check its initialization.
-                self.push_target(MemoryInitOp::CheckRef { operand: Operand::Copy(place.clone()) });
+            match place_operation {
+                PlaceOperation::Write => {
+                    // If we are mutating the place, initialize it.
+                    self.push_target(MemoryInitOp::SetRef {
+                        operand: Operand::Copy(place.clone()),
+                        value: true,
+                        position: InsertPosition::After,
+                    })
+                }
+                PlaceOperation::Deinit => {
+                    // If we are mutating the place, initialize it.
+                    self.push_target(MemoryInitOp::SetRef {
+                        operand: Operand::Copy(place.clone()),
+                        value: false,
+                        position: InsertPosition::After,
+                    })
+                }
+                PlaceOperation::Read => {
+                    // Otherwise, check its initialization.
+                    self.push_target(MemoryInitOp::CheckRef {
+                        operand: Operand::Copy(place.clone()),
+                    });
+                }
+                PlaceOperation::Noop => {}
             }
         }
-        self.super_place(place, ptx, location)
     }
 }

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

@@ -13,12 +13,10 @@ use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
 use stable_mir::{
     mir::{
-        mono::Instance, visit::Location, AggregateKind, BasicBlock, Body, ConstOperand, Mutability,
-        Operand, Place, Rvalue, Statement, StatementKind, Terminator, TerminatorKind, UnwindAction,
-    },
-    ty::{
-        FnDef, GenericArgKind, GenericArgs, MirConst, RigidTy, Span, Ty, TyConst, TyKind, UintTy,
+        mono::Instance, AggregateKind, BasicBlock, Body, ConstOperand, Mutability, Operand, Place,
+        Rvalue, Statement, StatementKind, Terminator, TerminatorKind, UnwindAction,
     },
+    ty::{FnDef, GenericArgKind, GenericArgs, MirConst, RigidTy, Ty, TyConst, TyKind, UintTy},
     CrateDef,
 };
 use std::collections::HashMap;
@@ -39,12 +37,6 @@ pub trait TargetFinder {
         body: &MutableBody,
         source: &SourceInstruction,
     ) -> Option<InitRelevantInstruction>;
-
-    /// TODO: This is just a temporary fix to unblock me because I couldn't create Location
-    /// directly. Perhaps we need to make a change in StableMIR to allow creating Location.
-    fn __location_hack_remove_before_merging(&self, span: Span) -> Location {
-        unsafe { std::mem::transmute(span) }
-    }
 }
 
 // Function bodies of those functions will not be instrumented as not to cause infinite recursion.