From a17fb802555c2ab36f65c541301b62a461c74fe1 Mon Sep 17 00:00:00 2001
From: varun-doshi <doshivarun202@gmail.com>
Date: Thu, 13 Mar 2025 11:22:33 +0530
Subject: [PATCH] feat: ProofOptionsBuilder

---
 air/src/air/tests.rs | 65 ++++++++++++++++++-----------------
 air/src/options.rs   | 82 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 31 deletions(-)

diff --git a/air/src/air/tests.rs b/air/src/air/tests.rs
index 3300dcc78..e3a2f83a8 100644
--- a/air/src/air/tests.rs
+++ b/air/src/air/tests.rs
@@ -12,7 +12,10 @@ use super::{
     Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
     TransitionConstraintDegree,
 };
-use crate::{options::BatchingMethod, FieldExtension};
+use crate::{
+    options::{BatchingMethod, ProofOptionsBuilder},
+    FieldExtension,
+};
 
 // PERIODIC COLUMNS
 // ================================================================================================
@@ -205,16 +208,16 @@ impl MockAir {
         let mut result = Self::new(
             TraceInfo::with_meta(4, trace_length, vec![1]),
             (),
-            ProofOptions::new(
-                32,
-                8,
-                0,
-                FieldExtension::None,
-                4,
-                31,
-                BatchingMethod::Linear,
-                BatchingMethod::Linear,
-            ),
+            ProofOptionsBuilder::new()
+                .num_queries(32)
+                .blowup_factor(8)
+                .grinding_factor(0)
+                .field_extension(FieldExtension::None)
+                .fri_folding_factor(4)
+                .fri_remainder_max_degree(31)
+                .batching_constraints(BatchingMethod::Linear)
+                .batching_deep(BatchingMethod::Linear)
+                .build(),
         );
         result.periodic_columns = column_values;
         result
@@ -224,16 +227,16 @@ impl MockAir {
         let mut result = Self::new(
             TraceInfo::with_meta(4, trace_length, vec![assertions.len() as u8]),
             (),
-            ProofOptions::new(
-                32,
-                8,
-                0,
-                FieldExtension::None,
-                4,
-                31,
-                BatchingMethod::Linear,
-                BatchingMethod::Linear,
-            ),
+            ProofOptionsBuilder::new()
+                .num_queries(32)
+                .blowup_factor(8)
+                .grinding_factor(0)
+                .field_extension(FieldExtension::None)
+                .fri_folding_factor(4)
+                .fri_remainder_max_degree(31)
+                .batching_constraints(BatchingMethod::Linear)
+                .batching_deep(BatchingMethod::Linear)
+                .build(),
         );
         result.assertions = assertions;
         result
@@ -283,16 +286,16 @@ pub fn build_context<B: StarkField>(
     trace_width: usize,
     num_assertions: usize,
 ) -> AirContext<B> {
-    let options = ProofOptions::new(
-        32,
-        8,
-        0,
-        FieldExtension::None,
-        4,
-        31,
-        BatchingMethod::Linear,
-        BatchingMethod::Linear,
-    );
+    let options = ProofOptionsBuilder::new()
+        .num_queries(32)
+        .blowup_factor(8)
+        .grinding_factor(0)
+        .field_extension(FieldExtension::None)
+        .fri_folding_factor(4)
+        .fri_remainder_max_degree(31)
+        .batching_constraints(BatchingMethod::Linear)
+        .batching_deep(BatchingMethod::Linear)
+        .build();
     let t_degrees = vec![TransitionConstraintDegree::new(2)];
     let trace_info = TraceInfo::new(trace_width, trace_length);
     AirContext::new(trace_info, t_degrees, num_assertions, options)
diff --git a/air/src/options.rs b/air/src/options.rs
index 4d1274250..e7675fb85 100644
--- a/air/src/options.rs
+++ b/air/src/options.rs
@@ -291,6 +291,88 @@ impl ProofOptions {
     }
 }
 
+pub struct ProofOptionsBuilder {
+    num_queries: u8,
+    blowup_factor: u8,
+    grinding_factor: u8,
+    field_extension: FieldExtension,
+    fri_folding_factor: u8,
+    fri_remainder_max_degree: u8,
+    batching_constraints: BatchingMethod,
+    batching_deep: BatchingMethod,
+    partition_options: PartitionOptions,
+}
+
+impl ProofOptionsBuilder {
+    pub fn new() -> Self {
+        Self {
+            num_queries: 0,
+            blowup_factor: 0,
+            grinding_factor: 0,
+            field_extension: FieldExtension::None,
+            fri_folding_factor: 0,
+            fri_remainder_max_degree: 0,
+            batching_constraints: BatchingMethod::Linear,
+            batching_deep: BatchingMethod::Linear,
+            partition_options: PartitionOptions::default(),
+        }
+    }
+
+    pub fn num_queries(mut self, value: u8) -> Self {
+        self.num_queries = value;
+        self
+    }
+
+    pub fn blowup_factor(mut self, value: u8) -> Self {
+        self.blowup_factor = value;
+        self
+    }
+
+    pub fn grinding_factor(mut self, value: u8) -> Self {
+        self.grinding_factor = value;
+        self
+    }
+
+    pub fn field_extension(mut self, value: FieldExtension) -> Self {
+        self.field_extension = value;
+        self
+    }
+
+    pub fn fri_folding_factor(mut self, value: u8) -> Self {
+        self.fri_folding_factor = value;
+        self
+    }
+
+    pub fn batching_constraints(mut self, value: BatchingMethod) -> Self {
+        self.batching_constraints = value;
+        self
+    }
+
+    pub fn fri_remainder_max_degree(mut self, value: u8) -> Self {
+        self.fri_remainder_max_degree = value;
+        self
+    }
+    
+    pub fn batching_deep(mut self, value: BatchingMethod) -> Self {
+        self.batching_deep = value;
+        self
+    }
+
+    pub fn build(self) -> ProofOptions {
+        ProofOptions {
+            num_queries: self.num_queries,
+            blowup_factor: self.blowup_factor,
+            grinding_factor: self.grinding_factor,
+            field_extension: self.field_extension,
+            fri_folding_factor: self.fri_folding_factor,
+            fri_remainder_max_degree: self.fri_remainder_max_degree,
+            batching_constraints: self.batching_constraints,
+            batching_deep: self.batching_deep,
+            partition_options: self.partition_options,
+        }
+    }
+}
+
 impl<E: StarkField> ToElements<E> for ProofOptions {
     /// Encodes these proof options into 3 field elements.
     fn to_elements(&self) -> Vec<E> {