1use ark_ff::ToConstraintField;
3use ark_r1cs_std::{prelude::*, uint64::UInt64};
4use ark_relations::r1cs::{ConstraintSystemRef, SynthesisError};
5
6use decaf377::{r1cs::FqVar, Fq};
7
8use crate::{internal::hash::DOMAIN_SEPARATOR, Position, Proof, StateCommitment};
9
10impl ToConstraintField<Fq> for Position {
11 fn to_field_elements(&self) -> Option<Vec<Fq>> {
12 let mut field_elements = Vec::<Fq>::new();
17 let value: u64 = u64::from(*self);
18 for i in 0..64 {
19 let bit = ((value >> i) & 1) != 0;
20 field_elements
21 .push(bool::to_field_elements(&bit).expect("can convert bit to field element")[0]);
22 }
23 Some(field_elements)
24 }
25}
26
27#[derive(Clone, Debug)]
28pub struct PositionVar {
30 pub position: FqVar,
32 pub bits: [Boolean<Fq>; 48],
34}
35
36impl AllocVar<Position, Fq> for PositionVar {
37 fn new_variable<T: std::borrow::Borrow<Position>>(
38 cs: impl Into<ark_relations::r1cs::Namespace<Fq>>,
39 f: impl FnOnce() -> Result<T, SynthesisError>,
40 mode: ark_r1cs_std::prelude::AllocationMode,
41 ) -> Result<Self, SynthesisError> {
42 let ns = cs.into();
43 let cs = ns.cs();
44 let inner: Position = *f()?.borrow();
45
46 let position = UInt64::new_variable(cs, || Ok(u64::from(inner)), mode)?;
47 let bits = position.to_bits_le();
48 for bit in &bits[48..] {
49 bit.enforce_equal(&Boolean::Constant(false))?;
50 }
51 let inner = Boolean::<Fq>::le_bits_to_fp_var(&bits[0..48])?;
52
53 Ok(Self {
54 bits: bits[0..48]
55 .to_vec()
56 .try_into()
57 .expect("should be able to fit in 48 bits"),
58 position: inner,
59 })
60 }
61}
62
63impl ToBitsGadget<Fq> for PositionVar {
64 fn to_bits_le(&self) -> Result<Vec<Boolean<Fq>>, SynthesisError> {
65 Ok(self.bits.to_vec())
66 }
67}
68
69impl PositionVar {
70 pub fn commitment(&self) -> Result<FqVar, SynthesisError> {
72 Boolean::<Fq>::le_bits_to_fp_var(&self.bits[0..16])
73 }
74
75 pub fn block(&self) -> Result<FqVar, SynthesisError> {
77 Boolean::<Fq>::le_bits_to_fp_var(&self.bits[16..32])
78 }
79
80 pub fn epoch(&self) -> Result<FqVar, SynthesisError> {
82 Boolean::<Fq>::le_bits_to_fp_var(&self.bits[32..48])
83 }
84}
85
86impl R1CSVar<Fq> for PositionVar {
87 type Value = Position;
88
89 fn cs(&self) -> ark_relations::r1cs::ConstraintSystemRef<Fq> {
90 self.position.cs()
91 }
92
93 fn value(&self) -> Result<Self::Value, SynthesisError> {
94 let inner_fq = self.position.value()?;
95 let inner_bytes = &inner_fq.to_bytes()[0..8];
96 let position_bytes: [u8; 8] = inner_bytes
97 .try_into()
98 .expect("should be able to fit in 16 bytes");
99 Ok(Position::from(u64::from_le_bytes(position_bytes)))
100 }
101}
102
103pub struct MerkleAuthPathVar {
105 inner: [[FqVar; 3]; 24],
106}
107
108impl AllocVar<Proof, Fq> for MerkleAuthPathVar {
109 fn new_variable<T: std::borrow::Borrow<Proof>>(
110 cs: impl Into<ark_relations::r1cs::Namespace<Fq>>,
111 f: impl FnOnce() -> Result<T, SynthesisError>,
112 mode: ark_r1cs_std::prelude::AllocationMode,
113 ) -> Result<Self, SynthesisError> {
114 let ns = cs.into();
115 let cs = ns.cs();
116 let inner1 = f()?;
117 let inner: &Proof = inner1.borrow();
118 let mut auth_path = Vec::<[FqVar; 3]>::new();
120 for depth in inner.auth_path() {
121 let mut nodes = [FqVar::zero(), FqVar::zero(), FqVar::zero()];
122 for (i, node) in depth.iter().enumerate() {
123 nodes[i] = FqVar::new_variable(cs.clone(), || Ok(Fq::from(*node)), mode)?;
124 }
125 auth_path.push(nodes);
126 }
127
128 Ok(Self {
129 inner: auth_path
130 .try_into()
131 .expect("TCT auth path should have depth 24"),
132 })
133 }
134}
135
136impl MerkleAuthPathVar {
137 pub fn hash_node(
139 cs: ConstraintSystemRef<Fq>,
140 height_var: FqVar,
141 a: FqVar,
142 b: FqVar,
143 c: FqVar,
144 d: FqVar,
145 ) -> Result<FqVar, SynthesisError> {
146 let domain_separator = FqVar::new_constant(cs.clone(), *DOMAIN_SEPARATOR)?;
147 poseidon377::r1cs::hash_4(cs, &(domain_separator + height_var), (a, b, c, d))
148 }
149
150 pub fn verify(
152 &self,
153 cs: ConstraintSystemRef<Fq>,
154 enforce: &Boolean<Fq>,
155 position_bits: &[Boolean<Fq>],
156 anchor_var: FqVar,
157 commitment_var: FqVar,
158 ) -> Result<(), SynthesisError> {
159 let domain_separator = FqVar::new_constant(cs.clone(), *DOMAIN_SEPARATOR)?;
162 let leaf_var = poseidon377::r1cs::hash_1(cs.clone(), &domain_separator, commitment_var)?;
163
164 let mut previous_level = leaf_var;
166
167 for height_value in 1..=24 {
170 let height_var = FqVar::new_constant(cs.clone(), Fq::from(height_value as u64))?;
171 let which_way_var = WhichWayVar::at(height_value, position_bits)?;
172 let siblings = &self.inner[(24 - height_value) as usize];
173 let [leftmost, left, right, rightmost] =
174 which_way_var.insert(previous_level.clone(), siblings.clone())?;
175
176 let parent = MerkleAuthPathVar::hash_node(
177 cs.clone(),
178 height_var,
179 leftmost,
180 left,
181 right,
182 rightmost,
183 )?;
184
185 previous_level = parent;
186 }
187
188 anchor_var.conditional_enforce_equal(&previous_level, enforce)?;
189
190 Ok(())
191 }
192}
193
194pub struct WhichWayVar {
198 inner: FqVar,
200}
201
202impl WhichWayVar {
203 pub fn at(height: u8, position_bits: &[Boolean<Fq>]) -> Result<WhichWayVar, SynthesisError> {
206 let shift = 2 * (height - 1);
207 let bit_1 = position_bits[shift as usize].clone();
208 let bit_2 = position_bits[(shift + 1) as usize].clone();
209
210 let inner = FqVar::from(bit_1) + FqVar::constant(Fq::from(2u64)) * FqVar::from(bit_2);
217
218 Ok(WhichWayVar { inner })
219 }
220
221 pub fn insert(&self, node: FqVar, siblings: [FqVar; 3]) -> Result<[FqVar; 4], SynthesisError> {
223 let is_leftmost = self.inner.is_eq(&FqVar::zero())?;
225 let is_left = self.inner.is_eq(&FqVar::one())?;
227 let is_right = self.inner.is_eq(&FqVar::constant(Fq::from(2u128)))?;
229 let is_rightmost = self.inner.is_eq(&FqVar::constant(Fq::from(3u128)))?;
231
232 let leftmost = FqVar::conditionally_select(&is_leftmost, &node, &siblings[0])?;
238
239 let is_left_or_leftmost_case = is_leftmost.or(&is_left)?;
245 let left_first_two_cases = FqVar::conditionally_select(&is_left, &node, &siblings[0])?;
246 let left = FqVar::conditionally_select(
247 &is_left_or_leftmost_case,
248 &left_first_two_cases,
249 &siblings[1],
250 )?;
251
252 let is_right_or_rightmost_case = is_right.or(&is_rightmost)?;
258 let right_last_two_cases = FqVar::conditionally_select(&is_right, &node, &siblings[2])?;
259 let right = FqVar::conditionally_select(
260 &is_right_or_rightmost_case,
261 &right_last_two_cases,
262 &siblings[1],
263 )?;
264
265 let rightmost = FqVar::conditionally_select(&is_rightmost, &node, &siblings[2])?;
271
272 Ok([leftmost, left, right, rightmost])
273 }
274}
275
276pub struct StateCommitmentVar {
278 pub inner: FqVar,
280}
281
282impl StateCommitmentVar {
283 pub fn inner(&self) -> FqVar {
285 self.inner.clone()
286 }
287}
288
289impl AllocVar<StateCommitment, Fq> for StateCommitmentVar {
290 fn new_variable<T: std::borrow::Borrow<StateCommitment>>(
291 cs: impl Into<ark_relations::r1cs::Namespace<Fq>>,
292 f: impl FnOnce() -> Result<T, SynthesisError>,
293 mode: ark_r1cs_std::prelude::AllocationMode,
294 ) -> Result<Self, SynthesisError> {
295 let ns = cs.into();
296 let cs = ns.cs();
297 match mode {
298 AllocationMode::Constant => unimplemented!(),
299 AllocationMode::Input => {
300 let note_commitment1 = f()?;
301 let note_commitment: StateCommitment = *note_commitment1.borrow();
302 let inner = FqVar::new_input(cs, || Ok(note_commitment.0))?;
303
304 Ok(Self { inner })
305 }
306 AllocationMode::Witness => {
307 let note_commitment1 = f()?;
308 let note_commitment: StateCommitment = *note_commitment1.borrow();
309 let inner = FqVar::new_witness(cs, || Ok(note_commitment.0))?;
310
311 Ok(Self { inner })
312 }
313 }
314 }
315}
316
317impl R1CSVar<Fq> for StateCommitmentVar {
318 type Value = StateCommitment;
319
320 fn cs(&self) -> ark_relations::r1cs::ConstraintSystemRef<Fq> {
321 self.inner.cs()
322 }
323
324 fn value(&self) -> Result<Self::Value, SynthesisError> {
325 let inner = self.inner.value()?;
326 Ok(StateCommitment(inner))
327 }
328}
329
330impl EqGadget<Fq> for StateCommitmentVar {
331 fn is_eq(&self, other: &Self) -> Result<Boolean<Fq>, SynthesisError> {
332 self.inner.is_eq(&other.inner)
333 }
334}