Skip to content

Commit

Permalink
Basic skeleton: CPU, added E2E test
Browse files Browse the repository at this point in the history
  • Loading branch information
supragya committed Jul 8, 2024
1 parent 57d14ce commit 14645c8
Show file tree
Hide file tree
Showing 4 changed files with 285 additions and 17 deletions.
88 changes: 88 additions & 0 deletions src/e2e_tests.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
#[cfg(test)]
mod tests {
use std::collections::HashMap;

use plonky2::{
field::goldilocks_field::GoldilocksField,
plonk::config::{
GenericConfig,
PoseidonGoldilocksConfig,
},
util::timing::TimingTree,
};
use starky::{
config::StarkConfig,
proof::StarkProofWithPublicInputs,
prover::prove,
verifier::verify_stark_proof,
};

use crate::{
stark_program_instructions::ProgramInstructionsStark,
vm_specs::{
Instruction,
MemoryLocation,
Program,
Register,
},
};

#[test]
fn test_add_program() {
let instructions = vec![
Instruction::Lb(Register::R0, MemoryLocation(0x40)),
Instruction::Lb(Register::R1, MemoryLocation(0x41)),
Instruction::Add(Register::R0, Register::R1),
Instruction::Sb(Register::R0, MemoryLocation(0x42)),
Instruction::Halt,
];

let code = instructions
.into_iter()
.enumerate()
.map(|(idx, inst)| (idx as u8, inst))
.collect::<HashMap<u8, Instruction>>();

let memory_init: HashMap<u8, u8> =
HashMap::from_iter(vec![(0x40, 0x20), (0x41, 0x45)]);

let program = Program {
entry_point: 0,
code,
memory_init,
};

// Generate the static part of the proof
let program_proof = {
const D: usize = 2;
type C = PoseidonGoldilocksConfig;
type F = <C as GenericConfig<D>>::F;
type S = ProgramInstructionsStark<F, D>;
type PR = StarkProofWithPublicInputs<GoldilocksField, C, 2>;

let stark = S::new();
let mut config = StarkConfig::standard_fast_config();
// Need to do this since our table is small. Need atleast 1<<5
// sized table to not affect this
config
.fri_config
.cap_height = 1;
let program = Program::default();
let trace =
ProgramInstructionsStark::<F, D>::generate_trace(&program);
let proof: Result<PR, anyhow::Error> = prove(
stark.clone(),
&config,
trace,
&[],
&mut TimingTree::default(),
);
assert!(proof.is_ok());
let proof = proof.unwrap();
let verification =
verify_stark_proof(stark, proof.clone(), &config);
assert!(verification.is_ok());
proof
};
}
}
6 changes: 6 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,5 +17,11 @@ mod preflight_simulator;
mod vm_specs;

// STARK tables -------------
//#[allow(dead_code)]
//mod stark_cpu;
#[allow(dead_code)]
mod stark_program_instructions;

// END TO END TEST ----------
#[allow(dead_code)]
mod e2e_tests;
182 changes: 182 additions & 0 deletions src/stark_cpu.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,182 @@
//! This file is an encoding of all the execution seen at the "CPU"
//! It is dynamic and changes depending on the memory_init. It has
//! to be linked to the static code "Program" by having a cross-table
//! -lookup with `ProgramInstructionsStark`.

use core::marker::PhantomData;
use plonky2::{
field::{
extension::{
Extendable,
FieldExtension,
},
packed::PackedField,
polynomial::PolynomialValues,
},
hash::hash_types::RichField,
iop::ext_target::ExtensionTarget,
plonk::circuit_builder::CircuitBuilder,
};
use starky::{
constraint_consumer::{
ConstraintConsumer,
RecursiveConstraintConsumer,
},
evaluation_frame::{
StarkEvaluationFrame,
StarkFrame,
},
stark::Stark,
util::trace_rows_to_poly_values,
};

use crate::vm_specs::Program;

// Table description:
// +-----+----+--------+--------+--------------+---------+
// | Clk | PC | Reg R0 | Reg R1 | MemoryAddr | Opcode* |
// +-----+----+--------+--------+--------------+---------+
// | .. | .. | ... | ... | .... | ... |
// +-----+----+--------+--------+--------------+---------+
//
// `Opcode*` means `Opcode` that is one-hot encoded
// 5 Columns for `Clk`, `PC`, `Reg R0`, `Reg R1`, `MemoryAccess`
// 10 Columns for opcodes. See `Instruction::get_opcode`.
const NUMBER_OF_COLS: usize = 5 + 10;
const PUBLIC_INPUTS: usize = 0;

#[derive(Clone, Copy)]
pub struct CPUStark<F, const D: usize> {
pub _f: PhantomData<F>,
}

