Tmr triplication diff

compiler/rustc_attr_parsing/src/attributes/rad_protected.rs

@@ -8,7 +8,7 @@ use crate::target_checking::{ALL_TARGETS, AllowedTargets};
 
 pub(crate) struct RadProtectedParser;
 impl<S: Stage> NoArgsAttributeParser<S> for RadProtectedParser {
-    const PATH: &[Symbol] = &[sym::rad_protected];
+    const PATH: &[Symbol] = &[sym::rad_protected_mir];
     const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Ignore;
     const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(ALL_TARGETS);
 

compiler/rustc_builtin_macros/src/lib.rs

@@ -45,6 +45,7 @@ mod global_allocator;
 mod iter;
 mod log_syntax;
 mod pattern_type;
+mod rad_protected;
 mod source_util;
 mod test;
 mod trace_macros;
@@ -118,6 +119,7 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
         eii_declaration: eii::eii_declaration,
         eii_shared_macro: eii::eii_shared_macro,
         global_allocator: global_allocator::expand,
+        rad_protected: rad_protected::triplicate,
         test: test::expand_test,
         test_case: test::expand_test_case,
         unsafe_eii: eii::unsafe_eii,

compiler/rustc_builtin_macros/src/rad_protected/mod.rs

@@ -0,0 +1,4 @@
+mod triplicate;
+mod parse_attr_opts;
+
+pub(crate) use triplicate::triplicate;

compiler/rustc_builtin_macros/src/rad_protected/parse_attr_opts.rs

@@ -0,0 +1,43 @@
+use rustc_ast as ast;
+use rustc_ast::MetaItemInner;
+use rustc_expand::base::ExtCtxt;
+use thin_vec::{thin_vec, ThinVec};
+use rustc_span::sym;
+
+pub(super) struct AttrOpts {
+    triplicate_unsafe: bool,
+}
+
+impl AttrOpts {
+    pub(super) fn triplicate_unsafe(&self) -> bool {
+        self.triplicate_unsafe
+    }
+}
+
+pub(super) fn parse_attr_opts(cx: &ExtCtxt<'_>, meta_item: &ast::MetaItem) -> Option<AttrOpts> {
+
+    let attr_opts: ThinVec<MetaItemInner> = match meta_item.kind {
+        ast::MetaItemKind::List(ref vec) => vec.clone(),
+        ast::MetaItemKind::Word => thin_vec![],
+        _ => {
+            cx.dcx().span_err(meta_item.span, "unsupported options kind in `#[rad_protected]`");
+            return None;
+        }
+    };
+
+    let mut triplicate_unsafe = false;
+
+    for opt in attr_opts {
+        match opt {
+            MetaItemInner::MetaItem(opt) if opt.has_name(sym::triplicate_unsafe) => {
+                triplicate_unsafe = true;
+            }
+            _ => {
+                cx.dcx().span_err(opt.span(), "unsupported option in `#[rad_protected]`");
+               return None;
+            }
+        }
+    }
+
+    Some(AttrOpts { triplicate_unsafe })
+}

compiler/rustc_builtin_macros/src/rad_protected/triplicate.rs

@@ -0,0 +1,85 @@
+use rustc_ast as ast;
+use rustc_expand::base::{Annotatable, ExtCtxt};
+use rustc_span::{Span, symbol::Ident, sym, DUMMY_SP};
+use thin_vec::thin_vec;
+use super::parse_attr_opts::parse_attr_opts;
+
+pub(crate) fn triplicate(
+    cx: &mut ExtCtxt<'_>,
+    span: Span,
+    meta_item: &ast::MetaItem,
+    mut item: Annotatable,
+) -> Vec<Annotatable> {
+
+    let Some(opts) = parse_attr_opts(cx, meta_item) else {
+        return vec![item];
+    };
+
+    if opts.triplicate_unsafe() {
+        let valid = match &mut item {
+            Annotatable::Expr(expr)
+                if matches!(&expr.kind, ast::ExprKind::Block(block, _)
+                    if matches!(block.rules, ast::BlockCheckMode::Unsafe(_))
+                ) => {
+                    expr.attrs.push(cx.attr_nested_word(
+                        sym::rad_protected_mir,
+                        sym::triplicate_unsafe,
+                        DUMMY_SP,
+                    ));
+                    true
+                }
+            _ => false,
+        };
+
+        if !valid {
+            cx.dcx().span_err(
+                span,
+                "`#[rad_protected(triplicate_unsafe)]` can only be applied to `unsafe` blocks",
+            );
+        }
+
+        return vec![item];
+    }
+
+    let Annotatable::Item(mut item) = item else {
+        cx.dcx().span_err(span, "`#[rad_protected]` can only be applied to functions");
+        return vec![item];
+    };
+
+    let ast::Item {
+        kind: ast::ItemKind::Fn(func),
+        ..
+    } = &mut *item
+    else {
+        cx.dcx().span_err(span, "`#[rad_protected]` can only be applied to functions");
+        return vec![Annotatable::Item(item)];
+    };
+
+    let func_body = match &mut func.body {
+        Some(b) => b,
+        None => {
+            cx.dcx().span_err(span, "`#[rad_protected]` can only be applied to functions with a body");
+            return vec![Annotatable::Item(item)];
+        }
+    };
+
+    func_body.stmts.insert(0, cx.stmt_let(
+        DUMMY_SP,
+        false,
+        Ident::new(sym::__guard, DUMMY_SP),
+        cx.expr_call_global(
+            DUMMY_SP,
+            vec![
+                Ident::new(sym::std, DUMMY_SP),
+                Ident::new(sym::RadRustRuntime, DUMMY_SP),
+                Ident::new(sym::triplicate_process, DUMMY_SP),
+            ], 
+            thin_vec![cx.expr_usize(DUMMY_SP, 0usize)]
+        )
+    ));
+
+    let mir_attr = cx.attr_word(sym::rad_protected_mir, DUMMY_SP);
+    item.attrs.push(mir_attr);
+
+    vec![Annotatable::Item(item)]
+}

compiler/rustc_expand/src/build.rs

@@ -769,6 +769,21 @@ impl<'a> ExtCtxt<'a> {
         )
     }
 
