Pick the largest niche even if the largest niche is wrapped around

5 files changed, 91 insertions, 35 deletions

diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs
index 716bb716cdb..90c63bc9db3 100644
--- a/compiler/rustc_abi/src/layout.rs
+++ b/compiler/rustc_abi/src/layout.rs
@@ -1,3 +1,4 @@
+use std::collections::BTreeSet;
 use std::fmt::{self, Write};
 use std::ops::{Bound, Deref};
 use std::{cmp, iter};
@@ -5,7 +6,7 @@ use std::{cmp, iter};
 use rustc_hashes::Hash64;
 use rustc_index::Idx;
 use rustc_index::bit_set::BitMatrix;
-use tracing::debug;
+use tracing::{debug, trace};
 
 use crate::{
     AbiAlign, Align, BackendRepr, FieldsShape, HasDataLayout, IndexSlice, IndexVec, Integer,
@@ -766,30 +767,63 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
 
         let niche_filling_layout = calculate_niche_filling_layout();
 
-        let (mut min, mut max) = (i128::MAX, i128::MIN);
         let discr_type = repr.discr_type();
-        let bits = Integer::from_attr(dl, discr_type).size().bits();
-        for (i, mut val) in discriminants {
-            if !repr.c() && variants[i].iter().any(|f| f.is_uninhabited()) {
-                continue;
-            }
-            if discr_type.is_signed() {
-                // sign extend the raw representation to be an i128
-                val = (val << (128 - bits)) >> (128 - bits);
-            }
-            if val < min {
-                min = val;
-            }
-            if val > max {
-                max = val;
-            }
-        }
-        // We might have no inhabited variants, so pretend there's at least one.
-        if (min, max) == (i128::MAX, i128::MIN) {
-            min = 0;
-            max = 0;
-        }
-        assert!(min <= max, "discriminant range is {min}...{max}");
+        let discr_int = Integer::from_attr(dl, discr_type);
+        let bits = discr_int.size().bits();
+        // Because we can only represent one range of valid values, we'll look for the
+        // largest range of invalid values and pick everything else as the range of valid
+        // values.
+
+        // First we need to sort the possible discriminant values so that we can look for the largest gap:
+        let valid_discriminants: BTreeSet<i128> = discriminants
+            .filter(|&(i, _)| repr.c() || variants[i].iter().all(|f| !f.is_uninhabited()))
+            .map(|(_, val)| {
+                if discr_type.is_signed() {
+                    // sign extend the raw representation to be an i128
+                    (val << (128 - bits)) >> (128 - bits)
+                } else {
+                    val
+                }
+            })
+            .collect();
+        trace!(?valid_discriminants);
+        let discriminants = valid_discriminants.iter().copied();
+        //let next_discriminants = discriminants.clone().cycle().skip(1);
+        let next_discriminants =
+            discriminants.clone().chain(valid_discriminants.first().copied()).skip(1);
+        // Iterate over pairs of each discriminant together with the next one.
+        // Since they were sorted, we can now compute the niche sizes and pick the largest.
+        let discriminants = discriminants.zip(next_discriminants);
+        let largest_niche = discriminants.max_by_key(|&(start, end)| {
+            trace!(?start, ?end);
+            // If this is a wraparound range, the niche size is `MAX - abs(diff)`, as the diff between
+            // the two end points is actually the size of the valid discriminants.
+            let dist = if start > end {
+                // Overflow can happen for 128 bit discriminants if `end` is negative.
+                // But in that case casting to `u128` still gets us the right value,
+                // as the distance must be positive if the lhs of the subtraction is larger than the rhs.
+                let dist = start.wrapping_sub(end);
+                if discr_type.is_signed() {
+                    discr_int.signed_max().wrapping_sub(dist) as u128
+                } else {
+                    discr_int.size().unsigned_int_max() - dist as u128
+                }
+            } else {
+                // Overflow can happen for 128 bit discriminants if `start` is negative.
+                // But in that case casting to `u128` still gets us the right value,
+                // as the distance must be positive if the lhs of the subtraction is larger than the rhs.
+                end.wrapping_sub(start) as u128
+            };
+            trace!(?dist);
+            dist
+        });
+        trace!(?largest_niche);
+
+        // `max` is the last valid discriminant before the largest niche
+        // `min` is the first valid discriminant after the largest niche
+        let (max, min) = largest_niche
+            // We might have no inhabited variants, so pretend there's at least one.
+            .unwrap_or((0, 0));
         let (min_ity, signed) = discr_range_of_repr(min, max); //Integer::repr_discr(tcx, ty, &repr, min, max);
 
         let mut align = dl.aggregate_align;
diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
index 8e346706877..14e256b8045 100644
--- a/compiler/rustc_abi/src/lib.rs
+++ b/compiler/rustc_abi/src/lib.rs
@@ -1205,6 +1205,19 @@ impl Integer {
         }
     }
 
