Skip to content

Commit

Permalink
refactor: Rewrite add_u128 to implement modern traits
Browse files Browse the repository at this point in the history 
Todo for this snippet: Add negative tests!
  • Loading branch information
@Sword-Smith
Sword-Smith committed Jul 10, 2024
1 parent 990acea commit 1e3bc49
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 157 deletions.
4 changes: 2 additions & 2 deletions tasm-lib/benchmarks/tasmlib_arithmetic_u128_add.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
"benchmark_result": {
"clock_cycle_count": 29,
"hash_table_height": 30,
"u32_table_height": 67,
"u32_table_height": 125,
"op_stack_table_height": 16,
"ram_table_height": 0
},
Expand All @@ -15,7 +15,7 @@
"benchmark_result": {
"clock_cycle_count": 29,
"hash_table_height": 30,
"u32_table_height": 67,
"u32_table_height": 127,
"op_stack_table_height": 16,
"ram_table_height": 0
},
Expand Down
256 changes: 101 additions & 155 deletions tasm-lib/src/arithmetic/u128/add_u128.rs
View file
Original file line number Diff line number Diff line change
@@ -1,65 +1,34 @@
use std::collections::HashMap;

use num::Zero;
use rand::prelude::*;
use triton_vm::prelude::*;
use triton_vm::twenty_first::prelude::U32s;

use crate::data_type::DataType;
use crate::empty_stack;
use crate::library::Library;
use crate::push_encodable;
use crate::traits::deprecated_snippet::DeprecatedSnippet;
use crate::InitVmState;
use crate::traits::basic_snippet::BasicSnippet;

#[derive(Clone, Debug)]
pub struct AddU128;