+    // Builds `#[unsafe(name = val)]`.
+    //
+    // Note: `span` is used for both the identifier and the value.
+    pub fn attr_name_value_str_unsafe(&self, name: Symbol, val: Symbol, span: Span) -> ast::Attribute {
+        let g = &self.sess.psess.attr_id_generator;
+        attr::mk_attr_name_value_str(
+            g,
+            ast::AttrStyle::Outer,
+            ast::Safety::Unsafe(span),
+            name,
+            val,
+            span,
+        )
+    }
+
     // Builds `#[outer(inner)]`.
     pub fn attr_nested_word(&self, outer: Symbol, inner: Symbol, span: Span) -> ast::Attribute {
         let g = &self.sess.psess.attr_id_generator;

compiler/rustc_expand/src/lib.rs

@@ -23,6 +23,8 @@ pub mod expand;
 pub mod module;
 pub mod proc_macro;
 
+pub mod patch_unsafe;
+
 pub fn provide(providers: &mut rustc_middle::query::Providers) {
     providers.derive_macro_expansion = proc_macro::provide_derive_macro_expansion;
 }

compiler/rustc_expand/src/patch_unsafe.rs

@@ -0,0 +1,115 @@
+use rustc_ast as ast;
+use rustc_ast::mut_visit::{self, MutVisitor};
+use crate::base::ExtCtxt;
+use rustc_span::{symbol::Ident, Symbol, sym, DUMMY_SP};
+use thin_vec::{ThinVec, thin_vec};
+use rustc_ast::MetaItemInner;
+
+pub fn patch_unsafe_blocks(cx: &mut ExtCtxt<'_>, krate: &mut ast::Crate) {
+    if cx.sess.opts.unstable_opts.force_unstable_if_unmarked {
+        return;
+    }
+    let mut visitor = UnsafeBlockRewriter { cx };
+    visitor.visit_crate(krate);
+}
+
+struct DummyIdAssigner<'a, 'cx> {
+    cx: &'a mut ExtCtxt<'cx>,
+}
+
+impl MutVisitor for DummyIdAssigner<'_, '_> {
+    fn visit_id(&mut self, id: &mut ast::NodeId) {
+        if *id == ast::DUMMY_NODE_ID {
+            *id = self.cx.resolver.next_node_id();
+        }
+    }
+}
+
+struct UnsafeBlockRewriter<'a, 'cx> {
+    cx: &'a mut ExtCtxt<'cx>,
+}
+
+impl MutVisitor for UnsafeBlockRewriter<'_, '_> {
+    fn visit_expr(&mut self, expr: &mut ast::Expr) {
+
+        if let ast::ExprKind::Block(block, _) = &mut expr.kind {
+            if matches!(block.rules, ast::BlockCheckMode::Unsafe(_)) {
+
+                if !skip_patch(&mut expr.attrs) {
+                    patch_unsafe_block(self.cx, block);
+                }
+                return;
+            }
+        }
+
+        mut_visit::walk_expr(self, expr);
+    }
+}
+
+fn patch_unsafe_block(cx: &mut ExtCtxt<'_>, block: &mut ast::Block) {
+    let runtime_method_call = |cx: &mut ExtCtxt<'_>, name: Symbol| {
+        cx.expr_call_global(
+            DUMMY_SP,
+            vec![
+                Ident::new(sym::std, DUMMY_SP),
+                Ident::new(sym::RadRustRuntime, DUMMY_SP),
+                Ident::new(name, DUMMY_SP),
+            ],
+            thin_vec![],
+        )
+    };
+
+    let enter_call = runtime_method_call(cx, sym::enter_critical_section);
+    let exit_call = runtime_method_call(cx, sym::exit_critical_section);
+
+    let inner_block = cx.block(block.span, block.stmts.clone());
+    let if_stmt = cx.stmt_expr(cx.expr_if(DUMMY_SP, 
+        enter_call, 
+        cx.expr_block(inner_block), 
+        None
+    ));
+    let exit_call_stmt = cx.stmt_expr(exit_call);
+
+    let mut assigner = DummyIdAssigner { cx };
+    let if_stmt = assign_stmt(&mut assigner, if_stmt);
+    let exit_call_stmt = assign_stmt(&mut assigner, exit_call_stmt);
+
+    block.stmts = thin_vec![if_stmt, exit_call_stmt];
+}
+
+fn skip_patch(attrs: &mut ThinVec<ast::Attribute>) -> bool {
+    let mut removed = false;
+
+    attrs.retain(|attr| {
+        let keep = !attr.meta().is_some_and(is_skip_attr);
+
+        if !keep {
+            removed = true;
+        }
+        keep
+    });
+
+    removed
+}
+
+fn is_skip_attr(meta: ast::MetaItem) -> bool {
+    if !meta.has_name(sym::rad_protected_mir) {
+        return false;
+    }
+
+    let ast::MetaItemKind::List(items) = &meta.kind else {
+        return false;
+    };
+
+    items.iter().any(|item| {
+        matches!(item, MetaItemInner::MetaItem(mi) if mi.has_name(sym::triplicate_unsafe))
+    })
+}
+
+fn assign_stmt(assigner: &mut DummyIdAssigner<'_, '_>, stmt: ast::Stmt) -> ast::Stmt {
+    assigner
+        .flat_map_stmt(stmt)
+        .into_iter()
+        .next()
+        .expect("statement unexpectedly removed")
+}

