krabcake: Successfully send intrinsics::asume data to Valgrind

This patch adds support to send the `bool` in `std::intrinsics::assume`
to Valgrind by inserting a Valgrind Client Request using a MirPass.

Author: bryangarza

compiler/rustc_hir/src/lang_items.rs

@@ -147,6 +147,8 @@ pub fn extract(attrs: &[ast::Attribute]) -> Option<(Symbol, Span)> {
 
 language_item_table! {
 //  Variant name,            Name,                     Getter method name,         Target                  Generic requirements;
+    KrabcakeRequest,         sym::KrabcakeRequest,     krabcake_request,           Target::Enum,           GenericRequirement::None;
+    InsertKrabcakeRequest,   sym::insert_krabcake_request, insert_krabcake_request_fn, Target::Fn,         GenericRequirement::None;
     Sized,                   sym::sized,               sized_trait,                Target::Trait,          GenericRequirement::Exact(0);
     Unsize,                  sym::unsize,              unsize_trait,               Target::Trait,          GenericRequirement::Minimum(1);
     /// Trait injected by `#[derive(PartialEq)]`, (i.e. "Partial EQ").

compiler/rustc_interface/src/tests.rs

@@ -754,6 +754,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(inline_mir_hint_threshold, Some(123));
     tracked!(inline_mir_threshold, Some(123));
     tracked!(instrument_coverage, Some(InstrumentCoverage::All));
+    tracked!(instrument_krabcake, false);
     tracked!(instrument_mcount, true);
     tracked!(instrument_xray, Some(InstrumentXRay::default()));
     tracked!(link_directives, false);

compiler/rustc_middle/src/mir/patch.rs

@@ -179,4 +179,16 @@ impl<'tcx> MirPatch<'tcx> {
         };
         Self::source_info_for_index(data, loc)
     }
+
+    pub fn terminator_for_location(
+        &self,
+        body: &Body<'tcx>,
+        loc: Location,
+    ) -> Option<Terminator<'tcx>> {
+        let data = match loc.block.index().checked_sub(body.basic_blocks.len()) {
+            Some(new) => &self.new_blocks[new],
+            None => &body[loc.block],
+        };
+        data.terminator.clone()
+    }
 }

compiler/rustc_middle/src/ty/layout.rs

@@ -680,7 +680,11 @@ where
                 | ty::GeneratorWitnessMIR(..)
                 | ty::Foreign(..)
                 | ty::Dynamic(_, _, ty::Dyn) => {
-                    bug!("TyAndLayout::field({:?}): not applicable", this)
+                    bug!(
+                        "TyAndLayout::field({:?}): not applicable for kind {:?}",
+                        this,
+                        this.ty.kind()
+                    )
                 }
 
                 // Potentially-fat pointers.

compiler/rustc_mir_transform/src/lib.rs

@@ -98,6 +98,7 @@ mod simplify_comparison_integral;
 mod sroa;
 mod uninhabited_enum_branching;
 mod unreachable_prop;
+mod valgrind_client_request;
 
 use rustc_const_eval::transform::check_consts::{self, ConstCx};
 use rustc_const_eval::transform::promote_consts;
@@ -588,6 +589,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &large_enums::EnumSizeOpt { discrepancy: 128 },
             // Some cleanup necessary at least for LLVM and potentially other codegen backends.
             &add_call_guards::CriticalCallEdges,
+            &valgrind_client_request::ValgrindClientRequest,
             // Dump the end result for testing and debugging purposes.
             &dump_mir::Marker("PreCodegen"),
         ],

compiler/rustc_mir_transform/src/valgrind_client_request.rs

