From bd2e37e9369238ae569aa5d6f9cc32b59b85bf89 Mon Sep 17 00:00:00 2001
From: erabinov <erabinovich92@gmail.com>
Date: Mon, 17 Jun 2024 17:13:36 -0700
Subject: [PATCH] First attempt at fix.

---
 recursion/compiler/src/ir/bits.rs     |  9 +++++++++
 recursion/program/src/machine/core.rs | 15 +++++++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/recursion/compiler/src/ir/bits.rs b/recursion/compiler/src/ir/bits.rs
index 396fb92be..b08a7e29e 100644
--- a/recursion/compiler/src/ir/bits.rs
+++ b/recursion/compiler/src/ir/bits.rs
@@ -47,6 +47,15 @@ impl<C: Config> Builder<C> {
         output
     }
 
+    /// Range checks a felt to a certain number of bits.
+    pub fn range_check_f(&mut self, num: Felt<C::F>, num_bits: usize) {
+        let bits = self.num2bits_f(num);
+        self.range(num_bits, bits.len()).for_each(|i, builder| {
+            let bit = builder.get(&bits, i);
+            builder.assert_var_eq(bit, C::N::zero());
+        });
+    }
+
     /// Converts a felt to bits.
     pub fn num2bits_f(&mut self, num: Felt<C::F>) -> Array<C, Var<C::N>> {
         let output = self.dyn_array::<Var<_>>(NUM_BITS);
diff --git a/recursion/program/src/machine/core.rs b/recursion/program/src/machine/core.rs
index e8a547cc0..88158f21b 100644
--- a/recursion/program/src/machine/core.rs
+++ b/recursion/program/src/machine/core.rs
@@ -104,17 +104,17 @@ where
     ///
     /// See [SP1Prover::verify] for the verification algorithm of a complete SP1 proof. In this
     /// function, we are aggregating several shard proofs and attesting to an aggregated state which
-    /// reprersents all the shards. The consistency conditions of the aggregated state are
+    /// represents all the shards. The consistency conditions of the aggregated state are
     /// asserted in the following way:
     ///
-    /// - Start pc for every shardf should be what the next pc declared in the previous shard was.
+    /// - Start pc for every shard should be what the next pc declared in the previous shard was.
     /// - Public input, deferred proof digests, and exit code should be the same in all shards.
     ///
     /// ## The leaf challenger.
     /// A key difference between the recursive tree verification and the complete one in
-    /// [SP1Prover::verify] is that the recursive verifier has no way of reconstructiing the
-    /// chanllenger only from a part of the shard proof. Therefoee, the value of the leaf challenger
-    /// is witnessed in the program and the verifier assertds correctness given this challenger.
+    /// [SP1Prover::verify] is that the recursive verifier has no way of reconstructing the
+    /// chanllenger only from a part of the shard proof. Therefore, the value of the leaf challenger
+    /// is witnessed in the program and the verifier asserts correctness given this challenger.
     /// In the course of the recursive verification, the challenger is reconstructed by observing
     /// the commitments one by one, and in the final step, the challenger is asserted to be the same
     /// as the one witnessed here.
@@ -223,7 +223,7 @@ where
                 builder.assign(exit_code, public_values.exit_code);
             });
 
-            // If the shard is zero, verify the global initial conditions hold on challenger and pc.
+            // If shard is one, verify the global initial conditions hold on challenger and pc.
             let shard = felt2var(builder, public_values.shard);
             builder.if_eq(shard, C::N::one()).then(|builder| {
                 // This should be the first proof as well
@@ -279,6 +279,9 @@ where
                 builder.assert_felt_eq(*digest, *current_digest);
             }
 
+            // Range check the shard count to be less than 1<<16.
+            builder.range_check_f(current_shard, 16);
+
             // Update the loop variables: the reconstruct challenger, cumulative sum, shard number,
             // and program counter.