From 9a74dfa0312d57b1069f8b8e35a70b90ccc550f7 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Mon, 25 Nov 2024 14:12:26 +0000
Subject: [PATCH] chore: remove temporary allocations from `num_bits` (#6600)

---
 acvm-repo/acir_field/src/field_element.rs | 63 +++++++++++++++--------
 1 file changed, 42 insertions(+), 21 deletions(-)

diff --git a/acvm-repo/acir_field/src/field_element.rs b/acvm-repo/acir_field/src/field_element.rs
index 47ceb903111..01b9bf8881d 100644
--- a/acvm-repo/acir_field/src/field_element.rs
+++ b/acvm-repo/acir_field/src/field_element.rs
@@ -158,23 +158,6 @@ impl<F: PrimeField> FieldElement<F> {
         let fr = F::from_str(input).ok()?;
         Some(FieldElement(fr))
     }
-
-    fn bits(&self) -> Vec<bool> {
-        fn byte_to_bit(byte: u8) -> Vec<bool> {
-            let mut bits = Vec::with_capacity(8);
-            for index in (0..=7).rev() {
-                bits.push((byte & (1 << index)) >> index == 1);
-            }
-            bits
-        }
-
-        let bytes = self.to_be_bytes();
-        let mut bits = Vec::with_capacity(bytes.len() * 8);
-        for byte in bytes {
-            bits.extend(byte_to_bit(byte));
-        }
-        bits
-    }
 }
 
 impl<F: PrimeField> AcirField for FieldElement<F> {
@@ -224,12 +207,26 @@ impl<F: PrimeField> AcirField for FieldElement<F> {
 
     /// This is the number of bits required to represent this specific field element
     fn num_bits(&self) -> u32 {
-        let bits = self.bits();
-        // Iterate the number of bits and pop off all leading zeroes
-        let iter = bits.iter().skip_while(|x| !(**x));
+        let bytes = self.to_be_bytes();
+
+        // Iterate through the byte decomposition and pop off all leading zeroes
+        let mut iter = bytes.iter().skip_while(|x| (**x) == 0);
+
+        // The first non-zero byte in the decomposition may have some leading zero-bits.
+        let Some(head_byte) = iter.next() else {
+            // If we don't have a non-zero byte then the field element is zero,
+            // which we consider to require a single bit to represent.
+            return 1;
+        };
+        let num_bits_for_head_byte = head_byte.ilog2();
+
+        // Each remaining byte in the byte decomposition requires 8 bits.
+        //
         // Note: count will panic if it goes over usize::MAX.
         // This may not be suitable for devices whose usize < u16
-        iter.count() as u32
+        let tail_length = iter.count() as u32;
+
+        8 * tail_length + num_bits_for_head_byte + 1
     }
 
     fn to_u128(self) -> u128 {
@@ -374,6 +371,30 @@ mod tests {
     use super::{AcirField, FieldElement};
     use proptest::prelude::*;
 
+    #[test]
+    fn requires_one_bit_to_hold_zero() {
+        let field = FieldElement::<ark_bn254::Fr>::zero();
+        assert_eq!(field.num_bits(), 1);
+    }
+
+    proptest! {
+        #[test]
+        fn num_bits_agrees_with_ilog2(num in 1u128..) {
+            let field = FieldElement::<ark_bn254::Fr>::from(num);
+            prop_assert_eq!(field.num_bits(), num.ilog2() + 1);
+        }
+    }
+
+    #[test]
+    fn test_fits_in_u128() {
+        let field = FieldElement::<ark_bn254::Fr>::from(u128::MAX);
+        assert_eq!(field.num_bits(), 128);
+        assert!(field.fits_in_u128());
+        let big_field = field + FieldElement::one();
+        assert_eq!(big_field.num_bits(), 129);
+        assert!(!big_field.fits_in_u128());
+    }
+
     #[test]
     fn serialize_fixed_test_vectors() {
         // Serialized field elements from of 0, -1, -2, -3