Add new intrinsic `is_constant` and optimize `pow`

Fix overflow check

Make MIRI choose the path randomly and rename the intrinsic

Add back test

Add miri test and make it operate on `ptr`

Define `llvm.is.constant` for primitives

Update MIRI comment and fix test in stage2

Add const eval test

Clarify that both branches must have the same side effects

guaranteed non guarantee

use immediate type instead

Co-Authored-By: Ralf Jung <post@ralfj.de>

compiler/rustc_codegen_llvm/src/context.rs

@@ -908,6 +908,20 @@ impl<'ll> CodegenCx<'ll, '_> {
         ifn!("llvm.lifetime.start.p0i8", fn(t_i64, ptr) -> void);
         ifn!("llvm.lifetime.end.p0i8", fn(t_i64, ptr) -> void);
 
+        // FIXME: This is an infinitesimally small portion of the types you can
+        // pass to this intrinsic, if we can ever lazily register intrinsics we
+        // should register these when they're used, that way any type can be
+        // passed.
+        ifn!("llvm.is.constant.i1", fn(i1) -> i1);
+        ifn!("llvm.is.constant.i8", fn(t_i8) -> i1);
+        ifn!("llvm.is.constant.i16", fn(t_i16) -> i1);
+        ifn!("llvm.is.constant.i32", fn(t_i32) -> i1);
+        ifn!("llvm.is.constant.i64", fn(t_i64) -> i1);
+        ifn!("llvm.is.constant.i128", fn(t_i128) -> i1);
+        ifn!("llvm.is.constant.isize", fn(t_isize) -> i1);
+        ifn!("llvm.is.constant.f32", fn(t_f32) -> i1);
+        ifn!("llvm.is.constant.f64", fn(t_f64) -> i1);
+
         ifn!("llvm.expect.i1", fn(i1, i1) -> i1);
         ifn!("llvm.eh.typeid.for", fn(ptr) -> t_i32);
         ifn!("llvm.localescape", fn(...) -> void);

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -119,6 +119,10 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> {
             sym::likely => {
                 self.call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(true)])
             }
+            sym::is_val_statically_known => self.call_intrinsic(
+                &format!("llvm.is.constant.{:?}", args[0].layout.immediate_llvm_type(self.cx)),
+                &[args[0].immediate()],
+            ),
             sym::unlikely => self
                 .call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(false)]),
             kw::Try => {

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -531,6 +531,11 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
                     )?;
                 }
             }
