MainPod implementation (#168)

* Initial circuit op work

* Fix copy op

* Add more ops

* add mainpod boilerplate

* pass basic test of mainpod

* fix duplicate imports

* WIP

* fixes

* wip

* fix test

* wip

* clean up

* address feedback from @ax0

* oops

---------

Co-authored-by: Ahmad <root@ahmadafuni.com>
This commit is contained in:
Eduard S. 2025-04-01 11:23:45 -07:00 committed by GitHub
parent 4a94b34792
commit ce26a316a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 530 additions and 104 deletions

View file

@ -4,10 +4,9 @@ use crate::backends::plonky2::basetypes::D;
use crate::backends::plonky2::mock::mainpod::Statement;
use crate::backends::plonky2::mock::mainpod::{Operation, OperationArg};
use crate::middleware::{
NativeOperation, NativePredicate, Params, Predicate, StatementArg, ToFields, Value, F,
HASH_SIZE, VALUE_SIZE,
NativeOperation, NativePredicate, Params, Predicate, StatementArg, ToFields, Value,
EMPTY_VALUE, F, HASH_SIZE, OPERATION_ARG_F_LEN, STATEMENT_ARG_F_LEN, VALUE_SIZE,
};
use crate::middleware::{OPERATION_ARG_F_LEN, STATEMENT_ARG_F_LEN};
use anyhow::Result;
use plonky2::field::extension::Extendable;
use plonky2::field::types::{Field, PrimeField64};
@ -51,10 +50,55 @@ impl ValueTarget {
}
}
#[derive(Clone)]
pub struct StatementArgTarget {
pub elements: [Target; STATEMENT_ARG_F_LEN],
}
impl StatementArgTarget {
pub fn set_targets(
&self,
pw: &mut PartialWitness<F>,
params: &Params,
arg: &StatementArg,
) -> Result<()> {
pw.set_target_arr(&self.elements, &arg.to_fields(params))
}
fn new(first: ValueTarget, second: ValueTarget) -> Self {
let elements: Vec<_> = first
.elements
.into_iter()
.chain(second.elements.into_iter())
.collect();
StatementArgTarget {
elements: elements.try_into().expect("size STATEMENT_ARG_F_LEN"),
}
}
pub fn none(builder: &mut CircuitBuilder<F, D>) -> Self {
let empty = builder.constant_value(EMPTY_VALUE);
Self::new(empty.clone(), empty)
}
pub fn literal(builder: &mut CircuitBuilder<F, D>, value: &ValueTarget) -> Self {
let empty = builder.constant_value(EMPTY_VALUE);
Self::new(value.clone(), empty)
}
pub fn anchored_key(
_builder: &mut CircuitBuilder<F, D>,
pod_id: &ValueTarget,
key: &ValueTarget,
) -> Self {
Self::new(pod_id.clone(), key.clone())
}
}
#[derive(Clone)]
pub struct StatementTarget {
pub predicate: [Target; Params::predicate_size()],
pub args: Vec<[Target; STATEMENT_ARG_F_LEN]>,
pub args: Vec<StatementArgTarget>,
}
impl StatementTarget {
@ -62,18 +106,16 @@ impl StatementTarget {
builder: &mut CircuitBuilder<F, D>,
params: &Params,
predicate: NativePredicate,
args: &[[Target; STATEMENT_ARG_F_LEN]],
args: &[StatementArgTarget],
) -> Self {
let predicate_vec = builder.constants(&Predicate::Native(predicate).to_fields(params));
Self {
predicate: array::from_fn(|i| predicate_vec[i]),
args: args
.iter()
.map(|arg| *arg)
.chain(
iter::repeat([builder.zero(); STATEMENT_ARG_F_LEN])
.take(params.max_statement_args - args.len()),
)
.cloned()
.chain(iter::repeat_with(|| StatementArgTarget::none(builder)))
.take(params.max_statement_args)
.collect(),
}
}
@ -92,7 +134,7 @@ impl StatementTarget {
.take(params.max_statement_args)
.enumerate()
{
pw.set_target_arr(&self.args[i], &arg.to_fields(params))?;
self.args[i].set_targets(pw, params, arg)?;
}
Ok(())
}
@ -159,7 +201,7 @@ impl Flattenable for StatementTarget {
fn flatten(&self) -> Vec<Target> {
self.predicate
.iter()
.chain(self.args.iter().flatten())
.chain(self.args.iter().flat_map(|a| &a.elements))
.cloned()
.collect()
}
@ -172,7 +214,11 @@ impl Flattenable for StatementTarget {
);
let predicate: [Target; Params::predicate_size()] = array::from_fn(|i| v[i]);
let args = (0..num_args)
.map(|i| array::from_fn(|j| v[Params::predicate_size() + i * STATEMENT_ARG_F_LEN + j]))
.map(|i| StatementArgTarget {
elements: array::from_fn(|j| {
v[Params::predicate_size() + i * STATEMENT_ARG_F_LEN + j]
}),
})
.collect();
Self { predicate, args }
@ -192,7 +238,7 @@ pub trait CircuitBuilderPod<F: RichField + Extendable<D>, const D: usize> {
// Convenience methods for checking values.
/// Checks whether `xs` is right-padded with 0s so as to represent a `Value`.
fn statement_arg_is_value(&mut self, xs: &[Target]) -> BoolTarget;
fn statement_arg_is_value(&mut self, arg: &StatementArgTarget) -> BoolTarget;
/// Checks whether `x < y` if `b` is true. This involves checking
/// that `x` and `y` each consist of two `u32` limbs.
fn assert_less_if(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget);
@ -231,7 +277,9 @@ impl CircuitBuilderPod<F, D> for CircuitBuilder<F, D> {
StatementTarget {
predicate: self.add_virtual_target_arr(),
args: (0..params.max_statement_args)
.map(|_| self.add_virtual_target_arr())
.map(|_| StatementArgTarget {
elements: self.add_virtual_target_arr(),
})
.collect(),
}
}
@ -272,11 +320,11 @@ impl CircuitBuilderPod<F, D> for CircuitBuilder<F, D> {
})
}
fn statement_arg_is_value(&mut self, xs: &[Target]) -> BoolTarget {
fn statement_arg_is_value(&mut self, arg: &StatementArgTarget) -> BoolTarget {
let zeros = iter::repeat(self.zero())
.take(STATEMENT_ARG_F_LEN - VALUE_SIZE)
.collect::<Vec<_>>();
self.is_equal_slice(&xs[VALUE_SIZE..], &zeros)
self.is_equal_slice(&arg.elements[VALUE_SIZE..], &zeros)
}
fn assert_less_if(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget) {

View file

@ -1,28 +1,16 @@
use anyhow::Result;
use itertools::Itertools;
use itertools::zip_eq;
use plonky2::{
field::types::Field,
hash::{
hash_types::{HashOut, HashOutTarget},
poseidon::PoseidonHash,
},
iop::{
target::{BoolTarget, Target},
witness::{PartialWitness, WitnessWrite},
},
hash::{hash_types::HashOutTarget, poseidon::PoseidonHash},
iop::{target::BoolTarget, witness::PartialWitness},
plonk::circuit_builder::CircuitBuilder,
};
use std::collections::HashMap;
use std::iter;
use crate::backends::plonky2::basetypes::{Hash, Value, D, EMPTY_HASH, EMPTY_VALUE, F, VALUE_SIZE};
use crate::backends::plonky2::basetypes::{Value, D, EMPTY_HASH, F, VALUE_SIZE};
use crate::backends::plonky2::circuits::common::{
CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget,
};
use crate::backends::plonky2::primitives::merkletree::MerkleTree;
use crate::backends::plonky2::primitives::merkletree::{
MerkleProofExistenceGadget, MerkleProofExistenceTarget,
};
use crate::backends::plonky2::mock::mainpod;
use crate::backends::plonky2::signedpod::SignedPod;
use crate::middleware::{
hash_str, AnchoredKey, NativeOperation, NativePredicate, Params, PodType, Statement,
@ -118,18 +106,20 @@ impl OperationVerifyGadget {
// The values embedded in the op args must match, the last
// `STATEMENT_ARG_F_LEN - VALUE_SIZE` slots of each being 0.
let arg1_value = resolved_op_args[0].args[1];
let arg2_value = resolved_op_args[1].args[1];
let arg1_value = &resolved_op_args[0].args[1];
let arg2_value = &resolved_op_args[1].args[1];
let op_arg_range_checks = [
builder.statement_arg_is_value(&arg1_value),
builder.statement_arg_is_value(&arg2_value),
builder.statement_arg_is_value(arg1_value),
builder.statement_arg_is_value(arg2_value),
];
let op_arg_range_ok = builder.all(op_arg_range_checks);
let op_args_eq =
builder.is_equal_slice(&arg1_value[..VALUE_SIZE], &arg2_value[..VALUE_SIZE]);
let op_args_eq = builder.is_equal_slice(
&arg1_value.elements[..VALUE_SIZE],
&arg2_value.elements[..VALUE_SIZE],
);
let arg1_key = resolved_op_args[0].args[0];
let arg2_key = resolved_op_args[1].args[0];
let arg1_key = resolved_op_args[0].args[0].clone();
let arg2_key = resolved_op_args[1].args[0].clone();
let expected_statement = StatementTarget::new_native(
builder,
&self.params,
@ -167,21 +157,21 @@ impl OperationVerifyGadget {
// The values embedded in the op args must satisfy `<`, the
// last `STATEMENT_ARG_F_LEN - VALUE_SIZE` slots of each being
// 0.
let arg1_value = resolved_op_args[0].args[1];
let arg2_value = resolved_op_args[1].args[1];
let op_arg_range_checks = [&arg1_value, &arg2_value]
let arg1_value = &resolved_op_args[0].args[1];
let arg2_value = &resolved_op_args[1].args[1];
let op_arg_range_checks = [arg1_value, arg2_value]
.into_iter()
.map(|x| builder.statement_arg_is_value(x))
.collect::<Vec<_>>();
let op_arg_range_ok = builder.all(op_arg_range_checks);
builder.assert_less_if(
op_code_ok,
ValueTarget::from_slice(&arg1_value[..VALUE_SIZE]),
ValueTarget::from_slice(&arg2_value[..VALUE_SIZE]),
ValueTarget::from_slice(&arg1_value.elements[..VALUE_SIZE]),
ValueTarget::from_slice(&arg2_value.elements[..VALUE_SIZE]),
);
let arg1_key = resolved_op_args[0].args[0];
let arg2_key = resolved_op_args[1].args[0];
let arg1_key = resolved_op_args[0].args[0].clone();
let arg2_key = resolved_op_args[1].args[0].clone();
let expected_statement = StatementTarget::new_native(
builder,
&self.params,
@ -222,14 +212,17 @@ impl OperationVerifyGadget {
let expected_arg_prefix = builder.constants(
&StatementArg::Key(AnchoredKey(SELF, EMPTY_HASH)).to_fields(&self.params)[..VALUE_SIZE],
);
let arg_prefix_ok = builder.is_equal_slice(&st.args[0][..VALUE_SIZE], &expected_arg_prefix);
let arg_prefix_ok =
builder.is_equal_slice(&st.args[0].elements[..VALUE_SIZE], &expected_arg_prefix);
let dupe_check = {
let individual_checks = prev_statements
.into_iter()
.map(|ps| {
.enumerate()
.map(|(i, ps)| {
let same_predicate = builder.is_equal_slice(&st.predicate, &ps.predicate);
let same_anchored_key = builder.is_equal_slice(&st.args[0], &ps.args[0]);
let same_anchored_key =
builder.is_equal_slice(&st.args[0].elements, &ps.args[0].elements);
builder.and(same_predicate, same_anchored_key)
})
.collect::<Vec<_>>();
@ -292,7 +285,22 @@ impl MainPodVerifyGadget {
// Build the statement array
let mut statements = Vec::new();
for signed_pod in &signed_pods {
statements.extend_from_slice(signed_pod.pub_statements().as_slice());
statements.extend_from_slice(signed_pod.pub_statements(builder, false).as_slice());
}
debug_assert_eq!(
statements.len(),
self.params.max_input_signed_pods * self.params.max_signed_pod_values
);
// TODO: Fill with input main pods
for _main_pod in 0..self.params.max_input_main_pods {
for _statement in 0..self.params.max_public_statements {
statements.push(StatementTarget::new_native(
builder,
&self.params,
NativePredicate::None,
&[],
))
}
}
// Add the input (private and public) statements and corresponding operations
@ -304,19 +312,22 @@ impl MainPodVerifyGadget {
}
let input_statements = &statements[input_statements_offset..];
let pub_statements = &input_statements[statements.len() - params.max_public_statements..];
let pub_statements =
&input_statements[input_statements.len() - params.max_public_statements..];
// 2. Calculate the Pod Id from the public statements
let pub_statements_flattened = pub_statements
.iter()
.map(|s| s.predicate.iter().chain(s.args.iter().flatten()))
.map(|s| {
s.predicate
.iter()
.chain(s.args.iter().flat_map(|a| &a.elements))
})
.flatten()
.cloned()
.collect();
let id = builder.hash_n_to_hash_no_pad::<PoseidonHash>(pub_statements_flattened);
// 3. TODO check that all `input_statements` of type `ValueOf` with origin=SELF have unique keys (no duplicates). Maybe we can do this via the NewEntry operation (check that the key doesn't exist in a previous statement with ID=SELF)
// 4. Verify type
let type_statement = &pub_statements[0];
// TODO: Store this hash in a global static with lazy init so that we don't have to
@ -330,10 +341,12 @@ impl MainPodVerifyGadget {
);
builder.connect_flattenable(type_statement, &expected_type_statement);
// 3. check that all `input_statements` of type `ValueOf` with origin=SELF have unique keys
// (no duplicates). We do this in the verification of NewEntry operation.
// 5. Verify input statements
let mut op_verifications = Vec::new();
for (i, (st, op)) in input_statements.iter().zip(operations.iter()).enumerate() {
let prev_statements = &statements[..input_statements_offset + i - 1];
let prev_statements = &statements[..input_statements_offset + i];
let op_verification = OperationVerifyGadget {
params: params.clone(),
}
@ -352,7 +365,7 @@ impl MainPodVerifyGadget {
}
}
struct MainPodVerifyTarget {
pub struct MainPodVerifyTarget {
params: Params,
id: HashOutTarget,
signed_pods: Vec<SignedPodVerifyTarget>,
@ -362,25 +375,33 @@ struct MainPodVerifyTarget {
op_verifications: Vec<OperationVerifyTarget>,
}
struct MainPodVerifyInput {
signed_pods: Vec<SignedPod>,
pub struct MainPodVerifyInput {
pub signed_pods: Vec<SignedPod>,
pub statements: Vec<mainpod::Statement>,
pub operations: Vec<mainpod::Operation>,
}
impl MainPodVerifyTarget {
fn set_targets(&self, pw: &mut PartialWitness<F>, input: &MainPodVerifyInput) -> Result<()> {
pub fn set_targets(
&self,
pw: &mut PartialWitness<F>,
input: &MainPodVerifyInput,
) -> Result<()> {
assert!(input.signed_pods.len() <= self.params.max_input_signed_pods);
for (i, signed_pod) in input.signed_pods.iter().enumerate() {
self.signed_pods[i].set_targets(pw, signed_pod)?;
}
// Padding
// TODO: Instead of using an input for padding, use a canonical minimal SignedPod
let pad_pod = &input.signed_pods[0];
for i in input.signed_pods.len()..self.params.max_input_signed_pods {
// TODO: We need to disable the verification for the unused slots.
// self.signed_pods[i].set_targets(pw, signed_pod)?;
self.signed_pods[i].set_targets(pw, pad_pod)?;
}
assert_eq!(input.statements.len(), self.params.max_statements);
for (i, (st, op)) in zip_eq(&input.statements, &input.operations).enumerate() {
self.statements[i].set_targets(pw, &self.params, st)?;
self.operations[i].set_targets(pw, &self.params, op)?;
}
// TODO: set_targets for:
// - statements
// - operations
// - op_verifications
Ok(())
}
}
@ -470,30 +491,38 @@ mod tests {
// NewEntry
let st1: mainpod::Statement =
Statement::ValueOf(AnchoredKey(SELF, "hello".into()), 55.into()).into();
let op = mainpod::Operation(
OperationType::Native(NativeOperation::NewEntry),
vec![],
OperationAux::None,
);
operation_verify(st1.clone(), op, vec![])?;
// Copy
let op = mainpod::Operation(
OperationType::Native(NativeOperation::CopyStatement),
vec![OperationArg::Index(0)],
OperationAux::None,
);
operation_verify(st, op, prev_statements)?;
// Eq
let st2: mainpod::Statement = Statement::ValueOf(
AnchoredKey(PodId(Value::from(75).into()), "hello".into()),
55.into(),
)
.into();
let prev_statements = vec![st2];
let op = mainpod::Operation(
OperationType::Native(NativeOperation::NewEntry),
vec![],
OperationAux::None,
);
operation_verify(st1.clone(), op, prev_statements.clone())?;
// Copy
let st: mainpod::Statement = Statement::None.into();
let op = mainpod::Operation(
OperationType::Native(NativeOperation::CopyStatement),
vec![OperationArg::Index(0)],
OperationAux::None,
);
let prev_statements = vec![Statement::None.into()];
operation_verify(st, op, prev_statements)?;
// Eq
let st2: mainpod::Statement = Statement::ValueOf(
AnchoredKey(PodId(Value::from(75).into()), "world".into()),
55.into(),
)
.into();
let st: mainpod::Statement = Statement::Equal(
AnchoredKey(SELF, "hello".into()),
AnchoredKey(PodId(Value::from(75).into()), "hello".into()),
AnchoredKey(PodId(Value::from(75).into()), "world".into()),
)
.into();
let op = mainpod::Operation(

View file

@ -2,20 +2,24 @@ use anyhow::Result;
use itertools::Itertools;
use plonky2::{
hash::hash_types::{HashOut, HashOutTarget},
iop::target::Target,
iop::witness::{PartialWitness, WitnessWrite},
plonk::circuit_builder::CircuitBuilder,
};
use std::iter;
use crate::backends::plonky2::{
basetypes::{Value, D, EMPTY_VALUE, F},
circuits::common::{CircuitBuilderPod, StatementTarget, ValueTarget},
circuits::common::{CircuitBuilderPod, StatementArgTarget, StatementTarget, ValueTarget},
primitives::{
merkletree::{MerkleProof, MerkleProofExistenceGadget, MerkleProofExistenceTarget},
signature::{PublicKey, SignatureVerifyGadget, SignatureVerifyTarget},
},
signedpod::SignedPod,
};
use crate::middleware::{hash_str, Params, PodType, KEY_SIGNER, KEY_TYPE};
use crate::middleware::{
hash_str, NativePredicate, Params, PodType, Predicate, ToFields, KEY_SIGNER, KEY_TYPE, SELF,
};
pub struct SignedPodVerifyGadget {
pub params: Params,
@ -73,9 +77,39 @@ pub struct SignedPodVerifyTarget {
}
impl SignedPodVerifyTarget {
pub fn pub_statements(&self) -> Vec<StatementTarget> {
// TODO: Here we need to use the self.id in the ValueOf statements
todo!()
pub fn pub_statements(
&self,
builder: &mut CircuitBuilder<F, D>,
self_id: bool,
) -> Vec<StatementTarget> {
let mut statements = Vec::new();
let predicate: [Target; Params::predicate_size()] = builder
.constants(&Predicate::Native(NativePredicate::ValueOf).to_fields(&self.params))
.try_into()
.expect("size predicate_size");
let pod_id = if self_id {
builder.constant_value(SELF.0.into())
} else {
ValueTarget {
elements: self.id.elements,
}
};
for mt_proof in &self.mt_proofs {
let args = [
StatementArgTarget::anchored_key(builder, &pod_id, &mt_proof.key),
StatementArgTarget::literal(builder, &mt_proof.value),
]
.into_iter()
.chain(iter::repeat_with(|| StatementArgTarget::none(builder)))
.take(self.params.max_statement_args)
.collect();
let statement = StatementTarget {
predicate: predicate.clone(),
args,
};
statements.push(statement);
}
statements
}
pub fn set_targets(&self, pw: &mut PartialWitness<F>, pod: &SignedPod) -> Result<()> {