From 4d66fbc4b9b95a652636e3723937c3accec85d65 Mon Sep 17 00:00:00 2001
From: Andreas Liljeqvist <bonega@gmail.com>
Date: Sun, 8 Aug 2021 13:34:46 +0200
Subject: [PATCH] enum niche allocation grows toward zero if possible

---
 compiler/rustc_target/src/abi/mod.rs          | 52 +++++++++++++++----
 src/test/assembly/niche-prefer-zero.rs        | 25 +++++++++
 .../ui/enum-discriminant/niche-prefer-zero.rs | 14 +++++
 3 files changed, 82 insertions(+), 9 deletions(-)
 create mode 100644 src/test/assembly/niche-prefer-zero.rs
 create mode 100644 src/test/ui/enum-discriminant/niche-prefer-zero.rs

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) });
+}