Merge pull request #16 from Lagrange-Labs/distributed-query-integration-modifications

Modifications for integration with `distributed-query`

Author: stjepangolemac

src/eth.rs

@@ -5,6 +5,7 @@ use eth_trie::{EthTrie, MemoryDB, Node, Trie};
 use ethers::{
     providers::{Http, Middleware, Provider},
     types::{Block, BlockId, Bytes, Transaction, TransactionReceipt, U64},
+    utils::hex,
 };
 use rlp::{Encodable, Rlp, RlpStream};
 #[cfg(feature = "ci")]
@@ -14,7 +15,7 @@ use std::sync::Arc;
 use crate::utils::keccak256;
 /// A wrapper around a transaction and its receipt. The receipt is used to filter
 /// bad transactions, so we only compute over valid transactions.
-pub(crate) struct TxAndReceipt(Transaction, TransactionReceipt);
+pub struct TxAndReceipt(Transaction, TransactionReceipt);
 
 impl TxAndReceipt {
     pub fn tx(&self) -> &Transaction {
@@ -188,7 +189,7 @@ pub(crate) fn rlp_opt<T: rlp::Encodable>(rlp: &mut rlp::RlpStream, opt: &Option<
 }
 
 /// Extract the hash in case of Extension node, or all the hashes in case of a Branch node.
-pub(crate) fn extract_child_hashes(rlp_data: &[u8]) -> Vec<Vec<u8>> {
+pub fn extract_child_hashes(rlp_data: &[u8]) -> Vec<Vec<u8>> {
     let rlp = Rlp::new(rlp_data);
     let mut hashes = Vec::new();
 

src/lib.rs

@@ -16,9 +16,10 @@ use serde::{Deserialize, Serialize};
 
 mod benches;
 mod eth;
+
 mod hash;
 mod rlp;
-mod transaction;
+pub mod transaction;
 mod utils;
 
 /// Bundle containing the raw proof, the verification key, and some common data

src/transaction/mod.rs

@@ -1,8 +1,7 @@
 mod header;
 mod mpt;
-mod proof;
-#[cfg(test)]
-mod prover;
+pub mod proof;
+pub mod prover;
 
 /// Length of a hash in bytes.
 const HASH_LEN: usize = 32;

src/transaction/mpt.rs

@@ -43,7 +43,7 @@ const HASH_LENGTH: usize = 32;
 
 /// There are different ways to extract values from a transaction. This enum
 /// list some.
-pub(crate) enum ExtractionMethod {
+pub enum ExtractionMethod {
     /// RLPBased decodes each header consecutively and extract the gas value
     /// TODO: Currently hardcode that the gas value is 3rd item in the tx list
     /// because we use const generics and can't pass the index as a parameter.

src/transaction/proof.rs

@@ -61,7 +61,11 @@ pub struct TransactionMPT {
 }
 
 impl TransactionMPT {
-    fn compute_proof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
+    pub fn compute_proof<
+        F: RichField + Extendable<D>,
+        C: GenericConfig<D, F = F>,
+        const D: usize,
+    >(
         self,
     ) -> Result<ProofTuple<F, C, D>> {
         let config = CircuitConfig::standard_recursion_config();
@@ -138,8 +142,8 @@ impl RootMPTHeader {
     }
 }
 
-struct HeaderAggregation {
-    header_proofs: Vec<Vec<u8>>,
+pub struct HeaderAggregation {
+    pub header_proofs: Vec<Vec<u8>>,
 }
 
 impl HeaderAggregation {

src/transaction/prover.rs

@@ -1,8 +1,5 @@
 use crate::eth::RLPBlock;
-use crate::hash::hash_to_fields;
 use crate::transaction::proof::{IntermediateMPT, ProofType, RootMPTHeader, TransactionMPT};
-use crate::transaction::PACKED_HASH_LEN;
-use crate::ByteProofTuple;
 use crate::{
     eth::{extract_child_hashes, BlockData},
     utils::keccak256,
@@ -16,7 +13,7 @@ use rlp::Encodable;
 use std::collections::HashMap;
 
 #[derive(Debug, Default)]
-pub(crate) enum TxFilter {
+pub enum TxFilter {
     // only prove tx that are less than this size in bytes
     BySize(usize),
     // only prove the tx which have these hashes
@@ -41,28 +38,29 @@ impl TxFilter {
         }
     }
 }
-struct TxBlockProver {
-    data: BlockData,
-    policy: TxFilter,
+pub struct TxBlockProver {
+    pub data: BlockData,
+    pub policy: TxFilter,
     #[cfg(test)]
-    nb_nodes: usize,
+    pub nb_nodes: usize,
 }
 
-struct MPTNode {
-    node_bytes: Vec<u8>, // RLP byte representation of the node
-    hash: Vec<u8>,       // its hash
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct MPTNode {
+    pub node_bytes: Vec<u8>, // RLP byte representation of the node
+    pub hash: Vec<u8>,       // its hash
     // ALL the children of this node if any (zero if leaf for example)
-    total_nb_children: usize,
+    pub total_nb_children: usize,
     // expected children hashes - will be filled in with only the hashes that we
     // traverse in the proofs of all tx we are looking at. For the nodes that we don't
     // traverse, their hash may be in children_hashes, and we'll only provide a "null" proof
     // for them.
-    exp_children: Vec<Vec<u8>>,
+    pub exp_children: Vec<Vec<u8>>,
     // child i : (key, proof) - key needed locate where is the hash of the child in the node
-    rcvd_children_proofs: Vec<ProverOutput>, // will be filled in
-    parent_hash: Option<Vec<u8>>, // indicator to go up one level when this node has been "proven"
+    pub rcvd_children_proofs: Vec<ProverOutput>, // will be filled in
+    pub parent_hash: Option<Vec<u8>>, // indicator to go up one level when this node has been "proven"
     /// Used for information about tx in the leaf node proving phase
-    transaction: Option<Transaction>,
+    pub transaction: Option<Transaction>,
 }
 impl MPTNode {
     fn is_provable(&self) -> bool {
@@ -93,7 +91,8 @@ impl MPTNode {
     }
 }
 
-struct ProverOutput {
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct ProverOutput {
     parent_hash: Option<Vec<u8>>,
     // hash of this node
     hash: Vec<u8>,
@@ -116,6 +115,7 @@ impl TxBlockProver {
 
     pub fn prove(&mut self) -> Result<Vec<u8>> {
         let (mut trie, leaves_hashes) = self.init_proofs_trie();
+
         println!(
             "[+] Built internal trie with {} leaves, and {} nodes in total",
             leaves_hashes.len(),
@@ -204,7 +204,9 @@ impl TxBlockProver {
             intermediate_node: node_bytes,
             children_proofs: inner_proofs,
         });
+
         let proof = prover.compute_proof()?;
+
         println!(
             "[+] OK Valid recursive proof for node hash {}",
             hex::encode(&node_hash)
@@ -232,10 +234,6 @@ impl TxBlockProver {
             transaction,
         });
         let proof = prover.compute_proof()?;
-        println!(
-            "[+] OK Valid proof for leaf - hash {}",
-            hex::encode(&leaf_hash)
-        );
         Ok(ProverOutput {
             parent_hash,
             proof,
@@ -246,7 +244,7 @@ impl TxBlockProver {
     // Returns the hashmap filled with the trie info
     // and returns the initial list of nodes's hash, which happen to be leaves, to prove
     #[allow(clippy::type_complexity)]
-    fn init_proofs_trie(&mut self) -> (HashMap<Vec<u8>, MPTNode>, Vec<Vec<u8>>) {
+    pub fn init_proofs_trie(&mut self) -> (HashMap<Vec<u8>, MPTNode>, Vec<Vec<u8>>) {
         // H(node) => { MPTNode() }
         let mut tree = HashMap::new();
         let mut leaves = Vec::new();
@@ -264,10 +262,12 @@ impl TxBlockProver {
             let idx = txr.receipt().transaction_index;
             let key = idx.rlp_bytes().to_vec();
             let proof_bytes = self.data.tx_trie.get_proof(&key).unwrap();
+
             let mut child_hash = None;
             for (i, node_bytes) in proof_bytes.iter().rev().enumerate() {
                 let idx_in_path = proof_bytes.len() - 1 - i;
                 let hash = keccak256(node_bytes);
+
                 if i == 0 {
                     leaves.push(hash.clone());
                 }
@@ -306,6 +306,7 @@ impl TxBlockProver {
     }
 }
 
+#[cfg(test)]
 mod test {
     use std::time::Instant;