3 files changed, 187 insertions, 0 deletions
diff --git a/compiler/rustc_hir_analysis/src/errors.rs b/compiler/rustc_hir_analysis/src/errors.rs
index 03311fed396..1bb4db9cbcf 100644
--- a/compiler/rustc_hir_analysis/src/errors.rs
+++ b/compiler/rustc_hir_analysis/src/errors.rs
@@ -1673,3 +1673,30 @@ pub struct InvalidReceiverTy<'tcx> {
 #[note]
 #[help]
 pub struct EffectsWithoutNextSolver;
+
+#[derive(Diagnostic)]
+#[diag(hir_analysis_cmse_call_inputs_stack_spill, code = E0798)]
+#[note]
+pub struct CmseCallInputsStackSpill {
+    #[primary_span]
+    #[label]
+    pub span: Span,
+    pub plural: bool,
+}
+
+#[derive(Diagnostic)]
+#[diag(hir_analysis_cmse_call_output_stack_spill, code = E0798)]
+#[note(hir_analysis_note1)]
+#[note(hir_analysis_note2)]
+pub struct CmseCallOutputStackSpill {
+    #[primary_span]
+    #[label]
+    pub span: Span,
+}
+
+#[derive(Diagnostic)]
+#[diag(hir_analysis_cmse_call_generic, code = E0798)]
+pub struct CmseCallGeneric {
+    #[primary_span]
+    pub span: Span,
+}
diff --git a/compiler/rustc_hir_analysis/src/hir_ty_lowering/cmse.rs b/compiler/rustc_hir_analysis/src/hir_ty_lowering/cmse.rs
new file mode 100644
index 00000000000..e99717ce00f
--- /dev/null
+++ b/compiler/rustc_hir_analysis/src/hir_ty_lowering/cmse.rs
@@ -0,0 +1,156 @@
+use rustc_errors::DiagCtxtHandle;
+use rustc_hir as hir;
+use rustc_hir::HirId;
+use rustc_middle::ty::layout::LayoutError;
+use rustc_middle::ty::{self, ParamEnv, TyCtxt};
+use rustc_span::Span;
+use rustc_target::spec::abi;
+
+use crate::errors;
+
+/// Check conditions on inputs and outputs that the cmse ABIs impose: arguments and results MUST be
+/// returned via registers (i.e. MUST NOT spill to the stack). LLVM will also validate these
+/// conditions, but by checking them here rustc can emit nicer error messages.
+pub fn validate_cmse_abi<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    dcx: DiagCtxtHandle<'_>,
+    hir_id: HirId,
+    abi: abi::Abi,
+    fn_sig: ty::PolyFnSig<'tcx>,
+) {
+    if let abi::Abi::CCmseNonSecureCall = abi {
+        let hir_node = tcx.hir_node(hir_id);
+        let hir::Node::Ty(hir::Ty {
+            span: bare_fn_span,
+            kind: hir::TyKind::BareFn(bare_fn_ty),
+            ..
+        }) = hir_node
+        else {
+            // might happen when this ABI is used incorrectly. That will be handled elsewhere
+            return;
+        };
+
+        match is_valid_cmse_inputs(tcx, fn_sig) {
+            Ok(Ok(())) => {}
+            Ok(Err(index)) => {
+                // fn(x: u32, u32, u32, u16, y: u16) -> u32,
+                //                           ^^^^^^
+                let span = bare_fn_ty.param_names[index]
+                    .span
+                    .to(bare_fn_ty.decl.inputs[index].span)
+                    .to(bare_fn_ty.decl.inputs.last().unwrap().span);
+                let plural = bare_fn_ty.param_names.len() - index != 1;
+                dcx.emit_err(errors::CmseCallInputsStackSpill { span, plural });
+            }
+            Err(layout_err) => {
+                if let Some(err) = cmse_layout_err(layout_err, *bare_fn_span) {
+                    dcx.emit_err(err);
+                }
+            }
+        }
+
+        match is_valid_cmse_output(tcx, fn_sig) {
+            Ok(true) => {}
+            Ok(false) => {
+                let span = bare_fn_ty.decl.output.span();
+                dcx.emit_err(errors::CmseCallOutputStackSpill { span });
+            }
+            Err(layout_err) => {
+                if let Some(err) = cmse_layout_err(layout_err, *bare_fn_span) {
+                    dcx.emit_err(err);
+                }
+            }
+        };
+    }
+}
+
+/// Returns whether the inputs will fit into the available registers
+fn is_valid_cmse_inputs<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    fn_sig: ty::PolyFnSig<'tcx>,
+) -> Result<Result<(), usize>, &'tcx LayoutError<'tcx>> {
+    let mut span = None;
+    let mut accum = 0u64;
+
+    for (index, arg_def) in fn_sig.inputs().iter().enumerate() {
+        let layout = tcx.layout_of(ParamEnv::reveal_all().and(*arg_def.skip_binder()))?;
+
+        let align = layout.layout.align().abi.bytes();
+        let size = layout.layout.size().bytes();
+
+        accum += size;
+        accum = accum.next_multiple_of(Ord::max(4, align));
+
+        // i.e. exceeds 4 32-bit registers
+        if accum > 16 {
+            span = span.or(Some(index));
+        }
+    }
+
+    match span {
+        None => Ok(Ok(())),
+        Some(span) => Ok(Err(span)),
+    }
+}
+
+/// Returns whether the output will fit into the available registers
+fn is_valid_cmse_output<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    fn_sig: ty::PolyFnSig<'tcx>,
+) -> Result<bool, &'tcx LayoutError<'tcx>> {
+    let mut ret_ty = fn_sig.output().skip_binder();
+    let layout = tcx.layout_of(ParamEnv::reveal_all().and(ret_ty))?;
+    let size = layout.layout.size().bytes();
+
+    if size <= 4 {
+        return Ok(true);
+    } else if size > 8 {
+        return Ok(false);
+    }
+
+    // next we need to peel any repr(transparent) layers off
+    'outer: loop {
+        let ty::Adt(adt_def, args) = ret_ty.kind() else {
+            break;
+        };
+
+        if !adt_def.repr().transparent() {
+            break;
+        }
+
+        // the first field with non-trivial size and alignment must be the data
+        for variant_def in adt_def.variants() {
+            for field_def in variant_def.fields.iter() {
+                let ty = field_def.ty(tcx, args);
+                let layout = tcx.layout_of(ParamEnv::reveal_all().and(ty))?;
+
+                if !layout.layout.is_1zst() {
+                    ret_ty = ty;
+                    continue 'outer;
+                }
+            }
+        }
+    }
+
+    Ok(ret_ty == tcx.types.i64 || ret_ty == tcx.types.u64 || ret_ty == tcx.types.f64)
+}
+
+fn cmse_layout_err<'tcx>(
+    layout_err: &'tcx LayoutError<'tcx>,
+    span: Span,
+) -> Option<crate::errors::CmseCallGeneric> {
+    use LayoutError::*;
+
+    match layout_err {
+        Unknown(ty) => {
+            if ty.is_impl_trait() {
+                None // prevent double reporting of this error
+            } else {
+                Some(errors::CmseCallGeneric { span })
+            }
+        }
+        SizeOverflow(..) | NormalizationFailure(..) | ReferencesError(..) | Cycle(..) => {
+            None // not our job to report these
+        }
+    }
+}
diff --git a/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs b/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
index 7d1ab232c71..d6eb1a66902 100644
--- a/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
+++ b/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
@@ -14,6 +14,7 @@
 //! trait references and bounds.
 
 mod bounds;
+mod cmse;
 pub mod errors;
 pub mod generics;
 mod lint;
@@ -2378,6 +2379,9 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
         let fn_ty = tcx.mk_fn_sig(input_tys, output_ty, decl.c_variadic, safety, abi);
         let bare_fn_ty = ty::Binder::bind_with_vars(fn_ty, bound_vars);
 
+        // reject function types that violate cmse ABI requirements
+        cmse::validate_cmse_abi(self.tcx(), self.dcx(), hir_id, abi, bare_fn_ty);
+
         // Find any late-bound regions declared in return type that do
         // not appear in the arguments. These are not well-formed.
         //