Move mux into cond_swap

halo2_gadgets/src/sinsemilla/merkle/chip.rs

@@ -441,6 +441,18 @@ where
         let chip = CondSwapChip::<pallas::Base>::construct(config);
         chip.swap(layouter, pair, swap)
     }
+
+    fn mux(
+        &self,
+        layouter: &mut impl Layouter<pallas::Base>,
+        choice: Self::Var,
+        left: Self::Var,
+        right: Self::Var,
+    ) -> Result<Self::Var, Error> {
+        let config = self.config().cond_swap_config.clone();
+        let chip = CondSwapChip::<pallas::Base>::construct(config);
+        chip.mux(layouter, choice, left, right)
+    }
 }
 
 impl<Hash, Commit, F> SinsemillaInstructions<pallas::Affine, { sinsemilla::K }, { sinsemilla::C }>

halo2_gadgets/src/utilities.rs

@@ -12,7 +12,6 @@ use std::ops::Range;
 pub mod cond_swap;
 pub mod decompose_running_sum;
 pub mod lookup_range_check;
-pub mod mux;
 
 /// A type that has a value at either keygen or proving time.
 pub trait FieldValue<F: Field> {

halo2_gadgets/src/utilities/cond_swap.rs

@@ -2,12 +2,14 @@
 
 use super::{bool_check, ternary, UtilitiesInstructions};
 
+use crate::ecc::chip::{EccPoint, NonIdentityEccPoint};
 use group::ff::{Field, PrimeField};
 use halo2_proofs::{
     circuit::{AssignedCell, Chip, Layouter, Value},
-    plonk::{Advice, Column, ConstraintSystem, Constraints, Error, Selector},
+    plonk::{self, Advice, Column, ConstraintSystem, Constraints, Error, Selector},
     poly::Rotation,
 };
+use pasta_curves::pallas;
 use std::marker::PhantomData;
 
 /// Instructions for a conditional swap gadget.
@@ -24,6 +26,17 @@ pub trait CondSwapInstructions<F: Field>: UtilitiesInstructions<F> {
         pair: (Self::Var, Value<F>),
         swap: Value<bool>,
     ) -> Result<(Self::Var, Self::Var), Error>;
+
+    /// Given an input `(choice, left, right)`, returns `left` if choice=0 and `right` otherwise.
+    ///
+    /// `choice` must be constrained to {0, 1} separately.
+    fn mux(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        choice: Self::Var,
+        left: Self::Var,
+        right: Self::Var,
+    ) -> Result<Self::Var, Error>;
 }
 
 /// A chip implementing a conditional swap.
@@ -121,6 +134,101 @@ impl<F: PrimeField> CondSwapInstructions<F> for CondSwapChip<F> {
             },
         )
     }
+
+    fn mux(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        choice: Self::Var,
+        left: Self::Var,
+        right: Self::Var,
+    ) -> Result<Self::Var, Error> {
+        let config = self.config();
+
+        layouter.assign_region(
+            || "mux",
+            |mut region| {
+                // Enable `q_swap` selector
+                config.q_swap.enable(&mut region, 0)?;
+
+                // Copy in `a` value
+                let left = left.copy_advice(|| "copy left", &mut region, config.a, 0)?;
+
+                // Copy in `b` value
+                let right = right.copy_advice(|| "copy right", &mut region, config.b, 0)?;
+
+                // Copy `choice` value
+                let choice = choice.copy_advice(|| "copy choice", &mut region, config.swap, 0)?;
+
+                let a_swapped = left
+                    .value()
+                    .zip(right.value())
+                    .zip(choice.value())
+                    .map(|((left, right), choice)| {
+                        if *choice == F::from(1_u64) {
+                            right
+                        } else {
+                            left
+                        }
+                    })
+                    .cloned();
+                let b_swapped = left
+                    .value()
+                    .zip(right.value())
+                    .zip(choice.value())
+                    .map(|((left, right), choice)| {
+                        if *choice == F::from(1_u64) {
+                            left
+                        } else {
+                            right
+                        }
+                    })
+                    .cloned();
+
+                region.assign_advice(|| "out b_swap", self.config.b_swapped, 0, || b_swapped)?;
+                region.assign_advice(|| "out a_swap", self.config.a_swapped, 0, || a_swapped)
+            },
+        )
+    }
+}
+
+impl CondSwapChip<pallas::Base> {
+    /// Given an input `(choice, left, right)`, returns `left` if choice=0 and `right` otherwise.
+    ///
+    /// `left` and `right` are `EccPoint`
+    /// `choice` must be constrained to {0, 1} separately.
+    pub fn mux_on_points(
+        &self,
+        mut layouter: impl Layouter<pallas::Base>,
+        choice: &AssignedCell<pallas::Base, pallas::Base>,
+        left: &EccPoint,
+        right: &EccPoint,
+    ) -> Result<EccPoint, plonk::Error> {
+        let x_cell = self.mux(&mut layouter, choice.clone(), left.x(), right.x())?;
+        let y_cell = self.mux(&mut layouter, choice.clone(), left.y(), right.y())?;
+        Ok(EccPoint::from_coordinates_unchecked(
+            x_cell.into(),
+            y_cell.into(),
+        ))
+    }
+
+    /// Given an input `(choice, left, right)`, returns `left` if choice=0 and `right` otherwise.
+    ///
+    /// `left` and `right` are `NonIdentityEccPoint`
+    /// `choice` must be constrained to {0, 1} separately.
+    pub fn mux_on_non_identity_points(
+        &self,
+        mut layouter: impl Layouter<pallas::Base>,
+        choice: &AssignedCell<pallas::Base, pallas::Base>,
+        left: &NonIdentityEccPoint,
+        right: &NonIdentityEccPoint,
+    ) -> Result<NonIdentityEccPoint, plonk::Error> {
+        let x_cell = self.mux(&mut layouter, choice.clone(), left.x(), right.x())?;
+        let y_cell = self.mux(&mut layouter, choice.clone(), left.y(), right.y())?;
+        Ok(NonIdentityEccPoint::from_coordinates_unchecked(
+            x_cell.into(),
+            y_cell.into(),
+        ))
+    }
 }
 
 impl<F: PrimeField> CondSwapChip<F> {