1use plonky2::field::field_types::Field;
  2use plonky2::hash::hash_types::{HashOutTarget, MerkleCapTarget};
  3use plonky2::hash::merkle_proofs::MerkleProofTarget;
  4use plonky2::hash::poseidon::PoseidonHash;
  5use plonky2::iop::target::Target;
  6use plonky2::iop::witness::{PartialWitness, Witness};
  7use plonky2::plonk::circuit_builder::CircuitBuilder;
  8
  9use crate::access_set::AccessSet;
 10use crate::signal::{Digest, F};
 11
 12pub struct SemaphoreTargets {
 13    merkle_root: HashOutTarget,
 14    topic: [Target; 4],
 15    merkle_proof: MerkleProofTarget,
 16    private_key: [Target; 4],
 17    public_key_index: Target,
 18}
 19
 20impl AccessSet {
 21    pub fn tree_height(&self) -> usize {
 22        self.0.leaves.len().trailing_zeros() as usize
 23    }
 24
 25    pub fn semaphore_circuit(&self, builder: &mut CircuitBuilder<F, 2>) -> SemaphoreTargets {
 26        // Register public inputs.
 27        let merkle_root = builder.add_virtual_hash();
 28        builder.register_public_inputs(&merkle_root.elements);
 29        let nullifier = builder.add_virtual_hash();
 30        builder.register_public_inputs(&nullifier.elements);
 31        let topic: [Target; 4] = builder.add_virtual_targets(4).try_into().unwrap();
 32        builder.register_public_inputs(&topic);
 33
 34        // Merkle proof
 35        let merkle_proof = MerkleProofTarget {
 36            siblings: builder.add_virtual_hashes(self.tree_height()),
 37        };
 38
 39        // Verify public key Merkle proof.
 40        let private_key: [Target; 4] = builder.add_virtual_targets(4).try_into().unwrap();
 41        let public_key_index = builder.add_virtual_target();
 42        let public_key_index_bits = builder.split_le(public_key_index, self.tree_height());
 43        let zero = builder.zero();
 44        builder.verify_merkle_proof::<PoseidonHash>(
 45            [private_key, [zero; 4]].concat(),
 46            &public_key_index_bits,
 47            &MerkleCapTarget(vec![merkle_root]),
 48            &merkle_proof,
 49        );
 50
 51        // Check nullifier.
 52        let should_be_nullifier =
 53            builder.hash_n_to_hash_no_pad::<PoseidonHash>([private_key, topic].concat());
 54        for i in 0..4 {
 55            builder.connect(nullifier.elements[i], should_be_nullifier.elements[i]);
 56        }
 57
 58        SemaphoreTargets {
 59            merkle_root,
 60            topic,
 61            merkle_proof,
 62            private_key,
 63            public_key_index,
 64        }
 65    }
 66
 67    pub fn fill_semaphore_targets(
 68        &self,
 69        pw: &mut PartialWitness<F>,
 70        private_key: Digest,
 71        topic: Digest,
 72        public_key_index: usize,
 73        targets: SemaphoreTargets,
 74    ) {
 75        let SemaphoreTargets {
 76            merkle_root,
 77            topic: topic_target,
 78            merkle_proof: merkle_proof_target,
 79            private_key: private_key_target,
 80            public_key_index: public_key_index_target,
 81        } = targets;
 82
 83        pw.set_hash_target(merkle_root, self.0.cap.0[0]);
 84        pw.set_targets(&private_key_target, &private_key);
 85        pw.set_targets(&topic_target, &topic);
 86        pw.set_target(
 87            public_key_index_target,
 88            F::from_canonical_usize(public_key_index),
 89        );
 90
 91        let merkle_proof = self.0.prove(public_key_index);
 92        for (ht, h) in merkle_proof_target
 93            .siblings
 94            .into_iter()
 95            .zip(merkle_proof.siblings)
 96        {
 97            pw.set_hash_target(ht, h);
 98        }
 99    }
100}