cleanup & refactor

iosis-tech · Sep 2, 2024 · 127f0b1 · 127f0b1
1 parent ff79189
commit 127f0b1
Show file tree

Hide file tree

Showing 13 changed files with 112 additions and 69 deletions.
diff --git a/crates/air/src/diluted.rs b/crates/air/src/diluted.rs
@@ -1,3 +1,4 @@
+use crate::consts::{FELT_0, FELT_1, FELT_2};
 use starknet_crypto::Felt;
 
 // The cumulative value is defined using the next recursive formula:
@@ -24,15 +25,15 @@ use starknet_crypto::Felt;
 //
 // Now we can compute p_{n_bits} and q_{n_bits} in just n_bits recursive steps and we are done.
 pub fn get_diluted_product(n_bits: Felt, spacing: Felt, z: Felt, alpha: Felt) -> Felt {
-    let diff_multiplier = Felt::TWO.pow_felt(&spacing);
-    let mut diff_x: Felt = diff_multiplier - 2;
-    let mut x: Felt = Felt::ONE;
-    let mut p: Felt = z + 1;
-    let mut q: Felt = Felt::ONE;
+    let diff_multiplier = FELT_2.pow_felt(&spacing);
+    let mut diff_x: Felt = diff_multiplier - FELT_2;
+    let mut x: Felt = FELT_1;
+    let mut p: Felt = z + FELT_1;
+    let mut q: Felt = FELT_1;
 
-    let mut i = Felt::ZERO;
+    let mut i = FELT_0;
     loop {
-        if i == n_bits - 1 {
+        if i == n_bits - FELT_1 {
             break p + q * alpha;
         }
 
@@ -43,6 +44,6 @@ pub fn get_diluted_product(n_bits: Felt, spacing: Felt, z: Felt, alpha: Felt) ->
         q = q * y + x * x_p + q;
         p *= y;
 
-        i = i + 1;
+        i += FELT_1;
     }
 }
diff --git a/crates/air/src/domains.rs b/crates/air/src/domains.rs
@@ -31,12 +31,11 @@ impl StarkDomains {
             log_eval_domain_size: log_trace_domain_size + log_n_cosets,
             eval_domain_size: Felt::TWO.pow_felt(&log_eval_domain_size),
             eval_generator: FIELD_GENERATOR.pow_felt(
-                &STARK_PRIME_MINUS_ONE
-                    .field_div(&NonZeroFelt::from_felt_unchecked(eval_domain_size)),
+                &STARK_PRIME_MINUS_ONE.field_div(&NonZeroFelt::try_from(eval_domain_size).unwrap()),
             ),
             trace_generator: FIELD_GENERATOR.pow_felt(
                 &STARK_PRIME_MINUS_ONE
-                    .field_div(&NonZeroFelt::from_felt_unchecked(trace_domain_size)),
+                    .field_div(&NonZeroFelt::try_from(trace_domain_size).unwrap()),
             ),
             trace_domain_size,
             log_trace_domain_size,

diff --git a/crates/air/src/layout/dex/mod.rs b/crates/air/src/layout/dex/mod.rs
@@ -275,6 +275,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -358,20 +363,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);
@@ -387,9 +392,9 @@ impl LayoutTrait for Layout {
             .collect::<Vec<Felt>>();
 
         ensure!(initial_pc == INITIAL_PC, PublicInputError::MaxSteps);
-        ensure!(final_pc == INITIAL_PC + 4, PublicInputError::MaxSteps);
+        ensure!(final_pc == INITIAL_PC + FELT_4, PublicInputError::MaxSteps);
 
-        let program_end_pc = initial_fp - 2;
+        let program_end_pc = initial_fp - FELT_2;
 
         let program: Vec<&Felt> = memory
             .iter()

diff --git a/crates/air/src/layout/dynamic/mod.rs b/crates/air/src/layout/dynamic/mod.rs
@@ -473,6 +473,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -744,20 +749,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);

diff --git a/crates/air/src/layout/mod.rs b/crates/air/src/layout/mod.rs
@@ -17,7 +17,7 @@ pub mod starknet_with_keccak;
 pub mod stark_curve {
     use starknet_crypto::Felt;
 
-    pub const ALPHA: Felt = Felt::from_hex_unchecked("1");
+    pub const ALPHA: Felt = Felt::from_hex_unchecked("0x1");
     pub const BETA: Felt = Felt::from_hex_unchecked(
         "0x6f21413efbe40de150e596d72f7a8c5609ad26c15c915c1f4cdfcb99cee9e89",
     );
@@ -147,6 +147,9 @@ pub enum PublicInputError {
     #[error("invalid number of builtin copies")]
     CopiesInvalid,
 
+    #[error("invalid number of segments")]
+    InvalidSegments,
+
     #[error("dynamic params missing")]
     DynamicParamsMissing,
 
@@ -210,6 +213,9 @@ pub enum PublicInputError {
     #[error("invalid number of builtin copies")]
     CopiesInvalid,
 
+    #[error("invalid number of segments")]
+    InvalidSegments,
+
     #[error("dynamic params missing")]
     DynamicParamsMissing,
 

diff --git a/crates/air/src/layout/recursive/mod.rs b/crates/air/src/layout/recursive/mod.rs
@@ -276,6 +276,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -359,20 +364,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);

diff --git a/crates/air/src/layout/recursive_with_poseidon/mod.rs b/crates/air/src/layout/recursive_with_poseidon/mod.rs
@@ -318,6 +318,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -416,20 +421,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);

diff --git a/crates/air/src/layout/small/mod.rs b/crates/air/src/layout/small/mod.rs
@@ -275,6 +275,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -358,20 +363,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);

diff --git a/crates/air/src/layout/starknet/mod.rs b/crates/air/src/layout/starknet/mod.rs
@@ -360,6 +360,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -488,20 +493,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);

diff --git a/crates/air/src/layout/starknet_with_keccak/mod.rs b/crates/air/src/layout/starknet_with_keccak/mod.rs
@@ -394,6 +394,11 @@ impl LayoutTrait for Layout {
             PublicInputError::TraceLengthInvalid
         );
 
+        ensure!(
+            public_input.segments.len() == segments::N_SEGMENTS,
+            PublicInputError::InvalidSegments
+        );
+
         ensure!(FELT_0 <= public_input.range_check_min, PublicInputError::RangeCheckInvalid);
         ensure!(
             public_input.range_check_min < public_input.range_check_max,
@@ -537,20 +542,20 @@ impl LayoutTrait for Layout {
             .stop_ptr;
         let initial_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .begin_addr;
         let initial_fp = initial_ap;
         let final_ap = public_segments
             .get(segments::EXECUTION)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::EXECUTION })?
             .stop_ptr;
         let output_start = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .begin_addr;
         let output_stop = public_segments
             .get(segments::OUTPUT)
-            .ok_or(PublicInputError::SegmentMissing { segment: segments::PROGRAM })?
+            .ok_or(PublicInputError::SegmentMissing { segment: segments::OUTPUT })?
             .stop_ptr;
 
         ensure!(initial_ap < MAX_ADDRESS, PublicInputError::MaxSteps);