cairo_proof_parser/
output.rs

1use starknet_crypto::{poseidon_hash_many, FieldElement};
2use std::collections::HashMap;
3use std::convert::TryInto;
4
5use crate::parse_raw;
6
7pub struct ExtractOutputResult {
8    pub program_output: Vec<FieldElement>,
9    pub program_output_hash: FieldElement,
10}
11
12pub fn extract_output(input: String) -> anyhow::Result<ExtractOutputResult> {
13    // Parse the input string into a proof structure
14    let proof = parse_raw(&input)?;
15
16    // Retrieve the output segment from the proof
17    let output_segment = proof
18        .public_input
19        .segments
20        .get(2)
21        .ok_or_else(|| anyhow::Error::msg("Output segment not found"))?;
22
23    // Construct a map for the main page elements
24    let mut main_page_map = HashMap::new();
25    for element in &proof.public_input.main_page {
26        let value_bytes = element.value.to_bytes_be();
27        let padded_value = vec![0u8; 32 - value_bytes.len()]
28            .iter()
29            .chain(value_bytes.iter())
30            .copied()
31            .collect::<Vec<u8>>();
32        let field_element = FieldElement::from_bytes_be(
33            &padded_value.try_into().expect("Failed to convert to array"),
34        )
35        .expect("Failed to convert to FieldElement");
36
37        main_page_map.insert(element.address, field_element);
38    }
39
40    // Extract program output using the address range in the output segment
41    let program_output: Vec<FieldElement> = (output_segment.begin_addr..output_segment.stop_ptr)
42        .map(|addr| {
43            *main_page_map
44                .get(&addr)
45                .expect("Address not found in main page map")
46        })
47        .collect();
48
49    // Calculate the Poseidon hash of the program output
50    let program_output_hash = poseidon_hash_many(&program_output);
51
52    Ok(ExtractOutputResult {
53        program_output,
54        program_output_hash,
55    })
56}