impl DeprecatedSnippet for AddU128 {
fn entrypoint_name(&self) -> String {
impl BasicSnippet for AddU128 {
fn entrypoint(&self) -> String {
"tasmlib_arithmetic_u128_add".to_string()
}

fn input_field_names(&self) -> Vec<String> {
vec![
"rhs_3".to_string(),
"rhs_2".to_string(),
"rhs_1".to_string(),
"rhs_0".to_string(),
"lhs_3".to_string(),
"lhs_2".to_string(),
"lhs_1".to_string(),
"lhs_0".to_string(),
]
}

fn input_types(&self) -> Vec<DataType> {
vec![DataType::U128, DataType::U128]
}

fn output_field_names(&self) -> Vec<String> {
fn inputs(&self) -> Vec<(DataType, String)> {
vec![
"(lhs + rhs)_3".to_string(),
"(lhs + rhs)_2".to_string(),
"(lhs + rhs)_1".to_string(),
"(lhs + rhs)_0".to_string(),
(DataType::U128, "lhs".to_owned()),
(DataType::U128, "rhs".to_owned()),
]
}

fn output_types(&self) -> Vec<DataType> {
vec![DataType::U128]
}

fn stack_diff(&self) -> isize {
-4
fn outputs(&self) -> Vec<(DataType, String)> {
vec![(DataType::U128, "sum".to_owned())]
}

/// Four top elements of stack are assumed to be valid u32s. So to have
/// a value that's less than 2^32.
fn function_code(&self, _library: &mut Library) -> String {
let entrypoint = self.entrypoint_name();
format!(
"
fn code(&self, _library: &mut Library) -> Vec<LabelledInstruction> {
let entrypoint = self.entrypoint();

triton_asm!(
// BEFORE: _ rhs_3 rhs_2 rhs_1 rhs_0 lhs_3 lhs_2 lhs_1 lhs_0
// AFTER: _ sum_3 sum_2 sum_1 sum_0
{entrypoint}:
Expand Down Expand Up @@ -128,150 +97,127 @@ impl DeprecatedSnippet for AddU128 {
// _ sum_3 sum_2 sum_1 sum_0

return
"
)
}

fn crash_conditions(&self) -> Vec<String> {
vec!["if (lhs + rhs) overflows u128".to_string()]
}

fn gen_input_states(&self) -> Vec<InitVmState> {
let mut rng = rand::thread_rng();

let mut states = vec![];
let zero = U32s::<4>::zero();

for _ in 0..20 {
let small_a = U32s::<4>::try_from(rng.gen::<u64>()).unwrap();
let small_b = U32s::<4>::try_from(rng.gen::<u64>()).unwrap();
let mut random_bytes: [u32; 4] = [0, 0, 0, 0];
rng.fill(&mut random_bytes);
let large_a = U32s::<4>::new(random_bytes);

// 0. one zero, one large
states.push({
let mut stack = empty_stack();
push_encodable(&mut stack, &zero);
push_encodable(&mut stack, &large_a);
InitVmState::with_stack(stack)
});

// 1. two small
states.push({
let mut stack = empty_stack();
push_encodable(&mut stack, &small_a);
push_encodable(&mut stack, &small_b);
InitVmState::with_stack(stack)
});
}

states
}

fn common_case_input_state(&self) -> InitVmState {
InitVmState::with_stack(
[
empty_stack(),
vec![BFieldElement::zero(), BFieldElement::new(1 << 31)],
vec![BFieldElement::zero(), BFieldElement::new(1 << 30)],
vec![BFieldElement::zero(), BFieldElement::new(1 << 30)],
vec![BFieldElement::zero(), BFieldElement::new(1 << 30)],
]
.concat(),
)
}

fn worst_case_input_state(&self) -> InitVmState {
InitVmState::with_stack(
[
empty_stack(),
vec![BFieldElement::zero(), BFieldElement::new(1 << 31)],
vec![BFieldElement::zero(), BFieldElement::new(1 << 30)],
vec![BFieldElement::zero(), BFieldElement::new(1 << 30)],
vec![BFieldElement::zero(), BFieldElement::new(1 << 30)],
]
.concat(),
)
}

fn rust_shadowing(
&self,
stack: &mut Vec<BFieldElement>,
_std_in: Vec<BFieldElement>,
_secret_in: Vec<BFieldElement>,
_memory: &mut HashMap<BFieldElement, BFieldElement>,
) {
// top element on stack
let a0: u32 = stack.pop().unwrap().try_into().unwrap();
let b0: u32 = stack.pop().unwrap().try_into().unwrap();
let c0: u32 = stack.pop().unwrap().try_into().unwrap();
let d0: u32 = stack.pop().unwrap().try_into().unwrap();
let ab0 = U32s::<4>::new([a0, b0, c0, d0]);

// second element on stack
let a1: u32 = stack.pop().unwrap().try_into().unwrap();
let b1: u32 = stack.pop().unwrap().try_into().unwrap();
let c1: u32 = stack.pop().unwrap().try_into().unwrap();
let d1: u32 = stack.pop().unwrap().try_into().unwrap();
let ab1 = U32s::<4>::new([a1, b1, c1, d1]);
let ab0_plus_ab1 = ab0 + ab1;
let mut res = ab0_plus_ab1.encode();
for _ in 0..res.len() {
stack.push(res.pop().unwrap());
}
}
}

#[cfg(test)]
mod tests {
use crate::test_helpers::test_rust_equivalence_given_input_values_deprecated;
use crate::test_helpers::test_rust_equivalence_multiple_deprecated;
use itertools::Itertools;
use rand::rngs::StdRng;
use rand::Rng;
use rand::SeedableRng;

use crate::snippet_bencher::BenchmarkCase;
use crate::test_helpers::test_rust_equivalence_given_complete_state;
use crate::traits::closure::Closure;
use crate::traits::closure::ShadowedClosure;
use crate::traits::rust_shadow::RustShadow;

use super::*;

#[test]
fn add_u128_test() {
test_rust_equivalence_multiple_deprecated(&AddU128, true);
ShadowedClosure::new(AddU128).test()
}

#[test]
fn add_u128_unit_test() {
let mut expected = empty_stack();
let snippet = AddU128;
let mut expected = snippet.init_stack_for_isolated_run();
expected.push(BFieldElement::new(0));
expected.push(BFieldElement::new(1 << 4));
expected.push(BFieldElement::new(0));
expected.push(BFieldElement::new(0));
prop_add(1u128 << 67, 1u128 << 67, Some(&expected))
}
snippet.prop_add(1u128 << 67, 1u128 << 67, Some(&expected))
}

impl AddU128 {
fn prop_add(&self, lhs: u128, rhs: u128, expected: Option<&[BFieldElement]>) {
let mut init_stack = self.init_stack_for_isolated_run();
for elem in rhs.encode().into_iter().rev() {
init_stack.push(elem);
}
for elem in lhs.encode().into_iter().rev() {
init_stack.push(elem);
}

test_rust_equivalence_given_complete_state(
&ShadowedClosure::new(AddU128),
&init_stack,
&[],
&NonDeterminism::default(),
&None,
expected,
);
}

fn prop_add(lhs: u128, rhs: u128, expected: Option<&[BFieldElement]>) {
let mut init_stack = empty_stack();
for elem in rhs.encode().into_iter().rev() {
init_stack.push(elem);
fn setup_init_stack(&self, lhs: u128, rhs: u128) -> Vec<BFieldElement> {
[
self.init_stack_for_isolated_run(),
lhs.encode().into_iter().rev().collect_vec(),
rhs.encode().into_iter().rev().collect_vec(),
]
.concat()
}
for elem in lhs.encode().into_iter().rev() {
init_stack.push(elem);
}

impl Closure for AddU128 {
fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
fn to_u128(a: u32, b: u32, c: u32, d: u32) -> u128 {
a as u128
+ b as u128 * (1u128 << 32)
+ c as u128 * (1u128 << 64)
+ d as u128 * (1u128 << 96)
}

// top element on stack
let a0: u32 = stack.pop().unwrap().try_into().unwrap();
let b0: u32 = stack.pop().unwrap().try_into().unwrap();
let c0: u32 = stack.pop().unwrap().try_into().unwrap();
let d0: u32 = stack.pop().unwrap().try_into().unwrap();
let ab0 = to_u128(a0, b0, c0, d0);

// second element on stack
let a1: u32 = stack.pop().unwrap().try_into().unwrap();
let b1: u32 = stack.pop().unwrap().try_into().unwrap();
let c1: u32 = stack.pop().unwrap().try_into().unwrap();
let d1: u32 = stack.pop().unwrap().try_into().unwrap();
let ab1 = to_u128(a1, b1, c1, d1);

let ab0_plus_ab1 = ab0.checked_add(ab1).unwrap();

let mut res = ab0_plus_ab1.encode();
for _ in 0..res.len() {
stack.push(res.pop().unwrap());
}
}

test_rust_equivalence_given_input_values_deprecated::<AddU128>(
&AddU128,
&init_stack,
&[],
HashMap::default(),
expected,
);
fn pseudorandom_initial_state(
&self,
seed: [u8; 32],
_bench_case: Option<BenchmarkCase>,
) -> Vec<BFieldElement> {
let mut rng: StdRng = SeedableRng::from_seed(seed);
let lhs: u128 = rng.gen();
let lhs = lhs / 2;
let rhs: u128 = rng.gen();
let rhs = rhs / 2;

self.setup_init_stack(lhs, rhs)
}
}
}

#[cfg(test)]
mod benches {
use crate::snippet_bencher::bench_and_write;
use crate::traits::closure::ShadowedClosure;
use crate::traits::rust_shadow::RustShadow;

use super::*;

#[test]
fn add_u128_benchmark() {
bench_and_write(AddU128);
ShadowedClosure::new(AddU128).bench()
}
}

0 comments on commit 1e3bc49

Please sign in to comment.