compiler/rustc_feature/src/builtin_attrs.rs

@@ -903,9 +903,10 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         EncodeCrossCrate::Yes, pin_ergonomics, experimental!(pin_v2),
     ),
 
-    // Radshield protection attribute (ungated for easy testing)
+    // Radshield MIR pass protection attribute (ungated for easy testing)
+    // This attribute is used internally by rad_protected, and is not intended for use by the user
     ungated!(
-        rad_protected, Normal, template!(Word), WarnFollowing,
+        rad_protected_mir, Normal, template!(Word), WarnFollowing,
         EncodeCrossCrate::No,
     ),
 

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -1234,7 +1234,7 @@ pub enum AttributeKind {
     /// Represents `#[profiler_runtime]`
     ProfilerRuntime,
 
-    /// Represents `#[rad_protected]`
+    /// Represents `#[rad_protected_mir]`
     RadProtected(Span),
 
     /// Represents [`#[recursion_limit]`](https://doc.rust-lang.org/reference/attributes/limits.html#the-recursion_limit-attribute)

compiler/rustc_interface/src/passes.rs

@@ -15,6 +15,7 @@ use rustc_data_structures::sync::{AppendOnlyIndexVec, FreezeLock, WorkerLocal, p
 use rustc_data_structures::thousands;
 use rustc_errors::timings::TimingSection;
 use rustc_expand::base::{ExtCtxt, LintStoreExpand};
+use rustc_expand::patch_unsafe::patch_unsafe_blocks;
 use rustc_feature::Features;
 use rustc_fs_util::try_canonicalize;
 use rustc_hir::attrs::AttributeKind;
@@ -213,12 +214,14 @@ fn configure_and_expand(
         let mut ecx = ExtCtxt::new(sess, cfg, resolver, Some(&lint_store));
         ecx.num_standard_library_imports = num_standard_library_imports;
         // Expand macros now!
-        let krate = sess.time("expand_crate", || ecx.monotonic_expander().expand_crate(krate));
+        let mut krate = sess.time("expand_crate", || ecx.monotonic_expander().expand_crate(krate));
 
         if ecx.nb_macro_errors > 0 {
             sess.dcx().abort_if_errors();
         }
 
+        patch_unsafe_blocks(&mut ecx, &mut krate);
+
         // The rest is error reporting and stats
 
         sess.psess.buffered_lints.with_lock(|buffered_lints: &mut Vec<BufferedEarlyLint>| {

compiler/rustc_mir_build/src/builder/checkpoint.rs

@@ -0,0 +1,44 @@
+use rustc_hir::def_id::LocalDefId;
+use rustc_middle::mir::{TerminatorKind, Operand, Place, UnwindAction, CallSource, BasicBlock};
+use rustc_middle::ty::TyCtxt;
+use rustc_span::{sym, Span};
+use rustc_hir::find_attr;
+
+use super::Builder;
+
+impl<'a, 'tcx> Builder<'a, 'tcx> {
+    pub(super) fn inject_checkpoint_marker(&mut self, block: BasicBlock, span: Span) -> BasicBlock {
+        if !Self::is_checkpoint(self.tcx, self.def_id) {
+            return block;
+        }
+
+        let Some(marker_def_id) = self.tcx.get_diagnostic_item(sym::__checkpoint) else {
+            return block;
+        };
+
+        let next = self.cfg.start_new_block();
+        let source_info = self.source_info(span);
+        let func = Operand::function_handle(self.tcx, marker_def_id, [], span);
+        let destination = Place::from(self.temp(self.tcx.types.unit, span));
+
+        self.cfg.terminate(
+            block,
+            source_info,
+            TerminatorKind::Call {
+                func,
+                args: Box::new([]),
+                destination,
+                target: Some(next),
+                unwind: UnwindAction::Continue,
+                call_source: CallSource::Misc,
+                fn_span: span,
+            },
+        );
+
+        next
+    }
+
+    fn is_checkpoint(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
+        find_attr!(tcx, def_id, RadProtected(_))
+    }
+}