Skip to content

Commit da30d1c

Browse files
committed
wip sha256
1 parent 061453c commit da30d1c

24 files changed

Lines changed: 1447 additions & 32 deletions

File tree

Cargo.lock

Lines changed: 13 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ system-info = { path = "crates/backend/system-info" }
6464

6565
# External
6666
sha3 = "0.11.0"
67+
sha2 = "0.11.0"
6768
clap = { version = "4.5.59", features = ["derive"] }
6869
rand = "0.10.0"
6970
rayon = "1.11.0"

crates/lean_compiler/src/a_simplify_lang/mod.rs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2280,6 +2280,32 @@ fn simplify_lines(
22802280
continue;
22812281
}
22822282

2283+
// Special handling for SHA256 compression precompile
2284+
if function_name == Table::sha256_compress().name() {
2285+
if !targets.is_empty() {
2286+
return Err(format!(
2287+
"Precompile {function_name} should not return values, at {location}"
2288+
));
2289+
}
2290+
if args.len() != 3 {
2291+
return Err(format!(
2292+
"Precompile {function_name} expects 3 arguments (state_ptr, block_ptr, out_ptr), got {}, at {location}",
2293+
args.len()
2294+
));
2295+
}
2296+
let simplified_args = args
2297+
.iter()
2298+
.map(|arg| simplify_expr(ctx, state, const_malloc, arg, &mut res))
2299+
.collect::<Result<Vec<_>, _>>()?;
2300+
res.push(SimpleLine::Precompile(PrecompileArgs {
2301+
arg_0: simplified_args[0].clone(),
2302+
arg_1: simplified_args[1].clone(),
2303+
res: simplified_args[2].clone(),
2304+
data: PrecompileCompTimeArgs::Sha256Compress,
2305+
}));
2306+
continue;
2307+
}
2308+
22832309
// Special handling for custom hints
22842310
if let Some(hint) = CustomHint::find_by_name(function_name) {
22852311
if !targets.is_empty() {

crates/lean_compiler/src/instruction_encoder.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ pub fn field_representation(instr: &Instruction) -> [F; N_INSTRUCTION_COLUMNS] {
4949
Instruction::Precompile(precompile) => {
5050
let precompile_data = match &precompile.data {
5151
PrecompileCompTimeArgs::Poseidon16 => POSEIDON_PRECOMPILE_DATA,
52+
PrecompileCompTimeArgs::Sha256Compress => SHA256_PRECOMPILE_DATA,
5253
PrecompileCompTimeArgs::ExtensionOp { size, mode } => {
5354
assert!(*size >= 1, "invalid extension_op size={size}");
5455
mode.flag_encoding() + EXT_OP_LEN_MULTIPLIER * size

crates/lean_compiler/src/parser/parsers/function.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::{
88
grammar::{ParsePair, Rule},
99
},
1010
};
11-
use lean_vm::{CUSTOM_HINTS, ExtensionOpMode, POSEIDON16_NAME};
11+
use lean_vm::{CUSTOM_HINTS, ExtensionOpMode, POSEIDON16_NAME, SHA256_COMPRESS_NAME};
1212

1313
/// Reserved function names that users cannot define.
1414
pub const RESERVED_FUNCTION_NAMES: &[&str] = &[
@@ -33,8 +33,8 @@ fn is_reserved_function_name(name: &str) -> bool {
3333
if RESERVED_FUNCTION_NAMES.contains(&name) || CUSTOM_HINTS.iter().any(|hint| hint.name() == name) {
3434
return true;
3535
}
36-
// Check precompile names (poseidon16, extension_op functions)
37-
if name == POSEIDON16_NAME {
36+
// Check precompile names (poseidon16, sha256, extension_op functions)
37+
if name == POSEIDON16_NAME || name == SHA256_COMPRESS_NAME {
3838
return true;
3939
}
4040
if ExtensionOpMode::from_name(name).is_some() {

crates/lean_compiler/tests/test_compiler.rs

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::time::Instant;
22

3-
use backend::BasedVectorSpace;
3+
use backend::{BasedVectorSpace, PrimeCharacteristicRing};
44
use lean_compiler::*;
55
use lean_vm::*;
66
use rand::{RngExt, SeedableRng, rngs::StdRng};
@@ -26,6 +26,30 @@ def main():
2626
let _ = dbg!(poseidon16_compress(public_input));
2727
}
2828

29+
#[test]
30+
fn test_sha256_compress() {
31+
let program = r#"
32+
def main():
33+
state = 0
34+
block = 16
35+
expected = 48
36+
out = Array(16)
37+
sha256_compress(state, block, out)
38+
39+
for i in unroll(0, 16):
40+
assert out[i] == expected[i]
41+
return
42+
"#;
43+
44+
let mut public_input = vec![F::ZERO; 64];
45+
public_input[0..16].copy_from_slice(&words_to_field_limbs_le(SHA256_IV));
46+
public_input[16..48].copy_from_slice(&words_to_field_limbs_le(SHA256_ABC_BLOCK));
47+
let expected = words_to_field_limbs_le(sha256_compress_words(SHA256_IV, SHA256_ABC_BLOCK));
48+
public_input[48..64].copy_from_slice(&expected);
49+
50+
compile_and_run(&ProgramSource::Raw(program.to_string()), &public_input, false);
51+
}
52+
2953
#[test]
3054
fn test_div_extension_field() {
3155
let program = r#"

crates/lean_prover/src/test_zkvm.rs

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,114 @@ use lean_vm::*;
55
use rand::{RngExt, SeedableRng, rngs::StdRng};
66
use utils::{init_tracing, poseidon16_compress};
77

8+
#[test]
9+
#[ignore = "benchmark; run with `cargo test --release -p lean_prover bench_poseidon -- --ignored --nocapture`"]
10+
fn bench_poseidon() {
11+
utils::init_tracing();
12+
let n_poseidon_calls = std::env::var("POSEIDON_BENCH_CALLS")
13+
.ok()
14+
.map(|raw| raw.parse::<usize>().expect("POSEIDON_BENCH_CALLS must be a usize"))
15+
.unwrap_or(1);
16+
let program_str = format!(
17+
r#"
18+
N_POSEIDON_CALLS = {n_poseidon_calls}
19+
DIGEST_LEN = 8
20+
21+
def main():
22+
input_left = 0
23+
input_right = DIGEST_LEN
24+
outputs = Array(N_POSEIDON_CALLS * DIGEST_LEN)
25+
for i in dynamic_unroll(0, N_POSEIDON_CALLS, 20):
26+
out = outputs + i * DIGEST_LEN
27+
poseidon16_compress(input_left, input_right, out)
28+
return
29+
"#
30+
);
31+
32+
let public_input: Vec<F> = (0..16).map(F::new).collect();
33+
let bytecode = compile_program(&ProgramSource::Raw(program_str));
34+
let witness = ExecutionWitness::default();
35+
let starting_log_inv_rate = 1;
36+
37+
let time = std::time::Instant::now();
38+
let proof = prove_execution(
39+
&bytecode,
40+
&public_input,
41+
&witness,
42+
&default_whir_config(starting_log_inv_rate),
43+
false,
44+
);
45+
let proof_time = time.elapsed();
46+
let proof_size_kib = proof.proof.proof_size_fe() * F::bits() / (8 * 1024);
47+
48+
println!("{}", proof.metadata.display());
49+
println!("Proof time: {:.3} s", proof_time.as_secs_f32());
50+
println!("Proof size: {proof_size_kib} KiB");
51+
52+
verify_execution(&bytecode, &public_input, proof.proof).unwrap();
53+
}
54+
55+
#[test]
56+
#[ignore = "benchmark; run with `cargo test --release -p lean_prover bench_sha256_compress -- --ignored --nocapture`"]
57+
fn bench_sha256_compress() {
58+
utils::init_tracing();
59+
let n_sha_calls = std::env::var("SHA256_BENCH_CALLS")
60+
.ok()
61+
.map(|raw| raw.parse::<usize>().expect("SHA256_BENCH_CALLS must be a usize"))
62+
.unwrap_or(1);
63+
const SHA_FIXTURE_STRIDE: usize = SHA256_STATE_LIMBS + SHA256_BLOCK_LIMBS + SHA256_STATE_LIMBS;
64+
let program_str = format!(
65+
r#"
66+
N_SHA_CALLS = {n_sha_calls}
67+
SHA_FIXTURE_STRIDE = 64
68+
69+
def main():
70+
for j in unroll(0, N_SHA_CALLS):
71+
base = j * SHA_FIXTURE_STRIDE
72+
state = base
73+
block = base + 16
74+
expected = base + 48
75+
out = Array(16)
76+
sha256_compress(state, block, out)
77+
78+
for i in unroll(0, 16):
79+
assert out[i] == expected[i]
80+
return
81+
"#
82+
);
83+
84+
let mut public_input = vec![F::ZERO; n_sha_calls * SHA_FIXTURE_STRIDE];
85+
let expected = words_to_field_limbs_le(sha256_compress_words(SHA256_IV, SHA256_ABC_BLOCK));
86+
for j in 0..n_sha_calls {
87+
let base = j * SHA_FIXTURE_STRIDE;
88+
public_input[base..base + SHA256_STATE_LIMBS].copy_from_slice(&words_to_field_limbs_le(SHA256_IV));
89+
public_input[base + 16..base + 16 + SHA256_BLOCK_LIMBS]
90+
.copy_from_slice(&words_to_field_limbs_le(SHA256_ABC_BLOCK));
91+
public_input[base + 48..base + 48 + SHA256_STATE_LIMBS].copy_from_slice(&expected);
92+
}
93+
94+
let bytecode = compile_program(&ProgramSource::Raw(program_str));
95+
let witness = ExecutionWitness::default();
96+
let starting_log_inv_rate = 1;
97+
98+
let time = std::time::Instant::now();
99+
let proof = prove_execution(
100+
&bytecode,
101+
&public_input,
102+
&witness,
103+
&default_whir_config(starting_log_inv_rate),
104+
false,
105+
);
106+
let proof_time = time.elapsed();
107+
let proof_size_kib = proof.proof.proof_size_fe() * F::bits() / (8 * 1024);
108+
109+
println!("{}", proof.metadata.display());
110+
println!("Proof time: {:.3} s", proof_time.as_secs_f32());
111+
println!("Proof size: {proof_size_kib} KiB");
112+
113+
verify_execution(&bytecode, &public_input, proof.proof).unwrap();
114+
}
115+
8116
#[test]
9117
fn test_zk_vm_all_precompiles() {
10118
let program_str = r#"
@@ -17,6 +125,15 @@ def main():
17125
pub_start = 0
18126
poseidon16_compress(pub_start + 4 * DIGEST_LEN, pub_start + 5 * DIGEST_LEN, pub_start + 6 * DIGEST_LEN)
19127
128+
# Keep the SHA fixture away from the extension-op fixture ranges below.
129+
sha_state = pub_start + 1400
130+
sha_block = sha_state + 16
131+
sha_expected = sha_block + 32
132+
sha_out = Array(16)
133+
sha256_compress(sha_state, sha_block, sha_out)
134+
for i in unroll(0, 16):
135+
assert sha_out[i] == sha_expected[i]
136+
20137
base_ptr = pub_start + 88
21138
ext_a_ptr = pub_start + 88 + N
22139
ext_b_ptr = pub_start + 88 + N * (DIM + 1)
@@ -62,6 +179,19 @@ def main():
62179
let poseidon_24_input: [F; 24] = rng.random();
63180
public_input[56..80].copy_from_slice(&poseidon_24_input);
64181

182+
// SHA256 compression test data: IV + padded "abc" block.
183+
// This mirrors the program's pub_start + 1400 offset; public_input is 2^13 cells,
184+
// so the state, block, and expected digest all fit in the public memory prefix.
185+
let sha_state_ptr = 1400;
186+
let sha_block_ptr = sha_state_ptr + SHA256_STATE_LIMBS;
187+
let sha_expected_ptr = sha_block_ptr + SHA256_BLOCK_LIMBS;
188+
public_input[sha_state_ptr..sha_state_ptr + SHA256_STATE_LIMBS]
189+
.copy_from_slice(&words_to_field_limbs_le(SHA256_IV));
190+
public_input[sha_block_ptr..sha_block_ptr + SHA256_BLOCK_LIMBS]
191+
.copy_from_slice(&words_to_field_limbs_le(SHA256_ABC_BLOCK));
192+
let sha_expected = words_to_field_limbs_le(sha256_compress_words(SHA256_IV, SHA256_ABC_BLOCK));
193+
public_input[sha_expected_ptr..sha_expected_ptr + SHA256_STATE_LIMBS].copy_from_slice(&sha_expected);
194+
65195
// Extension op operands: base[N], ext_a[N], ext_b[N]
66196
let base_slice: [F; N] = rng.random();
67197
let ext_a_slice: [EF; N] = rng.random();

crates/lean_prover/src/trace_gen.rs

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,24 @@ pub fn get_execution_trace(bytecode: &Bytecode, execution_result: ExecutionResul
9999
let null_poseidon_16_hash_ptr = memory_padded.len();
100100
memory_padded.extend_from_slice(get_poseidon_16_of_zero());
101101

102+
let sha256_padding_state_ptr = memory_padded.len();
103+
memory_padded.extend(words_to_field_limbs_le(SHA256_IV));
104+
let sha256_padding_block_ptr = memory_padded.len();
105+
memory_padded.extend(words_to_field_limbs_le(SHA256_ZERO_BLOCK));
106+
let sha256_padding_out_ptr = memory_padded.len();
107+
memory_padded.extend(words_to_field_limbs_le(sha256_compress_words(
108+
SHA256_IV,
109+
SHA256_ZERO_BLOCK,
110+
)));
111+
112+
let padding_memory = PaddingMemory {
113+
zero_vec_ptr: padding_zero_vec_ptr,
114+
null_poseidon_16_hash_ptr,
115+
sha256_state_ptr: sha256_padding_state_ptr,
116+
sha256_block_ptr: sha256_padding_block_ptr,
117+
sha256_out_ptr: sha256_padding_out_ptr,
118+
};
119+
102120
// IMPORTANT: memory size should always be >= number of VM cycles
103121
let padded_memory_len = (memory_padded.len().max(n_cycles).max(1 << MIN_LOG_N_ROWS_PER_TABLE)).next_power_of_two();
104122
memory_padded.resize(padded_memory_len, F::ZERO);
@@ -120,7 +138,7 @@ pub fn get_execution_trace(bytecode: &Bytecode, execution_result: ExecutionResul
120138
},
121139
);
122140
for table in traces.keys().copied().collect::<Vec<_>>() {
123-
pad_table(&table, &mut traces, padding_zero_vec_ptr, null_poseidon_16_hash_ptr);
141+
pad_table(&table, &mut traces, &padding_memory);
124142
}
125143

126144
ExecutionTrace {
@@ -131,12 +149,7 @@ pub fn get_execution_trace(bytecode: &Bytecode, execution_result: ExecutionResul
131149
}
132150
}
133151

134-
fn pad_table(
135-
table: &Table,
136-
traces: &mut BTreeMap<Table, TableTrace>,
137-
zero_vec_ptr: usize,
138-
null_poseidon_16_hash_ptr: usize,
139-
) {
152+
fn pad_table(table: &Table, traces: &mut BTreeMap<Table, TableTrace>, padding_memory: &PaddingMemory) {
140153
let trace = traces.get_mut(table).unwrap();
141154
let h = trace.columns[0].len();
142155
trace
@@ -148,7 +161,7 @@ fn pad_table(
148161
trace.non_padded_n_rows = h;
149162
trace.log_n_rows = log2_ceil_usize(h + 1).max(MIN_LOG_N_ROWS_PER_TABLE);
150163
let n_rows = 1 << trace.log_n_rows;
151-
let padding_row = table.padding_row(zero_vec_ptr, null_poseidon_16_hash_ptr);
164+
let padding_row = table.padding_row(padding_memory);
152165
trace.columns.par_iter_mut().enumerate().for_each(|(i, col)| {
153166
assert!(col.len() <= h); // potentially some columns have not been filled (in Poseidon -> we fill it later with SIMD + parallelism), but the first one should always be representative
154167
col.resize(n_rows, padding_row[i]);

crates/lean_vm/Cargo.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,6 @@ rand.workspace = true
1515
tracing.workspace = true
1616
backend.workspace = true
1717
itertools.workspace = true
18+
19+
[dev-dependencies]
20+
sha2.workspace = true

crates/lean_vm/src/core/constants.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,13 @@ pub const MIN_BYTECODE_LOG_SIZE: usize = 8;
2121

2222
/// Minimum and maximum number of rows per table (as powers of two), both inclusive
2323
pub const MIN_LOG_N_ROWS_PER_TABLE: usize = 8; // Zero padding will be added to each at least, if this minimum is not reached, (ensuring AIR / GKR work fine, with SIMD, without too much edge cases). Long term, we should find a more elegant solution.
24-
pub const MAX_LOG_N_ROWS_PER_TABLE: [(Table, usize); 3] = [
24+
pub const MAX_LOG_N_ROWS_PER_TABLE: [(Table, usize); 4] = [
2525
(Table::execution(), 25),
2626
(Table::extension_op(), 20),
2727
(Table::poseidon16(), 21),
28+
// Direct Plonky3-style SHA256 has 7524 columns. 2^13 rows already exceeds
29+
// the current commitment-surface guard; 2^12 is the largest safe cap today.
30+
(Table::sha256_compress(), 12),
2831
];
2932

3033
/// Starting program counter

0 commit comments

Comments
 (0)