Stub panics during MIR transformation (#4169)

Modifies the `RustcIntrinsicsPass` to stub out panic functions during
body transformation at the MIR level. This avoids having to bring in
(and do transformation passes on) the large amounts of code they use for
string formatting. Our panic hook has been modified to then detect these
stubs during codegen and still generate the same Goto code for them.

The results of this change on the MIR bloat of `panic!()` and its
callers is shown below:
| test | lines of MIR (before) | lines of MIR (w/ panic stubbing) |
| - | - | - |
| `Option<usize>::unwrap()` | 292 | 217 |
| `Option<usize>::expect()` | 8664 | 225 | 
| `Result<usize, usize>::unwrap()` | 20461 | 20374 |
| `Result<usize, usize>::expect()` | 20463 | 20376 | 
| `panic!("hello!")` | 8498 | 41 |
| `panic!("hello {x}!")` (x being a `&'static str`) | 8502 | 45 |

The lack of significant improvement on the two `Result` functions is
elaborated on in #4168 and will take additional consideration to fix.

Resolves #2835

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

Author: AlexanderPortland

kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs

@@ -326,7 +326,12 @@ struct Panic;
 impl GotocHook for Panic {
     fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
         let def_id = rustc_internal::internal(tcx, instance.def.def_id());
-        Some(def_id) == tcx.lang_items().panic_fn()
+        let kani_tool_attr = attributes::fn_marker(instance.def);
+
+        // we check the attributes to make sure this hook applies to
+        // panic functions we've stubbed too
+        kani_tool_attr.is_some_and(|kani| kani.contains("PanicStub"))
+            || Some(def_id) == tcx.lang_items().panic_fn()
             || tcx.has_attr(def_id, rustc_span::sym::rustc_const_panic_str)
             || Some(def_id) == tcx.lang_items().panic_fmt()
             || Some(def_id) == tcx.lang_items().begin_panic_fn()

kani-compiler/src/kani_middle/kani_functions.rs

@@ -99,6 +99,8 @@ pub enum KaniModel {
     SetSliceChunkPtrInitialized,
     #[strum(serialize = "SetSlicePtrInitializedModel")]
     SetSlicePtrInitialized,
+    #[strum(serialize = "PanicStub")]
+    PanicStub,
     #[strum(serialize = "SetStrPtrInitializedModel")]
     SetStrPtrInitialized,
     #[strum(serialize = "SizeOfDynObjectModel")]

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

@@ -52,15 +52,24 @@ impl TransformPass for RustcIntrinsicsPass {
     /// For every unsafe dereference or a transmute operation, we check all values are valid.
     fn transform(&mut self, tcx: TyCtxt, body: Body, instance: Instance) -> (bool, Body) {
         debug!(function=?instance.name(), "transform");
+
         let mut new_body = MutableBody::from(body);
         let mut visitor =
-            ReplaceIntrinsicCallVisitor::new(&self.models, new_body.locals().to_vec());
+            ReplaceIntrinsicCallVisitor::new(&self.models, new_body.locals().to_vec(), tcx);
         visitor.visit_body(&mut new_body);
         let changed = self.replace_lowered_intrinsics(tcx, &mut new_body);
         (visitor.changed || changed, new_body.into())
     }
 }
 
+fn is_panic_function(tcx: &TyCtxt, def_id: rustc_smir::stable_mir::DefId) -> bool {
+    let def_id = rustc_internal::internal(*tcx, def_id);
+    Some(def_id) == tcx.lang_items().panic_fn()
+        || tcx.has_attr(def_id, rustc_span::sym::rustc_const_panic_str)
+        || Some(def_id) == tcx.lang_items().panic_fmt()
+        || Some(def_id) == tcx.lang_items().begin_panic_fn()
+}
+
 impl RustcIntrinsicsPass {
     pub fn new(queries: &QueryDb) -> Self {
         let models = queries
@@ -132,19 +141,24 @@ impl RustcIntrinsicsPass {
     }
 }
 
-struct ReplaceIntrinsicCallVisitor<'a> {
+struct ReplaceIntrinsicCallVisitor<'a, 'tcx> {
     models: &'a HashMap<KaniModel, FnDef>,
     locals: Vec<LocalDecl>,
+    tcx: TyCtxt<'tcx>,
     changed: bool,
 }
 
-impl<'a> ReplaceIntrinsicCallVisitor<'a> {
-    fn new(models: &'a HashMap<KaniModel, FnDef>, locals: Vec<LocalDecl>) -> Self {
-        ReplaceIntrinsicCallVisitor { models, locals, changed: false }
+impl<'a, 'tcx> ReplaceIntrinsicCallVisitor<'a, 'tcx> {
+    fn new(
+        models: &'a HashMap<KaniModel, FnDef>,
+        locals: Vec<LocalDecl>,
+        tcx: TyCtxt<'tcx>,
+    ) -> Self {
+        ReplaceIntrinsicCallVisitor { models, locals, changed: false, tcx }
     }
 }
 
-impl MutMirVisitor for ReplaceIntrinsicCallVisitor<'_> {
+impl MutMirVisitor for ReplaceIntrinsicCallVisitor<'_, '_> {
     /// Replace the terminator for some rustc's intrinsics.
     ///
     /// In some cases, we replace a function call to a rustc intrinsic by a call to the
@@ -159,22 +173,35 @@ impl MutMirVisitor for ReplaceIntrinsicCallVisitor<'_> {
         if let TerminatorKind::Call { func, .. } = &mut term.kind
             && let TyKind::RigidTy(RigidTy::FnDef(def, args)) =
                 func.ty(&self.locals).unwrap().kind()
-            && def.is_intrinsic()
         {
-            let instance = Instance::resolve(def, &args).unwrap();
-            let intrinsic = Intrinsic::from_instance(&instance);
-            debug!(?intrinsic, "handle_terminator");
-            let model = match intrinsic {
-                Intrinsic::SizeOfVal => self.models[&KaniModel::SizeOfVal],
-                Intrinsic::MinAlignOfVal => self.models[&KaniModel::AlignOfVal],
-                Intrinsic::PtrOffsetFrom => self.models[&KaniModel::PtrOffsetFrom],
-                Intrinsic::PtrOffsetFromUnsigned => self.models[&KaniModel::PtrOffsetFromUnsigned],
-                // The rest is handled in codegen.
-                _ => {
-                    return self.super_terminator(term);
+            // Get the model we should use to replace this function call, if any.
+            let replacement_model = if def.is_intrinsic() {
+                let instance = Instance::resolve(def, &args).unwrap();
+                let intrinsic = Intrinsic::from_instance(&instance);
+                debug!(?intrinsic, "handle_terminator");
+                match intrinsic {
+                    Intrinsic::SizeOfVal => self.models[&KaniModel::SizeOfVal],
+                    Intrinsic::MinAlignOfVal => self.models[&KaniModel::AlignOfVal],
+                    Intrinsic::PtrOffsetFrom => self.models[&KaniModel::PtrOffsetFrom],
+                    Intrinsic::PtrOffsetFromUnsigned => {
+                        self.models[&KaniModel::PtrOffsetFromUnsigned]
+                    }
+                    // The rest is handled in codegen.
+                    _ => {
+                        return self.super_terminator(term);
+                    }
                 }
+            } else if is_panic_function(&self.tcx, def.0) {
+                // If we find a panic function, we replace it with our stub.
+                self.models[&KaniModel::PanicStub]
+            } else {
+                return self.super_terminator(term);
             };
-            let new_instance = Instance::resolve(model, &args).unwrap();
+
+            let new_instance = Instance::resolve(replacement_model, &args).unwrap();
+
+            // Construct the wrapper types needed to insert our resolved model [Instance]
+            // back into the MIR as an operand.
             let literal = MirConst::try_new_zero_sized(new_instance.ty()).unwrap();
             let span = term.span;
             let new_func = ConstOperand { span, user_ty: None, const_: literal };

library/kani_core/src/models.rs

@@ -32,6 +32,13 @@ macro_rules! generate_models {
                 }
             }
 
+            #[kanitool::fn_marker = "PanicStub"]
+            pub fn panic_stub(t: &str) -> ! {
+                // Using an infinite loop here to have the function return the never (`!`) type.
+                // We could also use `exit()` / `abort()` but both require depending on std::process.
+                loop {}
+            }
+
             #[kanitool::fn_marker = "AlignOfValRawModel"]
             pub fn align_of_val_raw<T: ?Sized>(ptr: *const T) -> usize {
                 if let Some(size) = kani::mem::checked_align_of_raw(ptr) {

tests/script-based-pre/mir_stub_panic/config.yml

@@ -0,0 +1,5 @@
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+script: run.sh
+expected: expected
+exit_code: 0

tests/script-based-pre/mir_stub_panic/expected

@@ -0,0 +1,52 @@
+--VERIFYING panic.rs--
+Manual Harness Summary: \
+Verification failed for - main \
+Complete - 0 successfully verified harnesses, 1 failures, 1 total.
+
+--READING MIR for panic.rs--
+ panic | fn kani::rustc_intrinsics::panic_stub(_1: &str) -> ! { \
+ panic |     let mut _0: !; \
+ panic |     debug t => _1; \
+ panic |     bb0: { \
+ panic |         goto -> bb1; \
+ panic |     } \
+ panic |     bb1: { \
+ panic |         goto -> bb1; \
+ panic |     } \
+ panic | }
+
+
+--VERIFYING option.rs--
+Manual Harness Summary: \
+Verification failed for - main \
+Complete - 0 successfully verified harnesses, 1 failures, 1 total.
+
+--READING MIR for option.rs--
+ option | fn kani::rustc_intrinsics::panic_stub(_1: &str) -> ! { \
+ option |     let mut _0: !; \
+ option |     debug t => _1; \
+ option |     bb0: { \
+ option |         goto -> bb1; \
+ option |     } \
+ option |     bb1: { \
+ option |         goto -> bb1; \
+ option |     } \
+ option | }
+
+
+--VERIFYING result.rs--
+Manual Harness Summary: \
+Verification failed for - main \
+Complete - 0 successfully verified harnesses, 1 failures, 1 total.
+
+--READING MIR for result.rs--
+ result | fn kani::rustc_intrinsics::panic_stub(_1: &str) -> ! { \
+ result |     let mut _0: !; \
+ result |     debug t => _1; \
+ result |     bb0: { \
+ result |         goto -> bb1; \
+ result |     } \
+ result |     bb1: { \
+ result |         goto -> bb1; \
+ result |     } \
+ result | }

tests/script-based-pre/mir_stub_panic/option.rs

@@ -0,0 +1,13 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// Ensure the panic!() internal to `Option::unwrap()` is stubbed.
+#[kani::proof]
+fn main() {
+    foo();
+}
+
+fn foo() -> usize {
+    let a: Option<usize> = kani::any();
+    a.unwrap()
+}

tests/script-based-pre/mir_stub_panic/panic.rs

@@ -0,0 +1,8 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// Ensure that the panic!() macro itself gets stubbed.
+#[kani::proof]
+fn main() {
+    panic!("hello!");
+}

tests/script-based-pre/mir_stub_panic/result.rs

@@ -0,0 +1,13 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// Ensure the panic!() internal to `Result::unwrap()` is stubbed.
+#[kani::proof]
+fn main() {
+    foo();
+}
+
+fn foo() -> usize {
+    let a: Result<usize, usize> = kani::any();
+    a.unwrap()
+}

tests/script-based-pre/mir_stub_panic/run.sh

@@ -0,0 +1,17 @@
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+
+echo "--VERIFYING panic.rs--"
+RUSTFLAGS="--emit mir" kani panic.rs
+echo "--READING MIR for panic.rs--"
+cat panic__* | sed 's/^/ panic | /'
+
+echo "--VERIFYING option.rs--"
+RUSTFLAGS="--emit mir" kani option.rs
+echo "--READING MIR for option.rs--"
+cat option__* | sed 's/^/ option | /'
+
+echo "--VERIFYING result.rs--"
+RUSTFLAGS="--emit mir" kani result.rs
+echo "--READING MIR for result.rs--"
+cat result__* | sed 's/^/ result | /'