From c43c91b7965920f3f84c13a7c843b8e8953bf172 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Mon, 21 Jun 2021 23:16:56 +0800
Subject: [PATCH] gadget::utilities: Adjustments to utilities gadgets and
 helpers.

---
 src/circuit/gadget/utilities.rs           | 12 ++++-----
 src/circuit/gadget/utilities/cond_swap.rs | 30 ++++++++++++++++-------
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/src/circuit/gadget/utilities.rs b/src/circuit/gadget/utilities.rs
index d7d05a43..ce01eb22 100644
--- a/src/circuit/gadget/utilities.rs
+++ b/src/circuit/gadget/utilities.rs
@@ -4,19 +4,19 @@ use halo2::{
 };
 use pasta_curves::arithmetic::FieldExt;
 
-mod cond_swap;
-mod enable_flag;
-mod lookup_range_check;
-mod plonk;
+pub(crate) mod cond_swap;
+pub(crate) mod enable_flag;
+pub(crate) mod lookup_range_check;
+pub(crate) mod plonk;
 
-/// A variable representing a number.
+/// A variable representing a field element.
 #[derive(Copy, Clone, Debug)]
 pub struct CellValue<F: FieldExt> {
     cell: Cell,
     value: Option<F>,
 }
 
-pub trait Var<F: FieldExt> {
+pub trait Var<F: FieldExt>: Copy + Clone + std::fmt::Debug {
     fn new(cell: Cell, value: Option<F>) -> Self;
     fn cell(&self) -> Cell;
     fn value(&self) -> Option<F>;
diff --git a/src/circuit/gadget/utilities/cond_swap.rs b/src/circuit/gadget/utilities/cond_swap.rs
index 8c3e2f90..c8be9e2d 100644
--- a/src/circuit/gadget/utilities/cond_swap.rs
+++ b/src/circuit/gadget/utilities/cond_swap.rs
@@ -1,4 +1,4 @@
-use super::{copy, CellValue, UtilitiesInstructions};
+use super::{copy, CellValue, UtilitiesInstructions, Var};
 use halo2::{
     circuit::{Chip, Layouter},
     plonk::{Advice, Column, ConstraintSystem, Error, Expression, Permutation, Selector},
@@ -11,10 +11,14 @@ pub trait CondSwapInstructions<F: FieldExt>: UtilitiesInstructions<F> {
     #[allow(clippy::type_complexity)]
     /// Given an input pair (a,b) and a `swap` boolean flag, returns
     /// (b,a) if `swap` is set, else (a,b) if `swap` is not set.
+    ///
+    ///
+    /// The second element of the pair is required to be a witnessed
+    /// value, not a variable that already exists in the circuit.
     fn swap(
         &self,
         layouter: impl Layouter<F>,
-        pair: (Self::Var, Self::Var),
+        pair: (Self::Var, Option<F>),
         swap: Option<bool>,
     ) -> Result<(Self::Var, Self::Var), Error>;
 }
@@ -59,7 +63,7 @@ impl<F: FieldExt> CondSwapInstructions<F> for CondSwapChip<F> {
     fn swap(
         &self,
         mut layouter: impl Layouter<F>,
-        pair: (Self::Var, Self::Var),
+        pair: (Self::Var, Option<F>),
         swap: Option<bool>,
     ) -> Result<(Self::Var, Self::Var), Error> {
         let config = self.config();
@@ -73,8 +77,16 @@ impl<F: FieldExt> CondSwapInstructions<F> for CondSwapChip<F> {
                 // Copy in `a` value
                 let a = copy(&mut region, || "copy a", config.a, 0, &pair.0, &config.perm)?;
 
-                // Copy in `b` value
-                let b = copy(&mut region, || "copy b", config.b, 0, &pair.1, &config.perm)?;
+                // Witness `b` value
+                let b = {
+                    let cell = region.assign_advice(
+                        || "witness b",
+                        config.b,
+                        0,
+                        || pair.1.ok_or(Error::SynthesisError),
+                    )?;
+                    CellValue::new(cell, pair.1)
+                };
 
                 // Witness `swap` value
                 let swap_val = swap.map(|swap| F::from_u64(swap as u64));
@@ -245,19 +257,19 @@ mod tests {
 
                 // Load the pair and the swap flag into the circuit.
                 let a = chip.load_private(layouter.namespace(|| "a"), config.a, self.a)?;
-                let b = chip.load_private(layouter.namespace(|| "b"), config.b, self.b)?;
                 // Return the swapped pair.
-                let swapped_pair = chip.swap(layouter.namespace(|| "swap"), (a, b), self.swap)?;
+                let swapped_pair =
+                    chip.swap(layouter.namespace(|| "swap"), (a, self.b), self.swap)?;
 
                 if let Some(swap) = self.swap {
                     if swap {
                         // Check that `a` and `b` have been swapped
-                        assert_eq!(swapped_pair.0.value.unwrap(), b.value.unwrap());
+                        assert_eq!(swapped_pair.0.value.unwrap(), self.b.unwrap());
                         assert_eq!(swapped_pair.1.value.unwrap(), a.value.unwrap());
                     } else {
                         // Check that `a` and `b` have not been swapped
                         assert_eq!(swapped_pair.0.value.unwrap(), a.value.unwrap());
-                        assert_eq!(swapped_pair.1.value.unwrap(), b.value.unwrap());
+                        assert_eq!(swapped_pair.1.value.unwrap(), self.b.unwrap());
                     }
                 }