plonk::circuit::ConstraintSystem: compress_dynamic_table_tags()

Co-authored-by: Avi Dessauer <avi.dessauer@platonic.systems>
This commit is contained in:
therealyingtong 2023-02-23 10:35:35 +08:00
parent 1c9424d166
commit cc26963337
3 changed files with 327 additions and 19 deletions

View File

@ -0,0 +1,7 @@
# Seeds for failure cases proptest has generated in the past. It is
# automatically read and these particular cases re-run before any
# novel cases are generated.
#
# It is recommended to check this file in to source control so that
# everyone who runs the test benefits from these saved cases.
cc 88c88a92a8c1799fc2b2dac1a5b9b434fc9ac5e1218e5104a85c345124d8a937 # shrinks to tables = [TableDescription { index: 0, activations: [false] }, TableDescription { index: 1, activations: [true] }]

View File

@ -1,8 +1,6 @@
use core::cmp::max;
use core::ops::{Add, Mul};
use ff::Field;
#[cfg(feature = "unstable-dynamic-lookups")]
use ff::PrimeField;
use std::{
convert::TryFrom,
ops::{Neg, Sub},
@ -15,6 +13,8 @@ use crate::{
};
mod compress_selectors;
#[cfg(feature = "unstable-dynamic-lookups")]
mod compress_table_tags;
/// A column type
pub trait ColumnType:
@ -757,7 +757,7 @@ impl<F: std::fmt::Debug> std::fmt::Debug for Expression<F> {
Expression::Constant(scalar) => f.debug_tuple("Constant").field(scalar).finish(),
Expression::Selector(selector) => f.debug_tuple("Selector").field(selector).finish(),
#[cfg(feature = "unstable-dynamic-lookups")]
Expression::TableTag(tag) => f.debug_tuple("TableTag").field(tag).finish(),
Expression::TableTag(vc) => f.debug_tuple("TableTag").field(vc).finish(),
// Skip enum variant and print query struct directly to maintain backwards compatibility.
Expression::Fixed(FixedQuery {
index,
@ -1188,28 +1188,23 @@ impl<F: Field> ConstraintSystem<F> {
) -> (Selector, Vec<(Expression<F>, Column<Any>)>),
) -> usize
where
F: PrimeField,
F: From<u64>,
{
let mut cells = VirtualCells::new(self);
let (selector, table_map) = table_map(&mut cells);
let selector = cells.query_selector(selector);
let non_table_columns: Vec<_> = table_map
.iter()
.map(|(_, c)| c)
.filter(|col| !table.columns.contains(col))
.collect();
if !non_table_columns.is_empty() {
panic!(
"{:?} does not contain {:?}. Try adding these columns to the dynamic table.",
table,
non_table_columns.as_slice()
);
}
let mut table_map: Vec<_> = table_map
.into_iter()
.map(|(input, table)| {
.map(|(input, table_col)| {
if !table.columns.contains(&table_col) {
panic!(
"{:?} does not contain {:?}. Try adding this column to the dynamic table.",
table,
table_col
);
}
if selector.contains_simple_selector() {
panic!("selector expression containing simple selector supplied to lookup argument");
}
@ -1219,7 +1214,7 @@ impl<F: Field> ConstraintSystem<F> {
panic!("input expression containing simple selector supplied to lookup argument");
}
let table_query = cells.query_any(table, Rotation::cur());
let table_query = cells.query_any(table_col, Rotation::cur());
(selector.clone() * input, selector.clone() * table_query)
})
.collect();
@ -1380,6 +1375,80 @@ impl<F: Field> ConstraintSystem<F> {
});
}
#[cfg(feature = "unstable-dynamic-lookups")]
pub(crate) fn compress_dynamic_table_tags(
mut self,
dynamic_tables: Vec<Vec<bool>>,
) -> (Self, Vec<Vec<F>>)
where
F: From<u64>,
{
assert!(self.dynamic_table_tag_map.is_empty());
assert_eq!(self.dynamic_tables.len(), dynamic_tables.len());
let mut new_columns = vec![];
let (polys, table_tag_assignment) = compress_table_tags::process(
dynamic_tables
.into_iter()
.enumerate()
.map(
|(index, activations)| compress_table_tags::TableDescription {
index,
activations,
},
)
.collect(),
|| {
let column = self.fixed_column();
new_columns.push(column);
Expression::Fixed(FixedQuery {
index: self.query_fixed_index(column),
column_index: column.index,
rotation: Rotation::cur(),
})
},
);
let mut table_tag_map = vec![None; table_tag_assignment.len()];
let mut table_tag_replacements = vec![None; table_tag_assignment.len()];
for assignment in table_tag_assignment {
table_tag_replacements[assignment.index] = Some(assignment.expression);
table_tag_map[assignment.index] = Some(new_columns[assignment.combination_index]);
}
self.dynamic_table_tag_map = table_tag_map
.into_iter()
.map(|a| a.unwrap())
.collect::<Vec<_>>();
let table_tag_replacements = table_tag_replacements
.into_iter()
.map(|a| a.unwrap())
.collect::<Vec<_>>();
// Substitute virtual tag columns for the real fixed columns in all lookup expressions
for expr in self.lookups.iter_mut().flat_map(|lookup| {
lookup
.input_expressions
.iter_mut()
.chain(lookup.table_expressions.iter_mut())
}) {
*expr = expr.evaluate(
&|constant| Expression::Constant(constant),
&|selector| Expression::Selector(selector),
&|table| table_tag_replacements[table.0].clone(),
&|query| Expression::Fixed(query),
&|query| Expression::Advice(query),
&|query| Expression::Instance(query),
&|a| -a,
&|a, b| a + b,
&|a, b| a * b,
&|a, f| a * f,
);
}
(self, polys)
}
/// This will compress selectors together depending on their provided
/// assignments. This `ConstraintSystem` will then be modified to add new
/// fixed columns (representing the actual selectors) and will return the

View File

@ -0,0 +1,232 @@
use super::{compress_selectors::exclusion_matrix, Expression};
use ff::Field;
/// This describes a table and where it is activated.
#[derive(Debug, Clone)]
pub struct TableDescription {
/// The table that this description references, by index.
pub index: usize,
/// The vector of booleans defining which rows are active for this table.
pub activations: Vec<bool>,
}
/// This describes the assigned combination of a particular table as well as
/// the expression it should be substituted with.
#[derive(Debug, Clone)]
pub struct TableAssignment<F> {
/// The table that this structure references, by index.
pub index: usize,
/// The combination this table was assigned to
pub combination_index: usize,
/// The expression we wish to substitute with
pub expression: Expression<F>,
}
/// This function takes a vector that defines each table as well as a closure
/// used to allocate new fixed columns, and returns the assignment of each
/// combination as well as details about each table assignment.
///
/// This function takes
/// * `tables`, a vector of `TableDescription`s that describe each table
/// * `allocate_fixed_columns`, a closure that constructs a new fixed column and
/// queries it at Rotation::cur(), returning the expression
///
/// and returns `Vec<Vec<F>>` containing the assignment of each new fixed column
/// (which each correspond to a combination) as well as a vector of
/// `TableAssignment` that the caller can use to perform the necessary
/// substitutions to the constraint system.
///
/// This function is completely deterministic.
pub fn process<F: Field + From<u64>, E>(
tables: Vec<TableDescription>,
mut allocate_fixed_column: E,
) -> (Vec<Vec<F>>, Vec<TableAssignment<F>>)
where
E: FnMut() -> Expression<F>,
{
if tables.is_empty() {
// There is nothing to optimize.
return (vec![], vec![]);
}
// The length of all provided table tags must be the same.
let n = tables[0].activations.len();
assert!(tables.iter().all(|a| a.activations.len() == n));
let mut combination_assignments = vec![];
let mut table_tag_assignments = vec![];
let exclusion_matrix = exclusion_matrix(&tables, |table| table.activations.iter().cloned());
// Virtual tag columns that we've added to combinations already.
let mut added = vec![false; tables.len()];
for (i, table) in tables.iter().enumerate() {
if added[i] {
continue;
}
added[i] = true;
let mut combination: Vec<_> = vec![table];
let mut combination_added = vec![i];
// Try to find other virtual tag columns that can join this one.
'try_columns: for (j, table) in tables.iter().enumerate().skip(i + 1) {
// Skip columns that have been added to previous combinations
if added[j] {
continue 'try_columns;
}
// Is this virtual tag column excluded from co-existing in the same
// combination with any of the other virtual tag column so far?
for i in combination_added.iter() {
if exclusion_matrix[j][*i] {
continue 'try_columns;
}
}
combination.push(table);
combination_added.push(j);
added[j] = true;
}
let mut combination_assignment = vec![F::ZERO; n];
let combination_index = combination_assignments.len();
let query = allocate_fixed_column();
table_tag_assignments.extend(combination.into_iter().map(|table| {
// Update the combination assignment
for (combination, active) in combination_assignment
.iter_mut()
.zip(table.activations.iter())
{
// This will not overwrite another table tag's activations because
// we have ensured that table tags are disjoint.
if *active {
*combination = F::from(table.index as u64 + 1);
}
}
TableAssignment {
index: table.index,
combination_index,
expression: query.clone(),
}
}));
combination_assignments.push(combination_assignment);
}
(combination_assignments, table_tag_assignments)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{plonk::FixedQuery, poly::Rotation};
use pasta_curves::Fp;
use proptest::collection::{vec, SizeRange};
use proptest::prelude::*;
prop_compose! {
fn arb_table(assignment_size: usize)
(assignment in vec(any::<bool>(), assignment_size))
-> Vec<bool> {
assignment
}
}
prop_compose! {
fn arb_table_list(assignment_size: usize, num_tables: impl Into<SizeRange>)
(list in vec(arb_table(assignment_size), num_tables))
-> Vec<TableDescription>
{
list.into_iter().enumerate().map(|(i, activations)| {
TableDescription {
index: i,
activations,
}
}).collect()
}
}
prop_compose! {
fn arb_instance(max_assignment_size: usize,
max_tables: usize)
(assignment_size in 1..max_assignment_size,
num_tables in 1..max_tables)
(list in arb_table_list(assignment_size, num_tables))
-> Vec<TableDescription>
{
list
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(10000))]
#[test]
fn test_table_combination(tables in arb_instance(10, 15)) {
let mut query = 0;
let (combination_assignments, table_assignments) =
process::<Fp, _>(tables.clone(), || {
let tmp = Expression::Fixed(FixedQuery {
index: query,
column_index: query,
rotation: Rotation::cur(),
});
query += 1;
tmp
});
{
let mut tables_seen = vec![];
assert_eq!(tables.len(), table_assignments.len());
for table in &table_assignments {
// Every table should be assigned to a combination
assert!(table.combination_index < combination_assignments.len());
assert!(!tables_seen.contains(&table.index));
tables_seen.push(table.index);
}
}
// Test that, for each table, the provided expression evaluates to
// the table tag on rows where the table's activation is on
for table in table_assignments {
assert_eq!(
tables[table.index].activations.len(),
combination_assignments[table.combination_index].len()
);
for (&activation, &assignment) in tables[table.index]
.activations
.iter()
.zip(combination_assignments[table.combination_index].iter())
{
let eval = table.expression.evaluate(
&|c| c,
&|_| panic!("should not occur in returned expressions"),
#[cfg(feature = "unstable-dynamic-lookups")]
&|_| panic!("should not occur in returned expressions"),
&|query| {
// Should be the correct combination in the expression
assert_eq!(table.combination_index, query.index);
assignment
},
&|_| panic!("should not occur in returned expressions"),
&|_| panic!("should not occur in returned expressions"),
&|a| -a,
&|a, b| a + b,
&|a, b| a * b,
&|a, f| a * f,
);
if activation {
assert_eq!(eval, Fp::from(table.index as u64 + 1));
}
}
}
}
}
}