From 34c6cba5375d08bd88d19ca1a268c9596e8510b2 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Fri, 6 Nov 2020 11:36:54 +0800
Subject: [PATCH] Add generic query_any_index() and get_any_query_index methods

---
 src/plonk/circuit.rs | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/plonk/circuit.rs b/src/plonk/circuit.rs
index 793e6d7..17ecd6a 100644
--- a/src/plonk/circuit.rs
+++ b/src/plonk/circuit.rs
@@ -399,6 +399,17 @@ impl<F: Field> ConstraintSystem<F> {
         index
     }
 
+    /// Query an Any column at a relative position
+    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),
+            Any::Fixed => self.query_fixed_index(Column::<Fixed>::from(column), at),
+            Any::Aux => self.query_aux_index(Column::<Aux>::from(column), at),
+        };
+
+        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))
@@ -412,7 +423,7 @@ impl<F: Field> ConstraintSystem<F> {
             }
         }
 
-        panic!("get_advice_query_index called for non-existant query");
+        panic!("get_advice_query_index called for non-existent query");
     }
 
     pub(crate) fn get_fixed_query_index(&self, column: Column<Fixed>, at: i32) -> usize {
@@ -426,6 +437,27 @@ impl<F: Field> ConstraintSystem<F> {
         panic!("get_fixed_query_index called for non-existent query");
     }
 
+    pub(crate) fn get_aux_query_index(&self, column: Column<Aux>, at: i32) -> usize {
+        let at = Rotation(at);
+        for (index, aux_query) in self.aux_queries.iter().enumerate() {
+            if aux_query == &(column, at) {
+                return index;
+            }
+        }
+
+        panic!("get_aux_query_index called for non-existent query");
+    }
+
+    pub(crate) fn get_any_query_index(&self, column: Column<Any>, at: i32) -> usize {
+        let index = match column.column_type {
+            Any::Advice => self.get_advice_query_index(Column::<Advice>::from(column), at),
+            Any::Fixed => self.get_fixed_query_index(Column::<Fixed>::from(column), at),
+            Any::Aux => self.get_aux_query_index(Column::<Aux>::from(column), at),
+        };
+
+        index
+    }
+
     /// Create a new gate
     pub fn create_gate(&mut self, f: impl FnOnce(&mut Self) -> Expression<F>) {
         let poly = f(self);