axiom_query/utils/
codec.rs

1use axiom_codec::{
2    encoder::field_elements::{
3        FIELD_ENCODED_SOLIDITY_NESTED_MAPPING_DEPTH_IDX, NUM_FE_ANY,
4        NUM_FE_SOLIDITY_NESTED_MAPPING_WITHOUT_KEYS,
5    },
6    types::{field_elements::*, native::SubqueryType},
7};
8use axiom_eth::halo2_base::{
9    gates::GateInstructions,
10    utils::ScalarField,
11    AssignedValue, Context,
12    QuantumCell::{Constant, Existing},
13};
14
15pub type AssignedSubqueryKey<F> = SubqueryKey<AssignedValue<F>>;
16pub type AssignedSubqueryOutput<F> = SubqueryOutput<AssignedValue<F>>;
17
18pub type AssignedHeaderSubquery<F> = FieldHeaderSubquery<AssignedValue<F>>;
19pub type AssignedAccountSubquery<F> = FieldAccountSubquery<AssignedValue<F>>;
20pub type AssignedStorageSubquery<F> = FieldStorageSubquery<AssignedValue<F>>;
21pub type AssignedTxSubquery<F> = FieldTxSubquery<AssignedValue<F>>;
22pub type AssignedReceiptSubquery<F> = FieldReceiptSubquery<AssignedValue<F>>;
23pub type AssignedSolidityNestedMappingSubquery<F> =
24    FieldSolidityNestedMappingSubquery<AssignedValue<F>>;
25
26pub type AssignedHeaderSubqueryResult<F> = FieldHeaderSubqueryResult<AssignedValue<F>>;
27pub type AssignedAccountSubqueryResult<F> = FieldAccountSubqueryResult<AssignedValue<F>>;
28pub type AssignedStorageSubqueryResult<F> = FieldStorageSubqueryResult<AssignedValue<F>>;
29pub type AssignedTxSubqueryResult<F> = FieldTxSubqueryResult<AssignedValue<F>>;
30pub type AssignedReceiptSubqueryResult<F> = FieldReceiptSubqueryResult<AssignedValue<F>>;
31pub type AssignedSolidityNestedMappingSubqueryResult<F> =
32    FieldSolidityNestedMappingSubqueryResult<AssignedValue<F>>;
33
34pub type AssignedSubqueryResult<F> = FlattenedSubqueryResult<AssignedValue<F>>;
35
36pub fn assign_flattened_subquery_result<F: ScalarField>(
37    ctx: &mut Context<F>,
38    f: &FlattenedSubqueryResult<F>,
39) -> AssignedSubqueryResult<F> {
40    let key = f.key.0.map(|x| ctx.load_witness(x));
41    let value = f.value.0.map(|x| ctx.load_witness(x));
42    FlattenedSubqueryResult::new(SubqueryKey(key), SubqueryOutput(value))
43}
44
45/// Parses the subquery type and determines the actual number of field elements in the
46/// field element encoding, **including** the subquery type.
47///
48/// This is `1` more than the flattened length of `FieldSubquery**` when the subquery type is not SolidityNestedMapping. When the subquery type is SolidityNestedMapping, the number of field elements is `1 + NUM_FE_SOLIDITY_NESTED_MAPPING_WITHOUT_KEYS + 2 * mapping_depth`.
49pub fn get_num_fe_from_subquery_key<F: ScalarField>(
50    ctx: &mut Context<F>,
51    gate: &impl GateInstructions<F>,
52    key: &AssignedSubqueryKey<F>,
53) -> AssignedValue<F> {
54    // num_fe_by_type will include 1 field element for the type itself, while NUM_FE_ANY does not
55    let num_fe_by_type: Vec<_> = NUM_FE_ANY
56        .iter()
57        .enumerate()
58        .map(|(subtype, &num_fe)| {
59            if subtype == SubqueryType::SolidityNestedMapping as usize {
60                let num_without_keys =
61                    Constant(F::from(NUM_FE_SOLIDITY_NESTED_MAPPING_WITHOUT_KEYS as u64 + 1));
62                let mapping_depth = key.0[FIELD_ENCODED_SOLIDITY_NESTED_MAPPING_DEPTH_IDX];
63                Existing(gate.mul_add(ctx, mapping_depth, Constant(F::from(2)), num_without_keys))
64            } else {
65                Constant(F::from(num_fe as u64 + 1))
66            }
67        })
68        .collect();
69    gate.select_from_idx(ctx, num_fe_by_type, key.0[0])
70}