Merge pull request #34 from HerodotusDev/feat/public-input-verify

Public input verification

src/air/constants.cairo

@@ -41,3 +41,15 @@ mod segments {
     const PROGRAM: usize = 0;
     const RANGE_CHECK: usize = 4;
 }
+
+// Pedersen builtin
+const SHIFT_POINT_X: felt252 = 0x49ee3eba8c1600700ee1b87eb599f16716b0b1022947733551fde4050ca6804;
+const SHIFT_POINT_Y: felt252 = 0x3ca0cfe4b3bc6ddf346d49d06ea0ed34e621062c0e056c1d0405d266e10268a;
+
+const SECURITY_BITS: felt252 = 128;
+const MAX_ADDRESS: felt252 = 0xffffffffffffffff;
+const INITIAL_PC: felt252 = 1;
+
+fn get_builtins() -> Array<felt252> {
+    array!['output', 'pedersen', 'range_check', 'bitwise']
+}

src/air/public_input.cairo

@@ -9,11 +9,15 @@ use cairo_verifier::{
         public_memory::{
             Page, PageTrait, ContinuousPageHeader, get_continuous_pages_product, AddrValueSize
         },
-        constants
+        constants::{
+            segments, MAX_ADDRESS, get_builtins, INITIAL_PC, MAX_LOG_N_STEPS, CPU_COMPONENT_HEIGHT,
+            MAX_RANGE_CHECK, LAYOUT_CODE, PEDERSEN_BUILTIN_RATIO, RC_BUILTIN_RATIO, BITWISE_RATIO
+        }
     },
-    domains::StarkDomains
+    domains::StarkDomains, common::hash::hash_felts
 };
 
+
 #[derive(Drop, Copy, PartialEq)]
 struct SegmentInfo {
     // Start address of the memory segment.
@@ -155,35 +159,120 @@ impl PublicInputImpl of PublicInputTrait {
         (prod, total_length)
     }
 
-    fn validate(self: @PublicInput, domains: @StarkDomains) {
-        assert_range_u128_le(*self.log_n_steps, constants::MAX_LOG_N_STEPS);
-        let n_steps = pow(2, *self.log_n_steps);
+    fn verify(self: @PublicInput) -> (felt252, felt252) {
+        let public_segments = self.segments;
+
+        let initial_pc = *public_segments.at(segments::PROGRAM).begin_addr;
+        let final_pc = *public_segments.at(segments::PROGRAM).stop_ptr;
+        let initial_ap = *public_segments.at(segments::EXECUTION).begin_addr;
+        let initial_fp = initial_ap;
+        let final_ap = *public_segments.at(segments::EXECUTION).stop_ptr;
+        let output_start = *public_segments.at(segments::OUTPUT).begin_addr;
+        let output_stop = *public_segments.at(segments::OUTPUT).stop_ptr;
+
+        assert(initial_ap < MAX_ADDRESS, 'Invalid initial_ap');
+        assert(final_ap < MAX_ADDRESS, 'Invalid final_ap');
+
+        // TODO support more pages?
+        assert(self.continuous_page_headers.len() == 0, 'Invalid continuous_page_headers');
+
+        let builtins = get_builtins();
+        let memory = self.main_page;
+
+        // 1. Program segment
+        assert(initial_pc == INITIAL_PC, 'Invalid initial_pc');
+        assert(final_pc == INITIAL_PC + 4, 'Invalid final_pc');
+
+        let mut memory_index: usize = 0;
+
+        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();
+
         assert(
-            n_steps * constants::CPU_COMPONENT_HEIGHT == *domains.trace_domain_size,
-            'Wrong trace size'
+            *program[0] == 0x40780017fff7fff, 'Invalid program'
+        ); // Instruction: ap += N_BUILTINS.
+        assert(*program[1] == builtins.len().into(), 'Invalid program');
+        assert(*program[2] == 0x1104800180018000, 'Invalid program'); // Instruction: call rel ?.
+        assert(*program[4] == 0x10780017fff7fff, 'Invalid program'); // Instruction: jmp rel 0.
+        assert(*program[5] == 0x0, 'Invalid program');
+
+        let program_hash = hash_felts(program);
+
+        // 2. Execution segment
+        // 2.1 Initial_fp, initial_pc
+        let fp2 = *memory.at(memory_index);
+        assert(fp2.address == initial_fp - 2, 'Invalid fp2 addr');
+        assert(fp2.value == initial_fp, 'Invalid fp2 val');
+
+        let fp1 = *memory.at(memory_index + 1);
+        assert(fp1.address == initial_fp - 1, 'Invalid fp1 addr');
+        assert(fp1.value == 0, 'Invalid fp1 val');
+        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();
+
+        memory
+            .verify_stack(
+                final_ap - builtins.len().into(),
+                *public_segments.at(2).stop_ptr,
+                builtins.span(),
+                memory_index.into()
+            );
+        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()
+            );
+        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),
+            'Invalid main page len'
         );
 