@@ -0,0 +1,112 @@
+use crate::MirPass;
+
+use rustc_index::vec::Idx;
+use rustc_middle::mir::patch::MirPatch;
+use rustc_middle::mir::{
+    AggregateKind, BasicBlockData, Body, Constant, ConstantKind, Location, NonDivergingIntrinsic,
+    Operand, Place, Rvalue, StatementKind, TerminatorKind,
+};
+use rustc_middle::ty::InternalSubsts;
+use rustc_middle::ty::TyCtxt;
+use rustc_target::abi::VariantIdx;
+
+pub struct ValgrindClientRequest;
+
+impl<'tcx> MirPass<'tcx> for ValgrindClientRequest {
+    #[instrument(skip(self, tcx, body))]
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        if !tcx.sess.opts.unstable_opts.instrument_krabcake {
+            info!("Not instrumenting for krabcake");
+            return;
+        }
+        info!("Instrumenting for krabcake now...");
+        let mut patch = MirPatch::new(body);
+
+        for (block_index, block_data) in body.basic_blocks.iter_enumerated() {
+            for (stmt_index, stmt) in block_data.statements.iter().enumerate() {
+                match &stmt.kind {
+                    StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(operand)) => {
+                        let loc = Location { block: block_index, statement_index: stmt_index };
+                        info!("Found assume intrinsic (operand={operand:?}. At {loc:?}");
+                        patch = call_ikr(tcx, patch, body, loc, operand);
+                    }
+                    _ => (),
+                }
+            }
+        }
+
+        patch.apply(body);
+    }
+}
+
+fn call_ikr<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    mut patch: MirPatch<'tcx>,
+    body: &Body<'tcx>,
+    loc: Location,
+    _operand: &Operand<'tcx>,
+) -> MirPatch<'tcx> {
+    let span = patch.source_info_for_location(body, loc).span;
+
+    let op = |flag: bool| {
+        Operand::Constant(Box::new(Constant {
+            span,
+            user_ty: None,
+            literal: ConstantKind::from_bool(tcx, flag),
+        }))
+    };
+
+    let (place, rvalue) = {
+        let krabcake_req_did = tcx.lang_items().krabcake_request().unwrap();
+        let krabcake_req_substs = InternalSubsts::identity_for_item(tcx, krabcake_req_did);
+        let krabcake_req_def = tcx.adt_def(krabcake_req_did);
+        let krabcake_req_ty = tcx.mk_adt(krabcake_req_def, krabcake_req_substs);
+        let rvalue = Rvalue::Aggregate(
+            Box::new(AggregateKind::Adt(
+                krabcake_req_did,
+                VariantIdx::new(0),
+                &krabcake_req_substs,
+                None,
+                None,
+            )),
+            vec![op(true)],
+        );
+        let temp = patch.new_temp(krabcake_req_ty, span);
+        let place = Place::from(temp);
+        (place, rvalue)
+    };
+
+    patch.add_assign(loc, place, rvalue);
+
+    let krabcake_req_operand = Operand::Copy(place);
+
+    let orig_terminator = patch.terminator_for_location(body, loc);
+    let ikr_did = tcx.lang_items().insert_krabcake_request_fn().unwrap();
+    let ikr_substs = InternalSubsts::identity_for_item(tcx, ikr_did);
+    let ikr_ty = tcx.mk_fn_def(ikr_did, ikr_substs);
+
+    let func = Operand::Constant(Box::new(Constant {
+        span,
+        user_ty: None,
+        literal: ConstantKind::zero_sized(ikr_ty),
+    }));
+    let storage = patch.new_temp(tcx.mk_mut_ptr(tcx.types.unit), span);
+    let storage = Place::from(storage);
+    let fn_call_terminator_kind = TerminatorKind::Call {
+        func,
+        args: vec![krabcake_req_operand],
+        destination: storage,
+        target: Some(loc.block + 1),
+        cleanup: None,
+        from_hir_call: false,
+        fn_span: span,
+    };
+
+    patch.patch_terminator(loc.block, fn_call_terminator_kind);
+
+    let new_bb =
+        BasicBlockData { statements: vec![], terminator: orig_terminator, is_cleanup: false };
+
+    patch.new_block(new_bb);
+    patch
+}

compiler/rustc_session/src/options.rs

