p3_util/
lib.rs

1//! Various simple utilities.
2
3#![no_std]
4
5extern crate alloc;
6
7use core::hint::unreachable_unchecked;
8
9pub mod array_serialization;
10pub mod linear_map;
11
12/// Computes `ceil(a / b)`. Assumes `a + b` does not overflow.
13#[must_use]
14pub const fn ceil_div_usize(a: usize, b: usize) -> usize {
15    (a + b - 1) / b
16}
17
18/// Computes `ceil(log_2(n))`.
19#[must_use]
20pub const fn log2_ceil_usize(n: usize) -> usize {
21    (usize::BITS - n.saturating_sub(1).leading_zeros()) as usize
22}
23
24#[must_use]
25pub fn log2_ceil_u64(n: u64) -> u64 {
26    (u64::BITS - n.saturating_sub(1).leading_zeros()).into()
27}
28
29/// Computes `log_2(n)`
30///
31/// # Panics
32/// Panics if `n` is not a power of two.
33#[must_use]
34#[inline]
35pub fn log2_strict_usize(n: usize) -> usize {
36    let res = n.trailing_zeros();
37    assert_eq!(n.wrapping_shr(res), 1, "Not a power of two: {n}");
38    res as usize
39}
40
41/// Returns `[0, ..., N - 1]`.
42#[must_use]
43pub const fn indices_arr<const N: usize>() -> [usize; N] {
44    let mut indices_arr = [0; N];
45    let mut i = 0;
46    while i < N {
47        indices_arr[i] = i;
48        i += 1;
49    }
50    indices_arr
51}
52
53#[inline]
54pub const fn reverse_bits(x: usize, n: usize) -> usize {
55    reverse_bits_len(x, n.trailing_zeros() as usize)
56}
57
58#[inline]
59pub const fn reverse_bits_len(x: usize, bit_len: usize) -> usize {
60    // NB: The only reason we need overflowing_shr() here as opposed
61    // to plain '>>' is to accommodate the case n == num_bits == 0,
62    // which would become `0 >> 64`. Rust thinks that any shift of 64
63    // bits causes overflow, even when the argument is zero.
64    x.reverse_bits()
65        .overflowing_shr(usize::BITS - bit_len as u32)
66        .0
67}
68
69pub fn reverse_slice_index_bits<F>(vals: &mut [F]) {
70    let n = vals.len();
71    if n == 0 {
72        return;
73    }
74    let log_n = log2_strict_usize(n);
75
76    for i in 0..n {
77        let j = reverse_bits_len(i, log_n);
78        if i < j {
79            vals.swap(i, j);
80        }
81    }
82}
83
84#[inline(always)]
85pub fn assume(p: bool) {
86    debug_assert!(p);
87    if !p {
88        unsafe {
89            unreachable_unchecked();
90        }
91    }
92}
93
94/// Try to force Rust to emit a branch. Example:
95///
96/// ```no_run
97/// let x = 100;
98/// if x > 20 {
99///     println!("x is big!");
100///     p3_util::branch_hint();
101/// } else {
102///     println!("x is small!");
103/// }
104/// ```
105///
106/// This function has no semantics. It is a hint only.
107#[inline(always)]
108pub fn branch_hint() {
109    // NOTE: These are the currently supported assembly architectures. See the
110    // [nightly reference](https://doc.rust-lang.org/nightly/reference/inline-assembly.html) for
111    // the most up-to-date list.
112    #[cfg(any(
113        target_arch = "aarch64",
114        target_arch = "arm",
115        target_arch = "riscv32",
116        target_arch = "riscv64",
117        target_arch = "x86",
118        target_arch = "x86_64",
119    ))]
120    unsafe {
121        core::arch::asm!("", options(nomem, nostack, preserves_flags));
122    }
123}
124
125/// Convenience methods for Vec.
126pub trait VecExt<T> {
127    /// Push `elem` and return a reference to it.
128    fn pushed_ref(&mut self, elem: T) -> &T;
129    /// Push `elem` and return a mutable reference to it.
130    fn pushed_mut(&mut self, elem: T) -> &mut T;
131}
132
133impl<T> VecExt<T> for alloc::vec::Vec<T> {
134    fn pushed_ref(&mut self, elem: T) -> &T {
135        self.push(elem);
136        self.last().unwrap()
137    }
138    fn pushed_mut(&mut self, elem: T) -> &mut T {
139        self.push(elem);
140        self.last_mut().unwrap()
141    }
142}
p3_util/lib.rs

p3_util/
lib.rs