+        (program_hash, output_hash)
+    }
+
+    fn validate(self: @PublicInput, domains: @StarkDomains) {
+        assert_range_u128_le(*self.log_n_steps, MAX_LOG_N_STEPS);
+        let n_steps = pow(2, *self.log_n_steps);
+        assert(n_steps * CPU_COMPONENT_HEIGHT == *domains.trace_domain_size, 'Wrong trace size');
+
         assert(0 <= *self.rc_min, 'wrong rc_min');
         assert(*self.rc_min < *self.rc_max, 'wrong rc range');
-        assert(*self.rc_max <= constants::MAX_RANGE_CHECK, 'wrong rc_max');
+        assert(*self.rc_max <= MAX_RANGE_CHECK, 'wrong rc_max');
 
-        assert(*self.layout == constants::LAYOUT_CODE, 'wrong layout code');
+        assert(*self.layout == LAYOUT_CODE, 'wrong layout code');
 
-        let pedersen_copies = n_steps / constants::PEDERSEN_BUILTIN_RATIO;
-        let pedersen_uses = (*self.segments.at(constants::segments::PEDERSEN).stop_ptr
-            - *self.segments.at(constants::segments::PEDERSEN).begin_addr)
+        let pedersen_copies = n_steps / PEDERSEN_BUILTIN_RATIO;
+        let pedersen_uses = (*self.segments.at(segments::PEDERSEN).stop_ptr
+            - *self.segments.at(segments::PEDERSEN).begin_addr)
             / 3;
         assert_range_u128_le(pedersen_uses, pedersen_copies);
 
-        let range_check_copies = n_steps / constants::RC_BUILTIN_RATIO;
-        let range_check_uses = *self.segments.at(constants::segments::RANGE_CHECK).stop_ptr
-            - *self.segments.at(constants::segments::RANGE_CHECK).begin_addr;
+        let range_check_copies = n_steps / RC_BUILTIN_RATIO;
+        let range_check_uses = *self.segments.at(segments::RANGE_CHECK).stop_ptr
+            - *self.segments.at(segments::RANGE_CHECK).begin_addr;
         assert_range_u128_le(range_check_uses, range_check_copies);
 
-        let bitwise_copies = n_steps / constants::BITWISE_RATIO;
-        let bitwise_uses = (*self.segments.at(constants::segments::BITWISE).stop_ptr
-            - *self.segments.at(constants::segments::BITWISE).begin_addr)
+        let bitwise_copies = n_steps / BITWISE_RATIO;
+        let bitwise_uses = (*self.segments.at(segments::BITWISE).stop_ptr
+            - *self.segments.at(segments::BITWISE).begin_addr)
             / 5;
         assert_range_u128_le(bitwise_uses, bitwise_copies);
     }
 }
+

src/air/public_memory.cairo

@@ -44,6 +44,48 @@ impl PageImpl of PageTrait {
             i += 1;
         }
     }
+
+    fn extract_range(self: @Page, addr: u32, len: usize) -> Span<felt252> {
+        let mut arr = ArrayTrait::new();
+        let mut i = 0;
+
+        loop {
+            if i == len {
+                break arr.span();
+            }
+
+            let current = *self.at(addr + i);
+
+            // 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;
+        }
+    }
+
+    fn verify_stack(
+        self: @Page,
+        start_ap: felt252,
+        segment_address: felt252,
+        builtins: Span<felt252>,
+        memory_index: felt252
+    ) {
+        let mut i = 0;
+
+        // TODO size of SegmentInfo
+        let size = 2;
+        loop {
+            if i == builtins.len() {
+                break;
+            }
+
+            let current = *self.at(memory_index.try_into().unwrap() + i);
+
+            assert(current.address == start_ap + i.into(), 'Invalid address');
+            assert(current.value == segment_address + size * (i.into() + 1), 'Invalid builtin');
+            i += 1;
+        };
+    }
 }
 
 fn get_continuous_pages_product(page_headers: Span<ContinuousPageHeader>) -> (felt252, felt252) {

src/common.cairo

@@ -11,6 +11,7 @@ mod flip_endianness;
 mod from_span;
 mod horner_eval;
 mod math;
+mod hash;
 mod merge_sort;
 mod powers_array;
 

src/common/hash.cairo

@@ -0,0 +1,15 @@
+use pedersen::PedersenTrait;
+use hash::HashStateTrait;
+
+fn hash_felts(s: Span<felt252>) -> felt252 {
+    let mut hash_state = PedersenTrait::new(0);
+    let mut i = 0;
+
+    loop {
+        if i == s.len() {
+            break hash_state.finalize();
+        }
+
+        hash_state.update(*s.at(i));
+    }
+}