+            // The intrinsic represents whether the value is known to the optimizer (LLVM).
+            // We're not doing any optimizations here, so there is no optimizer that could know the value.
+            // (We know the value here in the machine of course, but this is the runtime of that code,
+            // not the optimization stage.)
+            sym::is_val_statically_known => ecx.write_scalar(Scalar::from_bool(false), dest)?,
             _ => {
                 throw_unsup_format!(
                     "intrinsic `{intrinsic_name}` is not supported at compile-time"

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -453,6 +453,8 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
 
             sym::black_box => (1, vec![param(0)], param(0)),
 
+            sym::is_val_statically_known => (1, vec![param(0)], tcx.types.bool),
+
             sym::const_eval_select => (4, vec![param(0), param(1), param(2)], param(3)),
 
             sym::vtable_size | sym::vtable_align => {

compiler/rustc_span/src/symbol.rs

@@ -907,6 +907,7 @@ symbols! {
         io_stderr,
         io_stdout,
         irrefutable_let_patterns,
+        is_val_statically_known,
         isa_attribute,
         isize,
         issue,

library/core/src/intrinsics.rs

@@ -2511,6 +2511,52 @@ extern "rust-intrinsic" {
     where
         G: FnOnce<ARG, Output = RET>,
         F: FnOnce<ARG, Output = RET>;
+
+    /// Returns whether the argument's value is statically known at
+    /// compile-time.
+    ///
+    /// This is useful when there is a way of writing the code that will
+    /// be *faster* when some variables have known values, but *slower*
+    /// in the general case: an `if is_val_statically_known(var)` can be used
+    /// to select between these two variants. The `if` will be optimized away
+    /// and only the desired branch remains.
+    ///
+    /// Formally speaking, this function non-deterministically returns `true`
+    /// or `false`, and the caller has to ensure sound behavior for both cases.
+    /// In other words, the following code has *Undefined Behavior*:
+    ///
+    /// ```rust
+    /// if !is_val_statically_known(0) { unreachable_unchecked(); }
+    /// ```
+    ///
+    /// This also means that the following code's behavior is unspecified; it
+    /// may panic, or it may not:
+    ///
+    /// ```rust,no_run
+    /// assert_eq!(is_val_statically_known(0), black_box(is_val_statically_known(0)))
+    /// ```
+    ///
+    /// Unsafe code may not rely on `is_val_statically_known` returning any
+    /// particular value, ever. However, the compiler will generally make it
+    /// return `true` only if the value of the argument is actually known.
+    ///
+    /// When calling this in a `const fn`, both paths must be semantically
+    /// equivalent, that is, the result of the `true` branch and the `false`
+    /// branch must return the same value and have the same side-effects *no
+    /// matter what*.
+    #[rustc_const_unstable(feature = "is_val_statically_known", issue = "none")]
+    #[rustc_nounwind]
+    #[cfg(not(bootstrap))]
+    pub fn is_val_statically_known<T>(arg: T) -> bool;
+}
+
+// FIXME: Seems using `unstable` here completely ignores `rustc_allow_const_fn_unstable`
+// and thus compiling stage0 core doesn't work.
+#[rustc_const_stable(feature = "is_val_statically_known", since = "never")]
+#[cfg(bootstrap)]
+pub const unsafe fn is_val_statically_known<T>(t: T) -> bool {
+    mem::forget(t);
+    false
 }
 
 // Some functions are defined here because they accidentally got made

library/core/src/lib.rs

@@ -197,6 +197,7 @@
 //
 // Language features:
 // tidy-alphabetical-start
+#![cfg_attr(not(bootstrap), feature(is_val_statically_known))]
 #![feature(abi_unadjusted)]
 #![feature(adt_const_params)]
 #![feature(allow_internal_unsafe)]

library/core/src/num/int_macros.rs

@@ -2088,26 +2088,49 @@ macro_rules! int_impl {
                       without modifying the original"]
         #[inline]
         #[rustc_inherit_overflow_checks]
+        #[rustc_allow_const_fn_unstable(is_val_statically_known)]
         pub const fn pow(self, mut exp: u32) -> Self {
-            if exp == 0 {
-                return 1;
-            }
-            let mut base = self;
-            let mut acc = 1;
-
-            while exp > 1 {
-                if (exp & 1) == 1 {
-                    acc = acc * base;
+            // SAFETY: This path has the same behavior as the other.
+            if unsafe { intrinsics::is_val_statically_known(self) }
+                && self > 0
+                && (self & (self - 1) == 0)
+            {
+                let power_used = match self.checked_ilog2() {
+                    Some(v) => v,
+                    // SAFETY: We just checked this is a power of two. and above zero.
+                    None => unsafe { core::hint::unreachable_unchecked() },
+                };
+                // So it panics. Have to use `overflowing_mul` to efficiently set the
+                // result to 0 if not.
+                #[cfg(debug_assertions)]
+                {
+                    _ = power_used * exp;
                 }
-                exp /= 2;
-                base = base * base;
-            }
+                let (num_shl, overflowed) = power_used.overflowing_mul(exp);
+                let fine = !overflowed
+                    & (num_shl < (mem::size_of::<Self>() * 8) as u32);
+                (1 << num_shl) * fine as Self
+            } else {
+                if exp == 0 {
+                    return 1;
+                }
+                let mut base = self;
+                let mut acc = 1;
 
-            // since exp!=0, finally the exp must be 1.
-            // Deal with the final bit of the exponent separately, since
-            // squaring the base afterwards is not necessary and may cause a
-            // needless overflow.
-            acc * base
+                while exp > 1 {
+                    if (exp & 1) == 1 {
+                        acc = acc * base;
+                    }
+                    exp /= 2;
+                    base = base * base;
+                }
+
+                // since exp!=0, finally the exp must be 1.
+                // Deal with the final bit of the exponent separately, since
+                // squaring the base afterwards is not necessary and may cause a
+                // needless overflow.
+                acc * base
+            }
         }
 
         /// Returns the square root of the number, rounded down.

library/core/src/num/uint_macros.rs

@@ -1973,26 +1973,60 @@ macro_rules! uint_impl {
                       without modifying the original"]
         #[inline]
         #[rustc_inherit_overflow_checks]
+        #[rustc_allow_const_fn_unstable(is_val_statically_known)]
         pub const fn pow(self, mut exp: u32) -> Self {
-            if exp == 0 {
-                return 1;
-            }
-            let mut base = self;
-            let mut acc = 1;
-
-            while exp > 1 {
-                if (exp & 1) == 1 {
-                    acc = acc * base;
+            // LLVM now knows that `self` is a constant value, but not a
+            // constant in Rust. This allows us to compute the power used at
+            // compile-time.
+            //
+            // This will likely add a branch in debug builds, but this should
+            // be ok.
+            //
+            // This is a massive performance boost in release builds as you can
+            // get the power of a power of two and the exponent through a `shl`
+            // instruction, but we must add a couple more checks for parity with
+            // our own `pow`.
+            // SAFETY: This path has the same behavior as the other.
+            if unsafe { intrinsics::is_val_statically_known(self) }
+                && self.is_power_of_two()
+            {
+                let power_used = match self.checked_ilog2() {
+                    Some(v) => v,
+                    // SAFETY: We just checked this is a power of two. `0` is not a
+                    // power of two.
+                    None => unsafe { core::hint::unreachable_unchecked() },
+                };
+                // So it panics. Have to use `overflowing_mul` to efficiently set the
+                // result to 0 if not.
+                #[cfg(debug_assertions)]
+                {
+                    _ = power_used * exp;
                 }
-                exp /= 2;
-                base = base * base;
-            }
+                let (num_shl, overflowed) = power_used.overflowing_mul(exp);
+                let fine = !overflowed
+                    & (num_shl < (mem::size_of::<Self>() * 8) as u32);
+                (1 << num_shl) * fine as Self
+            } else {
+                if exp == 0 {
+                    return 1;
+                }
+                let mut base = self;
+                let mut acc = 1;
 
-            // since exp!=0, finally the exp must be 1.
-            // Deal with the final bit of the exponent separately, since
-            // squaring the base afterwards is not necessary and may cause a
-            // needless overflow.
-            acc * base
+                while exp > 1 {
+                    if (exp & 1) == 1 {
+                        acc = acc * base;
+                    }
+                    exp /= 2;
+                    base = base * base;
+                }
+
+                // since exp!=0, finally the exp must be 1.
+                // Deal with the final bit of the exponent separately, since
+                // squaring the base afterwards is not necessary and may cause a
+                // needless overflow.
+                acc * base
+            }
         }
 
         /// Returns the square root of the number, rounded down.

src/tools/miri/src/shims/intrinsics/mod.rs

@@ -5,6 +5,7 @@ use std::iter;
 
 use log::trace;
 
+use rand::Rng;
 use rustc_apfloat::{Float, Round};
 use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::{
@@ -141,6 +142,17 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 this.write_pointer(Pointer::new(ptr.provenance, masked_addr), dest)?;
             }
 
+            // We want to return either `true` or `false` at random, or else something like
+            // ```
+            // if !is_val_statically_known(0) { unreachable_unchecked(); }
+            // ```
+            // Would not be considered UB, or the other way around (`is_val_statically_known(0)`).
+            "is_val_statically_known" => {
+                let [_] = check_arg_count(args)?;
+                let branch: bool = this.machine.rng.get_mut().gen();
+                this.write_scalar(Scalar::from_bool(branch), dest)?;
+            }
+
             // Floating-point operations
             "fabsf32" => {
                 let [f] = check_arg_count(args)?;

src/tools/miri/tests/pass/intrinsics.rs

@@ -33,6 +33,21 @@ fn main() {
     assert_eq!(intrinsics::likely(false), false);
     assert_eq!(intrinsics::unlikely(true), true);
 
+    let mut saw_true = false;
+    let mut saw_false = false;
+
+    for _ in 0..50 {
+        if unsafe { intrinsics::is_val_statically_known(0) } {
+            saw_true = true;
+        } else {
+            saw_false = true;
+        }
+    }
+    assert!(
+        saw_true && saw_false,
+        "`is_val_statically_known` failed to return both true and false. Congrats, you won the lottery!"
+    );
+
     intrinsics::forget(Bomb);
 
     let _v = intrinsics::discriminant_value(&Some(()));

tests/codegen/is_val_statically_known.rs

@@ -0,0 +1,50 @@
+// #[cfg(bootstrap)]
+// ignore-stage1
+// compile-flags: --crate-type=lib -Zmerge-functions=disabled
+
+#![feature(core_intrinsics)]
+
+use std::intrinsics::is_val_statically_known;
+
+pub struct A(u32);
+pub enum B {
+    Ye(u32),
+}
+
+#[inline]
+pub fn _u32(a: u32) -> i32 {
+    if unsafe { is_val_statically_known(a) } { 1 } else { 0 }
+}
+
+// CHECK-LABEL: @_u32_true(
+#[no_mangle]
+pub fn _u32_true() -> i32 {
+    // CHECK: ret i32 1
+    _u32(1)
+}
+
+// CHECK-LABEL: @_u32_false(
+#[no_mangle]
+pub fn _u32_false(a: u32) -> i32 {
+    // CHECK: ret i32 0
+    _u32(a)
+}
+
+#[inline]
+pub fn _bool(b: bool) -> i32 {
+    if unsafe { is_val_statically_known(b) } { 3 } else { 2 }
+}
+
+// CHECK-LABEL: @_bool_true(
+#[no_mangle]
+pub fn _bool_true() -> i32 {
+    // CHECK: ret i32 3
+    _bool(true)
+}
+
+// CHECK-LABEL: @_bool_false(
+#[no_mangle]
+pub fn _bool_false(b: bool) -> i32 {
+    // CHECK: ret i32 2
+    _bool(b)
+}

tests/codegen/pow_of_two.rs

@@ -0,0 +1,68 @@
+// #[cfg(bootstrap)]
+// ignore-stage1
+// compile-flags: --crate-type=lib -Zmerge-functions=disabled
+
+// CHECK-LABEL: @a(
+#[no_mangle]
+pub fn a(exp: u32) -> u64 {
+    // CHECK: %[[R:.+]] = and i32 %exp, 63
+    // CHECK: %[[R:.+]] = zext i32 %[[R:.+]] to i64
+    // CHECK: %[[R:.+]] = shl nuw i64 %[[R:.+]].i, %[[R:.+]]
+    // CHECK: ret i64 %[[R:.+]]
+    2u64.pow(exp)
+}
+
+#[no_mangle]
+pub fn b(exp: u32) -> i64 {
+    // CHECK: %[[R:.+]] = and i32 %exp, 63
+    // CHECK: %[[R:.+]] = zext i32 %[[R:.+]] to i64
+    // CHECK: %[[R:.+]] = shl nuw i64 %[[R:.+]].i, %[[R:.+]]
+    // CHECK: ret i64 %[[R:.+]]
+    2i64.pow(exp)
+}
+
+// CHECK-LABEL: @c(
+#[no_mangle]
+pub fn c(exp: u32) -> u32 {
+    // CHECK: %[[R:.+]].0.i = shl i32 %exp, 1
+    // CHECK: %[[R:.+]].1.i = icmp sgt i32 %exp, -1
+    // CHECK: %[[R:.+]].i = icmp ult i32 %[[R:.+]].0.i, 32
+    // CHECK: %fine.i = and i1 %[[R:.+]].1.i, %[[R:.+]].i
+    // CHECK: %0 = and i32 %[[R:.+]].0.i, 30
+    // CHECK: %[[R:.+]].i = zext i1 %fine.i to i32
+    // CHECK: %[[R:.+]] = shl nuw nsw i32 %[[R:.+]].i, %0
+    // CHECK: ret i32 %[[R:.+]]
+    4u32.pow(exp)
+}
+
+// CHECK-LABEL: @d(
+#[no_mangle]
+pub fn d(exp: u32) -> u32 {
+    // CHECK: tail call { i32, i1 } @llvm.umul.with.overflow.i32(i32 %exp, i32 5)
+    // CHECK: %[[R:.+]].0.i = extractvalue { i32, i1 } %[[R:.+]], 0
+    // CHECK: %[[R:.+]].1.i = extractvalue { i32, i1 } %[[R:.+]], 1
+    // CHECK: %[[R:.+]].i = xor i1 %[[R:.+]].1.i, true
+    // CHECK: %[[R:.+]].i = icmp ult i32 %[[R:.+]].0.i, 32
+    // CHECK: %fine.i = and i1 %[[R:.+]].i, %[[R:.+]].i
+    // CHECK: %[[R:.+]] = and i32 %[[R:.+]].0.i, 31
+    // CHECK: %[[R:.+]].i = zext i1 %fine.i to i32
+    // CHECK: %[[R:.+]] = shl nuw i32 %[[R:.+]].i, %1
+    // CHECK: ret i32 %[[R:.+]]
+    32u32.pow(exp)
+}
+
+// CHECK-LABEL: @e(
+#[no_mangle]
+pub fn e(exp: u32) -> i32 {
+    // CHECK: tail call { i32, i1 } @llvm.umul.with.overflow.i32(i32 %exp, i32 5)
+    // CHECK: %[[R:.+]].0.i = extractvalue { i32, i1 } %[[R:.+]], 0
+    // CHECK: %[[R:.+]].i = icmp ult i32 %[[R:.+]].0.i, 32
+    // CHECK: %[[R:.+]].1.i = extractvalue { i32, i1 } %[[R:.+]], 1
+    // CHECK: %[[R:.+]].i = xor i1 %[[R:.+]].1.i, true
+    // CHECK: %fine.i = and i1 %[[R:.+]].i, %[[R:.+]].i
+    // CHECK: %[[R:.+]].i = zext i1 %fine.i to i32
+    // CHECK: %[[R:.+]] = and i32 %[[R:.+]].0.i, 31
+    // CHECK: %[[R:.+]] = shl nuw i32 %[[R:.+]].i, %1
+    // CHECK: ret i32 %[[R:.+]]
+    32i32.pow(exp)
+}

tests/ui/consts/is_val_statically_known.rs

@@ -0,0 +1,15 @@
+// run-pass
+
+#![feature(core_intrinsics)]
+#![feature(is_val_statically_known)]
+
+use std::intrinsics::is_val_statically_known;
+
+const CONST_TEST: bool = unsafe { is_val_statically_known(0) };
+
+fn main() {
+    if CONST_TEST {
+        unreachable!("currently expected to return false during const eval");
+        // but note that this is not a guarantee!
+    }
+}