diff --git a/.gitignore b/.gitignore index e6fde4272..0f3dffb54 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ bins *.s *.txt *.codex +picus_out # Proofs **/proof-with-pis.bin diff --git a/Cargo.lock b/Cargo.lock index 63ce05462..c2dc05cf2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6136,7 +6136,7 @@ dependencies = [ [[package]] name = "test-artifacts" -version = "1.2.5" +version = "1.2.6" dependencies = [ "zkm-build", ] @@ -7718,7 +7718,7 @@ dependencies = [ [[package]] name = "zkm-build" -version = "1.2.5" +version = "1.2.6" dependencies = [ "anyhow", "cargo_metadata", @@ -7728,7 +7728,7 @@ dependencies = [ [[package]] name = "zkm-cli" -version = "1.2.5" +version = "1.2.6" dependencies = [ "anyhow", "cargo_metadata", @@ -7741,7 +7741,7 @@ dependencies = [ [[package]] name = "zkm-core-executor" -version = "1.2.5" +version = "1.2.6" dependencies = [ "anyhow", "bincode", @@ -7777,13 +7777,13 @@ dependencies = [ "vec_map", "zkm-curves", "zkm-instruction-test-defs", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-stark", ] [[package]] name = "zkm-core-machine" -version = "1.2.5" +version = "1.2.6" dependencies = [ "bincode", "cbindgen", @@ -7834,13 +7834,13 @@ dependencies = [ "zkm-curves", "zkm-derive", "zkm-instruction-test-defs", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-stark", ] [[package]] name = "zkm-cuda" -version = "1.2.5" +version = "1.2.6" dependencies = [ "bincode", "ctrlc", @@ -7857,7 +7857,7 @@ dependencies = [ [[package]] name = "zkm-curves" -version = "1.2.5" +version = "1.2.6" dependencies = [ "cfg-if", "curve25519-dalek", @@ -7876,13 +7876,13 @@ dependencies = [ "thiserror 1.0.69", "tracing", "typenum", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-stark", ] [[package]] name = "zkm-derive" -version = "1.2.5" +version = "1.2.6" dependencies = [ "proc-macro2", "quote", @@ -7891,7 +7891,7 @@ dependencies = [ [[package]] name = "zkm-instruction-test-defs" -version = "1.2.5" +version = "1.2.6" dependencies = [ "zkm-core-executor", ] @@ -7911,19 +7911,19 @@ dependencies = [ [[package]] name = "zkm-lib" -version = "1.2.5" +version = "1.2.6" dependencies = [ "bincode", "cfg-if", "elliptic-curve", "serde", "sha2", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", ] [[package]] name = "zkm-picus" -version = "1.2.5" +version = "1.2.6" dependencies = [ "clap", "p3-air", @@ -7955,7 +7955,7 @@ dependencies = [ [[package]] name = "zkm-primitives" -version = "1.2.5" +version = "1.2.6" dependencies = [ "bincode", "hex", @@ -7972,7 +7972,7 @@ dependencies = [ [[package]] name = "zkm-prover" -version = "1.2.5" +version = "1.2.6" dependencies = [ "anyhow", "bincode", @@ -8000,7 +8000,7 @@ dependencies = [ "tracing-subscriber 0.3.23", "zkm-core-executor", "zkm-core-machine", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-recursion-circuit", "zkm-recursion-compiler", "zkm-recursion-core", @@ -8010,7 +8010,7 @@ dependencies = [ [[package]] name = "zkm-recursion-circuit" -version = "1.2.5" +version = "1.2.6" dependencies = [ "ff 0.13.1", "hashbrown 0.14.5", @@ -8038,7 +8038,7 @@ dependencies = [ "zkm-core-executor", "zkm-core-machine", "zkm-derive", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-recursion-compiler", "zkm-recursion-core", "zkm-recursion-gnark-ffi", @@ -8047,7 +8047,7 @@ dependencies = [ [[package]] name = "zkm-recursion-compiler" -version = "1.2.5" +version = "1.2.6" dependencies = [ "backtrace", "itertools 0.13.0", @@ -8063,7 +8063,7 @@ dependencies = [ "tracing", "vec_map", "zkm-core-machine", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-recursion-core", "zkm-recursion-derive", "zkm-stark", @@ -8071,7 +8071,7 @@ dependencies = [ [[package]] name = "zkm-recursion-core" -version = "1.2.5" +version = "1.2.6" dependencies = [ "backtrace", "cbindgen", @@ -8107,13 +8107,13 @@ dependencies = [ "zkhash", "zkm-core-machine", "zkm-derive", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-stark", ] [[package]] name = "zkm-recursion-derive" -version = "1.2.5" +version = "1.2.6" dependencies = [ "quote", "syn 1.0.109", @@ -8121,7 +8121,7 @@ dependencies = [ [[package]] name = "zkm-recursion-gnark-ffi" -version = "1.2.5" +version = "1.2.6" dependencies = [ "anyhow", "bincode", @@ -8144,7 +8144,7 @@ dependencies = [ [[package]] name = "zkm-sdk" -version = "1.2.5" +version = "1.2.6" dependencies = [ "alloy-primitives", "alloy-signer", @@ -8184,14 +8184,14 @@ dependencies = [ "zkm-core-executor", "zkm-core-machine", "zkm-cuda", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-prover", "zkm-stark", ] [[package]] name = "zkm-stark" -version = "1.2.5" +version = "1.2.6" dependencies = [ "arrayref", "hashbrown 0.14.5", @@ -8226,13 +8226,13 @@ dependencies = [ "tracing-forest", "tracing-subscriber 0.3.23", "zkm-derive", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-zkvm", ] [[package]] name = "zkm-verifier" -version = "1.2.5" +version = "1.2.6" dependencies = [ "anyhow", "ark-bn254", @@ -8262,7 +8262,7 @@ dependencies = [ "thiserror 2.0.18", "zkm-core-executor", "zkm-core-machine", - "zkm-primitives 1.2.5", + "zkm-primitives 1.2.6", "zkm-prover", "zkm-recursion-core", "zkm-sdk", @@ -8271,7 +8271,7 @@ dependencies = [ [[package]] name = "zkm-zkvm" -version = "1.2.5" +version = "1.2.6" dependencies = [ "bincode", "cfg-if", @@ -8285,8 +8285,8 @@ dependencies = [ "rand 0.8.5", "serde", "sha2", - "zkm-lib 1.2.5", - "zkm-primitives 1.2.5", + "zkm-lib 1.2.6", + "zkm-primitives 1.2.6", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index a92151346..38a4917fd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace.package] -version = "1.2.5" +version = "1.2.6" edition = "2021" license = "MIT OR Apache-2.0" rust-version = "1.80" diff --git a/crates/core/executor/src/artifacts/mips_costs.json b/crates/core/executor/src/artifacts/mips_costs.json index 889a4894f..98aee564e 100644 --- a/crates/core/executor/src/artifacts/mips_costs.json +++ b/crates/core/executor/src/artifacts/mips_costs.json @@ -4,11 +4,11 @@ "CloClz": 41, "Bls12831Fp2AddSubAssign": 2070, "SyscallInstrs": 97, - "DivRem": 162, + "DivRem": 170, "ShiftRight": 131, "Secp256r1Decompress": 2686, "Secp256k1Decompress": 2686, - "KeccakSponge": 102216, + "KeccakSponge": 112104, "Bn254AddAssign": 4013, "Bitwise": 42, "ShiftLeft": 68, @@ -16,11 +16,11 @@ "Program": 31, "Global": 115, "Secp256k1AddAssign": 4013, - "BooleanCircuitGarble": 588, + "BooleanCircuitGarble": 456, "AddSub": 47, "Jump": 82, "Bn254FpOpAssign": 704, - "Poseidon2Permute": 1117, + "Poseidon2Permute": 1245, "Mul": 110, "ShaExtend": 15936, "Bls12381AddAssign": 6045, diff --git a/crates/core/machine/Cargo.toml b/crates/core/machine/Cargo.toml index 8f9204b55..2a601f889 100644 --- a/crates/core/machine/Cargo.toml +++ b/crates/core/machine/Cargo.toml @@ -81,6 +81,7 @@ glob = { version = "0.3.1", optional = true } [features] default = [] +picus = ["zkm-stark/picus"] debug = [] bigint-rug = ["zkm-curves/bigint-rug", "zkm-core-executor/bigint-rug"] pre-alloc = ["zkm-core-executor/pre-alloc"] diff --git a/crates/core/machine/build.rs b/crates/core/machine/build.rs index d5d94af8e..e075da60e 100644 --- a/crates/core/machine/build.rs +++ b/crates/core/machine/build.rs @@ -11,7 +11,7 @@ fn main() { #[cfg(feature = "sys")] mod sys { use std::{ - env, fs, os, + env, fs, io, os, path::{Path, PathBuf}, }; @@ -31,6 +31,44 @@ mod sys { const AUTOGEN_WARNING: &str = "/* Automatically generated by `cbindgen`. Not intended for manual editing. */"; + fn patch_generated_cbindgen_header(header_path: &Path) -> io::Result<()> { + if !header_path.exists() { + return Ok(()); + } + + let original = fs::read_to_string(header_path)?; + if original.contains("constexpr static const uintptr_t U64_LIMBS = 4;") { + return Ok(()); + } + + let mut inserted = false; + let mut out: Vec = Vec::new(); + for line in original.lines() { + if line.contains("constexpr static const uintptr_t U64_LIMBS") { + continue; + } + + out.push(line.to_owned()); + if !inserted && line.contains("namespace zkm_core_machine_sys") { + out.push(String::new()); + out.push("constexpr static const uintptr_t U64_LIMBS = 4;".to_owned()); + out.push(String::new()); + inserted = true; + } + } + + if !inserted { + return Ok(()); + } + + let patched = out.join("\n") + "\n"; + if patched != original { + fs::write(header_path, patched)?; + } + + Ok(()) + } + pub fn build_ffi() { // The name of the header generated by `cbindgen`. let cbindgen_hpp = &format!("{LIB_NAME}-cbindgen.hpp"); @@ -171,6 +209,7 @@ mod sys { // Write the bindings to the target include directory. let header_path = target_include_dir.join(cbindgen_hpp); if bindings.write_to_file(&header_path) { + patch_generated_cbindgen_header(&header_path).unwrap(); if let Some(ref target_include_dir_fixed) = target_include_dir_fixed { // Symlink the header to the fixed include directory. rel_symlink_file(header_path, target_include_dir_fixed.join(cbindgen_hpp)); diff --git a/crates/core/machine/src/air/memory.rs b/crates/core/machine/src/air/memory.rs index 8b3f2d1bd..883d5423a 100644 --- a/crates/core/machine/src/air/memory.rs +++ b/crates/core/machine/src/air/memory.rs @@ -4,7 +4,7 @@ use p3_air::AirBuilder; use p3_field::FieldAlgebra; use zkm_core_executor::ByteOpcode; use zkm_stark::{ - air::{AirLookup, BaseAirBuilder, ByteAirBuilder, LookupScope}, + air::{AirLookup, BaseAirBuilder, ByteAirBuilder, LookupScope, OperationSummaryAirBuilder}, LookupKind, }; @@ -22,7 +22,9 @@ pub trait MemoryAirBuilder: BaseAirBuilder { addr: impl Into, memory_access: &impl MemoryCols, do_check: impl Into, - ) { + ) where + Self: OperationSummaryAirBuilder, + { let do_check: Self::Expr = do_check.into(); let shard: Self::Expr = shard.into(); let clk: Self::Expr = clk.into(); @@ -69,7 +71,9 @@ pub trait MemoryAirBuilder: BaseAirBuilder { initial_addr: impl Into + Clone, memory_access_slice: &[impl MemoryCols], verify_memory_access: impl Into + Copy, - ) { + ) where + Self: OperationSummaryAirBuilder, + { for (i, access_slice) in memory_access_slice.iter().enumerate() { self.eval_memory_access( shard, @@ -93,24 +97,40 @@ pub trait MemoryAirBuilder: BaseAirBuilder { do_check: impl Into, shard: impl Into + Clone, clk: impl Into, - ) { + ) where + Self: OperationSummaryAirBuilder, + { let do_check: Self::Expr = do_check.into(); let compare_clk: Self::Expr = mem_access.compare_clk.clone().into(); let shard: Self::Expr = shard.clone().into(); let prev_shard: Self::Expr = mem_access.prev_shard.clone().into(); + let prev_clk: Self::Expr = mem_access.prev_clk.clone().into(); + let clk: Self::Expr = clk.into(); + let diff_16bit_limb: Self::Expr = mem_access.diff_16bit_limb.clone().into(); + let diff_8bit_limb: Self::Expr = mem_access.diff_8bit_limb.clone().into(); + + if self.try_emit_memory_timestamp_summary( + do_check.clone(), + shard.clone(), + clk.clone(), + prev_shard.clone(), + prev_clk.clone(), + compare_clk.clone(), + diff_16bit_limb.clone(), + diff_8bit_limb.clone(), + ) { + return; + } // First verify that compare_clk's value is correct. self.when(do_check.clone()).assert_bool(compare_clk.clone()); self.when(do_check.clone()).when(compare_clk.clone()).assert_eq(shard.clone(), prev_shard); // Get the comparison timestamp values for the current and previous memory access. - let prev_comp_value = self.if_else( - mem_access.compare_clk.clone(), - mem_access.prev_clk.clone(), - mem_access.prev_shard.clone(), - ); + let prev_comp_value = + self.if_else(mem_access.compare_clk.clone(), prev_clk, mem_access.prev_shard.clone()); - let current_comp_val = self.if_else(compare_clk.clone(), clk.into(), shard.clone()); + let current_comp_val = self.if_else(compare_clk.clone(), clk, shard.clone()); // Assert `current_comp_val > prev_comp_val`. We check this by asserting that // `0 <= current_comp_val-prev_comp_val-1 < 2^24`. @@ -125,12 +145,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder { // Verify that mem_access.ts_diff = mem_access.ts_diff_16bit_limb // + mem_access.ts_diff_8bit_limb * 2^16. - self.eval_range_check_24bits( - diff_minus_one, - mem_access.diff_16bit_limb.clone(), - mem_access.diff_8bit_limb.clone(), - do_check, - ); + self.eval_range_check_24bits(diff_minus_one, diff_16bit_limb, diff_8bit_limb, do_check); } /// Verifies the inputted value is within 24 bits. diff --git a/crates/core/machine/src/alu/add_sub/mod.rs b/crates/core/machine/src/alu/add_sub/mod.rs index 2bb1cf0b0..f3a794b73 100644 --- a/crates/core/machine/src/alu/add_sub/mod.rs +++ b/crates/core/machine/src/alu/add_sub/mod.rs @@ -13,9 +13,13 @@ use zkm_core_executor::{ events::{AluEvent, ByteLookupEvent, ByteRecord}, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ - air::{MachineAir, PicusInfo, ZKMAirBuilder}, + air::{MachineAir, ZKMAirBuilder}, Word, }; @@ -38,7 +42,8 @@ pub const NUM_ADD_SUB_COLS: usize = size_of::>(); pub struct AddSubChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)] +#[derive(AlignedBorrow, Default, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct AddSubCols { /// The current/next pc, used for instruction lookup table. @@ -56,11 +61,11 @@ pub struct AddSubCols { pub operand_2: Word, /// Flag indicating whether the opcode is `ADD`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_add: T, /// Flag indicating whether the opcode is `SUB`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sub: T, } @@ -84,6 +89,7 @@ impl MachineAir for AddSubChip { Some(nb_rows) } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { AddSubCols::::picus_info() } diff --git a/crates/core/machine/src/alu/bitwise/mod.rs b/crates/core/machine/src/alu/bitwise/mod.rs index f0f44c1b8..1e50b8f3a 100644 --- a/crates/core/machine/src/alu/bitwise/mod.rs +++ b/crates/core/machine/src/alu/bitwise/mod.rs @@ -13,9 +13,13 @@ use zkm_core_executor::{ events::{AluEvent, ByteLookupEvent, ByteRecord}, ByteOpcode, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ - air::{MachineAir, PicusInfo, ZKMAirBuilder}, + air::{MachineAir, ZKMAirBuilder}, Word, }; @@ -32,7 +36,8 @@ pub const NUM_BITWISE_COLS: usize = size_of::>(); pub struct BitwiseChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)] +#[derive(AlignedBorrow, Default, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct BitwiseCols { /// The current/next pc, used for instruction lookup table. @@ -49,19 +54,19 @@ pub struct BitwiseCols { pub c: Word, /// If the opcode is NOR. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_nor: T, /// If the opcode is XOR. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_xor: T, // If the opcode is OR. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_or: T, /// If the opcode is AND. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_and: T, } @@ -76,6 +81,7 @@ impl MachineAir for BitwiseChip { "Bitwise".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { BitwiseCols::::picus_info() } diff --git a/crates/core/machine/src/alu/clo_clz/mod.rs b/crates/core/machine/src/alu/clo_clz/mod.rs index f3d028ce8..5953478d4 100644 --- a/crates/core/machine/src/alu/clo_clz/mod.rs +++ b/crates/core/machine/src/alu/clo_clz/mod.rs @@ -22,8 +22,12 @@ use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord}, ByteOpcode, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; -use zkm_stark::{air::MachineAir, PicusInfo, Word}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ air::ZKMCoreAirBuilder, @@ -42,7 +46,8 @@ pub struct CloClzChip; /// /// Optimized: `sr1` removed (hardcoded as 1 in SRL lookup since we always verify sr1 == 1), /// `is_clo` removed (derived as `is_real - is_clz`). -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct CloClzCols { /// The current/next pc, used for instruction lookup table. @@ -63,7 +68,7 @@ pub struct CloClzCols { pub is_bb_zero: T, /// Flag to indicate whether the opcode is CLZ. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_clz: T, /// Selector to know whether this row is enabled. @@ -81,6 +86,11 @@ impl MachineAir for CloClzChip { "CloClz".to_string() } + fn local_only(&self) -> bool { + true + } + + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { CloClzCols::::picus_info() } diff --git a/crates/core/machine/src/alu/divrem/mod.rs b/crates/core/machine/src/alu/divrem/mod.rs index 792d77cc5..22aedb94d 100644 --- a/crates/core/machine/src/alu/divrem/mod.rs +++ b/crates/core/machine/src/alu/divrem/mod.rs @@ -76,12 +76,13 @@ use zkm_core_executor::{ }; use crate::{memory::MemoryReadWriteCols, CoreChipError}; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_primitives::consts::WORD_SIZE; -use zkm_stark::{ - air::{MachineAir, PicusInfo}, - Word, -}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ air::{WordAirBuilder, ZKMCoreAirBuilder}, @@ -104,7 +105,8 @@ const LONG_WORD_SIZE: usize = 2 * WORD_SIZE; pub struct DivRemChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct DivRemCols { /// The current/next pc, used for instruction lookup table. @@ -142,19 +144,19 @@ pub struct DivRemCols { pub is_c_0: IsZeroWordOperation, /// Flag to indicate whether the opcode is DIV. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_div: T, /// Flag to indicate whether the opcode is DIVU. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_divu: T, /// Flag to indicate whether the opcode is MOD. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_mod: T, /// Flag to indicate whether the opcode is MODU. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_modu: T, /// Flag to indicate whether the division operation overflows. @@ -214,6 +216,7 @@ impl MachineAir for DivRemChip { "DivRem".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { DivRemCols::::picus_info() } @@ -353,6 +356,8 @@ impl MachineAir for DivRemChip { // Range check. { + output.add_u8_range_checks(&event.b.to_le_bytes()); + output.add_u8_range_checks(&event.c.to_le_bytes()); output.add_u8_range_checks("ient.to_le_bytes()); output.add_u8_range_checks(&remainder.to_le_bytes()); output.add_u8_range_checks(&c_times_quotient); @@ -665,6 +670,10 @@ where // Range check all the bytes. { + // Constrain operands to byte limbs so extracted standalone modules + // cannot pick non-byte witness values for word inputs. + builder.slice_range_check_u8(&local.b.0, is_real.clone()); + builder.slice_range_check_u8(&local.c.0, is_real.clone()); builder.slice_range_check_u8(&local.quotient.0, is_real.clone()); builder.slice_range_check_u8(&local.remainder.0, is_real.clone()); diff --git a/crates/core/machine/src/alu/lt/mod.rs b/crates/core/machine/src/alu/lt/mod.rs index 8c872b05b..cababa267 100644 --- a/crates/core/machine/src/alu/lt/mod.rs +++ b/crates/core/machine/src/alu/lt/mod.rs @@ -13,10 +13,14 @@ use zkm_core_executor::{ events::{AluEvent, ByteLookupEvent, ByteRecord}, ByteOpcode, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ air::{BaseAirBuilder, MachineAir, ZKMAirBuilder}, - PicusInfo, Word, + Word, }; use crate::{ @@ -32,7 +36,8 @@ pub const NUM_LT_COLS: usize = size_of::>(); pub struct LtChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)] +#[derive(AlignedBorrow, Default, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct LtCols { /// The current/next pc, used for instruction lookup table. @@ -40,11 +45,11 @@ pub struct LtCols { pub next_pc: T, /// If the opcode is SLT. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_slt: T, /// If the opcode is SLTU. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sltu: T, /// The output operand. @@ -113,6 +118,7 @@ impl MachineAir for LtChip { Some(nb_rows) } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { LtCols::::picus_info() } diff --git a/crates/core/machine/src/alu/mul/mod.rs b/crates/core/machine/src/alu/mul/mod.rs index ed2145184..af9c42dc6 100644 --- a/crates/core/machine/src/alu/mul/mod.rs +++ b/crates/core/machine/src/alu/mul/mod.rs @@ -44,9 +44,13 @@ use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord, CompAluEvent, MemoryAccessPosition, MemoryRecordEnum}, ByteOpcode, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_primitives::consts::WORD_SIZE; -use zkm_stark::{air::MachineAir, PicusInfo, Word}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ air::{WordAirBuilder, ZKMCoreAirBuilder}, @@ -74,11 +78,12 @@ pub const BYTE_MASK: u8 = 0xff; pub struct MulChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct MulCols { /// The current/next pc, used for instruction lookup table. - #[picus(input)] + #[cfg_attr(feature = "picus", picus(input))] pub pc: T, pub next_pc: T, @@ -113,15 +118,15 @@ pub struct MulCols { pub c_sign_extend: T, /// Flag indicating whether the opcode is `MUL`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_mul: T, /// Flag indicating whether the opcode is `MULT`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_mult: T, /// Flag indicating whether the opcode is `MULTU`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_multu: T, pub is_real: T, @@ -149,6 +154,7 @@ impl MachineAir for MulChip { "Mul".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { MulCols::::picus_info() } diff --git a/crates/core/machine/src/alu/sll/mod.rs b/crates/core/machine/src/alu/sll/mod.rs index 4ba299561..1751833c7 100644 --- a/crates/core/machine/src/alu/sll/mod.rs +++ b/crates/core/machine/src/alu/sll/mod.rs @@ -45,9 +45,13 @@ use zkm_core_executor::{ events::{AluEvent, ByteLookupEvent, ByteRecord}, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_primitives::consts::WORD_SIZE; -use zkm_stark::{air::MachineAir, PicusInfo, Word}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ air::ZKMCoreAirBuilder, @@ -66,7 +70,8 @@ pub const BYTE_SIZE: usize = 8; pub struct ShiftLeft; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct ShiftLeftCols { /// The current/next pc, used for instruction lookup table. @@ -114,6 +119,7 @@ impl MachineAir for ShiftLeft { "ShiftLeft".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { ShiftLeftCols::::picus_info() } diff --git a/crates/core/machine/src/alu/sr/mod.rs b/crates/core/machine/src/alu/sr/mod.rs index 2b10bd768..61eba3eed 100644 --- a/crates/core/machine/src/alu/sr/mod.rs +++ b/crates/core/machine/src/alu/sr/mod.rs @@ -57,9 +57,13 @@ use zkm_core_executor::{ events::{AluEvent, ByteLookupEvent, ByteRecord}, ByteOpcode, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_primitives::consts::WORD_SIZE; -use zkm_stark::{air::MachineAir, PicusInfo, Word}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ air::ZKMCoreAirBuilder, @@ -83,7 +87,8 @@ const BYTE_SIZE: usize = 8; pub struct ShiftRightChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct ShiftRightCols { /// The current/next pc, used for instruction lookup table. @@ -121,15 +126,15 @@ pub struct ShiftRightCols { pub c_least_sig_byte: [T; BYTE_SIZE], /// If the opcode is SRL. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_srl: T, /// If the opcode is ROR. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_ror: T, /// If the opcode is SRA. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sra: T, /// Selector to know whether this row is enabled. @@ -147,6 +152,7 @@ impl MachineAir for ShiftRightChip { "ShiftRight".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { ShiftRightCols::::picus_info() } @@ -225,6 +231,10 @@ impl MachineAir for ShiftRightChip { !shard.shift_right_events.is_empty() } } + + fn local_only(&self) -> bool { + true + } } impl ShiftRightChip { diff --git a/crates/core/machine/src/control_flow/branch/columns.rs b/crates/core/machine/src/control_flow/branch/columns.rs index 00d6d93a6..150b97414 100644 --- a/crates/core/machine/src/control_flow/branch/columns.rs +++ b/crates/core/machine/src/control_flow/branch/columns.rs @@ -1,13 +1,18 @@ use std::mem::size_of; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; -use zkm_stark::{PicusInfo, Word}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +use zkm_stark::Word; use crate::operations::KoalaBearWordRangeChecker; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_BRANCH_COLS: usize = size_of::>(); /// The column layout for branching. -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct BranchColumns { /// The current program counter. @@ -36,17 +41,17 @@ pub struct BranchColumns { pub op_c_value: Word, /// Branch Instructions Selectors. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_beq: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_bne: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_bltz: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_blez: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_bgtz: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_bgez: T, /// The branching column is equal to: diff --git a/crates/core/machine/src/control_flow/branch/trace.rs b/crates/core/machine/src/control_flow/branch/trace.rs index cee8006a8..b469526a7 100644 --- a/crates/core/machine/src/control_flow/branch/trace.rs +++ b/crates/core/machine/src/control_flow/branch/trace.rs @@ -9,7 +9,9 @@ use zkm_core_executor::{ events::{BranchEvent, ByteLookupEvent, ByteRecord}, ExecutionRecord, Opcode, Program, }; -use zkm_stark::{air::MachineAir, PicusInfo, Word}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ utils::{next_power_of_two, zeroed_f_vec}, @@ -29,6 +31,7 @@ impl MachineAir for BranchChip { "Branch".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { BranchColumns::::picus_info() } diff --git a/crates/core/machine/src/control_flow/jump/columns.rs b/crates/core/machine/src/control_flow/jump/columns.rs index f155ecfe9..7deb55620 100644 --- a/crates/core/machine/src/control_flow/jump/columns.rs +++ b/crates/core/machine/src/control_flow/jump/columns.rs @@ -1,12 +1,17 @@ use std::mem::size_of; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; -use zkm_stark::{PicusInfo, Word}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +use zkm_stark::Word; use crate::operations::KoalaBearWordRangeChecker; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_JUMP_COLS: usize = size_of::>(); -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct JumpColumns { /// The current program counter. @@ -28,11 +33,11 @@ pub struct JumpColumns { pub op_c_value: Word, /// Jump Instructions Selectors. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_jump: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_jumpi: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_jumpdirect: T, // A range checker for `op_a` which may contain `next_pc + 4`. diff --git a/crates/core/machine/src/control_flow/jump/trace.rs b/crates/core/machine/src/control_flow/jump/trace.rs index 74299aa4b..627de3e4d 100644 --- a/crates/core/machine/src/control_flow/jump/trace.rs +++ b/crates/core/machine/src/control_flow/jump/trace.rs @@ -9,7 +9,9 @@ use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord, JumpEvent}, ExecutionRecord, Opcode, Program, }; -use zkm_stark::{air::MachineAir, PicusInfo, Word}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::{air::MachineAir, Word}; use crate::{ utils::{next_power_of_two, zeroed_f_vec}, @@ -29,6 +31,7 @@ impl MachineAir for JumpChip { "Jump".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { JumpColumns::::picus_info() } diff --git a/crates/core/machine/src/cpu/air/mod.rs b/crates/core/machine/src/cpu/air/mod.rs index c7db873a0..02fa3557c 100644 --- a/crates/core/machine/src/cpu/air/mod.rs +++ b/crates/core/machine/src/cpu/air/mod.rs @@ -82,6 +82,9 @@ where // Check public values constraints. self.eval_pc(builder, local, next, public_values); + // Check control flag consistency. + self.eval_control_flags(builder, local); + // Check that the is_real flag is correct. self.eval_is_real(builder, local, next); @@ -89,10 +92,30 @@ where builder.when(not_real.clone()).assert_zero(AB::Expr::one() - local.instruction.imm_b); builder.when(not_real.clone()).assert_zero(AB::Expr::one() - local.instruction.imm_c); builder.when(not_real.clone()).assert_zero(AB::Expr::one() - local.is_rw_a); + builder.when(not_real.clone()).assert_zero(local.is_check_memory); + builder.when(not_real.clone()).assert_zero(local.is_halt); + builder.when(not_real.clone()).assert_zero(local.is_sequential); + builder.when(not_real.clone()).assert_zero(local.next_pc); + builder.when(not_real).assert_zero(local.next_next_pc); } } impl CpuChip { + /// Constraints for control flags carried in the CPU row. + pub(crate) fn eval_control_flags( + &self, + builder: &mut AB, + local: &CpuCols, + ) { + builder.when(local.is_real).assert_bool(local.is_rw_a); + builder.when(local.is_real).assert_bool(local.is_check_memory); + builder.when(local.is_real).assert_bool(local.is_halt); + builder.when(local.is_real).assert_bool(local.is_sequential); + + // Halting instructions are not sequential. + builder.when(local.is_real).assert_zero(local.is_halt * local.is_sequential); + } + /// Constraints related to the shard and clk. /// /// This method ensures that all of the shard values are the same and that the clk starts at 0 @@ -169,7 +192,6 @@ impl CpuChip { .assert_eq(local.next_next_pc, next.next_pc); builder - .when_transition() .when(local.is_real) .when(local.is_sequential) .assert_eq(local.next_next_pc, local.next_pc + AB::Expr::from_canonical_u32(4)); diff --git a/crates/core/machine/src/cpu/columns/mod.rs b/crates/core/machine/src/cpu/columns/mod.rs index 5244d5027..4b96ed504 100644 --- a/crates/core/machine/src/cpu/columns/mod.rs +++ b/crates/core/machine/src/cpu/columns/mod.rs @@ -4,9 +4,13 @@ pub use instruction::*; use p3_util::indices_arr; use std::mem::{size_of, transmute}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::Word; use crate::memory::{MemoryCols, MemoryReadCols, MemoryReadWriteCols}; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_CPU_COLS: usize = size_of::>(); @@ -14,6 +18,7 @@ pub const CPU_COL_MAP: CpuCols = make_col_map(); /// The column layout for the CPU. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct CpuCols { /// The current shard. diff --git a/crates/core/machine/src/cpu/trace.rs b/crates/core/machine/src/cpu/trace.rs index a117e4b80..2f35e37d2 100644 --- a/crates/core/machine/src/cpu/trace.rs +++ b/crates/core/machine/src/cpu/trace.rs @@ -28,6 +28,11 @@ impl MachineAir for CpuChip { self.id().to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> zkm_stark::PicusInfo { + CpuCols::::picus_info() + } + fn num_rows(&self, input: &Self::Record) -> Option { let n_real_rows = input.cpu_events.len(); let padded_nb_rows = if let Some(shape) = &input.shape { diff --git a/crates/core/machine/src/global/mod.rs b/crates/core/machine/src/global/mod.rs index ac4f0f104..734cd080a 100644 --- a/crates/core/machine/src/global/mod.rs +++ b/crates/core/machine/src/global/mod.rs @@ -13,6 +13,8 @@ use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord, GlobalLookupEvent}, ExecutionRecord, Program, }; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ air::{AirLookup, LookupScope, MachineAir}, septic_curve::{SepticCurve, SepticCurveComplete}, @@ -27,6 +29,8 @@ use crate::{ CoreChipError, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; const NUM_GLOBAL_COLS: usize = size_of::>(); @@ -51,9 +55,11 @@ pub struct Ghost { pub struct GlobalChip; #[derive(AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct GlobalCols { pub message: [T; 7], + #[cfg_attr(feature = "picus", picus(output))] pub kind: T, pub lookup: GlobalLookupOperation, pub is_receive: T, @@ -74,6 +80,11 @@ impl MachineAir for GlobalChip { "Global".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + GlobalCols::::picus_info() + } + fn generate_dependencies( &self, input: &Self::Record, @@ -233,6 +244,13 @@ where // For a syscall global lookup, `kind = LookupKind::Syscall` is used. // Therefore, `is_send`, `is_receive` are already known to be boolean, and `kind` is also known to be a `u8` value. // Note that `local.is_real` is constrained to be boolean in `eval_single_digest`. + // Enforce local consistency for direction selectors: + // - each selector is boolean; + // - real rows choose exactly one direction; padded rows choose none. + builder.assert_bool(local.is_receive); + builder.assert_bool(local.is_send); + builder.assert_eq(local.is_receive + local.is_send, local.is_real.into()); + builder.receive( AirLookup::new( vec![ diff --git a/crates/core/machine/src/memory/global.rs b/crates/core/machine/src/memory/global.rs index 254d5f3dc..440b7fcc3 100644 --- a/crates/core/machine/src/memory/global.rs +++ b/crates/core/machine/src/memory/global.rs @@ -11,6 +11,10 @@ use p3_maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator}; use zkm_core_executor::events::{GlobalLookupEvent, MemoryInitializeFinalizeEvent}; use zkm_core_executor::{ExecutionRecord, Program}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; use zkm_stark::{ air::{ AirLookup, BaseAirBuilder, LookupScope, MachineAir, PublicValues, ZKMAirBuilder, @@ -59,6 +63,11 @@ impl MachineAir for MemoryGlobalChip { } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> zkm_stark::PicusInfo { + MemoryInitCols::::picus_info() + } + fn generate_dependencies( &self, input: &ExecutionRecord, @@ -208,15 +217,19 @@ impl MachineAir for MemoryGlobalChip { } #[derive(AlignedBorrow, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct MemoryInitCols { /// The shard number of the memory access. + #[cfg_attr(feature = "picus", picus(input, transition_input))] pub shard: T, /// The timestamp of the memory access. + #[cfg_attr(feature = "picus", picus(input, transition_input))] pub timestamp: T, /// The address of the memory access. + #[cfg_attr(feature = "picus", picus(input, transition_input))] pub addr: T, /// Comparison assertions for address to be strictly increasing. @@ -226,6 +239,7 @@ pub struct MemoryInitCols { pub addr_bits: KoalaBearBitDecomposition, /// The value of the memory access. + #[cfg_attr(feature = "picus", picus(transition_input))] pub value: [T; 32], /// Whether the memory access is a real access. @@ -261,6 +275,35 @@ where for i in 0..32 { builder.assert_bool(local.value[i]); } + // Canonicalize padded rows to the default zero trace shape so witness columns cannot + // drift in extraction modules. + builder.when_not(local.is_real).assert_zero(local.shard); + builder.when_not(local.is_real).assert_zero(local.timestamp); + builder.when_not(local.is_real).assert_zero(local.addr); + builder.when_not(local.is_real).assert_zero(local.is_next_comp); + for i in 0..32 { + builder.when_not(local.is_real).assert_zero(local.value[i]); + builder.when_not(local.is_real).assert_zero(local.addr_bits.bits[i]); + builder.when_not(local.is_real).assert_zero(local.lt_cols.bit_flags[i]); + } + builder + .when_not(local.is_real) + .assert_zero(local.addr_bits.and_most_sig_byte_decomp_0_to_2); + builder + .when_not(local.is_real) + .assert_zero(local.addr_bits.and_most_sig_byte_decomp_0_to_3); + builder + .when_not(local.is_real) + .assert_zero(local.addr_bits.and_most_sig_byte_decomp_0_to_4); + builder + .when_not(local.is_real) + .assert_zero(local.addr_bits.and_most_sig_byte_decomp_0_to_5); + builder + .when_not(local.is_real) + .assert_zero(local.addr_bits.and_most_sig_byte_decomp_0_to_6); + builder + .when_not(local.is_real) + .assert_zero(local.addr_bits.and_most_sig_byte_decomp_0_to_7); let mut byte1 = AB::Expr::zero(); let mut byte2 = AB::Expr::zero(); @@ -279,8 +322,8 @@ where builder.send( AirLookup::new( vec![ - AB::Expr::zero(), - AB::Expr::zero(), + AB::Expr::zero(), // shard + AB::Expr::zero(), // timestamp local.addr.into(), value[0].clone(), value[1].clone(), @@ -339,8 +382,19 @@ where // - In the last real row, we need to assert that addr = last_finalize_addr. // Assert that addr < addr' when the next row is real. - builder.when_transition().assert_eq(next.is_next_comp, next.is_real); - next.lt_cols.eval(builder, &local.addr_bits.bits, &next.addr_bits.bits, next.is_next_comp); + // + // Keep these constraints transition-scoped so boundary modules don't introduce unconstrained + // `next`-row helper witnesses. + { + let mut transition = builder.when_transition(); + transition.assert_eq(next.is_next_comp, next.is_real); + next.lt_cols.eval( + &mut transition, + &local.addr_bits.bits, + &next.addr_bits.bits, + next.is_next_comp, + ); + } // Assert that the real rows are all padded to the top. builder.when_transition().when_not(local.is_real).assert_zero(next.is_real); @@ -378,11 +432,50 @@ where let is_first_row = builder.is_first_row(); IsZeroOperation::::eval(builder, prev_addr, local.is_prev_addr_zero, is_first_row); + // Outside the first row, `is_prev_addr_zero` is a pure witness helper and should stay at + // its default trace value. Constrain this through transition `next` columns so the single + // row case does not accidentally force first-row helpers to zero in boundary extraction. + builder.when_transition().assert_zero(next.is_prev_addr_zero.inverse); + builder.when_transition().assert_zero(next.is_prev_addr_zero.result); + + // When prev_addr == 0 in the first row, canonicalize the helper witness to match trace + // population (inverse = 0). + builder + .when_first_row() + .when(local.is_prev_addr_zero.result) + .assert_zero(local.is_prev_addr_zero.inverse); + // Constrain the is_first_comp column. builder.assert_bool(local.is_first_comp); builder .when_first_row() .assert_eq(local.is_first_comp, AB::Expr::one() - local.is_prev_addr_zero.result); + // In the degenerate single-row case, force the row to be the `%x0` address case. + // This removes a Picus-only underconstrained branch where `addr` can drift without inputs. + let is_single_row = builder.is_first_row() * builder.is_last_row(); + builder.when(is_single_row.clone()).assert_zero(local.is_first_comp); + // In the degenerate single-row finalize case, the only real finalized + // address is `%x0`, whose routing metadata is canonical in the executor + // trace population (`shard = 0`, `timestamp = 1`). + // Pin these to avoid underconstrained single-row witnesses in extraction. + if self.kind == MemoryChipType::Finalize { + builder.when(is_single_row.clone()).assert_zero(local.shard); + builder.when(is_single_row).assert_eq(local.timestamp, AB::Expr::one()); + } + builder.when_transition().assert_zero(next.is_first_comp); + // For all non-first real rows (`is_next_comp = 1` in this trace), first-row-only helper + // columns must be zero. + builder.when(local.is_next_comp).assert_zero(local.is_prev_addr_zero.inverse); + builder.when(local.is_next_comp).assert_zero(local.is_prev_addr_zero.result); + builder.when(local.is_next_comp).assert_zero(local.is_first_comp); + + // Canonicalize local less-than helper flags when no local comparison is requested. + // This is exactly the case `is_first_comp = 0` and `is_next_comp = 0`. + let no_local_lt_check = + (AB::Expr::one() - local.is_first_comp) * (AB::Expr::one() - local.is_next_comp); + for flag in local.lt_cols.bit_flags.iter() { + builder.assert_zero(no_local_lt_check.clone() * (*flag)); + } // Ensure at least one real row. builder.when_first_row().assert_one(local.is_real); @@ -404,6 +497,7 @@ where if self.kind == MemoryChipType::Initialize { builder.when(local.is_real).assert_eq(local.timestamp, AB::F::ONE); + builder.when(local.is_real).assert_eq(local.shard, AB::F::ONE); } // Constraints related to register %x0. @@ -431,6 +525,8 @@ where // - The flag `is_real` is set to one and the next `is_real` is set to zero. // Constrain the `is_last_addr` flag. + builder.assert_bool(local.is_last_addr); + builder.when_last_row().assert_eq(local.is_last_addr, local.is_real); builder .when_transition() .assert_eq(local.is_last_addr, local.is_real * (AB::Expr::one() - next.is_real)); diff --git a/crates/core/machine/src/memory/instructions/columns.rs b/crates/core/machine/src/memory/instructions/columns.rs index a688dd1ba..8dfdc88ee 100644 --- a/crates/core/machine/src/memory/instructions/columns.rs +++ b/crates/core/machine/src/memory/instructions/columns.rs @@ -1,16 +1,21 @@ use std::mem::size_of; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; -use zkm_stark::{PicusInfo, Word}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +use zkm_stark::Word; use crate::{ memory::MemoryReadWriteCols, operations::{IsZeroOperation, KoalaBearWordRangeChecker}, }; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_MEMORY_INSTRUCTIONS_COLUMNS: usize = size_of::>(); /// The column layout for memory. -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct MemoryInstructionsColumns { /// The current/next program counter of the instruction. @@ -30,46 +35,46 @@ pub struct MemoryInstructionsColumns { pub op_c_value: Word, /// Whether this is a load byte instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lb: T, /// Whether this is a load byte unsigned instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lbu: T, /// Whether this is a load half instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lh: T, /// Whether this is a load half unsigned instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lhu: T, /// Whether this is a load word instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lw: T, /// Whether this is a lwl instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lwl: T, /// Whether this is a lwr instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_lwr: T, /// Whether this is a ll instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_ll: T, /// Whether this is a store byte instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sb: T, /// Whether this is a store half instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sh: T, /// Whether this is a store word instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sw: T, /// Whether this is a swl instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_swl: T, /// Whether this is a swr instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_swr: T, /// Whether this is a sc instruction. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sc: T, /// The relationships among addr_word, addr_aligned, and addr_offset is as follows: diff --git a/crates/core/machine/src/memory/instructions/trace.rs b/crates/core/machine/src/memory/instructions/trace.rs index af61cddc7..41a59735a 100644 --- a/crates/core/machine/src/memory/instructions/trace.rs +++ b/crates/core/machine/src/memory/instructions/trace.rs @@ -33,6 +33,7 @@ impl MachineAir for MemoryInstructionsChip { "MemoryInstrs".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> zkm_stark::PicusInfo { MemoryInstructionsColumns::::picus_info() } diff --git a/crates/core/machine/src/mips/mod.rs b/crates/core/machine/src/mips/mod.rs index 4f6fd6a54..03a9518cf 100644 --- a/crates/core/machine/src/mips/mod.rs +++ b/crates/core/machine/src/mips/mod.rs @@ -18,8 +18,9 @@ use zkm_core_executor::{ events::PrecompileLocalMemory, syscalls::SyscallCode, ExecutionRecord, MipsAirId, Program, }; use zkm_curves::weierstrass::{bls12_381::Bls12381BaseField, bn254::Bn254BaseField}; +use zkm_stark::PicusInfo; use zkm_stark::{ - air::{LookupScope, MachineAir, PicusInfo, ZKM_PROOF_NUM_PV_ELTS}, + air::{LookupScope, MachineAir, ZKM_PROOF_NUM_PV_ELTS}, Chip, LookupKind, StarkGenericConfig, StarkMachine, }; diff --git a/crates/core/machine/src/misc/mov_cond/mod.rs b/crates/core/machine/src/misc/mov_cond/mod.rs index f4f08ef8a..45dd6a393 100644 --- a/crates/core/machine/src/misc/mov_cond/mod.rs +++ b/crates/core/machine/src/misc/mov_cond/mod.rs @@ -13,10 +13,14 @@ use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord, MovCondEvent}, ExecutionRecord, Opcode, Program, }; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ air::{BaseAirBuilder, MachineAir, ZKMAirBuilder}, - {PicusInfo, Word}, + Word, }; use crate::{air::WordAirBuilder, CoreChipError}; @@ -33,7 +37,8 @@ pub const NUM_MOV_COND_COLS: usize = size_of::>(); pub struct MovCondChip; /// The column layout for the chip. -#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)] +#[derive(AlignedBorrow, Default, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct MovCondCols { /// The current/next pc, used for instruction lookup table. @@ -52,15 +57,15 @@ pub struct MovCondCols { pub c_eq_0: IsZeroWordOperation, /// Flag indicating whether the opcode is `MNE`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_mne: T, /// Flag indicating whether the opcode is `MEQ`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_meq: T, /// Flag indicating whether the opcode is `WSBH`. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_wsbh: T, } @@ -75,6 +80,7 @@ impl MachineAir for MovCondChip { "MovCond".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { MovCondCols::::picus_info() } diff --git a/crates/core/machine/src/misc/others/columns/mod.rs b/crates/core/machine/src/misc/others/columns/mod.rs index 31b35e58d..c7430e6d4 100644 --- a/crates/core/machine/src/misc/others/columns/mod.rs +++ b/crates/core/machine/src/misc/others/columns/mod.rs @@ -11,12 +11,17 @@ pub use misc_specific::*; pub use sext::*; use std::mem::size_of; -use zkm_derive::{AlignedBorrow, PicusAnnotations}; -use zkm_stark::{PicusInfo, Word}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +use zkm_stark::Word; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_MISC_INSTR_COLS: usize = size_of::>(); -#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)] +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct MiscInstrColumns { /// The shard number. @@ -39,20 +44,20 @@ pub struct MiscInstrColumns { pub misc_specific_columns: MiscSpecificCols, /// Misc Instruction Selectors. - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_sext: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_ins: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_ext: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_maddu: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_msubu: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_madd: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_msub: T, - #[picus(selector)] + #[cfg_attr(feature = "picus", picus(selector))] pub is_teq: T, } diff --git a/crates/core/machine/src/misc/others/trace.rs b/crates/core/machine/src/misc/others/trace.rs index e7bb81ee6..680b8555b 100644 --- a/crates/core/machine/src/misc/others/trace.rs +++ b/crates/core/machine/src/misc/others/trace.rs @@ -32,10 +32,15 @@ impl MachineAir for MiscInstrsChip { "MiscInstrs".to_string() } + #[cfg(feature = "picus")] fn picus_info(&self) -> zkm_stark::PicusInfo { MiscInstrColumns::::picus_info() } + fn local_only(&self) -> bool { + true + } + fn num_rows(&self, input: &Self::Record) -> Option { let nb_rows = next_power_of_two( input.misc_events.len(), diff --git a/crates/core/machine/src/operations/add.rs b/crates/core/machine/src/operations/add.rs index 30074abe7..6d626557c 100644 --- a/crates/core/machine/src/operations/add.rs +++ b/crates/core/machine/src/operations/add.rs @@ -1,11 +1,22 @@ +#[cfg(feature = "picus")] +use core::borrow::Borrow; +#[cfg(feature = "picus")] +use core::mem::{size_of, transmute}; + use zkm_core_executor::events::ByteRecord; +#[cfg(feature = "picus")] +use zkm_primitives::consts::WORD_SIZE; use zkm_stark::{air::ZKMAirBuilder, Word}; use p3_air::AirBuilder; use p3_field::{Field, FieldAlgebra}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusProjection; use crate::air::WordAirBuilder; +#[cfg(feature = "picus")] +use crate::utils::indices_arr; /// A set of columns needed to compute the add of two words. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] @@ -18,6 +29,51 @@ pub struct AddOperation { pub carry: [T; 3], } +#[cfg(feature = "picus")] +const NUM_ADD_OPERATION_SUMMARY_COLS: usize = size_of::>(); + +#[cfg(feature = "picus")] +const ADD_OPERATION_SUMMARY_COL_MAP: AddOperationSummaryCols = + make_add_operation_summary_col_map(); + +#[cfg(feature = "picus")] +const fn make_add_operation_summary_col_map() -> AddOperationSummaryCols { + let indices_arr = indices_arr::(); + unsafe { + transmute::<[usize; NUM_ADD_OPERATION_SUMMARY_COLS], AddOperationSummaryCols>( + indices_arr, + ) + } +} + +/// Hidden witness layout used when Picus emits the exact two-word add AIR as a +/// local auxiliary module. +#[derive(AlignedBorrow, Clone, Copy)] +#[repr(C)] +#[cfg(feature = "picus")] +struct AddOperationSummaryCols { + pub a: Word, + pub b: Word, + pub is_real: T, + pub cols: AddOperation, +} + +#[cfg(feature = "picus")] +#[cfg_attr(feature = "picus", derive(PicusProjection))] +#[cfg_attr(feature = "picus", picus_projection( + source = AddOperationSummaryCols, + col_map = ADD_OPERATION_SUMMARY_COL_MAP +))] +#[allow(dead_code)] +struct AddOperationSummaryProjection { + #[cfg_attr(feature = "picus", picus(input, path = a))] + pub a: Word, + #[cfg_attr(feature = "picus", picus(input, path = b))] + pub b: Word, + #[cfg_attr(feature = "picus", picus(output, path = cols.value))] + pub value: Word, +} + impl AddOperation { #[allow(unused_assignments)] pub fn populate(&mut self, record: &mut impl ByteRecord, a_u32: u32, b_u32: u32) -> u32 { @@ -54,7 +110,7 @@ impl AddOperation { expected } - pub fn eval( + fn eval_exact( builder: &mut AB, a: Word, b: Word, @@ -98,4 +154,51 @@ impl AddOperation { builder.slice_range_check_u8(&cols.value.0, is_real); } } + + pub fn eval( + builder: &mut AB, + a: Word, + b: Word, + cols: AddOperation, + is_real: AB::Expr, + ) { + #[cfg(feature = "picus")] + { + let mut current_inputs: Vec = Vec::with_capacity(WORD_SIZE * 2); + for limb in a.0 { + current_inputs.push(limb.into()); + } + for limb in b.0 { + current_inputs.push(limb.into()); + } + + let current_outputs: Vec = + cols.value.0.iter().map(|limb| (*limb).into()).collect(); + + if builder.is_known_one(&is_real) + && builder.try_emit_projected_summary_with_hidden_consts( + "AddOperation", + &AddOperationSummaryProjection::picus_projection_info(), + ¤t_inputs, + ¤t_outputs, + size_of::>(), + &[(ADD_OPERATION_SUMMARY_COL_MAP.is_real, 1)], + |builder, source_row| { + let source: &AddOperationSummaryCols = (*source_row).borrow(); + Self::eval_exact( + builder, + source.a, + source.b, + source.cols, + source.is_real.into(), + ); + }, + ) + { + return; + } + } + + Self::eval_exact(builder, a, b, cols, is_real); + } } diff --git a/crates/core/machine/src/operations/add5.rs b/crates/core/machine/src/operations/add5.rs index 3d0db30d5..bcdfe8d89 100644 --- a/crates/core/machine/src/operations/add5.rs +++ b/crates/core/machine/src/operations/add5.rs @@ -1,12 +1,21 @@ +#[cfg(feature = "picus")] +use core::borrow::Borrow; +#[cfg(feature = "picus")] +use core::mem::{size_of, transmute}; + use p3_air::AirBuilder; use p3_field::{Field, FieldAlgebra}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusProjection; use zkm_core_executor::events::ByteRecord; use zkm_primitives::consts::WORD_SIZE; use zkm_stark::{air::ZKMAirBuilder, Word}; use crate::air::WordAirBuilder; +#[cfg(feature = "picus")] +use crate::utils::indices_arr; /// A set of columns needed to compute the sum of five words. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] @@ -34,6 +43,48 @@ pub struct Add5Operation { pub carry: Word, } +#[cfg(feature = "picus")] +const NUM_ADD5_OPERATION_SUMMARY_COLS: usize = size_of::>(); + +#[cfg(feature = "picus")] +const ADD5_OPERATION_SUMMARY_COL_MAP: Add5OperationSummaryCols = + make_add5_operation_summary_col_map(); + +#[cfg(feature = "picus")] +const fn make_add5_operation_summary_col_map() -> Add5OperationSummaryCols { + let indices_arr = indices_arr::(); + unsafe { + transmute::<[usize; NUM_ADD5_OPERATION_SUMMARY_COLS], Add5OperationSummaryCols>( + indices_arr, + ) + } +} + +/// Hidden witness layout used when Picus emits the exact five-word add AIR as +/// a local auxiliary module. +#[derive(AlignedBorrow, Clone, Copy)] +#[repr(C)] +#[cfg(feature = "picus")] +struct Add5OperationSummaryCols { + pub words: [Word; 5], + pub is_real: T, + pub cols: Add5Operation, +} + +#[cfg(feature = "picus")] +#[cfg_attr(feature = "picus", derive(PicusProjection))] +#[cfg_attr(feature = "picus", picus_projection( + source = Add5OperationSummaryCols, + col_map = ADD5_OPERATION_SUMMARY_COL_MAP +))] +#[allow(dead_code)] +struct Add5OperationSummaryProjection { + #[cfg_attr(feature = "picus", picus(input, path = words))] + pub words: [Word; 5], + #[cfg_attr(feature = "picus", picus(output, path = cols.value))] + pub value: Word, +} + impl Add5Operation { #[allow(clippy::too_many_arguments)] pub fn populate( @@ -87,7 +138,7 @@ impl Add5Operation { expected } - pub fn eval( + fn eval_exact( builder: &mut AB, words: &[Word; 5], is_real: AB::Var, @@ -157,4 +208,44 @@ impl Add5Operation { } } } + + pub fn eval( + builder: &mut AB, + words: &[Word; 5], + is_real: AB::Var, + cols: Add5Operation, + ) { + #[cfg(feature = "picus")] + { + let is_real_expr = AB::Expr::zero() + is_real; + let mut current_inputs: Vec = Vec::with_capacity(WORD_SIZE * 5); + for word in words { + for limb in word.0 { + current_inputs.push(limb.into()); + } + } + + let current_outputs: Vec = + cols.value.0.iter().map(|limb| (*limb).into()).collect(); + + if builder.is_known_one(&is_real_expr) + && builder.try_emit_projected_summary_with_hidden_consts( + "Add5Operation", + &Add5OperationSummaryProjection::picus_projection_info(), + ¤t_inputs, + ¤t_outputs, + size_of::>(), + &[(ADD5_OPERATION_SUMMARY_COL_MAP.is_real, 1)], + |builder, source_row| { + let source: &Add5OperationSummaryCols = (*source_row).borrow(); + Self::eval_exact(builder, &source.words, source.is_real, source.cols); + }, + ) + { + return; + } + } + + Self::eval_exact(builder, words, is_real, cols); + } } diff --git a/crates/core/machine/src/operations/and.rs b/crates/core/machine/src/operations/and.rs index b3e4843be..36a70c4a3 100644 --- a/crates/core/machine/src/operations/and.rs +++ b/crates/core/machine/src/operations/and.rs @@ -1,5 +1,12 @@ +#[cfg(feature = "picus")] +use core::borrow::Borrow; +#[cfg(feature = "picus")] +use core::mem::{size_of, transmute}; + use p3_field::{Field, FieldAlgebra}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusProjection; use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord}, @@ -8,6 +15,9 @@ use zkm_core_executor::{ use zkm_primitives::consts::WORD_SIZE; use zkm_stark::{air::ZKMAirBuilder, Word}; +#[cfg(feature = "picus")] +use crate::utils::indices_arr; + /// A set of columns needed to compute the and of two words. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] #[repr(C)] @@ -16,6 +26,53 @@ pub struct AndOperation { pub value: Word, } +#[cfg(feature = "picus")] +const NUM_AND_OPERATION_SUMMARY_COLS: usize = size_of::>(); + +#[cfg(feature = "picus")] +const AND_OPERATION_SUMMARY_COL_MAP: AndOperationSummaryCols = + make_and_operation_summary_col_map(); + +#[cfg(feature = "picus")] +const fn make_and_operation_summary_col_map() -> AndOperationSummaryCols { + let indices_arr = indices_arr::(); + unsafe { + transmute::<[usize; NUM_AND_OPERATION_SUMMARY_COLS], AndOperationSummaryCols>( + indices_arr, + ) + } +} + +/// Hidden witness layout used when Picus emits the exact AND AIR as a local +/// auxiliary module. +#[derive(AlignedBorrow, Clone, Copy)] +#[repr(C)] +#[cfg(feature = "picus")] +struct AndOperationSummaryCols { + pub a: Word, + pub b: Word, + pub is_real: T, + pub cols: AndOperation, +} + +#[cfg(feature = "picus")] +#[cfg_attr(feature = "picus", derive(PicusProjection))] +#[cfg_attr(feature = "picus", picus_projection( + source = AndOperationSummaryCols, + col_map = AND_OPERATION_SUMMARY_COL_MAP +))] +#[allow(dead_code)] +struct AndOperationSummaryProjection { + #[cfg_attr(feature = "picus", picus(input, path = a))] + pub a: Word, + #[cfg_attr(feature = "picus", picus(input, path = b))] + pub b: Word, + #[cfg_attr(feature = "picus", picus(input, path = is_real))] + pub is_real: u8, + #[cfg_attr(feature = "picus", picus(output, path = cols.value))] + pub value: Word, +} + impl AndOperation { pub fn populate(&mut self, record: &mut impl ByteRecord, x: u32, y: u32) -> u32 { let expected = x & y; @@ -37,8 +94,7 @@ impl AndOperation { expected } - #[allow(unused_variables)] - pub fn eval( + fn eval_exact( builder: &mut AB, a: Word, b: Word, @@ -55,4 +111,45 @@ impl AndOperation { ); } } + + #[allow(unused_variables)] + pub fn eval( + builder: &mut AB, + a: Word, + b: Word, + cols: AndOperation, + is_real: AB::Var, + ) { + let is_real_expr = AB::Expr::zero() + is_real; + let mut current_inputs: Vec = Vec::with_capacity(WORD_SIZE * 2 + 1); + for limb in a.0 { + current_inputs.push(limb.into()); + } + for limb in b.0 { + current_inputs.push(limb.into()); + } + current_inputs.push(is_real_expr.clone()); + + let current_outputs: Vec = + cols.value.0.iter().map(|limb| (*limb).into()).collect(); + + #[cfg(feature = "picus")] + if builder.is_known_one(&is_real_expr) + && builder.try_emit_projected_summary( + "AndOperation", + &AndOperationSummaryProjection::picus_projection_info(), + ¤t_inputs, + ¤t_outputs, + size_of::>(), + |builder, source_row| { + let source: &AndOperationSummaryCols = (*source_row).borrow(); + Self::eval_exact(builder, source.a, source.b, source.cols, source.is_real); + }, + ) + { + return; + } + + Self::eval_exact(builder, a, b, cols, is_real); + } } diff --git a/crates/core/machine/src/operations/cmp.rs b/crates/core/machine/src/operations/cmp.rs index cabb00c07..be213d571 100644 --- a/crates/core/machine/src/operations/cmp.rs +++ b/crates/core/machine/src/operations/cmp.rs @@ -9,7 +9,7 @@ use zkm_core_executor::{ ByteOpcode, }; use zkm_derive::AlignedBorrow; -use zkm_stark::air::{BaseAirBuilder, ZKMAirBuilder}; +use zkm_stark::air::ZKMAirBuilder; use zkm_stark::Word; /// Operation columns for verifying that an element is within the range `[0, modulus)`. @@ -320,11 +320,7 @@ impl AssertLtColsBits { } impl AssertLtColsBits { - pub fn eval< - AB: ZKMAirBuilder, - Ea: Into + Clone, - Eb: Into + Clone, - >( + pub fn eval, Ea: Into + Clone, Eb: Into + Clone>( &self, builder: &mut AB, a: &[Ea], @@ -377,10 +373,9 @@ impl AssertLtColsBits { a_comparison_bit = a_comparison_bit.clone() + a_bit.clone() * flag; b_comparison_bit = b_comparison_bit.clone() + b_bit.clone() * flag; - builder - .when(is_real.clone()) - .when_not(is_inequality_visited.clone()) - .assert_eq(a_bit.clone(), b_bit.clone()); + builder.when(is_real.clone()).assert_zero( + (AB::Expr::one() - is_inequality_visited.clone()) * (a_bit.clone() - b_bit.clone()), + ); } builder.when(is_real.clone()).assert_eq(a_comparison_bit, AB::F::zero()); diff --git a/crates/core/machine/src/operations/field/harness.rs b/crates/core/machine/src/operations/field/harness.rs new file mode 100644 index 000000000..23caf5098 --- /dev/null +++ b/crates/core/machine/src/operations/field/harness.rs @@ -0,0 +1,256 @@ +use core::{borrow::Borrow, mem::size_of}; +use std::fmt::Debug; + +use p3_air::{Air, BaseAir}; +use p3_field::{Field, PrimeField32}; +use p3_matrix::{dense::RowMajorMatrix, Matrix}; +use zkm_core_executor::{events::FieldOperation, ExecutionRecord, Program}; +use zkm_curves::{edwards::ed25519::Ed25519BaseField, params::Limbs}; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; +use zkm_stark::air::{MachineAir, ZKMAirBuilder}; + +use crate::{ + operations::field::{ + field_den::FieldDenCols, field_inner_product::FieldInnerProductCols, field_op::FieldOpCols, + }, + CoreChipError, +}; + +#[derive(AlignedBorrow, Debug, Clone)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] +#[repr(C)] +pub struct FieldOpHarnessCols { + pub is_real: T, + #[cfg_attr(feature = "picus", picus(input))] + pub a: Limbs::Limbs>, + #[cfg_attr(feature = "picus", picus(input))] + pub b: Limbs::Limbs>, + #[cfg_attr(feature = "picus", picus(output))] + pub op: FieldOpCols, +} + +#[derive(AlignedBorrow, Debug, Clone)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] +#[repr(C)] +pub struct FieldDenHarnessCols { + pub is_real: T, + #[cfg_attr(feature = "picus", picus(input))] + pub a: Limbs::Limbs>, + #[cfg_attr(feature = "picus", picus(input))] + pub b: Limbs::Limbs>, + #[cfg_attr(feature = "picus", picus(output))] + pub den: FieldDenCols, +} + +#[derive(AlignedBorrow, Debug, Clone)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] +#[repr(C)] +pub struct FieldInnerProductHarnessCols { + pub is_real: T, + #[cfg_attr(feature = "picus", picus(input))] + pub a: [Limbs::Limbs>; 2], + #[cfg_attr(feature = "picus", picus(input))] + pub b: [Limbs::Limbs>; 2], + #[cfg_attr(feature = "picus", picus(output))] + pub inner_product: FieldInnerProductCols, +} + +pub const NUM_FIELD_OP_HARNESS_COLS: usize = size_of::>(); +pub const NUM_FIELD_DEN_HARNESS_COLS: usize = size_of::>(); +pub const NUM_FIELD_INNER_PRODUCT_HARNESS_COLS: usize = + size_of::>(); + +pub struct FieldOpHarnessChip { + pub operation: FieldOperation, + pub name: &'static str, +} + +impl FieldOpHarnessChip { + pub const fn new(operation: FieldOperation, name: &'static str) -> Self { + Self { operation, name } + } +} + +impl MachineAir for FieldOpHarnessChip { + type Record = ExecutionRecord; + type Program = Program; + type Error = CoreChipError; + + fn name(&self) -> String { + self.name.to_string() + } + + fn generate_trace( + &self, + _: &ExecutionRecord, + _: &mut ExecutionRecord, + ) -> Result, Self::Error> { + Ok(RowMajorMatrix::new(vec![F::ZERO; NUM_FIELD_OP_HARNESS_COLS], NUM_FIELD_OP_HARNESS_COLS)) + } + + fn included(&self, _: &Self::Record) -> bool { + true + } + + fn local_only(&self) -> bool { + true + } + + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + FieldOpHarnessCols::::picus_info() + } +} + +impl BaseAir for FieldOpHarnessChip { + fn width(&self) -> usize { + NUM_FIELD_OP_HARNESS_COLS + } +} + +impl Air for FieldOpHarnessChip +where + AB: ZKMAirBuilder, + Limbs::Limbs>: Copy, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0); + let local: &FieldOpHarnessCols = (*local).borrow(); + builder.assert_bool(local.is_real); + local.op.eval(builder, &local.a, &local.b, self.operation, local.is_real); + } +} + +pub struct FieldDenHarnessChip { + pub sign: bool, + pub name: &'static str, +} + +impl FieldDenHarnessChip { + pub const fn new(sign: bool, name: &'static str) -> Self { + Self { sign, name } + } +} + +impl MachineAir for FieldDenHarnessChip { + type Record = ExecutionRecord; + type Program = Program; + type Error = CoreChipError; + + fn name(&self) -> String { + self.name.to_string() + } + + fn generate_trace( + &self, + _: &ExecutionRecord, + _: &mut ExecutionRecord, + ) -> Result, Self::Error> { + Ok(RowMajorMatrix::new( + vec![F::ZERO; NUM_FIELD_DEN_HARNESS_COLS], + NUM_FIELD_DEN_HARNESS_COLS, + )) + } + + fn included(&self, _: &Self::Record) -> bool { + true + } + + fn local_only(&self) -> bool { + true + } + + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + FieldDenHarnessCols::::picus_info() + } +} + +impl BaseAir for FieldDenHarnessChip { + fn width(&self) -> usize { + NUM_FIELD_DEN_HARNESS_COLS + } +} + +impl Air for FieldDenHarnessChip +where + AB: ZKMAirBuilder, + Limbs::Limbs>: Copy, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0); + let local: &FieldDenHarnessCols = (*local).borrow(); + builder.assert_bool(local.is_real); + local.den.eval(builder, &local.a, &local.b, self.sign, local.is_real); + } +} + +pub struct FieldInnerProductHarnessChip { + pub name: &'static str, +} + +impl FieldInnerProductHarnessChip { + pub const fn new(name: &'static str) -> Self { + Self { name } + } +} + +impl MachineAir for FieldInnerProductHarnessChip { + type Record = ExecutionRecord; + type Program = Program; + type Error = CoreChipError; + + fn name(&self) -> String { + self.name.to_string() + } + + fn generate_trace( + &self, + _: &ExecutionRecord, + _: &mut ExecutionRecord, + ) -> Result, Self::Error> { + Ok(RowMajorMatrix::new( + vec![F::ZERO; NUM_FIELD_INNER_PRODUCT_HARNESS_COLS], + NUM_FIELD_INNER_PRODUCT_HARNESS_COLS, + )) + } + + fn included(&self, _: &Self::Record) -> bool { + true + } + + fn local_only(&self) -> bool { + true + } + + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + FieldInnerProductHarnessCols::::picus_info() + } +} + +impl BaseAir for FieldInnerProductHarnessChip { + fn width(&self) -> usize { + NUM_FIELD_INNER_PRODUCT_HARNESS_COLS + } +} + +impl Air for FieldInnerProductHarnessChip +where + AB: ZKMAirBuilder, + Limbs::Limbs>: Copy, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0); + let local: &FieldInnerProductHarnessCols = (*local).borrow(); + builder.assert_bool(local.is_real); + local.inner_product.eval(builder, &local.a, &local.b, local.is_real); + } +} diff --git a/crates/core/machine/src/operations/field/mod.rs b/crates/core/machine/src/operations/field/mod.rs index 5a4f1b4a5..388b69189 100644 --- a/crates/core/machine/src/operations/field/mod.rs +++ b/crates/core/machine/src/operations/field/mod.rs @@ -2,6 +2,7 @@ pub mod field_den; pub mod field_inner_product; pub mod field_op; pub mod field_sqrt; +pub mod harness; // pub mod params; pub mod range; pub mod util; diff --git a/crates/core/machine/src/operations/is_zero.rs b/crates/core/machine/src/operations/is_zero.rs index a899ace96..e44a213cb 100644 --- a/crates/core/machine/src/operations/is_zero.rs +++ b/crates/core/machine/src/operations/is_zero.rs @@ -39,7 +39,7 @@ impl IsZeroOperation { (a == F::ZERO) as u32 } - pub fn eval( + fn eval_exact( builder: &mut AB, a: AB::Expr, cols: IsZeroOperation, @@ -64,4 +64,17 @@ impl IsZeroOperation { // If the result is 1, then the input is 0. builder.when(is_real).when(cols.result).assert_zero(a); } + + pub fn eval( + builder: &mut AB, + a: AB::Expr, + cols: IsZeroOperation, + is_real: AB::Expr, + ) { + if builder.try_emit_is_zero_summary(a.clone(), cols.result.into(), is_real.clone()) { + return; + } + + Self::eval_exact(builder, a, cols, is_real); + } } diff --git a/crates/core/machine/src/operations/is_zero_word.rs b/crates/core/machine/src/operations/is_zero_word.rs index 10753e57d..cd4c600e4 100644 --- a/crates/core/machine/src/operations/is_zero_word.rs +++ b/crates/core/machine/src/operations/is_zero_word.rs @@ -46,7 +46,7 @@ impl IsZeroWordOperation { is_zero as u32 } - pub fn eval( + fn eval_exact( builder: &mut AB, a: Word, cols: IsZeroWordOperation, @@ -80,4 +80,23 @@ impl IsZeroWordOperation { ); builder_is_real.assert_eq(cols.result, cols.is_lower_half_zero * cols.is_upper_half_zero); } + + pub fn eval( + builder: &mut AB, + a: Word, + cols: IsZeroWordOperation, + is_real: AB::Expr, + ) { + if builder.try_emit_is_zero_word_summary( + a.clone(), + cols.is_lower_half_zero.into(), + cols.is_upper_half_zero.into(), + cols.result.into(), + is_real.clone(), + ) { + return; + } + + Self::eval_exact(builder, a, cols, is_real); + } } diff --git a/crates/core/machine/src/operations/koala_bear_word.rs b/crates/core/machine/src/operations/koala_bear_word.rs index 0f42c7c0d..b4c262b24 100644 --- a/crates/core/machine/src/operations/koala_bear_word.rs +++ b/crates/core/machine/src/operations/koala_bear_word.rs @@ -48,7 +48,7 @@ impl KoalaBearWordRangeChecker { self.and_most_sig_byte_decomp_0_to_6 * self.most_sig_byte_decomp[6]; } - pub fn range_check( + fn range_check_exact( builder: &mut AB, value: Word, cols: KoalaBearWordRangeChecker, @@ -103,4 +103,20 @@ impl KoalaBearWordRangeChecker { .when(cols.and_most_sig_byte_decomp_0_to_7) .assert_zero(value[0] + value[1] + value[2]); } + + pub fn range_check( + builder: &mut AB, + value: Word, + cols: KoalaBearWordRangeChecker, + is_real: AB::Expr, + ) { + if builder.try_emit_koala_bear_word_range_summary( + value.map(|limb| AB::Expr::zero() + limb), + is_real.clone(), + ) { + return; + } + + Self::range_check_exact(builder, value, cols, is_real); + } } diff --git a/crates/core/machine/src/operations/poseidon2/air.rs b/crates/core/machine/src/operations/poseidon2/air.rs index f5b84ade1..7ceffafbf 100644 --- a/crates/core/machine/src/operations/poseidon2/air.rs +++ b/crates/core/machine/src/operations/poseidon2/air.rs @@ -1,12 +1,16 @@ use std::array; +#[cfg(feature = "picus")] +use std::mem::size_of; use p3_air::PairBuilder; use p3_field::{FieldAlgebra, PrimeField32}; use p3_koala_bear::KoalaBear; use p3_poseidon2::matmul_internal; use zkm_primitives::RC_16_30_U32; -use zkm_stark::air::MachineAirBuilder; +use zkm_stark::air::{MachineAirBuilder, OperationSummaryAirBuilder}; +#[cfg(feature = "picus")] +use super::permutation::{permutation, Poseidon2Degree3Cols, Poseidon2Degree3Projection}; use super::{permutation::Poseidon2Cols, NUM_EXTERNAL_ROUNDS, NUM_INTERNAL_ROUNDS, WIDTH}; const INTERNAL_DIAG_MONTY_16: [KoalaBear; 16] = KoalaBear::new_array([ @@ -156,3 +160,48 @@ where builder.assert_eq(external_state[i], state[i].clone()) } } + +/// Evaluate the degree-3 Poseidon2 permutation, optionally replacing the exact +/// inlined constraints with a semantic projected submodule. +/// +/// The projected interface exposes only the caller-visible permutation +/// boundary: +/// - input: the first external-round state +/// - output: the final permutation state +/// +/// All intermediate round-state and S-box witness columns remain existential +/// inside the submodule. +pub fn eval_degree3(builder: &mut AB, local_row: &dyn Poseidon2Cols) +where + AB: MachineAirBuilder + PairBuilder + OperationSummaryAirBuilder, +{ + #[cfg(feature = "picus")] + let current_inputs: Vec = + local_row.external_rounds_state()[0].iter().map(|value| (*value).into()).collect(); + #[cfg(feature = "picus")] + let current_outputs: Vec = + local_row.perm_output().iter().map(|value| (*value).into()).collect(); + + #[cfg(feature = "picus")] + if builder.try_emit_projected_summary( + "Poseidon2Degree3Permutation", + &Poseidon2Degree3Projection::picus_projection_info(), + ¤t_inputs, + ¤t_outputs, + size_of::>(), + |builder, source_row| { + let projected = permutation::<_, 3>(source_row); + for r in 0..NUM_EXTERNAL_ROUNDS { + eval_external_round(builder, projected.as_ref(), r); + } + eval_internal_rounds(builder, projected.as_ref()); + }, + ) { + return; + } + + for r in 0..NUM_EXTERNAL_ROUNDS { + eval_external_round(builder, local_row, r); + } + eval_internal_rounds(builder, local_row); +} diff --git a/crates/core/machine/src/operations/poseidon2/permutation.rs b/crates/core/machine/src/operations/poseidon2/permutation.rs index 300384e29..ba356acce 100644 --- a/crates/core/machine/src/operations/poseidon2/permutation.rs +++ b/crates/core/machine/src/operations/poseidon2/permutation.rs @@ -5,6 +5,8 @@ use std::{ }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusProjection; use crate::operations::poseidon2::{NUM_EXTERNAL_ROUNDS, NUM_INTERNAL_ROUNDS, WIDTH}; use crate::utils::indices_arr; @@ -45,6 +47,28 @@ pub struct Poseidon2Degree3Cols { pub sbox_state: Poseidon2SBoxCols, } +/// Semantic Picus projection for the observable input/output contract of the +/// degree-3 Poseidon2 permutation witness. +/// +/// The full witness layout contains many intermediate round columns that are +/// internal to the permutation and should remain existential when Poseidon2 is +/// eventually emitted as an operation-level submodule. This projection keeps +/// only the caller-visible boundary: +/// - `state_in`: the first external-round state +/// - `state_out`: the final permutation output +/// +/// Projection `path = ...` points at the semantic source slice. The projected +/// field type determines the width, while the derive recursively takes the +/// first source column from the path. +#[cfg_attr(feature = "picus", derive(PicusProjection))] +#[cfg_attr(feature = "picus", picus_projection(source = Poseidon2Degree3Cols, col_map = POSEIDON2_DEGREE3_COL_MAP))] +pub struct Poseidon2Degree3Projection { + #[cfg_attr(feature = "picus", picus(input, path = state.external_rounds_state[0]))] + pub state_in: [u8; WIDTH], + #[cfg_attr(feature = "picus", picus(output, path = state.output_state))] + pub state_out: [u8; WIDTH], +} + /// A column layout for a poseidon2 permutation with degree 9 constraints. #[derive(AlignedBorrow, Clone, Copy)] #[repr(C)] diff --git a/crates/core/machine/src/operations/xor.rs b/crates/core/machine/src/operations/xor.rs index f62d6b4fb..4ce7bd8e5 100644 --- a/crates/core/machine/src/operations/xor.rs +++ b/crates/core/machine/src/operations/xor.rs @@ -1,12 +1,22 @@ +#[cfg(feature = "picus")] +use core::borrow::Borrow; +#[cfg(feature = "picus")] +use core::mem::{size_of, transmute}; + use p3_field::{Field, FieldAlgebra}; use zkm_core_executor::{ events::{ByteLookupEvent, ByteRecord}, ByteOpcode, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusProjection; use zkm_primitives::consts::WORD_SIZE; use zkm_stark::{air::ZKMAirBuilder, Word}; +#[cfg(feature = "picus")] +use crate::utils::indices_arr; + /// A set of columns needed to compute the xor of two words. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] #[repr(C)] @@ -15,6 +25,57 @@ pub struct XorOperation { pub value: Word, } +#[cfg(feature = "picus")] +const NUM_XOR_OPERATION_SUMMARY_COLS: usize = size_of::>(); + +#[cfg(feature = "picus")] +const XOR_OPERATION_SUMMARY_COL_MAP: XorOperationSummaryCols = + make_xor_operation_summary_col_map(); + +#[cfg(feature = "picus")] +const fn make_xor_operation_summary_col_map() -> XorOperationSummaryCols { + let indices_arr = indices_arr::(); + unsafe { + transmute::<[usize; NUM_XOR_OPERATION_SUMMARY_COLS], XorOperationSummaryCols>( + indices_arr, + ) + } +} + +/// Hidden witness layout used when Picus emits the exact XOR AIR as a local +/// auxiliary module. +/// +/// The caller should reason only about the semantic interface `(a, b, is_real) +/// -> value`; the full internal witness row remains existential inside the +/// summarized module. +#[derive(AlignedBorrow, Clone, Copy)] +#[repr(C)] +#[cfg(feature = "picus")] +struct XorOperationSummaryCols { + pub a: Word, + pub b: Word, + pub is_real: T, + pub cols: XorOperation, +} + +#[cfg(feature = "picus")] +#[cfg_attr(feature = "picus", derive(PicusProjection))] +#[cfg_attr(feature = "picus", picus_projection( + source = XorOperationSummaryCols, + col_map = XOR_OPERATION_SUMMARY_COL_MAP +))] +#[allow(dead_code)] +struct XorOperationSummaryProjection { + #[cfg_attr(feature = "picus", picus(input, path = a))] + pub a: Word, + #[cfg_attr(feature = "picus", picus(input, path = b))] + pub b: Word, + #[cfg_attr(feature = "picus", picus(input, path = is_real))] + pub is_real: u8, + #[cfg_attr(feature = "picus", picus(output, path = cols.value))] + pub value: Word, +} + impl XorOperation { pub fn populate(&mut self, record: &mut impl ByteRecord, x: u32, y: u32) -> u32 { let expected = x ^ y; @@ -36,8 +97,7 @@ impl XorOperation { expected } - #[allow(unused_variables)] - pub fn eval( + fn eval_exact( builder: &mut AB, a: Word, b: Word, @@ -54,4 +114,51 @@ impl XorOperation { ); } } + + #[allow(unused_variables)] + pub fn eval( + builder: &mut AB, + a: Word, + b: Word, + cols: XorOperation, + is_real: AB::Var, + ) { + let is_real_expr = AB::Expr::zero() + is_real; + let mut current_inputs: Vec = Vec::with_capacity(WORD_SIZE * 2 + 1); + for limb in a.0 { + current_inputs.push(limb.into()); + } + for limb in b.0 { + current_inputs.push(limb.into()); + } + current_inputs.push(is_real_expr.clone()); + + let current_outputs: Vec = + cols.value.0.iter().map(|limb| (*limb).into()).collect(); + + // Keep the exact byte-lookup AIR, but hide it behind a local projected + // submodule so callers only see the semantic word-level boundary. + // + // This operation is only functional when `is_real = 1`; otherwise the + // exact AIR leaves `cols.value` unconstrained. Only outline it once the + // guard has already been specialized to one. + #[cfg(feature = "picus")] + if builder.is_known_one(&is_real_expr) + && builder.try_emit_projected_summary( + "XorOperation", + &XorOperationSummaryProjection::picus_projection_info(), + ¤t_inputs, + ¤t_outputs, + size_of::>(), + |builder, source_row| { + let source: &XorOperationSummaryCols = (*source_row).borrow(); + Self::eval_exact(builder, source.a, source.b, source.cols, source.is_real); + }, + ) + { + return; + } + + Self::eval_exact(builder, a, b, cols, is_real); + } } diff --git a/crates/core/machine/src/syscall/chip.rs b/crates/core/machine/src/syscall/chip.rs index b61debf21..413eeab41 100644 --- a/crates/core/machine/src/syscall/chip.rs +++ b/crates/core/machine/src/syscall/chip.rs @@ -14,7 +14,11 @@ use p3_maybe_rayon::prelude::ParallelIterator; use zkm_core_executor::events::{ByteRecord, GlobalLookupEvent, PrecompileEvent}; use zkm_core_executor::{events::SyscallEvent, ByteOpcode, ExecutionRecord, Program}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::air::AirLookup; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{LookupScope, MachineAir, ZKMAirBuilder}; use zkm_stark::LookupKind; @@ -59,24 +63,20 @@ impl SyscallChip { } /// The column layout for the chip. -/// -/// `arg1` and `arg2` are NOT stored as columns. They are derived inline as -/// `arg1_lo + arg1_hi * 65536` to avoid redundant columns while keeping the -/// reduced field element available for local `send_syscall`/`receive_syscall`. -/// -/// **Soundness**: `arg1_lo/hi` and `arg2_lo/hi` are U16Range-checked inside -/// `send_syscall_result_packed` (see `crates/stark/src/air/builder.rs`). -/// Any chip using this interaction gets range-checked half-words automatically. #[derive(AlignedBorrow, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct SyscallCols { /// The shard number of the syscall. + #[cfg_attr(feature = "picus", picus(input))] pub shard: T, /// The clk of the syscall. + #[cfg_attr(feature = "picus", picus(input))] pub clk: T, /// The syscall_id of the syscall. + #[cfg_attr(feature = "picus", picus(input))] pub syscall_id: T, /// Half-word packed arg1: low 16 bits (byte0 + byte1 * 256). @@ -85,20 +85,27 @@ pub struct SyscallCols { /// arg1/arg2 through receive_syscall and don't use the half-words. /// If a new precompile needs byte-level argument access, it should use /// receive_syscall_result_packed to get these half-words. + #[cfg_attr(feature = "picus", picus(input))] pub arg1_lo: T, /// Half-word packed arg1: high 16 bits (byte2 + byte3 * 256). + #[cfg_attr(feature = "picus", picus(input))] pub arg1_hi: T, /// Half-word packed arg2: low 16 bits. + #[cfg_attr(feature = "picus", picus(input))] pub arg2_lo: T, /// Half-word packed arg2: high 16 bits. + #[cfg_attr(feature = "picus", picus(input))] pub arg2_hi: T, /// Half-word packed result (lo = byte0 + byte1*256, hi = byte2 + byte3*256). + #[cfg_attr(feature = "picus", picus(output))] pub result_lo: T, + #[cfg_attr(feature = "picus", picus(output))] pub result_hi: T, /// Whether the syscall is a linux syscall. + #[cfg_attr(feature = "picus", picus(input))] pub is_linux: T, pub is_real: T, @@ -115,6 +122,15 @@ impl MachineAir for SyscallChip { format!("Syscall{}", self.shard_kind).to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + SyscallCols::::picus_info() + } + + fn local_only(&self) -> bool { + true + } + fn generate_dependencies( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/instructions/air.rs b/crates/core/machine/src/syscall/instructions/air.rs index 5f1108854..88e2f0c13 100644 --- a/crates/core/machine/src/syscall/instructions/air.rs +++ b/crates/core/machine/src/syscall/instructions/air.rs @@ -162,14 +162,20 @@ impl SyscallInstrsChip { builder.when(AB::Expr::one() - local.is_real).assert_zero(send_to_table.clone()); // KoalaBear range checks on op_b and op_c, activated by stored flags. - // op_b_check = 1 when send_to_table || is_halt (covers both syscall bridge and exit code). - // op_c_check = 1 when send_to_table || is_commit_deferred_proofs (covers bridge and digest). + // Only required on the precompile bridge (where args travel as a single reduced field + // element), for `is_halt` (exit code is reduced), and for `is_commit_deferred_proofs` + // (digest element is reduced). Linux syscalls travel via half-word packed columns in + // `SyscallChip`, which are U16-range-checked there — reduce() collision is impossible, + // so the KoalaBear range check is not needed (and would reject legal u32 args like + // AT_FDCWD = 0xFFFFFF9C). + let send_to_precompile: AB::Expr = get_send_table::(local).into(); + let op_b_check_active: AB::Expr = send_to_precompile.clone() + local.is_halt.into(); + let op_c_check_active: AB::Expr = + send_to_precompile + local.is_commit_deferred_proofs.result.into(); builder.assert_bool(local.op_b_check); builder.assert_bool(local.op_c_check); - builder.when(send_to_table.clone()).assert_one(local.op_b_check); - builder.when(local.is_halt).assert_one(local.op_b_check); - builder.when(send_to_table.clone()).assert_one(local.op_c_check); - builder.when(local.is_commit_deferred_proofs.result).assert_one(local.op_c_check); + builder.when(op_b_check_active).assert_one(local.op_b_check); + builder.when(op_c_check_active).assert_one(local.op_c_check); builder.when_not(local.is_real).assert_zero(local.op_b_check); builder.when_not(local.is_real).assert_zero(local.op_c_check); diff --git a/crates/core/machine/src/syscall/instructions/columns.rs b/crates/core/machine/src/syscall/instructions/columns.rs index 17fa19ad0..358961be4 100644 --- a/crates/core/machine/src/syscall/instructions/columns.rs +++ b/crates/core/machine/src/syscall/instructions/columns.rs @@ -1,12 +1,17 @@ use std::mem::size_of; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::{air::PV_DIGEST_NUM_WORDS, Word}; use crate::operations::{IsZeroOperation, KoalaBearWordRangeChecker}; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_SYSCALL_INSTR_COLS: usize = size_of::>(); #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct SyscallInstrColumns { pub pc: T, diff --git a/crates/core/machine/src/syscall/instructions/trace.rs b/crates/core/machine/src/syscall/instructions/trace.rs index 5cc52c7da..3e4d2ae91 100644 --- a/crates/core/machine/src/syscall/instructions/trace.rs +++ b/crates/core/machine/src/syscall/instructions/trace.rs @@ -11,6 +11,8 @@ use zkm_core_executor::{ ExecutionRecord, Program, }; use zkm_stark::air::MachineAir; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use crate::{ utils::{next_power_of_two, zeroed_f_vec}, @@ -33,6 +35,11 @@ impl MachineAir for SyscallInstrsChip { "SyscallInstrs".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + SyscallInstrColumns::::picus_info() + } + fn num_rows(&self, input: &Self::Record) -> Option { let nb_rows = next_power_of_two( input.syscall_events.len(), @@ -115,7 +122,6 @@ impl SyscallInstrsChip { cols.is_sys_linux = F::from_bool(event.a_record.prev_value & 0x0ff00 != 0); let prev_a_bytes = event.a_record.prev_value.to_le_bytes(); - let send_to_table = (prev_a_bytes[1] != 0) || (prev_a_bytes[2] == 1); let is_halt_val = cols.is_halt == F::ONE; // Populate is_prev_a1_zero for bidirectional is_sys_linux constraint. @@ -161,10 +167,17 @@ impl SyscallInstrsChip { } // Populate unified KoalaBear range check flags and columns. + // Activate the KoalaBear range check only when the value travels as a single reduced + // field element: the precompile bridge (prev_a_bytes[2] == 1), `is_halt` (exit code), + // and `is_commit_deferred_proofs` (digest). Linux syscall args travel via half-word + // packed columns in `SyscallChip` (U16-range-checked), so reduce() collisions are + // impossible there and the KoalaBear constraint must not be applied — otherwise legal + // u32 args like AT_FDCWD = 0xFFFFFF9C fail the check. + let send_to_precompile = prev_a_bytes[2] == 1; let is_commit_deferred = syscall_id == F::from_canonical_u32(SyscallCode::COMMIT_DEFERRED_PROOFS.syscall_id()); - let op_b_needs_check = send_to_table || is_halt_val; - let op_c_needs_check = send_to_table || is_commit_deferred; + let op_b_needs_check = send_to_precompile || is_halt_val; + let op_c_needs_check = send_to_precompile || is_commit_deferred; if op_b_needs_check { cols.op_b_check = F::ONE; diff --git a/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/air.rs b/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/air.rs index acc3966ec..1f9b0128a 100644 --- a/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/air.rs +++ b/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/air.rs @@ -52,10 +52,20 @@ impl BooleanCircuitGarbleChip { builder: &mut AB, local: &BooleanCircuitGarbleCols, ) { + // In a true single-row trace, this chip only has the prelude row + // (num_gates + delta read), never a gate row. + let single_row_phase = builder.is_first_row() * builder.is_last_row(); + builder.when(single_row_phase).assert_one(local.is_first_row); + builder.assert_bool(local.is_real); + builder.assert_bool(local.is_first_row); builder.assert_bool(local.is_first_gate); builder.assert_bool(local.not_last_gate); builder.assert_bool(local.is_gate); + builder.assert_bool(local.checks_acc); + builder.assert_eq(local.is_first_gate * local.is_gate, local.is_first_gate); + builder.assert_eq(local.is_last_gate * local.is_gate, local.is_last_gate); + builder.assert_eq(local.not_last_gate * local.is_gate, local.not_last_gate); builder.assert_zero(local.is_last_gate * local.is_first_gate); builder.when(local.is_gate).assert_one(local.is_last_gate + local.not_last_gate); builder.assert_bool(local.gate_type[0]); @@ -107,6 +117,18 @@ impl BooleanCircuitGarbleChip { &local.result_mem, local.is_last_gate, ); + + // The syscall writes a boolean result (as u32) at the final gate. + builder + .when(local.is_last_gate) + .assert_eq(local.result_mem.access.value[0], local.checks_acc * local.checks[2]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.access.value[1]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.access.value[2]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.access.value[3]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.prev_value[0]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.prev_value[1]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.prev_value[2]); + builder.when(local.is_last_gate).assert_zero(local.result_mem.prev_value[3]); } fn eval_logic_check( @@ -115,6 +137,7 @@ impl BooleanCircuitGarbleChip { local: &BooleanCircuitGarbleCols, next: &BooleanCircuitGarbleCols, ) { + let transition_continuation = local.not_last_gate * local.is_gate; // eval XOR operations for i in 0..4 { let h0_id = 1 + i; @@ -178,9 +201,13 @@ impl BooleanCircuitGarbleChip { local.checks[2], local.is_equal_words[3].is_diff_zero.result * local.checks[1], ); + builder.when(local.is_first_row).assert_zero(local.checks[0]); + builder.when(local.is_first_row).assert_zero(local.checks[1]); + builder.when(local.is_first_row).assert_zero(local.checks[2]); + builder.when(local.is_first_row).assert_one(local.checks_acc); builder - .when(local.not_last_gate) - .assert_eq(next.checks[3], local.checks[3] * next.checks[2]); + .when(transition_continuation) + .assert_eq(next.checks_acc, local.checks_acc * local.checks[2]); } fn eval_transition( @@ -189,42 +216,68 @@ impl BooleanCircuitGarbleChip { local: &BooleanCircuitGarbleCols, next: &BooleanCircuitGarbleCols, ) { + let transition_continuation = local.not_last_gate * local.is_gate; let bytes_shift = AB::F::from_canonical_u32(256); let num_gates = local.gates_input_mem[0].access.value.0[0] + local.gates_input_mem[0].access.value.0[1] * bytes_shift + local.gates_input_mem[0].access.value.0[2] * bytes_shift * bytes_shift + local.gates_input_mem[0].access.value.0[3] * bytes_shift * bytes_shift * bytes_shift; - builder.when_first_row().assert_eq(local.gates_num, num_gates.clone()); + builder.when(local.is_first_row).assert_eq(local.gates_num, num_gates.clone()); for i in 0..4 { let delta_i = local.gates_input_mem[i + 1].access.value; for j in 0..4 { - builder.when_first_row().assert_eq(local.delta[i][j], delta_i[j]); + builder.when(local.is_first_row).assert_eq(local.delta[i][j], delta_i[j]); } } - let gate_type_value = local.gate_type[0] * AB::Expr::zero() + local.gate_type[1]; - builder - .when(local.is_gate) - .assert_eq(gate_type_value * AB::Expr::from_canonical_u32(OR_GATE_ID), num_gates); + let gate_type_value = + local.gate_type[0] + local.gate_type[1] * AB::Expr::from_canonical_u32(OR_GATE_ID); + builder.when(local.is_gate).assert_eq(gate_type_value, num_gates); builder.when(local.is_first_gate).assert_zero(local.gate_id); builder.when(local.is_last_gate).assert_eq(local.gates_num - AB::F::ONE, local.gate_id); - builder.when(local.not_last_gate).assert_eq(local.gate_id + AB::F::ONE, next.gate_id); + builder + .when(transition_continuation.clone()) + .assert_eq(local.gate_id + AB::F::ONE, next.gate_id); - builder.when(local.not_last_gate * local.is_gate).assert_eq( + // Bridge the prelude row (num_gates + delta read) to the first gate row. + builder + .when(local.is_first_row) + .assert_eq(next.input_address, local.input_address + AB::F::from_canonical_u32(20)); + builder.when(local.is_first_row).assert_eq(next.output_address, local.output_address); + builder.when(local.is_first_row).assert_eq(next.shard, local.shard); + builder.when(local.is_first_row).assert_eq(next.clk, local.clk); + builder.when(local.is_first_row).assert_eq(next.gates_num, local.gates_num); + builder.when(local.is_first_row).assert_zero(next.is_first_row); + builder.when(local.is_first_row).assert_eq(next.is_gate, next.is_first_gate); + builder.when(local.is_first_row).assert_eq(next.checks_acc, next.is_gate); + // Continue with next gate row only when explicitly in same-event continuation. + builder.when(transition_continuation.clone()).assert_one(next.is_gate); + builder.when(transition_continuation.clone()).assert_zero(next.is_first_row); + builder.when(transition_continuation.clone()).assert_zero(next.is_first_gate); + builder + .when(transition_continuation.clone()) + .assert_eq(next.output_address, local.output_address); + builder.when(transition_continuation.clone()).assert_eq(next.shard, local.shard); + builder.when(transition_continuation.clone()).assert_eq(next.clk, local.clk); + for i in 0..4 { + for j in 0..4 { + builder.when(local.is_first_row).assert_eq(local.delta[i][j], next.delta[i][j]); + } + } + + builder.when(transition_continuation.clone()).assert_eq( local.input_address + AB::F::from_canonical_usize(GATE_INFO_BYTES * 4), next.input_address, ); - builder - .when(local.not_last_gate * local.is_gate) - .assert_eq(local.gates_num, next.gates_num); + builder.when(transition_continuation.clone()).assert_eq(local.gates_num, next.gates_num); for i in 0..4 { for j in 0..4 { builder - .when(local.not_last_gate * local.is_gate) + .when(transition_continuation.clone()) .assert_eq(local.delta[i][j], next.delta[i][j]); } } diff --git a/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/columns.rs b/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/columns.rs index 149cd23e0..7c5711426 100644 --- a/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/columns.rs +++ b/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/columns.rs @@ -1,34 +1,52 @@ use crate::memory::{MemoryReadCols, MemoryWriteCols}; use crate::operations::{IsEqualWordOperation, XorOperation}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::Word; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; /// BooleanCircuitGarbleCols is the column layout for the Boolean Circuit Garble. /// The number of rows equal to the number of gates #[derive(AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct BooleanCircuitGarbleCols { + #[cfg_attr(feature = "picus", picus(transition_input))] pub shard: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub clk: T, pub is_real: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub input_address: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub output_address: T, - pub is_first_row: T, // The first row contains gates_num and delta + pub is_first_row: T, // The first row contains gates_num and delt + #[cfg_attr(feature = "picus", picus(transition_input))] pub is_gate: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub is_first_gate: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub is_last_gate: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub not_last_gate: T, // from first gate -> (last - 1)-th gate pub gate_type: [T; 2], + #[cfg_attr(feature = "picus", picus(transition_input))] pub gate_id: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub gates_num: T, - pub delta: [Word; 4], // [u8; 16] + #[cfg_attr(feature = "picus", picus(transition_input))] + pub delta: [Word; 4], // [u8; 16] pub gates_input_mem: [MemoryReadCols; 17], // gate_type, h0, h1, label_b, expected_ciphertext pub result_mem: MemoryWriteCols, pub aux1: [XorOperation; 4], // h1 ^ h0 pub aux2: [XorOperation; 4], // h1 ^ h0 ^ label_b pub aux3: [XorOperation; 4], // h1 ^ h0 ^ label_b ^ delta pub is_equal_words: [IsEqualWordOperation; 4], // computed ciphertext == expected_ciphertext - pub checks: [T; 4], // check result for each pair of is_equal_words + pub checks: [T; 3], // row-local chaining result + #[cfg_attr(feature = "picus", picus(input, transition_input))] + pub checks_acc: T, // cross-row accumulated check state } pub const NUM_BOOLEAN_CIRCUIT_GARBLE_COLS: usize = size_of::>(); diff --git a/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/trace.rs b/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/trace.rs index cffd6687c..7e558f626 100644 --- a/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/trace.rs +++ b/crates/core/machine/src/syscall/precompiles/boolean_circuit_garble/trace.rs @@ -18,6 +18,8 @@ use zkm_core_executor::events::{ }; use zkm_core_executor::syscalls::SyscallCode; use zkm_core_executor::{ExecutionRecord, Program}; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::MachineAir; impl MachineAir for BooleanCircuitGarbleChip { @@ -29,6 +31,11 @@ impl MachineAir for BooleanCircuitGarbleChip { "BooleanCircuitGarble".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + BooleanCircuitGarbleCols::::picus_info() + } + fn generate_dependencies( &self, input: &Self::Record, @@ -120,6 +127,7 @@ impl BooleanCircuitGarbleChip { cols.input_address = F::from_canonical_u32(input_address); cols.output_address = F::from_canonical_u32(event.output_addr); cols.gates_num = F::from_canonical_u32(gates_num as u32); + cols.checks_acc = F::ONE; for i in 0..4 { let delta_i_bytes = event.delta[i].to_le_bytes(); cols.delta[i] @@ -208,7 +216,7 @@ impl BooleanCircuitGarbleChip { cols.checks[0] = F::from_canonical_u32(check_u32s[1]); cols.checks[1] = F::from_canonical_u32(check_u32s[2]); cols.checks[2] = F::from_canonical_u32(check_u32s[3]); - cols.checks[3] = F::from_canonical_u32(check_u32s[3] * (pre_check as u32)); + cols.checks_acc = F::from_bool(pre_check); pre_check = pre_check && (check_u32s[3] == 1); // if this is the last gate, write result diff --git a/crates/core/machine/src/syscall/precompiles/edwards/ed_add.rs b/crates/core/machine/src/syscall/precompiles/edwards/ed_add.rs index 4f37df705..86eb09415 100644 --- a/crates/core/machine/src/syscall/precompiles/edwards/ed_add.rs +++ b/crates/core/machine/src/syscall/precompiles/edwards/ed_add.rs @@ -24,6 +24,10 @@ use zkm_curves::{ AffinePoint, EllipticCurve, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{BaseAirBuilder, LookupScope, MachineAir, ZKMAirBuilder}; use crate::{ @@ -40,6 +44,7 @@ pub const NUM_ED_ADD_COLS: usize = size_of::>(); /// Right now the number of limbs is assumed to be a constant, although this could be macro-ed /// or made generic in the future. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct EdAddAssignCols { pub is_real: T, @@ -111,6 +116,11 @@ impl MachineAir for Ed "EdAddAssign".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + EdAddAssignCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/precompiles/edwards/ed_decompress.rs b/crates/core/machine/src/syscall/precompiles/edwards/ed_decompress.rs index a50a7b238..14422b9e9 100644 --- a/crates/core/machine/src/syscall/precompiles/edwards/ed_decompress.rs +++ b/crates/core/machine/src/syscall/precompiles/edwards/ed_decompress.rs @@ -25,6 +25,10 @@ use zkm_curves::{ CurveError, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{BaseAirBuilder, LookupScope, MachineAir, ZKMAirBuilder}; use crate::{ @@ -41,6 +45,7 @@ pub const NUM_ED_DECOMPRESS_COLS: usize = size_of::>(); /// /// After `EdDecompress`, the first 32 bytes of the slice are overwritten with the decompressed X. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct EdDecompressCols { pub is_real: T, @@ -210,6 +215,11 @@ impl MachineAir for EdDecompressChip PicusInfo { + EdDecompressCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/precompiles/fptower/fp.rs b/crates/core/machine/src/syscall/precompiles/fptower/fp.rs index e53e06b8d..147bcb55f 100644 --- a/crates/core/machine/src/syscall/precompiles/fptower/fp.rs +++ b/crates/core/machine/src/syscall/precompiles/fptower/fp.rs @@ -21,7 +21,11 @@ use zkm_curves::{ weierstrass::{FieldType, FpOpField}, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::air::{BaseAirBuilder, LookupScope, MachineAir, Polynomial, ZKMAirBuilder}; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; use crate::{ memory::{value_as_limbs, MemoryReadCols, MemoryWriteCols}, @@ -39,13 +43,17 @@ pub struct FpOpChip

