From 22b6d5bd704611beff941267c28906d457b9c2a9 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Sat, 7 Nov 2020 13:22:33 +0800
Subject: [PATCH] Cleanups in circuit.rs

---
 src/plonk/circuit.rs | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/plonk/circuit.rs b/src/plonk/circuit.rs
index 17ecd6a9..327c8e94 100644
--- a/src/plonk/circuit.rs
+++ b/src/plonk/circuit.rs
@@ -284,7 +284,7 @@ pub struct ConstraintSystem<F> {
     // Mapping from a witness vector rotation to the index in the point vector.
     pub(crate) rotations: BTreeMap<Rotation, PointIndex>,
 
-    // Vector of permutation arguments, where each corresponds to a set of columns
+    // Vector of permutation arguments, where each corresponds to a sequence of columns
     // that are involved in a permutation argument.
     pub(crate) permutations: Vec<Vec<Column<Advice>>>,
 }
@@ -399,7 +399,11 @@ impl<F: Field> ConstraintSystem<F> {
         index
     }
 
-    /// Query an Any column at a relative position
+    /// Query an auxiliary column at a relative position
+    pub fn query_aux(&mut self, column: Column<Aux>, at: i32) -> Expression<F> {
+        Expression::Aux(self.query_aux_index(column, at))
+    }
+
     fn query_any_index(&mut self, column: Column<Any>, at: i32) -> usize {
         let index = match column.column_type {
             Any::Advice => self.query_advice_index(Column::<Advice>::from(column), at),
@@ -410,9 +414,17 @@ impl<F: Field> ConstraintSystem<F> {
         index
     }
 
-    /// Query an auxiliary column at a relative position
-    pub fn query_aux(&mut self, column: Column<Aux>, at: i32) -> Expression<F> {
-        Expression::Aux(self.query_aux_index(column, at))
+    /// Query an Any column at a relative position
+    pub fn query_any(&mut self, column: Column<Any>, at: i32) -> Expression<F> {
+        match column.column_type {
+            Any::Advice => {
+                Expression::Advice(self.query_advice_index(Column::<Advice>::from(column), at))
+            }
+            Any::Fixed => {
+                Expression::Fixed(self.query_fixed_index(Column::<Fixed>::from(column), at))
+            }
+            Any::Aux => Expression::Aux(self.query_aux_index(Column::<Aux>::from(column), at)),
+        }
     }
 
     pub(crate) fn get_advice_query_index(&self, column: Column<Advice>, at: i32) -> usize {