+    /// Returns the smallest signed value that can be represented by this Integer.
+    #[inline]
+    pub fn signed_min(self) -> i128 {
+        use Integer::*;
+        match self {
+            I8 => i8::MIN as i128,
+            I16 => i16::MIN as i128,
+            I32 => i32::MIN as i128,
+            I64 => i64::MIN as i128,
+            I128 => i128::MIN,
+        }
+    }
+
     /// Finds the smallest Integer type which can represent the signed value.
     #[inline]
     pub fn fit_signed(x: i128) -> Integer {
diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index a5123576fc6..aed94f9aa04 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -107,8 +107,8 @@ impl abi::Integer {
             abi::Integer::I8
         };
 
-        // If there are no negative values, we can use the unsigned fit.
-        if min >= 0 {
+        // Pick the smallest fit.
+        if unsigned_fit <= signed_fit {
             (cmp::max(unsigned_fit, at_least), false)
         } else {
             (cmp::max(signed_fit, at_least), true)
diff --git a/tests/ui/enum-discriminant/wrapping_niche.stderr b/tests/ui/enum-discriminant/wrapping_niche.stderr
index 76cafa43e90..e3e1755e14d 100644
--- a/tests/ui/enum-discriminant/wrapping_niche.stderr
+++ b/tests/ui/enum-discriminant/wrapping_niche.stderr
@@ -9,7 +9,7 @@ error: layout_of(UnsignedAroundZero) = Layout {
                        I16,
                        false,
                    ),
-                   valid_range: 0..=65535,
+                   valid_range: (..=1) | (65535..),
                },
            ),
            fields: Arbitrary {
@@ -20,7 +20,16 @@ error: layout_of(UnsignedAroundZero) = Layout {
                    0,
                ],
            },
-           largest_niche: None,
+           largest_niche: Some(
+               Niche {
+                   offset: Size(0 bytes),
+                   value: Int(
+                       I16,
+                       false,
+                   ),
+                   valid_range: (..=1) | (65535..),
+               },
+           ),
            uninhabited: false,
            variants: Multiple {
                tag: Initialized {
@@ -28,7 +37,7 @@ error: layout_of(UnsignedAroundZero) = Layout {
                        I16,
                        false,
                    ),
-                   valid_range: 0..=65535,
+                   valid_range: (..=1) | (65535..),
                },
                tag_encoding: Direct,
                tag_field: 0,
diff --git a/tests/ui/transmutability/enums/niche_optimization.rs b/tests/ui/transmutability/enums/niche_optimization.rs
index 2436be50027..316a857662a 100644
--- a/tests/ui/transmutability/enums/niche_optimization.rs
+++ b/tests/ui/transmutability/enums/niche_optimization.rs
@@ -75,8 +75,8 @@ fn one_niche() {
 
     assert::is_transmutable::<OptionLike, u8>();
     assert::is_transmutable::<V0, OptionLike>();
+    assert::is_transmutable::<V1, OptionLike>();
     assert::is_transmutable::<V254, OptionLike>();
-    assert::is_transmutable::<V255, OptionLike>();
 }
 
 fn one_niche_alt() {
@@ -97,9 +97,9 @@ fn one_niche_alt() {
     };
 
     assert::is_transmutable::<OptionLike, u8>();
-    assert::is_transmutable::<V0, OptionLike>();
+    assert::is_transmutable::<V1, OptionLike>();
+    assert::is_transmutable::<V2, OptionLike>();
     assert::is_transmutable::<V254, OptionLike>();
-    assert::is_transmutable::<V255, OptionLike>();
 }
 
 fn two_niche() {
@@ -121,9 +121,9 @@ fn two_niche() {
 
     assert::is_transmutable::<OptionLike, u8>();
     assert::is_transmutable::<V0, OptionLike>();
+    assert::is_transmutable::<V1, OptionLike>();
+    assert::is_transmutable::<V2, OptionLike>();
     assert::is_transmutable::<V253, OptionLike>();
-    assert::is_transmutable::<V254, OptionLike>();
-    assert::is_transmutable::<V255, OptionLike>();
 }
 
 fn no_niche() {
@@ -142,7 +142,7 @@ fn no_niche() {
     }
 
     const _: () = {
-        assert!(std::mem::size_of::<OptionLike>() == 2);
+        assert!(std::mem::size_of::<OptionLike>() == 1);
     };
 
     #[repr(C)]