From c75d374e65304dd30509795fe20236a74e4803e0 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Thu, 14 Nov 2024 13:05:24 -0800
Subject: [PATCH 1/8] Get rid of most kani usage of rustc's diagnostic

Still missing hooks
---
 .../compiler_interface.rs                     |  11 +-
 kani-compiler/src/kani_middle/attributes.rs   |  30 +++-
 .../src/kani_middle/kani_functions.rs         | 154 ++++++++++++++++++
 kani-compiler/src/kani_middle/mod.rs          |   1 +
 .../kani_middle/transform/kani_intrinsics.rs  |  63 +++----
 .../src/kani_middle/transform/mod.rs          |   6 +-
 kani-compiler/src/kani_queries/mod.rs         |  31 +++-
 7 files changed, 253 insertions(+), 43 deletions(-)
 create mode 100644 kani-compiler/src/kani_middle/kani_functions.rs

diff --git a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
index 3935be784311..d959f327e748 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
@@ -237,10 +237,19 @@ impl CodegenBackend for GotocCodegenBackend {
         let ret_val = rustc_internal::run(tcx, || {
             super::utils::init();
 
-            // Queries shouldn't change today once codegen starts.
+            // Any changes to queries from this point on is just related to caching information
+            // for efficiency purpose that should not outlive the stable-mir `run` scope.
             let queries = self.queries.lock().unwrap().clone();
+
             check_target(tcx.sess);
             check_options(tcx.sess);
+            if queries.args().reachability_analysis != ReachabilityType::None
+                && queries.kani_functions().is_empty()
+            {
+                tcx.sess.dcx().err(
+                    "Failed to detect Kani functions. Please check your installation is correct.",
+                );
+            }
 
             // Codegen all items that need to be processed according to the selected reachability mode:
             //
diff --git a/kani-compiler/src/kani_middle/attributes.rs b/kani-compiler/src/kani_middle/attributes.rs
index 368e0ec91495..10d2c227240f 100644
--- a/kani-compiler/src/kani_middle/attributes.rs
+++ b/kani-compiler/src/kani_middle/attributes.rs
@@ -15,17 +15,17 @@ use rustc_middle::ty::{Instance, TyCtxt, TyKind};
 use rustc_session::Session;
 use rustc_smir::rustc_internal;
 use rustc_span::{Span, Symbol};
+use stable_mir::crate_def::Attribute as AttributeStable;
 use stable_mir::mir::mono::Instance as InstanceStable;
-use stable_mir::{CrateDef, DefId as StableDefId};
+use stable_mir::{CrateDef, DefId as StableDefId, Symbol as SymbolStable};
 use std::str::FromStr;
 use strum_macros::{AsRefStr, EnumString};
 use syn::parse::Parser;
 use syn::punctuated::Punctuated;
-use syn::{PathSegment, TypePath};
-
-use tracing::{debug, trace};
+use syn::{Expr, ExprLit, Lit, PathSegment, TypePath};
 
 use super::resolve::{FnResolution, ResolveError, resolve_fn, resolve_fn_path};
+use tracing::{debug, trace};
 
 #[derive(Debug, Clone, Copy, AsRefStr, EnumString, PartialEq, Eq, PartialOrd, Ord)]
 #[strum(serialize_all = "snake_case")]
@@ -1030,6 +1030,12 @@ fn syn_attr(attr: &Attribute) -> syn::Attribute {
     parser.parse_str(&attr_str).unwrap().pop().unwrap()
 }
 
+/// Parse a stable attribute using `syn`.
+fn syn_attr_stable(attr: &AttributeStable) -> syn::Attribute {
+    let parser = syn::Attribute::parse_outer;
+    parser.parse_str(&attr.as_str()).unwrap().pop().unwrap()
+}
+
 /// Return a more user-friendly string for path by trying to remove unneeded whitespace.
 ///
 /// `quote!()` and `TokenString::to_string()` introduce unnecessary space around separators.
@@ -1071,3 +1077,19 @@ fn pretty_type_path(path: &TypePath) -> String {
         format!("{leading}{}", segments_str(&path.path.segments))
     }
 }
+
+pub(crate) fn fn_marker<T: CrateDef>(def: T) -> Option<String> {
+    let fn_marker: [SymbolStable; 2] = ["kanitool".into(), "fn_marker".into()];
+    let marker = def.attrs_by_path(&fn_marker).pop()?;
+    let attribute = syn_attr_stable(&marker);
+    let meta_name = attribute.meta.require_name_value().unwrap_or_else(|_| {
+        panic!("Expected name value attribute for `kanitool::fn_marker`, but found: `{:?}`", marker)
+    });
+    let Expr::Lit(ExprLit { lit: Lit::Str(lit_str), .. }) = &meta_name.value else {
+        panic!(
+            "Expected string literal for `kanitool::fn_marker`, but found: `{:?}`",
+            meta_name.value
+        );
+    };
+    Some(lit_str.value())
+}
diff --git a/kani-compiler/src/kani_middle/kani_functions.rs b/kani-compiler/src/kani_middle/kani_functions.rs
new file mode 100644
index 000000000000..961727bec41c
--- /dev/null
+++ b/kani-compiler/src/kani_middle/kani_functions.rs
@@ -0,0 +1,154 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//! This module contains code for handling special functions from the Kani library.
+//!
+//! There are three types of functions today:
+//!    1. Kani intrinsics: These are functions whose body is generated during
+//!       compilation time. Their body usually require some extra knowledge about the given types
+//!       that's only available during compilation.
+//!    2. Kani models: These are functions that model a specific behavior but that cannot be used
+//!       directly by the user. For example, retrieving information about memory initialization.
+//!       Kani compiler determines when and where to use these models, but the logic itself is
+//!       encoded in Rust.
+//!    3. Kani hooks: These are similar to Kani intrinsics but their code generation depends
+//!       on some backend specific logic. From a Kani library perspective, there is no difference
+//!       between hooks and intrinsics. This is a compiler implementation detail.
+//!
+//! Functions #1 and #2 have a `kanitool::fn_marker` attribute attached to them.
+//! The marker value will contain "Intrinsic" or "Model" suffix, indicating which category they
+//! fit in.
+//!
+//! We are transitioning the third category. It mostly uses rustc diagnostics. See how they are
+//! handled in [crate::codegen_cprover_gotoc::overrides::hooks].
+//!
+//! Note that we still need to migrate some of 1 and 2 to this new structure which are currently
+//! using rustc's diagnostic infrastructure.
+
+use crate::kani_middle::attributes;
+use stable_mir::mir::mono::Instance;
+use stable_mir::ty::FnDef;
+use std::collections::HashMap;
+use std::str::FromStr;
+use strum::IntoEnumIterator;
+use strum_macros::{EnumIter, EnumString, IntoStaticStr};
+use tracing::debug;
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+pub enum KaniFunction {
+    Model(KaniModel),
+    Intrinsic(KaniIntrinsic),
+}
+
+/// Kani intrinsics are functions generated by the compiler.
+///
+/// These functions usually depend on information that require extra knowledge about the type
+/// or extra Kani instrumentation.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, IntoStaticStr, EnumIter, EnumString, Hash)]
+pub enum KaniIntrinsic {
+    #[strum(serialize = "ValidValueIntrinsic")]
+    ValidValue,
+    #[strum(serialize = "IsInitializedIntrinsic")]
+    IsInitialized,
+}
+
+/// Kani models are Rust functions that model some runtime behavior used by Kani instrumentation.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, IntoStaticStr, EnumIter, EnumString, Hash)]
+pub enum KaniModel {
+    #[strum(serialize = "IsPtrInitializedModel")]
+    IsPtrInitialized,
+    #[strum(serialize = "IsStrPtrInitializedModel")]
+    IsStrPtrInitialized,
+    #[strum(serialize = "IsSlicePtrInitializedModel")]
+    IsSlicePtrInitialized,
+}
+
+impl From<KaniIntrinsic> for KaniFunction {
+    fn from(value: KaniIntrinsic) -> Self {
+        KaniFunction::Intrinsic(value)
+    }
+}
+
+impl From<KaniModel> for KaniFunction {
+    fn from(value: KaniModel) -> Self {
+        KaniFunction::Model(value)
+    }
+}
+
+impl TryFrom<FnDef> for KaniFunction {
+    type Error = ();
+
+    fn try_from(def: FnDef) -> Result<Self, Self::Error> {
+        let value = attributes::fn_marker(def).ok_or(())?;
+        if let Ok(intrisic) = KaniIntrinsic::from_str(&value) {
+            Ok(intrisic.into())
+        } else if let Ok(model) = KaniModel::from_str(&value) {
+            Ok(model.into())
+        } else {
+            Err(())
+        }
+    }
+}
+
+impl TryFrom<Instance> for KaniFunction {
+    type Error = ();
+
+    fn try_from(instance: Instance) -> Result<Self, Self::Error> {
+        let value = attributes::fn_marker(instance.def).ok_or(())?;
+        if let Ok(intrisic) = KaniIntrinsic::from_str(&value) {
+            Ok(intrisic.into())
+        } else if let Ok(model) = KaniModel::from_str(&value) {
+            Ok(model.into())
+        } else {
+            Err(())
+        }
+    }
+}
+
+/// Find all Kani functions.
+///
+/// First try to find `kani` crate. If that exists, look for the items there.
+/// If there's no Kani crate, look for the items in `core` since we could be using `kani_core`.
+/// Note that users could have other `kani` crates, so we look in all the ones we find.
+pub fn find_kani_functions() -> HashMap<KaniFunction, FnDef> {
+    let mut kani = stable_mir::find_crates("kani");
+    if kani.is_empty() {
+        // In case we are using `kani_core`.
+        kani.extend(stable_mir::find_crates("core"));
+    }
+    let fns = kani
+        .into_iter()
+        .find_map(|krate| {
+            let kani_funcs: HashMap<_, _> = krate
+                .fn_defs()
+                .into_iter()
+                .filter_map(|fn_def| {
+                    KaniFunction::try_from(fn_def).ok().map(|kani_function| {
+                        debug!(?kani_function, ?fn_def, "Found kani function");
+                        (kani_function, fn_def)
+                    })
+                })
+                .collect();
+            // All definitions should live in the same crate, so we can return the first one.
+            // If there are no definitions, return `None` to indicate that.
+            (!kani_funcs.is_empty()).then_some(kani_funcs)
+        })
+        .unwrap_or_default();
+    if cfg!(debug_assertions) {
+        validate_kani_functions(&fns);
+    }
+    fns
+}
+
+/// Ensure we have the valid definitions.
+pub fn validate_kani_functions(kani_funcs: &HashMap<KaniFunction, FnDef>) {
+    let mut missing = 0u8;
+    for func in KaniIntrinsic::iter().map(|i| i.into()).chain(KaniModel::iter().map(|m| m.into())) {
+        if let Some(fn_def) = kani_funcs.get(&func) {
+            assert_eq!(KaniFunction::try_from(*fn_def), Ok(func), "Unexpected function marker");
+        } else {
+            tracing::error!(?func, "Missing kani function");
+            missing += 1;
+        }
+    }
+    assert_eq!(missing, 0, "Failed to find `{missing}` Kani functions");
+}
diff --git a/kani-compiler/src/kani_middle/mod.rs b/kani-compiler/src/kani_middle/mod.rs
index 3e1c7e3b4fca..7e1092f9950c 100644
--- a/kani-compiler/src/kani_middle/mod.rs
+++ b/kani-compiler/src/kani_middle/mod.rs
@@ -31,6 +31,7 @@ pub mod attributes;
 pub mod codegen_units;
 pub mod coercion;
 mod intrinsics;
+pub mod kani_functions;
 pub mod metadata;
 pub mod points_to;
 pub mod provide;
diff --git a/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs b/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
index 845916af5181..6b185d8bdeff 100644
--- a/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
+++ b/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
@@ -7,14 +7,15 @@
 //! information; thus, they are implemented as a transformation pass where their body get generated
 //! by the transformation.
 
-use crate::args::{Arguments, ExtraChecks};
-use crate::kani_middle::attributes::matches_diagnostic;
+use crate::args::ExtraChecks;
+use crate::kani_middle::attributes::KaniAttributes;
+use crate::kani_middle::kani_functions::{KaniFunction, KaniIntrinsic, KaniModel};
 use crate::kani_middle::transform::body::{
     CheckType, InsertPosition, MutableBody, SourceInstruction,
 };
 use crate::kani_middle::transform::check_uninit::PointeeInfo;
 use crate::kani_middle::transform::check_uninit::{
-    PointeeLayout, get_mem_init_fn_def, mk_layout_operand, resolve_mem_init_fn,
+    PointeeLayout, mk_layout_operand, resolve_mem_init_fn,
 };
 use crate::kani_middle::transform::check_values::{build_limits, ty_validity_per_offset};
 use crate::kani_middle::transform::{TransformPass, TransformationType};
@@ -29,17 +30,17 @@ use stable_mir::target::MachineInfo;
 use stable_mir::ty::{FnDef, MirConst, RigidTy, Ty, TyKind, UintTy};
 use std::collections::HashMap;
 use std::fmt::Debug;
-use strum_macros::AsRefStr;
-use tracing::trace;
+use std::str::FromStr;
+use tracing::{debug, trace};
 
 /// Generate the body for a few Kani intrinsics.
 #[derive(Debug)]
 pub struct IntrinsicGeneratorPass {
-    pub check_type: CheckType,
-    /// Used to cache FnDef lookups of injected memory initialization functions.
-    pub mem_init_fn_cache: HashMap<&'static str, FnDef>,
+    check_type: CheckType,
+    /// Used to cache FnDef lookups for models and Kani intrinsics.
+    kani_defs: HashMap<KaniFunction, FnDef>,
     /// Used to enable intrinsics depending on the flags passed.
-    pub arguments: Arguments,
+    enable_uninit: bool,
 }
 
 impl TransformPass for IntrinsicGeneratorPass {
@@ -47,7 +48,7 @@ impl TransformPass for IntrinsicGeneratorPass {
     where
         Self: Sized,
     {
-        TransformationType::Instrumentation
+        TransformationType::Stubbing
     }
 
     fn is_enabled(&self, _query_db: &QueryDb) -> bool
@@ -61,10 +62,14 @@ impl TransformPass for IntrinsicGeneratorPass {
     /// 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) {
         trace!(function=?instance.name(), "transform");
-        if matches_diagnostic(tcx, instance.def, Intrinsics::KaniValidValue.as_ref()) {
-            (true, self.valid_value_body(tcx, body))
-        } else if matches_diagnostic(tcx, instance.def, Intrinsics::KaniIsInitialized.as_ref()) {
-            (true, self.is_initialized_body(tcx, body))
+        let attributes = KaniAttributes::for_instance(tcx, instance);
+        if let Some(kani_intrinsic) =
+            attributes.fn_marker().and_then(|name| KaniIntrinsic::from_str(name.as_str()).ok())
+        {
+            match kani_intrinsic {
+                KaniIntrinsic::IsInitialized => (true, self.is_initialized_body(tcx, body)),
+                KaniIntrinsic::ValidValue => (true, self.valid_value_body(tcx, body)),
+            }
         } else {
             (false, body)
         }
@@ -72,6 +77,13 @@ impl TransformPass for IntrinsicGeneratorPass {
 }
 
 impl IntrinsicGeneratorPass {
+    pub fn new(check_type: CheckType, queries: &QueryDb) -> Self {
+        let enable_uninit = queries.args().ub_check.contains(&ExtraChecks::Uninit);
+        let kani_defs = queries.kani_functions().clone();
+        debug!(?kani_defs, ?enable_uninit, "IntrinsicGeneratorPass::new");
+        IntrinsicGeneratorPass { check_type, enable_uninit, kani_defs }
+    }
+
     /// Generate the body for valid value. Which should be something like:
     ///
     /// ```
@@ -168,7 +180,7 @@ impl IntrinsicGeneratorPass {
         let mut source = SourceInstruction::Terminator { bb: 0 };
 
         // Short-circut if uninitialized memory checks are not enabled.
-        if !self.arguments.ub_check.contains(&ExtraChecks::Uninit) {
+        if !self.enable_uninit {
             // Initialize return variable with True.
             let span = new_body.locals()[ret_var].span;
             let assign = StatementKind::Assign(
@@ -222,11 +234,7 @@ impl IntrinsicGeneratorPass {
                             return new_body.into();
                         }
                         let is_ptr_initialized_instance = resolve_mem_init_fn(
-                            get_mem_init_fn_def(
-                                tcx,
-                                "KaniIsPtrInitialized",
-                                &mut self.mem_init_fn_cache,
-                            ),
+                            *self.kani_defs.get(&KaniModel::IsPtrInitialized.into()).unwrap(),
                             layout.len(),
                             *pointee_info.ty(),
                         );
@@ -256,17 +264,17 @@ impl IntrinsicGeneratorPass {
                     }
                     PointeeLayout::Slice { element_layout } => {
                         // Since `str`` is a separate type, need to differentiate between [T] and str.
-                        let (slicee_ty, diagnostic) = match pointee_info.ty().kind() {
+                        let (slicee_ty, intrinsic) = match pointee_info.ty().kind() {
                             TyKind::RigidTy(RigidTy::Slice(slicee_ty)) => {
-                                (slicee_ty, "KaniIsSlicePtrInitialized")
+                                (slicee_ty, KaniModel::IsSlicePtrInitialized.into())
                             }
                             TyKind::RigidTy(RigidTy::Str) => {
-                                (Ty::unsigned_ty(UintTy::U8), "KaniIsStrPtrInitialized")
+                                (Ty::unsigned_ty(UintTy::U8), KaniModel::IsStrPtrInitialized.into())
                             }
                             _ => unreachable!(),
                         };
                         let is_ptr_initialized_instance = resolve_mem_init_fn(
-                            get_mem_init_fn_def(tcx, diagnostic, &mut self.mem_init_fn_cache),
+                            *self.kani_defs.get(&intrinsic).unwrap(),
                             element_layout.len(),
                             slicee_ty,
                         );
@@ -363,10 +371,3 @@ impl IntrinsicGeneratorPass {
         new_body.into()
     }
 }
-
-#[derive(Debug, Copy, Clone, Eq, PartialEq, AsRefStr)]
-#[strum(serialize_all = "PascalCase")]
-enum Intrinsics {
-    KaniValidValue,
-    KaniIsInitialized,
-}
diff --git a/kani-compiler/src/kani_middle/transform/mod.rs b/kani-compiler/src/kani_middle/transform/mod.rs
index 3f324d72c39e..e0d986afa3f4 100644
--- a/kani-compiler/src/kani_middle/transform/mod.rs
+++ b/kani-compiler/src/kani_middle/transform/mod.rs
@@ -86,11 +86,7 @@ impl BodyTransformation {
             check_type: CheckType::new_assert(tcx),
             mem_init_fn_cache: HashMap::new(),
         });
-        transformer.add_pass(queries, IntrinsicGeneratorPass {
-            check_type,
-            mem_init_fn_cache: HashMap::new(),
-            arguments: queries.args().clone(),
-        });
+        transformer.add_pass(queries, IntrinsicGeneratorPass::new(check_type, &queries));
         transformer.add_pass(queries, LoopContractPass::new(tcx, &unit));
         transformer
     }
diff --git a/kani-compiler/src/kani_queries/mod.rs b/kani-compiler/src/kani_queries/mod.rs
index bb28237248d3..83b636caacf1 100644
--- a/kani-compiler/src/kani_queries/mod.rs
+++ b/kani-compiler/src/kani_queries/mod.rs
@@ -2,9 +2,14 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 //! Define the communication between KaniCompiler and the codegen implementation.
 
-use std::sync::{Arc, Mutex};
-
 use crate::args::Arguments;
+use crate::kani_middle::kani_functions::{
+    KaniFunction, find_kani_functions, validate_kani_functions,
+};
+use stable_mir::ty::FnDef;
+use std::cell::OnceCell;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
 
 /// This structure should only be used behind a synchronized reference or a snapshot.
 ///
@@ -12,6 +17,7 @@ use crate::args::Arguments;
 #[derive(Debug, Default, Clone)]
 pub struct QueryDb {
     args: Option<Arguments>,
+    kani_functions: OnceCell<HashMap<KaniFunction, FnDef>>,
 }
 
 impl QueryDb {
@@ -26,4 +32,25 @@ impl QueryDb {
     pub fn args(&self) -> &Arguments {
         self.args.as_ref().expect("Arguments have not been initialized")
     }
+
+    /// Return a map with model and intrinsic functions defined in Kani's library.
+    ///
+    /// For `kani_core`, those definitions live in the `core` library.
+    ///
+    /// We cache these definitions to avoid doing the lookup every time it is needed.
+    /// The cache should be invalidated if the `stable_mir` context changes.
+    /// Since that doesn't occur today, we just run a sanity check to ensure all definitions
+    /// are still correct, and abort otherwise.
+    pub fn kani_functions(&self) -> &HashMap<KaniFunction, FnDef> {
+        if let Some(kani_functions) = self.kani_functions.get() {
+            // Sanity check the values stored in case someone misused this API.
+            validate_kani_functions(kani_functions);
+            kani_functions
+        } else {
+            self.kani_functions.get_or_init(|| {
+                // Find all kani functions
+                find_kani_functions()
+            })
+        }
+    }
 }

From 14eb1de1f6c83e777c17392e7648959e06cd54b6 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Thu, 14 Nov 2024 14:22:47 -0800
Subject: [PATCH 2/8] Move almost everything out of diagnostic

The only missing item is the SimdBitmaskModel
---
 .../codegen_cprover_gotoc/overrides/hooks.rs  |  96 ++++++-----
 kani-compiler/src/kani_middle/attributes.rs   |  11 --
 .../src/kani_middle/kani_functions.rs         | 104 +++++++++--
 kani-compiler/src/kani_middle/mod.rs          |  10 --
 .../src/kani_middle/transform/body.rs         |  88 +++-------
 .../transform/check_uninit/delayed_ub/mod.rs  |   8 +-
 .../kani_middle/transform/check_uninit/mod.rs | 163 ++++++++----------
 .../transform/check_uninit/ptr_uninit/mod.rs  |  14 +-
 .../src/kani_middle/transform/check_values.rs |  11 +-
 .../src/kani_middle/transform/contracts.rs    |  22 +--
 .../kani_middle/transform/kani_intrinsics.rs  |  14 +-
 .../kani_middle/transform/loop_contracts.rs   |  10 +-
 .../src/kani_middle/transform/mod.rs          |  15 +-
 library/kani_core/src/lib.rs                  |  40 ++---
 library/kani_core/src/mem.rs                  |  10 +-
 library/kani_core/src/mem_init.rs             |  28 +--
 16 files changed, 323 insertions(+), 321 deletions(-)

diff --git a/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs b/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs
index 74c39c45ed74..3570f38188ad 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs
@@ -10,23 +10,23 @@
 
 use crate::codegen_cprover_gotoc::GotocCtx;
 use crate::codegen_cprover_gotoc::codegen::{PropertyClass, bb_label};
-use crate::kani_middle::attributes::KaniAttributes;
-use crate::kani_middle::attributes::matches_diagnostic as matches_function;
+use crate::kani_middle::attributes;
+use crate::kani_middle::kani_functions::{KaniFunction, KaniHook};
 use crate::unwrap_or_return_codegen_unimplemented_stmt;
 use cbmc::goto_program::CIntType;
 use cbmc::goto_program::{BuiltinFn, Expr, Stmt, Type};
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
-use rustc_span::Symbol;
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::{BasicBlockIdx, Place};
 use stable_mir::{CrateDef, ty::Span};
+use std::collections::HashMap;
 use std::rc::Rc;
 use tracing::debug;
 
 pub trait GotocHook {
     /// if the hook applies, it means the codegen would do something special to it
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool;
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool;
     /// the handler for codegen
     fn handle(
         &self,
@@ -49,8 +49,8 @@ pub trait GotocHook {
 /// for more details.
 struct Cover;
 impl GotocHook for Cover {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniCover")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -84,8 +84,8 @@ impl GotocHook for Cover {
 
 struct Assume;
 impl GotocHook for Assume {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniAssume")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -108,8 +108,8 @@ impl GotocHook for Assume {
 
 struct Assert;
 impl GotocHook for Assert {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniAssert")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -148,8 +148,8 @@ impl GotocHook for Assert {
 
 struct Check;
 impl GotocHook for Check {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniCheck")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -184,8 +184,8 @@ impl GotocHook for Check {
 struct Nondet;
 
 impl GotocHook for Nondet {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniAnyRaw")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -229,7 +229,6 @@ impl GotocHook for Panic {
             || 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()
-            || matches_function(tcx, instance.def, "KaniPanic")
     }
 
     fn handle(
@@ -248,8 +247,8 @@ impl GotocHook for Panic {
 /// Encodes __CPROVER_r_ok(ptr, size)
 struct IsAllocated;
 impl GotocHook for IsAllocated {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniIsAllocated")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -285,8 +284,8 @@ impl GotocHook for IsAllocated {
 /// Encodes __CPROVER_pointer_object(ptr)
 struct PointerObject;
 impl GotocHook for PointerObject {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniPointerObject")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -321,8 +320,8 @@ impl GotocHook for PointerObject {
 /// Encodes __CPROVER_pointer_offset(ptr)
 struct PointerOffset;
 impl GotocHook for PointerOffset {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniPointerOffset")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -467,8 +466,8 @@ impl GotocHook for MemCmp {
 struct UntrackedDeref;
 
 impl GotocHook for UntrackedDeref {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniUntrackedDeref")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -515,8 +514,8 @@ struct InitContracts;
 /// free(NULL);
 /// ```
 impl GotocHook for InitContracts {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        matches_function(tcx, instance.def, "KaniInitContracts")
+    fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
+        unreachable!("Hooks from kani library handled as a map")
     }
 
     fn handle(
@@ -557,9 +556,9 @@ impl GotocHook for InitContracts {
 pub struct LoopInvariantRegister;
 
 impl GotocHook for LoopInvariantRegister {
-    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
-        KaniAttributes::for_instance(tcx, instance).fn_marker()
-            == Some(Symbol::intern("kani_register_loop_contract"))
+    fn hook_applies(&self, _tcx: TyCtxt, instance: Instance) -> bool {
+        attributes::fn_marker(instance.def)
+            .map_or(false, |marker| marker == "kani_register_loop_contract")
     }
 
     fn handle(
@@ -617,37 +616,46 @@ impl GotocHook for LoopInvariantRegister {
 }
 
 pub fn fn_hooks() -> GotocHooks {
+    let kani_hooks: [(KaniHook, Rc<dyn GotocHook>); 11] = [
+        (KaniHook::Assert, Rc::new(Assert)),
+        (KaniHook::Assume, Rc::new(Assume)),
+        (KaniHook::Panic, Rc::new(Panic)),
+        (KaniHook::Check, Rc::new(Check)),
+        (KaniHook::Cover, Rc::new(Cover)),
+        (KaniHook::AnyRaw, Rc::new(Nondet)),
+        (KaniHook::IsAllocated, Rc::new(IsAllocated)),
+        (KaniHook::PointerObject, Rc::new(PointerObject)),
+        (KaniHook::PointerOffset, Rc::new(PointerOffset)),
+        (KaniHook::UntrackedDeref, Rc::new(UntrackedDeref)),
+        (KaniHook::InitContracts, Rc::new(InitContracts)),
+    ];
     GotocHooks {
-        hooks: vec![
+        rustc_hooks: vec![
             Rc::new(Panic),
-            Rc::new(Assume),
-            Rc::new(Assert),
-            Rc::new(Check),
-            Rc::new(Cover),
-            Rc::new(Nondet),
-            Rc::new(IsAllocated),
-            Rc::new(PointerObject),
-            Rc::new(PointerOffset),
             Rc::new(RustAlloc),
             Rc::new(MemCmp),
-            Rc::new(UntrackedDeref),
-            Rc::new(InitContracts),
             Rc::new(LoopInvariantRegister),
         ],
+        kani_hooks: HashMap::from(kani_hooks),
     }
 }
 
 pub struct GotocHooks {
-    hooks: Vec<Rc<dyn GotocHook>>,
+    rustc_hooks: Vec<Rc<dyn GotocHook>>,
+    kani_hooks: HashMap<KaniHook, Rc<dyn GotocHook>>,
 }
 
 impl GotocHooks {
     pub fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> Option<Rc<dyn GotocHook>> {
-        for h in &self.hooks {
-            if h.hook_applies(tcx, instance) {
-                return Some(h.clone());
+        if let Ok(KaniFunction::Hook(hook)) = KaniFunction::try_from(instance) {
+            Some(self.kani_hooks[&hook].clone())
+        } else {
+            for h in &self.rustc_hooks {
+                if h.hook_applies(tcx, instance) {
+                    return Some(h.clone());
+                }
             }
+            None
         }
-        None
     }
 }
diff --git a/kani-compiler/src/kani_middle/attributes.rs b/kani-compiler/src/kani_middle/attributes.rs
index 10d2c227240f..86d89d7d86cd 100644
--- a/kani-compiler/src/kani_middle/attributes.rs
+++ b/kani-compiler/src/kani_middle/attributes.rs
@@ -1010,17 +1010,6 @@ fn attr_kind(tcx: TyCtxt, attr: &Attribute) -> Option<KaniAttributeKind> {
     }
 }
 
-pub fn matches_diagnostic<T: CrateDef>(tcx: TyCtxt, def: T, attr_name: &str) -> bool {
-    let attr_sym = rustc_span::symbol::Symbol::intern(attr_name);
-    if let Some(attr_id) = tcx.all_diagnostic_items(()).name_to_id.get(&attr_sym) {
-        if rustc_internal::internal(tcx, def.def_id()) == *attr_id {
-            debug!("matched: {:?} {:?}", attr_id, attr_sym);
-            return true;
-        }
-    }
-    false
-}
-
 /// Parse an attribute using `syn`.
 ///
 /// This provides a user-friendly interface to manipulate than the internal compiler AST.
diff --git a/kani-compiler/src/kani_middle/kani_functions.rs b/kani-compiler/src/kani_middle/kani_functions.rs
index 961727bec41c..23d127368ef6 100644
--- a/kani-compiler/src/kani_middle/kani_functions.rs
+++ b/kani-compiler/src/kani_middle/kani_functions.rs
@@ -12,17 +12,13 @@
 //!       encoded in Rust.
 //!    3. Kani hooks: These are similar to Kani intrinsics but their code generation depends
 //!       on some backend specific logic. From a Kani library perspective, there is no difference
-//!       between hooks and intrinsics. This is a compiler implementation detail.
+//!       between hooks and intrinsics.
+//!       This is a compiler implementation detail, but for simplicity we use the "Hook" suffix
+//!       in their names.
 //!
-//! Functions #1 and #2 have a `kanitool::fn_marker` attribute attached to them.
-//! The marker value will contain "Intrinsic" or "Model" suffix, indicating which category they
-//! fit in.
-//!
-//! We are transitioning the third category. It mostly uses rustc diagnostics. See how they are
-//! handled in [crate::codegen_cprover_gotoc::overrides::hooks].
-//!
-//! Note that we still need to migrate some of 1 and 2 to this new structure which are currently
-//! using rustc's diagnostic infrastructure.
+//! These special functions are marked with `kanitool::fn_marker` attribute attached to them.
+//! The marker value will contain "Intrinsic", "Model", or "Hook" suffix, indicating which category
+//! they fit in.
 
 use crate::kani_middle::attributes;
 use stable_mir::mir::mono::Instance;
@@ -37,6 +33,7 @@ use tracing::debug;
 pub enum KaniFunction {
     Model(KaniModel),
     Intrinsic(KaniIntrinsic),
+    Hook(KaniHook),
 }
 
 /// Kani intrinsics are functions generated by the compiler.
@@ -45,21 +42,87 @@ pub enum KaniFunction {
 /// or extra Kani instrumentation.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, IntoStaticStr, EnumIter, EnumString, Hash)]
 pub enum KaniIntrinsic {
-    #[strum(serialize = "ValidValueIntrinsic")]
-    ValidValue,
+    #[strum(serialize = "AnyModifiesIntrinsic")]
+    AnyModifies,
     #[strum(serialize = "IsInitializedIntrinsic")]
     IsInitialized,
+    #[strum(serialize = "ValidValueIntrinsic")]
+    ValidValue,
+    #[strum(serialize = "WriteAnyIntrinsic")]
+    WriteAny,
 }
 
 /// Kani models are Rust functions that model some runtime behavior used by Kani instrumentation.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, IntoStaticStr, EnumIter, EnumString, Hash)]
 pub enum KaniModel {
+    #[strum(serialize = "AnyModel")]
+    Any,
+    #[strum(serialize = "CopyInitStateModel")]
+    CopyInitState,
+    #[strum(serialize = "CopyInitStateSingleModel")]
+    CopyInitStateSingle,
+    #[strum(serialize = "LoadArgumentModel")]
+    LoadArgument,
+    #[strum(serialize = "InitializeMemoryInitializationStateModel")]
+    InitializeMemoryInitializationState,
     #[strum(serialize = "IsPtrInitializedModel")]
     IsPtrInitialized,
     #[strum(serialize = "IsStrPtrInitializedModel")]
     IsStrPtrInitialized,
+    #[strum(serialize = "IsSliceChunkPtrInitializedModel")]
+    IsSliceChunkPtrInitialized,
     #[strum(serialize = "IsSlicePtrInitializedModel")]
     IsSlicePtrInitialized,
+    #[strum(serialize = "RunContractModel")]
+    RunContract,
+    #[strum(serialize = "RunLoopContractModel")]
+    RunLoopContract,
+    #[strum(serialize = "SetPtrInitializedModel")]
+    SetPtrInitialized,
+    #[strum(serialize = "SetSliceChunkPtrInitializedModel")]
+    SetSliceChunkPtrInitialized,
+    #[strum(serialize = "SetSlicePtrInitializedModel")]
+    SetSlicePtrInitialized,
+    #[strum(serialize = "SetStrPtrInitializedModel")]
+    SetStrPtrInitialized,
+    #[strum(serialize = "StoreArgumentModel")]
+    StoreArgument,
+    #[strum(serialize = "WriteAnySliceModel")]
+    WriteAnySlice,
+    #[strum(serialize = "WriteAnySlimModel")]
+    WriteAnySlim,
+    #[strum(serialize = "WriteAnyStrModel")]
+    WriteAnyStr,
+}
+
+/// Kani hooks are Rust functions defined in the Kani library that exposes some semantics
+/// from our solvers / symbolic execution engines.
+///
+/// Thus, they get handled by the codegen stage.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, IntoStaticStr, EnumIter, EnumString, Hash)]
+pub enum KaniHook {
+    #[strum(serialize = "AnyRawHook")]
+    AnyRaw,
+    #[strum(serialize = "AssertHook")]
+    Assert,
+    #[strum(serialize = "AssumeHook")]
+    Assume,
+    #[strum(serialize = "CheckHook")]
+    Check,
+    #[strum(serialize = "CoverHook")]
+    Cover,
+    #[strum(serialize = "InitContractsHook")]
+    InitContracts,
+    #[strum(serialize = "IsAllocatedHook")]
+    IsAllocated,
+    #[strum(serialize = "PanicHook")]
+    Panic,
+    #[strum(serialize = "PointerObjectHook")]
+    PointerObject,
+    #[strum(serialize = "PointerOffsetHook")]
+    PointerOffset,
+    #[strum(serialize = "UntrackedDerefHook")]
+    UntrackedDeref,
 }
 
 impl From<KaniIntrinsic> for KaniFunction {
@@ -74,6 +137,12 @@ impl From<KaniModel> for KaniFunction {
     }
 }
 
+impl From<KaniHook> for KaniFunction {
+    fn from(value: KaniHook) -> Self {
+        KaniFunction::Hook(value)
+    }
+}
+
 impl TryFrom<FnDef> for KaniFunction {
     type Error = ();
 
@@ -83,6 +152,8 @@ impl TryFrom<FnDef> for KaniFunction {
             Ok(intrisic.into())
         } else if let Ok(model) = KaniModel::from_str(&value) {
             Ok(model.into())
+        } else if let Ok(hook) = KaniHook::from_str(&value) {
+            Ok(hook.into())
         } else {
             Err(())
         }
@@ -98,6 +169,8 @@ impl TryFrom<Instance> for KaniFunction {
             Ok(intrisic.into())
         } else if let Ok(model) = KaniModel::from_str(&value) {
             Ok(model.into())
+        } else if let Ok(hook) = KaniHook::from_str(&value) {
+            Ok(hook.into())
         } else {
             Err(())
         }
@@ -115,6 +188,7 @@ pub fn find_kani_functions() -> HashMap<KaniFunction, FnDef> {
         // In case we are using `kani_core`.
         kani.extend(stable_mir::find_crates("core"));
     }
+    debug!(?kani, "find_kani_functions");
     let fns = kani
         .into_iter()
         .find_map(|krate| {
@@ -142,7 +216,11 @@ pub fn find_kani_functions() -> HashMap<KaniFunction, FnDef> {
 /// Ensure we have the valid definitions.
 pub fn validate_kani_functions(kani_funcs: &HashMap<KaniFunction, FnDef>) {
     let mut missing = 0u8;
-    for func in KaniIntrinsic::iter().map(|i| i.into()).chain(KaniModel::iter().map(|m| m.into())) {
+    for func in KaniIntrinsic::iter()
+        .map(|i| i.into())
+        .chain(KaniModel::iter().map(|m| m.into()))
+        .chain(KaniHook::iter().map(|h| h.into()))
+    {
         if let Some(fn_def) = kani_funcs.get(&func) {
             assert_eq!(KaniFunction::try_from(*fn_def), Ok(func), "Unexpected function marker");
         } else {
diff --git a/kani-compiler/src/kani_middle/mod.rs b/kani-compiler/src/kani_middle/mod.rs
index 7e1092f9950c..185f9ef8947c 100644
--- a/kani-compiler/src/kani_middle/mod.rs
+++ b/kani-compiler/src/kani_middle/mod.rs
@@ -228,16 +228,6 @@ impl<'tcx> FnAbiOfHelpers<'tcx> for CompilerHelpers<'tcx> {
     }
 }
 
-/// Find an instance of a function from the given crate that has been annotated with `diagnostic`
-/// item.
-fn find_fn_def(tcx: TyCtxt, diagnostic: &str) -> Option<FnDef> {
-    let attr_id = tcx
-        .all_diagnostic_items(())
-        .name_to_id
-        .get(&rustc_span::symbol::Symbol::intern(diagnostic))?;
-    stable_fn_def(tcx, *attr_id)
-}
-
 /// Try to convert an internal `DefId` to a `FnDef`.
 pub fn stable_fn_def(tcx: TyCtxt, def_id: InternalDefId) -> Option<FnDef> {
     if let TyKind::RigidTy(RigidTy::FnDef(def, _)) =
diff --git a/kani-compiler/src/kani_middle/transform/body.rs b/kani-compiler/src/kani_middle/transform/body.rs
index 6cb104c7470a..0bdd4a1603b8 100644
--- a/kani-compiler/src/kani_middle/transform/body.rs
+++ b/kani-compiler/src/kani_middle/transform/body.rs
@@ -3,8 +3,8 @@
 //
 //! Utility functions that allow us to modify a function body.
 
-use crate::kani_middle::find_fn_def;
-use rustc_middle::ty::TyCtxt;
+use crate::kani_middle::kani_functions::KaniHook;
+use crate::kani_queries::QueryDb;
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::*;
 use stable_mir::ty::{GenericArgs, MirConst, Span, Ty, UintTy};
@@ -178,7 +178,6 @@ impl MutableBody {
     /// point to the new terminator.
     pub fn insert_check(
         &mut self,
-        tcx: TyCtxt,
         check_type: &CheckType,
         source: &mut SourceInstruction,
         position: InsertPosition,
@@ -192,37 +191,22 @@ impl MutableBody {
         );
         let new_bb = self.blocks.len();
         let span = source.span(&self.blocks);
-        match check_type {
-            CheckType::Assert(assert_fn) => {
-                let assert_op = Operand::Copy(Place::from(self.new_local(
-                    assert_fn.ty(),
-                    span,
-                    Mutability::Not,
-                )));
-                let msg_op = self.new_str_operand(msg, span);
-                let kind = TerminatorKind::Call {
-                    func: assert_op,
-                    args: vec![Operand::Move(Place::from(value)), msg_op],
-                    destination: Place {
-                        local: self.new_local(Ty::new_tuple(&[]), span, Mutability::Not),
-                        projection: vec![],
-                    },
-                    target: Some(new_bb),
-                    unwind: UnwindAction::Terminate,
-                };
-                let terminator = Terminator { kind, span };
-                self.insert_terminator(source, position, terminator);
-            }
-            CheckType::Panic | CheckType::NoCore => {
-                tcx.sess
-                    .dcx()
-                    .struct_err("Failed to instrument the code. Cannot find `kani::assert`")
-                    .with_note("Kani requires `kani` library in order to verify a crate.")
-                    .emit();
-                tcx.sess.dcx().abort_if_errors();
-                unreachable!();
-            }
-        }
+        let CheckType::Assert(assert_fn) = check_type;
+        let assert_op =
+            Operand::Copy(Place::from(self.new_local(assert_fn.ty(), span, Mutability::Not)));
+        let msg_op = self.new_str_operand(msg, span);
+        let kind = TerminatorKind::Call {
+            func: assert_op,
+            args: vec![Operand::Move(Place::from(value)), msg_op],
+            destination: Place {
+                local: self.new_local(Ty::new_tuple(&[]), span, Mutability::Not),
+                projection: vec![],
+            },
+            target: Some(new_bb),
+            unwind: UnwindAction::Terminate,
+        };
+        let terminator = Terminator { kind, span };
+        self.insert_terminator(source, position, terminator);
     }
 
     /// Add a new call to the basic block indicated by the given index.
@@ -452,32 +436,19 @@ impl MutableBody {
     }
 }
 
+// TODO: Remove-me
 #[derive(Clone, Debug)]
 pub enum CheckType {
     /// This is used by default when the `kani` crate is available.
     Assert(Instance),
-    /// When the `kani` crate is not available, we have to model the check as an `if { panic!() }`.
-    Panic,
-    /// When building non-core crate, such as `rustc-std-workspace-core`, we cannot
-    /// instrument code, but we can still compile them.
-    NoCore,
 }
 
 impl CheckType {
     /// This will create the type of check that is available in the current crate, attempting to
     /// create a check that generates an assertion following by an assumption of the same assertion.
-    ///
-    /// If `kani` crate is available, this will return [CheckType::Assert], and the instance will
-    /// point to `kani::assert`. Otherwise, we will collect the `core::panic_str` method and return
-    /// [CheckType::Panic].
-    pub fn new_assert_assume(tcx: TyCtxt) -> CheckType {
-        if let Some(instance) = find_instance(tcx, "KaniAssert") {
-            CheckType::Assert(instance)
-        } else if find_instance(tcx, "panic_str").is_some() {
-            CheckType::Panic
-        } else {
-            CheckType::NoCore
-        }
+    pub fn new_assert_assume(queries: &QueryDb) -> CheckType {
+        let fn_def = queries.kani_functions()[&KaniHook::Assert.into()];
+        CheckType::Assert(Instance::resolve(fn_def, &GenericArgs(vec![])).unwrap())
     }
 
     /// This will create the type of check that is available in the current crate, attempting to
@@ -486,14 +457,9 @@ impl CheckType {
     /// If `kani` crate is available, this will return [CheckType::Assert], and the instance will
     /// point to `kani::assert`. Otherwise, we will collect the `core::panic_str` method and return
     /// [CheckType::Panic].
-    pub fn new_assert(tcx: TyCtxt) -> CheckType {
-        if let Some(instance) = find_instance(tcx, "KaniCheck") {
-            CheckType::Assert(instance)
-        } else if find_instance(tcx, "panic_str").is_some() {
-            CheckType::Panic
-        } else {
-            CheckType::NoCore
-        }
+    pub fn new_assert(queries: &QueryDb) -> CheckType {
+        let fn_def = queries.kani_functions()[&KaniHook::Check.into()];
+        CheckType::Assert(Instance::resolve(fn_def, &GenericArgs(vec![])).unwrap())
     }
 }
 
@@ -519,10 +485,6 @@ impl SourceInstruction {
     }
 }
 
-fn find_instance(tcx: TyCtxt, diagnostic: &str) -> Option<Instance> {
-    Instance::resolve(find_fn_def(tcx, diagnostic)?, &GenericArgs(vec![])).ok()
-}
-
 /// Basic mutable body visitor.
 ///
 /// We removed many methods for simplicity.
diff --git a/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/mod.rs b/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/mod.rs
index b2a7ffd13bc1..19d625da6c53 100644
--- a/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/mod.rs
+++ b/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/mod.rs
@@ -7,6 +7,7 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 
 use crate::args::ExtraChecks;
+use crate::kani_middle::kani_functions::KaniFunction;
 use crate::kani_middle::{
     points_to::{MemLoc, PointsToGraph, run_points_to_analysis},
     reachability::CallGraph,
@@ -33,12 +34,12 @@ mod instrumentation_visitor;
 #[derive(Debug)]
 pub struct DelayedUbPass {
     pub check_type: CheckType,
-    pub mem_init_fn_cache: HashMap<&'static str, FnDef>,
+    pub mem_init_fn_cache: HashMap<KaniFunction, FnDef>,
 }
 
 impl DelayedUbPass {
-    pub fn new(check_type: CheckType) -> Self {
-        Self { check_type, mem_init_fn_cache: HashMap::new() }
+    pub fn new(check_type: CheckType, queries: &QueryDb) -> Self {
+        Self { check_type, mem_init_fn_cache: queries.kani_functions().clone() }
     }
 }
 
@@ -120,7 +121,6 @@ impl GlobalPass for DelayedUbPass {
                 );
                 let (instrumentation_added, body) = UninitInstrumenter::run(
                     body,
-                    tcx,
                     instance,
                     self.check_type.clone(),
                     &mut self.mem_init_fn_cache,
diff --git a/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs b/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
index 1ce21ef50669..892c2086c65c 100644
--- a/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
+++ b/kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
@@ -4,15 +4,11 @@
 //! Module containing multiple transformation passes that instrument the code to detect possible UB
 //! due to the accesses to uninitialized memory.
 
-use crate::kani_middle::{
-    find_fn_def,
-    transform::body::{CheckType, InsertPosition, MutableBody, SourceInstruction},
+use crate::kani_middle::transform::body::{
+    CheckType, InsertPosition, MutableBody, SourceInstruction,
 };
 use relevant_instruction::{InitRelevantInstruction, MemoryInitOp};
-use rustc_middle::ty::TyCtxt;
-use rustc_smir::rustc_internal;
 use stable_mir::{
-    CrateDef,
     mir::{
         AggregateKind, BasicBlock, Body, ConstOperand, Mutability, Operand, Place, Rvalue,
         Statement, StatementKind, Terminator, TerminatorKind, UnwindAction, mono::Instance,
@@ -21,6 +17,7 @@ use stable_mir::{
 };
 use std::collections::HashMap;
 
+use crate::kani_middle::kani_functions::{KaniFunction, KaniModel};
 pub use delayed_ub::DelayedUbPass;
 pub use ptr_uninit::UninitPass;
 pub use ty_layout::{PointeeInfo, PointeeLayout};
@@ -35,56 +32,59 @@ pub trait TargetFinder {
     fn find_all(self, body: &MutableBody) -> Vec<InitRelevantInstruction>;
 }
 
-const KANI_IS_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniIsPtrInitialized";
-const KANI_SET_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniSetPtrInitialized";
-const KANI_IS_SLICE_CHUNK_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniIsSliceChunkPtrInitialized";
-const KANI_SET_SLICE_CHUNK_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniSetSliceChunkPtrInitialized";
-const KANI_IS_SLICE_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniIsSlicePtrInitialized";
-const KANI_SET_SLICE_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniSetSlicePtrInitialized";
-const KANI_IS_STR_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniIsStrPtrInitialized";
-const KANI_SET_STR_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniSetStrPtrInitialized";
-const KANI_COPY_INIT_STATE_DIAGNOSTIC: &str = "KaniCopyInitState";
-const KANI_COPY_INIT_STATE_SINGLE_DIAGNOSTIC: &str = "KaniCopyInitStateSingle";
-const KANI_LOAD_ARGUMENT_DIAGNOSTIC: &str = "KaniLoadArgument";
-const KANI_STORE_ARGUMENT_DIAGNOSTIC: &str = "KaniStoreArgument";
-const KANI_RESET_ARGUMENT_BUFFER_DIAGNOSTIC: &str = "KaniResetArgumentBuffer";
+const KANI_IS_PTR_INITIALIZED: KaniFunction = KaniFunction::Model(KaniModel::IsPtrInitialized);
+const KANI_SET_PTR_INITIALIZED: KaniFunction = KaniFunction::Model(KaniModel::SetPtrInitialized);
+const KANI_IS_SLICE_CHUNK_PTR_INITIALIZED: KaniFunction =
+    KaniFunction::Model(KaniModel::IsSliceChunkPtrInitialized);
+const KANI_SET_SLICE_CHUNK_PTR_INITIALIZED: KaniFunction =
+    KaniFunction::Model(KaniModel::SetSliceChunkPtrInitialized);
+const KANI_IS_SLICE_PTR_INITIALIZED: KaniFunction =
+    KaniFunction::Model(KaniModel::IsSlicePtrInitialized);
+const KANI_SET_SLICE_PTR_INITIALIZED: KaniFunction =
+    KaniFunction::Model(KaniModel::SetSlicePtrInitialized);
+const KANI_IS_STR_PTR_INITIALIZED: KaniFunction =
+    KaniFunction::Model(KaniModel::IsStrPtrInitialized);
+const KANI_SET_STR_PTR_INITIALIZED: KaniFunction =
+    KaniFunction::Model(KaniModel::SetStrPtrInitialized);
+const KANI_COPY_INIT_STATE: KaniFunction = KaniFunction::Model(KaniModel::CopyInitState);
+const KANI_COPY_INIT_STATE_SINGLE: KaniFunction =
+    KaniFunction::Model(KaniModel::CopyInitStateSingle);
+const KANI_LOAD_ARGUMENT: KaniFunction = KaniFunction::Model(KaniModel::LoadArgument);
+const KANI_STORE_ARGUMENT: KaniFunction = KaniFunction::Model(KaniModel::StoreArgument);
 
 // Function bodies of those functions will not be instrumented as not to cause infinite recursion.
-const SKIPPED_DIAGNOSTIC_ITEMS: &[&str] = &[
-    KANI_IS_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_SET_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_IS_SLICE_CHUNK_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_SET_SLICE_CHUNK_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_IS_SLICE_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_SET_SLICE_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_IS_STR_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_SET_STR_PTR_INITIALIZED_DIAGNOSTIC,
-    KANI_COPY_INIT_STATE_DIAGNOSTIC,
-    KANI_COPY_INIT_STATE_SINGLE_DIAGNOSTIC,
-    KANI_LOAD_ARGUMENT_DIAGNOSTIC,
-    KANI_STORE_ARGUMENT_DIAGNOSTIC,
-    KANI_RESET_ARGUMENT_BUFFER_DIAGNOSTIC,
+const SKIPPED_ITEMS: &[KaniFunction] = &[
+    KANI_IS_PTR_INITIALIZED,
+    KANI_SET_PTR_INITIALIZED,
+    KANI_IS_SLICE_CHUNK_PTR_INITIALIZED,
+    KANI_SET_SLICE_CHUNK_PTR_INITIALIZED,
+    KANI_IS_SLICE_PTR_INITIALIZED,
+    KANI_SET_SLICE_PTR_INITIALIZED,
+    KANI_IS_STR_PTR_INITIALIZED,
+    KANI_SET_STR_PTR_INITIALIZED,
+    KANI_COPY_INIT_STATE,
+    KANI_COPY_INIT_STATE_SINGLE,
+    KANI_LOAD_ARGUMENT,
+    KANI_STORE_ARGUMENT,
 ];
 
 /// Instruments the code with checks for uninitialized memory, agnostic to the source of targets.
-pub struct UninitInstrumenter<'a, 'tcx> {
+pub struct UninitInstrumenter<'a> {
     check_type: CheckType,
     /// Used to cache FnDef lookups of injected memory initialization functions.
-    mem_init_fn_cache: &'a mut HashMap<&'static str, FnDef>,
-    tcx: TyCtxt<'tcx>,
+    mem_init_fn_cache: &'a mut HashMap<KaniFunction, FnDef>,
 }
 
-impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
+impl<'a> UninitInstrumenter<'a> {
     /// Create the instrumenter and run it with the given parameters.
     pub(crate) fn run(
         body: Body,
-        tcx: TyCtxt<'tcx>,
         instance: Instance,
         check_type: CheckType,
-        mem_init_fn_cache: &'a mut HashMap<&'static str, FnDef>,
+        mem_init_fn_cache: &'a mut HashMap<KaniFunction, FnDef>,
         target_finder: impl TargetFinder,
     ) -> (bool, Body) {
-        let mut instrumenter = Self { check_type, mem_init_fn_cache, tcx };
+        let mut instrumenter = Self { check_type, mem_init_fn_cache };
         let body = MutableBody::from(body);
         let (changed, new_body) = instrumenter.instrument(body, instance, target_finder);
         (changed, new_body.into())
@@ -100,14 +100,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
     ) -> (bool, MutableBody) {
         // Need to break infinite recursion when memory initialization checks are inserted, so the
         // internal functions responsible for memory initialization are skipped.
-        if self
-            .tcx
-            .get_diagnostic_name(rustc_internal::internal(self.tcx, instance.def.def_id()))
-            .map(|diagnostic_name| {
-                SKIPPED_DIAGNOSTIC_ITEMS.contains(&diagnostic_name.to_ident_string().as_str())
-            })
-            .unwrap_or(false)
-        {
+        if KaniFunction::try_from(instance).map(|f| SKIPPED_ITEMS.contains(&f)).unwrap_or(false) {
             return (false, body);
         }
 
@@ -146,7 +139,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
         {
             // If the operation is unsupported or trivially accesses uninitialized memory, encode
             // the check as `assert!(false)`.
-            self.inject_assert_false(self.tcx, body, source, operation.position(), reason);
+            self.inject_assert_false(body, source, operation.position(), reason);
             return;
         };
 
@@ -169,7 +162,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                     let reason = format!(
                         "Kani currently doesn't support checking memory initialization for pointers to `{pointee_ty}. {reason}",
                     );
-                    self.inject_assert_false(self.tcx, body, source, operation.position(), &reason);
+                    self.inject_assert_false(body, source, operation.position(), &reason);
                     return;
                 }
             }
@@ -224,12 +217,12 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                 // pass is as an argument.
                 let (diagnostic, args) = match &operation {
                     MemoryInitOp::Check { .. } | MemoryInitOp::CheckRef { .. } => {
-                        let diagnostic = KANI_IS_PTR_INITIALIZED_DIAGNOSTIC;
+                        let diagnostic = KANI_IS_PTR_INITIALIZED;
                         let args = vec![ptr_operand.clone(), layout_operand];
                         (diagnostic, args)
                     }
                     MemoryInitOp::CheckSliceChunk { .. } => {
-                        let diagnostic = KANI_IS_SLICE_CHUNK_PTR_INITIALIZED_DIAGNOSTIC;
+                        let diagnostic = KANI_IS_SLICE_CHUNK_PTR_INITIALIZED;
                         let args =
                             vec![ptr_operand.clone(), layout_operand, operation.expect_count()];
                         (diagnostic, args)
@@ -237,7 +230,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                     _ => unreachable!(),
                 };
                 let is_ptr_initialized_instance = resolve_mem_init_fn(
-                    get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+                    get_mem_init_fn_def(diagnostic, &mut self.mem_init_fn_cache),
                     layout.len(),
                     *pointee_info.ty(),
                 );
@@ -260,15 +253,15 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                 // Since `str`` is a separate type, need to differentiate between [T] and str.
                 let (slicee_ty, diagnostic) = match pointee_info.ty().kind() {
                     TyKind::RigidTy(RigidTy::Slice(slicee_ty)) => {
-                        (slicee_ty, KANI_IS_SLICE_PTR_INITIALIZED_DIAGNOSTIC)
+                        (slicee_ty, KANI_IS_SLICE_PTR_INITIALIZED)
                     }
                     TyKind::RigidTy(RigidTy::Str) => {
-                        (Ty::unsigned_ty(UintTy::U8), KANI_IS_STR_PTR_INITIALIZED_DIAGNOSTIC)
+                        (Ty::unsigned_ty(UintTy::U8), KANI_IS_STR_PTR_INITIALIZED)
                     }
                     _ => unreachable!(),
                 };
                 let is_ptr_initialized_instance = resolve_mem_init_fn(
-                    get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+                    get_mem_init_fn_def(diagnostic, &mut self.mem_init_fn_cache),
                     element_layout.len(),
                     slicee_ty,
                 );
@@ -291,7 +284,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
             }
             PointeeLayout::TraitObject => {
                 let reason = "Kani does not support reasoning about memory initialization of pointers to trait objects.";
-                self.inject_assert_false(self.tcx, body, source, operation.position(), reason);
+                self.inject_assert_false(body, source, operation.position(), reason);
                 return;
             }
             PointeeLayout::Union { .. } => {
@@ -299,7 +292,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                 // TODO: we perhaps need to check that the union at least contains an intersection
                 // of all layouts initialized.
                 let reason = "Interaction between raw pointers and unions is not yet supported.";
-                self.inject_assert_false(self.tcx, body, source, operation.position(), reason);
+                self.inject_assert_false(body, source, operation.position(), reason);
                 return;
             }
         };
@@ -316,7 +309,6 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
             _ => unreachable!(),
         };
         body.insert_check(
-            self.tcx,
             &self.check_type,
             source,
             operation.position(),
@@ -353,7 +345,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                 // pass is as an argument.
                 let (diagnostic, args) = match &operation {
                     MemoryInitOp::Set { .. } | MemoryInitOp::SetRef { .. } => {
-                        let diagnostic = KANI_SET_PTR_INITIALIZED_DIAGNOSTIC;
+                        let diagnostic = KANI_SET_PTR_INITIALIZED;
                         let args = vec![
                             ptr_operand,
                             layout_operand,
@@ -366,7 +358,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                         (diagnostic, args)
                     }
                     MemoryInitOp::SetSliceChunk { .. } => {
-                        let diagnostic = KANI_SET_SLICE_CHUNK_PTR_INITIALIZED_DIAGNOSTIC;
+                        let diagnostic = KANI_SET_SLICE_CHUNK_PTR_INITIALIZED;
                         let args = vec![
                             ptr_operand,
                             layout_operand,
@@ -382,7 +374,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                     _ => unreachable!(),
                 };
                 let set_ptr_initialized_instance = resolve_mem_init_fn(
-                    get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+                    get_mem_init_fn_def(diagnostic, &mut self.mem_init_fn_cache),
                     layout.len(),
                     *pointee_info.ty(),
                 );
@@ -405,15 +397,15 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                 // Since `str`` is a separate type, need to differentiate between [T] and str.
                 let (slicee_ty, diagnostic) = match pointee_info.ty().kind() {
                     TyKind::RigidTy(RigidTy::Slice(slicee_ty)) => {
-                        (slicee_ty, KANI_SET_SLICE_PTR_INITIALIZED_DIAGNOSTIC)
+                        (slicee_ty, KANI_SET_SLICE_PTR_INITIALIZED)
                     }
                     TyKind::RigidTy(RigidTy::Str) => {
-                        (Ty::unsigned_ty(UintTy::U8), KANI_SET_STR_PTR_INITIALIZED_DIAGNOSTIC)
+                        (Ty::unsigned_ty(UintTy::U8), KANI_SET_STR_PTR_INITIALIZED)
                     }
                     _ => unreachable!(),
                 };
                 let set_ptr_initialized_instance = resolve_mem_init_fn(
-                    get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+                    get_mem_init_fn_def(diagnostic, &mut self.mem_init_fn_cache),
                     element_layout.len(),
                     slicee_ty,
                 );
@@ -460,19 +452,13 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                     None => {
                         let reason =
                             "Interaction between raw pointers and unions is not yet supported.";
-                        self.inject_assert_false(
-                            self.tcx,
-                            body,
-                            source,
-                            operation.position(),
-                            reason,
-                        );
+                        self.inject_assert_false(body, source, operation.position(), reason);
                         return;
                     }
                 };
                 let layout = &field_layouts[union_field];
                 let layout_operand = mk_layout_operand(body, &mut statements, source, layout);
-                let diagnostic = KANI_SET_PTR_INITIALIZED_DIAGNOSTIC;
+                let diagnostic = KANI_SET_PTR_INITIALIZED;
                 let args = vec![
                     ptr_operand,
                     layout_operand,
@@ -483,7 +469,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                     }),
                 ];
                 let set_ptr_initialized_instance = resolve_mem_init_fn(
-                    get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+                    get_mem_init_fn_def(diagnostic, &mut self.mem_init_fn_cache),
                     layout.len(),
                     *pointee_info.ty(),
                 );
@@ -521,11 +507,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
         };
         let layout_size = pointee_info.layout().maybe_size().unwrap();
         let copy_init_state_instance = resolve_mem_init_fn(
-            get_mem_init_fn_def(
-                self.tcx,
-                KANI_COPY_INIT_STATE_DIAGNOSTIC,
-                &mut self.mem_init_fn_cache,
-            ),
+            get_mem_init_fn_def(KANI_COPY_INIT_STATE, &mut self.mem_init_fn_cache),
             layout_size,
             *pointee_info.ty(),
         );
@@ -558,12 +540,12 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
         let mut statements = vec![];
         let layout_size = pointee_info.layout().maybe_size().unwrap();
         let diagnostic = match operation {
-            MemoryInitOp::LoadArgument { .. } => KANI_LOAD_ARGUMENT_DIAGNOSTIC,
-            MemoryInitOp::StoreArgument { .. } => KANI_STORE_ARGUMENT_DIAGNOSTIC,
+            MemoryInitOp::LoadArgument { .. } => KANI_LOAD_ARGUMENT,
+            MemoryInitOp::StoreArgument { .. } => KANI_STORE_ARGUMENT,
             _ => unreachable!(),
         };
         let argument_operation_instance = resolve_mem_init_fn(
-            get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+            get_mem_init_fn_def(diagnostic, &mut self.mem_init_fn_cache),
             layout_size,
             *pointee_info.ty(),
         );
@@ -611,11 +593,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
         let mut statements = vec![];
         let layout_size = pointee_info.layout().maybe_size().unwrap();
         let copy_init_state_instance = resolve_mem_init_fn(
-            get_mem_init_fn_def(
-                self.tcx,
-                KANI_COPY_INIT_STATE_SINGLE_DIAGNOSTIC,
-                &mut self.mem_init_fn_cache,
-            ),
+            get_mem_init_fn_def(KANI_COPY_INIT_STATE_SINGLE, &mut self.mem_init_fn_cache),
             layout_size,
             *pointee_info.ty(),
         );
@@ -641,7 +619,6 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
 
     fn inject_assert_false(
         &self,
-        tcx: TyCtxt,
         body: &mut MutableBody,
         source: &mut SourceInstruction,
         position: InsertPosition,
@@ -654,7 +631,7 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
             user_ty: None,
         }));
         let result = body.insert_assignment(rvalue, source, position);
-        body.insert_check(tcx, &self.check_type, source, position, result, reason);
+        body.insert_check(&self.check_type, source, position, result, reason);
     }
 }
 
@@ -700,12 +677,10 @@ pub fn mk_layout_operand(
 
 /// Retrieve a function definition by diagnostic string, caching the result.
 pub fn get_mem_init_fn_def(
-    tcx: TyCtxt,
-    diagnostic: &'static str,
-    cache: &mut HashMap<&'static str, FnDef>,
+    diagnostic: KaniFunction,
+    cache: &mut HashMap<KaniFunction, FnDef>,
 ) -> FnDef {
-    let entry = cache.entry(diagnostic).or_insert_with(|| find_fn_def(tcx, diagnostic).unwrap());
-    *entry
+    cache[&diagnostic]
 }
 
 /// Resolves a given memory initialization function with passed type parameters.
diff --git a/kani-compiler/src/kani_middle/transform/check_uninit/ptr_uninit/mod.rs b/kani-compiler/src/kani_middle/transform/check_uninit/ptr_uninit/mod.rs
index 2408a3b50a6c..b8a8806ec48a 100644
--- a/kani-compiler/src/kani_middle/transform/check_uninit/ptr_uninit/mod.rs
+++ b/kani-compiler/src/kani_middle/transform/check_uninit/ptr_uninit/mod.rs
@@ -5,6 +5,7 @@
 //! uninitialized memory via raw pointers.
 
 use crate::args::ExtraChecks;
+use crate::kani_middle::kani_functions::{KaniFunction, KaniModel};
 use crate::kani_middle::transform::{
     TransformPass, TransformationType,
     body::{CheckType, InsertPosition, MutableBody, SourceInstruction},
@@ -30,7 +31,7 @@ mod uninit_visitor;
 #[derive(Debug)]
 pub struct UninitPass {
     pub check_type: CheckType,
-    pub mem_init_fn_cache: HashMap<&'static str, FnDef>,
+    pub mem_init_fn_cache: HashMap<KaniFunction, FnDef>,
 }
 
 impl TransformPass for UninitPass {
@@ -57,14 +58,13 @@ impl TransformPass for UninitPass {
 
         // Inject a call to set-up memory initialization state if the function is a harness.
         if is_harness(instance, tcx) {
-            inject_memory_init_setup(&mut new_body, tcx, &mut self.mem_init_fn_cache);
+            inject_memory_init_setup(&mut new_body, &mut self.mem_init_fn_cache);
             changed = true;
         }
 
         // Call a helper that performs the actual instrumentation.
         let (instrumentation_added, body) = UninitInstrumenter::run(
             new_body.into(),
-            tcx,
             instance,
             self.check_type.clone(),
             &mut self.mem_init_fn_cache,
@@ -95,8 +95,7 @@ fn is_harness(instance: Instance, tcx: TyCtxt) -> bool {
 /// Inject an initial call to set-up memory initialization tracking.
 fn inject_memory_init_setup(
     new_body: &mut MutableBody,
-    tcx: TyCtxt,
-    mem_init_fn_cache: &mut HashMap<&'static str, FnDef>,
+    mem_init_fn_cache: &mut HashMap<KaniFunction, FnDef>,
 ) {
     // First statement or terminator in the harness.
     let mut source = if !new_body.blocks()[0].statements.is_empty() {
@@ -117,7 +116,10 @@ fn inject_memory_init_setup(
 
     // Resolve the instance and inject a call to set-up the memory initialization state.
     let memory_initialization_init = Instance::resolve(
-        get_mem_init_fn_def(tcx, "KaniInitializeMemoryInitializationState", mem_init_fn_cache),
+        get_mem_init_fn_def(
+            KaniModel::InitializeMemoryInitializationState.into(),
+            mem_init_fn_cache,
+        ),
         &GenericArgs(vec![]),
     )
     .unwrap();
diff --git a/kani-compiler/src/kani_middle/transform/check_values.rs b/kani-compiler/src/kani_middle/transform/check_values.rs
index 04636f517c73..5d66d00cf094 100644
--- a/kani-compiler/src/kani_middle/transform/check_values.rs
+++ b/kani-compiler/src/kani_middle/transform/check_values.rs
@@ -72,14 +72,14 @@ impl TransformPass for ValidValuePass {
             else {
                 continue;
             };
-            self.build_check(tcx, &mut new_body, candidate);
+            self.build_check(&mut new_body, candidate);
         }
         (orig_len != new_body.blocks().len(), new_body.into())
     }
 }
 
 impl ValidValuePass {
-    fn build_check(&self, tcx: TyCtxt, body: &mut MutableBody, instruction: UnsafeInstruction) {
+    fn build_check(&self, body: &mut MutableBody, instruction: UnsafeInstruction) {
         debug!(?instruction, "build_check");
         let mut source = instruction.source;
         for operation in instruction.operations {
@@ -92,7 +92,6 @@ impl ValidValuePass {
                         let msg =
                             format!("Undefined Behavior: Invalid value of type `{target_ty}`",);
                         body.insert_check(
-                            tcx,
                             &self.check_type,
                             &mut source,
                             InsertPosition::Before,
@@ -107,7 +106,6 @@ impl ValidValuePass {
                         let msg =
                             format!("Undefined Behavior: Invalid value of type `{pointee_ty}`",);
                         body.insert_check(
-                            tcx,
                             &self.check_type,
                             &mut source,
                             InsertPosition::Before,
@@ -120,7 +118,7 @@ impl ValidValuePass {
                     let reason = format!(
                         "Kani currently doesn't support checking validity of `{check}` for `{ty}`",
                     );
-                    self.unsupported_check(tcx, body, &mut source, &reason);
+                    self.unsupported_check(body, &mut source, &reason);
                 }
             }
         }
@@ -128,7 +126,6 @@ impl ValidValuePass {
 
     fn unsupported_check(
         &self,
-        tcx: TyCtxt,
         body: &mut MutableBody,
         source: &mut SourceInstruction,
         reason: &str,
@@ -140,7 +137,7 @@ impl ValidValuePass {
             user_ty: None,
         }));
         let result = body.insert_assignment(rvalue, source, InsertPosition::Before);
-        body.insert_check(tcx, &self.check_type, source, InsertPosition::Before, result, reason);
+        body.insert_check(&self.check_type, source, InsertPosition::Before, result, reason);
     }
 }
 
diff --git a/kani-compiler/src/kani_middle/transform/contracts.rs b/kani-compiler/src/kani_middle/transform/contracts.rs
index 480c480a0aea..0281742d008f 100644
--- a/kani-compiler/src/kani_middle/transform/contracts.rs
+++ b/kani-compiler/src/kani_middle/transform/contracts.rs
@@ -3,7 +3,7 @@
 //! This module contains code related to the MIR-to-MIR pass to enable contracts.
 use crate::kani_middle::attributes::KaniAttributes;
 use crate::kani_middle::codegen_units::CodegenUnit;
-use crate::kani_middle::find_fn_def;
+use crate::kani_middle::kani_functions::{KaniIntrinsic, KaniModel};
 use crate::kani_middle::transform::body::{InsertPosition, MutableBody, SourceInstruction};
 use crate::kani_middle::transform::{TransformPass, TransformationType};
 use crate::kani_queries::QueryDb;
@@ -77,13 +77,14 @@ impl TransformPass for AnyModifiesPass {
 
 impl AnyModifiesPass {
     /// Build the pass with non-extern function stubs.
-    pub fn new(tcx: TyCtxt, unit: &CodegenUnit) -> AnyModifiesPass {
-        let kani_any = find_fn_def(tcx, "KaniAny");
-        let kani_any_modifies = find_fn_def(tcx, "KaniAnyModifies");
-        let kani_write_any = find_fn_def(tcx, "KaniWriteAny");
-        let kani_write_any_slim = find_fn_def(tcx, "KaniWriteAnySlim");
-        let kani_write_any_slice = find_fn_def(tcx, "KaniWriteAnySlice");
-        let kani_write_any_str = find_fn_def(tcx, "KaniWriteAnyStr");
+    pub fn new(tcx: TyCtxt, queries: &QueryDb, unit: &CodegenUnit) -> AnyModifiesPass {
+        let kani_fns = queries.kani_functions();
+        let kani_any = kani_fns.get(&KaniModel::Any.into()).copied();
+        let kani_any_modifies = kani_fns.get(&KaniIntrinsic::AnyModifies.into()).copied();
+        let kani_write_any = kani_fns.get(&KaniIntrinsic::WriteAny.into()).copied();
+        let kani_write_any_slim = kani_fns.get(&KaniModel::WriteAnySlim.into()).copied();
+        let kani_write_any_slice = kani_fns.get(&KaniModel::WriteAnySlice.into()).copied();
+        let kani_write_any_str = kani_fns.get(&KaniModel::WriteAnyStr.into()).copied();
         let target_fn = if let Some(harness) = unit.harnesses.first() {
             let attributes = KaniAttributes::for_instance(tcx, *harness);
             let target_fn =
@@ -330,7 +331,7 @@ impl TransformPass for FunctionWithContractPass {
 impl FunctionWithContractPass {
     /// Build the pass by collecting which functions we are stubbing and which ones we are
     /// verifying.
-    pub fn new(tcx: TyCtxt, unit: &CodegenUnit) -> FunctionWithContractPass {
+    pub fn new(tcx: TyCtxt, queries: &QueryDb, unit: &CodegenUnit) -> FunctionWithContractPass {
         if let Some(harness) = unit.harnesses.first() {
             let attrs = KaniAttributes::for_instance(tcx, *harness);
             let check_fn = attrs.interpret_for_contract_attribute().map(|(_, def_id, _)| def_id);
@@ -339,7 +340,8 @@ impl FunctionWithContractPass {
                 .iter()
                 .map(|(_, def_id, _)| *def_id)
                 .collect();
-            let run_contract_fn = find_fn_def(tcx, "KaniRunContract");
+            let run_contract_fn =
+                queries.kani_functions().get(&KaniModel::RunContract.into()).copied();
             assert!(run_contract_fn.is_some(), "Failed to find Kani run contract function");
             FunctionWithContractPass {
                 check_fn,
diff --git a/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs b/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
index 6b185d8bdeff..533c27445b84 100644
--- a/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
+++ b/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
@@ -67,8 +67,10 @@ impl TransformPass for IntrinsicGeneratorPass {
             attributes.fn_marker().and_then(|name| KaniIntrinsic::from_str(name.as_str()).ok())
         {
             match kani_intrinsic {
-                KaniIntrinsic::IsInitialized => (true, self.is_initialized_body(tcx, body)),
-                KaniIntrinsic::ValidValue => (true, self.valid_value_body(tcx, body)),
+                KaniIntrinsic::IsInitialized => (true, self.is_initialized_body(body)),
+                KaniIntrinsic::ValidValue => (true, self.valid_value_body(body)),
+                // This is handled in contracts pass for now.
+                KaniIntrinsic::WriteAny | KaniIntrinsic::AnyModifies => (false, body),
             }
         } else {
             (false, body)
@@ -96,7 +98,7 @@ impl IntrinsicGeneratorPass {
     ///     ret
     /// }
     /// ```
-    fn valid_value_body(&self, tcx: TyCtxt, body: Body) -> Body {
+    fn valid_value_body(&self, body: Body) -> Body {
         let mut new_body = MutableBody::from(body);
         new_body.clear_body(TerminatorKind::Return);
 
@@ -153,7 +155,6 @@ impl IntrinsicGeneratorPass {
                     "Kani currently doesn't support checking validity of `{target_ty}`. {msg}"
                 );
                 new_body.insert_check(
-                    tcx,
                     &self.check_type,
                     &mut terminator,
                     InsertPosition::Before,
@@ -173,7 +174,7 @@ impl IntrinsicGeneratorPass {
     ///     __kani_mem_init_sm_get(ptr, layout, len)
     /// }
     /// ```
-    fn is_initialized_body(&mut self, tcx: TyCtxt, body: Body) -> Body {
+    fn is_initialized_body(&mut self, body: Body) -> Body {
         let mut new_body = MutableBody::from(body);
         new_body.clear_body(TerminatorKind::Return);
         let ret_var = RETURN_LOCAL;
@@ -316,7 +317,6 @@ impl IntrinsicGeneratorPass {
                         let reason: &str = "Kani does not support reasoning about memory initialization of pointers to trait objects.";
 
                         new_body.insert_check(
-                            tcx,
                             &self.check_type,
                             &mut source,
                             InsertPosition::Before,
@@ -336,7 +336,6 @@ impl IntrinsicGeneratorPass {
                             "Kani does not yet support using initialization predicates on unions.";
 
                         new_body.insert_check(
-                            tcx,
                             &self.check_type,
                             &mut source,
                             InsertPosition::Before,
@@ -359,7 +358,6 @@ impl IntrinsicGeneratorPass {
                     "Kani currently doesn't support checking memory initialization for pointers to `{pointee_ty}. {reason}",
                 );
                 new_body.insert_check(
-                    tcx,
                     &self.check_type,
                     &mut source,
                     InsertPosition::Before,
diff --git a/kani-compiler/src/kani_middle/transform/loop_contracts.rs b/kani-compiler/src/kani_middle/transform/loop_contracts.rs
index 887c07ed4792..4c564ac10505 100644
--- a/kani-compiler/src/kani_middle/transform/loop_contracts.rs
+++ b/kani-compiler/src/kani_middle/transform/loop_contracts.rs
@@ -4,9 +4,10 @@
 //! This module contains code related to the MIR-to-MIR pass to enable loop contracts.
 //!
 
+use super::TransformPass;
 use crate::kani_middle::KaniAttributes;
 use crate::kani_middle::codegen_units::CodegenUnit;
-use crate::kani_middle::find_fn_def;
+use crate::kani_middle::kani_functions::KaniModel;
 use crate::kani_middle::transform::TransformationType;
 use crate::kani_middle::transform::body::{InsertPosition, MutableBody, SourceInstruction};
 use crate::kani_queries::QueryDb;
@@ -22,8 +23,6 @@ use stable_mir::ty::{FnDef, MirConst, RigidTy, UintTy};
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::fmt::Debug;
 
-use super::TransformPass;
-
 #[derive(Debug, Default)]
 pub struct LoopContractPass {
     /// Cache KaniRunContract function used to implement contracts.
@@ -146,9 +145,10 @@ impl TransformPass for LoopContractPass {
 }
 
 impl LoopContractPass {
-    pub fn new(tcx: TyCtxt, unit: &CodegenUnit) -> LoopContractPass {
+    pub fn new(_tcx: TyCtxt, queries: &QueryDb, unit: &CodegenUnit) -> LoopContractPass {
         if !unit.harnesses.is_empty() {
-            let run_contract_fn = find_fn_def(tcx, "KaniRunLoopContract");
+            let run_contract_fn =
+                queries.kani_functions().get(&KaniModel::RunLoopContract.into()).copied();
             assert!(run_contract_fn.is_some(), "Failed to find Kani run contract function");
             LoopContractPass { run_contract_fn, new_loop_latches: HashMap::new() }
         } else {
diff --git a/kani-compiler/src/kani_middle/transform/mod.rs b/kani-compiler/src/kani_middle/transform/mod.rs
index e0d986afa3f4..5a02d6e725f2 100644
--- a/kani-compiler/src/kani_middle/transform/mod.rs
+++ b/kani-compiler/src/kani_middle/transform/mod.rs
@@ -68,13 +68,13 @@ impl BodyTransformation {
             inst_passes: vec![],
             cache: Default::default(),
         };
-        let check_type = CheckType::new_assert_assume(tcx);
+        let check_type = CheckType::new_assert_assume(queries);
         transformer.add_pass(queries, FnStubPass::new(&unit.stubs));
         transformer.add_pass(queries, ExternFnStubPass::new(&unit.stubs));
-        transformer.add_pass(queries, FunctionWithContractPass::new(tcx, &unit));
+        transformer.add_pass(queries, FunctionWithContractPass::new(tcx, queries, &unit));
         // This has to come after the contract pass since we want this to only replace the closure
         // body that is relevant for this harness.
-        transformer.add_pass(queries, AnyModifiesPass::new(tcx, &unit));
+        transformer.add_pass(queries, AnyModifiesPass::new(tcx, queries, &unit));
         transformer.add_pass(queries, ValidValuePass { check_type: check_type.clone() });
         // Putting `UninitPass` after `ValidValuePass` makes sure that the code generated by
         // `UninitPass` does not get unnecessarily instrumented by valid value checks. However, it
@@ -83,11 +83,11 @@ impl BodyTransformation {
         // generated code for future instrumentation passes.
         transformer.add_pass(queries, UninitPass {
             // Since this uses demonic non-determinism under the hood, should not assume the assertion.
-            check_type: CheckType::new_assert(tcx),
-            mem_init_fn_cache: HashMap::new(),
+            check_type: CheckType::new_assert(queries),
+            mem_init_fn_cache: queries.kani_functions().clone(),
         });
         transformer.add_pass(queries, IntrinsicGeneratorPass::new(check_type, &queries));
-        transformer.add_pass(queries, LoopContractPass::new(tcx, &unit));
+        transformer.add_pass(queries, LoopContractPass::new(tcx, queries, &unit));
         transformer
     }
 
@@ -188,7 +188,8 @@ pub struct GlobalPasses {
 impl GlobalPasses {
     pub fn new(queries: &QueryDb, tcx: TyCtxt) -> Self {
         let mut global_passes = GlobalPasses { global_passes: vec![] };
-        global_passes.add_global_pass(queries, DelayedUbPass::new(CheckType::new_assert(tcx)));
+        global_passes
+            .add_global_pass(queries, DelayedUbPass::new(CheckType::new_assert(queries), queries));
         global_passes.add_global_pass(queries, DumpMirPass::new(tcx));
         global_passes
     }
diff --git a/library/kani_core/src/lib.rs b/library/kani_core/src/lib.rs
index 4919473bf0f1..a906f3030154 100644
--- a/library/kani_core/src/lib.rs
+++ b/library/kani_core/src/lib.rs
@@ -101,14 +101,14 @@ macro_rules! kani_intrinsics {
         /// kani::assume(i > 10);
         /// ```
         #[inline(never)]
-        #[rustc_diagnostic_item = "KaniAssume"]
+        #[kanitool::fn_marker = "AssumeHook"]
         #[cfg(not(feature = "concrete_playback"))]
         pub fn assume(cond: bool) {
             let _ = cond;
         }
 
         #[inline(never)]
-        #[rustc_diagnostic_item = "KaniAssume"]
+        #[kanitool::fn_marker = "AssumeHook"]
         #[cfg(feature = "concrete_playback")]
         pub fn assume(cond: bool) {
             assert!(cond, "`kani::assume` should always hold");
@@ -125,7 +125,7 @@ macro_rules! kani_intrinsics {
         /// ```
         #[cfg(not(feature = "concrete_playback"))]
         #[inline(never)]
-        #[rustc_diagnostic_item = "KaniAssert"]
+        #[kanitool::fn_marker = "AssertHook"]
         pub const fn assert(cond: bool, msg: &'static str) {
             let _ = cond;
             let _ = msg;
@@ -133,7 +133,7 @@ macro_rules! kani_intrinsics {
 
         #[cfg(feature = "concrete_playback")]
         #[inline(never)]
-        #[rustc_diagnostic_item = "KaniAssert"]
+        #[kanitool::fn_marker = "AssertHook"]
         pub const fn assert(cond: bool, msg: &'static str) {
             assert!(cond, "{}", msg);
         }
@@ -162,7 +162,7 @@ macro_rules! kani_intrinsics {
         /// convenient to use.
         ///
         #[inline(never)]
-        #[rustc_diagnostic_item = "KaniCover"]
+        #[kanitool::fn_marker = "CoverHook"]
         pub const fn cover(_cond: bool, _msg: &'static str) {}
 
         /// This creates an symbolic *valid* value of type `T`. You can assign the return value of this
@@ -185,7 +185,7 @@ macro_rules! kani_intrinsics {
         /// Note: This is a safe construct and can only be used with types that implement the `Arbitrary`
         /// trait. The Arbitrary trait is used to build a symbolic value that represents all possible
         /// valid values for type `T`.
-        #[rustc_diagnostic_item = "KaniAny"]
+        #[kanitool::fn_marker = "AnyModel"]
         #[inline(always)]
         pub fn any<T: Arbitrary>() -> T {
             T::any()
@@ -195,7 +195,7 @@ macro_rules! kani_intrinsics {
         /// It behaves exaclty like `kani::any<T>()`, except it will check for the trait bounds
         /// at compilation time. It allows us to avoid type checking errors while using function
         /// contracts only for verification.
-        #[rustc_diagnostic_item = "KaniAnyModifies"]
+        #[kanitool::fn_marker = "AnyModifiesIntrinsic"]
         #[inline(never)]
         #[doc(hidden)]
         pub fn any_modifies<T>() -> T {
@@ -264,7 +264,7 @@ macro_rules! kani_intrinsics {
         use concrete_playback::{any_raw_array, any_raw_internal};
 
         /// This low-level function returns nondet bytes of size T.
-        #[rustc_diagnostic_item = "KaniAnyRaw"]
+        #[kanitool::fn_marker = "AnyRawHook"]
         #[inline(never)]
         #[allow(dead_code)]
         fn any_raw<T: Copy>() -> T {
@@ -278,7 +278,7 @@ macro_rules! kani_intrinsics {
         /// invoke the regular `core::panic!()` function. This function is used by our standard library
         /// overrides, but not the other way around.
         #[inline(never)]
-        #[rustc_diagnostic_item = "KaniPanic"]
+        #[kanitool::fn_marker = "PanicHook"]
         #[doc(hidden)]
         pub const fn panic(message: &'static str) -> ! {
             panic!("{}", message)
@@ -350,7 +350,7 @@ macro_rules! kani_intrinsics {
             /// guarantee it is done safely.
             #[inline(never)]
             #[doc(hidden)]
-            #[rustc_diagnostic_item = "KaniUntrackedDeref"]
+            #[kanitool::fn_marker = "UntrackedDerefHook"]
             pub fn untracked_deref<T>(_: &T) -> T {
                 todo!()
             }
@@ -366,7 +366,7 @@ macro_rules! kani_intrinsics {
             /// ```
             #[inline(never)]
             #[doc(hidden)]
-            #[rustc_diagnostic_item = "KaniInitContracts"]
+            #[kanitool::fn_marker = "InitContractsHook"]
             pub fn init_contracts() {}
 
             /// This should only be used within contracts. The intent is to
@@ -379,7 +379,7 @@ macro_rules! kani_intrinsics {
             /// Recieves a reference to a pointer-like object and assigns kani::any_modifies to that object.
             /// Only for use within function contracts and will not be replaced if the recursive or function stub
             /// replace contracts are not used.
-            #[rustc_diagnostic_item = "KaniWriteAny"]
+            #[kanitool::fn_marker = "WriteAnyIntrinsic"]
             #[inline(never)]
             #[doc(hidden)]
             pub unsafe fn write_any<T: ?Sized>(_pointer: *mut T) {
@@ -387,12 +387,12 @@ macro_rules! kani_intrinsics {
                 // Users must include `#[kani::recursion]` in any function contracts for recursive functions;
                 // otherwise, this might not be properly instantiate. We mark this as unreachable to make
                 // sure Kani doesn't report any false positives.
-                unreachable!()
+                super::kani_intrinsic()
             }
 
             /// Fill in a slice with kani::any.
             /// Intended as a post compilation replacement for write_any
-            #[rustc_diagnostic_item = "KaniWriteAnySlice"]
+            #[kanitool::fn_marker = "WriteAnySliceModel"]
             #[inline(always)]
             pub unsafe fn write_any_slice<T: Arbitrary>(slice: *mut [T]) {
                 (*slice).fill_with(T::any)
@@ -400,7 +400,7 @@ macro_rules! kani_intrinsics {
 
             /// Fill in a pointer with kani::any.
             /// Intended as a post compilation replacement for write_any
-            #[rustc_diagnostic_item = "KaniWriteAnySlim"]
+            #[kanitool::fn_marker = "WriteAnySlimModel"]
             #[inline(always)]
             pub unsafe fn write_any_slim<T: Arbitrary>(pointer: *mut T) {
                 ptr::write(pointer, T::any())
@@ -409,7 +409,7 @@ macro_rules! kani_intrinsics {
             /// Fill in a str with kani::any.
             /// Intended as a post compilation replacement for write_any.
             /// Not yet implemented
-            #[rustc_diagnostic_item = "KaniWriteAnyStr"]
+            #[kanitool::fn_marker = "WriteAnyStrModel"]
             #[inline(always)]
             pub unsafe fn write_any_str(_s: *mut str) {
                 //TODO: strings introduce new UB
@@ -425,7 +425,7 @@ macro_rules! kani_intrinsics {
             /// so we swap the register body by this function body.
             #[doc(hidden)]
             #[allow(dead_code)]
-            #[rustc_diagnostic_item = "KaniRunContract"]
+            #[kanitool::fn_marker = "RunContractModel"]
             fn run_contract_fn<T, F: FnOnce() -> T>(func: F) -> T {
                 func()
             }
@@ -437,7 +437,7 @@ macro_rules! kani_intrinsics {
             /// so we swap the register body by this function body.
             #[doc(hidden)]
             #[allow(dead_code)]
-            #[rustc_diagnostic_item = "KaniRunLoopContract"]
+            #[kanitool::fn_marker = "RunLoopContractModel"]
             fn run_loop_contract_fn<F: Fn() -> bool>(func: &F, _transformed: usize) -> bool {
                 func()
             }
@@ -474,7 +474,7 @@ macro_rules! kani_intrinsics {
             ///
             #[cfg(not(feature = "concrete_playback"))]
             #[inline(never)]
-            #[rustc_diagnostic_item = "KaniCheck"]
+            #[kanitool::fn_marker = "CheckHook"]
             pub(crate) const fn check(cond: bool, msg: &'static str) {
                 let _ = cond;
                 let _ = msg;
@@ -482,7 +482,7 @@ macro_rules! kani_intrinsics {
 
             #[cfg(feature = "concrete_playback")]
             #[inline(never)]
-            #[rustc_diagnostic_item = "KaniCheck"]
+            #[kanitool::fn_marker = "CheckHook"]
             pub(crate) const fn check(cond: bool, msg: &'static str) {
                 assert!(cond, "{}", msg);
             }
diff --git a/library/kani_core/src/mem.rs b/library/kani_core/src/mem.rs
index 772483e88ad3..ef1e277d18d0 100644
--- a/library/kani_core/src/mem.rs
+++ b/library/kani_core/src/mem.rs
@@ -301,7 +301,7 @@ macro_rules! kani_mem {
         ///
         /// I.e.: This function always returns `true` if the pointer is valid.
         /// Otherwise, it returns non-det boolean.
-        #[rustc_diagnostic_item = "KaniIsAllocated"]
+        #[kanitool::fn_marker = "IsAllocatedHook"]
         #[inline(never)]
         unsafe fn is_allocated(_ptr: *const (), _size: usize) -> bool {
             kani_intrinsic()
@@ -312,14 +312,14 @@ macro_rules! kani_mem {
         /// # Safety
         ///
         /// - Users have to ensure that the pointer is aligned the pointed memory is allocated.
-        #[rustc_diagnostic_item = "KaniValidValue"]
+        #[kanitool::fn_marker = "ValidValueIntrinsic"]
         #[inline(never)]
         unsafe fn has_valid_value<T: ?Sized>(_ptr: *const T) -> bool {
             kani_intrinsic()
         }
 
         /// Check whether `len * size_of::<T>()` bytes are initialized starting from `ptr`.
-        #[rustc_diagnostic_item = "KaniIsInitialized"]
+        #[kanitool::fn_marker = "IsInitializedIntrinsic"]
         #[inline(never)]
         pub(crate) fn is_initialized<T: ?Sized>(_ptr: *const T) -> bool {
             kani_intrinsic()
@@ -353,7 +353,7 @@ macro_rules! kani_mem {
 
         /// Get the object ID of the given pointer.
         #[doc(hidden)]
-        #[rustc_diagnostic_item = "KaniPointerObject"]
+        #[kanitool::fn_marker = "PointerObjectHook"]
         #[inline(never)]
         pub(crate) fn pointer_object<T: ?Sized>(_ptr: *const T) -> usize {
             kani_intrinsic()
@@ -366,7 +366,7 @@ macro_rules! kani_mem {
             issue = 3184,
             reason = "experimental ghost state/shadow memory API"
         )]
-        #[rustc_diagnostic_item = "KaniPointerOffset"]
+        #[kanitool::fn_marker = "PointerOffsetHook"]
         #[inline(never)]
         pub fn pointer_offset<T: ?Sized>(_ptr: *const T) -> usize {
             kani_intrinsic()
diff --git a/library/kani_core/src/mem_init.rs b/library/kani_core/src/mem_init.rs
index 6475c62e31ae..ecafe5d8e861 100644
--- a/library/kani_core/src/mem_init.rs
+++ b/library/kani_core/src/mem_init.rs
@@ -23,7 +23,7 @@
 macro_rules! kani_mem_init {
     ($core:path) => {
         /// Global object for tracking memory initialization state.
-        #[rustc_diagnostic_item = "KaniMemoryInitializationState"]
+        #[kanitool::fn_marker = "MemoryInitializationStateModel"]
         static mut MEM_INIT_STATE: MemoryInitializationState = MemoryInitializationState::new();
 
         /// Global object for tracking union initialization state across function boundaries.
@@ -200,7 +200,7 @@ macro_rules! kani_mem_init {
 
         /// Set tracked object and tracked offset to a non-deterministic value.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniInitializeMemoryInitializationState"]
+        #[kanitool::fn_marker = "InitializeMemoryInitializationStateModel"]
         fn initialize_memory_initialization_state() {
             unsafe {
                 MEM_INIT_STATE.tracked_object_id = super::any();
@@ -211,7 +211,7 @@ macro_rules! kani_mem_init {
 
         /// Get initialization state of `num_elts` items laid out according to the `layout` starting at address `ptr`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniIsPtrInitialized"]
+        #[kanitool::fn_marker = "IsPtrInitializedModel"]
         fn is_ptr_initialized<const LAYOUT_SIZE: usize, T>(
             ptr: *const T,
             layout: Layout<LAYOUT_SIZE>,
@@ -225,7 +225,7 @@ macro_rules! kani_mem_init {
 
         /// Set initialization state to `value` for `num_elts` items laid out according to the `layout` starting at address `ptr`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniSetPtrInitialized"]
+        #[kanitool::fn_marker = "SetPtrInitializedModel"]
         fn set_ptr_initialized<const LAYOUT_SIZE: usize, T>(
             ptr: *const T,
             layout: Layout<LAYOUT_SIZE>,
@@ -242,7 +242,7 @@ macro_rules! kani_mem_init {
 
         /// Get initialization state of `num_elts` items laid out according to the `layout` starting at address `ptr`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniIsSliceChunkPtrInitialized"]
+        #[kanitool::fn_marker = "IsSliceChunkPtrInitializedModel"]
         fn is_slice_chunk_ptr_initialized<const LAYOUT_SIZE: usize, T>(
             ptr: *const T,
             layout: Layout<LAYOUT_SIZE>,
@@ -257,7 +257,7 @@ macro_rules! kani_mem_init {
 
         /// Set initialization state to `value` for `num_elts` items laid out according to the `layout` starting at address `ptr`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniSetSliceChunkPtrInitialized"]
+        #[kanitool::fn_marker = "SetSliceChunkPtrInitializedModel"]
         fn set_slice_chunk_ptr_initialized<const LAYOUT_SIZE: usize, T>(
             ptr: *const T,
             layout: Layout<LAYOUT_SIZE>,
@@ -275,7 +275,7 @@ macro_rules! kani_mem_init {
 
         /// Get initialization state of the slice, items of which are laid out according to the `layout` starting at address `ptr`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniIsSlicePtrInitialized"]
+        #[kanitool::fn_marker = "IsSlicePtrInitializedModel"]
         fn is_slice_ptr_initialized<const LAYOUT_SIZE: usize, T>(
             ptr: *const [T],
             layout: Layout<LAYOUT_SIZE>,
@@ -289,7 +289,7 @@ macro_rules! kani_mem_init {
 
         /// Set initialization state of the slice, items of which are laid out according to the `layout` starting at address `ptr` to `value`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniSetSlicePtrInitialized"]
+        #[kanitool::fn_marker = "SetSlicePtrInitializedModel"]
         fn set_slice_ptr_initialized<const LAYOUT_SIZE: usize, T>(
             ptr: *const [T],
             layout: Layout<LAYOUT_SIZE>,
@@ -306,7 +306,7 @@ macro_rules! kani_mem_init {
 
         /// Get initialization state of the string slice, items of which are laid out according to the `layout` starting at address `ptr`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniIsStrPtrInitialized"]
+        #[kanitool::fn_marker = "IsStrPtrInitializedModel"]
         fn is_str_ptr_initialized<const LAYOUT_SIZE: usize>(
             ptr: *const str,
             layout: Layout<LAYOUT_SIZE>,
@@ -320,7 +320,7 @@ macro_rules! kani_mem_init {
 
         /// Set initialization state of the string slice, items of which are laid out according to the `layout` starting at address `ptr` to `value`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniSetStrPtrInitialized"]
+        #[kanitool::fn_marker = "SetStrPtrInitializedModel"]
         fn set_str_ptr_initialized<const LAYOUT_SIZE: usize>(
             ptr: *const str,
             layout: Layout<LAYOUT_SIZE>,
@@ -338,7 +338,7 @@ macro_rules! kani_mem_init {
         /// Copy initialization state of `size_of::<T> * num_elts` bytes from one pointer to the other. Note
         /// that in this case `LAYOUT_SIZE == size_of::<T>`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniCopyInitState"]
+        #[kanitool::fn_marker = "CopyInitStateModel"]
         fn copy_init_state<const LAYOUT_SIZE: usize, T>(
             from: *const T,
             to: *const T,
@@ -361,7 +361,7 @@ macro_rules! kani_mem_init {
         /// Copy initialization state of `size_of::<T>` bytes from one pointer to the other. Note that in
         /// this case `LAYOUT_SIZE == size_of::<T>`.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniCopyInitStateSingle"]
+        #[kanitool::fn_marker = "CopyInitStateSingleModel"]
         fn copy_init_state_single<const LAYOUT_SIZE: usize, T>(from: *const T, to: *const T) {
             copy_init_state::<LAYOUT_SIZE, T>(from, to, 1);
         }
@@ -379,7 +379,7 @@ macro_rules! kani_mem_init {
         /// Non-deterministically store information about currently tracked argument in the argument
         /// buffer.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniStoreArgument"]
+        #[kanitool::fn_marker = "StoreArgumentModel"]
         fn store_argument<const LAYOUT_SIZE: usize, T>(from: *const T, selected_argument: usize) {
             let (from_ptr, _) = from.to_raw_parts();
             let should_store: bool = super::any();
@@ -399,7 +399,7 @@ macro_rules! kani_mem_init {
         /// callee. Otherwise, mark that the argument as initialized, as it will be checked by
         /// another non-deterministic branch. Reset the argument buffer after loading from it.
         #[kanitool::disable_checks(pointer)]
-        #[rustc_diagnostic_item = "KaniLoadArgument"]
+        #[kanitool::fn_marker = "LoadArgumentModel"]
         fn load_argument<const LAYOUT_SIZE: usize, T>(to: *const T, selected_argument: usize) {
             let (to_ptr, _) = to.to_raw_parts();
             if let Some(buffer) = unsafe { ARGUMENT_BUFFER } {

From 2613708b2608ab9bd43592853491c5b7cef77ce9 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Thu, 14 Nov 2024 16:59:37 -0800
Subject: [PATCH 3/8] Change how we test std codegen

Instead of running a script at the end, use our script based suite.
Also, inject `kani_core` so we can stop special casing Kani compiler
to handle this job.
---
 scripts/kani-regression.sh                    |  3 -
 scripts/std-lib-regression.sh                 | 86 -------------------
 .../std_codegen/codegen_std.expected          |  3 +
 .../std_codegen/codegen_std.sh                | 71 +++++++++++++++
 tests/script-based-pre/std_codegen/config.yml |  4 +
 .../std_codegen/dummy/Cargo.toml              |  9 ++
 .../std_codegen/dummy/src/lib.rs              |  7 ++
 7 files changed, 94 insertions(+), 89 deletions(-)
 delete mode 100755 scripts/std-lib-regression.sh
 create mode 100644 tests/script-based-pre/std_codegen/codegen_std.expected
 create mode 100755 tests/script-based-pre/std_codegen/codegen_std.sh
 create mode 100644 tests/script-based-pre/std_codegen/config.yml
 create mode 100644 tests/script-based-pre/std_codegen/dummy/Cargo.toml
 create mode 100644 tests/script-based-pre/std_codegen/dummy/src/lib.rs

diff --git a/scripts/kani-regression.sh b/scripts/kani-regression.sh
index 88215bdf0317..1837c4cc8a82 100755
--- a/scripts/kani-regression.sh
+++ b/scripts/kani-regression.sh
@@ -80,9 +80,6 @@ for testp in "${TESTS[@]}"; do
       --quiet --no-fail-fast
 done
 
-# Check codegen for the standard library
-time "$SCRIPT_DIR"/std-lib-regression.sh
-
 # We rarely benefit from re-using build artifacts in the firecracker test,
 # and we often end up with incompatible leftover artifacts:
 # "error[E0514]: found crate `serde_derive` compiled by an incompatible version of rustc"
diff --git a/scripts/std-lib-regression.sh b/scripts/std-lib-regression.sh
deleted file mode 100755
index b010da4581f6..000000000000
--- a/scripts/std-lib-regression.sh
+++ /dev/null
@@ -1,86 +0,0 @@
-#!/usr/bin/env bash
-# Copyright Kani Contributors
-# SPDX-License-Identifier: Apache-2.0 OR MIT
-
-set -eu
-
-# Test for platform
-PLATFORM=$(uname -sp)
-if [[ $PLATFORM == "Linux x86_64" ]]
-then
-  TARGET="x86_64-unknown-linux-gnu"
-  # 'env' necessary to avoid bash built-in 'time'
-  WRAPPER="env time -v"
-elif [[ $PLATFORM == "Darwin i386" ]]
-then
-  # Temporarily disabled (in CI) to keeps CI times down
-  # See https://github.com/model-checking/kani/issues/1578
-  exit 0
-
-  TARGET="x86_64-apple-darwin"
-  # mac 'time' doesn't have -v
-  WRAPPER=""
-elif [[ $PLATFORM == "Darwin arm" ]]
-then
-  TARGET="aarch64-apple-darwin"
-  # mac 'time' doesn't have -v
-  WRAPPER=""
-else
-  echo
-  echo "Std-Lib codegen regression only works on Linux or OSX x86 platforms, skipping..."
-  echo
-  exit 0
-fi
-
-# Get Kani root
-SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
-KANI_DIR=$(dirname "$SCRIPT_DIR")
-
-echo
-echo "Starting Kani codegen for the Rust standard library..."
-echo
-
-cd /tmp
-if [ -d std_lib_test ]
-then
-    rm -rf std_lib_test
-fi
-cargo new std_lib_test --lib
-cd std_lib_test
-
-# Add some content to the rust file including an std function that is non-generic.
-echo '
-pub fn main() {
-    assert!("2021".parse::<u32>().unwrap() == 2021);
-}
-' > src/lib.rs
-
-# Until we add support to this via our bundle, rebuild the kani library too.
-echo "
-kani = {path=\"${KANI_DIR}/library/kani\"}
-" >> Cargo.toml
-
-# Use same nightly toolchain used to build Kani
-cp ${KANI_DIR}/rust-toolchain.toml .
-
-echo "Starting cargo build with Kani"
-export RUST_BACKTRACE=1
-export RUSTC_LOG=error
-
-RUST_FLAGS=(
-    "--kani-compiler"
-    "-Cpanic=abort"
-    "-Zalways-encode-mir"
-    "-Cllvm-args=--backend=cprover"
-    "-Cllvm-args=--ignore-global-asm"
-    "-Cllvm-args=--reachability=pub_fns"
-    "-Cllvm-args=--build-std"
-)
-export RUSTFLAGS="${RUST_FLAGS[@]}"
-export RUSTC="$KANI_DIR/target/kani/bin/kani-compiler"
-# Compile rust to iRep
-$WRAPPER cargo build --verbose -Z build-std --lib --target $TARGET
-
-echo
-echo "Finished Kani codegen for the Rust standard library successfully..."
-echo
diff --git a/tests/script-based-pre/std_codegen/codegen_std.expected b/tests/script-based-pre/std_codegen/codegen_std.expected
new file mode 100644
index 000000000000..ae155f0ecf44
--- /dev/null
+++ b/tests/script-based-pre/std_codegen/codegen_std.expected
@@ -0,0 +1,3 @@
+[TEST] Copy standard library from the current toolchain
+[TEST] Modify library
+------ Build succeeded -------
diff --git a/tests/script-based-pre/std_codegen/codegen_std.sh b/tests/script-based-pre/std_codegen/codegen_std.sh
new file mode 100755
index 000000000000..c78b2a3c5f45
--- /dev/null
+++ b/tests/script-based-pre/std_codegen/codegen_std.sh
@@ -0,0 +1,71 @@
+#!/usr/bin/env bash
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+
+# Test that we can codegen the entire standard library.
+# 1. Make a copy of the rust standard library.
+# 2. Add a few Kani definitions to it
+# 3. Run Kani compiler
+
+set -e
+set -u
+
+KANI_DIR=$(git rev-parse --show-toplevel)
+TMP_DIR="tmp_dir"
+
+rm -rf ${TMP_DIR}
+mkdir ${TMP_DIR}
+
+cp -r dummy ${TMP_DIR}
+
+# Create a custom standard library.
+echo "[TEST] Copy standard library from the current toolchain"
+SYSROOT=$(rustc --print sysroot)
+STD_PATH="${SYSROOT}/lib/rustlib/src/rust/library"
+cp -r "${STD_PATH}" "${TMP_DIR}"
+
+# Insert Kani definitions.
+CORE_CODE='
+#[cfg(kani)]
+kani_core::kani_lib!(core);
+'
+
+echo "[TEST] Modify library"
+echo "${CORE_CODE}" >> ${TMP_DIR}/library/core/src/lib.rs
+
+# Note: Prepending with sed doesn't work on MacOs the same way it does in linux.
+# sed -i '1s/^/#![cfg_attr(kani, feature(kani))]\n/' ${TMP_DIR}/library/std/src/lib.rs
+cp ${TMP_DIR}/library/std/src/lib.rs ${TMP_DIR}/std_lib.rs
+echo '#![cfg_attr(kani, feature(kani))]' > ${TMP_DIR}/library/std/src/lib.rs
+cat ${TMP_DIR}/std_lib.rs >> ${TMP_DIR}/library/std/src/lib.rs
+
+export RUST_BACKTRACE=1
+export RUSTC_LOG=error
+export __CARGO_TESTS_ONLY_SRC_ROOT=$(readlink -f ${TMP_DIR})/library
+RUST_FLAGS=(
+    "--kani-compiler"
+    "-Cpanic=abort"
+    "-Zalways-encode-mir"
+    "-Cllvm-args=--backend=cprover"
+    "-Cllvm-args=--ignore-global-asm"
+    "-Cllvm-args=--reachability=pub_fns"
+    "-L${KANI_DIR}/target/kani/no_core/lib"
+    "--extern=kani_core"
+    "--cfg=kani"
+    "-Zcrate-attr=feature(register_tool)"
+    "-Zcrate-attr=register_tool(kanitool)"
+)
+export RUSTFLAGS="${RUST_FLAGS[@]}"
+export RUSTC="$KANI_DIR/target/kani/bin/kani-compiler"
+export KANI_LOGS=kani_compiler::kani_middle=debug
+TARGET=$(rustc -vV | awk '/^host/ { print $2 }')
+
+pushd ${TMP_DIR}/dummy > /dev/null
+# Compile the standard library to iRep
+cargo build --verbose -Z build-std --lib --target ${TARGET}
+popd > /dev/null
+
+echo "------ Build succeeded -------"
+
+# Cleanup
+rm -r ${TMP_DIR}
diff --git a/tests/script-based-pre/std_codegen/config.yml b/tests/script-based-pre/std_codegen/config.yml
new file mode 100644
index 000000000000..c5f720824c6e
--- /dev/null
+++ b/tests/script-based-pre/std_codegen/config.yml
@@ -0,0 +1,4 @@
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+script: codegen_std.sh
+expected: codegen_std.expected
diff --git a/tests/script-based-pre/std_codegen/dummy/Cargo.toml b/tests/script-based-pre/std_codegen/dummy/Cargo.toml
new file mode 100644
index 000000000000..a6e852d46050
--- /dev/null
+++ b/tests/script-based-pre/std_codegen/dummy/Cargo.toml
@@ -0,0 +1,9 @@
+# Copyright Kani Contributors
+# SPDX-License-Identifier: Apache-2.0 OR MIT
+
+[package]
+name = "dummy"
+version = "0.1.0"
+edition = "2021"
+
+[workspace]
diff --git a/tests/script-based-pre/std_codegen/dummy/src/lib.rs b/tests/script-based-pre/std_codegen/dummy/src/lib.rs
new file mode 100644
index 000000000000..9349fee3a608
--- /dev/null
+++ b/tests/script-based-pre/std_codegen/dummy/src/lib.rs
@@ -0,0 +1,7 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//! Just a dummy file. We are interested in the standard library only
+
+pub fn void() {
+    todo!()
+}

From 7d45dd910958a509e93112769c770a0306dca738 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Thu, 14 Nov 2024 21:59:07 -0800
Subject: [PATCH 4/8] Revert marker for shadow memory for uninit checks

Handle this later since we'll have to create a function to handle
static objects.
---
 library/kani_core/src/mem_init.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/library/kani_core/src/mem_init.rs b/library/kani_core/src/mem_init.rs
index ecafe5d8e861..935f14d71c5d 100644
--- a/library/kani_core/src/mem_init.rs
+++ b/library/kani_core/src/mem_init.rs
@@ -23,7 +23,7 @@
 macro_rules! kani_mem_init {
     ($core:path) => {
         /// Global object for tracking memory initialization state.
-        #[kanitool::fn_marker = "MemoryInitializationStateModel"]
+        #[rustc_diagnostic_item = "KaniMemoryInitializationState"]
         static mut MEM_INIT_STATE: MemoryInitializationState = MemoryInitializationState::new();
 
         /// Global object for tracking union initialization state across function boundaries.

From 6f60d4fd01ae6fa6d3a7e0573d0271a4c39be4f1 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 15 Nov 2024 11:21:58 -0800
Subject: [PATCH 5/8] Apply suggestions from code review

Co-authored-by: Carolyn Zech <cmzech@amazon.com>
---
 .../src/codegen_cprover_gotoc/compiler_interface.rs    |  4 ++--
 kani-compiler/src/kani_middle/kani_functions.rs        | 10 +++++-----
 .../src/kani_middle/transform/kani_intrinsics.rs       |  2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
index d959f327e748..16391216cd8b 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
@@ -237,8 +237,8 @@ impl CodegenBackend for GotocCodegenBackend {
         let ret_val = rustc_internal::run(tcx, || {
             super::utils::init();
 
-            // Any changes to queries from this point on is just related to caching information
-            // for efficiency purpose that should not outlive the stable-mir `run` scope.
+            // Any changes to queries from this point on is just related to caching;
+            // the cached information will not outlive the stable-mir `run` scope.
             let queries = self.queries.lock().unwrap().clone();
 
             check_target(tcx.sess);
diff --git a/kani-compiler/src/kani_middle/kani_functions.rs b/kani-compiler/src/kani_middle/kani_functions.rs
index 23d127368ef6..56098b76f36e 100644
--- a/kani-compiler/src/kani_middle/kani_functions.rs
+++ b/kani-compiler/src/kani_middle/kani_functions.rs
@@ -4,7 +4,7 @@
 //!
 //! There are three types of functions today:
 //!    1. Kani intrinsics: These are functions whose body is generated during
-//!       compilation time. Their body usually require some extra knowledge about the given types
+//!       compilation time. Their body usually requires some extra knowledge about the given types
 //!       that's only available during compilation.
 //!    2. Kani models: These are functions that model a specific behavior but that cannot be used
 //!       directly by the user. For example, retrieving information about memory initialization.
@@ -148,8 +148,8 @@ impl TryFrom<FnDef> for KaniFunction {
 
     fn try_from(def: FnDef) -> Result<Self, Self::Error> {
         let value = attributes::fn_marker(def).ok_or(())?;
-        if let Ok(intrisic) = KaniIntrinsic::from_str(&value) {
-            Ok(intrisic.into())
+        if let Ok(intrinsic) = KaniIntrinsic::from_str(&value) {
+            Ok(intrinsic.into())
         } else if let Ok(model) = KaniModel::from_str(&value) {
             Ok(model.into())
         } else if let Ok(hook) = KaniHook::from_str(&value) {
@@ -165,8 +165,8 @@ impl TryFrom<Instance> for KaniFunction {
 
     fn try_from(instance: Instance) -> Result<Self, Self::Error> {
         let value = attributes::fn_marker(instance.def).ok_or(())?;
-        if let Ok(intrisic) = KaniIntrinsic::from_str(&value) {
-            Ok(intrisic.into())
+        if let Ok(intrinsic) = KaniIntrinsic::from_str(&value) {
+            Ok(intrinsic.into())
         } else if let Ok(model) = KaniModel::from_str(&value) {
             Ok(model.into())
         } else if let Ok(hook) = KaniHook::from_str(&value) {
diff --git a/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs b/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
index 533c27445b84..9c541d1e1ee4 100644
--- a/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
+++ b/kani-compiler/src/kani_middle/transform/kani_intrinsics.rs
@@ -39,7 +39,7 @@ pub struct IntrinsicGeneratorPass {
     check_type: CheckType,
     /// Used to cache FnDef lookups for models and Kani intrinsics.
     kani_defs: HashMap<KaniFunction, FnDef>,
-    /// Used to enable intrinsics depending on the flags passed.
+    /// Whether the user enabled uninitialized memory checks when they invoked Kani.
     enable_uninit: bool,
 }
 

From 194776da0999951f4d101ae3d338f0d2b2742511 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 15 Nov 2024 12:05:02 -0800
Subject: [PATCH 6/8] Improve documentation and naming schema

---
 .../compiler_interface.rs                     |  6 ++-
 .../codegen_cprover_gotoc/overrides/hooks.rs  | 41 +++++++++++--------
 kani-compiler/src/kani_middle/attributes.rs   |  1 +
 .../src/kani_middle/kani_functions.rs         | 24 +++++++----
 kani-compiler/src/kani_queries/mod.rs         | 17 ++++++--
 5 files changed, 57 insertions(+), 32 deletions(-)

diff --git a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
index 16391216cd8b..13f0eddfab99 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
@@ -237,8 +237,10 @@ impl CodegenBackend for GotocCodegenBackend {
         let ret_val = rustc_internal::run(tcx, || {
             super::utils::init();
 
-            // Any changes to queries from this point on is just related to caching;
-            // the cached information will not outlive the stable-mir `run` scope.
+            // Any changes to queries from this point on is just related to caching information
+            // needed for generating code to the given crate.
+            // The cached information mut not outlive the stable-mir `run` scope.
+            // See [QueryDb::kani_functions] for more information.
             let queries = self.queries.lock().unwrap().clone();
 
             check_target(tcx.sess);
diff --git a/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs b/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs
index 3570f38188ad..96d7f4e42100 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs
@@ -48,9 +48,12 @@ pub trait GotocHook {
 /// <https://github.com/model-checking/kani/blob/main/rfc/src/rfcs/0003-cover-statement.md>
 /// for more details.
 struct Cover;
+
+const UNEXPECTED_CALL: &str = "Hooks from kani library handled as a map";
+
 impl GotocHook for Cover {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -85,7 +88,7 @@ impl GotocHook for Cover {
 struct Assume;
 impl GotocHook for Assume {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -109,7 +112,7 @@ impl GotocHook for Assume {
 struct Assert;
 impl GotocHook for Assert {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -149,7 +152,7 @@ impl GotocHook for Assert {
 struct Check;
 impl GotocHook for Check {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -185,7 +188,7 @@ struct Nondet;
 
 impl GotocHook for Nondet {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -248,7 +251,7 @@ impl GotocHook for Panic {
 struct IsAllocated;
 impl GotocHook for IsAllocated {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -285,7 +288,7 @@ impl GotocHook for IsAllocated {
 struct PointerObject;
 impl GotocHook for PointerObject {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -321,7 +324,7 @@ impl GotocHook for PointerObject {
 struct PointerOffset;
 impl GotocHook for PointerOffset {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -467,7 +470,7 @@ struct UntrackedDeref;
 
 impl GotocHook for UntrackedDeref {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -515,7 +518,7 @@ struct InitContracts;
 /// ```
 impl GotocHook for InitContracts {
     fn hook_applies(&self, _tcx: TyCtxt, _instance: Instance) -> bool {
-        unreachable!("Hooks from kani library handled as a map")
+        unreachable!("{UNEXPECTED_CALL}")
     }
 
     fn handle(
@@ -616,7 +619,7 @@ impl GotocHook for LoopInvariantRegister {
 }
 
 pub fn fn_hooks() -> GotocHooks {
-    let kani_hooks: [(KaniHook, Rc<dyn GotocHook>); 11] = [
+    let kani_lib_hooks: [(KaniHook, Rc<dyn GotocHook>); 11] = [
         (KaniHook::Assert, Rc::new(Assert)),
         (KaniHook::Assume, Rc::new(Assume)),
         (KaniHook::Panic, Rc::new(Panic)),
@@ -630,27 +633,31 @@ pub fn fn_hooks() -> GotocHooks {
         (KaniHook::InitContracts, Rc::new(InitContracts)),
     ];
     GotocHooks {
-        rustc_hooks: vec![
+        kani_lib_hooks: HashMap::from(kani_lib_hooks),
+        other_hooks: vec![
             Rc::new(Panic),
             Rc::new(RustAlloc),
             Rc::new(MemCmp),
             Rc::new(LoopInvariantRegister),
         ],
-        kani_hooks: HashMap::from(kani_hooks),
     }
 }
 
 pub struct GotocHooks {
-    rustc_hooks: Vec<Rc<dyn GotocHook>>,
-    kani_hooks: HashMap<KaniHook, Rc<dyn GotocHook>>,
+    /// Match functions that are unique and defined in the Kani library, which we can prefetch
+    /// using `KaniFunctions`.
+    kani_lib_hooks: HashMap<KaniHook, Rc<dyn GotocHook>>,
+    /// Match functions that are not defined in the Kani library, which we cannot prefetch
+    /// beforehand.
+    other_hooks: Vec<Rc<dyn GotocHook>>,
 }
 
 impl GotocHooks {
     pub fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> Option<Rc<dyn GotocHook>> {
         if let Ok(KaniFunction::Hook(hook)) = KaniFunction::try_from(instance) {
-            Some(self.kani_hooks[&hook].clone())
+            Some(self.kani_lib_hooks[&hook].clone())
         } else {
-            for h in &self.rustc_hooks {
+            for h in &self.other_hooks {
                 if h.hook_applies(tcx, instance) {
                     return Some(h.clone());
                 }
diff --git a/kani-compiler/src/kani_middle/attributes.rs b/kani-compiler/src/kani_middle/attributes.rs
index 86d89d7d86cd..3ee695e6f481 100644
--- a/kani-compiler/src/kani_middle/attributes.rs
+++ b/kani-compiler/src/kani_middle/attributes.rs
@@ -1067,6 +1067,7 @@ fn pretty_type_path(path: &TypePath) -> String {
     }
 }
 
+/// Retrieve the value of the `fn_marker` attribute for the given definition if it has one.
 pub(crate) fn fn_marker<T: CrateDef>(def: T) -> Option<String> {
     let fn_marker: [SymbolStable; 2] = ["kanitool".into(), "fn_marker".into()];
     let marker = def.attrs_by_path(&fn_marker).pop()?;
diff --git a/kani-compiler/src/kani_middle/kani_functions.rs b/kani-compiler/src/kani_middle/kani_functions.rs
index 56098b76f36e..10441fd00f1e 100644
--- a/kani-compiler/src/kani_middle/kani_functions.rs
+++ b/kani-compiler/src/kani_middle/kani_functions.rs
@@ -3,16 +3,16 @@
 //! This module contains code for handling special functions from the Kani library.
 //!
 //! There are three types of functions today:
-//!    1. Kani intrinsics: These are functions whose body is generated during
-//!       compilation time. Their body usually requires some extra knowledge about the given types
+//!    1. Kani intrinsics: These are functions whose MIR body is generated during compilation time.
+//!       Their body usually requires some extra knowledge about the given types
 //!       that's only available during compilation.
-//!    2. Kani models: These are functions that model a specific behavior but that cannot be used
-//!       directly by the user. For example, retrieving information about memory initialization.
-//!       Kani compiler determines when and where to use these models, but the logic itself is
-//!       encoded in Rust.
-//!    3. Kani hooks: These are similar to Kani intrinsics but their code generation depends
-//!       on some backend specific logic. From a Kani library perspective, there is no difference
-//!       between hooks and intrinsics.
+//!    2. Kani models: These are functions that model a specific behavior that is of interest of
+//!       the compiler. The body of the model is kept as is, and the compiler may invoke the model
+//!       to perform some operation, such as, retrieving information about memory initialization.
+//!    3. Kani hooks: These are similar to Kani intrinsics in a sense that their body will be
+//!       generated by the compiler. However, their body exposes some backend specific logic, thus,
+//!       their bodies are generated as part of the codegen stage.
+//!       From a Kani library perspective, there is no difference between hooks and intrinsics.
 //!       This is a compiler implementation detail, but for simplicity we use the "Hook" suffix
 //!       in their names.
 //!
@@ -214,6 +214,12 @@ pub fn find_kani_functions() -> HashMap<KaniFunction, FnDef> {
 }
 
 /// Ensure we have the valid definitions.
+///
+/// This is a utility function used to ensure that we have found all the expected functions, as well
+/// as to ensure that the cached IDs are up-to-date.
+///
+/// The second condition is to ensure the cached StableMIR `FnDef` objects are still valid, i.e.:,
+/// that we are not storing them past the StableMIR `run` context.
 pub fn validate_kani_functions(kani_funcs: &HashMap<KaniFunction, FnDef>) {
     let mut missing = 0u8;
     for func in KaniIntrinsic::iter()
diff --git a/kani-compiler/src/kani_queries/mod.rs b/kani-compiler/src/kani_queries/mod.rs
index 83b636caacf1..5eb94496b8e8 100644
--- a/kani-compiler/src/kani_queries/mod.rs
+++ b/kani-compiler/src/kani_queries/mod.rs
@@ -38,12 +38,21 @@ impl QueryDb {
     /// For `kani_core`, those definitions live in the `core` library.
     ///
     /// We cache these definitions to avoid doing the lookup every time it is needed.
-    /// The cache should be invalidated if the `stable_mir` context changes.
-    /// Since that doesn't occur today, we just run a sanity check to ensure all definitions
-    /// are still correct, and abort otherwise.
+    /// The cache should not be used after the `stable_mir` context ends.
+    /// For example, in the goto backend, we run the entire crate codegen under the same StableMIR
+    /// context, which is defined by the scope of the StableMIR `run` callback.
+    /// See the `codegen_crate` function in [crate::codegen_cprover_gotoc::compiler_interface].
+    /// It is OK to set the cache and use it inside the callback scope, however, the cache should
+    /// not be accessible after that.
+    ///
+    /// For that, users should create a new QueryDb that does not outlive the callback scope.
+    ///
+    /// To ensure we don't accidentally use the cache outside of the callback context, we run a
+    /// sanity check if we are reusing the cache.
     pub fn kani_functions(&self) -> &HashMap<KaniFunction, FnDef> {
         if let Some(kani_functions) = self.kani_functions.get() {
-            // Sanity check the values stored in case someone misused this API.
+            // Sanity check the values stored to ensure the cache is being within the StableMIR
+            // context used to populate the cache.
             validate_kani_functions(kani_functions);
             kani_functions
         } else {

From 9b8881026ea825dbc71047fbe38202cda5303d86 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 15 Nov 2024 15:57:35 -0800
Subject: [PATCH 7/8] Apply suggestions from code review

Co-authored-by: Carolyn Zech <cmzech@amazon.com>
---
 kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs | 2 +-
 kani-compiler/src/kani_middle/transform/body.rs               | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
index 13f0eddfab99..6b22b246718a 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs
@@ -239,7 +239,7 @@ impl CodegenBackend for GotocCodegenBackend {
 
             // Any changes to queries from this point on is just related to caching information
             // needed for generating code to the given crate.
-            // The cached information mut not outlive the stable-mir `run` scope.
+            // The cached information must not outlive the stable-mir `run` scope.
             // See [QueryDb::kani_functions] for more information.
             let queries = self.queries.lock().unwrap().clone();
 
diff --git a/kani-compiler/src/kani_middle/transform/body.rs b/kani-compiler/src/kani_middle/transform/body.rs
index 0bdd4a1603b8..911ccf17e55f 100644
--- a/kani-compiler/src/kani_middle/transform/body.rs
+++ b/kani-compiler/src/kani_middle/transform/body.rs
@@ -436,7 +436,7 @@ impl MutableBody {
     }
 }
 
-// TODO: Remove-me
+// TODO: Remove this enum, since we now only support kani's assert.
 #[derive(Clone, Debug)]
 pub enum CheckType {
     /// This is used by default when the `kani` crate is available.

From a466e725853160a1922d1bfa28039a2c09d58ffc Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 15 Nov 2024 18:08:27 -0800
Subject: [PATCH 8/8] Fix firecracker codegen regression test

---
 scripts/codegen-firecracker.sh | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/scripts/codegen-firecracker.sh b/scripts/codegen-firecracker.sh
index bffb27023412..2deb4bc02641 100755
--- a/scripts/codegen-firecracker.sh
+++ b/scripts/codegen-firecracker.sh
@@ -20,8 +20,7 @@ else
 fi
 
 # Get Kani root
-SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
-KANI_DIR=$SCRIPT_DIR/..
+KANI_DIR=$(git rev-parse --show-toplevel)
 
 echo
 echo "Starting Firecracker codegen regression..."
@@ -40,12 +39,16 @@ export RUST_BACKTRACE=1
 # Compile rust to iRep
 RUST_FLAGS=(
     "--kani-compiler"
-    "-Cpanic=abort"
-    "-Zalways-encode-mir"
-    "-Cllvm-args=--backend=cprover"
-    "-Cllvm-args=--ignore-global-asm"
-    "-Cllvm-args=--reachability=pub_fns"
-    "--sysroot=${KANI_DIR}/target/kani"
+    "-Cllvm-args=--assertion-reach-checks"
+    "-Zunstable-options"
+    "--sysroot"
+    "${KANI_DIR}/target/kani"
+    "-L"
+    "${KANI_DIR}/target/kani/lib"
+    "--extern"
+    "kani"
+    "--extern"
+    "noprelude:std=${KANI_DIR}/target/kani/lib/libstd.rlib"
 )
 export RUSTFLAGS="${RUST_FLAGS[@]}"
 export RUSTC="$KANI_DIR/target/kani/bin/kani-compiler"