{ /// A set of columns for the FpAdd operation. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct FpOpCols { pub is_real: T, pub shard: T, pub clk: T, + #[cfg_attr(feature = "picus", picus(selector))] pub is_add: T, + #[cfg_attr(feature = "picus", picus(selector))] pub is_sub: T, + #[cfg_attr(feature = "picus", picus(selector))] pub is_mul: T, pub x_ptr: T, pub y_ptr: T, @@ -87,6 +95,11 @@ impl MachineAir for FpOpChip

{ } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + FpOpCols::::picus_info() + } + fn generate_trace( &self, input: &Self::Record, diff --git a/crates/core/machine/src/syscall/precompiles/fptower/fp2_addsub.rs b/crates/core/machine/src/syscall/precompiles/fptower/fp2_addsub.rs index e6bbb265b..efa84446b 100644 --- a/crates/core/machine/src/syscall/precompiles/fptower/fp2_addsub.rs +++ b/crates/core/machine/src/syscall/precompiles/fptower/fp2_addsub.rs @@ -22,6 +22,10 @@ use zkm_curves::{ weierstrass::{FieldType, FpOpField}, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{BaseAirBuilder, LookupScope, MachineAir, Polynomial, ZKMAirBuilder}; use crate::{ @@ -36,6 +40,7 @@ pub const fn num_fp2_addsub_cols() -> usize { /// A set of columns for the Fp2AddSub operation. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct Fp2AddSubAssignCols { pub is_real: T, @@ -90,6 +95,11 @@ impl MachineAir for Fp2AddSubAssignChip

{ } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + Fp2AddSubAssignCols::::picus_info() + } + fn generate_trace( &self, input: &Self::Record, diff --git a/crates/core/machine/src/syscall/precompiles/fptower/fp2_mul.rs b/crates/core/machine/src/syscall/precompiles/fptower/fp2_mul.rs index 4445b93d6..031d77d4d 100644 --- a/crates/core/machine/src/syscall/precompiles/fptower/fp2_mul.rs +++ b/crates/core/machine/src/syscall/precompiles/fptower/fp2_mul.rs @@ -22,6 +22,10 @@ use zkm_curves::{ weierstrass::{FieldType, FpOpField}, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{BaseAirBuilder, LookupScope, MachineAir, Polynomial, ZKMAirBuilder}; use crate::{ @@ -36,6 +40,7 @@ pub const fn num_fp2_mul_cols() -> usize { /// A set of columns for the Fp2Mul operation. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct Fp2MulAssignCols { pub is_real: T, @@ -133,6 +138,11 @@ impl MachineAir for Fp2MulAssignChip

{ } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + Fp2MulAssignCols::::picus_info() + } + fn generate_trace( &self, input: &Self::Record, diff --git a/crates/core/machine/src/syscall/precompiles/keccak_sponge/air.rs b/crates/core/machine/src/syscall/precompiles/keccak_sponge/air.rs index b0eba2375..b10cd32a5 100644 --- a/crates/core/machine/src/syscall/precompiles/keccak_sponge/air.rs +++ b/crates/core/machine/src/syscall/precompiles/keccak_sponge/air.rs @@ -1,6 +1,8 @@ use crate::air::{MemoryAirBuilder, WordAirBuilder}; use crate::memory::MemoryCols; use crate::operations::XorOperation; +#[cfg(feature = "picus")] +use crate::syscall::precompiles::keccak_sponge::columns::KeccakPermutationProjection; use crate::syscall::precompiles::keccak_sponge::columns::{ KeccakSpongeCols, NUM_KECCAK_SPONGE_COLS, }; @@ -69,6 +71,7 @@ where builder.when_last_row().assert_zero(local.is_real); // Xor + let not_read_block = AB::Expr::one() - local.read_block; for i in 0..KECCAK_GENERAL_RATE_U32S { XorOperation::::eval( builder, @@ -77,13 +80,66 @@ where local.xored_general_rate[i], local.read_block, ); + for j in 0..4 { + builder + .when(not_read_block.clone()) + .assert_zero(local.xored_general_rate[i].value[j]); + builder + .when(not_read_block.clone()) + .assert_zero(local.block_mem[i].access.value[j]); + } + } + + // Range-constrain the sponge state bytes. + // + // `original_state` is interpreted as bytes when building u16/u64 limbs + // (including for `next.original_state` in absorbed-state transitions). + // Enforce each byte is in [0, 255] to prevent unconstrained field limbs + // from satisfying packed equalities spuriously. + let mut original_state_bytes = Vec::with_capacity(KECCAK_STATE_U32S * 4); + for i in 0..KECCAK_STATE_U32S { + for j in 0..4 { + original_state_bytes.push(local.original_state[i][j]); + } + } + builder.slice_range_check_u8(&original_state_bytes, local.is_real); + + // Range-constrain memory words that are interpreted as byte-packed u16/u64 limbs. + // `input_length_mem` is constrained against `input_len` on all real rows, + // so keep it byte-range constrained on all real rows as well. + // `block_mem` is only read on `read_block`, + // and `output_mem` is only used on `write_output`. + let mut input_len_bytes = Vec::with_capacity(4); + for j in 0..4 { + input_len_bytes.push(local.input_length_mem.value()[j]); + } + builder.slice_range_check_u8(&input_len_bytes, local.is_real); + + let mut block_mem_bytes = Vec::with_capacity(KECCAK_GENERAL_RATE_U32S * 4); + for i in 0..KECCAK_GENERAL_RATE_U32S { + for j in 0..4 { + block_mem_bytes.push(local.block_mem[i].access.value[j]); + } + } + builder.slice_range_check_u8(&block_mem_bytes, local.read_block); + + let mut output_mem_bytes = Vec::with_capacity(KECCAK_GENERAL_OUTPUT_U32S * 4); + for i in 0..KECCAK_GENERAL_OUTPUT_U32S { + for j in 0..4 { + output_mem_bytes.push(local.output_mem[i].value()[j]); + } } + builder.slice_range_check_u8(&output_mem_bytes, local.write_output); // Constrain the absorbed bytes builder .when_transition() - .when(not_final_step) + .when(not_final_step.clone()) .assert_eq(local.already_absorbed_u32s, next.already_absorbed_u32s); + builder + .when_transition() + .when(not_final_step) + .assert_eq(local.input_address, next.input_address); // If this is the first block, absorbed bytes should be 0 builder.when(first_block).assert_eq(local.already_absorbed_u32s, AB::Expr::zero()); // If this is the final block, absorbed bytes should be equal to the input length - KECCAK_GENERAL_RATE_U32S @@ -104,14 +160,103 @@ where next.input_address - AB::Expr::from_canonical_u32(KECCAK_GENERAL_RATE_U32S as u32 * 4), ); - // Eval the plonky3 keccak air - let mut sub_builder = - SubAirBuilder::::new(builder, 0..NUM_KECCAK_COLS); - self.p3_keccak.eval(&mut sub_builder); + // Outside the absorbed-edge transition, keep `original_state` stable + // across rows. Exclude the final output row because its successor is a + // padding row. + let keep_original_state = + (AB::Expr::one() - local.is_absorbed) * (AB::Expr::one() - local.write_output); + let mut keep_transition_builder = builder.when_transition(); + let mut keep_state_builder = keep_transition_builder.when(keep_original_state); + for i in 0..KECCAK_STATE_U32S { + keep_state_builder.assert_word_eq(local.original_state[i], next.original_state[i]); + } + + // Eval the plonky3 keccak air. Picus can hide the full sub-AIR behind a + // semantic boundary; other builders continue to inline the exact + // `SubAirBuilder` path. + #[cfg(feature = "picus")] + let current_inputs = self.keccak_summary_inputs::(local); + #[cfg(feature = "picus")] + let current_outputs = self.keccak_summary_outputs::(local); + #[cfg(feature = "picus")] + if !builder.try_emit_hidden_subair_summary( + "KeccakAir", + &KeccakPermutationProjection::picus_projection_info(), + ¤t_inputs, + ¤t_outputs, + NUM_KECCAK_COLS, + false, + |nested_builder| self.p3_keccak.eval(nested_builder), + ) { + let mut sub_builder = + SubAirBuilder::::new(builder, 0..NUM_KECCAK_COLS); + self.p3_keccak.eval(&mut sub_builder); + } + #[cfg(not(feature = "picus"))] + { + let mut sub_builder = + SubAirBuilder::::new(builder, 0..NUM_KECCAK_COLS); + self.p3_keccak.eval(&mut sub_builder); + } } } impl KeccakSpongeChip { + /// Flatten the visible Keccak-f input state from the embedded sub-AIR. + /// + /// The surrounding sponge AIR ties these lanes to memory/original-state + /// data, so they must remain visible caller inputs even when the full + /// Keccak round system is summarized as a hidden submodule. + #[cfg(feature = "picus")] + fn keccak_summary_inputs( + &self, + local: &KeccakSpongeCols, + ) -> Vec { + let mut inputs = Vec::with_capacity(25 * U64_LIMBS); + for y in 0..5 { + for x in 0..5 { + for limb in 0..U64_LIMBS { + inputs.push(local.keccak.a[y][x][limb].into()); + } + } + } + inputs + } + + /// Flatten the caller-visible outputs of the embedded Keccak-f sub-AIR. + /// + /// The sponge AIR depends on the round-position flags as well as the + /// post-round `A'''` state. The `(0, 0)` lane is stored separately from the + /// remaining 24 lanes, so the flattened order mirrors the projection + /// metadata exactly: + /// 1. `first_step` + /// 2. `final_step` + /// 3. `A'''[0, 0]` + /// 4. the remaining 24 `A'''` lanes in witness layout order + #[cfg(feature = "picus")] + fn keccak_summary_outputs( + &self, + local: &KeccakSpongeCols, + ) -> Vec { + let mut outputs = Vec::with_capacity(2 + 25 * U64_LIMBS); + outputs.push(local.keccak.step_flags[0].into()); + outputs.push(local.keccak.step_flags[NUM_ROUNDS - 1].into()); + for limb in 0..U64_LIMBS { + outputs.push(local.keccak.a_prime_prime_prime(0, 0, limb).into()); + } + for y in 0..5 { + for x in 0..5 { + if y == 0 && x == 0 { + continue; + } + for limb in 0..U64_LIMBS { + outputs.push(local.keccak.a_prime_prime_prime(y, x, limb).into()); + } + } + } + outputs + } + fn eval_flags(&self, builder: &mut AB, local: &KeccakSpongeCols) { let first_block = local.is_first_input_block; let final_block = local.is_final_input_block; @@ -120,9 +265,20 @@ impl KeccakSpongeChip { let first_step = local.keccak.step_flags[0]; let final_step = local.keccak.step_flags[NUM_ROUNDS - 1]; + // Defensive constraints for the summarized Keccak sub-AIR boundary: + // enforce booleanity and mutual exclusion of first/final step flags. + // This prevents degenerate witnesses where a single row is both + // first-round and final-round when summary internals are hidden. + builder.when(local.is_real).assert_bool(first_step); + builder.when(local.is_real).assert_bool(final_step); + builder.when(local.is_real).assert_zero(first_step * final_step); + // receive syscall builder.assert_eq(first_block * first_step * local.is_real, local.receive_syscall); + // Input block memory is only read on the first Keccak round of a real row. + builder.assert_eq(local.read_block, first_step * local.is_real); + // write output flag builder.assert_eq(final_block * final_step * local.is_real, local.write_output); @@ -143,6 +299,11 @@ impl KeccakSpongeChip { &local.input_length_mem, local.receive_syscall, ); + // Bind the scalar `input_len` to the memory-provided 4-byte word on + // the syscall row. + builder + .when(local.receive_syscall) + .assert_eq(local.input_len, local.input_length_mem.value().reduce::()); // Verify the input length has not changed builder .when(local.is_real) diff --git a/crates/core/machine/src/syscall/precompiles/keccak_sponge/columns.rs b/crates/core/machine/src/syscall/precompiles/keccak_sponge/columns.rs index aeee3e3d9..b72663261 100644 --- a/crates/core/machine/src/syscall/precompiles/keccak_sponge/columns.rs +++ b/crates/core/machine/src/syscall/precompiles/keccak_sponge/columns.rs @@ -1,35 +1,52 @@ -use core::mem::size_of; +use core::mem::{size_of, transmute}; use crate::memory::{MemoryReadCols, MemoryWriteCols}; use crate::operations::XorOperation; use crate::syscall::precompiles::keccak_sponge::{ KECCAK_GENERAL_OUTPUT_U32S, KECCAK_GENERAL_RATE_U32S, KECCAK_STATE_U32S, }; +use crate::utils::indices_arr; -use p3_keccak_air::KeccakCols; +use p3_keccak_air::{KeccakCols, NUM_KECCAK_COLS, U64_LIMBS}; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_derive::PicusProjection; use zkm_stark::Word; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; /// KeccakSpongeCols is the column layout for the keccak sponge. /// The number of rows equal to the number of block. #[derive(AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub(crate) struct KeccakSpongeCols { pub keccak: KeccakCols, pub block_mem: [MemoryReadCols; KECCAK_GENERAL_RATE_U32S], + #[cfg_attr(feature = "picus", picus(transition_input))] pub shard: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub clk: T, pub is_real: T, pub read_block: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub input_address: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub output_address: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub input_len: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub already_absorbed_u32s: T, pub is_absorbed: T, pub receive_syscall: T, pub write_output: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub is_first_input_block: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub is_final_input_block: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub original_state: [Word; KECCAK_STATE_U32S], pub xored_general_rate: [XorOperation; KECCAK_GENERAL_RATE_U32S], pub input_length_mem: MemoryReadCols, @@ -37,3 +54,50 @@ pub(crate) struct KeccakSpongeCols { } pub const NUM_KECCAK_SPONGE_COLS: usize = size_of::>(); + +#[allow(dead_code)] +/// The full Keccak-f state is 25 lanes, each stored as 4 u16 limbs. +pub const KECCAK_STATE_LIMBS: usize = 25 * U64_LIMBS; +#[allow(dead_code)] +/// The witness stores the final `(0, 0)` lane after iota in a dedicated field. +/// The remaining 24 lanes stay in `a_prime_prime`, so the semantic output +/// boundary needs a tail slice for those lanes. +pub const KECCAK_STATE_TAIL_LIMBS: usize = KECCAK_STATE_LIMBS - U64_LIMBS; + +#[allow(dead_code)] +pub const KECCAK_PICUS_COL_MAP: KeccakCols = make_keccak_picus_col_map(); + +#[allow(dead_code)] +const fn make_keccak_picus_col_map() -> KeccakCols { + let indices_arr = indices_arr::(); + unsafe { transmute::<[usize; NUM_KECCAK_COLS], KeccakCols>(indices_arr) } +} + +/// Semantic Picus projection for the observable input/output contract of the +/// embedded Keccak-f permutation witness. +/// +/// The full witness stores many intermediate round columns that should remain +/// internal to any future Keccak operation submodule. The semantic boundary is: +/// - `state_in`: the full 25-lane permutation input state +/// - `first_step` / `final_step`: the caller-visible round-position flags +/// - `state_out_0_0`: the `(0, 0)` output lane after iota +/// - `state_out_rest`: the remaining 24 output lanes +/// +/// The output is split because the witness layout stores the `(0, 0)` lane in +/// `a_prime_prime_prime_0_0_limbs`, while the other 24 lanes remain in +/// `a_prime_prime`. +#[allow(dead_code)] +#[cfg_attr(feature = "picus", derive(PicusProjection))] +#[cfg_attr(feature = "picus", picus_projection(source = KeccakCols, col_map = KECCAK_PICUS_COL_MAP))] +pub struct KeccakPermutationProjection { + #[cfg_attr(feature = "picus", picus(input, path = a))] + pub state_in: [[[u8; U64_LIMBS]; 5]; 5], + #[cfg_attr(feature = "picus", picus(output, path = step_flags[0]))] + pub first_step: u8, + #[cfg_attr(feature = "picus", picus(output, path = step_flags[23]))] + pub final_step: u8, + #[cfg_attr(feature = "picus", picus(output, path = a_prime_prime_prime_0_0_limbs))] + pub state_out_0_0: [u8; U64_LIMBS], + #[cfg_attr(feature = "picus", picus(output, path = a_prime_prime[0][1]))] + pub state_out_rest: [u8; KECCAK_STATE_TAIL_LIMBS], +} diff --git a/crates/core/machine/src/syscall/precompiles/keccak_sponge/trace.rs b/crates/core/machine/src/syscall/precompiles/keccak_sponge/trace.rs index 5607ab98f..68fb182da 100644 --- a/crates/core/machine/src/syscall/precompiles/keccak_sponge/trace.rs +++ b/crates/core/machine/src/syscall/precompiles/keccak_sponge/trace.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "picus")] +use crate::syscall::precompiles::keccak_sponge::columns::KeccakPermutationProjection; use crate::syscall::precompiles::keccak_sponge::columns::{ KeccakSpongeCols, NUM_KECCAK_SPONGE_COLS, }; @@ -29,6 +31,27 @@ impl MachineAir for KeccakSpongeChip { "KeccakSponge".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> zkm_stark::PicusInfo { + let mut info = KeccakSpongeCols::::picus_info(); + #[cfg(feature = "picus")] + { + let projection = KeccakPermutationProjection::picus_projection_info(); + + // Expose the embedded Keccak permutation input state on the parent + // module interface so Picus cannot vary it existentially when checking + // boundary/transition phases. + info.transition_input_ranges.extend( + projection + .input_ranges + .into_iter() + .map(|(start, end, name)| (start, end, format!("keccak_{name}"))), + ); + } + + info + } + fn generate_dependencies( &self, input: &Self::Record, @@ -131,7 +154,14 @@ impl KeccakSpongeChip { cols.shard = F::from_canonical_u32(event.shard); cols.clk = F::from_canonical_u32(event.clk); cols.is_real = F::ONE; - cols.input_len = F::from_canonical_u32(event.input.len() as u32); + // Keep `input_len` consistent with the dedicated memory read + // (`input_length_record`) used by the AIR constraints. + cols.input_len = F::from_canonical_u32(event.input_len_u32s); + // Keep `input_length_mem` populated on every real row so AIR + // constraints that bind `input_len`/range-check these bytes on + // all real rows are satisfied consistently. + cols.input_length_mem.access.value = Word::from(event.input_length_record.value); + blu.add_u8_range_checks(&event.input_length_record.value.to_le_bytes()); cols.already_absorbed_u32s = F::from_canonical_u32(already_absorbed_u32s); cols.is_absorbed = F::from_bool((round == (NUM_ROUNDS - 1)) && (i != (block_num - 1))); @@ -149,16 +179,18 @@ impl KeccakSpongeChip { // read the input if round == 0 { for j in 0..KECCAK_GENERAL_RATE_U32S { - cols.block_mem[j].populate( - event.input_read_records[i * KECCAK_GENERAL_RATE_U32S + j], - blu, - ); + let record = event.input_read_records[i * KECCAK_GENERAL_RATE_U32S + j]; + cols.block_mem[j].populate(record, blu); + blu.add_u8_range_checks(&record.value.to_le_bytes()); } } // original state for j in 0..KECCAK_STATE_U32S { cols.original_state[j] = Word::from(state_u32s[j]); + // Match AIR-side `slice_range_check_u8` constraints on + // `original_state` bytes for each real row. + blu.add_u8_range_checks(&state_u32s[j].to_le_bytes()); } // xor @@ -180,7 +212,9 @@ impl KeccakSpongeChip { // if this is the last row of the last block, populate writing output if i == (block_num - 1) && round == (NUM_ROUNDS - 1) { for j in 0..KECCAK_GENERAL_OUTPUT_U32S { - cols.output_mem[j].populate(event.output_write_records[j], blu); + let record = event.output_write_records[j]; + cols.output_mem[j].populate(record, blu); + blu.add_u8_range_checks(&record.value.to_le_bytes()); } } diff --git a/crates/core/machine/src/syscall/precompiles/poseidon2/air.rs b/crates/core/machine/src/syscall/precompiles/poseidon2/air.rs index f3f4dd19b..2ee61f56c 100644 --- a/crates/core/machine/src/syscall/precompiles/poseidon2/air.rs +++ b/crates/core/machine/src/syscall/precompiles/poseidon2/air.rs @@ -3,10 +3,11 @@ use std::borrow::Borrow; use p3_air::{Air, AirBuilder, BaseAir, PairBuilder}; use p3_field::FieldAlgebra; use p3_matrix::Matrix; +use zkm_core_executor::ByteOpcode; -use crate::operations::poseidon2::air::{eval_external_round, eval_internal_rounds}; +use crate::operations::poseidon2::air::eval_degree3; use crate::operations::poseidon2::permutation::Poseidon2Cols; -use crate::operations::poseidon2::{NUM_EXTERNAL_ROUNDS, WIDTH}; +use crate::operations::poseidon2::WIDTH; use crate::operations::KoalaBearWordRangeChecker; use crate::syscall::precompiles::poseidon2::{ columns::{Poseidon2MemCols, NUM_COLS}, @@ -44,6 +45,21 @@ where local.is_real.into(), ); let pre_state_word = local.state_mem[i].prev_value().0; + // Enforce that the low three limbs are bytes. + builder.send_byte( + AB::Expr::from_canonical_u8(ByteOpcode::U8Range as u8), + AB::Expr::zero(), + pre_state_word[0], + pre_state_word[1], + local.is_real, + ); + builder.send_byte( + AB::Expr::from_canonical_u8(ByteOpcode::U8Range as u8), + AB::Expr::zero(), + pre_state_word[2], + AB::Expr::zero(), + local.is_real, + ); let pre_state = pre_state_word .iter() .enumerate() @@ -58,10 +74,7 @@ where } // Constrain the permutation. - for r in 0..NUM_EXTERNAL_ROUNDS { - eval_external_round(builder, &local.poseidon2.permutation, r); - } - eval_internal_rounds(builder, &local.poseidon2.permutation); + eval_degree3(builder, &local.poseidon2.permutation); for i in 0..WIDTH { // Range check the current value of the state in memory. This ensures that each part @@ -73,6 +86,21 @@ where local.is_real.into(), ); let post_state_word = local.state_mem[i].value().0; + // Enforce that the low three limbs are bytes. + builder.send_byte( + AB::Expr::from_canonical_u8(ByteOpcode::U8Range as u8), + AB::Expr::zero(), + post_state_word[0], + post_state_word[1], + local.is_real, + ); + builder.send_byte( + AB::Expr::from_canonical_u8(ByteOpcode::U8Range as u8), + AB::Expr::zero(), + post_state_word[2], + AB::Expr::zero(), + local.is_real, + ); let post_state = post_state_word .iter() .enumerate() diff --git a/crates/core/machine/src/syscall/precompiles/poseidon2/columns.rs b/crates/core/machine/src/syscall/precompiles/poseidon2/columns.rs index 08d59b877..593df1868 100644 --- a/crates/core/machine/src/syscall/precompiles/poseidon2/columns.rs +++ b/crates/core/machine/src/syscall/precompiles/poseidon2/columns.rs @@ -1,13 +1,17 @@ use core::mem::size_of; -use zkm_derive::AlignedBorrow; - use crate::memory::MemoryWriteCols; use crate::operations::poseidon2::{Poseidon2Operation, WIDTH}; use crate::operations::KoalaBearWordRangeChecker; +use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; /// Poseidon2MemCols is the column layout for the poseidon2 permutation. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub(crate) struct Poseidon2MemCols { pub poseidon2: Poseidon2Operation, diff --git a/crates/core/machine/src/syscall/precompiles/poseidon2/trace.rs b/crates/core/machine/src/syscall/precompiles/poseidon2/trace.rs index 50eed112d..36f7072cb 100644 --- a/crates/core/machine/src/syscall/precompiles/poseidon2/trace.rs +++ b/crates/core/machine/src/syscall/precompiles/poseidon2/trace.rs @@ -27,6 +27,15 @@ impl MachineAir for Poseidon2PermuteChip { "Poseidon2Permute".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> zkm_stark::PicusInfo { + Poseidon2MemCols::::picus_info() + } + + fn local_only(&self) -> bool { + true + } + fn generate_trace( &self, input: &ExecutionRecord, @@ -125,6 +134,10 @@ impl Poseidon2PermuteChip { cols.state_mem[i].populate(event.state_records[i], blu); cols.pre_state_range_check_cols[i].populate(event.pre_state[i]); cols.post_state_range_check_cols[i].populate(event.post_state[i]); + + // Match AIR-side U8Range lookups for low three limbs of each word. + blu.add_u8_range_checks(&event.pre_state[i].to_le_bytes()[..3]); + blu.add_u8_range_checks(&event.post_state[i].to_le_bytes()[..3]); } } } diff --git a/crates/core/machine/src/syscall/precompiles/sha256/compress/columns.rs b/crates/core/machine/src/syscall/precompiles/sha256/compress/columns.rs index d1ae9d3b6..5d3c56b4f 100644 --- a/crates/core/machine/src/syscall/precompiles/sha256/compress/columns.rs +++ b/crates/core/machine/src/syscall/precompiles/sha256/compress/columns.rs @@ -1,6 +1,8 @@ use std::mem::size_of; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::Word; use crate::{ @@ -10,6 +12,8 @@ use crate::{ XorOperation, }, }; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_SHA_COMPRESS_COLS: usize = size_of::>(); @@ -21,12 +25,17 @@ pub const NUM_SHA_COMPRESS_COLS: usize = size_of::>(); /// compression cycle, one iteration of sha compress is computed. During finalize, the columns are /// combined and written back to memory. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct ShaCompressCols { /// Inputs. + #[cfg_attr(feature = "picus", picus(transition_input))] pub shard: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub clk: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub w_ptr: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub h_ptr: T, pub start: T, @@ -47,13 +56,21 @@ pub struct ShaCompressCols { /// compression, this is w[i] being read only. pub mem_addr: T, + #[cfg_attr(feature = "picus", picus(transition_input))] pub a: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub b: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub c: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub d: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub e: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub f: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub g: Word, + #[cfg_attr(feature = "picus", picus(transition_input))] pub h: Word, /// Current value of K[i]. This is a constant array that loops around every 64 iterations. @@ -101,8 +118,11 @@ pub struct ShaCompressCols { pub finalized_operand: Word, pub finalize_add: AddOperation, + #[cfg_attr(feature = "picus", picus(selector))] pub is_initialize: T, + #[cfg_attr(feature = "picus", picus(selector))] pub is_compression: T, + #[cfg_attr(feature = "picus", picus(selector))] pub is_finalize: T, pub is_last_row: T, diff --git a/crates/core/machine/src/syscall/precompiles/sha256/compress/trace.rs b/crates/core/machine/src/syscall/precompiles/sha256/compress/trace.rs index 4967aebfe..82aa66fcb 100644 --- a/crates/core/machine/src/syscall/precompiles/sha256/compress/trace.rs +++ b/crates/core/machine/src/syscall/precompiles/sha256/compress/trace.rs @@ -10,6 +10,8 @@ use zkm_core_executor::{ syscalls::SyscallCode, ExecutionRecord, Program, }; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{air::MachineAir, Word}; use super::{ @@ -29,6 +31,24 @@ impl MachineAir for ShaCompressChip { "ShaCompress".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + ShaCompressCols::::picus_info() + } + + fn selectors_partition_real_rows(&self) -> bool { + true + } + + fn picus_selector_specialization_allowed(&self, phase: &str, selector_name: &str) -> bool { + match phase { + "first_row" => selector_name == "is_initialize", + "boundary" => selector_name == "is_finalize", + "last_row" => false, + _ => true, + } + } + fn generate_trace( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/precompiles/sha256/extend/columns.rs b/crates/core/machine/src/syscall/precompiles/sha256/extend/columns.rs index 904ba2209..cec917a53 100644 --- a/crates/core/machine/src/syscall/precompiles/sha256/extend/columns.rs +++ b/crates/core/machine/src/syscall/precompiles/sha256/extend/columns.rs @@ -1,6 +1,10 @@ use std::mem::size_of; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; use crate::{ memory::{MemoryReadCols, MemoryWriteCols}, @@ -13,6 +17,7 @@ use crate::{ pub const NUM_SHA_EXTEND_COLS: usize = size_of::>(); #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct ShaExtendCols { /// Inputs. diff --git a/crates/core/machine/src/syscall/precompiles/sha256/extend/trace.rs b/crates/core/machine/src/syscall/precompiles/sha256/extend/trace.rs index 372e20ec0..baebc399c 100644 --- a/crates/core/machine/src/syscall/precompiles/sha256/extend/trace.rs +++ b/crates/core/machine/src/syscall/precompiles/sha256/extend/trace.rs @@ -10,6 +10,8 @@ use zkm_core_executor::{ ExecutionRecord, Program, }; use zkm_stark::air::MachineAir; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use crate::CoreChipError; @@ -26,6 +28,11 @@ impl MachineAir for ShaExtendChip { "ShaExtend".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + ShaExtendCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/precompiles/sys_linux/columns.rs b/crates/core/machine/src/syscall/precompiles/sys_linux/columns.rs index fc9b90840..e38b1fe93 100644 --- a/crates/core/machine/src/syscall/precompiles/sys_linux/columns.rs +++ b/crates/core/machine/src/syscall/precompiles/sys_linux/columns.rs @@ -1,21 +1,22 @@ use std::mem::size_of; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; use zkm_stark::Word; use crate::{ memory::MemoryReadWriteCols, operations::{AddOperation, GtColsBytes, IsZeroOperation}, }; +#[cfg(feature = "picus")] +use zkm_stark::PicusInfo; pub const NUM_SYS_LINUX_COLS: usize = size_of::>(); -/// Linux Syscall AIR columns. -/// -/// All branch selectors are **derived** from `syscall_id` / `a0` / `a1` via `IsZeroOperation`. -/// Intermediate values (`page_offset`, `upper_address`, `is_offset_0`) are computed inline -/// from byte decompositions, not stored. +/// A set of columns needed to compute the Linux Syscall. #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct SysLinuxCols { // ── Common inputs (15 cols) ──────────────────────────────────────── diff --git a/crates/core/machine/src/syscall/precompiles/sys_linux/trace.rs b/crates/core/machine/src/syscall/precompiles/sys_linux/trace.rs index 05f9ec27f..d83bbe906 100644 --- a/crates/core/machine/src/syscall/precompiles/sys_linux/trace.rs +++ b/crates/core/machine/src/syscall/precompiles/sys_linux/trace.rs @@ -29,6 +29,11 @@ impl MachineAir for SysLinuxChip { "SysLinux".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> zkm_stark::PicusInfo { + SysLinuxCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, @@ -100,6 +105,10 @@ impl MachineAir for SysLinuxChip { !shard.get_precompile_events(SyscallCode::SYS_LINUX).is_empty() } } + + fn local_only(&self) -> bool { + true + } } impl SysLinuxChip { diff --git a/crates/core/machine/src/syscall/precompiles/u256x2048_mul/air.rs b/crates/core/machine/src/syscall/precompiles/u256x2048_mul/air.rs index 5ec0c5ef3..57033999d 100644 --- a/crates/core/machine/src/syscall/precompiles/u256x2048_mul/air.rs +++ b/crates/core/machine/src/syscall/precompiles/u256x2048_mul/air.rs @@ -26,6 +26,10 @@ use zkm_curves::{ uint256::U256Field, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ air::{BaseAirBuilder, LookupScope, MachineAir, Polynomial, ZKMAirBuilder}, MachineRecord, @@ -49,6 +53,7 @@ const HI_REGISTER: u32 = Register::A3 as u32; /// A set of columns for the U256x2048Mul operation. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct U256x2048MulCols { /// The shard number of the syscall. @@ -96,6 +101,11 @@ impl MachineAir for U256x2048MulChip { "U256XU2048Mul".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + U256x2048MulCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, @@ -242,6 +252,10 @@ impl MachineAir for U256x2048MulChip { !shard.get_precompile_events(SyscallCode::U256XU2048_MUL).is_empty() } } + + fn local_only(&self) -> bool { + true + } } impl BaseAir for U256x2048MulChip { diff --git a/crates/core/machine/src/syscall/precompiles/uint256/air.rs b/crates/core/machine/src/syscall/precompiles/uint256/air.rs index ac2e4d1a6..8ecce39d0 100644 --- a/crates/core/machine/src/syscall/precompiles/uint256/air.rs +++ b/crates/core/machine/src/syscall/precompiles/uint256/air.rs @@ -33,6 +33,10 @@ use zkm_curves::{ uint256::U256Field, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::{ air::{BaseAirBuilder, LookupScope, MachineAir, Polynomial, ZKMAirBuilder}, MachineRecord, @@ -55,6 +59,7 @@ const WORDS_FIELD_ELEMENT: usize = WordsFieldElement::USIZE; /// A set of columns for the Uint256Mul operation. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct Uint256MulCols { /// The shard number of the syscall. @@ -99,6 +104,11 @@ impl MachineAir for Uint256MulChip { "Uint256MulMod".to_string() } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + Uint256MulCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_add.rs b/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_add.rs index bb67ec62f..4359ba1eb 100644 --- a/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_add.rs +++ b/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_add.rs @@ -26,6 +26,10 @@ use zkm_curves::{ AffinePoint, CurveError, CurveType, EllipticCurve, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{LookupScope, MachineAir, ZKMAirBuilder}; use crate::{ @@ -43,6 +47,7 @@ pub const fn num_weierstrass_add_cols() -> usize /// Right now the number of limbs is assumed to be a constant, although this could be macro-ed or /// made generic in the future. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct WeierstrassAddAssignCols { pub is_real: T, @@ -141,6 +146,11 @@ impl MachineAir } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + WeierstrassAddAssignCols::::picus_info() + } + fn generate_dependencies( &self, input: &Self::Record, diff --git a/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_decompress.rs b/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_decompress.rs index bb794f073..71decf706 100644 --- a/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_decompress.rs +++ b/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_decompress.rs @@ -30,6 +30,10 @@ use zkm_curves::{ CurveError, CurveType, EllipticCurve, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{BaseAirBuilder, LookupScope, MachineAir, Polynomial, ZKMAirBuilder}; use crate::{ @@ -48,6 +52,7 @@ pub const fn num_weierstrass_decompress_cols() -> /// A set of columns to compute `WeierstrassDecompress` that decompresses a point on a Weierstrass /// curve. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct WeierstrassDecompressCols { pub is_real: T, @@ -158,6 +163,11 @@ impl MachineAir } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + WeierstrassDecompressCols::::picus_info() + } + fn generate_trace( &self, input: &ExecutionRecord, diff --git a/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_double.rs b/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_double.rs index d12eba38f..436de0a34 100644 --- a/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_double.rs +++ b/crates/core/machine/src/syscall/precompiles/weierstrass/weierstrass_double.rs @@ -25,6 +25,10 @@ use zkm_curves::{ AffinePoint, CurveType, EllipticCurve, }; use zkm_derive::AlignedBorrow; +#[cfg(feature = "picus")] +use zkm_derive::PicusAnnotations; +#[cfg(feature = "picus")] +use zkm_stark::air::PicusInfo; use zkm_stark::air::{LookupScope, MachineAir, ZKMAirBuilder}; use crate::{ @@ -42,6 +46,7 @@ pub const fn num_weierstrass_double_cols() -> usi /// Right now the number of limbs is assumed to be a constant, although this could be macro-ed or /// made generic in the future. #[derive(Debug, Clone, AlignedBorrow)] +#[cfg_attr(feature = "picus", derive(PicusAnnotations))] #[repr(C)] pub struct WeierstrassDoubleAssignCols { pub is_real: T, @@ -156,6 +161,11 @@ impl MachineAir } } + #[cfg(feature = "picus")] + fn picus_info(&self) -> PicusInfo { + WeierstrassDoubleAssignCols::::picus_info() + } + fn generate_dependencies( &self, input: &Self::Record, diff --git a/crates/derive/src/lib.rs b/crates/derive/src/lib.rs index 077f79997..5e911deba 100644 --- a/crates/derive/src/lib.rs +++ b/crates/derive/src/lib.rs @@ -214,6 +214,21 @@ pub fn machine_air_derive(input: TokenStream) -> TokenStream { } }); + let selectors_partition_real_rows_arms = variants.iter().map(|(variant_name, field)| { + let field_ty = &field.ty; + quote! { + #name::#variant_name(x) => <#field_ty as zkm_stark::air::MachineAir>::selectors_partition_real_rows(x) + } + }); + + let picus_selector_specialization_allowed_arms = + variants.iter().map(|(variant_name, field)| { + let field_ty = &field.ty; + quote! { + #name::#variant_name(x) => <#field_ty as zkm_stark::air::MachineAir>::picus_selector_specialization_allowed(x, phase, selector_name) + } + }); + let machine_air = quote! { impl #impl_generics zkm_stark::air::MachineAir for #name #ty_generics #where_clause { type Record = #execution_record_path; @@ -286,6 +301,22 @@ pub fn machine_air_derive(input: TokenStream) -> TokenStream { #(#picus_info_arms,)* } } + + fn selectors_partition_real_rows(&self) -> bool { + match self { + #(#selectors_partition_real_rows_arms,)* + } + } + + fn picus_selector_specialization_allowed( + &self, + phase: &str, + selector_name: &str, + ) -> bool { + match self { + #(#picus_selector_specialization_allowed_arms,)* + } + } } }; @@ -401,6 +432,11 @@ fn find_eval_trait_bound(attrs: &[syn::Attribute]) -> Option { None } +#[proc_macro_derive(PicusProjection, attributes(picus, picus_projection))] +pub fn picus_projection_derive(input: TokenStream) -> TokenStream { + picus_annotations::picus_projection_derive(input) +} + #[proc_macro_derive(PicusAnnotations, attributes(picus))] pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { picus_annotations::picus_annotations_derive(input) diff --git a/crates/derive/src/picus_annotations.rs b/crates/derive/src/picus_annotations.rs index c5d19669e..ab9bb992c 100644 --- a/crates/derive/src/picus_annotations.rs +++ b/crates/derive/src/picus_annotations.rs @@ -1,5 +1,3 @@ -use std::collections::HashSet; - use proc_macro::TokenStream; use quote::quote; use syn::Generics; @@ -12,17 +10,23 @@ use syn::{ GenericArgument, Path, PathArguments, Result, Token, Type, TypeArray, TypeReference, TypeSlice, }; -#[derive(Default, Debug, Clone)] +#[derive(Default, Clone)] struct PicusArgs { input: bool, output: bool, + transition_input: bool, + transition_output: bool, selector: bool, + path: Option>, } enum Arg { Input, Output, + TransitionInput, + TransitionOutput, Selector, + Path(Box), } impl Parse for Arg { @@ -38,9 +42,19 @@ impl Parse for Arg { if is("output") { return Ok(Arg::Output); } + if is("transition_input") { + return Ok(Arg::TransitionInput); + } + if is("transition_output") { + return Ok(Arg::TransitionOutput); + } if is("selector") { return Ok(Arg::Selector); } + if is("path") { + input.parse::()?; + return Ok(Arg::Path(input.parse()?)); + } Err(syn::Error::new_spanned(key, "unknown key in #[picus(...)]")) } @@ -58,25 +72,94 @@ fn parse_picus_attr(attr: &syn::Attribute) -> syn::Result> { match it { Arg::Input => out.input = true, Arg::Output => out.output = true, + Arg::TransitionInput => out.transition_input = true, + Arg::TransitionOutput => out.transition_output = true, Arg::Selector => out.selector = true, + Arg::Path(expr) => out.path = Some(Box::new(*expr)), } } Ok(Some(out)) } -// ---------- type substitution: replace *type* params with `u8` ---------- -fn type_params_set(gens: &Generics) -> HashSet { - gens.type_params().map(|tp| tp.ident.clone()).collect() +#[derive(Clone)] +struct ProjectionStructArgs { + source: Type, + col_map: syn::Expr, +} + +enum ProjectionStructArg { + Source(Type), + ColMap(syn::Expr), +} + +impl Parse for ProjectionStructArg { + fn parse(input: ParseStream<'_>) -> Result { + let key: Path = input.parse()?; + if key.is_ident("source") { + input.parse::()?; + return Ok(Self::Source(input.parse()?)); + } + if key.is_ident("col_map") { + input.parse::()?; + return Ok(Self::ColMap(input.parse()?)); + } + Err(syn::Error::new_spanned(key, "unknown key in #[picus_projection(...)]")) + } +} + +fn parse_picus_projection_attr( + attrs: &[syn::Attribute], +) -> syn::Result> { + let mut source: Option = None; + let mut col_map: Option = None; + let mut found = false; + for attr in attrs { + if !attr.path.is_ident("picus_projection") { + continue; + } + found = true; + let items = + attr.parse_args_with(Punctuated::::parse_terminated)?; + for it in items { + match it { + ProjectionStructArg::Source(ty) => source = Some(ty), + ProjectionStructArg::ColMap(expr) => col_map = Some(expr), + } + } + } + + if !found { + return Ok(None); + } + let source = source.ok_or_else(|| { + syn::Error::new( + proc_macro2::Span::call_site(), + "missing `source = ...` in #[picus_projection(...)]", + ) + })?; + let col_map = col_map.ok_or_else(|| { + syn::Error::new( + proc_macro2::Span::call_site(), + "missing `col_map = ...` in #[picus_projection(...)]", + ) + })?; + Ok(Some(ProjectionStructArgs { source, col_map })) +} + +// ---------- type substitution: replace the primary column element type with `u8` ---------- +fn first_type_param_ident(gens: &Generics) -> Option { + gens.type_params().next().map(|tp| tp.ident.clone()) } // column values are determined by computing the offset of the ColStruct when instantiated -// with the u8 parameter. This utility substitutes a type parameter with `u8` so we can calculate offsets. -fn ty_sub_u8(mut ty: Type, type_params: &HashSet) -> Type { +// with the u8 parameter. Only the leading column element type should be rewritten; structural +// generics such as curve/field parameter types must be preserved. +fn ty_sub_u8(mut ty: Type, first_type_param: &Option) -> Type { match ty { Type::Path(ref mut tp) => { if tp.qself.is_none() && tp.path.segments.len() == 1 { let seg = &tp.path.segments[0]; - if type_params.contains(&seg.ident) { + if first_type_param.as_ref().is_some_and(|ident| *ident == seg.ident) { return parse_quote!(u8); } } @@ -84,7 +167,7 @@ fn ty_sub_u8(mut ty: Type, type_params: &HashSet) -> Type { if let PathArguments::AngleBracketed(ref mut ab) = seg.arguments { for arg in ab.args.iter_mut() { if let GenericArgument::Type(ref mut inner) = arg { - *inner = ty_sub_u8(inner.clone(), type_params); + *inner = ty_sub_u8(inner.clone(), first_type_param); } } } @@ -92,17 +175,17 @@ fn ty_sub_u8(mut ty: Type, type_params: &HashSet) -> Type { ty } Type::Reference(TypeReference { ref mut elem, .. }) => { - **elem = ty_sub_u8((**elem).clone(), type_params); + **elem = ty_sub_u8((**elem).clone(), first_type_param); ty } Type::Array(TypeArray { ref mut elem, .. }) | Type::Slice(TypeSlice { ref mut elem, .. }) => { - **elem = ty_sub_u8((**elem).clone(), type_params); + **elem = ty_sub_u8((**elem).clone(), first_type_param); ty } Type::Tuple(ref mut tup) => { for el in tup.elems.iter_mut() { - *el = ty_sub_u8(el.clone(), type_params); + *el = ty_sub_u8(el.clone(), first_type_param); } ty } @@ -110,10 +193,19 @@ fn ty_sub_u8(mut ty: Type, type_params: &HashSet) -> Type { } } -// Build Self actual type args; keep lifetimes/consts as-is. +// Build Self actual type args; keep non-leading type params, lifetimes and consts as-is. fn concrete_type_args(gens: &Generics) -> proc_macro2::TokenStream { + let mut seen_type_param = false; let args = gens.params.iter().map(|p| match p { - syn::GenericParam::Type(_) => quote!(u8), + syn::GenericParam::Type(tp) => { + if !seen_type_param { + seen_type_param = true; + quote!(u8) + } else { + let ident = &tp.ident; + quote!(#ident) + } + } syn::GenericParam::Lifetime(lt) => { let lt = <.lifetime; quote!(#lt) @@ -126,20 +218,22 @@ fn concrete_type_args(gens: &Generics) -> proc_macro2::TokenStream { quote!(<#(#args),*>) } -// impl generics = lifetimes + consts only (type params fixed to u8) -fn impl_generics_without_type_params(gens: &Generics) -> proc_macro2::TokenStream { - let lifetimes = gens.lifetimes().map(|d| d.lifetime.clone()); - let consts = gens.const_params().map(|c| { - let id = &c.ident; - let ty = &c.ty; - quote!(const #id: #ty) - }); +// impl generics = all generics except the leading element type parameter. +fn impl_generics_without_primary_type_param(gens: &Generics) -> proc_macro2::TokenStream { let mut parts: Vec = Vec::new(); - for lt in lifetimes { - parts.push(quote!(#lt)); - } - for c in consts { - parts.push(c); + let mut skipped_primary_type = false; + for param in gens.params.iter() { + match param { + syn::GenericParam::Type(tp) => { + if !skipped_primary_type { + skipped_primary_type = true; + } else { + parts.push(quote!(#tp)); + } + } + syn::GenericParam::Lifetime(lt) => parts.push(quote!(#lt)), + syn::GenericParam::Const(c) => parts.push(quote!(#c)), + } } if parts.is_empty() { quote!() @@ -170,10 +264,11 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { } }; - let type_params = type_params_set(&gens); - let impl_gens = impl_generics_without_type_params(&gens); + let first_type_param = first_type_param_ident(&gens); + let impl_gens = impl_generics_without_primary_type_param(&gens); let self_args = concrete_type_args(&gens); let self_conc = quote!(#ident #self_args); + let where_clause = &gens.where_clause; // Per-field code let mut steps = Vec::new(); @@ -188,6 +283,8 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { Ok(Some(a)) => { flags.input |= a.input; flags.output |= a.output; + flags.transition_input |= a.transition_input; + flags.transition_output |= a.transition_output; flags.selector |= a.selector; } Ok(None) => {} @@ -197,7 +294,7 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { } // Field type with all *type* params → u8 - let conc_ty: Type = ty_sub_u8(field.ty.clone(), &type_params); + let conc_ty: Type = ty_sub_u8(field.ty.clone(), &first_type_param); // Add name to id map let push_name = { @@ -222,6 +319,26 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { quote!() }; + let push_transition_in = if flags.transition_input { + quote! { + if width > 0 { + info.transition_input_ranges.push((cur, cur + width, #f_name.to_string())); + } + } + } else { + quote!() + }; + + let push_transition_out = if flags.transition_output { + quote! { + if width > 0 { + info.transition_output_ranges.push((cur, cur + width, #f_name.to_string())); + } + } + } else { + quote!() + }; + let push_sel = if flags.selector { quote! { if width > 0 { @@ -250,6 +367,8 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { #push_name #push_in #push_out + #push_transition_in + #push_transition_out #push_sel #push_is_real cur += width; @@ -258,7 +377,7 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { let expanded = quote! { // Implement on the concrete instantiation where *type* params are `u8` - impl #impl_gens #self_conc { + impl #impl_gens #self_conc #where_clause { pub fn picus_info() -> PicusInfo { let mut info = PicusInfo::default(); let mut cur: usize = 0; // 1 column == 1 byte @@ -269,3 +388,132 @@ pub fn picus_annotations_derive(input: TokenStream) -> TokenStream { }; expanded.into() } + +pub fn picus_projection_derive(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + let ident = input.ident.clone(); + let generics = input.generics.clone(); + let where_clause = &generics.where_clause; + let (impl_generics, ty_generics, _) = generics.split_for_impl(); + + let struct_args = match parse_picus_projection_attr(&input.attrs) { + Ok(Some(args)) => args, + Ok(None) => { + return syn::Error::new_spanned( + &input, + "PicusProjection requires #[picus_projection(source = ..., col_map = ...)]", + ) + .to_compile_error() + .into() + } + Err(e) => return e.to_compile_error().into(), + }; + + let data = match &input.data { + syn::Data::Struct(s) => s, + _ => { + return syn::Error::new_spanned(&input, "PicusProjection only supports structs") + .to_compile_error() + .into() + } + }; + let fields = match &data.fields { + syn::Fields::Named(f) => &f.named, + _ => { + return syn::Error::new_spanned(&input, "PicusProjection requires named fields") + .to_compile_error() + .into() + } + }; + + let source_ty = struct_args.source; + let col_map = struct_args.col_map; + + let mut steps = Vec::new(); + for field in fields.iter() { + let f_ident = field.ident.as_ref().unwrap(); + let f_name = f_ident.to_string(); + + let mut flags = PicusArgs::default(); + for attr in &field.attrs { + if attr.path.is_ident("picus") { + match parse_picus_attr(attr) { + Ok(Some(a)) => { + flags.input |= a.input; + flags.output |= a.output; + flags.transition_input |= a.transition_input; + flags.transition_output |= a.transition_output; + flags.selector |= a.selector; + if let Some(path) = a.path { + flags.path = Some(path); + } + } + Ok(None) => {} + Err(e) => return e.to_compile_error().into(), + } + } + } + + if flags.transition_input || flags.transition_output || flags.selector { + return syn::Error::new_spanned( + field, + "PicusProjection fields currently support only #[picus(input, ...)] and #[picus(output, ...)]", + ) + .to_compile_error() + .into(); + } + if !flags.input && !flags.output { + return syn::Error::new_spanned( + field, + "PicusProjection fields must be marked with #[picus(input, path = ...)] or #[picus(output, path = ...)]", + ) + .to_compile_error() + .into(); + } + let Some(path_expr) = flags.path.clone() else { + return syn::Error::new_spanned( + field, + "PicusProjection fields require `path = ...` to identify the source slice", + ) + .to_compile_error() + .into(); + }; + + let field_ty = &field.ty; + let push_in = if flags.input { + quote! { info.input_ranges.push((start, end, #f_name.to_string())); } + } else { + quote!() + }; + let push_out = if flags.output { + quote! { info.output_ranges.push((start, end, #f_name.to_string())); } + } else { + quote!() + }; + + steps.push(quote! {{ + let start: usize = + zkm_stark::PicusProjectionStart::projection_start(&((#col_map).#path_expr)); + let width: usize = ::core::mem::size_of::<#field_ty>(); + let end = start + width; + info.name_to_colrange.insert(#f_name.to_string(), (start, end)); + for x in start..end { + info.col_to_name.insert(x, format!("{}_{}", #f_name, x)); + } + #push_in + #push_out + }}); + } + + let expanded = quote! { + impl #impl_generics #ident #ty_generics #where_clause { + pub fn picus_projection_info() -> zkm_stark::PicusProjectionInfo { + let mut info = zkm_stark::PicusProjectionInfo::default(); + let _ = ::core::mem::size_of::<#source_ty>(); + #(#steps)* + info + } + } + }; + expanded.into() +} diff --git a/crates/go-runtime/zkvm_runtime/crypto.go b/crates/go-runtime/zkvm_runtime/crypto.go index 117807108..99b66c4a2 100644 --- a/crates/go-runtime/zkvm_runtime/crypto.go +++ b/crates/go-runtime/zkvm_runtime/crypto.go @@ -20,16 +20,18 @@ func Sha256(data []byte) [32]byte { msgLen := len(data) bitLen := uint64(msgLen) * 8 - // Append 0x80 byte - data = append(data, 0x80) - // Pad to 56 mod 64 - for len(data)%64 != 56 { - data = append(data, 0x00) - } - // Append original length in bits as big-endian uint64 - var lenBuf [8]byte - binary.BigEndian.PutUint64(lenBuf[:], bitLen) - data = append(data, lenBuf[:]...) + // Copy into a freshly-sized buffer so we never mutate the caller's slice + // (callers may pass sub-slices that share a backing array with adjacent data). + finalBlockLen := msgLen % 64 + paddedLen := msgLen - finalBlockLen + 64 + if finalBlockLen >= 56 { + paddedLen += 64 + } + padded := make([]byte, paddedLen) + copy(padded, data) + padded[msgLen] = 0x80 + binary.BigEndian.PutUint64(padded[paddedLen-8:], bitLen) + data = padded // Process each 64-byte block for offset := 0; offset < len(data); offset += 64 { diff --git a/crates/go-runtime/zkvm_runtime/deserialize.go b/crates/go-runtime/zkvm_runtime/deserialize.go index 1cf2113c9..c98b3fe54 100644 --- a/crates/go-runtime/zkvm_runtime/deserialize.go +++ b/crates/go-runtime/zkvm_runtime/deserialize.go @@ -78,18 +78,36 @@ func deserializeData(data []byte, v reflect.Value, index int) (int, error) { v.SetUint(a) return index + 8, nil case reflect.Slice: + const maxSliceLen = 1_000_000 b := []byte{data[index], data[index+1], data[index+2], data[index+3], data[index+4], data[index+5], data[index+6], data[index+7]} length := binary.LittleEndian.Uint64(b) index += 8 - switch v.Type().Elem().Kind() { - case reflect.Uint8: + elemKind := v.Type().Elem().Kind() + if elemKind == reflect.Uint8 { + if length > uint64(len(data)-index) { + return index, fmt.Errorf("deserialize failed: []byte length %d exceeds remaining %d", length, len(data)-index) + } bytes := data[index : index+int(length)] v.SetBytes(bytes) return index + int(length), nil } - return index, fmt.Errorf("unsupported type: %v, elem: %v", v.Kind(), v.Elem().Kind()) + + if length > maxSliceLen { + return index, fmt.Errorf("deserialize failed: slice length %d exceeds max %d", length, maxSliceLen) + } + l := int(length) + slice := reflect.MakeSlice(v.Type(), l, l) + for i := 0; i < l; i++ { + var err error + index, err = deserializeData(data, slice.Index(i), index) + if err != nil { + return index, err + } + } + v.Set(slice) + return index, nil case reflect.Array: for i := 0; i < v.Len(); i++ { var err error diff --git a/crates/go-runtime/zkvm_runtime/runtime.go b/crates/go-runtime/zkvm_runtime/runtime.go index 8bedd8e3d..2ddb90778 100644 --- a/crates/go-runtime/zkvm_runtime/runtime.go +++ b/crates/go-runtime/zkvm_runtime/runtime.go @@ -15,10 +15,12 @@ func SyscallWrite(fd int, write_buf []byte, nbytes int) int func SyscallHintLen() int func SyscallHintRead(ptr []byte, len int) func SyscallCommit(index int, word uint32) +func SyscallCommitDeferredProofs(index int, word uint32) func SyscallExit(code int) func SyscallEnterUnconstrained() int func SyscallExitUnconstrained() func SyscallKeccakSponge(input unsafe.Pointer, result unsafe.Pointer) +func SyscallVerifyZKMProof(vkDigest unsafe.Pointer, pvDigest unsafe.Pointer) // secp256k1 precompiles func SyscallSecp256k1Add(p unsafe.Pointer, q unsafe.Pointer) @@ -81,7 +83,7 @@ func HintSlice(data []byte) { func ReadHintVec() []byte { // Read length prefix (4 bytes LE) lenLen := SyscallHintLen() - lenBuf := make([]byte, ((lenLen + 3) / 4) * 4) + lenBuf := make([]byte, ((lenLen+3)/4)*4) SyscallHintRead(lenBuf, lenLen) dataLen := int(lenBuf[0]) | int(lenBuf[1])<<8 | int(lenBuf[2])<<16 | int(lenBuf[3])<<24 @@ -94,6 +96,7 @@ func ReadHintVec() []byte { } var PublicValuesHasher hash.Hash = sha256.New() +var DeferredProofsDigest [8]uint32 const EMBEDDED_RESERVED_INPUT_REGION_SIZE int = 1024 * 1024 * 1024 const MAX_MEMORY int = 0x7f000000 @@ -127,6 +130,17 @@ func Commit[T any](value T) { SyscallWrite(13, bytes, length) } +func VerifyZKMProof(vkDigest *[8]uint32, pvDigest *[32]byte) { + SyscallVerifyZKMProof(unsafe.Pointer(vkDigest), unsafe.Pointer(pvDigest)) +} + +func CommitDeferredProofsDigest(digest [8]uint32) { + DeferredProofsDigest = digest + for i, word := range digest { + SyscallCommitDeferredProofs(i, word) + } +} + //go:linkname RuntimeExit zkvm.RuntimeExit func RuntimeExit(code int) { hashBytes := PublicValuesHasher.Sum(nil) @@ -136,6 +150,7 @@ func RuntimeExit(code int) { word := binary.LittleEndian.Uint32(hashBytes[i*4 : (i+1)*4]) SyscallCommit(i, word) } + CommitDeferredProofsDigest(DeferredProofsDigest) SyscallExit(code) } diff --git a/crates/go-runtime/zkvm_runtime/serialize.go b/crates/go-runtime/zkvm_runtime/serialize.go index 49c815170..25eff2b61 100644 --- a/crates/go-runtime/zkvm_runtime/serialize.go +++ b/crates/go-runtime/zkvm_runtime/serialize.go @@ -56,21 +56,16 @@ func serializeData(v reflect.Value) ([]byte, error) { binary.LittleEndian.PutUint64(b, uint64(v.Uint())) return b, nil case reflect.Slice: - switch v.Type().Elem().Kind() { - case reflect.Uint8: - output := make([]byte, 8) - binary.LittleEndian.PutUint64(output, uint64(v.Len())) - - for i := 0; i < v.Len(); i++ { - d, err := serializeData(v.Index(i)) - if err != nil { - return nil, err - } - output = append(output, d...) + output := make([]byte, 8) + binary.LittleEndian.PutUint64(output, uint64(v.Len())) + for i := 0; i < v.Len(); i++ { + d, err := serializeData(v.Index(i)) + if err != nil { + return nil, err } - return output, nil + output = append(output, d...) } - return nil, fmt.Errorf("unsupported type: %v, elem: %v", v.Kind(), v.Elem().Kind()) + return output, nil case reflect.Array: switch v.Type().Elem().Kind() { case reflect.Uint8: diff --git a/crates/go-runtime/zkvm_runtime/syscall_mipsle.s b/crates/go-runtime/zkvm_runtime/syscall_mipsle.s index 8721cbf2c..28e8d47ce 100644 --- a/crates/go-runtime/zkvm_runtime/syscall_mipsle.s +++ b/crates/go-runtime/zkvm_runtime/syscall_mipsle.s @@ -38,6 +38,13 @@ TEXT ·SyscallCommit(SB), $0-8 SYSCALL RET +TEXT ·SyscallCommitDeferredProofs(SB), $0-8 + MOVW index+0(FP), R4 + MOVW word+4(FP), R5 + MOVW $0x1A, R2 + SYSCALL + RET + TEXT ·SyscallExit(SB), $0-4 MOVW code+0(FP), R4 // a0 = code MOVW $0, R2 // v0 = syscall 0 @@ -62,6 +69,13 @@ TEXT ·SyscallKeccakSponge(SB), $0-8 SYSCALL RET +TEXT ·SyscallVerifyZKMProof(SB), $0-8 + MOVW $0x1B, R2 + MOVW vkDigest+0(FP), R4 + MOVW pvDigest+4(FP), R5 + SYSCALL + RET + // secp256k1 elliptic curve precompiles TEXT ·SyscallSecp256k1Add(SB), $0-8 diff --git a/crates/picus/Cargo.toml b/crates/picus/Cargo.toml index 2cbd865ee..ba217857f 100644 --- a/crates/picus/Cargo.toml +++ b/crates/picus/Cargo.toml @@ -8,7 +8,7 @@ keywords = { workspace = true } categories = { workspace = true } [dependencies] -zkm-core-machine = { workspace = true } +zkm-core-machine = { workspace = true, features = ["picus"] } zkm-core-executor = {workspace = true} zkm-stark = {workspace = true} p3-air = { workspace = true } diff --git a/crates/picus/README.md b/crates/picus/README.md index 1ada8ee9f..43cf57f37 100644 --- a/crates/picus/README.md +++ b/crates/picus/README.md @@ -20,6 +20,198 @@ Relevant entry points: - Picus AST / serialization: `crates/picus/src/pcl/` - Instruction opcode routing spec: `crates/picus/src/opcode_spec.rs` +## Picus Annotations + +Picus annotations are metadata on AIR column structs. They do not change the AIR +semantics or the trace; they only tell the extractor which columns should become +Picus module inputs, outputs, transition state, or selectors. + +In the common case, Picus can infer a useful module interface directly from +lookup interactions such as instruction, memory, syscall, and global messages. +That is often enough for chips whose observable behavior is already exposed by +those interactions. + +However, interaction-based inference is not always sufficient. In particular, +chips that enforce important semantics across rows can carry state that never +appears directly in an interaction payload. For those chips, Picus annotations +tell the extractor which columns are part of the semantic interface and which +columns should be treated as carried transition state. + +The metadata is collected by deriving `PicusAnnotations` on the column struct and +then returning `Cols::::picus_info()` from the chip's `picus_info()` method. + +Supported annotations: + +- `#[picus(input)]` + - Marks the current-row field as a Picus module input. + - Use this when the extracted module should treat the field as externally supplied + state for the current row. + +- `#[picus(output)]` + - Marks the current-row field as a Picus module output. + - Use this when the field is a current-row result you want visible at the Picus + interface. + +- `#[picus(transition_input)]` + - Marks the current-row field as incoming carried state for transition-capable phases. + - In practice this means the field is exposed as an input for phases that reason + about cross-row behavior (`FirstRow`, `Transition`, `Boundary`, `LastRow`). + +- `#[picus(transition_output)]` + - Marks the immediate successor row's version of the field as an output in phases + that expose successor state (`FirstRow` and `Transition`). + - Use this when the next-row value is semantically important and you want Picus to + check determinism of that successor state. + +- `#[picus(selector)]` + - Marks a field as a selector column. + - The extractor uses selector columns to build selector-specialized modules and + the optional `top` module that constrains selector shape. + +Rule of thumb: + +- Start by relying on interaction-driven inference; do not annotate columns just + because they exist. +- Use `input` / `output` for values that are meaningful on the current row by + themselves. +- Use `transition_input` / `transition_output` for state that is intentionally + carried from one row to the next. +- If a field is both incoming carried state and outgoing successor state, mark it + with both `transition_input` and `transition_output`. +- Do not mark a field as transition state just because it appears in a transition + constraint. Mark it only when it is part of the semantic row-to-row interface of + the chip. +- Be conservative with `transition_output`: exporting next-row values makes Picus + check determinism of that successor state. If the next-row value is padding-only, + phase-local, or intentionally existential, do not expose it. + +Concrete example: + +- In `BooleanCircuitGarble`, `delta` and `checks` are carried across gate rows, so + they are marked with `#[picus(transition_input, transition_output)]`. +- That tells Picus to treat the current row's `delta` / `checks` as inputs and the + immediate next row's `delta` / `checks` as outputs in the phases where successor + state is part of the interface. + +## Picus Projections + +`PicusAnnotations` describe concrete trace storage fields on a chip row. +`PicusProjection` is different: it describes a smaller semantic interface +projected out of a larger witness layout. + +This is useful when a submodule has many internal witness columns, but Picus +should expose only the semantically relevant boundary to the caller. Typical +examples are operation summaries such as Poseidon2 or embedded sub-AIRs such as +Keccak. + +Use a projection when: + +- the full witness layout is too large or too internal to expose directly, +- the caller should see only a specific input/output contract, +- the remaining witness columns should stay existential inside the summarized + submodule. + +Do not use a projection as a replacement for `PicusAnnotations` on the chip's +main trace columns. Projections are for operation/submodule boundaries, not for +describing ordinary chip row I/O. + +### Derive shape + +Projection metadata is declared on a separate struct with: + +- `#[derive(PicusProjection)]` +- a struct-level `#[picus_projection(source = ..., col_map = ...)]` +- field-level `#[picus(input, path = ...)]` / `#[picus(output, path = ...)]` + +Example: + +```rust +use zkm_derive::PicusProjection; + +#[derive(PicusProjection)] +#[picus_projection( + source = Poseidon2Degree3Cols, + col_map = POSEIDON2_DEGREE3_COL_MAP +)] +pub struct Poseidon2Degree3Projection { + #[picus(input, path = state.external_rounds_state[0])] + pub state_in: [u8; WIDTH], + + #[picus(output, path = state.output_state)] + pub state_out: [u8; WIDTH], +} +``` + +This does not change the witness layout. It only says: + +- the semantic input starts at `state.external_rounds_state[0]` and spans + `WIDTH` columns, +- the semantic output starts at `state.output_state` and spans `WIDTH` columns. + +### How `path` works + +`path = ...` points at the source slice inside the `col_map`. + +Important detail: + +- `path` selects the start of the semantic slice, +- the projected field type determines the width, +- the derive recursively takes the first concrete source column from the path. + +So this: + +```rust +#[picus(input, path = state.external_rounds_state[0])] +pub state_in: [u8; WIDTH], +``` + +means: + +- start at the first column of `state.external_rounds_state[0]`, +- take `WIDTH` consecutive columns. + +You do not need to write `state.external_rounds_state[0][0]` unless you really +want to refer to a scalar source column. + +### Column map requirement + +The `col_map` argument should be a `usize`-instantiated version of the source +layout whose fields hold concrete column indices. + +For example: + +```rust +pub const POSEIDON2_DEGREE3_COL_MAP: Poseidon2Degree3Cols = make_col_map_degree3(); +``` + +This lets the derive resolve `path = ...` into concrete source ranges without +changing the actual runtime trace type. + +### Intended usage in Picus + +Today projections are consumed by summary hooks in `OperationSummaryAirBuilder`, +not by ordinary chip extraction. + +The supported pattern is: + +1. Keep the full exact AIR/witness layout unchanged. +2. Define a projection for the caller-visible semantic boundary. +3. Use a Picus summary hook to emit an auxiliary module whose interface is the + projected inputs/outputs. +4. Keep the rest of the source witness internal to that auxiliary module. + +That is how Picus can summarize a large operation while preserving exact +internal constraints. + +### Rule of thumb + +- Use `PicusAnnotations` for concrete chip row metadata. +- Use `PicusProjection` for summarized operation or sub-AIR boundaries. +- Keep projections semantic and minimal: expose only what the caller should + reason about. +- Leave intermediate round state, helper witnesses, and existential internals + out of the projection unless they are part of the intended contract. + ## OpcodeSpec `OpcodeSpec` defines how instruction lookups are routed during extraction. @@ -106,8 +298,10 @@ This example assumes you added a new machine chip named `MyChip` in `zkm-core-ma - Ensure the chip has a stable `name()`; this is what `--chip` matches. 2. Ensure the chip has Picus metadata. -- Derive/provide `PicusInfo` metadata for column naming and selectors. -- Mark selector columns (for case-splitting). +- Derive/provide `PicusInfo` metadata for column naming, row I/O, transition state, and selectors. +- Mark selector columns for case-splitting. +- Mark current-row inputs/outputs only when they are intended to be part of the extracted module interface. +- Mark transition inputs/outputs only for fields that are semantically carried across rows. Example (from `AddSubCols`): @@ -131,7 +325,8 @@ pub struct AddSubCols { ``` This annotation pattern is what allows extraction to discover selector columns and produce -selector-specialized Picus modules. +selector-specialized Picus modules. If the chip also carries row-to-row state, annotate those +fields with `transition_input` / `transition_output` as needed. Also implement `picus_info` on the chip. For example, here is what needs to be added for the `AddSub` chip. diff --git a/crates/picus/src/lib.rs b/crates/picus/src/lib.rs index df976d8f9..63b5af8c8 100644 --- a/crates/picus/src/lib.rs +++ b/crates/picus/src/lib.rs @@ -1,5 +1,6 @@ pub mod opcode_spec; pub mod pcl; pub mod picus_builder; +pub mod syscall_spec; use pcl::*; diff --git a/crates/picus/src/main.rs b/crates/picus/src/main.rs index e0a97918a..fa132e4c0 100644 --- a/crates/picus/src/main.rs +++ b/crates/picus/src/main.rs @@ -1,14 +1,20 @@ -use std::{collections::BTreeMap, path::PathBuf}; +use std::{ + collections::{BTreeMap, BTreeSet}, + path::PathBuf, +}; use clap::{Parser, ValueEnum, ValueHint}; use p3_air::{Air, BaseAir}; +use p3_matrix::{dense::RowMajorMatrix, Matrix}; use zkm_core_machine::MipsAir; use zkm_picus::{ pcl::{ - initialize_fresh_var_ctr, set_field_modulus, set_picus_names, Felt, PicusAtom, - PicusConstraint, PicusExpr, PicusModule, PicusProgram, + begin_expr_reify_scope, initialize_fresh_var_ctr, partial_evaluate_expr, set_field_modulus, + set_picus_names, Felt, PicusAtom, PicusConstraint, PicusExpr, PicusModule, PicusProgram, + }, + picus_builder::{ + ColumnOutputMode, ExtractionPhase, PicusBuilder, ShrCarrySummaryMode, SubmoduleMode, }, - picus_builder::{PicusBuilder, ShrCarrySummaryMode, SubmoduleMode}, }; use zkm_stark::{Chip, MachineAir, PicusInfo}; @@ -38,6 +44,14 @@ struct Args { /// How to summarize ByteOpcode::ShrCarry during extraction. #[arg(long = "shrcarry-summary", value_enum, default_value_t = ShrCarrySummaryModeArg::Abstract)] pub shrcarry_summary: ShrCarrySummaryModeArg, + + /// How aggressively to expose chip columns as Picus module outputs. + #[arg( + long = "column-output-mode", + value_enum, + default_value_t = ColumnOutputModeArg::InteractionsOnly + )] + pub column_output_mode: ColumnOutputModeArg, } #[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)] @@ -46,6 +60,12 @@ enum ShrCarrySummaryModeArg { Precise, } +#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)] +enum ColumnOutputModeArg { + InteractionsOnly, + AllNonInputsAreOutputs, +} + impl From for ShrCarrySummaryMode { fn from(value: ShrCarrySummaryModeArg) -> Self { match value { @@ -55,15 +75,37 @@ impl From for ShrCarrySummaryMode { } } -/// Analyze a single chip and process all its deferred sub-chip tasks. -/// This replaces direct recursion in `MessageBuilder::send()`. +impl From for ColumnOutputMode { + fn from(value: ColumnOutputModeArg) -> Self { + match value { + ColumnOutputModeArg::InteractionsOnly => ColumnOutputMode::InteractionsOnly, + ColumnOutputModeArg::AllNonInputsAreOutputs => ColumnOutputMode::AllNonInputsAreOutputs, + } + } +} + +/// Extracts one selector-specialized module for one chip/phase pair. +/// +/// The workflow is: +/// 1. Build a `PicusBuilder` for the requested phase and specialization env. +/// 2. Run the chip AIR once on the primary window. +/// 3. Optionally rerun the AIR on a shifted window so the exposed successor row +/// is itself locally feasible. +/// 4. Recursively inline any deferred sub-chip tasks produced by interactions. +/// 5. Expose inputs/outputs according to `PicusInfo` and the chosen output mode. +/// +/// This replaces direct recursion in `MessageBuilder::send()` so extraction can +/// keep one explicit worklist of deferred sub-chip analyses. +#[allow(clippy::too_many_arguments)] fn analyze_chip<'chips, A>( chip: &'chips Chip, chips: &'chips [Chip], picus_builder: Option<&mut PicusBuilder<'chips, A>>, specialization_env: Option>, + phase: ExtractionPhase, submodule_mode: SubmoduleMode, shr_carry_summary_mode: ShrCarrySummaryMode, + column_output_mode: ColumnOutputMode, ) -> (PicusModule, BTreeMap) where A: MachineAir + BaseAir + Air>, @@ -82,15 +124,21 @@ where } else { &mut PicusBuilder::new( chip, - PicusModule::new(chip.name()), + PicusModule::new(format!("{}__{}", chip.name(), phase.module_suffix())), chips, None, specialization_env, Some(submodule_mode), Some(shr_carry_summary_mode), + Some(phase), + None, ) }; - chip.air.eval(builder); + { + let _expr_reify_scope = begin_expr_reify_scope(None); + chip.air.eval(builder); + } + run_shifted_air_eval(chip, builder); // Process deferred tasks recursively while let Some(task) = builder.concrete_pending_tasks.pop() { @@ -118,12 +166,14 @@ where let mut sub_builder = PicusBuilder::new( target_chip, - PicusModule::new(task.chip_name.clone()), + PicusModule::new(format!("{}__{}", task.chip_name, builder.phase.module_suffix())), builder.chips, Some(task.main_vars.clone()), Some(env.clone()), Some(SubmoduleMode::Inline), Some(shr_carry_summary_mode), + Some(builder.phase), + Some(task.capture_interface), ); let (mut sub_module, aux_modules) = analyze_chip( @@ -131,8 +181,10 @@ where builder.chips, Some(&mut sub_builder), None, + builder.phase, SubmoduleMode::Inline, shr_carry_summary_mode, + column_output_mode, ); // Merge submodules builder.aux_modules.extend(aux_modules.into_iter()); @@ -142,11 +194,666 @@ where let updated_picus_module = sub_module.partial_eval(&env); builder.picus_module.constraints.extend_from_slice(&updated_picus_module.constraints); builder.picus_module.calls.extend_from_slice(&updated_picus_module.calls); + if task.capture_interface { + let propagated_global_outputs = sub_builder + .global_send_outputs + .iter() + .map(|expr| partial_evaluate_expr(expr, &env)) + .collect::>(); + builder.picus_module.outputs.extend_from_slice(&propagated_global_outputs); + builder.global_send_outputs.extend(propagated_global_outputs); + } builder.picus_module.postconditions.extend_from_slice(&sub_module.postconditions); } + let picus_info = chip.picus_info(); + if builder.capture_interface { + builder.expose_transition_inputs(&picus_info.transition_input_ranges); + builder.expose_annotated_primary_outputs(&picus_info.output_ranges); + builder.expose_transition_outputs(&picus_info.transition_output_ranges); + if column_output_mode == ColumnOutputMode::AllNonInputsAreOutputs { + builder.expose_primary_row_non_inputs_as_outputs(&picus_info.input_ranges); + builder.expose_full_next_row_as_outputs(); + } + } (builder.picus_module.clone(), builder.aux_modules.clone()) } +/// Re-evaluates the AIR on a shifted row window without exposing new interface ports. +/// +/// For phases such as `FirstRow` and `Transition`, the first AIR eval only +/// proves that `row1` is a valid successor of `row0`. This shifted pass reuses +/// the same builder with `(row1, row2)` as `(local, next)` so the exported +/// successor row is also checked as a valid local row. +fn run_shifted_air_eval<'chips, A>( + chip: &'chips Chip, + builder: &mut PicusBuilder<'chips, A>, +) where + A: MachineAir + BaseAir + Air>, +{ + let Some(shifted_phase) = builder.phase.shifted_eval_phase(builder.local_only) else { + return; + }; + + let width = builder.main.width(); + assert!( + builder.main.height() >= 3, + "shifted eval for phase {:?} requires a 3-row main trace", + builder.phase + ); + + // Reuse the existing builder by temporarily viewing rows `(1, 2)` as the + // active `(local, next)` pair. This adds the shifted constraints to the + // same module while keeping row2 existential. + let shifted_rows = builder + .main + .row_slice(1) + .iter() + .chain(builder.main.row_slice(2).iter()) + .copied() + .collect::>(); + let original_main = + std::mem::replace(&mut builder.main, RowMajorMatrix::new(shifted_rows, width)); + let original_phase = std::mem::replace(&mut builder.phase, shifted_phase); + let original_capture_interface = std::mem::replace(&mut builder.capture_interface, false); + let _expr_reify_scope = begin_expr_reify_scope(None); + chip.air.eval(builder); + builder.capture_interface = original_capture_interface; + builder.phase = original_phase; + builder.main = original_main; +} + +fn collect_expr_vars(expr: &PicusExpr, vars: &mut BTreeSet) { + match expr { + PicusExpr::Const(_) => {} + PicusExpr::Var(id) => { + vars.insert(*id); + } + PicusExpr::Add(left, right) + | PicusExpr::Sub(left, right) + | PicusExpr::Mul(left, right) + | PicusExpr::Div(left, right) => { + collect_expr_vars(left, vars); + collect_expr_vars(right, vars); + } + PicusExpr::Neg(expr) | PicusExpr::Pow(_, expr) => collect_expr_vars(expr, vars), + } +} + +fn collect_interface_vars(module: &PicusModule) -> BTreeSet { + let mut vars = BTreeSet::new(); + for expr in &module.inputs { + collect_expr_vars(expr, &mut vars); + } + for expr in &module.outputs { + collect_expr_vars(expr, &mut vars); + } + for call in &module.calls { + for expr in &call.inputs { + collect_expr_vars(expr, &mut vars); + } + for expr in &call.outputs { + collect_expr_vars(expr, &mut vars); + } + } + vars +} + +fn as_var(expr: &PicusExpr) -> Option { + match expr { + PicusExpr::Var(id) => Some(*id), + _ => None, + } +} + +fn is_var_minus_one(expr: &PicusExpr, var_id: usize) -> bool { + matches!( + expr, + PicusExpr::Sub(left, right) + if matches!(&**left, PicusExpr::Var(id) if *id == var_id) + && matches!(&**right, PicusExpr::Const(1)) + ) +} + +fn match_bit_expr(expr: &PicusExpr) -> Option { + match expr { + PicusExpr::Mul(left, right) => { + if let Some(var_id) = as_var(left) { + if is_var_minus_one(right, var_id) { + return Some(var_id); + } + } + if let Some(var_id) = as_var(right) { + if is_var_minus_one(left, var_id) { + return Some(var_id); + } + } + None + } + _ => None, + } +} + +fn is_boolean_guard_expr(expr: &PicusExpr, known_guards: &BTreeSet) -> bool { + match expr { + PicusExpr::Const(0 | 1) => true, + PicusExpr::Var(id) => known_guards.contains(id), + PicusExpr::Mul(left, right) => { + is_boolean_guard_expr(left, known_guards) && is_boolean_guard_expr(right, known_guards) + } + PicusExpr::Sub(left, right) if matches!(&**left, PicusExpr::Const(1)) => { + is_boolean_guard_expr(right, known_guards) + } + // `-(g - 1)` is syntactically common after simplification and is equivalent to `1 - g`. + PicusExpr::Neg(inner) => { + matches!(&**inner, PicusExpr::Sub(left, right) if matches!(&**right, PicusExpr::Const(1)) && is_boolean_guard_expr(left, known_guards)) + } + _ => false, + } +} + +fn guard_factor_count(expr: &PicusExpr, known_guards: &BTreeSet) -> Option { + match expr { + PicusExpr::Const(1) => Some(0), + PicusExpr::Const(0) => None, + PicusExpr::Var(id) if known_guards.contains(id) => Some(1), + PicusExpr::Mul(left, right) => { + Some(guard_factor_count(left, known_guards)? + guard_factor_count(right, known_guards)?) + } + // A complemented boolean still acts as a guard factor. + PicusExpr::Sub(left, right) if matches!(&**left, PicusExpr::Const(1)) => match &**right { + PicusExpr::Const(0) => Some(0), + PicusExpr::Const(1) => None, + _ if is_boolean_guard_expr(right, known_guards) => { + Some(guard_factor_count(right, known_guards).unwrap_or(0).max(1)) + } + _ => None, + }, + PicusExpr::Neg(inner) => match &**inner { + PicusExpr::Sub(left, right) if matches!(&**right, PicusExpr::Const(1)) => match &**left + { + PicusExpr::Const(0) => Some(0), + PicusExpr::Const(1) => None, + _ if is_boolean_guard_expr(left, known_guards) => { + Some(guard_factor_count(left, known_guards).unwrap_or(0).max(1)) + } + _ => None, + }, + _ => None, + }, + _ => None, + } +} + +fn match_guard_definition_var(expr: &PicusExpr, known_guards: &BTreeSet) -> Option { + match expr { + PicusExpr::Sub(left, right) => { + if let Some(id) = as_var(left) { + if is_boolean_guard_expr(right, known_guards) { + return Some(id); + } + } + if let Some(id) = as_var(right) { + if is_boolean_guard_expr(left, known_guards) { + return Some(id); + } + } + None + } + _ => None, + } +} + +fn collect_known_guard_vars(constraints: &[PicusConstraint]) -> BTreeSet { + let mut known_guards = BTreeSet::new(); + let mut changed = true; + while changed { + changed = false; + for constraint in constraints { + let inferred_guard = match constraint { + PicusConstraint::Eq(expr) => { + match_bit_expr(expr).or_else(|| match_guard_definition_var(expr, &known_guards)) + } + _ => None, + }; + if let Some(var_id) = inferred_guard { + changed |= known_guards.insert(var_id); + } + } + } + known_guards +} + +fn collect_expr_guarded_uses( + expr: &PicusExpr, + active_guard: bool, + known_guards: &BTreeSet, + guarded_vars: &mut BTreeSet, + unguarded_vars: &mut BTreeSet, +) { + match expr { + PicusExpr::Const(_) => {} + PicusExpr::Var(id) => { + if active_guard { + guarded_vars.insert(*id); + } else { + unguarded_vars.insert(*id); + } + } + PicusExpr::Add(left, right) | PicusExpr::Sub(left, right) | PicusExpr::Div(left, right) => { + collect_expr_guarded_uses( + left, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + collect_expr_guarded_uses( + right, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + PicusExpr::Neg(expr) | PicusExpr::Pow(_, expr) => { + collect_expr_guarded_uses( + expr, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + PicusExpr::Mul(left, right) => { + let left_guard_factor_count = guard_factor_count(left, known_guards); + let right_guard_factor_count = guard_factor_count(right, known_guards); + match (left_guard_factor_count, right_guard_factor_count) { + (Some(left_count), None) => { + collect_expr_guarded_uses( + left, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + collect_expr_guarded_uses( + right, + active_guard || left_count > 0, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + (None, Some(right_count)) => { + collect_expr_guarded_uses( + left, + active_guard || right_count > 0, + known_guards, + guarded_vars, + unguarded_vars, + ); + collect_expr_guarded_uses( + right, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + _ => { + collect_expr_guarded_uses( + left, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + collect_expr_guarded_uses( + right, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + } + } + } +} + +fn collect_constraint_guarded_uses( + constraint: &PicusConstraint, + active_guard: bool, + known_guards: &BTreeSet, + guarded_vars: &mut BTreeSet, + unguarded_vars: &mut BTreeSet, +) { + match constraint { + PicusConstraint::Lt(left, right) + | PicusConstraint::Leq(left, right) + | PicusConstraint::Gt(left, right) + | PicusConstraint::Geq(left, right) => { + collect_expr_guarded_uses( + left, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + collect_expr_guarded_uses( + right, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + PicusConstraint::Implies(left, right) + | PicusConstraint::Iff(left, right) + | PicusConstraint::And(left, right) + | PicusConstraint::Or(left, right) => { + collect_constraint_guarded_uses( + left, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + collect_constraint_guarded_uses( + right, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + PicusConstraint::Not(inner) => { + collect_constraint_guarded_uses( + inner, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + PicusConstraint::Eq(expr) | PicusConstraint::Det(expr) => { + collect_expr_guarded_uses( + expr, + active_guard, + known_guards, + guarded_vars, + unguarded_vars, + ); + } + } +} + +fn collect_dominated_vars( + module: &PicusModule, + protected_vars: &BTreeSet, + known_guards: &BTreeSet, +) -> BTreeSet { + let mut guarded_vars = BTreeSet::new(); + let mut unguarded_vars = BTreeSet::new(); + + for constraint in &module.constraints { + collect_constraint_guarded_uses( + constraint, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + for constraint in &module.postconditions { + collect_constraint_guarded_uses( + constraint, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + for expr in &module.assume_deterministic { + collect_expr_guarded_uses( + expr, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + for expr in &module.inputs { + collect_expr_guarded_uses( + expr, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + for expr in &module.outputs { + collect_expr_guarded_uses( + expr, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + for call in &module.calls { + for expr in &call.inputs { + collect_expr_guarded_uses( + expr, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + for expr in &call.outputs { + collect_expr_guarded_uses( + expr, + false, + known_guards, + &mut guarded_vars, + &mut unguarded_vars, + ); + } + } + + guarded_vars + .difference(&unguarded_vars) + .copied() + .filter(|var_id| !protected_vars.contains(var_id)) + .collect() +} + +fn match_guarded_bit_expr(expr: &PicusExpr, known_guards: &BTreeSet) -> Option { + match expr { + PicusExpr::Mul(left, right) => { + if guard_factor_count(left, known_guards).is_some_and(|count| count > 0) { + if let Some(var_id) = match_bit_expr(right) { + return Some(var_id); + } + } + if guard_factor_count(right, known_guards).is_some_and(|count| count > 0) { + if let Some(var_id) = match_bit_expr(left) { + return Some(var_id); + } + } + None + } + _ => None, + } +} + +fn match_guarded_var_product(expr: &PicusExpr, known_guards: &BTreeSet) -> Option { + match expr { + PicusExpr::Mul(left, right) => { + if guard_factor_count(left, known_guards).is_some_and(|count| count > 0) { + return as_var(right); + } + if guard_factor_count(right, known_guards).is_some_and(|count| count > 0) { + return as_var(left); + } + None + } + _ => None, + } +} + +fn split_guarded_var_product( + expr: &PicusExpr, + known_guards: &BTreeSet, +) -> Option<(PicusExpr, usize)> { + match expr { + PicusExpr::Mul(left, right) => { + if guard_factor_count(left, known_guards).is_some_and(|count| count > 0) { + if let Some(var_id) = as_var(right) { + return Some(((*left.clone()), var_id)); + } + } + if guard_factor_count(right, known_guards).is_some_and(|count| count > 0) { + if let Some(var_id) = as_var(left) { + return Some(((*right.clone()), var_id)); + } + } + None + } + _ => None, + } +} + +fn split_guarded_const_product( + expr: &PicusExpr, + known_guards: &BTreeSet, +) -> Option<(PicusExpr, u64)> { + match expr { + PicusExpr::Mul(left, right) => { + if guard_factor_count(left, known_guards).is_some_and(|count| count > 0) { + if let PicusExpr::Const(constant) = &**right { + return Some(((*left.clone()), *constant)); + } + } + if guard_factor_count(right, known_guards).is_some_and(|count| count > 0) { + if let PicusExpr::Const(constant) = &**left { + return Some(((*right.clone()), *constant)); + } + } + None + } + _ => None, + } +} + +fn normalize_constraint( + constraint: PicusConstraint, + known_guards: &BTreeSet, + dominated_vars: &BTreeSet, +) -> PicusConstraint { + match constraint { + PicusConstraint::Eq(expr) => { + let expr = *expr; + if let Some(var_id) = match_guarded_bit_expr(&expr, known_guards) { + if dominated_vars.contains(&var_id) { + return PicusConstraint::new_bit(PicusExpr::Var(var_id)); + } + } + PicusConstraint::Eq(Box::new(expr)) + } + PicusConstraint::Leq(left, right) => { + let left = *left; + let right = *right; + if let (Some((left_guard, var_id)), Some((right_guard, constant))) = ( + split_guarded_var_product(&left, known_guards), + split_guarded_const_product(&right, known_guards), + ) { + if left_guard == right_guard && dominated_vars.contains(&var_id) { + return PicusConstraint::new_leq( + PicusExpr::Var(var_id), + PicusExpr::Const(constant), + ); + } + } + if matches!(right, PicusExpr::Const(_)) { + if let Some(var_id) = match_guarded_var_product(&left, known_guards) { + if dominated_vars.contains(&var_id) { + return PicusConstraint::new_leq(PicusExpr::Var(var_id), right); + } + } + } + PicusConstraint::Leq(Box::new(left), Box::new(right)) + } + PicusConstraint::Lt(left, right) => { + let left = *left; + let right = *right; + if let (Some((left_guard, var_id)), Some((right_guard, constant))) = ( + split_guarded_var_product(&left, known_guards), + split_guarded_const_product(&right, known_guards), + ) { + if left_guard == right_guard && dominated_vars.contains(&var_id) { + return PicusConstraint::new_lt( + PicusExpr::Var(var_id), + PicusExpr::Const(constant), + ); + } + } + if matches!(right, PicusExpr::Const(_)) { + if let Some(var_id) = match_guarded_var_product(&left, known_guards) { + if dominated_vars.contains(&var_id) { + return PicusConstraint::new_lt(PicusExpr::Var(var_id), right); + } + } + } + PicusConstraint::Lt(Box::new(left), Box::new(right)) + } + other => other, + } +} + +fn postprocess_module(chip_name: &str, module: &mut PicusModule) { + if matches!(chip_name, "BooleanCircuitGarble" | "SysLinux" | "MemoryInstrs") { + let protected_vars = collect_interface_vars(module); + loop { + // The guarded-width simplifier is only sound when the target variable is + // boolean-guarded everywhere it matters. We therefore infer concrete + // guard vars from actual constraints, including aliases and simple + // complements, track which hidden vars are used only under those + // guards, and only then strip the guards off range/bit constraints. + // Re-running the analysis lets newly-exposed bit constraints create + // more guards in later rounds. + let known_guards = collect_known_guard_vars(&module.constraints); + let dominated_vars = collect_dominated_vars(module, &protected_vars, &known_guards); + + let mut changed = false; + let mut constraints = Vec::with_capacity(module.constraints.len()); + let mut seen = BTreeSet::new(); + for constraint in std::mem::take(&mut module.constraints) { + let original_rendered = constraint.to_string(); + let normalized = normalize_constraint(constraint, &known_guards, &dominated_vars); + let rendered = normalized.to_string(); + changed |= rendered != original_rendered; + if !seen.insert(rendered) { + changed = true; + continue; + } + constraints.push(normalized); + } + module.constraints = constraints; + if !changed { + break; + } + } + } +} + +fn postprocess_modules(chip_name: &str, modules: &mut BTreeMap) { + for module in modules.values_mut() { + postprocess_module(chip_name, module); + } +} + fn format_env(env: &BTreeMap) -> String { if env.is_empty() { return "{}".to_string(); @@ -159,11 +866,17 @@ fn format_env(env: &BTreeMap) -> String { fn build_selector_env( picus_info: &PicusInfo, selected_selector_col: Option, + phase: Option, ) -> BTreeMap { let mut env = BTreeMap::new(); - // Specialize to real rows for chips that carry an `is_real` column. + // Most extraction phases model an active computation row, so specializing `is_real = 1` + // removes vacuous gating. `LastRow` is the exception: some chips, such as `ShaCompress`, + // require the final trace row to be padding, so the AIR must be free to decide whether the + // last row is real. if let Some(id) = picus_info.is_real_index { - env.insert(id, 1); + if !matches!(phase, Some(ExtractionPhase::LastRow)) { + env.insert(id, 1); + } } // One-hot selector assignment for this extraction pass. if let Some(selected_col) = selected_selector_col { @@ -177,9 +890,127 @@ fn build_selector_env( env } +fn selector_specialization_allowed( + chip: &Chip, + phase: ExtractionPhase, + selector_name: &str, +) -> bool +where + A: MachineAir, +{ + chip.picus_selector_specialization_allowed(phase.module_suffix(), selector_name) +} + +fn add_selector_shape_postconditions( + top_module: &mut PicusModule, + picus_info: &PicusInfo, + assume_selectors_deterministic: bool, + selectors_partition_real_rows: bool, + real_row_only: bool, +) -> bool { + if picus_info.selector_indices.is_empty() { + return false; + } + + // These are emitted as postconditions rather than ordinary assertions so the selector + // contract remains visible even for shape-only top modules that do not carry an AIR body. + let mut one_hot_sum = PicusExpr::Const(0); + for (selector_col, _) in &picus_info.selector_indices { + let selector_var = PicusExpr::Var(*selector_col); + one_hot_sum += selector_var.clone(); + top_module.outputs.push(selector_var.clone()); + top_module.postconditions.push(PicusConstraint::new_bit(selector_var.clone())); + if assume_selectors_deterministic { + top_module.assume_deterministic.push(selector_var); + } + } + if selectors_partition_real_rows { + if real_row_only { + top_module.postconditions.push(PicusConstraint::new_equality(one_hot_sum, 1.into())); + } else { + let Some(is_real_index) = picus_info.is_real_index else { + panic!("selector partition requires is_real to be annotated"); + }; + let is_real_var = PicusExpr::Var(is_real_index); + top_module.outputs.push(is_real_var.clone()); + top_module.postconditions.push(PicusConstraint::new_bit(is_real_var.clone())); + top_module.postconditions.push(PicusConstraint::new_equality(one_hot_sum, is_real_var)); + } + } else { + top_module.postconditions.push(PicusConstraint::new_lt(one_hot_sum, 2.into())); + } + true +} + +/// Builds the program's `top` module and any auxiliary modules needed by it. +/// +/// `top` is not meant to model one concrete trace row. Its job is narrower: prove that the +/// selector columns satisfy the polynomial relationships that make them real selectors. +/// +/// To avoid phase-specific obligations like `when_first_row()` / `when_last_row()` and to keep +/// interactions out of the contract, we analyze the chip once under a dedicated neutral phase +/// with interaction lowering disabled. The resulting hidden witness still has to satisfy the AIR's +/// unconditional selector equations, but `top` is no longer over-constrained by row-position or +/// routing semantics. +/// +/// For chips whose selectors partition real rows, `top` is intentionally a real-row-only proof: +/// we specialize `is_real = 1` and ask Picus to prove that the selectors form a partition of one. +/// Padding-row behavior stays out of `top`. +fn build_top_module<'chips, A>( + chip: &'chips Chip, + chips: &'chips [Chip], + picus_info: &PicusInfo, + fresh_var_ctr_base: usize, + shr_carry_summary_mode: ShrCarrySummaryMode, + column_output_mode: ColumnOutputMode, + assume_selectors_deterministic: bool, +) -> Option<(PicusModule, BTreeMap)> +where + A: MachineAir + BaseAir + Air>, +{ + if picus_info.selector_indices.is_empty() { + return None; + } + + let mut top_env = BTreeMap::new(); + let real_row_only = chip.selectors_partition_real_rows() && picus_info.is_real_index.is_some(); + if real_row_only { + top_env.insert(picus_info.is_real_index.unwrap(), 1); + } + + initialize_fresh_var_ctr(fresh_var_ctr_base); + let (top_base_module, top_aux_modules) = analyze_chip( + chip, + chips, + None, + Some(top_env.clone()), + ExtractionPhase::Top, + SubmoduleMode::Ignore, + shr_carry_summary_mode, + column_output_mode, + ); + let mut top_module = top_base_module.partial_eval(&top_env); + top_module.inputs.clear(); + top_module.outputs.clear(); + // `top` should prove selector structure from constraints alone; any interaction-derived + // determinism assumptions belong to the specialized row modules instead. + top_module.assume_deterministic.clear(); + + top_module.name = "top".to_string(); + add_selector_shape_postconditions( + &mut top_module, + picus_info, + assume_selectors_deterministic, + chip.selectors_partition_real_rows(), + real_row_only, + ); + Some((top_module, top_aux_modules)) +} + fn main() { let args = Args::parse(); let shr_carry_summary_mode: ShrCarrySummaryMode = args.shrcarry_summary.into(); + let column_output_mode: ColumnOutputMode = args.column_output_mode.into(); if args.chip.is_none() { panic!("Chip name must be provided!"); @@ -210,9 +1041,18 @@ fn main() { // Build selector-specialized modules directly by running extraction once per // selector assignment. This lets opcodes fold to constants before send-dispatch. + // + // Conceptually, extraction proceeds one phase at a time: + // - `SingleRow` handles degenerate one-row traces. + // - `FirstRow` and `Transition` materialize an extra witness row and use a + // shifted AIR pass to prove the exported successor row is locally valid. + // - `Boundary` stops at the last real row before padding and does not expose + // padding-row outputs. + // - `LastRow` models the final trace row and only imports carried state. println!("Generating Picus program for {} chip.....", chip.name()); let mut selector_modules = BTreeMap::new(); let mut all_aux_modules = BTreeMap::new(); + let phases = ExtractionPhase::all(chip.local_only(), chip.local_only_row_sensitive()); if picus_info.selector_indices.is_empty() && picus_info.is_real_index.is_none() { panic!("PicusBuilder needs at least one selector to be enabled!") @@ -222,70 +1062,87 @@ fn main() { println!("selector indices: {:?}", picus_info.selector_indices); if picus_info.selector_indices.is_empty() { // No selector columns: still run one extraction pass (is_real specialized if present). - let env = build_selector_env(&picus_info, None); - initialize_fresh_var_ctr(fresh_var_ctr_base); - let (base_module, mut aux_modules) = analyze_chip( - chip, - &chips, - None, - Some(env.clone()), - SubmoduleMode::Inline, - shr_carry_summary_mode, - ); - all_aux_modules.append(&mut aux_modules); - let updated_module = base_module.partial_eval(&env); - selector_modules.insert(updated_module.name.clone(), updated_module); - } else { - for (selector_col, _) in &picus_info.selector_indices { - let env = build_selector_env(&picus_info, Some(*selector_col)); + for phase in &phases { + let env = build_selector_env(&picus_info, None, Some(*phase)); initialize_fresh_var_ctr(fresh_var_ctr_base); let (base_module, mut aux_modules) = analyze_chip( chip, &chips, None, Some(env.clone()), + *phase, SubmoduleMode::Inline, shr_carry_summary_mode, + column_output_mode, ); all_aux_modules.append(&mut aux_modules); let updated_module = base_module.partial_eval(&env); selector_modules.insert(updated_module.name.clone(), updated_module); } + } else { + for phase in &phases { + let allowed_selectors = picus_info + .selector_indices + .iter() + .filter(|(_, selector_name)| { + selector_specialization_allowed(chip, *phase, selector_name) + }) + .collect::>(); + + if allowed_selectors.is_empty() { + let env = build_selector_env(&picus_info, None, Some(*phase)); + initialize_fresh_var_ctr(fresh_var_ctr_base); + let (base_module, mut aux_modules) = analyze_chip( + chip, + &chips, + None, + Some(env.clone()), + *phase, + SubmoduleMode::Inline, + shr_carry_summary_mode, + column_output_mode, + ); + all_aux_modules.append(&mut aux_modules); + let updated_module = base_module.partial_eval(&env); + selector_modules.insert(updated_module.name.clone(), updated_module); + continue; + } + for (selector_col, _) in allowed_selectors { + let env = build_selector_env(&picus_info, Some(*selector_col), Some(*phase)); + initialize_fresh_var_ctr(fresh_var_ctr_base); + let (base_module, mut aux_modules) = analyze_chip( + chip, + &chips, + None, + Some(env.clone()), + *phase, + SubmoduleMode::Inline, + shr_carry_summary_mode, + column_output_mode, + ); + all_aux_modules.append(&mut aux_modules); + let updated_module = base_module.partial_eval(&env); + selector_modules.insert(updated_module.name.clone(), updated_module); + } + } } + postprocess_modules(&chip.name(), &mut all_aux_modules); + postprocess_modules(&chip.name(), &mut selector_modules); picus_program.add_modules(&mut all_aux_modules); picus_program.add_modules(&mut selector_modules); - // Build the top module from chip constraints but ignore instruction submodules. - // This keeps top focused on selector determinism while still retaining chip-local constraints. - let top_env = build_selector_env(&picus_info, None); - initialize_fresh_var_ctr(fresh_var_ctr_base); - let (top_base_module, mut top_aux_modules) = analyze_chip( + if let Some((top_module, mut top_aux_modules)) = build_top_module( chip, &chips, - None, - Some(top_env.clone()), - SubmoduleMode::Ignore, + &picus_info, + fresh_var_ctr_base, shr_carry_summary_mode, - ); - picus_program.add_modules(&mut top_aux_modules); - let mut top_module = top_base_module.partial_eval(&top_env); - top_module.name = "top".to_string(); - // Top exists only to prove selector properties, so expose only selectors as outputs. - top_module.outputs.clear(); - if !picus_info.selector_indices.is_empty() { - let mut one_hot_sum = PicusExpr::Const(0); - for (selector_col, _) in &picus_info.selector_indices { - let selector_var = PicusExpr::Var(*selector_col); - one_hot_sum += selector_var.clone(); - top_module.outputs.push(selector_var.clone()); - top_module.postconditions.push(PicusConstraint::new_bit(selector_var.clone())); - if args.assume_selectors_deterministic { - top_module.assume_deterministic.push(selector_var); - } - } - top_module.postconditions.push(PicusConstraint::new_lt(one_hot_sum, 2.into())) + column_output_mode, + args.assume_selectors_deterministic, + ) { + picus_program.add_modules(&mut top_aux_modules); + picus_program.add_module("top", top_module); } - picus_program.add_module("top", top_module); let res = picus_program.write_to_path(args.picus_out_dir.join(format!("{}.picus", chip.name()))); if res.is_err() { diff --git a/crates/picus/src/opcode_spec.rs b/crates/picus/src/opcode_spec.rs index 27552aa51..1038404f4 100644 --- a/crates/picus/src/opcode_spec.rs +++ b/crates/picus/src/opcode_spec.rs @@ -55,7 +55,7 @@ pub fn spec_for(kind: Opcode) -> OpcodeSpec { arg_to_colname: &[ (Single(2), "pc"), (Single(3), "next_pc"), - (Range { start: 7, end: 11 }, "a"), + (Range { start: 7, end: 11 }, "bit_shift_result"), (Range { start: 11, end: 15 }, "b"), (Range { start: 15, end: 19 }, "c"), ], @@ -78,7 +78,18 @@ pub fn spec_for(kind: Opcode) -> OpcodeSpec { arg_to_colname: &[ (Single(2), "pc"), (Single(3), "next_pc"), - (Range { start: 7, end: 11 }, "a"), + (Range { start: 7, end: 11 }, "bit_shift_result"), + (Range { start: 11, end: 15 }, "b"), + (Range { start: 15, end: 19 }, "c"), + ], + }, + Opcode::SRA => OpcodeSpec { + selector: "is_sra", + chip: "ShiftRight", + arg_to_colname: &[ + (Single(2), "pc"), + (Single(3), "next_pc"), + (Range { start: 7, end: 11 }, "bit_shift_result"), (Range { start: 11, end: 15 }, "b"), (Range { start: 15, end: 19 }, "c"), ], diff --git a/crates/picus/src/pcl/expr.rs b/crates/picus/src/pcl/expr.rs index 7ddd01494..3efe72283 100644 --- a/crates/picus/src/pcl/expr.rs +++ b/crates/picus/src/pcl/expr.rs @@ -1,4 +1,5 @@ use std::{ + cell::RefCell, collections::HashMap, fmt::{self, Display, Formatter}, iter::{Product, Sum}, @@ -14,10 +15,41 @@ static PICUS_NAMES_GLOBAL: OnceLock>> = OnceLock:: /// Maintains col indices for fresh variables during the course of extraction static FRESH_VAR_CTR: OnceLock = OnceLock::new(); + +/// Default AST-size threshold beyond which Picus reifies an expression into a +/// fresh variable during extraction. +/// +/// This is a pragmatic memory-control heuristic for very large AIRs such as +/// Poseidon2. The arithmetic overloads in this file build tree-shaped syntax and +/// clone aggressively, so letting expressions grow without bound can exhaust +/// memory before they ever reach a constraint sink. +const DEFAULT_EXPR_REIFY_THRESHOLD: usize = 128; + +#[derive(Default)] +struct ExprReifyContext { + threshold: usize, + pending_bindings: Vec<(usize, PicusExpr)>, +} + +thread_local! { + /// Optional extraction-local queue of oversized expressions that have been + /// replaced by fresh variables during `PicusExpr` construction. + /// + /// `PicusExpr` itself does not own a builder or module, so arithmetic + /// overloads cannot emit constraints directly. Instead they enqueue + /// `fresh_k = expr` bindings here, and the active `PicusBuilder` flushes + /// them before it appends the next real constraint. + static EXPR_REIFY_CONTEXT: RefCell> = const { RefCell::new(None) }; +} + pub fn set_picus_names(map: HashMap) { let _ = PICUS_NAMES_GLOBAL.set(RwLock::new(map)); } +pub(crate) fn picus_name_for(id: usize) -> Option { + PICUS_NAMES_GLOBAL.get()?.read().unwrap().get(&id).cloned() +} + // Get or initialize the fresh var counter fn ctr() -> &'static AtomicUsize { FRESH_VAR_CTR.get_or_init(|| AtomicUsize::new(0)) @@ -43,6 +75,56 @@ pub fn fresh_picus_expr() -> PicusExpr { PicusExpr::Var(fresh_picus_var_id()) } +/// Scope guard enabling threshold-based expression reification for the current +/// thread. +/// +/// While this guard is live, `PicusExpr` arithmetic will replace oversized +/// subexpressions with fresh variables and queue the defining equalities for the +/// active builder to flush later. +pub struct ExprReifyScope; + +impl Drop for ExprReifyScope { + fn drop(&mut self) { + EXPR_REIFY_CONTEXT.with(|ctx| { + let mut ctx = ctx.borrow_mut(); + let prev = ctx.take(); + if let Some(prev) = prev.as_ref() { + assert!( + prev.pending_bindings.is_empty(), + "dropping expression reification scope with unflushed pending bindings" + ); + } + assert!(prev.is_some(), "ExprReifyScope dropped without an active reification context"); + }); + } +} + +/// Enables threshold-based expression reification for the current thread. +pub fn begin_expr_reify_scope(threshold: Option) -> ExprReifyScope { + EXPR_REIFY_CONTEXT.with(|ctx| { + let mut ctx = ctx.borrow_mut(); + assert!(ctx.is_none(), "expression reification context is already active on this thread"); + *ctx = Some(ExprReifyContext { + threshold: threshold.unwrap_or(DEFAULT_EXPR_REIFY_THRESHOLD), + pending_bindings: Vec::new(), + }); + }); + ExprReifyScope +} + +/// Drains any pending `fresh = expr` bindings produced by oversized +/// subexpression reification during the active extraction pass. +pub fn drain_pending_expr_bindings() -> Vec<(usize, PicusExpr)> { + EXPR_REIFY_CONTEXT.with(|ctx| { + let mut ctx = ctx.borrow_mut(); + if let Some(ctx) = ctx.as_mut() { + std::mem::take(&mut ctx.pending_bindings) + } else { + Vec::new() + } + }) +} + use p3_field::{FieldAlgebra, PrimeField32}; /// Global, thread-safe holder for the PCL prime field modulus. @@ -118,12 +200,10 @@ impl Display for PicusAtom { match self { Self::Const(c) => write!(f, "{c}"), Self::Var(id) => { - if let Some(lock) = PICUS_NAMES_GLOBAL.get() { - if let Some(name) = lock.read().unwrap().get(id) { - return f.write_str(name); - } + if let Some(name) = picus_name_for(*id) { + return f.write_str(&name); } - write!(f, "v{id}") + write!(f, "fresh_{id}") } } } @@ -277,6 +357,22 @@ impl PicusExpr { pub fn is_const_zero(&self) -> bool { matches!(self, PicusExpr::Const(c) if *c == 0) } + + fn maybe_reify(self) -> Self { + EXPR_REIFY_CONTEXT.with(|ctx| { + let mut ctx = ctx.borrow_mut(); + let Some(ctx) = ctx.as_mut() else { + return self; + }; + if self.size() <= ctx.threshold { + return self; + } + + let fresh = fresh_picus_var_id(); + ctx.pending_bindings.push((fresh, self)); + PicusExpr::Var(fresh) + }) + } } macro_rules! impl_from_ints { @@ -308,17 +404,17 @@ impl Add for PicusExpr { if c == 0 { rhs } else { - PicusExpr::Add(Box::new(lhs), Box::new(rhs)) + PicusExpr::Add(Box::new(lhs), Box::new(rhs)).maybe_reify() } } (_, PicusExpr::Const(c)) => { if c == 0 { lhs } else { - PicusExpr::Add(Box::new(lhs), Box::new(rhs)) + PicusExpr::Add(Box::new(lhs), Box::new(rhs)).maybe_reify() } } - _ => PicusExpr::Add(Box::new(lhs), Box::new(rhs)), + _ => PicusExpr::Add(Box::new(lhs), Box::new(rhs)).maybe_reify(), } } } @@ -364,10 +460,10 @@ impl Sub for PicusExpr { if c == 0 { lhs } else { - PicusExpr::Sub(Box::new(self), Box::new(rhs)) + PicusExpr::Sub(Box::new(self), Box::new(rhs)).maybe_reify() } } - _ => PicusExpr::Sub(Box::new(self), Box::new(rhs)), + _ => PicusExpr::Sub(Box::new(self), Box::new(rhs)).maybe_reify(), } } } @@ -406,7 +502,7 @@ impl Neg for PicusExpr { let lhs = self.clone(); match lhs.clone() { PicusExpr::Const(c) => reduce_mod((current_modulus().unwrap() - c) as i64).into(), - _ => PicusExpr::Neg(Box::new(lhs)), + _ => PicusExpr::Neg(Box::new(lhs)).maybe_reify(), } } } @@ -422,7 +518,7 @@ impl Mul for PicusExpr { match (lhs.clone(), rhs.clone()) { (PicusExpr::Const(c), _) => rhs * c, (_, PicusExpr::Const(c)) => lhs * c, - _ => PicusExpr::Mul(Box::new(lhs), Box::new(rhs)), + _ => PicusExpr::Mul(Box::new(lhs), Box::new(rhs)).maybe_reify(), } } } @@ -469,7 +565,7 @@ impl Mul for PicusExpr { let lhs = self.clone(); match lhs { PicusExpr::Const(c_1) => reduce_mod((c_1 * rhs) as i64).into(), - _ => PicusExpr::Mul(Box::new(lhs), Box::new(rhs.into())), + _ => PicusExpr::Mul(Box::new(lhs), Box::new(rhs.into())).maybe_reify(), } } } @@ -549,6 +645,8 @@ pub enum PicusConstraint { Or(Box, Box), /// Canonical equality-to-zero form: `Eq(e)` represents `e = 0`. Eq(Box), + /// Assert expression is deterministic + Det(Box), } impl PicusConstraint { @@ -652,6 +750,10 @@ impl PicusConstraint { let new_e = multiplier.clone() * (*e.clone()); PicusConstraint::Eq(Box::new(new_e)) } + Det(e) => { + let new_expr = multiplier.clone() * (*e.clone()); + PicusConstraint::Det(Box::new(new_expr)) + } } } } diff --git a/crates/picus/src/pcl/partial_evaluator.rs b/crates/picus/src/pcl/partial_evaluator.rs index 672d84b28..32c56f81d 100644 --- a/crates/picus/src/pcl/partial_evaluator.rs +++ b/crates/picus/src/pcl/partial_evaluator.rs @@ -118,6 +118,14 @@ pub fn subst_constraint( let keep = |cc: PicusConstraint| Some(cc); match c { + Det(e) => { + let ee = subst_expr(e, env); + match ee { + PicusExpr::Const(_) => None, + _ => keep(Det(Box::new(ee))), + } + } + Eq(e) => { let ee = subst_expr(e, env); // Drop tautologies Eq(0); keep contradictions as Eq(1) diff --git a/crates/picus/src/pcl/program.rs b/crates/picus/src/pcl/program.rs index 1e42bfeb5..c721696fe 100644 --- a/crates/picus/src/pcl/program.rs +++ b/crates/picus/src/pcl/program.rs @@ -203,7 +203,7 @@ impl Display for PicusModule { /// Examples: /// /// - `Const(5)` → `5` -/// - `Var("x",1,0)` → `x_1_0` +/// - `Var(7)` → `fresh_7` (or the registered column name when available) /// - `Add(a,b)` → `(+ a b)` /// - `Neg(e)` → `(- e)` /// - `Pow(2, e)` → `(pow 2 e)` @@ -212,7 +212,13 @@ impl Display for PicusExpr { use PicusExpr::{Add, Const, Div, Mul, Neg, Pow, Sub, Var}; match self { Const(v) => write!(f, "{v}"), - Var(id) => write!(f, "x_{id}"), + Var(id) => { + if let Some(name) = super::picus_name_for(*id) { + f.write_str(&name) + } else { + write!(f, "fresh_{id}") + } + } Add(a, b) => write!(f, "(+ {a} {b})"), Sub(a, b) => write!(f, "(- {a} {b})"), Mul(a, b) => write!(f, "(* {a} {b})"), @@ -231,13 +237,14 @@ impl Display for PicusExpr { /// constraints/expressions. impl Display for PicusConstraint { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - use PicusConstraint::{And, Eq, Geq, Gt, Iff, Implies, Leq, Lt, Not, Or}; + use PicusConstraint::{And, Det, Eq, Geq, Gt, Iff, Implies, Leq, Lt, Not, Or}; match self { Lt(e1, e2) => write!(f, "(< {e1} {e2})"), Leq(e1, e2) => write!(f, "(<= {e1} {e2})"), Gt(e1, e2) => write!(f, "(> {e1} {e2})"), Geq(e1, e2) => write!(f, "(>= {e1} {e2})"), Eq(e) => write!(f, "(= {e} 0)"), + Det(e) => write!(f, "(det {e})"), Implies(c1, c2) => write!(f, "(=> {c1} {c2})"), Iff(c1, c2) => write!(f, "(<=> {c1} {c2})"), Not(c) => write!(f, "(! {c})"), diff --git a/crates/picus/src/picus_builder.rs b/crates/picus/src/picus_builder.rs index f8009c69a..ed5ee3ac9 100644 --- a/crates/picus/src/picus_builder.rs +++ b/crates/picus/src/picus_builder.rs @@ -3,15 +3,21 @@ use std::collections::BTreeMap; use crate::{ opcode_spec::{spec_for, IndexSlice}, pcl::{ - fresh_picus_expr, fresh_picus_var, fresh_picus_var_id, partial_evaluate_expr, Felt, - PicusAtom, PicusCall, PicusConstraint, PicusExpr, PicusModule, + drain_pending_expr_bindings, fresh_picus_expr, fresh_picus_var, fresh_picus_var_id, + partial_evaluate_expr, Felt, PicusAtom, PicusCall, PicusConstraint, PicusExpr, PicusModule, }, + syscall_spec::spec_for_sender, }; use p3_air::{AirBuilder, AirBuilderWithPublicValues, PairBuilder}; use p3_matrix::dense::{DenseMatrix, RowMajorMatrix}; +use p3_matrix::Matrix; use zkm_core_executor::{ByteOpcode, Opcode}; -use zkm_stark::{AirLookup, Chip, LookupKind, MachineAir, MessageBuilder, ZKM_PROOF_NUM_PV_ELTS}; +use zkm_stark::{ + AirLookup, Chip, LookupKind, MachineAir, MessageBuilder, OperationSummaryAirBuilder, Word, + ZKM_PROOF_NUM_PV_ELTS, +}; +/// Controls how instruction and syscall lookups are represented during extraction. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SubmoduleMode { /// Ignore instruction submodules entirely. @@ -24,6 +30,7 @@ pub enum SubmoduleMode { Submodule, } +/// Controls how `ByteOpcode::ShrCarry` is summarized in the extracted module. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum ShrCarrySummaryMode { /// Keep ShrCarry abstract as a module call. @@ -32,18 +39,157 @@ pub enum ShrCarrySummaryMode { Precise, } -/// Implementation `AirBuilder` which builds Picus programs +/// Controls which chip columns become explicit Picus module outputs. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ColumnOutputMode { + /// Only expose ports inferred from interactions, summaries, and explicit Picus annotations. + InteractionsOnly, + /// Expose every primary-row column as an output unless it is annotated as an input. + AllNonInputsAreOutputs, +} + +/// Selects which trace shape the extracted module is meant to model. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ExtractionPhase { + /// Selector-proof extraction with all row-position predicates disabled. + Top, + /// Trace of length 1. + SingleRow, + /// First row of a multi-row trace. + FirstRow, + /// Interior row of a multi-row trace. + Transition, + /// Last real row before padding begins. + Boundary, + /// Final trace row. + LastRow, +} + +impl ExtractionPhase { + /// Returns every extraction phase relevant for the chip's row model. + /// + /// `local_only` chips that also ignore absolute row position only need + /// `SingleRow`. Row-sensitive local-only chips still need the first/interior/last + /// split, but never the true cross-row `Boundary` phase. Non-local chips do + /// not get `SingleRow` by default: for them, that phase would mean "first + /// and last simultaneously", which is only meaningful for chips that truly + /// admit a one-real-row trace. + pub fn all(local_only: bool, local_only_row_sensitive: bool) -> Vec { + if local_only && !local_only_row_sensitive { + return vec![Self::SingleRow]; + } + + let mut phases = vec![Self::FirstRow, Self::Transition, Self::LastRow]; + if local_only { + phases.insert(0, Self::SingleRow); + } + if !local_only { + phases.insert(3, Self::Boundary); + } + phases + } + + /// Returns the stable suffix used in generated module names for this phase. + pub fn module_suffix(self) -> &'static str { + match self { + Self::Top => "top", + Self::SingleRow => "single_row", + Self::FirstRow => "first_row", + Self::Transition => "transition", + Self::Boundary => "boundary", + Self::LastRow => "last_row", + } + } + + fn is_first_row(self) -> bool { + matches!(self, Self::SingleRow | Self::FirstRow) + } + + fn is_last_row(self) -> bool { + matches!(self, Self::SingleRow | Self::LastRow) + } + + fn is_transition(self, local_only: bool) -> bool { + !local_only && matches!(self, Self::FirstRow | Self::Transition | Self::Boundary) + } + + /// Returns the number of concrete trace rows materialized for this phase. + /// + /// `FirstRow` and `Transition` use 3 rows so we can re-run the AIR on + /// `(row1, row2)` and prove that the exposed successor row is itself locally + /// feasible. `Boundary` stops at 2 rows because its successor is padding and + /// is not exposed. + pub fn row_count(self, local_only: bool) -> usize { + if local_only { + 1 + } else if matches!(self, Self::Top) { + 2 + } else if self.requires_shifted_eval(local_only) { + 3 + } else { + 2 + } + } + + /// Returns whether this phase needs a shifted second AIR evaluation. + /// + /// The shifted pass reinterprets the exposed successor row as the local row + /// of a fresh transition window to prove that successor is itself feasible. + pub fn requires_shifted_eval(self, local_only: bool) -> bool { + !local_only && matches!(self, Self::FirstRow | Self::Transition) + } + + /// Returns the phase semantics used by the shifted AIR pass, if any. + /// + /// The shifted pass treats the exposed successor row as an ordinary interior + /// row. This intentionally drops `FirstRow` semantics on the second eval. + pub fn shifted_eval_phase(self, local_only: bool) -> Option { + self.requires_shifted_eval(local_only).then_some(Self::Transition) + } + + fn exposes_transition_inputs(self, local_only: bool) -> bool { + !local_only + && matches!(self, Self::FirstRow | Self::Transition | Self::Boundary | Self::LastRow) + } + + /// Returns whether this phase should export the immediate successor row as outputs. + /// + /// `Boundary` deliberately returns `false`: its successor is padding and + /// should remain existential. + pub fn exposes_next_row_outputs(self, local_only: bool) -> bool { + !local_only && matches!(self, Self::FirstRow | Self::Transition) + } + + /// The synthesized second row is partially specialized only when the phase + /// already determines whether that row is real. + fn next_is_real(self) -> Option { + match self { + Self::Top => None, + Self::FirstRow | Self::Transition => Some(1), + Self::Boundary => Some(0), + Self::SingleRow | Self::LastRow => None, + } + } +} + +/// `AirBuilder` implementation that lowers one chip/phase view into a Picus module. #[derive(Clone)] pub struct PicusBuilder<'chips, A: MachineAir> { pub preprocessed: RowMajorMatrix, pub main: RowMajorMatrix, pub public_values: Vec, pub picus_module: PicusModule, + pub global_send_outputs: Vec, pub aux_modules: BTreeMap, pub chips: &'chips [Chip], pub extract_modularly: bool, pub submodule_mode: SubmoduleMode, pub shr_carry_summary_mode: ShrCarrySummaryMode, + pub phase: ExtractionPhase, + pub local_only: bool, + /// Hidden witness rows used by shifted re-evaluation should add + /// constraints, but must not leak new module interface ports. + pub capture_interface: bool, /// Fixed assignments used to specialize expressions during extraction. /// This allows opcodes/selectors to fold to constants before dispatch logic. pub specialization_env: BTreeMap, @@ -57,10 +203,15 @@ pub struct ConcretePendingTask { pub main_vars: Vec, pub multiplicity: PicusExpr, pub selector: String, + pub capture_interface: bool, } impl ConcretePendingTask { - // Gets the var number located at column `col_num` in the target chip + /// Returns the concrete Picus variable id assigned to a target chip column. + /// + /// Deferred sub-chip extraction stores a concrete main-row assignment for the + /// callee chip. This helper maps a column index in that callee back to the + /// variable id used in the caller's module. pub fn get_actual_var_num_for_col(&self, col_num: usize) -> usize { assert!(col_num < self.main_vars.len()); let main_var = self.main_vars[col_num]; @@ -75,10 +226,16 @@ impl ConcretePendingTask { pub struct SymbolicPendingTask { pub selector: PicusExpr, pub multiplicity: PicusExpr, + pub capture_interface: bool, } impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { - /// Constructor for the builder + /// Builds a Picus extraction builder for one chip under one phase/environment. + /// + /// The builder materializes the trace rows needed by `phase`, optionally + /// specializes selected columns to constants, and records whether this pass + /// should contribute interface ports or only additional constraints. + #[allow(clippy::too_many_arguments)] pub fn new( chip_to_analyze: &'chips Chip, picus_module: PicusModule, @@ -87,9 +244,14 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { specialization_env: Option>, submodule_mode: Option, shr_carry_summary_mode: Option, + phase: Option, + capture_interface: Option, ) -> Self { let width = chip_to_analyze.air.width(); let specialization_env = specialization_env.unwrap_or_default(); + let phase = phase.unwrap_or(ExtractionPhase::SingleRow); + let local_only = chip_to_analyze.local_only(); + let capture_interface = capture_interface.unwrap_or(true); // Initialize the public values. let public_values = (0..ZKM_PROOF_NUM_PV_ELTS).map(PicusAtom::new_var).collect(); // Initialize the preprocessed and main traces. @@ -102,6 +264,22 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { } else { (0..width).map(PicusAtom::new_var).collect() }; + if !local_only { + for row_idx in 1..phase.row_count(local_only) { + let mut next_row = (0..width).map(|_| fresh_picus_var()).collect::>(); + // Only the immediate successor row participates in the primary + // phase interface. Later rows exist solely to witness the shifted + // AIR pass and remain fully existential. + if row_idx == 1 { + if let Some(next_is_real) = phase.next_is_real() { + if let Some(is_real_idx) = chip_to_analyze.picus_info().is_real_index { + next_row[is_real_idx] = PicusAtom::Const(next_is_real); + } + } + } + main.extend(next_row); + } + } // Specialize main-row variables to constants for this extraction pass. // We key by variable id instead of column index so this also works for // sub-chip builders whose main vars may be remapped. @@ -118,12 +296,16 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { main: RowMajorMatrix::new(main, width), public_values, picus_module, + global_send_outputs: Vec::new(), aux_modules, chips, extract_modularly: false, submodule_mode: submodule_mode.unwrap_or(SubmoduleMode::Inline), shr_carry_summary_mode: shr_carry_summary_mode .unwrap_or(ShrCarrySummaryMode::AbstractModule), + phase, + local_only, + capture_interface, specialization_env, concrete_pending_tasks: Vec::new(), symbolic_pending_tasks: Vec::new(), @@ -142,6 +324,226 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { self.submodule_mode == SubmoduleMode::Ignore } + fn push_input_port(&mut self, expr: PicusExpr) { + if self.capture_interface { + self.flush_pending_expr_bindings(); + self.picus_module.inputs.push(expr); + } + } + + fn push_output_port(&mut self, expr: PicusExpr) { + if self.capture_interface { + self.flush_pending_expr_bindings(); + self.picus_module.outputs.push(expr); + } + } + + fn push_global_output_port(&mut self, expr: PicusExpr) { + if self.capture_interface { + self.flush_pending_expr_bindings(); + self.picus_module.outputs.push(expr.clone()); + self.global_send_outputs.push(expr); + } + } + + /// Emits any queued `fresh = expr` bindings produced by oversized-expression + /// reification before appending the next real module item. + /// + /// This keeps the thresholding logic in `PicusExpr` arithmetic small and + /// local while still making the resulting fresh variables explicit in the + /// extracted module. We flush eagerly at builder sinks so pending bindings do + /// not drift across unrelated constraints or module calls. + fn flush_pending_expr_bindings(&mut self) { + Self::flush_pending_expr_bindings_into(&mut self.picus_module); + } + + fn flush_pending_expr_bindings_into(module: &mut PicusModule) { + for (fresh, expr) in drain_pending_expr_bindings() { + module.constraints.push(PicusConstraint::Eq(Box::new(PicusExpr::Sub( + Box::new(PicusExpr::Var(fresh)), + Box::new(expr), + )))); + } + } + + fn push_constraint(&mut self, constraint: PicusConstraint) { + self.flush_pending_expr_bindings(); + self.picus_module.constraints.push(constraint); + } + + fn push_constraint_into(module: &mut PicusModule, constraint: PicusConstraint) { + Self::flush_pending_expr_bindings_into(module); + module.constraints.push(constraint); + } + + fn push_call(&mut self, call: PicusCall) { + self.flush_pending_expr_bindings(); + self.picus_module.calls.push(call); + } + + fn flatten_projection_ranges( + source_row: &[PicusAtom], + ranges: &[(usize, usize, String)], + ) -> Vec { + let mut exprs = Vec::new(); + for (start, end, _) in ranges { + assert!(*start <= *end && *end <= source_row.len()); + #[allow(clippy::needless_range_loop)] + for col_idx in *start..*end { + exprs.push(source_row[col_idx].into()); + } + } + exprs + } + + fn bind_projection_ports( + module: &mut PicusModule, + ports: &[PicusExpr], + source_exprs: &[PicusExpr], + ) { + assert_eq!(ports.len(), source_exprs.len()); + for (port, source_expr) in ports.iter().zip(source_exprs) { + Self::push_constraint_into( + module, + PicusConstraint::new_equality(port.clone(), source_expr.clone()), + ); + } + } + + fn add_guarded_constraint(&mut self, is_real: PicusExpr, constraint: PicusConstraint) { + match is_real { + PicusExpr::Const(0) => {} + PicusExpr::Const(1) => self.push_constraint(constraint), + _ => self.push_constraint(constraint.apply_multiplier(is_real)), + } + } + + fn and_constraints( + left: PicusConstraint, + right: PicusConstraint, + rest: impl IntoIterator, + ) -> PicusConstraint { + rest.into_iter().fold(PicusConstraint::And(Box::new(left), Box::new(right)), |acc, next| { + PicusConstraint::And(Box::new(acc), Box::new(next)) + }) + } + + fn expose_row_ranges_as_outputs(&mut self, row_idx: usize, ranges: &[(usize, usize, String)]) { + if !self.capture_interface { + return; + } + self.flush_pending_expr_bindings(); + let width = self.main.width(); + let row = self.main.row_slice(row_idx); + for (start, end, _) in ranges { + assert!(*start <= *end && *end <= width); + for col_idx in *start..*end { + let expr: PicusExpr = row[col_idx].into(); + if !self.picus_module.outputs.contains(&expr) { + self.picus_module.outputs.push(expr); + } + } + } + } + + fn expose_row_ranges_as_inputs(&mut self, row_idx: usize, ranges: &[(usize, usize, String)]) { + if !self.capture_interface { + return; + } + self.flush_pending_expr_bindings(); + let width = self.main.width(); + let row = self.main.row_slice(row_idx); + for (start, end, _) in ranges { + assert!(*start <= *end && *end <= width); + for col_idx in *start..*end { + let expr: PicusExpr = row[col_idx].into(); + if !self.picus_module.inputs.contains(&expr) { + self.picus_module.inputs.push(expr); + } + } + } + } + + /// Exposes the primary-row outputs explicitly annotated in `PicusInfo`. + pub fn expose_annotated_primary_outputs(&mut self, output_ranges: &[(usize, usize, String)]) { + self.expose_row_ranges_as_outputs(0, output_ranges); + } + + /// Exposes carried state that must be supplied as inputs to this phase. + /// + /// Transition inputs always come from the current row, even for phases that + /// also materialize successor rows. + pub fn expose_transition_inputs(&mut self, transition_input_ranges: &[(usize, usize, String)]) { + if !self.phase.exposes_transition_inputs(self.local_only) { + return; + } + // Transition inputs always refer to the current row's carried state. + self.expose_row_ranges_as_inputs(0, transition_input_ranges); + } + + /// Exposes the immediate successor row's annotated transition outputs. + /// + /// Hidden witness rows used by shifted evaluation stay existential and are + /// never surfaced through the module interface. + pub fn expose_transition_outputs( + &mut self, + transition_output_ranges: &[(usize, usize, String)], + ) { + if self.local_only || !self.phase.exposes_next_row_outputs(self.local_only) { + return; + } + // Only expose the immediate successor row. Any third row is existential + // support for the shifted AIR evaluation and should stay hidden. + self.expose_row_ranges_as_outputs(1, transition_output_ranges); + } + + /// Exposes every primary-row column except those explicitly marked as inputs. + /// + /// This is the broadest output mode and is intended for debugging or very + /// aggressive interface generation. + pub fn expose_primary_row_non_inputs_as_outputs( + &mut self, + input_ranges: &[(usize, usize, String)], + ) { + self.flush_pending_expr_bindings(); + let width = self.main.width(); + let mut is_input = vec![false; width]; + #[allow(clippy::needless_range_loop)] + for (start, end, _) in input_ranges { + assert!(*start <= *end && *end <= width); + for idx in *start..*end { + is_input[idx] = true; + } + } + + for (col_idx, atom) in self.main.row_slice(0).iter().enumerate() { + if is_input[col_idx] { + continue; + } + let expr: PicusExpr = (*atom).into(); + if !self.picus_module.outputs.contains(&expr) { + self.picus_module.outputs.push(expr); + } + } + } + + /// Exposes the entire immediate successor row as outputs. + /// + /// This is only used by the broad `AllNonInputsAreOutputs` mode and still + /// stops at the immediate successor row. + pub fn expose_full_next_row_as_outputs(&mut self) { + if self.local_only || !self.phase.exposes_next_row_outputs(self.local_only) { + return; + } + self.flush_pending_expr_bindings(); + for atom in self.main.row_slice(1).iter() { + let expr: PicusExpr = (*atom).into(); + if !self.picus_module.outputs.contains(&expr) { + self.picus_module.outputs.push(expr); + } + } + } + /// Precise summary for ByteOpcode::ShrCarry. /// /// Inputs/outputs follow byte interaction ordering: @@ -162,19 +564,19 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { let num_bits_to_shift = values[4].clone(); // Base range constraints. - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( out.clone() * multiplicity.clone(), PicusExpr::Const(255), )); - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( input.clone() * multiplicity.clone(), PicusExpr::Const(255), )); - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( carry.clone() * multiplicity.clone(), PicusExpr::Const(255), )); - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( num_bits_to_shift.clone() * multiplicity.clone(), PicusExpr::Const(7), )); @@ -201,9 +603,7 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { )), ) }; - self.picus_module - .constraints - .push(PicusConstraint::Implies(Box::new(cond), Box::new(consequence))); + self.push_constraint(PicusConstraint::Implies(Box::new(cond), Box::new(consequence))); } } @@ -220,11 +620,24 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { fn add_default_bitwise_byte_call(&mut self, values: &[PicusExpr]) { let byte_mod_name = "byte_interaction_mod".to_string(); if !self.aux_modules.contains_key(&byte_mod_name) { - let byte_mod = PicusModule::build_empty(byte_mod_name.clone(), 2, 1); + let mut byte_mod = PicusModule::build_empty(byte_mod_name.clone(), 2, 1); + // The abstract bitwise-byte helper intentionally omits the opcode semantics, but its + // interface still represents a byte operation. Keep those width guarantees explicit so + // callers cannot satisfy the module with out-of-range field elements. + for expr in [ + byte_mod.inputs[0].clone(), + byte_mod.inputs[1].clone(), + byte_mod.outputs[0].clone(), + ] { + Self::push_constraint_into( + &mut byte_mod, + PicusConstraint::new_leq(expr, PicusExpr::Const(255)), + ); + } self.aux_modules.insert(byte_mod_name.clone(), byte_mod); } assert!(values.len() == 5); - self.picus_module.calls.push(PicusCall::new(byte_mod_name, &values[1..2], &values[3..5])); + self.push_call(PicusCall::new(byte_mod_name, &values[1..2], &values[3..5])); } fn try_add_and_127_optimization(&mut self, values: &[PicusExpr]) -> bool { @@ -236,17 +649,19 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { } let var_hi = fresh_picus_expr(); - self.picus_module.constraints.push(PicusConstraint::new_lt(values[1].clone(), 128.into())); - self.picus_module.constraints.push(PicusConstraint::new_bit(var_hi.clone())); - self.picus_module.constraints.push(PicusConstraint::new_equality( + self.push_constraint(PicusConstraint::new_lt(values[1].clone(), 128.into())); + self.push_constraint(PicusConstraint::new_bit(var_hi.clone())); + self.push_constraint(PicusConstraint::new_equality( values[3].clone(), var_hi * 128 + values[1].clone(), )); true } - /// Gets a chip by name or panics if no chip is found. Kept as a slice since the number of chips is small - /// < 200 + /// Looks up a chip by name in the extraction universe. + /// + /// The extractor keeps chips in a flat slice because the total count is + /// small, so a linear scan is sufficient here. pub fn get_chip(&self, name: &str) -> &'chips Chip { self.chips .iter() @@ -265,7 +680,7 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { assert!(*v < 256); continue; } else { - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( val.clone() * multiplicity.clone(), 255.into(), )) @@ -277,7 +692,7 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { assert!(*v < 65536); continue; } else { - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( val.clone() * multiplicity.clone(), 65535.into(), )) @@ -292,27 +707,48 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { continue; } let fresh_picus_var: PicusExpr = fresh_picus_expr(); - self.picus_module.constraints.push(PicusConstraint::new_leq( + self.push_constraint(PicusConstraint::new_leq( fresh_picus_var.clone() * multiplicity.clone(), picus127_const.clone(), )); - self.picus_module.constraints.push(PicusConstraint::Eq(Box::new( + self.push_constraint(PicusConstraint::Eq(Box::new( multiplicity.clone() * msb.clone() * (msb.clone() - PicusExpr::Const(1)), ))); let decomp = byte.clone() - (msb.clone() * PicusExpr::Const(128) + fresh_picus_var); - self.picus_module - .constraints - .push(PicusConstraint::Eq(Box::new(multiplicity.clone() * decomp))); + self.push_constraint(PicusConstraint::Eq(Box::new( + multiplicity.clone() * decomp, + ))); } } else if v == (ByteOpcode::ShrCarry as u64) { match self.shr_carry_summary_mode { ShrCarrySummaryMode::AbstractModule => { if !self.aux_modules.contains_key("ShrCarry") { - let carry_module = + let mut carry_module = PicusModule::build_empty("ShrCarry".to_string(), 2, 2); + let first_input = carry_module.inputs[0].clone(); + let second_input = carry_module.inputs[1].clone(); + let output_exprs = carry_module.outputs.clone(); + // Keep the abstract helper byte-shaped even when we do not inline + // the precise `shr_carry` semantics. The first operand and both + // returned limbs are bytes, while the rotation amount is always in + // [0, 7]. + Self::push_constraint_into( + &mut carry_module, + PicusConstraint::new_leq(first_input, PicusExpr::Const(255)), + ); + Self::push_constraint_into( + &mut carry_module, + PicusConstraint::new_leq(second_input, PicusExpr::Const(7)), + ); + for expr in output_exprs { + Self::push_constraint_into( + &mut carry_module, + PicusConstraint::new_leq(expr, PicusExpr::Const(255)), + ); + } self.aux_modules.insert("ShrCarry".to_string(), carry_module); } let shrcarry = PicusCall::new( @@ -320,23 +756,26 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { &[values[1].clone(), values[2].clone()], &[values[3].clone(), values[4].clone()], ); - self.picus_module.calls.push(shrcarry); + self.push_call(shrcarry); } ShrCarrySummaryMode::Precise => { self.summarize_shr_carry_precise(multiplicity.clone(), values); } } } else if v == (ByteOpcode::LTU as u64) { - let lt_const = PicusConstraint::new_lt(values[2].clone(), values[3].clone()); + // Byte lookup values are laid out as: [opcode, a1, a2, b, c]. + // For LTU, a1 should encode (b < c), so compare values[3] and values[4]. + let lt_const = PicusConstraint::new_lt(values[3].clone(), values[4].clone()); if let PicusExpr::Const(1) = values[1] { - self.picus_module.constraints.push(lt_const); + self.push_constraint(lt_const); } else { let bit_const = PicusConstraint::new_bit(values[1].clone()); let eq_one = PicusConstraint::new_equality(values[1].clone(), 1.into()); - self.picus_module.constraints.extend_from_slice(&[ - PicusConstraint::Iff(Box::new(eq_one), Box::new(lt_const)), - bit_const, - ]); + self.push_constraint(PicusConstraint::Iff( + Box::new(eq_one), + Box::new(lt_const), + )); + self.push_constraint(bit_const); } } else if Self::is_default_bitwise_byte_opcode(v) && !self.try_add_and_127_optimization(values) @@ -391,17 +830,17 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { assert!(matches!(pc_val, PicusExpr::Const(_) | PicusExpr::Var(_))); if let PicusExpr::Var(_) = pc_val { // if the pc is a variable we mark it as an input - self.picus_module.inputs.push(pc_val.clone()); + self.push_input_port(pc_val.clone()); } // allocate a fresh var for pc out let next_pc_out = fresh_picus_expr(); // mark it as an output - self.picus_module.outputs.push(next_pc_out.clone()); + self.push_output_port(next_pc_out.clone()); // assign it conditionally to the corresponding element in the value array - self.picus_module.constraints.push(eq_mul(&multiplicity, &values[3], &next_pc_out)); + self.push_constraint(eq_mul(&multiplicity, &values[3], &next_pc_out)); // the cpu table should constrain next_pc = pc + 4 always due to delay-slot semantics of MIPS // so we add that constraint here - self.picus_module.constraints.push(PicusConstraint::new_equality( + self.push_constraint(PicusConstraint::new_equality( values[3].clone(), pc_val.clone() + PicusExpr::Const(4), )); @@ -425,8 +864,8 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { // Always expose `next_next_pc` as an output; it is often zero but still part of the // instruction interface. let next_next_pc_out = fresh_picus_expr(); - self.picus_module.constraints.push(eq_mul(&multiplicity, &values[4], &next_next_pc_out)); - self.picus_module.outputs.push(next_next_pc_out); + self.push_constraint(eq_mul(&multiplicity, &values[4], &next_next_pc_out)); + self.push_output_port(next_next_pc_out); // We need to mark some of the register values as inputs and other values as outputs. // In particular, the parameters `b` and `c` to `receive_instruction` are inputs and // parameter `a` is an output when `is_sequential` is 1. `b` and `c` are at indexes 11-14 and 15-18 in `values` whereas @@ -435,34 +874,34 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { for value in values.iter().take(11).skip(7) { let a_var = fresh_picus_expr(); if op_a_immutable { - self.picus_module.inputs.push(a_var.clone()); + self.push_input_port(a_var.clone()); } else { - self.picus_module.outputs.push(a_var.clone()); + self.push_output_port(a_var.clone()); } - self.picus_module.constraints.push(eq_mul(&multiplicity, value, &a_var)); + self.push_constraint(eq_mul(&multiplicity, value, &a_var)); // Mirrors CPU's limb range check: crates/core/machine/src/cpu/air/register.rs // (`builder.slice_range_check_u8(&local.op_a_access.access.value.0, local.is_real)`). - self.picus_module.constraints.push(u8_range(&a_var)); + self.push_constraint(u8_range(&a_var)); } for value in values.iter().take(15).skip(11) { let b_var = fresh_picus_expr(); - self.picus_module.inputs.push(b_var.clone()); - self.picus_module.constraints.push(eq_mul(&multiplicity, value, &b_var)); + self.push_input_port(b_var.clone()); + self.push_constraint(eq_mul(&multiplicity, value, &b_var)); } for value in values.iter().take(19).skip(15) { let c_var = fresh_picus_expr(); - self.picus_module.inputs.push(c_var.clone()); - self.picus_module.constraints.push(eq_mul(&multiplicity, value, &c_var)); + self.push_input_port(c_var.clone()); + self.push_constraint(eq_mul(&multiplicity, value, &c_var)); } // Route HI values by is_rw_a. for value in values.iter().take(23).skip(19) { let hi_var = fresh_picus_expr(); if is_rw_a { - self.picus_module.inputs.push(hi_var.clone()); + self.push_input_port(hi_var.clone()); } else { - self.picus_module.outputs.push(hi_var.clone()); + self.push_output_port(hi_var.clone()); } - self.picus_module.constraints.push(eq_mul(&multiplicity, value, &hi_var)); + self.push_constraint(eq_mul(&multiplicity, value, &hi_var)); } } @@ -487,40 +926,248 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { let addr_var = fresh_picus_expr(); if is_send { - self.picus_module.inputs.push(addr_var.clone()); + self.push_input_port(addr_var.clone()); } else { - self.picus_module.outputs.push(addr_var.clone()); + self.push_output_port(addr_var.clone()); } - self.picus_module.constraints.push(eq_mul(&multiplicity, &values[2], &addr_var)); + self.push_constraint(eq_mul(&multiplicity, &values[2], &addr_var)); for value in values.iter().skip(3) { let value_var = fresh_picus_expr(); if is_send { - self.picus_module.inputs.push(value_var.clone()); - self.picus_module - .constraints - .push(PicusConstraint::new_lt(value_var.clone(), PicusExpr::Const(255))); + self.push_input_port(value_var.clone()); + self.push_constraint(PicusConstraint::new_lt( + value_var.clone(), + PicusExpr::Const(255), + )); } else { - self.picus_module.outputs.push(value_var.clone()); + self.push_output_port(value_var.clone()); } - self.picus_module.constraints.push(eq_mul(&multiplicity, value, &value_var)); + self.push_constraint(eq_mul(&multiplicity, value, &value_var)); } } - fn get_main_vars_for_call(&mut self, message_values: &[PicusExpr]) -> Option> { - let opcode_spec = match message_values[6].clone() { - PicusExpr::Const(v) => { - assert!(v < Opcode::UNIMPL as u64); - spec_for(Opcode::try_from(v as u8).unwrap()) - } - _ => panic!("Opcode should be constant"), + // Program lookups are encoded as: + // [pc, instruction fields...] + // For Picus determinism extraction, we treat the fetched program row as fixed + // external context to the sending chip. That means `pc` and all instruction + // fields become inputs to the current module rather than a separate submodule call. + fn handle_program_send(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + if matches!(multiplicity, PicusExpr::Const(0)) { + return; + } + assert!( + values.len() >= 2, + "Expected program lookup to include at least pc and instruction fields" + ); + + let eq_mul = |multiplicity: &PicusExpr, val: &PicusExpr, var: &PicusExpr| { + PicusConstraint::new_equality(var.clone(), val.clone() * multiplicity.clone()) + }; + + for value in values { + let input_var = fresh_picus_expr(); + self.push_input_port(input_var.clone()); + self.push_constraint(eq_mul(&multiplicity, value, &input_var)); + } + } + + // When `send_instruction` comes from CPU with a symbolic opcode, we cannot dispatch into one + // concrete instruction chip yet. Instead we summarize the CPU/instruction-chip contract: + // + // - the instruction payload consumed by opcode chips is treated as input context to CPU + // - `next_pc` and `next_next_pc` stay outputs because they are the control-flow values the + // instruction chips are responsible for fixing + // - `opcode` is assumed deterministic so the symbolic dispatch itself is stable + // - non-sequential instructions must determine `next_next_pc`, which we encode directly as + // `is_sequential = 0 => det(next_next_pc)` + // + // We intentionally ignore `shard` and `clk` here. They are routing metadata for memory and + // syscall timing, not part of the opcode-level contract Picus is trying to prove. + fn handle_send_instruction_symbolic(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + if matches!(multiplicity, PicusExpr::Const(0)) { + return; + } + assert_eq!(values.len(), 28, "Expected instruction lookup to contain 28 values"); + + const PC_IDX: usize = 2; + const NEXT_PC_IDX: usize = 3; + const NEXT_NEXT_PC_IDX: usize = 4; + const NUM_EXTRA_CYCLES_IDX: usize = 5; + const OPCODE_IDX: usize = 6; + const A_START: usize = 7; + const HI_END: usize = 23; + const OP_A_IMMUTABLE_IDX: usize = 23; + const IS_RW_A_IDX: usize = 24; + const IS_CHECK_MEMORY_IDX: usize = 25; + const IS_HALT_IDX: usize = 26; + const IS_SEQUENTIAL_IDX: usize = 27; + + let eq_mul = |multiplicity: &PicusExpr, val: &PicusExpr, var: &PicusExpr| { + PicusConstraint::new_equality(var.clone(), val.clone() * multiplicity.clone()) + }; + + let bind_input = |builder: &mut Self, value: &PicusExpr| -> PicusExpr { + let input_var = fresh_picus_expr(); + builder.push_input_port(input_var.clone()); + builder.push_constraint(eq_mul(&multiplicity, value, &input_var)); + input_var + }; + let bind_output = |builder: &mut Self, value: &PicusExpr| -> PicusExpr { + let output_var = fresh_picus_expr(); + builder.push_output_port(output_var.clone()); + builder.push_constraint(eq_mul(&multiplicity, value, &output_var)); + output_var + }; + + bind_input(self, &values[PC_IDX]); + bind_input(self, &values[NUM_EXTRA_CYCLES_IDX]); + let opcode_in = bind_input(self, &values[OPCODE_IDX]); + bind_output(self, &values[NEXT_PC_IDX]); + let next_next_pc_out = bind_output(self, &values[NEXT_NEXT_PC_IDX]); + + for value in values.iter().take(HI_END).skip(A_START) { + bind_input(self, value); + } + + bind_input(self, &values[OP_A_IMMUTABLE_IDX]); + bind_input(self, &values[IS_RW_A_IDX]); + bind_input(self, &values[IS_CHECK_MEMORY_IDX]); + bind_input(self, &values[IS_HALT_IDX]); + let is_sequential_in = bind_input(self, &values[IS_SEQUENTIAL_IDX]); + + self.picus_module.assume_deterministic.push(opcode_in); + self.push_constraint(PicusConstraint::Implies( + Box::new(PicusConstraint::Eq(Box::new(is_sequential_in))), + Box::new(PicusConstraint::Det(Box::new(next_next_pc_out))), + )); + } + + fn handle_receive_syscall(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + if matches!(multiplicity, PicusExpr::Const(0)) { + return; + } + assert_eq!(values.len(), 5, "Expected syscall lookup to contain 5 values"); + + let eq_mul = |multiplicity: &PicusExpr, val: &PicusExpr, var: &PicusExpr| { + PicusConstraint::new_equality(var.clone(), val.clone() * multiplicity.clone()) }; - let target_chip = self.get_chip(opcode_spec.chip); + + let syscall_id_var = fresh_picus_expr(); + self.push_input_port(syscall_id_var.clone()); + self.push_constraint(eq_mul(&multiplicity, &values[2], &syscall_id_var)); + + for value in values.iter().skip(3) { + let input_var = fresh_picus_expr(); + self.push_input_port(input_var.clone()); + self.push_constraint(eq_mul(&multiplicity, value, &input_var)); + } + } + + // Syscall result bridges are encoded as: + // [shard, clk, result_lo, result_hi, arg1_lo, arg1_hi, arg2_lo, arg2_hi] + // + // For Picus we care about the functional contract, not the bridge metadata, so shard/clk stay + // hidden. The syscall result stays exposed in the same half-word form as the lookup + // (`result_lo`, `result_hi`), and the argument halves are inputs. + // + // We intentionally keep everything at half-word granularity here. The interaction itself is + // defined over packed u16 limbs, so Picus should not invent a wider `u32` result or a finer + // byte-level decomposition. We still assert that each argument half is a `u16`, because those + // bounds are part of the bridge contract and make the interface explicit in the extracted + // module. + fn handle_syscall_result_interaction(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + if matches!(multiplicity, PicusExpr::Const(0)) { + return; + } + assert_eq!(values.len(), 8, "Expected syscall result lookup to contain 8 values"); + + const RESULT_LO_IDX: usize = 2; + const RESULT_HI_IDX: usize = 3; + const ARG1_LO_IDX: usize = 4; + const ARG1_HI_IDX: usize = 5; + const ARG2_LO_IDX: usize = 6; + const ARG2_HI_IDX: usize = 7; + + let eq_mul = |multiplicity: &PicusExpr, val: &PicusExpr, expr: &PicusExpr| { + PicusConstraint::new_equality(expr.clone(), val.clone() * multiplicity.clone()) + }; + let u16_bound = |expr: PicusExpr| PicusConstraint::new_leq(expr, PicusExpr::Const(65535)); + + for result_half in [&values[RESULT_LO_IDX], &values[RESULT_HI_IDX]] { + let output_var = fresh_picus_expr(); + self.push_output_port(output_var.clone()); + self.push_constraint(eq_mul(&multiplicity, result_half, &output_var)); + } + + for halfword in + [&values[ARG1_LO_IDX], &values[ARG1_HI_IDX], &values[ARG2_LO_IDX], &values[ARG2_HI_IDX]] + { + let input_var = fresh_picus_expr(); + self.push_input_port(input_var.clone()); + self.push_constraint(eq_mul(&multiplicity, halfword, &input_var)); + self.push_constraint(u16_bound(input_var)); + } + } + + fn handle_receive_global(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + if matches!(multiplicity, PicusExpr::Const(0)) { + return; + } + + let eq_mul = |multiplicity: &PicusExpr, val: &PicusExpr, var: &PicusExpr| { + PicusConstraint::new_equality(var.clone(), val.clone() * multiplicity.clone()) + }; + + for value in values { + let input_var = fresh_picus_expr(); + self.push_input_port(input_var.clone()); + self.push_constraint(eq_mul(&multiplicity, value, &input_var)); + } + } + + fn handle_send_global(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + if matches!(multiplicity, PicusExpr::Const(0)) { + return; + } + + let eq_mul = |multiplicity: &PicusExpr, val: &PicusExpr, var: &PicusExpr| { + PicusConstraint::new_equality(var.clone(), val.clone() * multiplicity.clone()) + }; + + for value in values { + let output_var = fresh_picus_expr(); + self.push_global_output_port(output_var.clone()); + self.push_constraint(eq_mul(&multiplicity, value, &output_var)); + } + } + + fn add_abstract_syscall_call(&mut self, multiplicity: PicusExpr, values: &[PicusExpr]) { + let module_name = "SyscallLookup".to_string(); + if !self.aux_modules.contains_key(&module_name) { + // Picus requires modules to expose at least one output, even for abstract summaries. + let syscall_module = PicusModule::build_empty(module_name.clone(), values.len(), 1); + self.aux_modules.insert(module_name.clone(), syscall_module); + } + + let inputs = + values.iter().map(|value| value.clone() * multiplicity.clone()).collect::>(); + let dummy_output = fresh_picus_expr(); + self.push_call(PicusCall::new(module_name, &[dummy_output], &inputs)); + } + + fn get_main_vars_for_named_call( + &mut self, + chip_name: &str, + arg_to_colname: &[(IndexSlice, &'static str)], + message_values: &[PicusExpr], + ) -> Option> { + let target_chip = self.get_chip(chip_name); let mut target_main_vals: Vec = (0..target_chip.air.width()).map(|_| fresh_picus_var()).collect(); let target_picus_info = target_chip.picus_info(); - for (slice, name) in opcode_spec.arg_to_colname { + for (slice, name) in arg_to_colname { let colrange = target_picus_info.name_to_colrange.get(*name).unwrap(); match *slice { IndexSlice::Range { start, end } => { @@ -533,7 +1180,7 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { } else { let id = fresh_picus_var_id(); let fresh_var = PicusAtom::Var(id); - self.picus_module.constraints.push(PicusConstraint::new_equality( + self.push_constraint(PicusConstraint::new_equality( PicusExpr::Var(id), message_values[i].clone(), )); @@ -549,7 +1196,7 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { target_main_vals[colrange.0] = PicusAtom::Const(c); } else { let fresh_var = fresh_picus_var_id(); - self.picus_module.constraints.push(PicusConstraint::new_equality( + self.push_constraint(PicusConstraint::new_equality( PicusExpr::Var(fresh_var), message_values[col].clone(), )); @@ -560,11 +1207,26 @@ impl<'chips, A: MachineAir> PicusBuilder<'chips, A> { } Some(target_main_vals) } + + fn get_main_vars_for_call(&mut self, message_values: &[PicusExpr]) -> Option> { + let opcode_spec = match message_values[6].clone() { + PicusExpr::Const(v) => { + assert!(v < Opcode::UNIMPL as u64); + spec_for(Opcode::try_from(v as u8).unwrap()) + } + _ => panic!("Opcode should be constant"), + }; + self.get_main_vars_for_named_call( + opcode_spec.chip, + opcode_spec.arg_to_colname, + message_values, + ) + } } impl<'chips, A: MachineAir> PairBuilder for PicusBuilder<'chips, A> { fn preprocessed(&self) -> Self::M { - todo!() + self.preprocessed.clone() } } @@ -572,12 +1234,18 @@ impl<'chips, A: MachineAir> AirBuilderWithPublicValues for PicusBuilder<'c type PublicVar = PicusAtom; fn public_values(&self) -> &[Self::PublicVar] { - todo!() + &self.public_values } } impl<'chips, A: MachineAir> MessageBuilder> for PicusBuilder<'chips, A> { fn send(&mut self, message: AirLookup, _scope: zkm_stark::LookupScope) { + // The "top" extraction path is meant to preserve only polynomial constraints + // emitted directly by the chip AIR. Interaction lowering adds derived ports, + // helper calls, and sub-chip routing, all of which should be absent there. + if self.submodule_mode == SubmoduleMode::Ignore { + return; + } // Apply specialization first so opcode routing can see concrete values whenever // selector assignments make them decidable. let specialized_values: Vec = @@ -590,6 +1258,9 @@ impl<'chips, A: MachineAir> MessageBuilder> for Picus LookupKind::Memory => { self.handle_memory_interaction(specialized_multiplicity, &specialized_values, true); } + LookupKind::Program => { + self.handle_program_send(specialized_multiplicity, &specialized_values); + } LookupKind::Instruction => { if self.submodule_mode == SubmoduleMode::Ignore { return; @@ -599,10 +1270,13 @@ impl<'chips, A: MachineAir> MessageBuilder> for Picus assert!(v < Opcode::UNIMPL as u64); spec_for(Opcode::try_from(v as u8).unwrap()) } - _ => panic!( - "Expected opcode val to be a constant after specialization: Got: {}", - specialized_values[6] - ), + _ => { + self.handle_send_instruction_symbolic( + specialized_multiplicity, + &specialized_values, + ); + return; + } }; let target_chip = self.get_chip(opcode_spec.chip); let main_vars = self.get_main_vars_for_call(&specialized_values); @@ -612,30 +1286,78 @@ impl<'chips, A: MachineAir> MessageBuilder> for Picus main_vars: vars, multiplicity: specialized_multiplicity, selector: opcode_spec.selector.to_string(), + capture_interface: self.capture_interface, }); } else { self.symbolic_pending_tasks.push(SymbolicPendingTask { selector: specialized_values[6].clone(), multiplicity: specialized_multiplicity, + capture_interface: self.capture_interface, }) } } + LookupKind::Syscall => { + if matches!(specialized_multiplicity, PicusExpr::Const(0)) { + return; + } + + if self.submodule_mode == SubmoduleMode::Ignore { + return; + } + + if let Some(syscall_spec) = spec_for_sender(&self.picus_module.name) { + let main_vars = self.get_main_vars_for_named_call( + syscall_spec.chip, + syscall_spec.arg_to_colname, + &specialized_values, + ); + if let Some(vars) = main_vars { + self.concrete_pending_tasks.push(ConcretePendingTask { + chip_name: syscall_spec.chip.to_string(), + main_vars: vars, + multiplicity: specialized_multiplicity, + selector: syscall_spec.selector.to_string(), + capture_interface: self.capture_interface, + }); + return; + } + } + + self.add_abstract_syscall_call(specialized_multiplicity, &specialized_values); + } + LookupKind::SyscallResult => { + self.handle_syscall_result_interaction( + specialized_multiplicity, + &specialized_values, + ); + } + LookupKind::Global => { + if matches!(specialized_multiplicity, PicusExpr::Const(0)) + || self.submodule_mode == SubmoduleMode::Ignore + { + return; + } + + self.handle_send_global(specialized_multiplicity, &specialized_values); + } _ => todo!("handle send: {}", message.kind), } } fn receive(&mut self, message: AirLookup, _scope: zkm_stark::LookupScope) { - // initialize another chip - // call eval with builder? let specialized_values: Vec = message.values.iter().map(|expr| self.specialize_expr(expr)).collect(); let specialized_multiplicity = self.specialize_expr(&message.multiplicity); + + if self.submodule_mode == SubmoduleMode::Ignore { + if message.kind == LookupKind::Instruction { + self.picus_module.assume_deterministic.push(specialized_values[6].clone()); + } + return; + } + match message.kind { LookupKind::Instruction => { - if self.submodule_mode == SubmoduleMode::Ignore { - self.picus_module.assume_deterministic.push(specialized_values[6].clone()); - return; - } self.handle_receive_instruction(specialized_multiplicity, &specialized_values); } LookupKind::Memory => { @@ -645,11 +1367,462 @@ impl<'chips, A: MachineAir> MessageBuilder> for Picus false, ); } + LookupKind::Syscall => { + self.handle_receive_syscall(specialized_multiplicity, &specialized_values); + } + LookupKind::SyscallResult => { + self.handle_syscall_result_interaction( + specialized_multiplicity, + &specialized_values, + ); + } + LookupKind::Global => { + self.handle_receive_global(specialized_multiplicity, &specialized_values); + } _ => todo!("handle receive: {}", message.kind), } } } +impl<'chips, A: MachineAir> OperationSummaryAirBuilder for PicusBuilder<'chips, A> { + fn is_known_one(&self, expr: &Self::Expr) -> bool { + matches!(self.specialize_expr(expr), PicusExpr::Const(1)) + } + + fn try_emit_is_zero_summary( + &mut self, + input: Self::Expr, + result: Self::Expr, + is_real: Self::Expr, + ) -> bool { + if self.is_selector_module_builder() { + return false; + } + self.add_guarded_constraint(is_real.clone(), PicusConstraint::new_bit(result.clone())); + self.add_guarded_constraint( + is_real.clone(), + PicusConstraint::new_equality(result.clone() * input.clone(), PicusExpr::Const(0)), + ); + self.add_guarded_constraint( + is_real, + PicusConstraint::Implies( + Box::new(PicusConstraint::new_equality(input, PicusExpr::Const(0))), + Box::new(PicusConstraint::new_equality(result, PicusExpr::Const(1))), + ), + ); + true + } + + fn try_emit_is_zero_word_summary( + &mut self, + input: Word, + is_lower_half_zero: Self::Expr, + is_upper_half_zero: Self::Expr, + result: Self::Expr, + is_real: Self::Expr, + ) -> bool { + if self.is_selector_module_builder() { + return false; + } + for flag in [is_lower_half_zero.clone(), is_upper_half_zero.clone(), result.clone()] { + self.add_guarded_constraint(is_real.clone(), PicusConstraint::new_bit(flag)); + } + + for limb in [input[0].clone(), input[1].clone()] { + self.add_guarded_constraint( + is_real.clone(), + PicusConstraint::new_equality( + is_lower_half_zero.clone() * limb, + PicusExpr::Const(0), + ), + ); + } + for limb in [input[2].clone(), input[3].clone()] { + self.add_guarded_constraint( + is_real.clone(), + PicusConstraint::new_equality( + is_upper_half_zero.clone() * limb, + PicusExpr::Const(0), + ), + ); + } + + let lower_zero = Self::and_constraints( + PicusConstraint::new_equality(input[0].clone(), PicusExpr::Const(0)), + PicusConstraint::new_equality(input[1].clone(), PicusExpr::Const(0)), + [], + ); + let upper_zero = Self::and_constraints( + PicusConstraint::new_equality(input[2].clone(), PicusExpr::Const(0)), + PicusConstraint::new_equality(input[3].clone(), PicusExpr::Const(0)), + [], + ); + + self.add_guarded_constraint( + is_real.clone(), + PicusConstraint::Implies( + Box::new(lower_zero), + Box::new(PicusConstraint::new_equality( + is_lower_half_zero.clone(), + PicusExpr::Const(1), + )), + ), + ); + self.add_guarded_constraint( + is_real.clone(), + PicusConstraint::Implies( + Box::new(upper_zero), + Box::new(PicusConstraint::new_equality( + is_upper_half_zero.clone(), + PicusExpr::Const(1), + )), + ), + ); + + self.add_guarded_constraint( + is_real, + PicusConstraint::new_equality(result, is_lower_half_zero * is_upper_half_zero), + ); + true + } + + fn try_emit_koala_bear_word_range_summary( + &mut self, + input: Word, + is_real: Self::Expr, + ) -> bool { + if self.is_selector_module_builder() { + return false; + } + // This is the exact semantic collapse of the current AIR in + // `KoalaBearWordRangeChecker`: + // - the most-significant byte must be < 128 + // - if that byte is exactly 127, the lower three limbs must sum to 0 + // + // Intentionally, this does not add byte constraints for the lower + // three limbs because the exact AIR does not prove them here. + self.add_guarded_constraint( + is_real.clone(), + PicusConstraint::new_leq(input[3].clone(), 127.into()), + ); + self.add_guarded_constraint( + is_real, + PicusConstraint::Implies( + Box::new(PicusConstraint::new_equality(input[3].clone(), PicusExpr::Const(127))), + Box::new(PicusConstraint::new_equality( + input[0].clone() + input[1].clone() + input[2].clone(), + PicusExpr::Const(0), + )), + ), + ); + true + } + + fn try_emit_memory_timestamp_summary( + &mut self, + do_check: Self::Expr, + shard: Self::Expr, + clk: Self::Expr, + prev_shard: Self::Expr, + prev_clk: Self::Expr, + compare_clk: Self::Expr, + diff_16bit_limb: Self::Expr, + diff_8bit_limb: Self::Expr, + ) -> bool { + if self.is_selector_module_builder() { + return false; + } + let module_name = "MemoryTimestampCheck".to_string(); + if !self.aux_modules.contains_key(&module_name) { + // Picus currently requires every helper module to expose at least + // one output. This checker is semantically output-free, so we add a + // single dummy result and constrain it to the constant 0. + let mut timestamp_module = PicusModule::build_empty(module_name.clone(), 8, 1); + let do_check = timestamp_module.inputs[0].clone(); + let shard = timestamp_module.inputs[1].clone(); + let clk = timestamp_module.inputs[2].clone(); + let prev_shard = timestamp_module.inputs[3].clone(); + let prev_clk = timestamp_module.inputs[4].clone(); + let compare_clk = timestamp_module.inputs[5].clone(); + let diff_16bit_limb = timestamp_module.inputs[6].clone(); + let diff_8bit_limb = timestamp_module.inputs[7].clone(); + let dummy_output = timestamp_module.outputs[0].clone(); + + // Keep the synthetic output fixed so the helper remains a pure + // checker module from the caller's perspective. + Self::push_constraint_into( + &mut timestamp_module, + PicusConstraint::new_equality(dummy_output, PicusExpr::Const(0)), + ); + + // Exact `eval_memory_access_timestamp` semantics: + // - compare_clk is a guarded bit + // - if compare_clk = 1, we compare clks within the same shard + // - otherwise we compare shard indices + // - diff limbs form a guarded 24-bit decomposition of + // current_comp - prev_comp - 1 + Self::push_constraint_into( + &mut timestamp_module, + PicusConstraint::new_bit(compare_clk.clone()).apply_multiplier(do_check.clone()), + ); + Self::push_constraint_into( + &mut timestamp_module, + PicusConstraint::new_equality(shard.clone(), prev_shard.clone()) + .apply_multiplier(do_check.clone() * compare_clk.clone()), + ); + + let prev_comp_value = compare_clk.clone() * prev_clk.clone() + + (PicusExpr::Const(1) - compare_clk.clone()) * prev_shard.clone(); + let current_comp_value = compare_clk.clone() * clk.clone() + + (PicusExpr::Const(1) - compare_clk.clone()) * shard.clone(); + let diff_minus_one = current_comp_value - prev_comp_value - PicusExpr::Const(1); + + Self::push_constraint_into( + &mut timestamp_module, + PicusConstraint::new_equality( + diff_minus_one, + diff_16bit_limb.clone() + diff_8bit_limb.clone() * PicusExpr::Const(1 << 16), + ) + .apply_multiplier(do_check.clone()), + ); + Self::push_constraint_into( + &mut timestamp_module, + PicusConstraint::new_leq(diff_16bit_limb, PicusExpr::Const(65535)) + .apply_multiplier(do_check.clone()), + ); + Self::push_constraint_into( + &mut timestamp_module, + PicusConstraint::new_leq(diff_8bit_limb, PicusExpr::Const(255)) + .apply_multiplier(do_check.clone()), + ); + + self.aux_modules.insert(module_name.clone(), timestamp_module); + } + + let dummy_output = fresh_picus_expr(); + self.push_call(PicusCall::new( + module_name, + &[dummy_output], + &[ + do_check, + shard, + clk, + prev_shard, + prev_clk, + compare_clk, + diff_16bit_limb, + diff_8bit_limb, + ], + )); + true + } + + fn try_emit_projected_summary( + &mut self, + module_name: &str, + projection_info: &zkm_stark::PicusProjectionInfo, + current_inputs: &[Self::Expr], + current_outputs: &[Self::Expr], + source_width: usize, + build_exact: F, + ) -> bool + where + F: FnOnce(&mut Self, &[Self::Var]), + { + self.try_emit_projected_summary_with_hidden_consts( + module_name, + projection_info, + current_inputs, + current_outputs, + source_width, + &[], + build_exact, + ) + } + + fn try_emit_projected_summary_with_hidden_consts( + &mut self, + module_name: &str, + projection_info: &zkm_stark::PicusProjectionInfo, + current_inputs: &[Self::Expr], + current_outputs: &[Self::Expr], + source_width: usize, + hidden_consts: &[(usize, u64)], + build_exact: F, + ) -> bool + where + F: FnOnce(&mut Self, &[Self::Var]), + { + if self.is_selector_module_builder() { + return false; + } + let projected_input_len: usize = + projection_info.input_ranges.iter().map(|(start, end, _)| end - start).sum(); + let projected_output_len: usize = + projection_info.output_ranges.iter().map(|(start, end, _)| end - start).sum(); + assert_eq!(current_inputs.len(), projected_input_len); + assert_eq!(current_outputs.len(), projected_output_len); + + if !self.aux_modules.contains_key(module_name) { + let mut hidden_source_row = Vec::with_capacity(source_width); + for _ in 0..source_width { + hidden_source_row.push(fresh_picus_var()); + } + for (col_idx, value) in hidden_consts { + assert!(*col_idx < source_width); + hidden_source_row[*col_idx] = PicusAtom::Const(*value); + } + + let mut nested_module = PicusModule::build_empty( + module_name.to_string(), + current_inputs.len(), + current_outputs.len(), + ); + let projected_source_inputs = + Self::flatten_projection_ranges(&hidden_source_row, &projection_info.input_ranges); + let projected_source_outputs = + Self::flatten_projection_ranges(&hidden_source_row, &projection_info.output_ranges); + let formal_inputs = nested_module.inputs.clone(); + let formal_outputs = nested_module.outputs.clone(); + Self::bind_projection_ports( + &mut nested_module, + &formal_inputs, + &projected_source_inputs, + ); + Self::bind_projection_ports( + &mut nested_module, + &formal_outputs, + &projected_source_outputs, + ); + + let mut nested_builder = PicusBuilder { + preprocessed: self.preprocessed.clone(), + main: self.main.clone(), + public_values: self.public_values.clone(), + picus_module: nested_module, + global_send_outputs: Vec::new(), + aux_modules: BTreeMap::new(), + chips: self.chips, + extract_modularly: self.extract_modularly, + submodule_mode: self.submodule_mode, + shr_carry_summary_mode: self.shr_carry_summary_mode, + phase: self.phase, + local_only: self.local_only, + capture_interface: false, + specialization_env: BTreeMap::new(), + concrete_pending_tasks: Vec::new(), + symbolic_pending_tasks: Vec::new(), + }; + + build_exact(&mut nested_builder, &hidden_source_row); + + for (name, module) in nested_builder.aux_modules { + self.aux_modules.entry(name).or_insert(module); + } + self.aux_modules.entry(module_name.to_string()).or_insert(nested_builder.picus_module); + } + + self.push_call(PicusCall::new(module_name.to_string(), current_outputs, current_inputs)); + true + } + + /// Emit an auxiliary module for an exact sub-AIR whose internal witness + /// spans multiple phase rows. + /// + /// This mirrors `try_emit_projected_summary`, but instead of hiding a + /// single source row it hides a full phase-shaped trace matrix. The caller + /// still sees only the projected semantic boundary from the hidden local + /// row; all other hidden rows remain existential to the nested module. + fn try_emit_hidden_subair_summary( + &mut self, + module_name: &str, + projection_info: &zkm_stark::PicusProjectionInfo, + current_inputs: &[Self::Expr], + current_outputs: &[Self::Expr], + source_width: usize, + source_local_only: bool, + build_exact: F, + ) -> bool + where + F: FnOnce(&mut Self), + { + if self.is_selector_module_builder() { + return false; + } + let projected_input_len: usize = + projection_info.input_ranges.iter().map(|(start, end, _)| end - start).sum(); + let projected_output_len: usize = + projection_info.output_ranges.iter().map(|(start, end, _)| end - start).sum(); + assert_eq!(current_inputs.len(), projected_input_len); + assert_eq!(current_outputs.len(), projected_output_len); + + if !self.aux_modules.contains_key(module_name) { + let row_count = self.phase.row_count(source_local_only); + let mut hidden_main = Vec::with_capacity(source_width * row_count); + for _ in 0..(source_width * row_count) { + hidden_main.push(fresh_picus_var()); + } + // Projections are interpreted against the hidden local row. Any + // successor rows exist solely so the nested exact AIR can witness + // its own `next`-row constraints. + let hidden_local_row = hidden_main[..source_width].to_vec(); + + let mut nested_module = PicusModule::build_empty( + module_name.to_string(), + current_inputs.len(), + current_outputs.len(), + ); + let projected_source_inputs = + Self::flatten_projection_ranges(&hidden_local_row, &projection_info.input_ranges); + let projected_source_outputs = + Self::flatten_projection_ranges(&hidden_local_row, &projection_info.output_ranges); + let formal_inputs = nested_module.inputs.clone(); + let formal_outputs = nested_module.outputs.clone(); + Self::bind_projection_ports( + &mut nested_module, + &formal_inputs, + &projected_source_inputs, + ); + Self::bind_projection_ports( + &mut nested_module, + &formal_outputs, + &projected_source_outputs, + ); + + let mut nested_builder = PicusBuilder { + preprocessed: self.preprocessed.clone(), + main: RowMajorMatrix::new(hidden_main, source_width), + public_values: self.public_values.clone(), + picus_module: nested_module, + global_send_outputs: Vec::new(), + aux_modules: BTreeMap::new(), + chips: self.chips, + extract_modularly: self.extract_modularly, + submodule_mode: self.submodule_mode, + shr_carry_summary_mode: self.shr_carry_summary_mode, + phase: self.phase, + local_only: source_local_only, + capture_interface: false, + specialization_env: BTreeMap::new(), + concrete_pending_tasks: Vec::new(), + symbolic_pending_tasks: Vec::new(), + }; + + // Populate the nested module with the original exact sub-AIR over + // the hidden phase-shaped witness matrix. + build_exact(&mut nested_builder); + + for (name, module) in nested_builder.aux_modules { + self.aux_modules.entry(name).or_insert(module); + } + self.aux_modules.entry(module_name.to_string()).or_insert(nested_builder.picus_module); + } + + self.push_call(PicusCall::new(module_name.to_string(), current_outputs, current_inputs)); + true + } +} + impl<'chips, A: MachineAir> AirBuilder for PicusBuilder<'chips, A> { type F = Felt; type Var = PicusAtom; @@ -662,18 +1835,22 @@ impl<'chips, A: MachineAir> AirBuilder for PicusBuilder<'chips, A> { } fn is_first_row(&self) -> Self::Expr { - todo!() + PicusExpr::Const(self.phase.is_first_row().into()) } fn is_last_row(&self) -> Self::Expr { - todo!() + PicusExpr::Const(self.phase.is_last_row().into()) } - fn is_transition_window(&self, _size: usize) -> Self::Expr { - todo!() + fn is_transition_window(&self, size: usize) -> Self::Expr { + if size == 2 { + PicusExpr::Const(self.phase.is_transition(self.local_only).into()) + } else { + panic!("PicusBuilder only supports a transition window size of 2") + } } fn assert_zero>(&mut self, x: I) { - self.picus_module.constraints.push(PicusConstraint::Eq(Box::new(x.into()))) + self.push_constraint(PicusConstraint::Eq(Box::new(x.into()))) } } diff --git a/crates/picus/src/syscall_spec.rs b/crates/picus/src/syscall_spec.rs new file mode 100644 index 000000000..d7a0661a0 --- /dev/null +++ b/crates/picus/src/syscall_spec.rs @@ -0,0 +1,33 @@ +use crate::opcode_spec::IndexSlice; + +/// Picus specialization for syscall sends that can be routed to a concrete table. +#[derive(Clone, Debug, Default)] +pub struct SyscallSpec { + /// Selector used to specialize the callee during extraction. + pub selector: &'static str, + /// Chip that receives the syscall lookup. + pub chip: &'static str, + /// Maps the syscall lookup payload into callee columns. + pub arg_to_colname: &'static [(IndexSlice, &'static str)], +} + +/// Returns the syscall routing spec for a sender chip name, when the send can be concretely +/// lowered into another extracted chip. +pub fn spec_for_sender(sender: &str) -> Option { + use IndexSlice::Single; + + match sender { + "SyscallInstrs" => Some(SyscallSpec { + selector: "is_real", + chip: "SyscallCore", + arg_to_colname: &[ + (Single(0), "shard"), + (Single(1), "clk"), + (Single(2), "syscall_id"), + (Single(3), "arg1"), + (Single(4), "arg2"), + ], + }), + _ => None, + } +} diff --git a/crates/prover/vk_map.bin b/crates/prover/vk_map.bin index 7d1a6db88..12d52f3ea 100644 Binary files a/crates/prover/vk_map.bin and b/crates/prover/vk_map.bin differ diff --git a/crates/recursion/core/src/air/builder.rs b/crates/recursion/core/src/air/builder.rs index 08aee0234..526eb063a 100644 --- a/crates/recursion/core/src/air/builder.rs +++ b/crates/recursion/core/src/air/builder.rs @@ -2,7 +2,7 @@ use core::iter::once; use p3_air::{AirBuilder, AirBuilderWithPublicValues}; use p3_field::FieldAlgebra; use zkm_stark::{ - air::{AirLookup, BaseAirBuilder, LookupScope, MachineAirBuilder}, + air::{AirLookup, BaseAirBuilder, LookupScope, MachineAirBuilder, OperationSummaryAirBuilder}, LookupKind, }; @@ -13,11 +13,17 @@ use super::{ /// A trait which contains all helper methods for building Ziren recursion machine AIRs. pub trait ZKMRecursionAirBuilder: - MachineAirBuilder + RecursionMemoryAirBuilder + RecursionLookupAirBuilder + MachineAirBuilder + + RecursionMemoryAirBuilder + + RecursionLookupAirBuilder + + OperationSummaryAirBuilder { } -impl ZKMRecursionAirBuilder for AB {} +impl + ZKMRecursionAirBuilder for AB +{ +} impl RecursionMemoryAirBuilder for AB {} impl RecursionLookupAirBuilder for AB {} diff --git a/crates/recursion/core/src/builder.rs b/crates/recursion/core/src/builder.rs index 2257859e9..764c1db26 100644 --- a/crates/recursion/core/src/builder.rs +++ b/crates/recursion/core/src/builder.rs @@ -3,16 +3,22 @@ use std::iter::once; use p3_air::AirBuilderWithPublicValues; use p3_field::FieldAlgebra; use zkm_stark::{ - air::{AirLookup, BaseAirBuilder, LookupScope, MachineAirBuilder}, + air::{AirLookup, BaseAirBuilder, LookupScope, MachineAirBuilder, OperationSummaryAirBuilder}, LookupKind, }; use crate::{air::Block, Address}; /// A trait which contains all helper methods for building Ziren recursion machine AIRs. -pub trait ZKMRecursionAirBuilder: MachineAirBuilder + RecursionAirBuilder {} +pub trait ZKMRecursionAirBuilder: + MachineAirBuilder + RecursionAirBuilder + OperationSummaryAirBuilder +{ +} -impl ZKMRecursionAirBuilder for AB {} +impl + ZKMRecursionAirBuilder for AB +{ +} impl RecursionAirBuilder for AB {} pub trait RecursionAirBuilder: BaseAirBuilder { diff --git a/crates/recursion/core/src/chips/exp_reverse_bits.rs b/crates/recursion/core/src/chips/exp_reverse_bits.rs index 04dd3ebeb..fb78bfeb9 100644 --- a/crates/recursion/core/src/chips/exp_reverse_bits.rs +++ b/crates/recursion/core/src/chips/exp_reverse_bits.rs @@ -10,7 +10,10 @@ use std::borrow::BorrowMut; use tracing::instrument; use zkm_core_machine::utils::pad_rows_fixed; use zkm_derive::AlignedBorrow; -use zkm_stark::air::{BaseAirBuilder, ExtensionAirBuilder, MachineAir, ZKMAirBuilder}; +use zkm_stark::{ + air::{BaseAirBuilder, ExtensionAirBuilder, MachineAir}, + ZKMAirBuilder, +}; #[cfg(feature = "sys")] use crate::ExpReverseBitsEvent; diff --git a/crates/stark/Cargo.toml b/crates/stark/Cargo.toml index caebd0f8f..04facd03d 100644 --- a/crates/stark/Cargo.toml +++ b/crates/stark/Cargo.toml @@ -7,6 +7,9 @@ repository = { workspace = true } keywords = { workspace = true } categories = { workspace = true } +[features] +picus = [] + [dependencies] p3-air = { workspace = true } p3-field = { workspace = true } @@ -48,4 +51,4 @@ rayon-scan = "0.1.1" num-traits = "0.2.19" sysinfo = "0.30.13" -[dev-dependencies] \ No newline at end of file +[dev-dependencies] diff --git a/crates/stark/src/air/builder.rs b/crates/stark/src/air/builder.rs index 5a9abe0da..9edeb54f3 100644 --- a/crates/stark/src/air/builder.rs +++ b/crates/stark/src/air/builder.rs @@ -184,6 +184,173 @@ pub trait ByteAirBuilder: BaseAirBuilder { } } +/// Optional hooks for builders that want to replace exact operation AIR with a summary. +/// +/// Builders should return `true` only when they have emitted a semantically sound replacement for +/// the exact constraints. Returning `false` tells the caller to proceed with exact lowering. +pub trait OperationSummaryAirBuilder: AirBuilder { + /// Returns whether `expr` is known to be the constant one in the current + /// extraction context. + /// + /// This is useful for guarded local operations whose exact AIR is only + /// functional when their enable flag is active. Builders can use this to + /// decide whether an exact operation is safe to outline as a submodule, or + /// whether it must remain inlined under its original guard. + fn is_known_one(&self, _expr: &Self::Expr) -> bool { + false + } + + fn try_emit_is_zero_summary( + &mut self, + _input: Self::Expr, + _result: Self::Expr, + _is_real: Self::Expr, + ) -> bool { + false + } + + fn try_emit_is_zero_word_summary( + &mut self, + _input: Word, + _is_lower_half_zero: Self::Expr, + _is_upper_half_zero: Self::Expr, + _result: Self::Expr, + _is_real: Self::Expr, + ) -> bool { + false + } + + /// Optional hook for replacing the exact `KoalaBearWordRangeChecker` AIR + /// with an equivalent semantic summary. + /// + /// This summary should preserve the current operation semantics only. In + /// particular, it should not invent missing byte-range assumptions for the + /// lower three limbs; those must come from surrounding AIR if they are + /// required for soundness. + fn try_emit_koala_bear_word_range_summary( + &mut self, + _input: Word, + _is_real: Self::Expr, + ) -> bool { + false + } + + /// Optional hook for replacing the exact memory timestamp ordering AIR + /// with a compact checker-style summary. + /// + /// This hook is intentionally scoped to the arithmetic part of + /// `eval_memory_access_timestamp`; the surrounding memory send/receive + /// interactions remain the caller's responsibility. Builders that support + /// this hook may emit an auxiliary module with a dummy constant output when + /// their target IR requires every module call to produce at least one + /// result. + #[allow(clippy::too_many_arguments)] + fn try_emit_memory_timestamp_summary( + &mut self, + _do_check: Self::Expr, + _shard: Self::Expr, + _clk: Self::Expr, + _prev_shard: Self::Expr, + _prev_clk: Self::Expr, + _compare_clk: Self::Expr, + _diff_16bit_limb: Self::Expr, + _diff_8bit_limb: Self::Expr, + ) -> bool { + false + } + + /// Optional hook for replacing a large exact operation AIR with a semantic + /// module call that exposes only projected inputs/outputs. + /// + /// `projection_info` describes which ranges inside the hidden witness row + /// correspond to the caller-visible semantic boundary. Builders that + /// support this hook should: + /// - emit an auxiliary module whose interface is the flattened projected + /// inputs/outputs, + /// - keep the rest of the witness row existential/internal, and + /// - use `build_exact` to populate the auxiliary module with the original + /// exact constraints over a fresh hidden witness row of width + /// `source_width`. + /// + /// Returning `false` leaves the caller responsible for emitting the exact + /// inline constraints instead. + fn try_emit_projected_summary( + &mut self, + _module_name: &str, + _projection_info: &crate::air::PicusProjectionInfo, + _current_inputs: &[Self::Expr], + _current_outputs: &[Self::Expr], + _source_width: usize, + _build_exact: F, + ) -> bool + where + F: FnOnce(&mut Self, &[Self::Var]), + { + false + } + + /// Variant of [`Self::try_emit_projected_summary`] that lets the caller + /// pin selected hidden witness columns to constants inside the outlined + /// module. + /// + /// This is useful for guarded operations that are only outlined when their + /// enable flag is known to be one. In that case, the hidden witness should + /// reflect the specialized value directly rather than carrying an + /// additional symbolic guard through the nested module. + #[allow(clippy::too_many_arguments)] + fn try_emit_projected_summary_with_hidden_consts( + &mut self, + _module_name: &str, + _projection_info: &crate::air::PicusProjectionInfo, + _current_inputs: &[Self::Expr], + _current_outputs: &[Self::Expr], + _source_width: usize, + _hidden_consts: &[(usize, u64)], + _build_exact: F, + ) -> bool + where + F: FnOnce(&mut Self, &[Self::Var]), + { + false + } + + /// Optional hook for replacing an embedded exact sub-AIR with a semantic + /// module call whose boundary is still described by a projection. + /// + /// Unlike [`Self::try_emit_projected_summary`], this hook builds the hidden + /// witness as a full phase-shaped trace matrix rather than a single hidden + /// row. This is intended for large sub-AIRs that use `next` rows internally + /// but should still remain behind a compact caller-visible boundary. + /// + /// `projection_info` is interpreted on the hidden sub-AIR's local row + /// (row 0). Builders that support this hook should: + /// - emit an auxiliary module whose interface is the flattened projected + /// inputs/outputs taken from that hidden local row, + /// - materialize a hidden trace matrix of width `source_width` and the row + /// count implied by the current extraction phase plus `source_local_only`, + /// - run `build_exact` against that hidden trace so the full exact sub-AIR + /// remains internal to the auxiliary module. + /// + /// Returning `false` leaves the caller responsible for lowering the sub-AIR + /// inline, typically through `SubAirBuilder`. + #[allow(clippy::too_many_arguments)] + fn try_emit_hidden_subair_summary( + &mut self, + _module_name: &str, + _projection_info: &crate::air::PicusProjectionInfo, + _current_inputs: &[Self::Expr], + _current_outputs: &[Self::Expr], + _source_width: usize, + _source_local_only: bool, + _build_exact: F, + ) -> bool + where + F: FnOnce(&mut Self), + { + false + } +} + /// A trait which contains methods related to ALU lookups in an AIR. pub trait InstructionAirBuilder: BaseAirBuilder { /// Sends a MIPS instruction to be processed. @@ -567,7 +734,10 @@ pub trait MachineAirBuilder: } /// A trait which contains all helper methods for building Ziren machine AIRs. -pub trait ZKMAirBuilder: MachineAirBuilder + ByteAirBuilder + InstructionAirBuilder {} +pub trait ZKMAirBuilder: + MachineAirBuilder + ByteAirBuilder + InstructionAirBuilder + OperationSummaryAirBuilder +{ +} impl, M> MessageBuilder for FilteredAirBuilder<'_, AB> { fn send(&mut self, message: M, scope: LookupScope) { @@ -586,7 +756,11 @@ impl InstructionAirBuilder for AB {} impl ExtensionAirBuilder for AB {} impl SepticExtensionAirBuilder for AB {} impl MachineAirBuilder for AB {} -impl ZKMAirBuilder for AB {} +impl OperationSummaryAirBuilder for AB {} +impl ZKMAirBuilder + for AB +{ +} impl EmptyMessageBuilder for ProverConstraintFolder<'_, SC> {} impl EmptyMessageBuilder for VerifierConstraintFolder<'_, SC> {} diff --git a/crates/stark/src/air/machine.rs b/crates/stark/src/air/machine.rs index bed1c4ece..65bd7a675 100644 --- a/crates/stark/src/air/machine.rs +++ b/crates/stark/src/air/machine.rs @@ -78,14 +78,50 @@ pub trait MachineAir: BaseAir + 'static + Send + Sync { false } + /// Specifies whether a local-only AIR still depends on absolute row position. + /// + /// This is for chips that never read the next row, but do use predicates like + /// `when_first_row()` or `when_last_row()` to distinguish the beginning or end + /// of the trace. Such chips still need `FirstRow` / `Transition` / `LastRow` + /// extraction phases even though they are `local_only()`. + fn local_only_row_sensitive(&self) -> bool { + false + } + /// Returns information about Picus annotations on AIR columns. /// /// This includes: /// - Input ranges: columns marked with `#[picus(input)]` + /// - Output ranges: columns marked with `#[picus(output)]` + /// - Transition-input ranges: columns marked with `#[picus(transition_input)]` + /// - Transition-output ranges: columns marked with `#[picus(transition_output)]` /// - Selector indices: columns marked with `#[picus(selector)]` fn picus_info(&self) -> PicusInfo { PicusInfo::default() } + + /// Specifies whether the chip's Picus selectors partition the set of real rows. + /// + /// When this returns `true`, Picus may emit a stronger selector-shape top module asserting + /// that the sum of all selector columns is exactly `is_real`, rather than merely proving the + /// selectors are boolean and mutually exclusive. This is opt-in so existing chips keep their + /// current selector semantics unless they explicitly declare the stronger invariant. + fn selectors_partition_real_rows(&self) -> bool { + false + } + + /// Specifies whether Picus should generate a selector-specialized module for this + /// `(phase, selector)` pair. + /// + /// `phase` is the stable Picus extraction phase suffix, such as `first_row`, `transition`, + /// `boundary`, or `last_row`. `selector_name` is the user-facing selector annotation name + /// from `PicusInfo`. + /// + /// Most chips can keep the default `true`. Override this only when some phase/selector pairs + /// are impossible by trace construction and would therefore produce contradictory modules. + fn picus_selector_specialization_allowed(&self, _phase: &str, _selector_name: &str) -> bool { + true + } } /// A program that defines the control flow of a machine through a program counter. diff --git a/crates/stark/src/air/picus_info.rs b/crates/stark/src/air/picus_info.rs index 940f009e8..25a8e991a 100644 --- a/crates/stark/src/air/picus_info.rs +++ b/crates/stark/src/air/picus_info.rs @@ -1,4 +1,6 @@ use std::collections::HashMap; + +use crate::Word; /// Information about Picus annotations on AIR columns. #[derive(Debug, Clone, Default)] pub struct PicusInfo { @@ -21,6 +23,22 @@ pub struct PicusInfo { /// - `field_name` is the name of the field pub output_ranges: Vec<(usize, usize, String)>, + /// Ranges of columns whose current-row values should be exposed as inputs in transition-capable + /// extraction phases. + /// Each tuple contains (`start_index`, `end_index`, `field_name`) where: + /// - `start_index` is the first column index (inclusive) + /// - `end_index` is the last column index (exclusive) + /// - `field_name` is the name of the field + pub transition_input_ranges: Vec<(usize, usize, String)>, + + /// Ranges of columns whose next-row values should be exposed as outputs in transition-capable + /// extraction phases. + /// Each tuple contains (`start_index`, `end_index`, `field_name`) where: + /// - `start_index` is the first column index (inclusive) + /// - `end_index` is the last column index (exclusive) + /// - `field_name` is the name of the field + pub transition_output_ranges: Vec<(usize, usize, String)>, + /// Indices of columns marked as selectors. /// Each tuple contains (`column_index`, `field_name`) where: /// - `column_index` is the index of the selector column @@ -30,3 +48,51 @@ pub struct PicusInfo { /// Indices of columns marked as `is_real` pub is_real_index: Option, } + +/// Information about a semantic projection over a larger Picus-annotated witness layout. +/// +/// Unlike [`PicusInfo`], which describes whole storage fields in a concrete trace column +/// struct, a projection describes only the semantically relevant slices that should appear +/// on an operation/module boundary. This is intended for operation-level submodules where +/// most witness columns are internal and should remain existential. +#[derive(Debug, Clone, Default)] +pub struct PicusProjectionInfo { + /// Column to projected-name mapping for every byte covered by the projection. + pub col_to_name: HashMap, + /// Projected field names to concrete column ranges in the source layout. + pub name_to_colrange: HashMap, + /// Projected ranges that should be treated as module inputs. + pub input_ranges: Vec<(usize, usize, String)>, + /// Projected ranges that should be treated as module outputs. + pub output_ranges: Vec<(usize, usize, String)>, +} + +/// Helper trait for projection metadata: recover the first concrete source +/// column from either a scalar column id or a nested array of column ids. +/// +/// Projection annotations should be able to point at the semantic slice they +/// mean, for example `state.external_rounds_state[0]` rather than +/// `state.external_rounds_state[0][0]`. The derive computes the projected width +/// from the destination field type, and this trait supplies the starting source +/// column by recursively taking the first element of nested arrays. +pub trait PicusProjectionStart { + fn projection_start(&self) -> usize; +} + +impl PicusProjectionStart for usize { + fn projection_start(&self) -> usize { + *self + } +} + +impl PicusProjectionStart for [T; N] { + fn projection_start(&self) -> usize { + self[0].projection_start() + } +} + +impl PicusProjectionStart for Word { + fn projection_start(&self) -> usize { + self.0[0].projection_start() + } +} diff --git a/crates/stark/src/chip.rs b/crates/stark/src/chip.rs index a1093d576..adc0ad598 100644 --- a/crates/stark/src/chip.rs +++ b/crates/stark/src/chip.rs @@ -6,11 +6,13 @@ use p3_matrix::dense::RowMajorMatrix; use p3_uni_stark::{get_max_constraint_degree, SymbolicAirBuilder}; use p3_util::log2_ceil_usize; +#[cfg(feature = "picus")] +use crate::PicusInfo; use crate::{ air::{LookupScope, MachineAir, MultiTableAirBuilder, ZKMAirBuilder}, local_permutation_trace_width, lookup::{Lookup, LookupBuilder, LookupKind}, - scoped_lookups, PicusInfo, + scoped_lookups, }; use super::{eval_permutation_constraints, generate_permutation_trace, PROOF_MAX_NUM_PVS}; @@ -248,9 +250,22 @@ where self.air.local_only() } + fn local_only_row_sensitive(&self) -> bool { + self.air.local_only_row_sensitive() + } + + #[cfg(feature = "picus")] fn picus_info(&self) -> PicusInfo { self.air.picus_info() } + + fn selectors_partition_real_rows(&self) -> bool { + self.air.selectors_partition_real_rows() + } + + fn picus_selector_specialization_allowed(&self, phase: &str, selector_name: &str) -> bool { + self.air.picus_selector_specialization_allowed(phase, selector_name) + } } // Implement AIR directly on Chip, evaluating both execution and permutation constraints. diff --git a/crates/stark/src/lookup/builder.rs b/crates/stark/src/lookup/builder.rs index b20c6b6f2..f40be19cb 100644 --- a/crates/stark/src/lookup/builder.rs +++ b/crates/stark/src/lookup/builder.rs @@ -4,7 +4,7 @@ use p3_matrix::dense::RowMajorMatrix; use p3_uni_stark::{Entry, SymbolicExpression, SymbolicVariable}; use crate::{ - air::{AirLookup, LookupScope, MessageBuilder}, + air::{AirLookup, LookupScope, MessageBuilder, OperationSummaryAirBuilder}, PROOF_MAX_NUM_PVS, }; @@ -120,6 +120,8 @@ impl AirBuilderWithPublicValues for LookupBuilder { } } +impl OperationSummaryAirBuilder for LookupBuilder {} + fn symbolic_to_virtual_pair(expression: &SymbolicExpression) -> VirtualPairCol { if expression.degree_multiple() > 1 { panic!("degree multiple is too high"); diff --git a/crates/verifier/bn254-vk/groth16_vk.bin b/crates/verifier/bn254-vk/groth16_vk.bin index 285650560..915287056 100644 Binary files a/crates/verifier/bn254-vk/groth16_vk.bin and b/crates/verifier/bn254-vk/groth16_vk.bin differ diff --git a/crates/verifier/bn254-vk/history/v1.2.6_part_stark_vk.bin b/crates/verifier/bn254-vk/history/v1.2.6_part_stark_vk.bin new file mode 100644 index 000000000..4fed7470c Binary files /dev/null and b/crates/verifier/bn254-vk/history/v1.2.6_part_stark_vk.bin differ diff --git a/crates/verifier/bn254-vk/part_stark_vk.bin b/crates/verifier/bn254-vk/part_stark_vk.bin index 4d8c82ccf..4fed7470c 100644 Binary files a/crates/verifier/bn254-vk/part_stark_vk.bin and b/crates/verifier/bn254-vk/part_stark_vk.bin differ diff --git a/crates/verifier/bn254-vk/plonk_vk.bin b/crates/verifier/bn254-vk/plonk_vk.bin index b724c3130..54306e283 100644 Binary files a/crates/verifier/bn254-vk/plonk_vk.bin and b/crates/verifier/bn254-vk/plonk_vk.bin differ diff --git a/examples/Cargo.lock b/examples/Cargo.lock index 07c23cab2..bfa2b73df 100644 --- a/examples/Cargo.lock +++ b/examples/Cargo.lock @@ -9113,8 +9113,3 @@ dependencies = [ "cc", "pkg-config", ] - -[[patch.unused]] -name = "secp256k1" -version = "0.29.1" -source = "git+https://github.com/ziren-patches/rust-secp256k1?branch=patch-0.29.1#c7e39ee0732d0e6b5c33721038f8f1b789f77c36" diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 21430d6e4..f5e00ec45 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -76,7 +76,6 @@ tracing = "0.1.40" curve25519-dalek = { git = "https://github.com/ziren-patches/curve25519-dalek", branch = "patch-4.1.3" } ecdsa-core = { git = "https://github.com/ziren-patches/signatures", package = "ecdsa", branch = "patch-ecdsa-0.16.9" } rsa = { git = "https://github.com/ziren-patches/RustCrypto-RSA.git", branch = "patch-rsa-0.9.6" } -secp256k1 = { git = "https://github.com/ziren-patches/rust-secp256k1", branch = "patch-0.29.1" } sha2-v0-10-8 = { git = "https://github.com/ziren-patches/RustCrypto-hashes", package = "sha2", branch = "patch-sha2-0.10.8" } substrate-bn = { git = "https://github.com/ziren-patches/bn", branch = "patch-0.6.0" } k256 = { git = "https://github.com/ziren-patches/elliptic-curves", branch = "patch-k256-0.13.4" }