use alloc::vec;
use alloc::vec::Vec;
use core::fmt::Debug;
use core::marker::PhantomData;
use crate::field::extension::Extendable;
use crate::field::types::Field;
use crate::hash::hash_types::RichField;
use crate::iop::ext_target::ExtensionTarget;
use crate::iop::target::Target;
use crate::iop::wire::Wire;
use crate::iop::witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite};
use crate::plonk::circuit_data::{CommonCircuitData, ProverOnlyCircuitData};
use crate::plonk::config::GenericConfig;
pub(crate) fn generate_partial_witness<
'a,
F: RichField + Extendable<D>,
C: GenericConfig<D, F = F>,
const D: usize,
>(
inputs: PartialWitness<F>,
prover_data: &'a ProverOnlyCircuitData<F, C, D>,
common_data: &'a CommonCircuitData<F, D>,
) -> PartitionWitness<'a, F> {
let config = &common_data.config;
let generators = &prover_data.generators;
let generator_indices_by_watches = &prover_data.generator_indices_by_watches;
let mut witness = PartitionWitness::new(
config.num_wires,
common_data.degree(),
&prover_data.representative_map,
);
for (t, v) in inputs.target_values.into_iter() {
witness.set_target(t, v);
}
let mut pending_generator_indices: Vec<_> = (0..generators.len()).collect();
let mut generator_is_expired = vec![false; generators.len()];
let mut remaining_generators = generators.len();
let mut buffer = GeneratedValues::empty();
while !pending_generator_indices.is_empty() {
let mut next_pending_generator_indices = Vec::new();
for &generator_idx in &pending_generator_indices {
if generator_is_expired[generator_idx] {
continue;
}
let finished = generators[generator_idx].run(&witness, &mut buffer);
if finished {
generator_is_expired[generator_idx] = true;
remaining_generators -= 1;
}
let new_target_reps = buffer
.target_values
.drain(..)
.flat_map(|(t, v)| witness.set_target_returning_rep(t, v));
for watch in new_target_reps {
let opt_watchers = generator_indices_by_watches.get(&watch);
if let Some(watchers) = opt_watchers {
for &watching_generator_idx in watchers {
if !generator_is_expired[watching_generator_idx] {
next_pending_generator_indices.push(watching_generator_idx);
}
}
}
}
}
pending_generator_indices = next_pending_generator_indices;
}
assert_eq!(
remaining_generators, 0,
"{} generators weren't run",
remaining_generators,
);
witness
}
pub trait WitnessGenerator<F: Field>: 'static + Send + Sync + Debug {
fn watch_list(&self) -> Vec<Target>;
fn run(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) -> bool;
}
#[derive(Debug)]
pub struct GeneratedValues<F: Field> {
pub(crate) target_values: Vec<(Target, F)>,
}
impl<F: Field> From<Vec<(Target, F)>> for GeneratedValues<F> {
fn from(target_values: Vec<(Target, F)>) -> Self {
Self { target_values }
}
}
impl<F: Field> WitnessWrite<F> for GeneratedValues<F> {
fn set_target(&mut self, target: Target, value: F) {
self.target_values.push((target, value));
}
}
impl<F: Field> GeneratedValues<F> {
pub fn with_capacity(capacity: usize) -> Self {
Vec::with_capacity(capacity).into()
}
pub fn empty() -> Self {
Vec::new().into()
}
pub fn singleton_wire(wire: Wire, value: F) -> Self {
Self::singleton_target(Target::Wire(wire), value)
}
pub fn singleton_target(target: Target, value: F) -> Self {
vec![(target, value)].into()
}
pub fn singleton_extension_target<const D: usize>(
et: ExtensionTarget<D>,
value: F::Extension,
) -> Self
where
F: RichField + Extendable<D>,
{
let mut witness = Self::with_capacity(D);
witness.set_extension_target(et, value);
witness
}
}
pub trait SimpleGenerator<F: Field>: 'static + Send + Sync + Debug {
fn dependencies(&self) -> Vec<Target>;
fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>);
fn adapter(self) -> SimpleGeneratorAdapter<F, Self>
where
Self: Sized,
{
SimpleGeneratorAdapter {
inner: self,
_phantom: PhantomData,
}
}
}
#[derive(Debug)]
pub struct SimpleGeneratorAdapter<F: Field, SG: SimpleGenerator<F> + ?Sized> {
_phantom: PhantomData<F>,
inner: SG,
}
impl<F: Field, SG: SimpleGenerator<F>> WitnessGenerator<F> for SimpleGeneratorAdapter<F, SG> {
fn watch_list(&self) -> Vec<Target> {
self.inner.dependencies()
}
fn run(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) -> bool {
if witness.contains_all(&self.inner.dependencies()) {
self.inner.run_once(witness, out_buffer);
true
} else {
false
}
}
}
#[derive(Debug)]
pub(crate) struct CopyGenerator {
pub(crate) src: Target,
pub(crate) dst: Target,
}
impl<F: Field> SimpleGenerator<F> for CopyGenerator {
fn dependencies(&self) -> Vec<Target> {
vec![self.src]
}
fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
let value = witness.get_target(self.src);
out_buffer.set_target(self.dst, value);
}
}
#[derive(Debug)]
pub(crate) struct RandomValueGenerator {
pub(crate) target: Target,
}
impl<F: Field> SimpleGenerator<F> for RandomValueGenerator {
fn dependencies(&self) -> Vec<Target> {
Vec::new()
}
fn run_once(&self, _witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
let random_value = F::rand();
out_buffer.set_target(self.target, random_value);
}
}
#[derive(Debug)]
pub(crate) struct NonzeroTestGenerator {
pub(crate) to_test: Target,
pub(crate) dummy: Target,
}
impl<F: Field> SimpleGenerator<F> for NonzeroTestGenerator {
fn dependencies(&self) -> Vec<Target> {
vec![self.to_test]
}
fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
let to_test_value = witness.get_target(self.to_test);
let dummy_value = if to_test_value == F::ZERO {
F::ONE
} else {
to_test_value.inverse()
};
out_buffer.set_target(self.dummy, dummy_value);
}
}
#[derive(Debug, Clone)]
pub(crate) struct ConstantGenerator<F: Field> {
pub row: usize,
pub constant_index: usize,
pub wire_index: usize,
pub constant: F,
}
impl<F: Field> ConstantGenerator<F> {
pub fn set_constant(&mut self, c: F) {
self.constant = c;
}
}
impl<F: Field> SimpleGenerator<F> for ConstantGenerator<F> {
fn dependencies(&self) -> Vec<Target> {
vec![]
}
fn run_once(&self, _witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
out_buffer.set_target(Target::wire(self.row, self.wire_index), self.constant);
}
}