Skip to content

Commit

Permalink
add public selector into setup phase
Browse files Browse the repository at this point in the history
  • Loading branch information
@ShuangWu121
ShuangWu121 committed Feb 5, 2025
1 parent bf2998a commit 5cbf0ce
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 29 deletions.
23 changes: 22 additions & 1 deletion backend/src/stwo/circuit_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ impl PowdrEval {
.enumerate()
.map(|(index, (_, id))| (id, index))
.collect();

let publics_by_stage = analyzed.get_publics().into_iter().fold(
vec![vec![]; analyzed.stage_count()],
|mut acc, (name, column_name, id, row, stage)| {
Expand Down Expand Up @@ -254,6 +254,27 @@ impl FrameworkEval for PowdrEval {
challenges: &challenges,
poly_stage_map: &self.poly_stage_map,
};

// build selector columns and constraints for publics, for now I am using constant columns as selectors
self.publics_by_stage.iter().flatten().enumerate().for_each(
|(index, (name, _, poly_id, _))| {
let selector = eval.get_preprocessed_column(PreprocessedColumn::Plonk(
index
+ constant_eval.len()
+ self.preprocess_col_offset
+ constant_shifted_eval.len(),
));
let stage= self.poly_stage_map[poly_id];
let witness_col= match stage {
0 => self.stage0_witness_columns[poly_id],
1 => self.stage1_witness_columns[poly_id],
_ => unreachable!(),
};
// constraining s(i) * (pub[i] - x(i)) = 0
// eval.add_constraint(selector * (public_value.into() - witness_col));
},
);

let mut evaluator =
ExpressionEvaluator::new_with_custom_expr(&data, &intermediate_definitions, |v| {
E::F::from(into_stwo_field(v))
Expand Down
84 changes: 56 additions & 28 deletions backend/src/stwo/prover.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use itertools::Itertools;
use num_traits::Zero;
use num_traits::{One, Zero};
use powdr_ast::analyzed::{AlgebraicExpression, Analyzed, DegreeRange};
use powdr_ast::parsed::visitor::AllChildren;
use powdr_backend_utils::{machine_fixed_columns, machine_witness_columns};
Expand Down Expand Up @@ -29,14 +29,15 @@ use crate::stwo::proof::{
use stwo_prover::constraint_framework::TraceLocationAllocator;

use stwo_prover::core::air::{Component, ComponentProver};
use stwo_prover::core::backend::{Backend, BackendForChannel};
use stwo_prover::core::backend::{Backend, BackendForChannel, Col, Column};
use stwo_prover::core::channel::{Channel, MerkleChannel};
use stwo_prover::core::fields::m31::BaseField;
use stwo_prover::core::fields::qm31::SecureField;
use stwo_prover::core::fri::FriConfig;
use stwo_prover::core::pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig};
use stwo_prover::core::poly::circle::{CanonicCoset, CircleDomain, CircleEvaluation};
use stwo_prover::core::poly::BitReversedOrder;
use stwo_prover::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index};
use stwo_prover::core::ColumnVec;

const FRI_LOG_BLOWUP: usize = 1;
Expand Down Expand Up @@ -151,7 +152,7 @@ where
.iter()
.filter_map(|(namespace, pil)| {
// if we have no fixed columns, we don't need to commit to anything.
if pil.constant_count() + pil.publics_count()== 0 {
if pil.constant_count() + pil.publics_count() == 0 {
None
} else {
let fixed_columns = machine_fixed_columns(&self.fixed, pil);
Expand Down Expand Up @@ -198,6 +199,32 @@ where

constant_trace.extend(constant_shifted_trace);

// get selector columns for the public inputs, as closures
let publics_selector: ColumnVec<
CircleEvaluation<B, BaseField, BitReversedOrder>,
> = pil
.get_publics()
.into_iter()
.map(|(_, _, _, row_id, _)| {
let mut col = Col::<B, BaseField>::zeros(1 << size.ilog2());
col.set(
bit_reverse_index(
coset_index_to_circle_domain_index(
row_id,
size.ilog2(),
),
size.ilog2(),
),
BaseField::one(),
);
CircleEvaluation::<B, BaseField, BitReversedOrder>::new(
*domain_map.get(&(size.ilog2() as usize)).unwrap(),
col,
)
})
.collect();
constant_trace.extend(publics_selector);

(
size as usize,
TableProvingKey {
Expand Down Expand Up @@ -378,27 +405,27 @@ where
}

let tree_span_provider = &mut TraceLocationAllocator::default();
// Each column size in machines needs its own component, the components from different machines are stored in this vector
let mut components = Vec::new();

// Build the circuit. The circuit includes constraints of all the machines in both stage 0 and stage 1
let mut constant_cols_offset_acc = 0;
self.split.iter().zip_eq(machine_log_sizes.iter()).for_each(
|((machine_name, pil), (proof_machine_name, &machine_log_size))| {
assert_eq!(machine_name, proof_machine_name);

let component = PowdrComponent::new(
tree_span_provider,
PowdrEval::new(
(*pil).clone(),
constant_cols_offset_acc,
machine_log_size,
stage0_challenges.clone(),
),
(SecureField::zero(), None),
);
let components = self
.split
.iter()
.zip_eq(machine_log_sizes.iter())
.map(
|((machine_name, pil), (proof_machine_name, &machine_log_size))| {
assert_eq!(machine_name, proof_machine_name);

components.push(component);
let component = PowdrComponent::new(
tree_span_provider,
PowdrEval::new(
(*pil).clone(),
constant_cols_offset_acc,
machine_log_size,
stage0_challenges.clone(),
),
(SecureField::zero(), None),
);

constant_cols_offset_acc +=
pil.constant_count() + get_constant_with_next_list(pil).len();
Expand Down Expand Up @@ -459,7 +486,6 @@ where
})
.collect();


let stage_count = self.analyzed.stage_count();
// get public values
let mut instance_map: BTreeMap<String, Vec<Vec<M31>>> = self
Expand All @@ -479,19 +505,21 @@ where
instance_map.get_mut(namespace).unwrap()[*stage as usize].push(*value);
});

let public_inputs = instance_map
let public_inputs = instance_map
.into_iter()
.map(|(name, values)| {
(
name,
values
.into_iter()
.map(|values| gen_stwo_circle_column::<B, BaseField>(
*domain_map
.get(&(values.len().ilog2() as usize))
.expect("Domain not found for given size"),
&values,
))
.map(|values| {
gen_stwo_circle_column::<B, BaseField>(
*domain_map
.get(&(values.len().ilog2() as usize))
.expect("Domain not found for given size"),
&values,
)
})
.collect_vec(),
)
})
Expand Down

0 comments on commit 5cbf0ce

Please sign in to comment.