@@ -1473,6 +1473,8 @@ options! {
         `=except-unused-generics`
         `=except-unused-functions`
         `=off` (default)"),
+    instrument_krabcake: bool = (false, parse_bool, [TRACKED],
+        "insert Valgrind client requests to communicate with Krabcake"),
     instrument_mcount: bool = (false, parse_bool, [TRACKED],
         "insert function instrument code for mcount-based tracing (default: no)"),
     instrument_xray: Option<InstrumentXRay> = (None, parse_instrument_xray, [TRACKED],

compiler/rustc_span/src/symbol.rs

@@ -216,6 +216,7 @@ symbols! {
         ItemContext,
         Iterator,
         IteratorItem,
+        KrabcakeRequest,
         Layout,
         Left,
         LinkedList,
@@ -824,6 +825,7 @@ symbols! {
         inline_const,
         inline_const_pat,
         inout,
+        insert_krabcake_request,
         instruction_set,
         integer_: "integer",
         integral,

library/core/src/krabcake.rs

@@ -0,0 +1,95 @@
+use core::arch::asm;
+
+mod vg_c_compatible {
+    const fn vg_userreq_tool_base(a: u64, b: u64) -> u64 {
+        ((a) & 0xff) << 24 | ((b) & 0xff) << 16
+    }
+
+    /// blah
+    #[allow(dead_code)] // For currently unused variants
+    #[derive(Debug, Clone, Copy)]
+    #[repr(u64)]
+    pub(super) enum ValgrindClientRequestCode {
+        /// blah
+        BorrowMut = vg_userreq_tool_base('K' as u64, 'C' as u64),
+        /// blah
+        BorrowShr,
+        /// blah
+        AsRaw,
+        /// blah
+        AsBorrowMut,
+        /// blah
+        AsBorrowShr,
+        /// blah
+        RetagFnPrologue,
+        /// blah
+        RetagAssign,
+        /// blah
+        RetagRaw,
+        /// blah
+        IntrinsicsAssume,
+        /// blah
+        KrabcakeRecordOverlapError = vg_userreq_tool_base('K' as u64, 'C' as u64) + 256,
+    }
+
+    /// Ultimately, this should map to a [u64; 6], which is what Valgrind expects
+    #[derive(Debug, Clone, Copy)]
+    #[repr(C)]
+    struct ValgrindClientRequest<T>
+    where
+        T: Into<u64>,
+    {
+        req_code: ValgrindClientRequestCode,
+        args: [T; 5],
+    }
+}
+
+/// blah
+/// #[doc(hidden)]
+// There is no issue created for krabcake yet so just using a placeholder
+#[unstable(feature = "instrument_krabcake", issue = "1")]
+#[cfg_attr(not(bootstrap), lang = "KrabcakeRequest")]
+#[derive(Debug, Clone, Copy)]
+#[repr(C)]
+pub enum KrabcakeRequest {
+    /// blah
+    IntrinsicsAssume {
+        /// blah
+        flag: bool,
+    },
+}
+
+#[doc(hidden)]
+// There is no issue created for krabcake yet so just using a placeholder
+#[unstable(feature = "instrument_krabcake", issue = "1")]
+#[inline(never)]
+#[cfg_attr(not(bootstrap), lang = "insert_krabcake_request")]
+pub fn insert_krabcake_request(req: KrabcakeRequest) {
+    // `res` not used for IntrinsicsAssume request
+    let mut res = false as u64;
+    let args: [u64; 6] = match req {
+        KrabcakeRequest::IntrinsicsAssume { flag } => [
+            vg_c_compatible::ValgrindClientRequestCode::IntrinsicsAssume as u64,
+            flag.into(),
+            0,
+            0,
+            0,
+            0,
+        ],
+    };
+
+    // SAFETY:
+    // This is a valid safety comment
+    unsafe {
+        asm!(
+            "rol rdi, 3",
+            "rol rdi, 13",
+            "rol rdi, 61",
+            "rol rdi, 51",
+            "xchg rbx, rbx",
+            inout("di") res,
+            in("ax") &args,
+        );
+    }
+    let _ = res;
+}

library/core/src/lib.rs

@@ -376,6 +376,11 @@ pub mod task;
 #[allow(missing_docs)]
 pub mod alloc;
 
+/// blah
+// There is no issue created for krabcake yet so just using a placeholder
+#[unstable(feature = "instrument_krabcake", issue = "1")]
+pub mod krabcake;
+
 // note: does not need to be public
 mod bool;
 mod tuple;