Update circuit implementation

benches/circuit.rs

@@ -68,7 +68,7 @@ fn criterion_benchmark<FL: OrchardFlavorBench>(c: &mut Criterion) {
                 b.iter(|| {
                     bundle
                         .authorization()
-                        .create_proof(&pk, &instances, rng)
+                        .create_proof::<FL>(&pk, &instances, rng)
                         .unwrap()
                 });
             });

src/builder.rs

@@ -32,8 +32,8 @@ use {
     crate::{
         action::Action,
         bundle::derive_bvk,
-        circuit::{Circuit, Instance, ProvingKey, Witnesses},
-        orchard_flavor::{Flavor, OrchardFlavor, OrchardVanilla, OrchardZSA},
+        circuit::{Circuit, Instance, OrchardCircuit, ProvingKey, Witnesses},
+        orchard_flavor::OrchardFlavor,
     },
     nonempty::NonEmpty,
 };
@@ -500,10 +500,10 @@ impl ActionInfo {
     ///
     /// Panics if the asset types of the spent and output notes do not match.
     #[cfg(feature = "circuit")]
-    fn build<D: OrchardDomainCommon>(
+    fn build<FL: OrchardFlavor>(
         self,
         mut rng: impl RngCore,
-    ) -> (Action<SigningMetadata, D>, Witnesses) {
+    ) -> (Action<SigningMetadata, FL>, Witnesses) {
         assert_eq!(
             self.spend.note.asset(),
             self.output.asset,
@@ -528,7 +528,7 @@ impl ActionInfo {
                     parts: SigningParts { ak, alpha },
                 },
             ),
-            Witnesses::from_action_context_unchecked(self.spend, note, alpha, self.rcv),
+            Witnesses::from_action_context_unchecked::<FL>(self.spend, note, alpha, self.rcv),
         )
     }
 
@@ -931,10 +931,7 @@ pub fn bundle<V: TryFrom<i64>, FL: OrchardFlavor>(
                     burn_vec,
                     anchor,
                     InProgress {
-                        proof: Unproven {
-                            witnesses,
-                            circuit_flavor: FL::FLAVOR,
-                        },
+                        proof: Unproven { witnesses },
                         sigs: Unauthorized { bsk },
                     },
                 ),
@@ -1153,44 +1150,30 @@ impl<P: fmt::Debug, S: InProgressSignatures> Authorization for InProgress<P, S>
 #[derive(Clone, Debug)]
 pub struct Unproven {
     witnesses: Vec<Witnesses>,
-    circuit_flavor: Flavor,
 }
 
 #[cfg(feature = "circuit")]
 impl<S: InProgressSignatures> InProgress<Unproven, S> {
     /// Creates the proof for this bundle.
-    pub fn create_proof(
+    ///
+    /// The `OrchardCircuit` type parameter must match the circuit used when generating the witnesses
+    /// contained in this `Unproven` structure to ensure consistency and correctness of the proof.
+    pub fn create_proof<C: OrchardCircuit>(
         &self,
         pk: &ProvingKey,
         instances: &[Instance],
         rng: impl RngCore,
     ) -> Result<Proof, halo2_proofs::plonk::Error> {
-        match self.proof.circuit_flavor {
-            Flavor::OrchardVanillaFlavor => {
-                let circuits = self
-                    .proof
-                    .witnesses
-                    .iter()
-                    .map(|witnesses| Circuit::<OrchardVanilla> {
-                        witnesses: witnesses.clone(),
-                        phantom: core::marker::PhantomData,
-                    })
-                    .collect::<Vec<Circuit<OrchardVanilla>>>();
-                Proof::create(pk, &circuits, instances, rng)
-            }
-            Flavor::OrchardZSAFlavor => {
-                let circuits = self
-                    .proof
-                    .witnesses
-                    .iter()
-                    .map(|witnesses| Circuit::<OrchardZSA> {
-                        witnesses: witnesses.clone(),
-                        phantom: core::marker::PhantomData,
-                    })
-                    .collect::<Vec<Circuit<OrchardZSA>>>();
-                Proof::create(pk, &circuits, instances, rng)
-            }
-        }
+        let circuits = self
+            .proof
+            .witnesses
+            .iter()
+            .map(|witnesses| Circuit::<C> {
+                witnesses: witnesses.clone(),
+                phantom: core::marker::PhantomData,
+            })
+            .collect::<Vec<Circuit<C>>>();
+        Proof::create(pk, &circuits, instances, rng)
     }
 }
 
@@ -1211,7 +1194,7 @@ impl<S: InProgressSignatures, V, FL: OrchardFlavor> Bundle<InProgress<Unproven,
             &mut (),
             |_, _, a| Ok(a),
             |_, auth| {
-                let proof = auth.create_proof(pk, &instances, &mut rng)?;
+                let proof = auth.create_proof::<FL>(pk, &instances, &mut rng)?;
                 Ok(InProgress {
                     proof,
                     sigs: auth.sigs,

src/circuit.rs

@@ -6,12 +6,6 @@
 use alloc::vec::Vec;
 
 use group::{Curve, GroupEncoding};
-use halo2_gadgets::{
-    ecc::chip::EccConfig,
-    poseidon::Pow5Config as PoseidonConfig,
-    sinsemilla::{chip::SinsemillaConfig, merkle::chip::MerkleConfig},
-    utilities::lookup_range_check::PallasLookupRangeCheck,
-};
 use halo2_proofs::{
     circuit::{floor_planner, Layouter, Value},
     plonk::{
@@ -44,6 +38,12 @@ use crate::{
     tree::{Anchor, MerkleHashOrchard},
     value::{NoteValue, ValueCommitTrapdoor, ValueCommitment},
 };
+use halo2_gadgets::{
+    ecc::chip::EccConfig,
+    poseidon::Pow5Config as PoseidonConfig,
+    sinsemilla::{chip::SinsemillaConfig, merkle::chip::MerkleConfig},
+    utilities::lookup_range_check::PallasLookupRangeCheck,
+};
 
 mod circuit_vanilla;
 mod circuit_zsa;
@@ -110,6 +110,14 @@ pub trait OrchardCircuit: Sized + Default {
         config: Self::Config,
         layouter: impl Layouter<pallas::Base>,
     ) -> Result<(), plonk::Error>;
+
+    /// Builds the ZSA-specific witnesses for the circuit.
+    /// For OrchardVanilla circuits, it should return `Value::unknown()`.
+    fn build_additional_zsa_witnesses(
+        psi_nf: pallas::Base,
+        asset: AssetBase,
+        split_flag: bool,
+    ) -> Value<AdditionalZsaWitnesses>;
 }
 
 impl<C: OrchardCircuit> plonk::Circuit<pallas::Base> for Circuit<C> {
@@ -135,9 +143,27 @@ impl<C: OrchardCircuit> plonk::Circuit<pallas::Base> for Circuit<C> {
 
 /// The Orchard Action circuit.
 #[derive(Clone, Debug, Default)]
-pub struct Circuit<D> {
+pub struct Circuit<C: OrchardCircuit> {
     pub(crate) witnesses: Witnesses,
-    pub(crate) phantom: core::marker::PhantomData<D>,
+    pub(crate) phantom: core::marker::PhantomData<C>,
+}
+
+/// The ZSA-specific witnesses.
+#[derive(Clone, Debug)]
+pub struct AdditionalZsaWitnesses {
+    pub(crate) psi_nf: pallas::Base,
+    pub(crate) asset: AssetBase,
+    pub(crate) split_flag: bool,
+}
+
+fn unpack(
+    zsa_values: Value<AdditionalZsaWitnesses>,
+) -> (Value<pallas::Base>, Value<AssetBase>, Value<bool>) {
+    (
+        zsa_values.clone().map(|values| values.psi_nf),
+        zsa_values.clone().map(|values| values.asset),
+        zsa_values.map(|values| values.split_flag),
+    )
 }
 
 /// The Orchard Action witnesses
@@ -152,7 +178,6 @@ pub struct Witnesses {
     pub(crate) psi_old: Value<pallas::Base>,
     pub(crate) rcm_old: Value<NoteCommitTrapdoor>,
     pub(crate) cm_old: Value<NoteCommitment>,
-    pub(crate) psi_nf: Value<pallas::Base>,
     pub(crate) alpha: Value<pallas::Scalar>,
     pub(crate) ak: Value<SpendValidatingKey>,
     pub(crate) nk: Value<NullifierDerivingKey>,
@@ -163,8 +188,10 @@ pub struct Witnesses {
     pub(crate) psi_new: Value<pallas::Base>,
     pub(crate) rcm_new: Value<NoteCommitTrapdoor>,
     pub(crate) rcv: Value<ValueCommitTrapdoor>,
-    pub(crate) asset: Value<AssetBase>,
-    pub(crate) split_flag: Value<bool>,
+
+    // The ZSA-specific witnesses.
+    // For OrchardVanilla circuits, this field should be initialized to `Value::unknown()`.
+    pub(crate) additional_zsa_witnesses: Value<AdditionalZsaWitnesses>,
 }
 
 impl Witnesses {
@@ -183,17 +210,17 @@ impl Witnesses {
     ///
     /// [`SpendInfo`]: crate::builder::SpendInfo
     /// [`Builder`]: crate::builder::Builder
-    pub fn from_action_context(
+    pub fn from_action_context<C: OrchardCircuit>(
         spend: SpendInfo,
         output_note: Note,
         alpha: pallas::Scalar,
         rcv: ValueCommitTrapdoor,
     ) -> Option<Self> {
         (Rho::from_nf_old(spend.note.nullifier(&spend.fvk)) == output_note.rho())
-            .then(|| Self::from_action_context_unchecked(spend, output_note, alpha, rcv))
+            .then(|| Self::from_action_context_unchecked::<C>(spend, output_note, alpha, rcv))
     }
 
-    pub(crate) fn from_action_context_unchecked(
+    pub(crate) fn from_action_context_unchecked<C: OrchardCircuit>(
         spend: SpendInfo,
         output_note: Note,
         alpha: pallas::Scalar,
@@ -204,13 +231,15 @@ impl Witnesses {
         let psi_old = spend.note.rseed().psi(&rho_old);
         let rcm_old = spend.note.rseed().rcm(&rho_old);
 
-        let nf_rseed = spend.note.rseed_split_note().unwrap_or(*spend.note.rseed());
-        let psi_nf = nf_rseed.psi(&rho_old);
-
         let rho_new = output_note.rho();
         let psi_new = output_note.rseed().psi(&rho_new);
         let rcm_new = output_note.rseed().rcm(&rho_new);
 
+        let nf_rseed = spend.note.rseed_split_note().unwrap_or(*spend.note.rseed());
+        let psi_nf = nf_rseed.psi(&rho_old);
+        let additional_zsa_witnesses =
+            C::build_additional_zsa_witnesses(psi_nf, spend.note.asset(), spend.split_flag);
+
         Witnesses {
             path: Value::known(spend.merkle_path.auth_path()),
             pos: Value::known(spend.merkle_path.position()),
@@ -221,7 +250,6 @@ impl Witnesses {
             psi_old: Value::known(psi_old),
             rcm_old: Value::known(rcm_old),
             cm_old: Value::known(spend.note.commitment()),
-            psi_nf: Value::known(psi_nf),
             alpha: Value::known(alpha),
             ak: Value::known(spend.fvk.clone().into()),
             nk: Value::known(*spend.fvk.nk()),
@@ -232,8 +260,8 @@ impl Witnesses {
             psi_new: Value::known(psi_new),
             rcm_new: Value::known(rcm_new),
             rcv: Value::known(rcv),
-            asset: Value::known(spend.note.asset()),
-            split_flag: Value::known(spend.split_flag),
+
+            additional_zsa_witnesses,
         }
     }
 }

src/circuit/circuit_vanilla.rs

@@ -30,15 +30,16 @@ use crate::{
     circuit::value_commit_orchard::gadgets::value_commit_orchard,
     circuit::{Config, Witnesses},
     constants::{OrchardFixedBases, OrchardFixedBasesFull, OrchardHashDomains},
+    note::AssetBase,
     orchard_flavor::OrchardVanilla,
 };
 
 use super::{
     commit_ivk::CommitIvkChip,
     gadget::{add_chip::AddChip, assign_free_advice},
     note_commit::NoteCommitChip,
-    OrchardCircuit, ANCHOR, CMX, CV_NET_X, CV_NET_Y, ENABLE_OUTPUT, ENABLE_SPEND, NF_OLD, RK_X,
-    RK_Y,
+    AdditionalZsaWitnesses, OrchardCircuit, ANCHOR, CMX, CV_NET_X, CV_NET_Y, ENABLE_OUTPUT,
+    ENABLE_SPEND, NF_OLD, RK_X, RK_Y,
 };
 
 impl OrchardCircuit for OrchardVanilla {
@@ -379,7 +380,7 @@ impl OrchardCircuit for OrchardVanilla {
         // Nullifier integrity (https://p.z.cash/ZKS:action-nullifier-integrity).
         let nf_old = {
             let nf_old = derive_nullifier(
-                &mut layouter.namespace(|| "nf_old = DeriveNullifier_nk(rho_old, psi_old, cm_old)"),
+                layouter.namespace(|| "nf_old = DeriveNullifier_nk(rho_old, psi_old, cm_old)"),
                 config.poseidon_chip(),
                 config.add_chip(),
                 ecc_chip.clone(),
@@ -609,6 +610,24 @@ impl OrchardCircuit for OrchardVanilla {
 
         Ok(())
     }
+
+    /// For OrchardVanilla circuits, `build_additional_zsa_witnesses` returns `Value::unknown()`.
+    ///
+    /// # Panics
+    /// Panics if the asset is not a native asset or if `split_flag` is true.
+    fn build_additional_zsa_witnesses(
+        _: pallas::Base,
+        asset: AssetBase,
+        split_flag: bool,
+    ) -> Value<AdditionalZsaWitnesses> {
+        if !(bool::from(asset.is_native())) {
+            panic!("asset must be native asset in OrchardVanilla circuit");
+        }
+        if split_flag {
+            panic!("split_flag must be false in OrchardVanilla circuit");
+        }
+        Value::unknown()
+    }
 }
 
 #[cfg(test)]
@@ -632,9 +651,7 @@ mod tests {
         value::{ValueCommitTrapdoor, ValueCommitment},
     };
 
-    type OrchardCircuitVanilla = Circuit<OrchardVanilla>;
-
-    fn generate_circuit_instance<R: RngCore>(mut rng: R) -> (OrchardCircuitVanilla, Instance) {
+    fn generate_circuit_instance<R: RngCore>(mut rng: R) -> (Circuit<OrchardVanilla>, Instance) {
         let (_, fvk, spent_note) = Note::dummy(&mut rng, None, AssetBase::native());
 
         let sender_address = spent_note.recipient();
@@ -656,10 +673,8 @@ mod tests {
         let path = MerklePath::dummy(&mut rng);
         let anchor = path.root(spent_note.commitment().into());
 
-        let psi_old = spent_note.rseed().psi(&spent_note.rho());
-
         (
-            OrchardCircuitVanilla {
+            Circuit {
                 witnesses: Witnesses {
                     path: Value::known(path.auth_path()),
                     pos: Value::known(path.position()),
@@ -670,8 +685,6 @@ mod tests {
                     psi_old: Value::known(spent_note.rseed().psi(&spent_note.rho())),
                     rcm_old: Value::known(spent_note.rseed().rcm(&spent_note.rho())),
                     cm_old: Value::known(spent_note.commitment()),
-                    // For non split note, psi_nf is equal to psi_old
-                    psi_nf: Value::known(psi_old),
                     alpha: Value::known(alpha),
                     ak: Value::known(ak),
                     nk: Value::known(nk),
@@ -682,8 +695,8 @@ mod tests {
                     psi_new: Value::known(output_note.rseed().psi(&output_note.rho())),
                     rcm_new: Value::known(output_note.rseed().rcm(&output_note.rho())),
                     rcv: Value::known(rcv),
-                    asset: Value::known(spent_note.asset()),
-                    split_flag: Value::known(false),
+
+                    ..Witnesses::default()
                 },
                 phantom: core::marker::PhantomData,
             },
@@ -860,7 +873,7 @@ mod tests {
             .titled("Orchard Action Circuit", ("sans-serif", 60))
             .unwrap();
 
-        let circuit = OrchardCircuitVanilla {
+        let circuit = Circuit::<OrchardVanilla> {
             witnesses: Witnesses::default(),
             phantom: core::marker::PhantomData,
         };