1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

use std::sync::{Arc, Mutex};

use axiom_circuit::{
    axiom_codec::{constants::MAX_SOLIDITY_MAPPING_KEYS, HiLo},
    axiom_eth::halo2_base::{AssignedValue, Context},
    subquery::{caller::SubqueryCaller, types::AssignedSolidityNestedMappingSubquery},
};
use ethers::providers::Http;

use crate::Fr;

/// Solidity nested mapping subquery builder
pub struct SolidityMapping<'a> {
    pub block_number: AssignedValue<Fr>,
    pub addr: AssignedValue<Fr>,
    pub mapping_slot: HiLo<AssignedValue<Fr>>,
    ctx: &'a mut Context<Fr>,
    caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
}

pub(crate) fn get_mapping(
    ctx: &mut Context<Fr>,
    caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
    block_number: AssignedValue<Fr>,
    addr: AssignedValue<Fr>,
    mapping_slot: HiLo<AssignedValue<Fr>>,
) -> SolidityMapping {
    SolidityMapping {
        block_number,
        addr,
        mapping_slot,
        ctx,
        caller,
    }
}

impl<'a> SolidityMapping<'a> {
    /// Fetches the Solidity nested mapping subquery and returns the HiLo<AssignedValue<Fr>> result
    ///
    /// * `keys` - A vector of nested keys into the specified mapping
    pub fn nested(self, keys: Vec<HiLo<AssignedValue<Fr>>>) -> HiLo<AssignedValue<Fr>> {
        if keys.is_empty() || keys.len() > MAX_SOLIDITY_MAPPING_KEYS {
            panic!(
                "Invalid number of keys for solidity mapping: {}. Must be in range [1, 4]",
                keys.len()
            );
        }
        let mut subquery_caller = self.caller.lock().unwrap();
        let depth = self.ctx.load_constant(Fr::from(keys.len() as u64));
        let mut padded_keys = keys.clone();
        padded_keys.resize_with(MAX_SOLIDITY_MAPPING_KEYS, || {
            let zeros = self.ctx.load_constants(&[Fr::zero(), Fr::zero()]);
            HiLo::from_hi_lo([zeros[0], zeros[1]])
        });
        let mut keys = [keys[0]; MAX_SOLIDITY_MAPPING_KEYS];
        keys.copy_from_slice(&padded_keys);
        let subquery = AssignedSolidityNestedMappingSubquery {
            block_number: self.block_number,
            addr: self.addr,
            mapping_slot: self.mapping_slot,
            mapping_depth: depth,
            keys,
        };
        subquery_caller.call(self.ctx, subquery)
    }

    /// Fetches the Solidity mapping subquery and returns the HiLo<AssignedValue<Fr>> result
    ///
    /// * `key` - The key into the specified mapping
    pub fn key(self, key: HiLo<AssignedValue<Fr>>) -> HiLo<AssignedValue<Fr>> {
        self.nested(vec![key])
    }
}