Implement PublicValues retrieval from public inputs (0xPolygonZero#…

…1405)

* Implement PublicValues retrieval from public inputs

* Use utility method

* Remove generic argument

* Typo

evm/src/fixed_recursive_verifier.rs

@@ -1100,13 +1100,12 @@ where
             }
 
             // Initialize the genesis state trie digest.
-            let genesis_state_trie_keys = TrieRootsTarget::SIZE * 2
-                + BlockMetadataTarget::SIZE
-                + BlockHashesTarget::BLOCK_HASHES_SIZE
-                ..TrieRootsTarget::SIZE * 2
-                    + BlockMetadataTarget::SIZE
-                    + BlockHashesTarget::BLOCK_HASHES_SIZE
-                    + 8;
+            let genesis_state_trie_keys =
+                TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE + BlockHashesTarget::SIZE
+                    ..TrieRootsTarget::SIZE * 2
+                        + BlockMetadataTarget::SIZE
+                        + BlockHashesTarget::SIZE
+                        + 8;
             for (key, &value) in genesis_state_trie_keys.zip_eq(&h256_limbs::<F>(
                 public_values.extra_block_data.genesis_state_trie_root,
             )) {
@@ -1115,9 +1114,7 @@ where
 
             // Initialize block hashes.
             let block_hashes_keys = TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE
-                ..TrieRootsTarget::SIZE * 2
-                    + BlockMetadataTarget::SIZE
-                    + BlockHashesTarget::BLOCK_HASHES_SIZE
+                ..TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE + BlockHashesTarget::SIZE
                     - 8;
 
             for i in 0..public_values.block_hashes.prev_hashes.len() - 1 {
@@ -1126,10 +1123,8 @@ where
                     nonzero_pis.insert(block_hashes_keys.start + 8 * (i + 1) + j, targets[j]);
                 }
             }
-            let block_hashes_current_start = TrieRootsTarget::SIZE * 2
-                + BlockMetadataTarget::SIZE
-                + BlockHashesTarget::BLOCK_HASHES_SIZE
-                - 8;
+            let block_hashes_current_start =
+                TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE + BlockHashesTarget::SIZE - 8;
             let cur_targets = h256_limbs::<F>(public_values.block_hashes.prev_hashes[255]);
             for i in 0..8 {
                 nonzero_pis.insert(block_hashes_current_start + i, cur_targets[i]);

evm/src/proof.rs

@@ -1,4 +1,4 @@
-use ethereum_types::{Address, H256, U256};
+use ethereum_types::{Address, H160, H256, U256};
 use itertools::Itertools;
 use plonky2::field::extension::{Extendable, FieldExtension};
 use plonky2::fri::oracle::PolynomialBatch;
@@ -19,6 +19,7 @@ use serde::{Deserialize, Serialize};
 use crate::all_stark::NUM_TABLES;
 use crate::config::StarkConfig;
 use crate::cross_table_lookup::GrandProductChallengeSet;
+use crate::util::{get_h160, get_h256, h2u};
 
 /// A STARK proof for each table, plus some metadata used to create recursive wrapper proofs.
 #[derive(Debug, Clone)]
@@ -47,7 +48,7 @@ pub(crate) struct AllProofChallenges<F: RichField + Extendable<D>, const D: usiz
 }
 
 /// Memory values which are public.
