1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use crate::Address;
use ark_ff::{PrimeField, ToConstraintField};
use ark_r1cs_std::prelude::*;
use ark_relations::r1cs::SynthesisError;
use decaf377::{
    r1cs::{ElementVar, FqVar},
    Element, FieldExt, Fq,
};

#[derive(Clone)]
pub struct AddressVar {
    pub diversified_generator: ElementVar,
    pub transmission_key: ElementVar,
    pub clue_key: FqVar,
}

impl AddressVar {
    pub fn diversified_generator(&self) -> ElementVar {
        self.diversified_generator.clone()
    }

    pub fn transmission_key(&self) -> ElementVar {
        self.transmission_key.clone()
    }

    pub fn clue_key(&self) -> FqVar {
        self.clue_key.clone()
    }
}

impl AllocVar<Address, Fq> for AddressVar {
    fn new_variable<T: std::borrow::Borrow<Address>>(
        cs: impl Into<ark_relations::r1cs::Namespace<Fq>>,
        f: impl FnOnce() -> Result<T, SynthesisError>,
        mode: ark_r1cs_std::prelude::AllocationMode,
    ) -> Result<Self, SynthesisError> {
        let ns = cs.into();
        let cs = ns.cs();
        let address: Address = f()?.borrow().to_owned();

        let diversified_generator: ElementVar = AllocVar::<Element, Fq>::new_variable(
            cs.clone(),
            || Ok(address.diversified_generator()),
            mode,
        )?;
        // Check the diversified base is not identity.
        let identity = ElementVar::new_constant(cs.clone(), decaf377::Element::default())?;
        identity.enforce_not_equal(&diversified_generator)?;

        let element_transmission_key = decaf377::Encoding(address.transmission_key().0)
            .vartime_decompress()
            .map_err(|_| SynthesisError::AssignmentMissing)?;
        let transmission_key: ElementVar = AllocVar::<Element, Fq>::new_variable(
            cs.clone(),
            || Ok(element_transmission_key),
            mode,
        )?;
        let clue_key = FqVar::new_variable(
            cs,
            || Ok(Fq::from_le_bytes_mod_order(&address.clue_key().0[..])),
            mode,
        )?;

        Ok(Self {
            diversified_generator,
            transmission_key,
            clue_key,
        })
    }
}

// We do not implement the R1CSVar trait for AddressVar since we do not have the
// diversifier in-circuit in order to construct an Address.

impl ToConstraintField<Fq> for Address {
    fn to_field_elements(&self) -> Option<Vec<Fq>> {
        let mut elements = Vec::new();
        elements.extend(self.diversified_generator().to_field_elements()?);
        let transmission_key_fq = decaf377::Encoding(self.transmission_key().0)
            .vartime_decompress()
            .expect("transmission key is valid decaf377 Element");
        elements.extend([transmission_key_fq.vartime_compress_to_field()]);
        elements.extend(Fq::from_bytes(self.clue_key().0));
        Some(elements)
    }
}