librustzcash/zcash_proofs/src/circuit/sprout/input.rs

use bellman::gadgets::boolean::{AllocatedBit, Boolean};
use bellman::gadgets::sha256::sha256_block_no_padding;
use bellman::{ConstraintSystem, SynthesisError};
use group::ff::PrimeField;

use super::commitment::note_comm;
use super::prfs::*;
use super::*;

pub struct InputNote {
    pub nf: Vec<Boolean>,
    pub mac: Vec<Boolean>,
}

impl InputNote {
    #[allow(clippy::too_many_arguments)]
    pub fn compute<Scalar, CS>(
        mut cs: CS,
        a_sk: Option<SpendingKey>,
        rho: Option<UniqueRandomness>,
        r: Option<CommitmentRandomness>,
        value: &NoteValue,
        h_sig: &[Boolean],
        nonce: bool,
        auth_path: [Option<([u8; 32], bool)>; TREE_DEPTH],
        rt: &[Boolean],
    ) -> Result<InputNote, SynthesisError>
    where
        Scalar: PrimeField,
        CS: ConstraintSystem<Scalar>,
    {
        let a_sk = witness_u252(
            cs.namespace(|| "a_sk"),
            a_sk.as_ref().map(|a_sk| &a_sk.0[..]),
        )?;

        let rho = witness_u256(cs.namespace(|| "rho"), rho.as_ref().map(|rho| &rho.0[..]))?;

        let r = witness_u256(cs.namespace(|| "r"), r.as_ref().map(|r| &r.0[..]))?;

        let a_pk = prf_a_pk(cs.namespace(|| "a_pk computation"), &a_sk)?;

        let nf = prf_nf(cs.namespace(|| "nf computation"), &a_sk, &rho)?;

        let mac = prf_pk(cs.namespace(|| "mac computation"), &a_sk, h_sig, nonce)?;

        let cm = note_comm(
            cs.namespace(|| "cm computation"),
            &a_pk,
            &value.bits_le(),
            &rho,
            &r,
        )?;

        // Witness into the merkle tree
        let mut cur = cm;

        for (i, layer) in auth_path.iter().enumerate() {
            let cs = &mut cs.namespace(|| format!("layer {}", i));

            let cur_is_right = AllocatedBit::alloc(
                cs.namespace(|| "cur is right"),
                layer.as_ref().map(|&(_, p)| p),
            )?;

            let lhs = cur;
            let rhs = witness_u256(
                cs.namespace(|| "sibling"),
                layer.as_ref().map(|(sibling, _)| &sibling[..]),
            )?;

            // Conditionally swap if cur is right
            let preimage = conditionally_swap_u256(
                cs.namespace(|| "conditional swap"),
                &lhs[..],
                &rhs[..],
                &cur_is_right,
            )?;

            cur = sha256_block_no_padding(cs.namespace(|| "hash of this layer"), &preimage)?;
        }

        // enforce must be true if the value is nonzero
        let enforce = AllocatedBit::alloc(
            cs.namespace(|| "enforce"),
            value.get_value().map(|n| n != 0),
        )?;

        // value * (1 - enforce) = 0
        // If `value` is zero, `enforce` _can_ be zero.
        // If `value` is nonzero, `enforce` _must_ be one.
        cs.enforce(
            || "enforce validity",
            |_| value.lc(),
            |lc| lc + CS::one() - enforce.get_variable(),
            |lc| lc,
        );

        assert_eq!(cur.len(), rt.len());

        // Check that the anchor (exposed as a public input)
        // is equal to the merkle tree root that we calculated
        // for this note
        for (i, (cur, rt)) in cur.into_iter().zip(rt.iter()).enumerate() {
            // (cur - rt) * enforce = 0
            // if enforce is zero, cur and rt can be different
            // if enforce is one, they must be equal
            cs.enforce(
                || format!("conditionally enforce correct root for bit {}", i),
                |_| cur.lc(CS::one(), Scalar::ONE) - &rt.lc(CS::one(), Scalar::ONE),
                |lc| lc + enforce.get_variable(),
                |lc| lc,
            );
        }

        Ok(InputNote { mac, nf })
    }
}

/// Swaps two 256-bit blobs conditionally, returning the
/// 512-bit concatenation.
pub fn conditionally_swap_u256<Scalar, CS>(
    mut cs: CS,
    lhs: &[Boolean],
    rhs: &[Boolean],
    condition: &AllocatedBit,
) -> Result<Vec<Boolean>, SynthesisError>
where
    Scalar: PrimeField,
    CS: ConstraintSystem<Scalar>,
{
    assert_eq!(lhs.len(), 256);
    assert_eq!(rhs.len(), 256);

    let mut new_lhs = vec![];
    let mut new_rhs = vec![];

    for (i, (lhs, rhs)) in lhs.iter().zip(rhs.iter()).enumerate() {
        let cs = &mut cs.namespace(|| format!("bit {}", i));

        let x = Boolean::from(AllocatedBit::alloc(
            cs.namespace(|| "x"),
            condition
                .get_value()
                .and_then(|v| if v { rhs.get_value() } else { lhs.get_value() }),
        )?);

        // x = (1-condition)lhs + (condition)rhs
        // x = lhs - lhs(condition) + rhs(condition)
        // x - lhs = condition (rhs - lhs)
        // if condition is zero, we don't swap, so
        //   x - lhs = 0
        //   x = lhs
        // if condition is one, we do swap, so
        //   x - lhs = rhs - lhs
        //   x = rhs
        cs.enforce(
            || "conditional swap for x",
            |lc| lc + &rhs.lc(CS::one(), Scalar::ONE) - &lhs.lc(CS::one(), Scalar::ONE),
            |lc| lc + condition.get_variable(),
            |lc| lc + &x.lc(CS::one(), Scalar::ONE) - &lhs.lc(CS::one(), Scalar::ONE),
        );

        let y = Boolean::from(AllocatedBit::alloc(
            cs.namespace(|| "y"),
            condition
                .get_value()
                .and_then(|v| if v { lhs.get_value() } else { rhs.get_value() }),
        )?);

        // y = (1-condition)rhs + (condition)lhs
        // y - rhs = condition (lhs - rhs)
        cs.enforce(
            || "conditional swap for y",
            |lc| lc + &lhs.lc(CS::one(), Scalar::ONE) - &rhs.lc(CS::one(), Scalar::ONE),
            |lc| lc + condition.get_variable(),
            |lc| lc + &y.lc(CS::one(), Scalar::ONE) - &rhs.lc(CS::one(), Scalar::ONE),
        );

        new_lhs.push(x);
        new_rhs.push(y);
    }

    let mut f = new_lhs;
    f.extend(new_rhs);

    assert_eq!(f.len(), 512);

    Ok(f)
}