axiom_circuit/subquery/
caller.rs

1use std::collections::BTreeMap;
2
3use anyhow::Result;
4use axiom_codec::{
5    constants::MAX_SUBQUERY_INPUTS,
6    types::{field_elements::SUBQUERY_RESULT_LEN, native::AnySubquery},
7    HiLo,
8};
9use axiom_query::axiom_eth::{
10    halo2_base::{AssignedValue, Context, ContextTag},
11    keccak::promise::{KeccakFixLenCall, KeccakVarLenCall},
12    utils::encode_h256_to_hilo,
13    Field,
14};
15use ethers::{
16    providers::{JsonRpcClient, Provider},
17    types::H256,
18};
19use itertools::Itertools;
20
21use super::{
22    keccak::{KeccakSubquery, KeccakSubqueryTypes},
23    types::Subquery,
24};
25use crate::subquery::{types::RawSubquery, utils::get_subquery_type_from_any_subquery};
26
27pub trait FetchSubquery<F: Field>: Clone {
28    fn flatten(&self) -> Vec<AssignedValue<F>>;
29    fn fetch<P: JsonRpcClient>(&self, p: &Provider<P>) -> Result<H256>;
30    fn any_subquery(&self) -> AnySubquery;
31    fn call<P: JsonRpcClient>(
32        &self,
33        ctx: &mut Context<F>,
34        caller: &mut SubqueryCaller<P, F>,
35    ) -> HiLo<AssignedValue<F>> {
36        caller.call(ctx, self.clone())
37    }
38}
39
40pub struct SubqueryCaller<P: JsonRpcClient, F: Field> {
41    pub provider: Provider<P>,
42    pub subqueries: BTreeMap<ContextTag, Vec<(AnySubquery, H256)>>,
43    pub subquery_assigned_values: BTreeMap<ContextTag, Vec<AssignedValue<F>>>,
44    pub keccak_fix_len_calls: Vec<(KeccakFixLenCall<F>, HiLo<AssignedValue<F>>)>,
45    pub keccak_var_len_calls: Vec<(KeccakVarLenCall<F>, HiLo<AssignedValue<F>>)>,
46    // if true, the fetched subquery will always be H256::zero()
47    mock_subquery_call: bool,
48}
49
50impl<P: JsonRpcClient, F: Field> SubqueryCaller<P, F> {
51    pub fn new(provider: Provider<P>, mock: bool) -> Self {
52        Self {
53            provider,
54            subqueries: BTreeMap::new(),
55            subquery_assigned_values: BTreeMap::new(),
56            keccak_fix_len_calls: Vec::new(),
57            keccak_var_len_calls: Vec::new(),
58            mock_subquery_call: mock,
59        }
60    }
61
62    pub fn clear(&mut self) {
63        self.subqueries.clear();
64        self.subquery_assigned_values.clear();
65        self.keccak_fix_len_calls.clear();
66        self.keccak_var_len_calls.clear();
67    }
68
69    pub fn data_query(&self) -> Vec<Subquery> {
70        let subqueries: Vec<Subquery> = self
71            .subqueries
72            .values()
73            .flat_map(|val| {
74                val.iter()
75                    .map(|(any_subquery, result)| Subquery {
76                        subquery_type: get_subquery_type_from_any_subquery(&any_subquery.clone()),
77                        subquery_data: RawSubquery(any_subquery.clone()),
78                        val: *result,
79                    })
80                    .collect_vec()
81            })
82            .collect_vec();
83        subqueries
84    }
85
86    pub fn instances(&self) -> Vec<AssignedValue<F>> {
87        self.subquery_assigned_values
88            .values()
89            .flatten()
90            .cloned()
91            .collect_vec()
92    }
93
94    pub fn call<T: FetchSubquery<F>>(
95        &mut self,
96        ctx: &mut Context<F>,
97        subquery: T,
98    ) -> HiLo<AssignedValue<F>> {
99        let result = if self.mock_subquery_call {
100            H256::zero()
101        } else {
102            subquery.fetch(&self.provider).unwrap()
103        };
104        let any_subquery = subquery.any_subquery();
105        let val = (any_subquery.clone(), result);
106        self.subqueries
107            .entry(ctx.tag())
108            .and_modify(|thread| thread.push(val.clone()))
109            .or_insert(vec![val]);
110        let subquery_type = get_subquery_type_from_any_subquery(&any_subquery);
111        let hilo = encode_h256_to_hilo(&result);
112        let hi = ctx.load_witness(hilo.hi());
113        let lo = ctx.load_witness(hilo.lo());
114        let subquery_type_assigned_value = ctx.load_constant(F::from(subquery_type));
115        let hi_lo_vec = vec![hi, lo];
116        let mut input = subquery.flatten();
117        input.resize_with(MAX_SUBQUERY_INPUTS, || ctx.load_constant(F::ZERO));
118        let mut flattened_subquery = vec![subquery_type_assigned_value];
119        flattened_subquery.extend(input);
120        flattened_subquery.extend(hi_lo_vec);
121        assert_eq!(flattened_subquery.len(), SUBQUERY_RESULT_LEN);
122        self.subquery_assigned_values
123            .entry(ctx.tag())
124            .and_modify(|thread| thread.extend(flattened_subquery.clone()))
125            .or_insert(flattened_subquery);
126        HiLo::from_hi_lo([hi, lo])
127    }
128
129    pub fn keccak<T: KeccakSubquery<F>>(
130        &mut self,
131        ctx: &mut Context<F>,
132        subquery: T,
133    ) -> HiLo<AssignedValue<F>> {
134        let logic_input = subquery.to_logical_input();
135        let output_fe = logic_input.compute_output();
136        let hi = ctx.load_witness(output_fe.hash.hi());
137        let lo = ctx.load_witness(output_fe.hash.lo());
138        let hilo = HiLo::from_hi_lo([hi, lo]);
139        match subquery.subquery_type() {
140            KeccakSubqueryTypes::FixLen(call) => {
141                self.keccak_fix_len_calls.push((call, hilo));
142            }
143            KeccakSubqueryTypes::VarLen(call) => {
144                self.keccak_var_len_calls.push((call, hilo));
145            }
146        };
147        hilo
148    }
149}
axiom_circuit/subquery/caller.rs

axiom_circuit/subquery/
caller.rs
