Auto merge of #87794 - bonega:enum_niche_prefer_zero, r=nagisa

Enum should prefer discriminant zero for niche

Given an enum with unassigned zero-discriminant, rust should prefer it for niche selection.
Zero as discriminant for `Option<Enum>` makes it possible for LLVM to optimize resulting asm.

- Eliminate branch when expected value coincides.
- Use smaller instruction `test eax, eax` instead of `cmp eax, ?`
- Possible interaction with zeroed memory?

Example:
```rust

pub enum Size {
    One = 1,
    Two = 2,
    Three = 3,
}

pub fn handle(x: Option<Size>) -> u8 {
    match x {
        None => {0}
        Some(size) => {size as u8}
    }
}
```
In this case discriminant zero is available as a niche.

Above example on nightly:
```asm
 mov     eax, edi
 cmp     al, 4
 jne     .LBB0_2
 xor     eax, eax
.LBB0_2:
 ret
```

PR:
```asm
 mov     eax, edi
 ret
```

I created this PR because I had a performance regression when I tried to use an enum to represent legal grapheme byte-length for utf8.

Using an enum instead of `NonZeroU8` [here](https://github.com/bonega/yore/blob/d683304f5dfe2e99f769e6ab8adf8d60a0d1d9b3/src/internal/decoder_incomplete.rs#L90)
resulted in a performance regression of about 5%.
I consider this to be a somewhat realistic benchmark.

Thanks to `@ogoffart` for pointing me in the right direction!

Edit: Updated description

3 files changed, 82 insertions, 9 deletions

diff --git a/compiler/rustc_target/src/abi/mod.rs b/compiler/rustc_target/src/abi/mod.rs
index af75007b09a..61607159208 100644
--- a/compiler/rustc_target/src/abi/mod.rs
+++ b/compiler/rustc_target/src/abi/mod.rs
@@ -1099,19 +1099,53 @@ impl Niche {
         assert!(size.bits() <= 128);
         let max_value = size.unsigned_int_max();
 
-        if count > max_value {
+        let niche = v.end.wrapping_add(1)..v.start;
+        let available = niche.end.wrapping_sub(niche.start) & max_value;
+        if count > available {
             return None;
         }
 
-        // Compute the range of invalid values being reserved.
-        let start = v.end.wrapping_add(1) & max_value;
-        let end = v.end.wrapping_add(count) & max_value;
-
-        if v.contains(end) {
-            return None;
+        // Extend the range of valid values being reserved by moving either `v.start` or `v.end` bound.
+        // Given an eventual `Option<T>`, we try to maximize the chance for `None` to occupy the niche of zero.
+        // This is accomplished by prefering enums with 2 variants(`count==1`) and always taking the shortest path to niche zero.
+        // Having `None` in niche zero can enable some special optimizations.
+        //
+        // Bound selection criteria:
+        // 1. Select closest to zero given wrapping semantics.
+        // 2. Avoid moving past zero if possible.
+        //
+        // In practice this means that enums with `count > 1` are unlikely to claim niche zero, since they have to fit perfectly.
+        // If niche zero is already reserved, the selection of bounds are of little interest.
+        let move_start = |v: WrappingRange| {
+            let start = v.start.wrapping_sub(1) & max_value;
+            Some((start, Scalar { value, valid_range: v.with_start(start) }))
+        };
+        let move_end = |v: WrappingRange| {
+            let start = v.end.wrapping_add(1) & max_value;
+            let end = v.end.wrapping_add(count) & max_value;
+            Some((start, Scalar { value, valid_range: v.with_end(end) }))
+        };
+        let distance_end_zero = max_value - v.end;
+        if v.start > v.end {
+            // zero is unavailable because wrapping occurs
+            move_end(v)
+        } else if v.start <= distance_end_zero {
+            if count <= v.start {
+                move_start(v)
+            } else {
+                // moved past zero, use other bound
+                move_end(v)
+            }
+        } else {
+            let end = v.end.wrapping_add(count) & max_value;
+            let overshot_zero = (1..=v.end).contains(&end);
+            if overshot_zero {
+                // moved past zero, use other bound
+                move_start(v)
+            } else {
+                move_end(v)
+            }
         }
-
-        Some((start, Scalar { value, valid_range: v.with_end(end) }))
     }
 }
 
diff --git a/src/test/assembly/niche-prefer-zero.rs b/src/test/assembly/niche-prefer-zero.rs
new file mode 100644
index 00000000000..0ab37a618da
--- /dev/null
+++ b/src/test/assembly/niche-prefer-zero.rs
@@ -0,0 +1,25 @@
+// Check that niche selection prefers zero and that jumps are optimized away.
+// See https://github.com/rust-lang/rust/pull/87794
+// assembly-output: emit-asm
+// only-x86
+// compile-flags: -Copt-level=3
+
+#![crate_type = "lib"]
+
+#[repr(u8)]
+pub enum Size {
+    One = 1,
+    Two = 2,
+    Three = 3,
+}
+
+#[no_mangle]
+pub fn handle(x: Option<Size>) -> u8 {
+    match x {
+        None => 0,
+        Some(size) => size as u8,
+    }
+}
+
+// There should be no jumps in output
+// CHECK-NOT: j
diff --git a/src/test/ui/enum-discriminant/niche-prefer-zero.rs b/src/test/ui/enum-discriminant/niche-prefer-zero.rs
new file mode 100644
index 00000000000..f20607a8903
--- /dev/null
+++ b/src/test/ui/enum-discriminant/niche-prefer-zero.rs
@@ -0,0 +1,14 @@
+// Check that niche selection prefers zero.
+// See https://github.com/rust-lang/rust/pull/87794
+// run-pass
+#[repr(u8)]
+pub enum Size {
+    One = 1,
+    Two = 2,
+    Three = 3,
+}
+
+fn main() {
+    // check that `None` is zero
+    assert_eq!(0, unsafe { std::mem::transmute::<Option<Size>, u8>(None) });
+}