Auto merge of #114790 - taiki-e:asm-maybe-uninit, r=Amanieu

Allow MaybeUninit in input and output of inline assembly

**Motivation:**

As part of the work to remove UBs from crossbeam's AtomicCell, I'm writing a library to implement atomic operations on MaybeUnint using inline assembly ([atomic-maybe-uninit](https://github.com/taiki-e/atomic-maybe-uninit), crossbeam-rs/crossbeam#1015).

However, currently, MaybeUnint cannot be used in input&output of inline assembly, so when processing MaybeUninit, values must be [passed through memory](https://github.com/taiki-e/atomic-maybe-uninit/blob/main/src/arch/aarch64.rs#L121-L122). It is inefficient and microbenchmarks have [actually shown significant performance degradation](crossbeam-rs/crossbeam#1015 (comment)).

It would be nice if we could allow MaybeUninit in input and output of inline assembly.

---

This PR changed the type check in rustc_hir_analysis to allow `MaybeUnint<int | float | ptr | fn ptr | simd vector>` in input and output of inline assembly and added a simple test.

To be honest, I'm not sure that this is the correct way to do it, because this is like doing transmute to integers/floats/etc from MaybeUninit on the compiler side. EDIT: [this seems fine](https://rust-lang.zulipchat.com/#narrow/stream/216763-project-inline-asm/topic/MaybeUninit.20in.20asm!/near/384662900)

r? `@Amanieu`
cc `@thomcc` (because you [had previously proposed this](https://rust-lang.zulipchat.com/#narrow/stream/216763-project-inline-asm/topic/MaybeUninit.20in.20asm!))

compiler/rustc_hir_analysis/src/check/intrinsicck.rs

@@ -44,31 +44,15 @@ impl<'a, 'tcx> InlineAsmCtxt<'a, 'tcx> {
         false
     }
 
-    fn check_asm_operand_type(
-        &self,
-        idx: usize,
-        reg: InlineAsmRegOrRegClass,
-        expr: &'tcx hir::Expr<'tcx>,
-        template: &[InlineAsmTemplatePiece],
-        is_input: bool,
-        tied_input: Option<(&'tcx hir::Expr<'tcx>, Option<InlineAsmType>)>,
-        target_features: &FxIndexSet<Symbol>,
-    ) -> Option<InlineAsmType> {
-        let ty = (self.get_operand_ty)(expr);
-        if ty.has_non_region_infer() {
-            bug!("inference variable in asm operand ty: {:?} {:?}", expr, ty);
-        }
+    fn get_asm_ty(&self, ty: Ty<'tcx>) -> Option<InlineAsmType> {
         let asm_ty_isize = match self.tcx.sess.target.pointer_width {
             16 => InlineAsmType::I16,
             32 => InlineAsmType::I32,
             64 => InlineAsmType::I64,
             _ => unreachable!(),
         };
 
-        let asm_ty = match *ty.kind() {
-            // `!` is allowed for input but not for output (issue #87802)
-            ty::Never if is_input => return None,
-            _ if ty.references_error() => return None,
+        match *ty.kind() {
             ty::Int(IntTy::I8) | ty::Uint(UintTy::U8) => Some(InlineAsmType::I8),
             ty::Int(IntTy::I16) | ty::Uint(UintTy::U16) => Some(InlineAsmType::I16),
             ty::Int(IntTy::I32) | ty::Uint(UintTy::U32) => Some(InlineAsmType::I32),
@@ -99,7 +83,6 @@ impl<'a, 'tcx> InlineAsmCtxt<'a, 'tcx> {
                 };
 
                 match ty.kind() {
-                    ty::Never | ty::Error(_) => return None,
                     ty::Int(IntTy::I8) | ty::Uint(UintTy::U8) => Some(InlineAsmType::VecI8(size)),
                     ty::Int(IntTy::I16) | ty::Uint(UintTy::U16) => {
                         Some(InlineAsmType::VecI16(size))
@@ -128,6 +111,38 @@ impl<'a, 'tcx> InlineAsmCtxt<'a, 'tcx> {
             }
             ty::Infer(_) => unreachable!(),
             _ => None,
+        }
+    }
+
+    fn check_asm_operand_type(
+        &self,
+        idx: usize,
+        reg: InlineAsmRegOrRegClass,
+        expr: &'tcx hir::Expr<'tcx>,
+        template: &[InlineAsmTemplatePiece],
+        is_input: bool,
+        tied_input: Option<(&'tcx hir::Expr<'tcx>, Option<InlineAsmType>)>,
+        target_features: &FxIndexSet<Symbol>,
+    ) -> Option<InlineAsmType> {
+        let ty = (self.get_operand_ty)(expr);
+        if ty.has_non_region_infer() {
+            bug!("inference variable in asm operand ty: {:?} {:?}", expr, ty);
+        }
+
+        let asm_ty = match *ty.kind() {
+            // `!` is allowed for input but not for output (issue #87802)
+            ty::Never if is_input => return None,
+            _ if ty.references_error() => return None,
+            ty::Adt(adt, args) if Some(adt.did()) == self.tcx.lang_items().maybe_uninit() => {
+                let fields = &adt.non_enum_variant().fields;
+                let ty = fields[FieldIdx::from_u32(1)].ty(self.tcx, args);
+                let ty::Adt(ty, args) = ty.kind() else { unreachable!() };
+                assert!(ty.is_manually_drop());
+                let fields = &ty.non_enum_variant().fields;
+                let ty = fields[FieldIdx::from_u32(0)].ty(self.tcx, args);
+                self.get_asm_ty(ty)
+            }
+            _ => self.get_asm_ty(ty),
         };
         let Some(asm_ty) = asm_ty else {
             let msg = format!("cannot use value of type `{ty}` for inline assembly");

tests/codegen/asm-maybe-uninit.rs

@@ -0,0 +1,27 @@
+// compile-flags: -O
+// only-x86_64
+
+#![crate_type = "rlib"]
+#![allow(asm_sub_register)]
+
+use std::mem::MaybeUninit;
+use std::arch::asm;
+
+// CHECK-LABEL: @int
+#[no_mangle]
+pub unsafe fn int(x: MaybeUninit<i32>) -> MaybeUninit<i32> {
+    let y: MaybeUninit<i32>;
+    asm!("/*{}{}*/", in(reg) x, out(reg) y);
+    y
+}
+
+// CHECK-LABEL: @inout
+#[no_mangle]
+pub unsafe fn inout(mut x: i32) -> MaybeUninit<u32> {
+    let mut y: MaybeUninit<u32>;
+    asm!("/*{}*/", inout(reg) x => y);
+    asm!("/*{}*/", inout(reg) y => x);
+    asm!("/*{}*/", inlateout(reg) x => y);
+    asm!("/*{}*/", inlateout(reg) y => x);
+    y
+}