Rollup merge of #133964 - joboet:select_unpredictable, r=tgross35

core: implement `bool::select_unpredictable`

Tracking issue: #133962
ACP: https://github.com/rust-lang/libs-team/issues/468

4 files changed, 86 insertions, 3 deletions

diff --git a/library/core/src/bool.rs b/library/core/src/bool.rs
index 1590b9f29fc..3c589ca5dfa 100644
--- a/library/core/src/bool.rs
+++ b/library/core/src/bool.rs
@@ -61,4 +61,52 @@ impl bool {
     pub fn then<T, F: FnOnce() -> T>(self, f: F) -> Option<T> {
         if self { Some(f()) } else { None }
     }
+
+    /// Returns either `true_val` or `false_val` depending on the value of
+    /// `self`, with a hint to the compiler that `self` is unlikely
+    /// to be correctly predicted by a CPU’s branch predictor.
+    ///
+    /// This method is functionally equivalent to
+    /// ```ignore (this is just for illustrative purposes)
+    /// fn select_unpredictable<T>(b: bool, true_val: T, false_val: T) -> T {
+    ///     if b { true_val } else { false_val }
+    /// }
+    /// ```
+    /// but might generate different assembly. In particular, on platforms with
+    /// a conditional move or select instruction (like `cmov` on x86 or `csel`
+    /// on ARM) the optimizer might use these instructions to avoid branches,
+    /// which can benefit performance if the branch predictor is struggling
+    /// with predicting `condition`, such as in an implementation of  binary
+    /// search.
+    ///
+    /// Note however that this lowering is not guaranteed (on any platform) and
+    /// should not be relied upon when trying to write constant-time code. Also
+    /// be aware that this lowering might *decrease* performance if `condition`
+    /// is well-predictable. It is advisable to perform benchmarks to tell if
+    /// this function is useful.
+    ///
+    /// # Examples
+    ///
+    /// Distribute values evenly between two buckets:
+    /// ```
+    /// #![feature(select_unpredictable)]
+    ///
+    /// use std::hash::BuildHasher;
+    ///
+    /// fn append<H: BuildHasher>(hasher: &H, v: i32, bucket_one: &mut Vec<i32>, bucket_two: &mut Vec<i32>) {
+    ///     let hash = hasher.hash_one(&v);
+    ///     let bucket = (hash % 2 == 0).select_unpredictable(bucket_one, bucket_two);
+    ///     bucket.push(v);
+    /// }
+    /// # let hasher = std::collections::hash_map::RandomState::new();
+    /// # let mut bucket_one = Vec::new();
+    /// # let mut bucket_two = Vec::new();
+    /// # append(&hasher, 42, &mut bucket_one, &mut bucket_two);
+    /// # assert_eq!(bucket_one.len() + bucket_two.len(), 1);
+    /// ```
+    #[inline(always)]
+    #[unstable(feature = "select_unpredictable", issue = "133962")]
+    pub fn select_unpredictable<T>(self, true_val: T, false_val: T) -> T {
+        crate::intrinsics::select_unpredictable(self, true_val, false_val)
+    }
 }
diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs
index b5c31d82467..7d44ea7a66b 100644
--- a/library/core/src/intrinsics/mod.rs
+++ b/library/core/src/intrinsics/mod.rs
@@ -1545,7 +1545,7 @@ pub const fn unlikely(b: bool) -> bool {
 /// Therefore, implementations must not require the user to uphold
 /// any safety invariants.
 ///
-/// This intrinsic does not have a stable counterpart.
+/// The public form of this instrinsic is [`bool::select_unpredictable`].
 #[unstable(feature = "core_intrinsics", issue = "none")]
 #[rustc_intrinsic]
 #[rustc_nounwind]
diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs
index df9720698d3..073cca7daef 100644
--- a/library/core/src/slice/mod.rs
+++ b/library/core/src/slice/mod.rs
@@ -7,7 +7,7 @@
 #![stable(feature = "rust1", since = "1.0.0")]
 
 use crate::cmp::Ordering::{self, Equal, Greater, Less};
-use crate::intrinsics::{exact_div, select_unpredictable, unchecked_sub};
+use crate::intrinsics::{exact_div, unchecked_sub};
 use crate::mem::{self, SizedTypeProperties};
 use crate::num::NonZero;
 use crate::ops::{Bound, OneSidedRange, Range, RangeBounds, RangeInclusive};
@@ -2835,7 +2835,7 @@ impl<T> [T] {
             // Binary search interacts poorly with branch prediction, so force
             // the compiler to use conditional moves if supported by the target
             // architecture.
-            base = select_unpredictable(cmp == Greater, base, mid);
+            base = (cmp == Greater).select_unpredictable(base, mid);
 
             // This is imprecise in the case where `size` is odd and the
             // comparison returns Greater: the mid element still gets included
diff --git a/tests/codegen/bool-select-unpredictable.rs b/tests/codegen/bool-select-unpredictable.rs
new file mode 100644
index 00000000000..1562b177542
--- /dev/null
+++ b/tests/codegen/bool-select-unpredictable.rs
@@ -0,0 +1,35 @@
+//@ compile-flags: -O
+
+#![feature(select_unpredictable)]
+#![crate_type = "lib"]
+
+#[no_mangle]
+pub fn test_int(p: bool, a: u64, b: u64) -> u64 {
+    // CHECK-LABEL: define{{.*}} @test_int
+    // CHECK: select i1 %p, i64 %a, i64 %b, !unpredictable
+    p.select_unpredictable(a, b)
+}
+
+#[no_mangle]
+pub fn test_pair(p: bool, a: (u64, u64), b: (u64, u64)) -> (u64, u64) {
+    // CHECK-LABEL: define{{.*}} @test_pair
+    // CHECK: select i1 %p, {{.*}}, !unpredictable
+    p.select_unpredictable(a, b)
+}
+
+struct Large {
+    e: [u64; 100],
+}
+
+#[no_mangle]
+pub fn test_struct(p: bool, a: Large, b: Large) -> Large {
+    // CHECK-LABEL: define{{.*}} @test_struct
+    // CHECK: select i1 %p, {{.*}}, !unpredictable
+    p.select_unpredictable(a, b)
+}
+
+#[no_mangle]
+pub fn test_zst(p: bool, a: (), b: ()) -> () {
+    // CHECK-LABEL: define{{.*}} @test_zst
+    p.select_unpredictable(a, b)
+}