ff: refactor find_naf and add unit tests

ff/src/biginteger/arithmetic.rs

@@ -1,4 +1,3 @@
-use ark_std::Zero;
 use ark_std::{vec, vec::*};
 
 macro_rules! adc {
@@ -153,53 +152,30 @@ pub fn mac_with_carry(a: u64, b: u64, c: u64, carry: &mut u64) -> u64 {
 
 /// Compute the NAF (non-adjacent form) of num
 pub fn find_naf(num: &[u64]) -> Vec<i8> {
-    let is_zero = |num: &[u64]| num.iter().all(Zero::is_zero);
-    let is_odd = |num: &[u64]| num[0] & 1 == 1;
-    let sub_noborrow = |num: &mut [u64], z: u64| {
-        let mut other = vec![0u64; num.len()];
-        other[0] = z;
-        let mut borrow = 0;
-
-        for (a, b) in num.iter_mut().zip(other) {
-            borrow = sbb(a, b, borrow);
-        }
-    };
-    let add_nocarry = |num: &mut [u64], z: u64| {
-        let mut other = vec![0u64; num.len()];
-        other[0] = z;
-        let mut carry = 0;
-
-        for (a, b) in num.iter_mut().zip(other) {
-            carry = adc(a, b, carry);
-        }
-    };
-    let div2 = |num: &mut [u64]| {
-        let mut t = 0;
-        for i in num.iter_mut().rev() {
-            let t2 = *i << 63;
-            *i >>= 1;
-            *i |= t;
-            t = t2;
-        }
-    };
-
     let mut num = num.to_vec();
     let mut res = vec![];
 
-    while !is_zero(&num) {
-        let z: i8;
-        if is_odd(&num) {
-            z = 2 - (num[0] % 4) as i8;
-            if z >= 0 {
-                sub_noborrow(&mut num, z as u64)
-            } else {
-                add_nocarry(&mut num, (-z) as u64)
-            }
+    while num.iter().any(|&x| x != 0) {
+        let z = if num[0] & 1 == 1 {
+            let z = 2 - (num[0] % 4) as i8;
+            // Choose the appropriate operation: subtract (sbb) if z >= 0, otherwise add (adc).
+            let op = if z >= 0 { sbb } else { adc };
+            num.iter_mut()
+                .zip(ark_std::iter::once(&(z.abs() as u64)).chain(ark_std::iter::repeat(&0)))
+                .fold(0, |carry, (a, b)| op(a, *b, carry));
+            z
         } else {
-            z = 0;
-        }
+            0
+        };
+
         res.push(z);
-        div2(&mut num);
+
+        // Perform an in-place division of `num` by 2, using bit-shifting.
+        num.iter_mut().rev().fold(0, |carry, x| {
+            let next_carry = *x << 63;
+            *x = (*x >> 1) | carry;
+            next_carry
+        });
     }
 
     res
@@ -454,4 +430,47 @@ mod tests {
             assert_eq!(test, test_expected);
         }
     }
+
+    #[test]
+    fn test_find_naf_zero() {
+        // Test for zero input
+        let naf = find_naf(&[0]);
+        assert!(naf.is_empty());
+    }
+
+    #[test]
+    fn test_find_naf_single_digit() {
+        // Test for small numbers
+        assert_eq!(find_naf(&[1]), vec![1]);
+        assert_eq!(find_naf(&[2]), vec![0, 1]);
+        assert_eq!(find_naf(&[3]), vec![-1, 0, 1]);
+        assert_eq!(find_naf(&[4]), vec![0, 0, 1]);
+    }
+
+    #[test]
+    fn test_find_naf_large_number() {
+        // Test for a larger number
+        assert_eq!(find_naf(&[13]), vec![1, 0, -1, 0, 1]);
+    }
+
+    #[test]
+    fn test_find_naf_multiple_blocks() {
+        // Test multi-block number (simulate large numbers split across blocks)
+        let num = [0, 1];
+        assert_eq!(
+            find_naf(&num),
+            vec![
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 1
+            ]
+        );
+    }
+
+    #[test]
+    fn test_find_naf_edge_cases() {
+        // Test edge cases
+        let naf = find_naf(&[u64::MAX]);
+        assert!(naf.len() > 0);
+    }
 }