From c7af338e6fa95dd9877a19be74b94f471233d75b Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Fri, 30 Sep 2022 23:39:15 -0700
Subject: [PATCH] Tell LLVM that `partition_point` returns a valid fencepost

This was already done for a successful `binary_search`, but this way `partition_point` can get similar optimizations.
---
 library/core/src/slice/mod.rs                 | 15 +++++++++++---
 .../binary-search-index-no-bound-check.rs     | 20 +++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs
index cd04fa00442..b2289c6bed7 100644
--- a/library/core/src/slice/mod.rs
+++ b/library/core/src/slice/mod.rs
@@ -2426,15 +2426,20 @@ impl<T> [T] {
     where
         F: FnMut(&'a T) -> Ordering,
     {
+        // INVARIANTS:
+        // - 0 <= left <= left + size = right <= self.len()
+        // - f returns Less for everything in self[..left]
+        // - f returns Greater for everything in self[right..]
         let mut size = self.len();
         let mut left = 0;
         let mut right = size;
         while left < right {
             let mid = left + size / 2;
 
-            // SAFETY: the call is made safe by the following invariants:
-            // - `mid >= 0`
-            // - `mid < size`: `mid` is limited by `[left; right)` bound.
+            // SAFETY: the while condition means `size` is strictly positive, so
+            // `size/2 < size`.  Thus `left + size/2 < left + size`, which
+            // coupled with the `left + size <= self.len()` invariant means
+            // we have `left + size/2 < self.len()`, and this is in-bounds.
             let cmp = f(unsafe { self.get_unchecked(mid) });
 
             // The reason why we use if/else control flow rather than match
@@ -2452,6 +2457,10 @@ impl<T> [T] {
 
             size = right - left;
         }
+
+        // SAFETY: directly true from the overall invariant.
+        // Note that this is `<=`, unlike the assume in the `Ok` path.
+        unsafe { crate::intrinsics::assume(left <= self.len()) };
         Err(left)
     }
 
diff --git a/src/test/codegen/binary-search-index-no-bound-check.rs b/src/test/codegen/binary-search-index-no-bound-check.rs
index 2deabcaa6c2..c1766a4a44a 100644
--- a/src/test/codegen/binary-search-index-no-bound-check.rs
+++ b/src/test/codegen/binary-search-index-no-bound-check.rs
@@ -16,3 +16,23 @@ pub fn binary_search_index_no_bounds_check(s: &[u8]) -> u8 {
         42
     }
 }
+
+// Similarly, check that `partition_point` is known to return a valid fencepost.
+
+// CHECK-LABEL: @unknown_split
+#[no_mangle]
+pub fn unknown_split(x: &[i32], i: usize) -> (&[i32], &[i32]) {
+    // This just makes sure that the subsequent function is looking for the
+    // absence of something that might actually be there.
+
+    // CHECK: call core::panicking::panic
+    x.split_at(i)
+}
+
+// CHECK-LABEL: @partition_point_split_no_bounds_check
+#[no_mangle]
+pub fn partition_point_split_no_bounds_check(x: &[i32], needle: i32) -> (&[i32], &[i32]) {
+    // CHECK-NOT: call core::panicking::panic
+    let i = x.partition_point(|p| p < &needle);
+    x.split_at(i)
+}