From f07de9194d364f072d8b91b60b536eab194c713f Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Sun, 26 Nov 2023 19:38:08 -0700
Subject: [PATCH] Add `BoxedUint::bits` (#328)

Adds a function for counting the number of bits in a value, i.e. the
index of the highest set bit.

Since `BoxedUint` isn't `const fn`, this can be implemented in
constant-time using `subtle`.
---
 src/uint/boxed.rs             |  6 +----
 src/uint/boxed/bits.rs        | 49 +++++++++++++++++++++++++++++++++++
 tests/boxed_uint_proptests.rs |  2 +-
 3 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 src/uint/boxed/bits.rs

diff --git a/src/uint/boxed.rs b/src/uint/boxed.rs
index cfc85143..2f1bf6be 100644
--- a/src/uint/boxed.rs
+++ b/src/uint/boxed.rs
@@ -3,6 +3,7 @@
 mod add;
 mod add_mod;
 mod bit_and;
+mod bits;
 mod cmp;
 pub(crate) mod encoding;
 mod modular;
@@ -150,11 +151,6 @@ impl BoxedUint {
         self.limbs.len()
     }
 
-    /// Get the precision of this [`BoxedUint`] in bits.
-    pub fn bits(&self) -> usize {
-        self.limbs.len() * Limb::BITS
-    }
-
     /// Perform a carry chain-like operation over the limbs of the inputs,
     /// constructing a result from the returned limbs and carry which is
     /// widened to the same width as the widest input.
diff --git a/src/uint/boxed/bits.rs b/src/uint/boxed/bits.rs
new file mode 100644
index 00000000..c838726b
--- /dev/null
+++ b/src/uint/boxed/bits.rs
@@ -0,0 +1,49 @@
+//! Bit manipulation functions.
+
+use crate::{BoxedUint, Limb, Zero};
+use subtle::{ConditionallySelectable, ConstantTimeEq};
+
+impl BoxedUint {
+    /// Calculate the number of bits needed to represent this number, i.e. the index of the highest
+    /// set bit.
+    ///
+    /// Use [`BoxedUint::bits_precision`] to get the total capacity of this integer.
+    pub fn bits(&self) -> usize {
+        // Use `u32` because `subtle` can't select on `usize` and it matches what `core` uses for
+        // the return value of `leading_zeros`
+        let mut leading_zeros = 0u32;
+        let mut n = 0u32;
+
+        for limb in self.limbs.iter().rev() {
+            n.conditional_assign(&(n + 1), !limb.is_zero() | !n.ct_eq(&0));
+
+            // Set `leading_zeros` for the first nonzero limb we encounter
+            leading_zeros.conditional_assign(&(limb.leading_zeros() as u32), n.ct_eq(&1));
+        }
+
+        Limb::BITS * (n as usize) - (leading_zeros as usize)
+    }
+
+    /// Get the precision of this [`BoxedUint`] in bits.
+    pub fn bits_precision(&self) -> usize {
+        self.limbs.len() * Limb::BITS
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::BoxedUint;
+    use hex_literal::hex;
+
+    #[test]
+    fn bits() {
+        assert_eq!(0, BoxedUint::zero().bits());
+        assert_eq!(128, BoxedUint::max(128).bits());
+
+        let n1 = BoxedUint::from_be_slice(&hex!("000000000029ffffffffffffffffffff"), 128).unwrap();
+        assert_eq!(86, n1.bits());
+
+        let n2 = BoxedUint::from_be_slice(&hex!("00000000004000000000000000000000"), 128).unwrap();
+        assert_eq!(87, n2.bits());
+    }
+}
diff --git a/tests/boxed_uint_proptests.rs b/tests/boxed_uint_proptests.rs
index 4dba4e73..493bbd85 100644
--- a/tests/boxed_uint_proptests.rs
+++ b/tests/boxed_uint_proptests.rs
@@ -42,7 +42,7 @@ proptest! {
         match Option::<BoxedUint>::from(a.checked_add(&b)) {
             Some(actual) => prop_assert_eq!(expected, to_biguint(&actual)),
             None => {
-                let max = BoxedUint::max(a.bits());
+                let max = BoxedUint::max(a.bits_precision());
                 prop_assert!(expected > to_biguint(&max));
             }
         }