-#[derive(Debug, Clone, Default, Deserialize, Serialize)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
 pub struct PublicValues {
     /// Trie hashes before the execution of the local state transition
     pub trie_roots_before: TrieRoots,
@@ -61,8 +62,50 @@ pub struct PublicValues {
     pub extra_block_data: ExtraBlockData,
 }
 
+impl PublicValues {
+    /// Extracts public values from the given public inputs of a proof.
+    /// Public values are always the first public inputs added to the circuit,
+    /// so we can start extracting at index 0.
+    pub fn from_public_inputs<F: RichField>(pis: &[F]) -> Self {
+        assert!(
+            pis.len()
+                > TrieRootsTarget::SIZE * 2
+                    + BlockMetadataTarget::SIZE
+                    + BlockHashesTarget::SIZE
+                    + ExtraBlockDataTarget::SIZE
+                    - 1
+        );
+
+        let trie_roots_before = TrieRoots::from_public_inputs(&pis[0..TrieRootsTarget::SIZE]);
+        let trie_roots_after =
+            TrieRoots::from_public_inputs(&pis[TrieRootsTarget::SIZE..TrieRootsTarget::SIZE * 2]);
+        let block_metadata = BlockMetadata::from_public_inputs(
+            &pis[TrieRootsTarget::SIZE * 2..TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE],
+        );
+        let block_hashes = BlockHashes::from_public_inputs(
+            &pis[TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE
+                ..TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE + BlockHashesTarget::SIZE],
+        );
+        let extra_block_data = ExtraBlockData::from_public_inputs(
+            &pis[TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE + BlockHashesTarget::SIZE
+                ..TrieRootsTarget::SIZE * 2
+                    + BlockMetadataTarget::SIZE
+                    + BlockHashesTarget::SIZE
+                    + ExtraBlockDataTarget::SIZE],
+        );
+
+        Self {
+            trie_roots_before,
+            trie_roots_after,
+            block_metadata,
+            block_hashes,
+            extra_block_data,
+        }
+    }
+}
+
 /// Trie hashes.
-#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
 pub struct TrieRoots {
     /// State trie hash.
     pub state_root: H256,
@@ -72,6 +115,22 @@ pub struct TrieRoots {
     pub receipts_root: H256,
 }
 
+impl TrieRoots {
+    pub fn from_public_inputs<F: RichField>(pis: &[F]) -> Self {
+        assert!(pis.len() == TrieRootsTarget::SIZE);
+
+        let state_root = get_h256(&pis[0..8]);
+        let transactions_root = get_h256(&pis[8..16]);
+        let receipts_root = get_h256(&pis[16..24]);
+
+        Self {
+            state_root,
+            transactions_root,
+            receipts_root,
+        }
+    }
+}
+
 // There should be 256 previous hashes stored, so the default should also contain 256 values.
 impl Default for BlockHashes {
     fn default() -> Self {
@@ -88,7 +147,7 @@ impl Default for BlockHashes {
 ///
 /// When the block number is less than 256, dummy values, i.e. `H256::default()`,
 /// should be used for the additional block hashes.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct BlockHashes {
     /// The previous 256 hashes to the current block. The leftmost hash, i.e. `prev_hashes[0]`,
     /// is the oldest, and the rightmost, i.e. `prev_hashes[255]` is the hash of the parent block.
@@ -97,9 +156,23 @@ pub struct BlockHashes {
     pub cur_hash: H256,
 }
 
+impl BlockHashes {
+    pub fn from_public_inputs<F: RichField>(pis: &[F]) -> Self {
+        assert!(pis.len() == BlockHashesTarget::SIZE);
+
+        let prev_hashes: [H256; 256] = core::array::from_fn(|i| get_h256(&pis[8 * i..8 + 8 * i]));
+        let cur_hash = get_h256(&pis[2048..2056]);
+
+        Self {
+            prev_hashes: prev_hashes.to_vec(),
+            cur_hash,
+        }
+    }
+}
+
 /// Metadata contained in a block header. Those are identical between
 /// all state transition proofs within the same block.
-#[derive(Debug, Clone, Default, Deserialize, Serialize)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
 pub struct BlockMetadata {
     /// The address of this block's producer.
     pub block_beneficiary: Address,
@@ -123,9 +196,40 @@ pub struct BlockMetadata {
     pub block_bloom: [U256; 8],
 }
 
+impl BlockMetadata {
+    pub fn from_public_inputs<F: RichField>(pis: &[F]) -> Self {
+        assert!(pis.len() == BlockMetadataTarget::SIZE);
+
+        let block_beneficiary = get_h160(&pis[0..5]);
+        let block_timestamp = pis[5].to_canonical_u64().into();
+        let block_number = pis[6].to_canonical_u64().into();
+        let block_difficulty = pis[7].to_canonical_u64().into();
+        let block_random = get_h256(&pis[8..16]);
+        let block_gaslimit = pis[16].to_canonical_u64().into();
+        let block_chain_id = pis[17].to_canonical_u64().into();
+        let block_base_fee =
+            (pis[18].to_canonical_u64() + (pis[19].to_canonical_u64() << 32)).into();
+        let block_gas_used = pis[20].to_canonical_u64().into();
+        let block_bloom = core::array::from_fn(|i| h2u(get_h256(&pis[21 + 8 * i..29 + 8 * i])));
+
+        Self {
+            block_beneficiary,
+            block_timestamp,
+            block_number,
+            block_difficulty,
+            block_random,
+            block_gaslimit,
+            block_chain_id,
+            block_base_fee,
+            block_gas_used,
+            block_bloom,
+        }
+    }
+}
+
 /// Additional block data that are specific to the local transaction being proven,
 /// unlike `BlockMetadata`.
-#[derive(Debug, Clone, Default, Deserialize, Serialize)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
 pub struct ExtraBlockData {
     /// The state trie digest of the genesis block.
     pub genesis_state_trie_root: H256,
@@ -142,6 +246,26 @@ pub struct ExtraBlockData {
     pub gas_used_after: U256,
 }
 
+impl ExtraBlockData {
+    pub fn from_public_inputs<F: RichField>(pis: &[F]) -> Self {
+        assert!(pis.len() == ExtraBlockDataTarget::SIZE);
+
+        let genesis_state_trie_root = get_h256(&pis[0..8]);
+        let txn_number_before = pis[8].to_canonical_u64().into();
+        let txn_number_after = pis[9].to_canonical_u64().into();
+        let gas_used_before = pis[10].to_canonical_u64().into();
+        let gas_used_after = pis[11].to_canonical_u64().into();
+
+        Self {
+            genesis_state_trie_root,
+            txn_number_before,
+            txn_number_after,
+            gas_used_before,
+            gas_used_after,
+        }
+    }
+}
+
 /// Memory values which are public.
 /// Note: All the larger integers are encoded with 32-bit limbs in little-endian order.
 #[derive(Eq, PartialEq, Debug)]
@@ -285,7 +409,7 @@ impl PublicValuesTarget {
             pis.len()
                 > TrieRootsTarget::SIZE * 2
                     + BlockMetadataTarget::SIZE
-                    + BlockHashesTarget::BLOCK_HASHES_SIZE
+                    + BlockHashesTarget::SIZE
                     + ExtraBlockDataTarget::SIZE
                     - 1
         );
@@ -303,15 +427,13 @@ impl PublicValuesTarget {
                 &pis[TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE
                     ..TrieRootsTarget::SIZE * 2
                         + BlockMetadataTarget::SIZE
-                        + BlockHashesTarget::BLOCK_HASHES_SIZE],
+                        + BlockHashesTarget::SIZE],
             ),
             extra_block_data: ExtraBlockDataTarget::from_public_inputs(
-                &pis[TrieRootsTarget::SIZE * 2
-                    + BlockMetadataTarget::SIZE
-                    + BlockHashesTarget::BLOCK_HASHES_SIZE
+                &pis[TrieRootsTarget::SIZE * 2 + BlockMetadataTarget::SIZE + BlockHashesTarget::SIZE
                     ..TrieRootsTarget::SIZE * 2
                         + BlockMetadataTarget::SIZE
-                        + BlockHashesTarget::BLOCK_HASHES_SIZE
+                        + BlockHashesTarget::SIZE
                         + ExtraBlockDataTarget::SIZE],
             ),
         }
@@ -568,7 +690,7 @@ pub(crate) struct BlockHashesTarget {
 
 impl BlockHashesTarget {
     /// Number of `Target`s required for previous and current block hashes.
-    pub(crate) const BLOCK_HASHES_SIZE: usize = 2056;
+    pub(crate) const SIZE: usize = 2056;
 
     /// Extracts the previous and current block hash `Target`s from the public input `Target`s.
     /// The provided `pis` should start with the block hashes.

evm/src/util.rs

@@ -210,3 +210,25 @@ pub(crate) fn biguint_to_mem_vec(x: BigUint) -> Vec<U256> {
 pub(crate) fn h2u(h: H256) -> U256 {
     U256::from_big_endian(&h.0)
 }
+
+pub(crate) fn get_h160<F: RichField>(slice: &[F]) -> H160 {
+    H160::from_slice(
+        &slice
+            .iter()
+            .rev()
+            .map(|x| x.to_canonical_u64() as u32)
+            .flat_map(|limb| limb.to_be_bytes())
+            .collect_vec(),
+    )
+}
+
+pub(crate) fn get_h256<F: RichField>(slice: &[F]) -> H256 {
+    H256::from_slice(
+        &slice
+            .iter()
+            .rev()
+            .map(|x| x.to_canonical_u64() as u32)
+            .flat_map(|limb| limb.to_be_bytes())
+            .collect_vec(),
+    )
+}

evm/tests/empty_txn_list.rs

@@ -15,7 +15,7 @@ use plonky2_evm::all_stark::AllStark;
 use plonky2_evm::config::StarkConfig;
 use plonky2_evm::fixed_recursive_verifier::AllRecursiveCircuits;
 use plonky2_evm::generation::{GenerationInputs, TrieInputs};
-use plonky2_evm::proof::{BlockHashes, BlockMetadata, TrieRoots};
+use plonky2_evm::proof::{BlockHashes, BlockMetadata, PublicValues, TrieRoots};
 use plonky2_evm::Node;
 
 type F = GoldilocksField;
@@ -132,8 +132,12 @@ fn test_empty_txn_list() -> anyhow::Result<()> {
     timing.filter(Duration::from_millis(100)).print();
     all_circuits.verify_root(root_proof.clone())?;
 
+    // Test retrieved public values from the proof public inputs.
+    let retrieved_public_values = PublicValues::from_public_inputs(&root_proof.public_inputs);
+    assert_eq!(retrieved_public_values, public_values);
+
     // We can duplicate the proofs here because the state hasn't mutated.
-    let (agg_proof, public_values) = all_circuits.prove_aggregation(
+    let (agg_proof, agg_public_values) = all_circuits.prove_aggregation(
         false,
         &root_proof,
         public_values.clone(),
@@ -143,9 +147,18 @@ fn test_empty_txn_list() -> anyhow::Result<()> {
     )?;
     all_circuits.verify_aggregation(&agg_proof)?;
 
-    let (block_proof, _) = all_circuits.prove_block(None, &agg_proof, public_values)?;
+    // Test retrieved public values from the proof public inputs.
+    let retrieved_public_values = PublicValues::from_public_inputs(&agg_proof.public_inputs);
+    assert_eq!(retrieved_public_values, agg_public_values);
+
+    let (block_proof, block_public_values) =
+        all_circuits.prove_block(None, &agg_proof, agg_public_values)?;
     all_circuits.verify_block(&block_proof)?;
 
+    // Test retrieved public values from the proof public inputs.
+    let retrieved_public_values = PublicValues::from_public_inputs(&block_proof.public_inputs);
+    assert_eq!(retrieved_public_values, block_public_values);
+
     // Get the verifier associated to these preprocessed circuits, and have it verify the block_proof.
     let verifier = all_circuits.final_verifier_data();
     verifier.verify(block_proof)