wip: try iterator-based for loop

openvm-org · Jan 6, 2025 · 2ebc0de · 2ebc0de
1 parent f1e01a7
commit 2ebc0de
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git a/extensions/native/compiler/src/ir/builder.rs b/extensions/native/compiler/src/ir/builder.rs
@@ -362,6 +362,22 @@ impl<C: Config> Builder<C> {
         }
     }
 
+    pub fn range_with_step(
+        &mut self,
+        start: impl Into<RVar<C::N>>,
+        end: impl Into<RVar<C::N>>,
+        step_size: usize,
+    ) -> RangeBuilder<C> {
+        let start = start.into();
+        let end = end.into();
+        RangeBuilder {
+            start,
+            end,
+            builder: self,
+            step_size,
+        }
+    }
+
     /// Evaluate a block of operations repeatedly (until a break).
     pub fn do_loop(&mut self, mut f: impl FnMut(&mut Builder<C>) -> Result<(), BreakLoop>) {
         let mut loop_body_builder = self.create_sub_builder();

diff --git a/extensions/native/recursion/src/stark/mod.rs b/extensions/native/recursion/src/stark/mod.rs
@@ -4,7 +4,9 @@ use itertools::Itertools;
 use openvm_circuit::arch::instructions::program::Program;
 use openvm_native_compiler::{
     conversion::CompilerOptions,
-    ir::{Array, Builder, Config, Ext, ExtConst, Felt, SymbolicExt, Usize},
+    ir::{
+        Array, Builder, Config, Ext, ExtConst, Felt, MemIndex, MemVariable, Ptr, SymbolicExt, Usize,
+    },
     prelude::RVar,
 };
 use openvm_stark_backend::{
@@ -25,6 +27,7 @@ use openvm_stark_sdk::{
 use crate::{
     challenger::{duplex::DuplexChallengerVariable, ChallengerVariable},
     commit::{PcsVariable, PolynomialSpaceVariable},
+    digest::DigestVariable,
     folder::RecursiveVerifierConstraintFolder,
     fri::{
         types::{TwoAdicPcsMatsVariable, TwoAdicPcsRoundVariable},
@@ -202,13 +205,38 @@ where
             quotient: quotient_commit,
         } = commitments;
 
+        /*
         // Observe main trace commitments
         builder
             .range(0, main_trace_commits.len())
             .for_each(|i, builder| {
                 let main_commit = builder.get(main_trace_commits, i);
                 challenger.observe_digest(builder, main_commit);
             });
+        */
+
+        builder
+            .range_with_step(
+                main_trace_commits.ptr().address,
+                builder
+                    .get_ptr(main_trace_commits, main_trace_commits.len() - 1)
+                    .address,
+                DigestVariable::<C>::size_of(),
+            )
+            .for_each(|address, builder| {
+                let main_commit = builder.uninit();
+                builder.load(
+                    main_commit,
+                    Ptr { address },
+                    MemIndex {
+                        index: RVar::zero(),
+                        offset: 0,
+                        size: DigestVariable::<C>::size_of(),
+                    },
+                );
+                challenger.observe_digest(builder, main_commit);
+            });
+
         builder.range(0, air_proofs.len()).for_each(|i, builder| {
             let air_proof = builder.get(air_proofs, i);
             let log_degree = if builder.flags.static_only {