Skip to content

Commit

Permalink
style: Move AIR tests to air.rs
Browse files Browse the repository at this point in the history
  • Loading branch information
aszepieniec authored and jan-ferdinand committed Aug 12, 2024
1 parent 38a6a1d commit 14c270c
Show file tree
Hide file tree
Showing 2 changed files with 269 additions and 246 deletions.
268 changes: 268 additions & 0 deletions triton-vm/src/air.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,270 @@
pub mod memory_layout;
pub mod tasm_air_constraints;

#[cfg(test)]
mod test {
use itertools::Itertools;
use ndarray::Array1;
use proptest::collection::vec;
use proptest::prelude::*;
use proptest_arbitrary_interop::arb;
use std::collections::HashMap;
use test_strategy::proptest;

use crate::air::memory_layout::DynamicTasmConstraintEvaluationMemoryLayout;
use crate::air::memory_layout::IntegralMemoryLayout;
use crate::air::memory_layout::MemoryRegion;
use crate::air::memory_layout::StaticTasmConstraintEvaluationMemoryLayout;
use crate::air::tasm_air_constraints::dynamic_air_constraint_evaluation_tasm;
use crate::air::tasm_air_constraints::static_air_constraint_evaluation_tasm;
use crate::bfe;
use crate::instruction::AnInstruction;
use crate::instruction::LabelledInstruction;
use crate::table::challenges::Challenges;
use crate::table::extension_table::Evaluable;
use crate::table::extension_table::Quotientable;
use crate::table::master_table::MasterExtTable;
use crate::table::NUM_BASE_COLUMNS;
use crate::table::NUM_EXT_COLUMNS;
use crate::triton_instr;
use crate::twenty_first::prelude::x_field_element::EXTENSION_DEGREE;
use crate::BFieldElement;
use crate::NonDeterminism;
use crate::Program;
use crate::PublicInput;
use crate::VMState;
use crate::XFieldElement;

#[derive(Debug, Clone, test_strategy::Arbitrary)]
struct ConstraintEvaluationPoint {
#[strategy(vec(arb(), NUM_BASE_COLUMNS))]
#[map(Array1::from)]
curr_base_row: Array1<XFieldElement>,

#[strategy(vec(arb(), NUM_EXT_COLUMNS))]
#[map(Array1::from)]
curr_ext_row: Array1<XFieldElement>,

#[strategy(vec(arb(), NUM_BASE_COLUMNS))]
#[map(Array1::from)]
next_base_row: Array1<XFieldElement>,

#[strategy(vec(arb(), NUM_EXT_COLUMNS))]
#[map(Array1::from)]
next_ext_row: Array1<XFieldElement>,

#[strategy(arb())]
challenges: Challenges,

#[strategy(arb())]
#[filter(#static_memory_layout.is_integral())]
static_memory_layout: StaticTasmConstraintEvaluationMemoryLayout,
}

impl ConstraintEvaluationPoint {
fn evaluate_all_constraints_rust(&self) -> Vec<XFieldElement> {
let init = MasterExtTable::evaluate_initial_constraints(
self.curr_base_row.view(),
self.curr_ext_row.view(),
&self.challenges,
);
let cons = MasterExtTable::evaluate_consistency_constraints(
self.curr_base_row.view(),
self.curr_ext_row.view(),
&self.challenges,
);
let tran = MasterExtTable::evaluate_transition_constraints(
self.curr_base_row.view(),
self.curr_ext_row.view(),
self.next_base_row.view(),
self.next_ext_row.view(),
&self.challenges,
);
let term = MasterExtTable::evaluate_terminal_constraints(
self.curr_base_row.view(),
self.curr_ext_row.view(),
&self.challenges,
);

[init, cons, tran, term].concat()
}

fn evaluate_all_constraints_tasm_static(&self) -> Vec<XFieldElement> {
let program = self.tasm_static_constraint_evaluation_code();
let mut vm_state =
self.set_up_triton_vm_to_evaluate_constraints_in_tasm_static(&program);

vm_state.run().unwrap();

let output_list_ptr = vm_state.op_stack.pop().unwrap().value();
let num_quotients = MasterExtTable::NUM_CONSTRAINTS;
Self::read_xfe_list_at_address(vm_state.ram, output_list_ptr, num_quotients)
}

fn evaluate_all_constraints_tasm_dynamic(&self) -> Vec<XFieldElement> {
let program = self.tasm_dynamic_constraint_evaluation_code();
let mut vm_state =
self.set_up_triton_vm_to_evaluate_constraints_in_tasm_dynamic(&program);

vm_state.run().unwrap();

let output_list_ptr = vm_state.op_stack.pop().unwrap().value();
let num_quotients = MasterExtTable::NUM_CONSTRAINTS;
Self::read_xfe_list_at_address(vm_state.ram, output_list_ptr, num_quotients)
}

fn tasm_static_constraint_evaluation_code(&self) -> Program {
let mut source_code = static_air_constraint_evaluation_tasm(self.static_memory_layout);
source_code.push(triton_instr!(halt));
Program::new(&source_code)
}

fn tasm_dynamic_constraint_evaluation_code(&self) -> Program {
let dynamic_memory_layout = DynamicTasmConstraintEvaluationMemoryLayout {
free_mem_page_ptr: self.static_memory_layout.free_mem_page_ptr,
challenges_ptr: self.static_memory_layout.challenges_ptr,
};
let mut source_code = dynamic_air_constraint_evaluation_tasm(dynamic_memory_layout);
source_code.push(triton_instr!(halt));
Program::new(&source_code)
}

fn set_up_triton_vm_to_evaluate_constraints_in_tasm_static(
&self,
program: &Program,
) -> VMState {
let curr_base_row_ptr = self.static_memory_layout.curr_base_row_ptr;
let curr_ext_row_ptr = self.static_memory_layout.curr_ext_row_ptr;
let next_base_row_ptr = self.static_memory_layout.next_base_row_ptr;
let next_ext_row_ptr = self.static_memory_layout.next_ext_row_ptr;
let challenges_ptr = self.static_memory_layout.challenges_ptr;

let mut ram = HashMap::default();
Self::extend_ram_at_address(&mut ram, self.curr_base_row.to_vec(), curr_base_row_ptr);
Self::extend_ram_at_address(&mut ram, self.curr_ext_row.to_vec(), curr_ext_row_ptr);
Self::extend_ram_at_address(&mut ram, self.next_base_row.to_vec(), next_base_row_ptr);
Self::extend_ram_at_address(&mut ram, self.next_ext_row.to_vec(), next_ext_row_ptr);
Self::extend_ram_at_address(&mut ram, self.challenges.challenges, challenges_ptr);
let non_determinism = NonDeterminism::default().with_ram(ram);

VMState::new(program, PublicInput::default(), non_determinism)
}

fn set_up_triton_vm_to_evaluate_constraints_in_tasm_dynamic(
&self,
program: &Program,
) -> VMState {
// for convenience, reuse the (integral) static memory layout
let mut vm_state =
self.set_up_triton_vm_to_evaluate_constraints_in_tasm_static(program);
vm_state
.op_stack
.push(self.static_memory_layout.curr_base_row_ptr);
vm_state
.op_stack
.push(self.static_memory_layout.curr_ext_row_ptr);
vm_state
.op_stack
.push(self.static_memory_layout.next_base_row_ptr);
vm_state
.op_stack
.push(self.static_memory_layout.next_ext_row_ptr);
vm_state
}

fn extend_ram_at_address(
ram: &mut HashMap<BFieldElement, BFieldElement>,
list: impl IntoIterator<Item = impl Into<XFieldElement>>,
address: BFieldElement,
) {
let list = list.into_iter().flat_map(|xfe| xfe.into().coefficients);
let indexed_list = list.enumerate();
let offset_address = |i| bfe!(i as u64) + address;
let ram_extension = indexed_list.map(|(i, bfe)| (offset_address(i), bfe));
ram.extend(ram_extension);
}

fn read_xfe_list_at_address(
ram: HashMap<BFieldElement, BFieldElement>,
address: u64,
len: usize,
) -> Vec<XFieldElement> {
let mem_region_end = address + (len * EXTENSION_DEGREE) as u64;
(address..mem_region_end)
.map(BFieldElement::new)
.map(|i| ram[&i])
.chunks(EXTENSION_DEGREE)
.into_iter()
.map(|c| XFieldElement::try_from(c.collect_vec()).unwrap())
.collect()
}
}

#[proptest]
fn triton_constraints_and_assembly_constraints_agree(point: ConstraintEvaluationPoint) {
let all_constraints_rust = point.evaluate_all_constraints_rust();
let all_constraints_tasm_static = point.evaluate_all_constraints_tasm_static();
prop_assert_eq!(all_constraints_rust.clone(), all_constraints_tasm_static);

let all_constraints_tasm_dynamic = point.evaluate_all_constraints_tasm_dynamic();
prop_assert_eq!(all_constraints_rust, all_constraints_tasm_dynamic);
}

#[proptest]
fn triton_assembly_constraint_evaluators_do_not_write_outside_of_dedicated_memory_region(
point: ConstraintEvaluationPoint,
) {
let program = point.tasm_static_constraint_evaluation_code();
let mut initial_state =
point.set_up_triton_vm_to_evaluate_constraints_in_tasm_static(&program);
let mut terminal_state = initial_state.clone();
terminal_state.run().unwrap();

let free_mem_page_ptr = point.static_memory_layout.free_mem_page_ptr;
let mem_page_size = StaticTasmConstraintEvaluationMemoryLayout::MEM_PAGE_SIZE;
let mem_page = MemoryRegion::new(free_mem_page_ptr, mem_page_size);
let not_in_mem_page = |addr: &_| !mem_page.contains_address(addr);

initial_state.ram.retain(|k, _| not_in_mem_page(k));
terminal_state.ram.retain(|k, _| not_in_mem_page(k));
prop_assert_eq!(initial_state.ram, terminal_state.ram);
}

#[proptest]
fn triton_assembly_constraint_evaluators_declare_no_labels(
#[strategy(arb())] static_memory_layout: StaticTasmConstraintEvaluationMemoryLayout,
#[strategy(arb())] dynamic_memory_layout: DynamicTasmConstraintEvaluationMemoryLayout,
) {
for instruction in static_air_constraint_evaluation_tasm(static_memory_layout)
.into_iter()
.chain(dynamic_air_constraint_evaluation_tasm(
dynamic_memory_layout,
))
{
if let LabelledInstruction::Label(label) = instruction {
return Err(TestCaseError::Fail(format!("Found label: {label}").into()));
}
}
}

#[proptest]
fn triton_assembly_constraint_evaluators_are_straight_line_and_does_not_halt(
#[strategy(arb())] static_memory_layout: StaticTasmConstraintEvaluationMemoryLayout,
#[strategy(arb())] dynamic_memory_layout: DynamicTasmConstraintEvaluationMemoryLayout,
) {
type I = AnInstruction<String>;
let is_legal = |i| !matches!(i, I::Call(_) | I::Return | I::Recurse | I::Skiz | I::Halt);

for instruction in static_air_constraint_evaluation_tasm(static_memory_layout) {
if let LabelledInstruction::Instruction(instruction) = instruction {
prop_assert!(is_legal(instruction));
}
}

for instruction in dynamic_air_constraint_evaluation_tasm(dynamic_memory_layout) {
if let LabelledInstruction::Instruction(instruction) = instruction {
prop_assert!(is_legal(instruction));
}
}
}
}
Loading

0 comments on commit 14c270c

Please sign in to comment.