From f2c95021855f6a2a54238d0c90cc3a6420d66d69 Mon Sep 17 00:00:00 2001
From: hkalbasi <hamidrezakalbasi@protonmail.com>
Date: Thu, 27 Oct 2022 23:28:34 +0330
Subject: [PATCH] support nonzero* niche optimizations

---
 crates/hir-ty/src/layout/adt.rs   | 50 +++++++++++++++++++++++++++++++
 crates/hir-ty/src/layout/tests.rs | 13 ++++++++
 crates/test-utils/src/minicore.rs | 10 +++++++
 3 files changed, 73 insertions(+)

diff --git a/crates/hir-ty/src/layout/adt.rs b/crates/hir-ty/src/layout/adt.rs
index e353034eb9..9244353f3a 100644
--- a/crates/hir-ty/src/layout/adt.rs
+++ b/crates/hir-ty/src/layout/adt.rs
@@ -4,6 +4,7 @@ use std::{
     cmp::{self, Ordering},
     iter,
     num::NonZeroUsize,
+    ops::Bound,
 };
 
 use chalk_ir::TyKind;
@@ -18,6 +19,8 @@ use hir_def::{
 };
 use la_arena::{ArenaMap, RawIdx};
 
+struct X(Option<NonZeroUsize>);
+
 use crate::{
     db::HirDatabase,
     lang_items::is_unsafe_cell,
@@ -137,7 +140,38 @@ pub fn layout_of_adt_query(
                 Abi::Aggregate { sized: _ } => {}
             }
             st.largest_niche = None;
+            return Ok(st);
         }
+
+        let (start, end) = layout_scalar_valid_range(db, def);
+        match st.abi {
+            Abi::Scalar(ref mut scalar) | Abi::ScalarPair(ref mut scalar, _) => {
+                if let Bound::Included(start) = start {
+                    let valid_range = scalar.valid_range_mut();
+                    valid_range.start = start;
+                }
+                if let Bound::Included(end) = end {
+                    let valid_range = scalar.valid_range_mut();
+                    valid_range.end = end;
+                }
+                // Update `largest_niche` if we have introduced a larger niche.
+                let niche = Niche::from_scalar(dl, Size::ZERO, *scalar);
+                if let Some(niche) = niche {
+                    match st.largest_niche {
+                        Some(largest_niche) => {
+                            // Replace the existing niche even if they're equal,
+                            // because this one is at a lower offset.
+                            if largest_niche.available(dl) <= niche.available(dl) {
+                                st.largest_niche = Some(niche);
+                            }
+                        }
+                        None => st.largest_niche = Some(niche),
+                    }
+                }
+            }
+            _ => user_error!("nonscalar layout for layout_scalar_valid_range"),
+        }
+
         return Ok(st);
     }
 
@@ -591,6 +625,22 @@ pub fn layout_of_adt_query(
     Ok(best_layout.layout)
 }
 
+fn layout_scalar_valid_range(db: &dyn HirDatabase, def: AdtId) -> (Bound<u128>, Bound<u128>) {
+    let attrs = db.attrs(def.into());
+    let get = |name| {
+        let attr = attrs.by_key(name).tt_values();
+        for tree in attr {
+            if let Some(x) = tree.token_trees.first() {
+                if let Ok(x) = x.to_string().parse() {
+                    return Bound::Included(x);
+                }
+            }
+        }
+        Bound::Unbounded
+    };
+    (get("rustc_layout_scalar_valid_range_start"), get("rustc_layout_scalar_valid_range_end"))
+}
+
 pub fn layout_of_adt_recover(
     _: &dyn HirDatabase,
     _: &[String],
diff --git a/crates/hir-ty/src/layout/tests.rs b/crates/hir-ty/src/layout/tests.rs
index 9543b4dcbc..1cd6d4eae2 100644
--- a/crates/hir-ty/src/layout/tests.rs
+++ b/crates/hir-ty/src/layout/tests.rs
@@ -146,6 +146,19 @@ fn tuple() {
     }
 }
 
+#[test]
+fn non_zero() {
+    check_size_and_align(
+        r#"
+    //- minicore: non_zero, option
+    use core::num::NonZeroU8;
+    struct Goal(Option<NonZeroU8>);
+    "#,
+        1,
+        1,
+    );
+}
+
 #[test]
 fn niche_optimization() {
     check_size_and_align(
diff --git a/crates/test-utils/src/minicore.rs b/crates/test-utils/src/minicore.rs
index 69d2e62b25..af9efd2600 100644
--- a/crates/test-utils/src/minicore.rs
+++ b/crates/test-utils/src/minicore.rs
@@ -29,6 +29,7 @@
 //!     index: sized
 //!     iterator: option
 //!     iterators: iterator, fn
+//!     non_zero:
 //!     option:
 //!     ord: eq, option
 //!     pin:
@@ -680,6 +681,15 @@ mod macros {
 }
 // endregion:derive
 
+// region:non_zero
+pub mod num {
+    #[repr(transparent)]
+    #[rustc_layout_scalar_valid_range_start(1)]
+    #[rustc_nonnull_optimization_guaranteed]
+    pub struct NonZeroU8(u8);
+}
+// endregion:non_zero
+
 // region:bool_impl
 #[lang = "bool"]
 impl bool {