impl<F, const D: usize> CPUStark<F, D>
where
F: RichField + Extendable<D>,
{
pub fn new() -> Self {
Self { _f: PhantomData }
}

pub fn generate_trace(prog: ) -> Vec<PolynomialValues<F>>
where
F: RichField,
{
let mut trace = prog
.code
.iter()
.map(|(pc, inst)| {
[
// Program Counter (ID = 0)
F::from_canonical_u8(*pc),
// Instruction Opcode (ID = 1)
F::from_canonical_u8(inst.get_opcode()),
// Filter, true if actual instructions (ID = 2)
F::ONE,
]
})
.collect::<Vec<[F; NUMBER_OF_COLS]>>();

// Need to pad the trace to a len of some power of 2
let pow2_len = trace
.len()
.next_power_of_two();
trace.resize(pow2_len, [F::ZERO, F::ZERO, F::ZERO]);

// Convert into polynomial values
trace_rows_to_poly_values(trace)
}
}

impl<F, const D: usize> Stark<F, D> for CPUStark<F, D>
where
F: RichField + Extendable<D>,
{
type EvaluationFrame<FE, P, const D2: usize> = StarkFrame<P, P::Scalar, NUMBER_OF_COLS, PUBLIC_INPUTS>
where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>;
type EvaluationFrameTarget = StarkFrame<
ExtensionTarget<D>,
ExtensionTarget<D>,
NUMBER_OF_COLS,
PUBLIC_INPUTS,
>;

const COLUMNS: usize = NUMBER_OF_COLS;
const PUBLIC_INPUTS: usize = PUBLIC_INPUTS;

fn eval_packed_generic<FE, P, const D2: usize>(
&self,
vars: &Self::EvaluationFrame<FE, P, D2>,
yield_constr: &mut ConstraintConsumer<P>,
) where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
}

fn eval_ext_circuit(
&self,
_builder: &mut CircuitBuilder<F, D>,
_vars: &Self::EvaluationFrameTarget,
_yield_constr: &mut RecursiveConstraintConsumer<F, D>,
) {
unimplemented!()
}

fn constraint_degree(&self) -> usize {
3
}
}

#[cfg(test)]
mod tests {

use plonky2::{
field::goldilocks_field::GoldilocksField,
plonk::config::{
GenericConfig,
PoseidonGoldilocksConfig,
},
util::timing::TimingTree,
};
use starky::{
config::StarkConfig,
proof::StarkProofWithPublicInputs,
prover::prove,
verifier::verify_stark_proof,
};

use super::*;

#[test]
fn test_nil_program() {
const D: usize = 2;
type C = PoseidonGoldilocksConfig;
type F = <C as GenericConfig<D>>::F;
type S = CPUStark<F, D>;
type PR = StarkProofWithPublicInputs<GoldilocksField, C, 2>;

let stark = S::new();
let mut config = StarkConfig::standard_fast_config();
// Need to do this since our table is small. Need atleast 1<<5
// sized table to not affect this
config
.fri_config
.cap_height = 1;
let program = Program::default();
let trace =
CPUStark::<F, D>::generate_program_instructions_trace(&program);
let proof: Result<PR, anyhow::Error> = prove(
stark.clone(),
&config,
trace,
&[],
&mut TimingTree::default(),
);
assert!(proof.is_ok());
let verification = verify_stark_proof(stark, proof.unwrap(), &config);
assert!(verification.is_ok());
}
}
26 changes: 9 additions & 17 deletions src/stark_program_instructions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,16 +33,13 @@ use starky::{

use crate::vm_specs::Program;

/// Represents one row in a STARK table, contains `is_filter` which
/// should be set to `true` in case it represents an actual instruction
//#[repr(C)]
//#[derive(Default, Clone, Copy, PartialEq, Debug)]
//pub struct ProgramInstructions<T> {
// pub program_counter: T, // ID = 0
// pub opcode: T, // ID = 1
// pub filter: T, // ID = 2
//}

// Table description:
// +-----------------+--------------------+-------------+
// | Program Counter | Instruction Opcode | Is_Executed |
// +-----------------+--------------------+-------------+
// | .... | .... | .... |
// | .... | .... | .... |
// +-----------------+--------------------+-------------+
const NUMBER_OF_COLS: usize = 3;
const PUBLIC_INPUTS: usize = 0;

Expand All @@ -59,9 +56,7 @@ where
Self { _f: PhantomData }
}

pub fn generate_program_instructions_trace(
prog: &Program
) -> Vec<PolynomialValues<F>>
pub fn generate_trace(prog: &Program) -> Vec<PolynomialValues<F>>
where
F: RichField,
{
Expand Down Expand Up @@ -174,10 +169,7 @@ mod tests {
.fri_config
.cap_height = 1;
let program = Program::default();
let trace =
ProgramInstructionsStark::<F, D>::generate_program_instructions_trace(
&program,
);
let trace = ProgramInstructionsStark::<F, D>::generate_trace(&program);
let proof: Result<PR, anyhow::Error> = prove(
stark.clone(),
&config,
Expand Down

0 comments on commit 14645c8

Please sign in to comment.