axiom_circuit/subquery/
account.rs

1use anyhow::Result;
2use axiom_codec::types::native::{AccountSubquery, AnySubquery};
3use axiom_query::axiom_eth::{halo2_base::AssignedValue, Field};
4use ethers::{
5    providers::{JsonRpcClient, Middleware, Provider},
6    types::{BigEndianHash, BlockId, H256, U256},
7};
8use num_derive::FromPrimitive;
9use num_traits::FromPrimitive;
10use tokio::runtime::Runtime;
11
12use super::{caller::FetchSubquery, types::AssignedAccountSubquery};
13use crate::impl_fr_from;
14
15#[derive(FromPrimitive, Copy, Clone)]
16pub enum AccountField {
17    Nonce,
18    Balance,
19    StorageHash,
20    CodeHash,
21}
22impl_fr_from!(AccountField);
23
24pub async fn get_account_field_value<P: JsonRpcClient>(
25    provider: &Provider<P>,
26    query: AccountSubquery,
27) -> Result<H256> {
28    let block_id = Some(BlockId::from(query.block_number as u64));
29
30    let account_field = AccountField::from_u32(query.field_idx).expect("Invalid field index");
31    let val = match account_field {
32        AccountField::Nonce => {
33            let nonce = provider.get_transaction_count(query.addr, block_id).await;
34            if nonce.is_err() {
35                return Ok(get_nonexistent_account_field_value(
36                    query.field_idx as usize,
37                ));
38            }
39            H256::from_uint(&nonce.unwrap())
40        }
41        AccountField::Balance => {
42            let balance = provider.get_balance(query.addr, block_id).await;
43            if balance.is_err() {
44                return Ok(get_nonexistent_account_field_value(
45                    query.field_idx as usize,
46                ));
47            }
48            H256::from_uint(&balance.unwrap())
49        }
50        AccountField::StorageHash => {
51            let proof = provider.get_proof(query.addr, vec![], block_id).await;
52            if proof.is_err() {
53                return Ok(get_nonexistent_account_field_value(
54                    query.field_idx as usize,
55                ));
56            }
57            proof.unwrap().storage_hash
58        }
59        AccountField::CodeHash => {
60            let proof = provider.get_proof(query.addr, vec![], block_id).await;
61            if proof.is_err() {
62                return Ok(get_nonexistent_account_field_value(
63                    query.field_idx as usize,
64                ));
65            }
66            proof.unwrap().code_hash
67        }
68    };
69
70    Ok(val)
71}
72
73impl<F: Field> FetchSubquery<F> for AssignedAccountSubquery<F> {
74    fn fetch<P: JsonRpcClient>(&self, p: &Provider<P>) -> Result<H256> {
75        let rt = Runtime::new()?;
76        let res = rt.block_on(get_account_field_value(p, (*self).into()))?;
77        Ok(res)
78    }
79
80    fn any_subquery(&self) -> AnySubquery {
81        AnySubquery::Account((*self).into())
82    }
83
84    fn flatten(&self) -> Vec<AssignedValue<F>> {
85        vec![self.block_number, self.addr, self.field_idx]
86    }
87}
88
89fn get_nonexistent_account_field_value(field_idx: usize) -> H256 {
90    let account_field = AccountField::from_usize(field_idx).expect("Invalid field index");
91    match account_field {
92        AccountField::Nonce => H256::from_uint(&U256::from(0)),
93        AccountField::Balance => H256::from_uint(&U256::from(0)),
94        AccountField::StorageHash => H256::from_uint(
95            &U256::from_dec_str(
96                "0x56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421",
97            )
98            .unwrap(),
99        ),
100        AccountField::CodeHash => H256::from_uint(&U256::from(0)),
101    }
102}
axiom_circuit/subquery/account.rs

axiom_circuit/subquery/
account.rs
