Skip to main content

binius_utils/
checked_arithmetics.rs

1// Copyright 2024-2025 Irreducible Inc.
2// Copyright (c) 2024 The Plonky3 Authors
3
4/// Division implementation that fails in case when `a`` isn't divisible by `b`
5pub const fn checked_int_div(a: usize, b: usize) -> usize {
6	let result = a / b;
7	assert!(b * result == a);
8
9	result
10}
11
12/// Computes binary logarithm of `val`.
13/// If `val` is not a power of 2, returns `None`.
14#[inline]
15#[must_use]
16pub const fn strict_log_2(val: usize) -> Option<usize> {
17	if val == 0 {
18		return None;
19	}
20
21	let pow = val.trailing_zeros();
22	if val.wrapping_shr(pow) == 1 {
23		Some(pow as usize)
24	} else {
25		None
26	}
27}
28
29/// Computes the binary logarithm of `val`.
30///
31/// `#[track_caller]` puts the panic at the call site rather than inside this function.
32/// A `const fn` cannot format a panic message, so the offending value is not reported.
33///
34/// # Panics
35/// Panics if `val` is not a power of two, zero included.
36#[inline]
37#[must_use]
38#[track_caller]
39pub const fn checked_log_2(val: usize) -> usize {
40	strict_log_2(val).expect("value is not a power of two")
41}
42
43/// Computes the binary logarithm of $n$ rounded up to the nearest integer.
44///
45/// When $n$ is 0, this function returns 0. Otherwise, it returns $\lceil \log_2 n \rceil$.
46#[must_use]
47pub const fn log2_ceil_usize(n: usize) -> usize {
48	min_bits(n.saturating_sub(1))
49}
50
51/// Returns the number of bits needed to represent $n$.
52///
53/// When $n$ is 0, this function returns 0. Otherwise, it returns $\lfloor \log_2 n \rfloor + 1$.
54#[must_use]
55pub const fn min_bits(n: usize) -> usize {
56	(usize::BITS - n.leading_zeros()) as usize
57}
58
59#[cfg(test)]
60mod tests {
61	use super::*;
62
63	#[test]
64	fn test_checked_int_div_success() {
65		assert_eq!(checked_int_div(6, 1), 6);
66		assert_eq!(checked_int_div(6, 2), 3);
67		assert_eq!(checked_int_div(6, 6), 1);
68	}
69
70	#[test]
71	#[should_panic]
72	const fn test_checked_int_div_fail() {
73		_ = checked_int_div(5, 2);
74	}
75
76	// Number of bits in `n`, counted one shift at a time.
77	fn min_bits_ref(n: usize) -> usize {
78		let mut bits = 0;
79		let mut rest = n;
80		while rest > 0 {
81			bits += 1;
82			rest >>= 1;
83		}
84		bits
85	}
86
87	// Smallest `k` with `2^k >= n`, found by counting up.
88	fn log2_ceil_ref(n: usize) -> usize {
89		let mut k = 0;
90		while (1usize << k) < n {
91			k += 1;
92		}
93		k
94	}
95
96	#[test]
97	fn test_checked_log2_success() {
98		assert_eq!(checked_log_2(1), 0);
99		assert_eq!(checked_log_2(2), 1);
100		assert_eq!(checked_log_2(4), 2);
101		assert_eq!(checked_log_2(64), 6);
102		assert_eq!(checked_log_2(1 << 63), 63);
103	}
104
105	#[test]
106	#[should_panic]
107	const fn test_checked_log2_fail() {
108		_ = checked_log_2(6)
109	}
110
111	#[test]
112	#[should_panic(expected = "value is not a power of two")]
113	fn test_checked_log2_zero_panics() {
114		_ = checked_log_2(0);
115	}
116
117	// Callers put `checked_log_2` in associated-const initializers, so const-ness is load-bearing.
118	const _: () = assert!(checked_log_2(64) == 6);
119
120	#[test]
121	fn test_strict_log_2_boundaries() {
122		assert_eq!(strict_log_2(0), None);
123		assert_eq!(strict_log_2(1), Some(0));
124		assert_eq!(strict_log_2(2), Some(1));
125		assert_eq!(strict_log_2(3), None);
126		assert_eq!(strict_log_2(1 << 63), Some(63));
127		assert_eq!(strict_log_2(usize::MAX), None);
128	}
129
130	#[test]
131	fn test_min_bits_boundaries() {
132		assert_eq!(min_bits(0), 0);
133		assert_eq!(min_bits(1), 1);
134		assert_eq!(min_bits(2), 2);
135		assert_eq!(min_bits(3), 2);
136		assert_eq!(min_bits(1 << 63), 64);
137		assert_eq!(min_bits(usize::MAX), 64);
138	}
139
140	#[test]
141	fn test_log2_ceil_usize_boundaries() {
142		assert_eq!(log2_ceil_usize(0), 0);
143		assert_eq!(log2_ceil_usize(1), 0);
144		assert_eq!(log2_ceil_usize(2), 1);
145		assert_eq!(log2_ceil_usize(3), 2);
146		// The last exact power of two, then the first value that needs one more bit.
147		assert_eq!(log2_ceil_usize(1 << 63), 63);
148		assert_eq!(log2_ceil_usize((1 << 63) + 1), 64);
149		assert_eq!(log2_ceil_usize(usize::MAX), 64);
150	}
151
152	// Exhaustive over the low 2^20, which beats sampling on a domain this small.
153	#[test]
154	fn test_bit_counts_match_reference() {
155		for n in 0..1usize << 20 {
156			assert_eq!(min_bits(n), min_bits_ref(n), "min_bits({n})");
157			assert_eq!(log2_ceil_usize(n), log2_ceil_ref(n), "log2_ceil_usize({n})");
158		}
159	}
160
161	// The three functions agree where their domains overlap: rounding a power of two up is a no-op,
162	// and `min_bits` of a power of two is one more than its logarithm.
163	#[test]
164	fn test_powers_of_two_agree() {
165		for log in 0..usize::BITS as usize {
166			let n = 1usize << log;
167			assert_eq!(checked_log_2(n), log);
168			assert_eq!(log2_ceil_usize(n), log);
169			assert_eq!(min_bits(n), log + 1);
170		}
171	}
172}