Merge pull request #90 from HerodotusDev/test/refactor

Test/refactor

src/air/public_input.cairo

@@ -3,7 +3,7 @@ use cairo_verifier::{
     common::{
         flip_endianness::FlipEndiannessTrait, array_append::ArrayAppendTrait, blake2s::blake2s,
         math::{pow, Felt252PartialOrd, Felt252Div}, asserts::assert_range_u128_le,
-        array_print::SpanPrintTrait
+        array_print::SpanPrintTrait, hash::hash_felts,
     },
     air::{
         public_memory::{
@@ -14,10 +14,9 @@ use cairo_verifier::{
             MAX_RANGE_CHECK, LAYOUT_CODE, PEDERSEN_BUILTIN_RATIO, RC_BUILTIN_RATIO, BITWISE_RATIO
         }
     },
-    domains::StarkDomains, common::hash::hash_felts
+    domains::StarkDomains
 };
 
-
 #[derive(Drop, Copy, PartialEq)]
 struct SegmentInfo {
     // Start address of the memory segment.
@@ -188,8 +187,9 @@ impl PublicInputImpl of PublicInputTrait {
         let program_end_pc = initial_fp - 2;
         let program_len = program_end_pc - initial_pc;
         let program = memory
-            .extract_range(initial_pc.try_into().unwrap(), program_len.try_into().unwrap());
-        memory_index += program.len().into();
+            .extract_range(
+                initial_pc.try_into().unwrap(), program_len.try_into().unwrap(), ref memory_index
+            );
 
         assert(
             *program[0] == 0x40780017fff7fff, 'Invalid program'
@@ -213,33 +213,40 @@ impl PublicInputImpl of PublicInputTrait {
         memory_index += 2;
 
         // 2.2 Main arguments and return values
-        memory
-            .verify_stack(
-                initial_ap, *public_segments.at(2).begin_addr, builtins.span(), memory_index.into()
-            );
-        memory_index += builtins.len();
+        let mut begin_addresses = ArrayTrait::new();
+        let mut stop_addresses = ArrayTrait::new();
+        let mut i = 0;
+        let builtins_len = builtins.len();
+        loop {
+            if i == builtins_len {
+                break;
+            }
 
+            begin_addresses.append(*public_segments.at(2 + i).begin_addr);
+            stop_addresses.append(*public_segments.at(2 + i).stop_ptr);
+
+            i += 1;
+        };
+        memory.verify_stack(initial_ap, begin_addresses.span(), builtins_len, ref memory_index);
         memory
             .verify_stack(
-                final_ap - builtins.len().into(),
-                *public_segments.at(2).stop_ptr,
-                builtins.span(),
-                memory_index.into()
+                final_ap - builtins_len.into(),
+                stop_addresses.span(),
+                builtins_len,
+                ref memory_index
             );
-        memory_index += builtins.len();
 
         // 3. Output segment 
         let output_len = output_stop - output_start;
         let output = memory
             .extract_range(
-                memory_index + output_start.try_into().unwrap(), output_len.try_into().unwrap()
+                output_start.try_into().unwrap(), output_len.try_into().unwrap(), ref memory_index
             );
-        memory_index += output.len().into();
         let output_hash = hash_felts(output);
 
         // Check main page len
         assert(
-            *memory.at(memory_index) == *self.main_page.at(self.main_page.len() - 1),
+            *memory.at(memory_index - 1) == *self.main_page.at(self.main_page.len() - 1),
             'Invalid main page len'
         );
 

src/air/public_memory.cairo

@@ -45,7 +45,7 @@ impl PageImpl of PageTrait {
         }
     }
 
-    fn extract_range(self: @Page, addr: u32, len: usize) -> Span<felt252> {
+    fn extract_range(self: @Page, addr: u32, len: usize, ref offset: usize) -> Span<felt252> {
         let mut arr = ArrayTrait::new();
         let mut i = 0;
 
@@ -54,54 +54,56 @@ impl PageImpl of PageTrait {
                 break arr.span();
             }
 
-            let current = *self.at(addr + i);
+            let current = *self.at(offset);
 
             // TODO is this needed? If not we can just use slice directly 
             assert(current.address == (addr + i).into(), 'Invalid address');
             arr.append(current.value);
             i += 1;
+            offset += 1;
         }
     }
 
     fn verify_stack(
         self: @Page,
         start_ap: felt252,
-        segment_address: felt252,
-        builtins: Span<felt252>,
-        memory_index: felt252
+        segment_addresses: Span<felt252>,
+        builtins_len: usize,
+        ref offset: usize
     ) {
         let mut i = 0;
 
-        // TODO size of SegmentInfo
-        let size = 2;
         loop {
-            if i == builtins.len() {
+            if i == builtins_len {
                 break;
             }
 
-            let current = *self.at(memory_index.try_into().unwrap() + i);
+            let current = *self.at(offset);
 
             assert(current.address == start_ap + i.into(), 'Invalid address');
-            assert(current.value == segment_address + size * (i.into() + 1), 'Invalid builtin');
+            assert(current.value == *segment_addresses.at(i), 'Invalid builtin');
             i += 1;
+            offset += 1;
         };
     }
 }
 
-fn get_continuous_pages_product(page_headers: Span<ContinuousPageHeader>) -> (felt252, felt252) {
+fn get_continuous_pages_product(
+    mut page_headers: Span<ContinuousPageHeader>
+) -> (felt252, felt252) {
     let mut res = 1;
     let mut total_length = 0;
-    let mut i = 0;
+
     loop {
-        if i == page_headers.len() {
-            break (res, total_length);
+        match page_headers.pop_front() {
+            Option::Some(header) => {
+                res *= *header.prod;
+                total_length += *header.size;
+            },
+            Option::None => { break; }
         }
-        let current = page_headers.at(i);
-
-        res *= *current.prod;
-        total_length += *current.size;
+    };
 
-        i += 1;
-    }
+    (res, total_length)
 }
 

src/air/tests.cairo

@@ -1,7 +1,7 @@
 mod test_autogenerated;
+mod test_composition;
 mod test_diluted;
 mod test_pedersen;
 mod test_public_input;
-mod test_traces;
 mod test_public_memory_prod;
-mod test_composition;
+mod test_traces;