1use std::{
2 io::{Error, ErrorKind, Result, Write},
3 iter,
4};
5
6use byteorder::{BigEndian, WriteBytesExt};
7use ethers_core::{types::H256, utils::keccak256};
8
9use crate::{
10 constants::USER_PROOF_LEN_BYTES,
11 types::native::{
12 AccountSubquery, AnySubquery, AxiomV2ComputeQuery, AxiomV2ComputeSnark, AxiomV2DataQuery,
13 HeaderSubquery, ReceiptSubquery, SolidityNestedMappingSubquery, StorageSubquery, Subquery,
14 SubqueryResult, SubqueryType, TxSubquery,
15 },
16 utils::writer::{write_curve_compressed, write_u256},
17 VERSION,
18};
19
20pub fn get_query_hash_v2(
21 source_chain_id: u64,
22 data_query: &AxiomV2DataQuery,
23 compute_query: &AxiomV2ComputeQuery,
24) -> Result<H256> {
25 if source_chain_id != data_query.source_chain_id {
26 return Err(Error::new(ErrorKind::InvalidInput, "source_chain_id mismatch"));
27 }
28 let data_query_hash = data_query.keccak();
29 let encoded_compute_query = compute_query.encode()?;
30
31 let mut encoded = vec![];
32 encoded.write_u8(VERSION)?;
33 encoded.write_u64::<BigEndian>(source_chain_id)?;
34 encoded.write_all(data_query_hash.as_bytes())?;
35 encoded.write_all(&encoded_compute_query)?;
36 Ok(H256(keccak256(encoded)))
37}
38
39impl AxiomV2ComputeQuery {
40 pub fn encode(&self) -> Result<Vec<u8>> {
41 if self.k == 0 {
42 return Ok([&[0u8], &self.result_len.to_be_bytes()[..]].concat());
43 }
44 let encoded_query_schema = encode_query_schema(self.k, self.result_len, &self.vkey)?;
45 let proof_len = self.compute_proof.len() as u32;
46 let encoded_proof_len = proof_len.to_be_bytes();
47 assert_eq!(encoded_proof_len.len(), USER_PROOF_LEN_BYTES);
48 let encoded_compute_proof = [&encoded_proof_len[..], &self.compute_proof].concat();
49 Ok([encoded_query_schema, encoded_compute_proof].concat())
50 }
51
52 pub fn keccak(&self) -> Result<H256> {
53 Ok(H256(keccak256(self.encode()?)))
54 }
55}
56
57impl AxiomV2ComputeSnark {
58 pub fn encode_instances(&self) -> Result<Vec<u8>> {
60 let mut encoded = vec![];
61 if let Some((lhs, rhs)) = self.kzg_accumulator {
62 write_curve_compressed(&mut encoded, lhs)?;
63 write_curve_compressed(&mut encoded, rhs)?;
64 } else {
65 encoded = vec![0u8; 64];
66 }
67 for &output in &self.compute_results {
68 encoded.write_all(&output[..])?;
69 }
70 Ok(encoded)
71 }
72
73 pub fn encode(&self) -> Result<Vec<u8>> {
74 let mut encoded = self.encode_instances()?;
75 encoded.write_all(&self.proof_transcript)?;
76 Ok(encoded)
77 }
78}
79
80impl AxiomV2DataQuery {
81 pub fn keccak(&self) -> H256 {
82 get_data_query_hash(self.source_chain_id, &self.subqueries)
83 }
84}
85
86pub fn get_data_query_hash(source_chain_id: u64, subqueries: &[Subquery]) -> H256 {
87 let subquery_hashes = subqueries.iter().flat_map(|subquery| subquery.keccak().0);
88 let encoded: Vec<_> =
89 iter::empty().chain(source_chain_id.to_be_bytes()).chain(subquery_hashes).collect();
90 H256(keccak256(encoded))
91}
92
93pub fn get_query_schema_hash(k: u8, result_len: u16, vkey: &[H256]) -> Result<H256> {
94 if k == 0 {
95 return Ok(H256::zero());
96 }
97 let encoded_query_schema = encode_query_schema(k, result_len, vkey)?;
98 Ok(H256(keccak256(encoded_query_schema)))
99}
100
101pub fn encode_query_schema(k: u8, result_len: u16, vkey: &[H256]) -> Result<Vec<u8>> {
102 if k >= 28 {
103 return Err(Error::new(ErrorKind::InvalidInput, "k must be less than 28"));
104 }
105 if k == 0 {
106 unreachable!()
107 }
108 let mut encoded = Vec::with_capacity(3 + vkey.len() * 32);
109 encoded.write_u8(k)?;
110 encoded.write_u16::<BigEndian>(result_len)?;
111 let vkey_len: u8 = vkey
112 .len()
113 .try_into()
114 .map_err(|_| Error::new(ErrorKind::InvalidInput, "vkey len exceeds u8"))?;
115 encoded.write_u8(vkey_len)?;
116 for fe in vkey {
117 encoded.write_all(fe.as_bytes())?;
118 }
119 Ok(encoded)
120}
121
122impl Subquery {
123 pub fn encode(&self) -> Vec<u8> {
124 let sub_type = (self.subquery_type as u16).to_be_bytes();
125 let subquery_data = self.encoded_subquery_data.as_ref();
126 [&sub_type[..], subquery_data].concat()
127 }
128 pub fn keccak(&self) -> H256 {
129 H256(keccak256(self.encode()))
130 }
131}
132
133impl SubqueryResult {
134 pub fn encode(&self) -> Vec<u8> {
135 let subquery = self.subquery.encode();
136 let value = self.value.as_ref();
137 [&subquery[..], value].concat()
138 }
139 pub fn keccak(&self) -> H256 {
140 H256(keccak256(self.encode()))
141 }
142}
143
144pub fn encode_header_subquery(writer: &mut impl Write, subquery: HeaderSubquery) -> Result<()> {
145 let HeaderSubquery { block_number, field_idx } = subquery;
146 writer.write_u32::<BigEndian>(block_number)?;
147 writer.write_u32::<BigEndian>(field_idx)?;
148 Ok(())
149}
150
151pub fn encode_account_subquery(writer: &mut impl Write, subquery: AccountSubquery) -> Result<()> {
152 let AccountSubquery { block_number, addr, field_idx } = subquery;
153 writer.write_u32::<BigEndian>(block_number)?;
154 writer.write_all(&addr[..])?;
155 writer.write_u32::<BigEndian>(field_idx)?;
156 Ok(())
157}
158
159pub fn encode_storage_subquery(writer: &mut impl Write, subquery: StorageSubquery) -> Result<()> {
160 let StorageSubquery { block_number, addr, slot } = subquery;
161 writer.write_u32::<BigEndian>(block_number)?;
162 writer.write_all(&addr[..])?;
163 write_u256(writer, slot)?;
164 Ok(())
165}
166
167pub fn encode_tx_subquery(writer: &mut impl Write, subquery: TxSubquery) -> Result<()> {
168 let TxSubquery { block_number, tx_idx, field_or_calldata_idx } = subquery;
169 writer.write_u32::<BigEndian>(block_number)?;
170 writer.write_u16::<BigEndian>(tx_idx)?;
171 writer.write_u32::<BigEndian>(field_or_calldata_idx)?;
172 Ok(())
173}
174
175pub fn encode_receipt_subquery(writer: &mut impl Write, subquery: ReceiptSubquery) -> Result<()> {
176 let ReceiptSubquery {
177 block_number,
178 tx_idx,
179 field_or_log_idx,
180 topic_or_data_or_address_idx,
181 event_schema,
182 } = subquery;
183 writer.write_u32::<BigEndian>(block_number)?;
184 writer.write_u16::<BigEndian>(tx_idx)?;
185 writer.write_u32::<BigEndian>(field_or_log_idx)?;
186 writer.write_u32::<BigEndian>(topic_or_data_or_address_idx)?;
187 writer.write_all(&event_schema[..])?;
188 Ok(())
189}
190
191pub fn encode_solidity_nested_mapping_subquery(
192 writer: &mut impl Write,
193 subquery: SolidityNestedMappingSubquery,
194) -> Result<()> {
195 let SolidityNestedMappingSubquery { block_number, addr, mapping_slot, mapping_depth, mut keys } =
196 subquery;
197 writer.write_u32::<BigEndian>(block_number)?;
198 writer.write_all(&addr[..])?;
199 write_u256(writer, mapping_slot)?;
200 writer.write_u8(mapping_depth)?;
201 keys.resize(mapping_depth as usize, H256::zero());
202 for key in keys {
203 writer.write_all(&key[..])?;
204 }
205 Ok(())
206}
207
208impl From<HeaderSubquery> for Subquery {
209 fn from(value: HeaderSubquery) -> Self {
210 let mut bytes = vec![];
211 encode_header_subquery(&mut bytes, value).unwrap();
212 Self { subquery_type: SubqueryType::Header, encoded_subquery_data: bytes.into() }
213 }
214}
215
216impl From<AccountSubquery> for Subquery {
217 fn from(value: AccountSubquery) -> Self {
218 let mut bytes = vec![];
219 encode_account_subquery(&mut bytes, value).unwrap();
220 Self { subquery_type: SubqueryType::Account, encoded_subquery_data: bytes.into() }
221 }
222}
223
224impl From<StorageSubquery> for Subquery {
225 fn from(value: StorageSubquery) -> Self {
226 let mut bytes = vec![];
227 encode_storage_subquery(&mut bytes, value).unwrap();
228 Self { subquery_type: SubqueryType::Storage, encoded_subquery_data: bytes.into() }
229 }
230}
231
232impl From<TxSubquery> for Subquery {
233 fn from(value: TxSubquery) -> Self {
234 let mut bytes = vec![];
235 encode_tx_subquery(&mut bytes, value).unwrap();
236 Self { subquery_type: SubqueryType::Transaction, encoded_subquery_data: bytes.into() }
237 }
238}
239
240impl From<ReceiptSubquery> for Subquery {
241 fn from(value: ReceiptSubquery) -> Self {
242 let mut bytes = vec![];
243 encode_receipt_subquery(&mut bytes, value).unwrap();
244 Self { subquery_type: SubqueryType::Receipt, encoded_subquery_data: bytes.into() }
245 }
246}
247
248impl From<SolidityNestedMappingSubquery> for Subquery {
249 fn from(value: SolidityNestedMappingSubquery) -> Self {
250 let mut bytes = vec![];
251 encode_solidity_nested_mapping_subquery(&mut bytes, value).unwrap();
252 Self {
253 subquery_type: SubqueryType::SolidityNestedMapping,
254 encoded_subquery_data: bytes.into(),
255 }
256 }
257}
258
259impl From<AnySubquery> for Subquery {
260 fn from(value: AnySubquery) -> Self {
261 match value {
262 AnySubquery::Null => {
263 Self { subquery_type: SubqueryType::Null, encoded_subquery_data: vec![].into() }
264 }
265 AnySubquery::Header(subquery) => subquery.into(),
266 AnySubquery::Account(subquery) => subquery.into(),
267 AnySubquery::Storage(subquery) => subquery.into(),
268 AnySubquery::Transaction(subquery) => subquery.into(),
269 AnySubquery::Receipt(subquery) => subquery.into(),
270 AnySubquery::SolidityNestedMapping(subquery) => subquery.into(),
271 }
272 }
273}
274
275impl From<HeaderSubquery> for AnySubquery {
276 fn from(value: HeaderSubquery) -> Self {
277 AnySubquery::Header(value)
278 }
279}
280impl From<AccountSubquery> for AnySubquery {
281 fn from(value: AccountSubquery) -> Self {
282 AnySubquery::Account(value)
283 }
284}
285impl From<StorageSubquery> for AnySubquery {
286 fn from(value: StorageSubquery) -> Self {
287 AnySubquery::Storage(value)
288 }
289}
290impl From<TxSubquery> for AnySubquery {
291 fn from(value: TxSubquery) -> Self {
292 AnySubquery::Transaction(value)
293 }
294}
295impl From<ReceiptSubquery> for AnySubquery {
296 fn from(value: ReceiptSubquery) -> Self {
297 AnySubquery::Receipt(value)
298 }
299}
300impl From<SolidityNestedMappingSubquery> for AnySubquery {
301 fn from(value: SolidityNestedMappingSubquery) -> Self {
302 AnySubquery::SolidityNestedMapping(value)
303 }
304}