1use std::collections::HashMap;
10
11use alloy::primitives::{FixedBytes, U256};
12use ark_ed_on_bn254::EdwardsConfig as Config;
13use ark_ff::PrimeField;
14use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
15use jf_crhf::CRHF;
16use jf_rescue::{crhf::VariableLengthRescueCRHF, RescueError, RescueParameter};
17use jf_signature::schnorr;
18use jf_utils::to_bytes;
19use rand::SeedableRng;
20use rand_chacha::ChaCha20Rng;
21use serde::{Deserialize, Serialize};
22use tagged_base64::tagged;
23
24use crate::signature_key::BLSPubKey;
25
26pub const DEFAULT_STAKE_TABLE_CAPACITY: usize = 200;
28pub type CircuitField = ark_ed_on_bn254::Fq;
30pub type LightClientState = GenericLightClientState<CircuitField>;
32pub type LightClientStateMsg = GenericLightClientStateMsg<CircuitField>;
34pub type StakeTableState = GenericStakeTableState<CircuitField>;
36pub type StateSignatureScheme =
38 jf_signature::schnorr::SchnorrSignatureScheme<ark_ed_on_bn254::EdwardsConfig>;
39pub type StateSignature = schnorr::Signature<Config>;
41pub type StateVerKey = schnorr::VerKey<Config>;
43pub type StateSignKey = schnorr::SignKey<ark_ed_on_bn254::Fr>;
45#[derive(Debug, Default, Clone)]
47pub struct StateKeyPair(pub schnorr::KeyPair<Config>);
48
49#[derive(Clone, Debug, CanonicalSerialize, CanonicalDeserialize, Serialize, Deserialize)]
51pub struct LCV1StateSignatureRequestBody {
52 pub key: StateVerKey,
54 pub state: LightClientState,
56 pub signature: StateSignature,
58}
59
60impl std::fmt::Display for LCV1StateSignatureRequestBody {
61 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62 write!(
63 f,
64 "LCV1StateSignatureRequestBody {{ key: {}, state: {}, signature: {} }}",
65 self.key, self.state, self.signature
66 )
67 }
68}
69
70#[derive(Clone, Debug, CanonicalSerialize, CanonicalDeserialize, Serialize, Deserialize)]
72pub struct LCV2StateSignatureRequestBody {
73 pub key: StateVerKey,
75 pub state: LightClientState,
77 pub next_stake: StakeTableState,
79 pub signature: StateSignature,
81}
82
83impl std::fmt::Display for LCV2StateSignatureRequestBody {
84 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85 write!(
86 f,
87 "LCV2StateSignatureRequestBody {{ key: {}, state: {}, next_stake: {}, signature: {} }}",
88 self.key, self.state, self.next_stake, self.signature
89 )
90 }
91}
92
93#[derive(Clone, Debug, Serialize, Deserialize)]
95pub struct LCV3StateSignatureRequestBody {
96 pub key: StateVerKey,
98 pub state: LightClientState,
100 pub next_stake: StakeTableState,
102 pub auth_root: FixedBytes<32>,
104 pub signature: StateSignature,
106 pub v2_signature: StateSignature,
108}
109
110impl std::fmt::Display for LCV3StateSignatureRequestBody {
111 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112 write!(
113 f,
114 "LCV3StateSignatureRequestBody {{ key: {}, state: {}, next_stake: {}, auth_root: {}, \
115 signature: {}, v2_signature: {} }}",
116 self.key,
117 self.state,
118 self.next_stake,
119 self.auth_root,
120 self.signature,
121 self.v2_signature
122 )
123 }
124}
125
126impl From<LCV1StateSignatureRequestBody> for LCV2StateSignatureRequestBody {
127 fn from(value: LCV1StateSignatureRequestBody) -> Self {
128 Self {
129 key: value.key,
130 state: value.state,
131 next_stake: StakeTableState::default(),
133 signature: value.signature,
134 }
135 }
136}
137
138impl From<LCV2StateSignatureRequestBody> for LCV1StateSignatureRequestBody {
139 fn from(value: LCV2StateSignatureRequestBody) -> Self {
140 Self {
141 key: value.key,
142 state: value.state,
143 signature: value.signature,
144 }
145 }
146}
147
148impl From<LCV3StateSignatureRequestBody> for LCV2StateSignatureRequestBody {
149 fn from(value: LCV3StateSignatureRequestBody) -> Self {
150 Self {
151 key: value.key,
152 state: value.state,
153 next_stake: value.next_stake,
154 signature: value.v2_signature,
155 }
156 }
157}
158
159#[derive(Clone, Debug, Serialize, Deserialize)]
161pub struct LCV1StateSignaturesBundle {
162 pub state: LightClientState,
164 pub signatures: HashMap<StateVerKey, StateSignature>,
166 pub accumulated_weight: U256,
168}
169
170#[derive(Clone, Debug, Serialize, Deserialize)]
172pub struct LCV2StateSignaturesBundle {
173 pub state: LightClientState,
175 pub next_stake: StakeTableState,
177 pub signatures: HashMap<StateVerKey, StateSignature>,
179 pub accumulated_weight: U256,
181}
182
183impl LCV2StateSignaturesBundle {
184 pub fn from_v1(value: LCV1StateSignaturesBundle) -> Self {
186 Self {
187 state: value.state,
188 next_stake: StakeTableState::default(), signatures: value.signatures,
190 accumulated_weight: value.accumulated_weight,
191 }
192 }
193}
194
195#[derive(Clone, Debug, Serialize, Deserialize)]
197pub struct LCV3StateSignaturesBundle {
198 pub state: LightClientState,
200 pub next_stake: StakeTableState,
202 pub auth_root: FixedBytes<32>,
204 pub signatures: HashMap<StateVerKey, StateSignature>,
206 pub accumulated_weight: U256,
208}
209
210#[tagged("LIGHT_CLIENT_STATE")]
212#[derive(
213 Clone,
214 Debug,
215 CanonicalSerialize,
216 CanonicalDeserialize,
217 Default,
218 Eq,
219 PartialEq,
220 PartialOrd,
221 Ord,
222 Hash,
223 Copy,
224)]
225pub struct GenericLightClientState<F: PrimeField> {
226 pub view_number: u64,
228 pub block_height: u64,
230 pub block_comm_root: F,
232}
233
234pub type GenericLightClientStateMsg<F> = [F; 3];
235
236impl<F: PrimeField> From<GenericLightClientState<F>> for GenericLightClientStateMsg<F> {
237 fn from(state: GenericLightClientState<F>) -> Self {
238 [
239 F::from(state.view_number),
240 F::from(state.block_height),
241 state.block_comm_root,
242 ]
243 }
244}
245
246impl<F: PrimeField> From<&GenericLightClientState<F>> for GenericLightClientStateMsg<F> {
247 fn from(state: &GenericLightClientState<F>) -> Self {
248 [
249 F::from(state.view_number),
250 F::from(state.block_height),
251 state.block_comm_root,
252 ]
253 }
254}
255
256impl<F: PrimeField + RescueParameter> GenericLightClientState<F> {
257 pub fn new(
258 view_number: u64,
259 block_height: u64,
260 block_comm_root: &[u8],
261 ) -> anyhow::Result<Self> {
262 Ok(Self {
263 view_number,
264 block_height,
265 block_comm_root: hash_bytes_to_field(block_comm_root)?,
266 })
267 }
268}
269
270#[tagged("STAKE_TABLE_STATE")]
272#[derive(
273 Clone,
274 Debug,
275 CanonicalSerialize,
276 CanonicalDeserialize,
277 Default,
278 Eq,
279 PartialEq,
280 PartialOrd,
281 Ord,
282 Hash,
283 Copy,
284)]
285pub struct GenericStakeTableState<F: PrimeField> {
286 pub bls_key_comm: F,
288 pub schnorr_key_comm: F,
290 pub amount_comm: F,
292 pub threshold: F,
294}
295
296impl<F: PrimeField> From<GenericStakeTableState<F>> for [F; 4] {
297 fn from(state: GenericStakeTableState<F>) -> Self {
298 [
299 state.bls_key_comm,
300 state.schnorr_key_comm,
301 state.amount_comm,
302 state.threshold,
303 ]
304 }
305}
306
307impl std::ops::Deref for StateKeyPair {
308 type Target = schnorr::KeyPair<Config>;
309
310 fn deref(&self) -> &Self::Target {
311 &self.0
312 }
313}
314
315impl StateKeyPair {
316 #[must_use]
318 pub fn from_sign_key(sk: StateSignKey) -> Self {
319 Self(schnorr::KeyPair::<Config>::from(sk))
320 }
321
322 #[must_use]
324 pub fn generate() -> StateKeyPair {
325 schnorr::KeyPair::generate(&mut rand::thread_rng()).into()
326 }
327
328 #[must_use]
330 pub fn generate_from_seed(seed: [u8; 32]) -> StateKeyPair {
331 schnorr::KeyPair::generate(&mut ChaCha20Rng::from_seed(seed)).into()
332 }
333
334 #[must_use]
336 pub fn generate_from_seed_indexed(seed: [u8; 32], index: u64) -> StateKeyPair {
337 let mut hasher = blake3::Hasher::new();
338 hasher.update(&seed);
339 hasher.update(&index.to_le_bytes());
340 let new_seed = *hasher.finalize().as_bytes();
341 Self::generate_from_seed(new_seed)
342 }
343}
344
345impl From<schnorr::KeyPair<Config>> for StateKeyPair {
346 fn from(value: schnorr::KeyPair<Config>) -> Self {
347 StateKeyPair(value)
348 }
349}
350
351pub fn hash_bytes_to_field<F: RescueParameter>(bytes: &[u8]) -> Result<F, RescueError> {
352 let bytes_len = (<F as PrimeField>::MODULUS_BIT_SIZE.div_ceil(8) - 1) as usize;
354 let elem = bytes
355 .chunks(bytes_len)
356 .map(F::from_le_bytes_mod_order)
357 .collect::<Vec<_>>();
358 Ok(VariableLengthRescueCRHF::<_, 1>::evaluate(elem)?[0])
359}
360
361pub trait ToFieldsLightClientCompat {
364 const SIZE: usize;
365 fn to_fields(&self) -> Vec<CircuitField>;
366}
367
368impl ToFieldsLightClientCompat for StateVerKey {
369 const SIZE: usize = 2;
370 fn to_fields(&self) -> Vec<CircuitField> {
372 let p = self.to_affine();
373 vec![p.x, p.y]
374 }
375}
376
377impl ToFieldsLightClientCompat for BLSPubKey {
378 const SIZE: usize = 3;
379 fn to_fields(&self) -> Vec<CircuitField> {
381 match to_bytes!(&self.to_affine()) {
382 Ok(bytes) => {
383 vec![
384 CircuitField::from_le_bytes_mod_order(&bytes[..31]),
385 CircuitField::from_le_bytes_mod_order(&bytes[31..62]),
386 CircuitField::from_le_bytes_mod_order(&bytes[62..]),
387 ]
388 },
389 Err(_) => unreachable!(),
390 }
391 }
392}