axiom_circuit/subquery/
mapping.rs

1use anyhow::Result;
2use axiom_codec::types::native::{AnySubquery, SolidityNestedMappingSubquery, StorageSubquery};
3use axiom_query::axiom_eth::{halo2_base::AssignedValue, Field};
4use ethers::{
5    providers::{JsonRpcClient, Provider},
6    types::{BigEndianHash, H256},
7    utils::keccak256,
8};
9use tokio::runtime::Runtime;
10
11use super::{
12    caller::FetchSubquery, storage::get_storage_field_value,
13    types::AssignedSolidityNestedMappingSubquery,
14};
15
16pub async fn get_solidity_nested_mapping_field_value<P: JsonRpcClient>(
17    provider: &Provider<P>,
18    query: SolidityNestedMappingSubquery,
19) -> Result<H256> {
20    let mut slot = H256::from_uint(&query.mapping_slot);
21    for i in 0..query.mapping_depth {
22        let key = query.keys.get(i as usize).unwrap();
23        let concat_h256 = key
24            .as_bytes()
25            .iter()
26            .copied()
27            .chain(slot.as_bytes().to_vec())
28            .collect::<Vec<u8>>();
29        slot = H256::from(keccak256(concat_h256));
30    }
31
32    let storage_query = StorageSubquery {
33        block_number: query.block_number,
34        addr: query.addr,
35        slot: slot.into_uint(),
36    };
37
38    get_storage_field_value(provider, storage_query).await
39}
40
41impl<F: Field> FetchSubquery<F> for AssignedSolidityNestedMappingSubquery<F> {
42    fn fetch<P: JsonRpcClient>(&self, p: &Provider<P>) -> Result<H256> {
43        let rt = Runtime::new()?;
44        let val = rt.block_on(get_solidity_nested_mapping_field_value(p, (*self).into()))?;
45        Ok(val)
46    }
47
48    fn any_subquery(&self) -> AnySubquery {
49        AnySubquery::SolidityNestedMapping((*self).into())
50    }
51
52    fn flatten(&self) -> Vec<AssignedValue<F>> {
53        let mut flattened = vec![
54            self.block_number,
55            self.addr,
56            self.mapping_slot.hi(),
57            self.mapping_slot.lo(),
58            self.mapping_depth,
59        ];
60        let hi_lo_keys = self
61            .keys
62            .iter()
63            .flat_map(|k| [k.hi(), k.lo()])
64            .collect::<Vec<_>>();
65        flattened.extend(hi_lo_keys);
66        flattened
67    }
68}
axiom_circuit/subquery/mapping.rs

axiom_circuit/subquery/
mapping.rs
