penumbra_sdk_tct/
r1cs.rs

1//! This module defines how to verify TCT auth paths in a rank-1 constraint system.
2use 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        // The variable created in AllocVar<Position, Fq> is a UInt64, which is a
13        // Vec of 64 Boolean<Fq> constraints. To construct the corresponding
14        // public input, we need to convert the u64 into 64 bits, and then
15        // convert each bit into an individual Fq element.
16        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)]
28/// Represents the position of a leaf in the TCT represented in R1CS.
29pub struct PositionVar {
30    /// The FqVar representing the leaf.
31    pub position: FqVar,
32    /// Bits
33    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    /// Witness the commitment index by taking the last 16 bits of the position.
71    pub fn commitment(&self) -> Result<FqVar, SynthesisError> {
72        Boolean::<Fq>::le_bits_to_fp_var(&self.bits[0..16])
73    }
74
75    /// Witness the block.
76    pub fn block(&self) -> Result<FqVar, SynthesisError> {
77        Boolean::<Fq>::le_bits_to_fp_var(&self.bits[16..32])
78    }
79
80    /// Witness the epoch by taking the first 16 bits of the position.
81    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
103/// This represents the TCT's auth path in R1CS.
104pub 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        // This adds one FqVar per sibling and keeps them grouped together by height.
119        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    /// Hash a node given the children at the given height.
138    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    /// Certify an auth path given a provided anchor, position, and leaf.
151    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        // We need to compute the root using the provided auth path, position,
160        // and leaf.
161        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        // Height 0 is the leaf.
165        let mut previous_level = leaf_var;
166
167        // Start hashing from height 1, first hashing the leaf and its three siblings together,
168        // then the next level and so on, until we reach the root of the quadtree.
169        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
194/// Represents the different paths a quadtree node can take.
195///
196/// A bundle of boolean R1CS constraints representing the path.
197pub struct WhichWayVar {
198    /// This FqVar has been constructed from two bits of the position.
199    inner: FqVar,
200}
201
202impl WhichWayVar {
203    /// Given a height and an index of a leaf, determine which direction the path down to that leaf
204    /// should branch at the node at that height. Allocates a `WhichWayVar`.
205    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        // Convert last two bits back to a field element.
211        //
212        // The below is effectively ensuring that the inner FqVar is constrained to be within
213        // the range [0, 3] via the equation `inner = bit_1 + 2 * bit_2`
214        // For example, for the maximum values: bit_1 = 1, bit_2 = 1
215        // inner = 1 + 2 * 1 = 3
216        let inner = FqVar::from(bit_1) + FqVar::constant(Fq::from(2u64)) * FqVar::from(bit_2);
217
218        Ok(WhichWayVar { inner })
219    }
220
221    /// Insert the provided node into the quadtree at the provided height.
222    pub fn insert(&self, node: FqVar, siblings: [FqVar; 3]) -> Result<[FqVar; 4], SynthesisError> {
223        // The node is the leftmost (0th) child.
224        let is_leftmost = self.inner.is_eq(&FqVar::zero())?;
225        // The node is the left (1st) child.
226        let is_left = self.inner.is_eq(&FqVar::one())?;
227        // The node is the right (2nd) child.
228        let is_right = self.inner.is_eq(&FqVar::constant(Fq::from(2u128)))?;
229        // The node is the rightmost (3rd) child.
230        let is_rightmost = self.inner.is_eq(&FqVar::constant(Fq::from(3u128)))?;
231
232        // Cases:
233        // * `is_leftmost`: the leftmost should be the node
234        // * `is_left`: the leftmost should be the first sibling (`siblings[0]`)
235        // * `is_right`: the leftmost should be the first sibling (`siblings[0]`)
236        // * `is_rightmost`: the leftmost should be the first sibling (`siblings[0]`)
237        let leftmost = FqVar::conditionally_select(&is_leftmost, &node, &siblings[0])?;
238
239        // Cases:
240        // * `is_leftmost`: the left should be the first sibling (`siblings[0]`)
241        // * `is_left`: the left should be the node
242        // * `is_right`: the left should be the second sibling (`siblings[1]`)
243        // * `is_rightmost`: the left should be the second sibling (`siblings[1]`)
244        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        // Cases:
253        // * `is_leftmost`: the right should be the second sibling (`siblings[1]`)
254        // * `is_left`: the right should be the second sibling (`siblings[1]`)
255        // * `is_right`: the right should be the node
256        // * `is_rightmost`: the right should be the last sibling (`siblings[2]`)
257        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        // Cases:
266        // * `is_leftmost`: the rightmost should be the last sibling (`siblings[2]`)
267        // * `is_left`: the rightmost should be the last sibling (`siblings[2]`)
268        // * `is_right`: the rightmost should be the last sibling (`siblings[2]`)
269        // * `is_rightmost`: the rightmost should be the node
270        let rightmost = FqVar::conditionally_select(&is_rightmost, &node, &siblings[2])?;
271
272        Ok([leftmost, left, right, rightmost])
273    }
274}
275
276/// Represents a state commitment in R1CS.
277pub struct StateCommitmentVar {
278    /// The `FqVar` representing the state commitment.
279    pub inner: FqVar,
280}
281
282impl StateCommitmentVar {
283    /// Access the inner `FqVar`.
284    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}