Files
darkfi/src/vm.rs
2021-09-23 17:44:07 +02:00

444 lines
17 KiB
Rust

use bellman::{gadgets::Assignment, groth16, Circuit, ConstraintSystem, SynthesisError};
use bls12_381::Bls12;
use bls12_381::Scalar;
use ff::{Field, PrimeField};
use rand::rngs::OsRng;
use std::ops::{AddAssign, MulAssign, SubAssign};
use std::time::Instant;
use crate::error::Result;
pub struct ZkVirtualMachine {
pub constants: Vec<Scalar>,
pub alloc: Vec<(AllocType, VariableIndex)>,
pub ops: Vec<CryptoOperation>,
pub constraints: Vec<ConstraintInstruction>,
pub aux: Vec<Scalar>,
pub params: Option<groth16::Parameters<Bls12>>,
pub verifying_key: Option<groth16::PreparedVerifyingKey<Bls12>>,
}
pub type VariableIndex = usize;
pub enum VariableRef {
Aux(VariableIndex),
Local(VariableIndex),
}
pub enum CryptoOperation {
Set(VariableRef, VariableRef),
Mul(VariableRef, VariableRef),
Add(VariableRef, VariableRef),
Sub(VariableRef, VariableRef),
Load(VariableRef, VariableIndex),
Divide(VariableRef, VariableRef),
Double(VariableRef),
Square(VariableRef),
Invert(VariableRef),
UnpackBits(VariableRef, VariableRef, VariableRef),
Local,
Debug(String, VariableRef),
DumpAlloc,
DumpLocal,
}
#[derive(Clone)]
pub enum AllocType {
Private,
Public,
}
#[derive(Debug, Clone)]
pub enum ConstraintInstruction {
Lc0Add(VariableIndex),
Lc1Add(VariableIndex),
Lc2Add(VariableIndex),
Lc0Sub(VariableIndex),
Lc1Sub(VariableIndex),
Lc2Sub(VariableIndex),
Lc0AddOne,
Lc1AddOne,
Lc2AddOne,
Lc0SubOne,
Lc1SubOne,
Lc2SubOne,
Lc0AddCoeff(VariableIndex, VariableIndex),
Lc1AddCoeff(VariableIndex, VariableIndex),
Lc2AddCoeff(VariableIndex, VariableIndex),
Lc0AddConstant(VariableIndex),
Lc1AddConstant(VariableIndex),
Lc2AddConstant(VariableIndex),
Enforce,
LcCoeffReset,
LcCoeffDouble,
}
#[derive(Debug)]
pub enum ZkVmError {
DivisionByZero,
MalformedRange,
}
impl ZkVirtualMachine {
pub fn initialize(
&mut self,
params: &Vec<(VariableIndex, Scalar)>,
) -> std::result::Result<(), ZkVmError> {
// Resize array
self.aux = vec![Scalar::zero(); self.alloc.len()];
// Copy over the parameters
for (index, value) in params {
//println!("Setting {} to {:?}", index, value);
self.aux[*index] = *value;
}
let mut local_stack: Vec<Scalar> = Vec::new();
for op in &self.ops {
match op {
CryptoOperation::Set(self_, other) => {
let other = match other {
VariableRef::Aux(index) => self.aux[*index],
VariableRef::Local(index) => local_stack[*index],
};
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
*self_ = other;
}
CryptoOperation::Mul(self_, other) => {
let other = match other {
VariableRef::Aux(index) => self.aux[*index],
VariableRef::Local(index) => local_stack[*index],
};
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
self_.mul_assign(other);
}
CryptoOperation::Add(self_, other) => {
let other = match other {
VariableRef::Aux(index) => self.aux[*index],
VariableRef::Local(index) => local_stack[*index],
};
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
self_.add_assign(other);
}
CryptoOperation::Sub(self_, other) => {
let other = match other {
VariableRef::Aux(index) => self.aux[*index],
VariableRef::Local(index) => local_stack[*index],
};
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
self_.sub_assign(other);
}
CryptoOperation::Load(self_, const_index) => {
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
*self_ = self.constants[*const_index];
}
CryptoOperation::Divide(self_, other) => {
let other = match other {
VariableRef::Aux(index) => self.aux[*index],
VariableRef::Local(index) => local_stack[*index],
};
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
let ret = other.invert().map(|other| *self_ * other);
if bool::from(ret.is_some()) {
*self_ = ret.unwrap();
} else {
return Err(ZkVmError::DivisionByZero);
}
}
CryptoOperation::Double(self_) => {
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
*self_ = self_.double();
}
CryptoOperation::Square(self_) => {
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
*self_ = self_.square();
}
CryptoOperation::Invert(self_) => {
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
if self_.is_zero() {
return Err(ZkVmError::DivisionByZero);
} else {
*self_ = self_.invert().unwrap();
}
}
CryptoOperation::UnpackBits(value, start, end) => {
let value = match value {
VariableRef::Aux(index) => self.aux[*index],
VariableRef::Local(index) => local_stack[*index],
};
let (self_, start_index, end_index) = match start {
VariableRef::Aux(start_index) => match end {
VariableRef::Aux(end_index) => (&mut self.aux, start_index, end_index),
VariableRef::Local(_) => {
return Err(ZkVmError::MalformedRange);
}
},
VariableRef::Local(start_index) => match end {
VariableRef::Aux(_) => {
return Err(ZkVmError::MalformedRange);
}
VariableRef::Local(end_index) => {
(&mut local_stack, start_index, end_index)
}
},
};
if start_index > end_index {
return Err(ZkVmError::MalformedRange);
}
if (end_index + 1) - start_index != 256 {
return Err(ZkVmError::MalformedRange);
}
if *end_index >= self_.len() {
return Err(ZkVmError::MalformedRange);
}
for (i, bit) in value.to_le_bits().into_iter().cloned().enumerate() {
match bit {
true => self_[start_index + i] = Scalar::one(),
false => self_[start_index + i] = Scalar::zero(),
}
}
}
CryptoOperation::Local => {
local_stack.push(Scalar::zero());
}
CryptoOperation::Debug(debug_str, self_) => {
let self_ = match self_ {
VariableRef::Aux(index) => &mut self.aux[*index],
VariableRef::Local(index) => &mut local_stack[*index],
};
println!("{}", debug_str);
println!("value = {:?}", self_);
}
CryptoOperation::DumpAlloc => {
println!("-------------------");
println!("alloc");
println!("-------------------");
for (i, value) in self.aux.iter().enumerate() {
println!("{}: {:?}", i, value);
}
println!("-------------------");
}
CryptoOperation::DumpLocal => {
println!("-------------------");
println!("local");
println!("-------------------");
for (i, value) in local_stack.iter().enumerate() {
println!("{}: {:?}", i, value);
}
println!("-------------------");
}
}
}
Ok(())
}
pub fn public(&self) -> Vec<(VariableIndex, Scalar)> {
let mut publics = Vec::new();
for (alloc_type, index) in &self.alloc {
match alloc_type {
AllocType::Private => {}
AllocType::Public => {
let scalar = self.aux[*index];
publics.push((*index, scalar));
}
}
}
publics
}
pub fn setup(&mut self) -> Result<()> {
let start = Instant::now();
// Create parameters for our circuit. In a production deployment these would
// be generated securely using a multiparty computation.
self.params = Some({
let circuit = ZkVmCircuit {
aux: vec![None; self.aux.len()],
alloc: self.alloc.clone(),
constraints: self.constraints.clone(),
constants: self.constants.clone(),
};
groth16::generate_random_parameters::<Bls12, _, _>(circuit, &mut OsRng)?
});
println!("Setup: [{:?}]", start.elapsed());
self.verifying_key = Some(groth16::prepare_verifying_key(
&self.params.as_ref().unwrap().vk,
));
Ok(())
}
pub fn prove(&self) -> groth16::Proof<Bls12> {
let aux = self.aux.iter().map(|scalar| Some(*scalar)).collect();
// Create an instance of our circuit (with the preimage as a witness).
let circuit = ZkVmCircuit {
aux,
alloc: self.alloc.clone(),
constraints: self.constraints.clone(),
constants: self.constants.clone(),
};
let start = Instant::now();
// Create a Groth16 proof with our parameters.
let proof =
groth16::create_random_proof(circuit, self.params.as_ref().unwrap(), &mut OsRng)
.unwrap();
println!("Prove: [{:?}]", start.elapsed());
proof
}
pub fn verify(&self, proof: &groth16::Proof<Bls12>, public_values: &[Scalar]) -> bool {
let start = Instant::now();
let is_passed =
groth16::verify_proof(self.verifying_key.as_ref().unwrap(), proof, public_values)
.is_ok();
println!("Verify: [{:?}]", start.elapsed());
is_passed
}
}
pub struct ZkVmCircuit {
aux: Vec<Option<bls12_381::Scalar>>,
alloc: Vec<(AllocType, VariableIndex)>,
constraints: Vec<ConstraintInstruction>,
constants: Vec<Scalar>,
}
impl Circuit<bls12_381::Scalar> for ZkVmCircuit {
fn synthesize<CS: ConstraintSystem<bls12_381::Scalar>>(
self,
cs: &mut CS,
) -> std::result::Result<(), SynthesisError> {
let mut variables = Vec::new();
for (alloc_type, index) in &self.alloc {
match alloc_type {
AllocType::Private => {
let var = cs.alloc(|| "private alloc", || Ok(*self.aux[*index].get()?))?;
variables.push(var);
}
AllocType::Public => {
let var = cs.alloc_input(|| "public alloc", || Ok(*self.aux[*index].get()?))?;
variables.push(var);
}
}
}
let mut coeff = bls12_381::Scalar::one();
let mut lc0 = bellman::LinearCombination::<Scalar>::zero();
let mut lc1 = bellman::LinearCombination::<Scalar>::zero();
let mut lc2 = bellman::LinearCombination::<Scalar>::zero();
for constraint in self.constraints {
match constraint {
ConstraintInstruction::Lc0Add(index) => {
lc0 = lc0 + (coeff, variables[index]);
}
ConstraintInstruction::Lc1Add(index) => {
lc1 = lc1 + (coeff, variables[index]);
}
ConstraintInstruction::Lc2Add(index) => {
lc2 = lc2 + (coeff, variables[index]);
}
ConstraintInstruction::Lc0Sub(index) => {
lc0 = lc0 - (coeff, variables[index]);
}
ConstraintInstruction::Lc1Sub(index) => {
lc1 = lc1 - (coeff, variables[index]);
}
ConstraintInstruction::Lc2Sub(index) => {
lc2 = lc2 - (coeff, variables[index]);
}
ConstraintInstruction::Lc0AddOne => {
lc0 = lc0 + CS::one();
}
ConstraintInstruction::Lc1AddOne => {
lc1 = lc1 + CS::one();
}
ConstraintInstruction::Lc2AddOne => {
lc2 = lc2 + CS::one();
}
ConstraintInstruction::Lc0SubOne => {
lc0 = lc0 - CS::one();
}
ConstraintInstruction::Lc1SubOne => {
lc1 = lc1 - CS::one();
}
ConstraintInstruction::Lc2SubOne => {
lc2 = lc2 - CS::one();
}
ConstraintInstruction::Lc0AddCoeff(const_index, index) => {
lc0 = lc0 + (self.constants[const_index], variables[index]);
}
ConstraintInstruction::Lc1AddCoeff(const_index, index) => {
lc1 = lc1 + (self.constants[const_index], variables[index]);
}
ConstraintInstruction::Lc2AddCoeff(const_index, index) => {
lc2 = lc2 + (self.constants[const_index], variables[index]);
}
ConstraintInstruction::Lc0AddConstant(const_index) => {
lc0 = lc0 + (self.constants[const_index], CS::one());
}
ConstraintInstruction::Lc1AddConstant(const_index) => {
lc1 = lc1 + (self.constants[const_index], CS::one());
}
ConstraintInstruction::Lc2AddConstant(const_index) => {
lc2 = lc2 + (self.constants[const_index], CS::one());
}
ConstraintInstruction::Enforce => {
cs.enforce(
|| "constraint",
|_| lc0.clone(),
|_| lc1.clone(),
|_| lc2.clone(),
);
coeff = bls12_381::Scalar::one();
lc0 = bellman::LinearCombination::<Scalar>::zero();
lc1 = bellman::LinearCombination::<Scalar>::zero();
lc2 = bellman::LinearCombination::<Scalar>::zero();
}
ConstraintInstruction::LcCoeffReset => {
coeff = bls12_381::Scalar::one();
}
ConstraintInstruction::LcCoeffDouble => {
coeff = coeff.double();
}
}
}
Ok(())
}
}