Replace ZExt and SExt flags with ArgExtension enum

Both flags are mutually exclusive

5 files changed, 53 insertions, 9 deletions

diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/compiler/rustc_codegen_llvm/src/abi.rs
index 7857ccb613b..3ca0a9ce8d9 100644
--- a/compiler/rustc_codegen_llvm/src/abi.rs
+++ b/compiler/rustc_codegen_llvm/src/abi.rs
@@ -36,7 +36,7 @@ impl ArgAttributeExt for ArgAttribute {
     where
         F: FnMut(llvm::Attribute),
     {
-        for_each_kind!(self, f, NoAlias, NoCapture, NonNull, ReadOnly, SExt, StructRet, ZExt, InReg)
+        for_each_kind!(self, f, NoAlias, NoCapture, NonNull, ReadOnly, StructRet, InReg)
     }
 }
 
@@ -65,6 +65,15 @@ impl ArgAttributesExt for ArgAttributes {
                 llvm::LLVMRustAddByValAttr(llfn, idx.as_uint(), ty.unwrap());
             }
             regular.for_each_kind(|attr| attr.apply_llfn(idx, llfn));
+            match self.arg_ext {
+                ArgExtension::None => {}
+                ArgExtension::Zext => {
+                    llvm::Attribute::ZExt.apply_llfn(idx, llfn);
+                }
+                ArgExtension::Sext => {
+                    llvm::Attribute::SExt.apply_llfn(idx, llfn);
+                }
+            }
         }
     }
 
@@ -95,6 +104,15 @@ impl ArgAttributesExt for ArgAttributes {
                 llvm::LLVMRustAddByValCallSiteAttr(callsite, idx.as_uint(), ty.unwrap());
             }
             regular.for_each_kind(|attr| attr.apply_callsite(idx, callsite));
+            match self.arg_ext {
+                ArgExtension::None => {}
+                ArgExtension::Zext => {
+                    llvm::Attribute::ZExt.apply_callsite(idx, callsite);
+                }
+                ArgExtension::Sext => {
+                    llvm::Attribute::SExt.apply_callsite(idx, callsite);
+                }
+            }
         }
     }
 }
diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index 5626c864fe1..32bd9651ece 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -2619,7 +2619,7 @@ where
                                       is_return: bool| {
             // Booleans are always an i1 that needs to be zero-extended.
             if scalar.is_bool() {
-                attrs.set(ArgAttribute::ZExt);
+                attrs.zext();
                 return;
             }
 
diff --git a/compiler/rustc_target/src/abi/call/mips64.rs b/compiler/rustc_target/src/abi/call/mips64.rs
index 917dd104d14..739ec6112ad 100644
--- a/compiler/rustc_target/src/abi/call/mips64.rs
+++ b/compiler/rustc_target/src/abi/call/mips64.rs
@@ -1,4 +1,4 @@
-use crate::abi::call::{ArgAbi, ArgAttribute, CastTarget, FnAbi, PassMode, Reg, RegKind, Uniform};
+use crate::abi::call::{ArgAbi, CastTarget, FnAbi, PassMode, Reg, RegKind, Uniform};
 use crate::abi::{self, HasDataLayout, LayoutOf, Size, TyAndLayout, TyAndLayoutMethods};
 
 fn extend_integer_width_mips<Ty>(arg: &mut ArgAbi<'_, Ty>, bits: u64) {
@@ -7,7 +7,7 @@ fn extend_integer_width_mips<Ty>(arg: &mut ArgAbi<'_, Ty>, bits: u64) {
         if let abi::Int(i, signed) = scalar.value {
             if !signed && i.size().bits() == 32 {
                 if let PassMode::Direct(ref mut attrs) = arg.mode {
-                    attrs.set(ArgAttribute::SExt);
+                    attrs.sext();
                     return;
                 }
             }
diff --git a/compiler/rustc_target/src/abi/call/mod.rs b/compiler/rustc_target/src/abi/call/mod.rs
index 429a3375cd8..9d657250c7c 100644
--- a/compiler/rustc_target/src/abi/call/mod.rs
+++ b/compiler/rustc_target/src/abi/call/mod.rs
@@ -57,19 +57,28 @@ mod attr_impl {
             const NoCapture = 1 << 2;
             const NonNull   = 1 << 3;
             const ReadOnly  = 1 << 4;
-            const SExt      = 1 << 5;
             const StructRet = 1 << 6;
-            const ZExt      = 1 << 7;
             const InReg     = 1 << 8;
         }
     }
 }
 
+/// Sometimes an ABI requires small integers to be extended to a full or partial register. This enum
+/// defines if this extension should be zero-extension or sign-extension when necssary. When it is
+/// not necesary to extend the argument, this enum is ignored.
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub enum ArgExtension {
+    None,
+    Zext,
+    Sext,
+}
+
 /// A compact representation of LLVM attributes (at least those relevant for this module)
 /// that can be manipulated without interacting with LLVM's Attribute machinery.
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct ArgAttributes {
     pub regular: ArgAttribute,
+    pub arg_ext: ArgExtension,
     /// The minimum size of the pointee, guaranteed to be valid for the duration of the whole call
     /// (corresponding to LLVM's dereferenceable and dereferenceable_or_null attributes).
     pub pointee_size: Size,
@@ -80,11 +89,24 @@ impl ArgAttributes {
     pub fn new() -> Self {
         ArgAttributes {
             regular: ArgAttribute::default(),
+            arg_ext: ArgExtension::None,
             pointee_size: Size::ZERO,
             pointee_align: None,
         }
     }
 
+    pub fn zext(&mut self) -> &mut Self {
+        assert_ne!(self.arg_ext, ArgExtension::Sext);
+        self.arg_ext = ArgExtension::Zext;
+        self
+    }
+
+    pub fn sext(&mut self) -> &mut Self {
+        assert_ne!(self.arg_ext, ArgExtension::Zext);
+        self.arg_ext = ArgExtension::Sext;
+        self
+    }
+
     pub fn set(&mut self, attr: ArgAttribute) -> &mut Self {
         self.regular |= attr;
         self
@@ -457,7 +479,11 @@ impl<'a, Ty> ArgAbi<'a, Ty> {
             if let abi::Int(i, signed) = scalar.value {
                 if i.size().bits() < bits {
                     if let PassMode::Direct(ref mut attrs) = self.mode {
-                        attrs.set(if signed { ArgAttribute::SExt } else { ArgAttribute::ZExt });
+                        if signed {
+                            attrs.sext()
+                        } else {
+                            attrs.zext()
+                        };
                     }
                 }
             }
diff --git a/compiler/rustc_target/src/abi/call/riscv.rs b/compiler/rustc_target/src/abi/call/riscv.rs
index 782c661c31f..9c4ac85ce61 100644
--- a/compiler/rustc_target/src/abi/call/riscv.rs
+++ b/compiler/rustc_target/src/abi/call/riscv.rs
@@ -4,7 +4,7 @@
 // Reference: Clang RISC-V ELF psABI lowering code
 // https://github.com/llvm/llvm-project/blob/8e780252a7284be45cf1ba224cabd884847e8e92/clang/lib/CodeGen/TargetInfo.cpp#L9311-L9773
 
-use crate::abi::call::{ArgAbi, ArgAttribute, CastTarget, FnAbi, PassMode, Reg, RegKind, Uniform};
+use crate::abi::call::{ArgAbi, CastTarget, FnAbi, PassMode, Reg, RegKind, Uniform};
 use crate::abi::{
     self, Abi, FieldsShape, HasDataLayout, LayoutOf, Size, TyAndLayout, TyAndLayoutMethods,
 };
@@ -308,7 +308,7 @@ fn extend_integer_width<'a, Ty>(arg: &mut ArgAbi<'a, Ty>, xlen: u64) {
             // 32-bit integers are always sign-extended
             if i.size().bits() == 32 && xlen > 32 {
                 if let PassMode::Direct(ref mut attrs) = arg.mode {
-                    attrs.set(ArgAttribute::SExt);
+                    attrs.sext();
                     return;
                 }
             }