Rollup merge of #138024 - reitermarkus:unicode-panic-optimization, r=ibraheemdev

Allow optimizing out `panic_bounds_check` in Unicode checks.

Allow optimizing out `panic_bounds_check` in Unicode checks.

For context, see https://github.com/japaric/ufmt/issues/52#issuecomment-2699207241.

2 files changed, 98 insertions, 39 deletions

diff --git a/src/tools/unicode-table-generator/src/range_search.rs b/src/tools/unicode-table-generator/src/range_search.rs
index 9a51979a2f0..02f9cf16d4d 100644
--- a/src/tools/unicode-table-generator/src/range_search.rs
+++ b/src/tools/unicode-table-generator/src/range_search.rs
@@ -45,45 +45,78 @@ const fn bitset_search<
     (word & (1 << (needle % 64) as u64)) != 0
 }
 
-fn decode_prefix_sum(short_offset_run_header: u32) -> u32 {
-    short_offset_run_header & ((1 << 21) - 1)
-}
+#[repr(transparent)]
+struct ShortOffsetRunHeader(u32);
+
+impl ShortOffsetRunHeader {
+    const fn new(start_index: usize, prefix_sum: u32) -> Self {
+        assert!(start_index < (1 << 11));
+        assert!(prefix_sum < (1 << 21));
 
-fn decode_length(short_offset_run_header: u32) -> usize {
-    (short_offset_run_header >> 21) as usize
+        Self((start_index as u32) << 21 | prefix_sum)
+    }
+
+    #[inline]
+    const fn start_index(&self) -> usize {
+        (self.0 >> 21) as usize
+    }
+
+    #[inline]
+    const fn prefix_sum(&self) -> u32 {
+        self.0 & ((1 << 21) - 1)
+    }
 }
 
+/// # Safety
+///
+/// - The last element of `short_offset_runs` must be greater than `std::char::MAX`.
+/// - The start indices of all elements in `short_offset_runs` must be less than `OFFSETS`.
 #[inline(always)]
-fn skip_search<const SOR: usize, const OFFSETS: usize>(
-    needle: u32,
-    short_offset_runs: &[u32; SOR],
+unsafe fn skip_search<const SOR: usize, const OFFSETS: usize>(
+    needle: char,
+    short_offset_runs: &[ShortOffsetRunHeader; SOR],
     offsets: &[u8; OFFSETS],
 ) -> bool {
-    // Note that this *cannot* be past the end of the array, as the last
-    // element is greater than std::char::MAX (the largest possible needle).
-    //
-    // So, we cannot have found it (i.e. Ok(idx) + 1 != length) and the correct
-    // location cannot be past it, so Err(idx) != length either.
-    //
-    // This means that we can avoid bounds checking for the accesses below, too.
+    let needle = needle as u32;
+
     let last_idx =
-        match short_offset_runs.binary_search_by_key(&(needle << 11), |header| header << 11) {
+        match short_offset_runs.binary_search_by_key(&(needle << 11), |header| (header.0 << 11)) {
             Ok(idx) => idx + 1,
             Err(idx) => idx,
         };
+    // SAFETY: `last_idx` *cannot* be past the end of the array, as the last
+    // element is greater than `std::char::MAX` (the largest possible needle)
+    // as guaranteed by the caller.
+    //
+    // So, we cannot have found it (i.e. `Ok(idx) => idx + 1 != length`) and the
+    // correct location cannot be past it, so `Err(idx) => idx != length` either.
+    //
+    // This means that we can avoid bounds checking for the accesses below, too.
+    //
+    // We need to use `intrinsics::assume` since the `panic_nounwind` contained
+    // in `hint::assert_unchecked` may not be optimized out.
+    unsafe { crate::intrinsics::assume(last_idx < SOR) };
 
-    let mut offset_idx = decode_length(short_offset_runs[last_idx]);
+    let mut offset_idx = short_offset_runs[last_idx].start_index();
     let length = if let Some(next) = short_offset_runs.get(last_idx + 1) {
-        decode_length(*next) - offset_idx
+        (*next).start_index() - offset_idx
     } else {
         offsets.len() - offset_idx
     };
+
     let prev =
-        last_idx.checked_sub(1).map(|prev| decode_prefix_sum(short_offset_runs[prev])).unwrap_or(0);
+        last_idx.checked_sub(1).map(|prev| short_offset_runs[prev].prefix_sum()).unwrap_or(0);
 
     let total = needle - prev;
     let mut prefix_sum = 0;
     for _ in 0..(length - 1) {
+        // SAFETY: It is guaranteed that `length <= OFFSETS - offset_idx`,
+        // so it follows that `length - 1 + offset_idx < OFFSETS`, therefore
+        // `offset_idx < OFFSETS` is always true in this loop.
+        //
+        // We need to use `intrinsics::assume` since the `panic_nounwind` contained
+        // in `hint::assert_unchecked` may not be optimized out.
+        unsafe { crate::intrinsics::assume(offset_idx < OFFSETS) };
         let offset = offsets[offset_idx];
         prefix_sum += offset as u32;
         if prefix_sum > total {
diff --git a/src/tools/unicode-table-generator/src/skiplist.rs b/src/tools/unicode-table-generator/src/skiplist.rs
index 8ca18ddc91a..34c9802e122 100644
--- a/src/tools/unicode-table-generator/src/skiplist.rs
+++ b/src/tools/unicode-table-generator/src/skiplist.rs
@@ -1,26 +1,23 @@
-use std::fmt::Write as _;
+use std::fmt::{self, Write as _};
 use std::ops::Range;
 
 use crate::fmt_list;
 use crate::raw_emitter::RawEmitter;
 
 /// This will get packed into a single u32 before inserting into the data set.
-#[derive(Debug, PartialEq)]
+#[derive(PartialEq)]
 struct ShortOffsetRunHeader {
-    /// Note, we only allow for 21 bits here.
-    prefix_sum: u32,
-
     /// Note, we actually only allow for 11 bits here. This should be enough --
     /// our largest sets are around ~1400 offsets long.
-    start_idx: u16,
-}
+    start_index: u16,
 
-impl ShortOffsetRunHeader {
-    fn pack(&self) -> u32 {
-        assert!(self.start_idx < (1 << 11));
-        assert!(self.prefix_sum < (1 << 21));
+    /// Note, we only allow for 21 bits here.
+    prefix_sum: u32,
+}
 
-        (self.start_idx as u32) << 21 | self.prefix_sum
+impl fmt::Debug for ShortOffsetRunHeader {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "ShortOffsetRunHeader::new({}, {})", self.start_index, self.prefix_sum)
     }
 }
 
@@ -53,7 +50,7 @@ impl RawEmitter {
                     coded_offsets.push(offset);
                 } else {
                     short_offset_runs.push(ShortOffsetRunHeader {
-                        start_idx: start.try_into().unwrap(),
+                        start_index: start.try_into().unwrap(),
                         prefix_sum,
                     });
                     // This is just needed to maintain indices even/odd
@@ -71,11 +68,12 @@ impl RawEmitter {
             assert!(inserted);
         }
 
+        writeln!(&mut self.file, "use super::ShortOffsetRunHeader;\n").unwrap();
         writeln!(
             &mut self.file,
-            "static SHORT_OFFSET_RUNS: [u32; {}] = [{}];",
+            "static SHORT_OFFSET_RUNS: [ShortOffsetRunHeader; {}] = [{}];",
             short_offset_runs.len(),
-            fmt_list(short_offset_runs.iter().map(|v| v.pack()))
+            fmt_list(short_offset_runs.iter())
         )
         .unwrap();
         self.bytes_used += 4 * short_offset_runs.len();
@@ -104,15 +102,43 @@ impl RawEmitter {
             writeln!(&mut self.file, "    (c as u32) >= {first_code_point:#04x} && lookup_slow(c)")
                 .unwrap();
             writeln!(&mut self.file, "}}").unwrap();
+            writeln!(&mut self.file).unwrap();
+            writeln!(&mut self.file, "#[inline(never)]").unwrap();
             writeln!(&mut self.file, "fn lookup_slow(c: char) -> bool {{").unwrap();
         } else {
             writeln!(&mut self.file, "pub fn lookup(c: char) -> bool {{").unwrap();
         }
-        writeln!(&mut self.file, "    super::skip_search(",).unwrap();
-        writeln!(&mut self.file, "        c as u32,").unwrap();
-        writeln!(&mut self.file, "        &SHORT_OFFSET_RUNS,").unwrap();
-        writeln!(&mut self.file, "        &OFFSETS,").unwrap();
-        writeln!(&mut self.file, "    )").unwrap();
+        writeln!(&mut self.file, "    const {{").unwrap();
+        writeln!(
+            &mut self.file,
+            "        assert!(SHORT_OFFSET_RUNS.last().unwrap().0 > char::MAX as u32);",
+        )
+        .unwrap();
+        writeln!(&mut self.file, "        let mut i = 0;").unwrap();
+        writeln!(&mut self.file, "        while i < SHORT_OFFSET_RUNS.len() {{").unwrap();
+        writeln!(
+            &mut self.file,
+            "            assert!(SHORT_OFFSET_RUNS[i].start_index() < OFFSETS.len());",
+        )
+        .unwrap();
+        writeln!(&mut self.file, "            i += 1;").unwrap();
+        writeln!(&mut self.file, "        }}").unwrap();
+        writeln!(&mut self.file, "    }}").unwrap();
+        writeln!(
+            &mut self.file,
+            "    // SAFETY: We just ensured the last element of `SHORT_OFFSET_RUNS` is greater than `std::char::MAX`",
+        )
+        .unwrap();
+        writeln!(
+            &mut self.file,
+            "    // and the start indices of all elements in `SHORT_OFFSET_RUNS` are smaller than `OFFSETS.len()`.",
+        )
+        .unwrap();
+        writeln!(
+            &mut self.file,
+            "    unsafe {{ super::skip_search(c, &SHORT_OFFSET_RUNS, &OFFSETS) }}"
+        )
+        .unwrap();
         writeln!(&mut self.file, "}}").unwrap();
     }
 }