From eaba5238a6400448e3aacb65c969e76b0519c006 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Thu, 17 Jun 2021 09:49:41 +0200
Subject: [PATCH] Change PoW to checking leading zeros

---
 src/fri/prover.rs             |  4 ++--
 src/fri/recursive_verifier.rs |  2 +-
 src/fri/verifier.rs           |  3 ++-
 src/gadgets/split_base.rs     | 36 ++++-------------------------------
 4 files changed, 9 insertions(+), 36 deletions(-)

diff --git a/src/fri/prover.rs b/src/fri/prover.rs
index a919b559..f937ee55 100644
--- a/src/fri/prover.rs
+++ b/src/fri/prover.rs
@@ -106,8 +106,8 @@ fn fri_proof_of_work<F: Field>(current_hash: Hash<F>, config: &FriConfig) -> F {
                 false,
             )
             .to_canonical_u64()
-            .trailing_zeros()
-                >= config.proof_of_work_bits
+            .leading_zeros()
+                >= config.proof_of_work_bits + F::ORDER.leading_zeros()
         })
         .map(F::from_canonical_u64)
         .expect("Proof of work failed.")
diff --git a/src/fri/recursive_verifier.rs b/src/fri/recursive_verifier.rs
index aa9e1d0c..4a92172d 100644
--- a/src/fri/recursive_verifier.rs
+++ b/src/fri/recursive_verifier.rs
@@ -57,7 +57,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         inputs.push(proof.pow_witness);
 
         let hash = self.hash_n_to_m(inputs, 1, false)[0];
-        self.assert_trailing_zeros::<2>(hash, config.proof_of_work_bits);
+        self.assert_leading_zeros(hash, config.proof_of_work_bits + F::ORDER.leading_zeros());
     }
 
     pub fn verify_fri_proof(
diff --git a/src/fri/verifier.rs b/src/fri/verifier.rs
index 45c94173..04432ed4 100644
--- a/src/fri/verifier.rs
+++ b/src/fri/verifier.rs
@@ -55,7 +55,8 @@ fn fri_verify_proof_of_work<F: Field + Extendable<D>, const D: usize>(
         false,
     );
     ensure!(
-        hash.to_canonical_u64().trailing_zeros() >= config.proof_of_work_bits,
+        hash.to_canonical_u64().leading_zeros()
+            >= config.proof_of_work_bits + F::ORDER.leading_zeros(),
         "Invalid proof of work witness."
     );
 
diff --git a/src/gadgets/split_base.rs b/src/gadgets/split_base.rs
index bf2c158c..9aefe6da 100644
--- a/src/gadgets/split_base.rs
+++ b/src/gadgets/split_base.rs
@@ -22,18 +22,9 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         )
     }
 
-    /// Asserts that `x`'s big-endian bit representation has at least `trailing_zeros` trailing zeros.
-    pub(crate) fn assert_trailing_zeros<const B: usize>(&mut self, x: Target, trailing_zeros: u32) {
-        let num_limbs = num_limbs(64, B);
-        let num_limbs_to_check = num_limbs_to_check(trailing_zeros, B);
-        let limbs = self.split_le_base::<B>(x, num_limbs);
-        assert!(
-            num_limbs_to_check <= self.config.num_routed_wires,
-            "Not enough routed wires."
-        );
-        for i in 0..num_limbs_to_check {
-            self.assert_zero(limbs[i]);
-        }
+    /// Asserts that `x`'s big-endian bit representation has at least `leading_zeros` leading zeros.
+    pub(crate) fn assert_leading_zeros(&mut self, x: Target, leading_zeros: u32) {
+        self.range_check(x, (64 - leading_zeros) as usize);
     }
 
     pub(crate) fn reverse_bits<const B: usize>(&mut self, x: Target, num_limbs: usize) -> Target {
@@ -45,23 +36,6 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     }
 }
 
-/// Returns `k` such that any number with `k` trailing zeros in base `base` has at least
-/// `n` trailing zeros in base 2.
-#[allow(unconditional_panic)]
-const fn num_limbs_to_check(n: u32, base: usize) -> usize {
-    if base % 2 == 1 {
-        // Dirty trick to panic if the base is odd.
-        // TODO: replace with `assert_eq!(base % 2, 0, "Base should be even.")` when stable.
-        [][0]
-    } else {
-        ceil_div_usize(n as usize, base.trailing_zeros() as usize)
-    }
-}
-
-fn num_limbs(n_log: u32, base: usize) -> usize {
-    ((n_log as f64) * 2.0_f64.log(base as f64)).ceil() as usize
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -90,9 +64,7 @@ mod tests {
         let revt = builder.reverse_bits::<2>(xt, 9);
         builder.route(revt, rev);
 
-        builder.assert_trailing_zeros::<6>(xt, 4);
-        builder.assert_trailing_zeros::<4>(xt, 5);
-        builder.assert_trailing_zeros::<12>(xt, 5);
+        builder.assert_leading_zeros(xt, 64 - 9);
         let data = builder.build();
 
         let proof = data.prove(PartialWitness::new());