axiom_eth/utils/component/
types.rs

1use std::{hash::Hash, marker::PhantomData};
2
3use crate::Field;
4use halo2_base::{AssignedValue, Context};
5use itertools::Itertools;
6use serde::{Deserialize, Serialize};
7
8use super::param::{POSEIDON_RATE, POSEIDON_T};
9use super::{ComponentType, ComponentTypeId, LogicalInputValue, LogicalResult};
10
11pub type PoseidonHasher<F> =
12    halo2_base::poseidon::hasher::PoseidonHasher<F, POSEIDON_T, POSEIDON_RATE>;
13
14/// Flatten represents a flatten fixed-len logical input/output/public instances.
15#[derive(Debug, Clone, Hash, PartialEq, Eq)]
16pub struct Flatten<T> {
17    pub fields: Vec<T>,
18    pub field_size: &'static [usize],
19}
20
21impl<F: Field> From<Flatten<AssignedValue<F>>> for Flatten<F> {
22    fn from(f: Flatten<AssignedValue<F>>) -> Self {
23        Flatten::<F> {
24            fields: f.fields.into_iter().map(|v| *v.value()).collect_vec(),
25            field_size: f.field_size,
26        }
27    }
28}
29
30impl<F: Field> Flatten<F> {
31    /// Assign Flatten<F>.
32    pub fn assign(&self, ctx: &mut Context<F>) -> Flatten<AssignedValue<F>> {
33        Flatten::<AssignedValue<F>> {
34            fields: ctx.assign_witnesses(self.fields.clone()),
35            field_size: self.field_size,
36        }
37    }
38}
39
40impl<T: Copy> From<Flatten<T>> for Vec<T> {
41    fn from(val: Flatten<T>) -> Self {
42        val.fields
43    }
44}
45
46/// A logical input/output should be able to convert to a flatten logical input/ouptut.
47pub trait FixLenLogical<T: Copy>:
48    TryFrom<Flatten<T>, Error = anyhow::Error> + Into<Flatten<T>> + Clone
49{
50    /// Get field size of this logical.
51    fn get_field_size() -> &'static [usize];
52    /// Get number of fields of this logical.
53    fn get_num_fields() -> usize {
54        Self::get_field_size().len()
55    }
56    /// From raw vec to logical.
57    fn try_from_raw(fields: Vec<T>) -> anyhow::Result<Self> {
58        // TODO: we should auto generate this as Into<Vec<T>>
59        let flatten = Flatten::<T> { fields, field_size: Self::get_field_size() };
60        Self::try_from(flatten)
61    }
62    /// Into raw vec.
63    fn into_raw(self) -> Vec<T> {
64        // TODO: we should auto generate this as Into<Vec<T>>
65        self.into().fields
66    }
67}
68
69#[derive(Clone, Debug)]
70pub struct ComponentPublicInstances<T: Copy> {
71    pub output_commit: T,
72    pub promise_result_commit: T,
73    pub other: Vec<T>,
74}
75
76type V<T> = Vec<T>;
77impl<T: Copy> From<ComponentPublicInstances<T>> for V<T> {
78    fn from(val: ComponentPublicInstances<T>) -> Self {
79        [vec![val.output_commit, val.promise_result_commit], val.other].concat()
80    }
81}
82
83impl<T: Copy> TryFrom<V<T>> for ComponentPublicInstances<T> {
84    type Error = anyhow::Error;
85    fn try_from(val: V<T>) -> anyhow::Result<Self> {
86        if val.len() < 2 {
87            return Err(anyhow::anyhow!("invalid length"));
88        }
89        Ok(Self { output_commit: val[0], promise_result_commit: val[1], other: val[2..].to_vec() })
90    }
91}
92
93impl<F: Field> From<ComponentPublicInstances<AssignedValue<F>>> for ComponentPublicInstances<F> {
94    fn from(f: ComponentPublicInstances<AssignedValue<F>>) -> Self {
95        Self {
96            output_commit: *f.output_commit.value(),
97            promise_result_commit: *f.promise_result_commit.value(),
98            other: f.other.into_iter().map(|v| *v.value()).collect_vec(),
99        }
100    }
101}
102
103const FIELD_SIZE_EMPTY: [usize; 0] = [];
104/// Type for empty public instance
105#[derive(Default, Clone, Hash, PartialEq, Eq, Debug, Serialize, Deserialize)]
106pub struct LogicalEmpty<T>(PhantomData<T>);
107
108impl<T: Copy> TryFrom<Flatten<T>> for LogicalEmpty<T> {
109    type Error = anyhow::Error;
110
111    fn try_from(value: Flatten<T>) -> std::result::Result<Self, Self::Error> {
112        if value.field_size != FIELD_SIZE_EMPTY {
113            return Err(anyhow::anyhow!("invalid field size"));
114        }
115        if value.field_size.len() != value.fields.len() {
116            return Err(anyhow::anyhow!("field length doesn't match"));
117        }
118        Ok(LogicalEmpty::<T>(PhantomData::<T>))
119    }
120}
121// This should be done by marco.
122impl<T: Copy> From<LogicalEmpty<T>> for Flatten<T> {
123    fn from(_val: LogicalEmpty<T>) -> Self {
124        Flatten::<T> { fields: vec![], field_size: &FIELD_SIZE_EMPTY }
125    }
126}
127// This should be done by marco.
128impl<T: Copy> FixLenLogical<T> for LogicalEmpty<T> {
129    fn get_field_size() -> &'static [usize] {
130        &FIELD_SIZE_EMPTY
131    }
132}
133
134impl<F: Field> From<LogicalEmpty<F>> for Vec<LogicalEmpty<F>> {
135    fn from(val: LogicalEmpty<F>) -> Self {
136        vec![val]
137    }
138}
139
140impl<F: Field> LogicalInputValue<F> for LogicalEmpty<F> {
141    fn get_capacity(&self) -> usize {
142        1
143    }
144}
145
146/// Empty component type.
147#[derive(Debug, Clone)]
148pub struct EmptyComponentType<F: Field>(PhantomData<F>);
149impl<F: Field> ComponentType<F> for EmptyComponentType<F> {
150    type InputValue = LogicalEmpty<F>;
151    type InputWitness = LogicalEmpty<AssignedValue<F>>;
152    type OutputValue = LogicalEmpty<F>;
153    type OutputWitness = LogicalEmpty<AssignedValue<F>>;
154    type LogicalInput = LogicalEmpty<F>;
155
156    fn get_type_id() -> ComponentTypeId {
157        "axiom-eth:EmptyComponentType".to_string()
158    }
159
160    fn logical_result_to_virtual_rows_impl(
161        _ins: &LogicalResult<F, Self>,
162    ) -> Vec<(Self::InputValue, Self::OutputValue)> {
163        unreachable!()
164    }
165
166    fn logical_input_to_virtual_rows_impl(_li: &Self::LogicalInput) -> Vec<Self::InputValue> {
167        unreachable!()
168    }
169}
170
171// ================== Macro for conversion to/from Flatten ==================
172#[macro_export]
173macro_rules! impl_flatten_conversion {
174    ($struct_name:ident, $bits_per_fe:ident) => {
175        impl<T: Copy> TryFrom<$crate::utils::component::types::Flatten<T>> for $struct_name<T> {
176            type Error = anyhow::Error;
177
178            fn try_from(
179                value: $crate::utils::component::types::Flatten<T>,
180            ) -> anyhow::Result<Self> {
181                if &value.field_size != &$bits_per_fe {
182                    anyhow::bail!("invalid field size");
183                }
184                if value.field_size.len() != value.fields.len() {
185                    anyhow::bail!("field length doesn't match");
186                }
187                let res = value.fields.try_into()?;
188                Ok(res)
189            }
190        }
191
192        impl<T: Copy> From<$struct_name<T>> for $crate::utils::component::types::Flatten<T> {
193            fn from(value: $struct_name<T>) -> Self {
194                $crate::utils::component::types::Flatten::<T> {
195                    fields: value.flatten().to_vec(),
196                    field_size: &$bits_per_fe,
197                }
198            }
199        }
200
201        impl<T: Copy> $crate::utils::component::types::FixLenLogical<T> for $struct_name<T> {
202            fn get_field_size() -> &'static [usize] {
203                &$bits_per_fe
204            }
205        }
206    };
207}
208
209#[macro_export]
210macro_rules! impl_logical_input {
211    ($struct_name:ident, $capacity:expr) => {
212        impl<F: Field<Repr = [u8; 32]>> $crate::utils::component::LogicalInputValue<F>
213            for $struct_name<F>
214        {
215            fn get_capacity(&self) -> usize {
216                $capacity
217            }
218        }
219    };
220}
221
222#[macro_export]
223macro_rules! impl_fix_len_call_witness {
224    ($call_name:ident, $fix_len_logical_name:ident, $component_type_name:ident) => {
225        #[derive(Clone, Copy, Debug)]
226        pub struct $call_name<F: Field>(pub $fix_len_logical_name<AssignedValue<F>>);
227        impl<F: Field> $crate::utils::component::PromiseCallWitness<F> for $call_name<F> {
228            fn get_component_type_id(&self) -> $crate::utils::component::ComponentTypeId {
229                $component_type_name::<F>::get_type_id()
230            }
231            fn get_capacity(&self) -> usize {
232                1
233            }
234            fn to_rlc(
235                &self,
236                (_, rlc_ctx): (&mut $crate::halo2_base::Context<F>, &mut $crate::halo2_base::Context<F>),
237                _range_chip: &$crate::halo2_base::gates::RangeChip<F>,
238                rlc_chip: &$crate::rlc::chip::RlcChip<F>,
239            ) -> AssignedValue<F> {
240                $crate::utils::component::promise_loader::flatten_witness_to_rlc(
241                    rlc_ctx,
242                    &rlc_chip,
243                    &self.0.clone().into(),
244                )
245            }
246            fn to_typeless_logical_input(
247                &self,
248            ) -> $crate::utils::component::TypelessLogicalInput {
249                let f_a: $crate::utils::component::types::Flatten<AssignedValue<F>> =
250                    self.0.clone().into();
251                let f_v: $crate::utils::component::types::Flatten<F> = f_a.into();
252                let l_v: <$component_type_name<F> as ComponentType<F>>::LogicalInput =
253                    f_v.try_into().unwrap();
254                $crate::utils::component::utils::into_key(l_v)
255            }
256            fn get_mock_output(&self) -> $crate::utils::component::types::Flatten<F> {
257                let output_val: <$component_type_name<F> as $crate::utils::component::ComponentType<F>>::OutputValue =
258                    Default::default();
259                output_val.into()
260            }
261            fn as_any(&self) -> &dyn std::any::Any {
262                self
263            }
264        }
265    };
266}