feat: poseidon2 input pack (byte-input-padding-only)

Cargo.toml

@@ -29,6 +29,7 @@ members = [
   "signatures",
   "state",
   "wasm-demo",
+  "poseidon2",
 ]
 resolver = "2"
 

circuits/Cargo.toml

@@ -22,6 +22,7 @@ mozak-runner = { path = "../runner" }
 mozak-sdk = { path = "../sdk" }
 plonky2 = { version = "0", default-features = false }
 plonky2_maybe_rayon = { version = "0", default-features = false }
+poseidon2 = { path = "../poseidon2" }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 starky = { version = "0", default-features = false, features = ["std"] }

circuits/src/generation/mod.rs

@@ -11,6 +11,7 @@ pub mod io_memory;
 pub mod memory;
 pub mod memory_zeroinit;
 pub mod memoryinit;
+pub mod poseidon2_preimage_pack;
 pub mod xor;
 use std::borrow::Borrow;
 use std::fmt::Display;
@@ -49,6 +50,7 @@ use crate::generation::memory_zeroinit::generate_memory_zero_init_trace;
 use crate::generation::memoryinit::{
     generate_elf_memory_init_trace, generate_mozak_memory_init_trace,
 };
+use crate::generation::poseidon2_preimage_pack::generate_poseidon2_preimage_pack_trace;
 use crate::poseidon2::generation::generate_poseidon2_trace;
 use crate::poseidon2_output_bytes::generation::generate_poseidon2_output_bytes_trace;
 use crate::poseidon2_sponge::generation::generate_poseidon2_sponge_trace;
@@ -101,6 +103,8 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     let cast_list_commitment_tape_rows = generate_cast_list_commitment_tape_trace(&record.executed);
     let poseiden2_sponge_rows = generate_poseidon2_sponge_trace(&record.executed);
     let poseidon2_output_bytes_rows = generate_poseidon2_output_bytes_trace(&poseiden2_sponge_rows);
+    let poseidon2_preimage_pack_rows =
+        generate_poseidon2_preimage_pack_trace(&poseiden2_sponge_rows);
     let poseidon2_rows = generate_poseidon2_trace(&record.executed);
 
     let memory_rows = generate_memory_trace(
@@ -166,6 +170,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         poseidon2_stark: trace_rows_to_poly_values(poseidon2_rows),
         poseidon2_sponge_stark: trace_rows_to_poly_values(poseiden2_sponge_rows),
         poseidon2_output_bytes_stark: trace_rows_to_poly_values(poseidon2_output_bytes_rows),
+        poseidon2_preimage_pack: trace_rows_to_poly_values(poseidon2_preimage_pack_rows),
         tape_commitments_stark: trace_rows_to_poly_values(tape_commitments_rows),
     }
     .build()

circuits/src/generation/poseidon2_preimage_pack.rs

@@ -0,0 +1,55 @@
+use plonky2::hash::hash_types::RichField;
+
+use crate::poseidon2_preimage_pack::columns::Poseidon2PreimagePack;
+use crate::poseidon2_sponge::columns::Poseidon2Sponge;
+use crate::utils::pad_trace_with_default;
+
+pub fn generate_poseidon2_preimage_pack_trace<F: RichField>(
+    poseidon2_sponge_rows: &[Poseidon2Sponge<F>],
+) -> Vec<Poseidon2PreimagePack<F>> {
+    let trace: Vec<Poseidon2PreimagePack<F>> = poseidon2_sponge_rows
+        .iter()
+        .flat_map(Into::<Vec<Poseidon2PreimagePack<F>>>::into)
+        .collect();
+    pad_trace_with_default(trace)
+}
+
+#[cfg(test)]
+mod tests {
+    use mozak_runner::vm::Row;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+
+    use crate::generation::MIN_TRACE_LENGTH;
+    use crate::poseidon2_sponge::generation::generate_poseidon2_sponge_trace;
+    use crate::test_utils::{create_poseidon2_test, Poseidon2Test};
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::F;
+    #[test]
+    fn generate_poseidon2_preimage_pack_trace() {
+        let data = "😇 Mozak is knowledge arguments based technology".to_string();
+        let input_start_addr = 1024;
+        let output_start_addr = 2048;
+        let (_program, record) = create_poseidon2_test(&[Poseidon2Test {
+            data,
+            input_start_addr,
+            output_start_addr,
+        }]);
+
+        let step_rows = record.executed;
+
+        let sponge_trace = generate_poseidon2_sponge_trace(&step_rows);
+        let trace = super::generate_poseidon2_preimage_pack_trace(&sponge_trace);
+        // for one sponge construct we have one row with gen_output = 1.
+        // So we expect other padding data to make trace of len MIN_TRACE_LENGTH.
+        assert_eq!(trace.len(), MIN_TRACE_LENGTH);
+    }
+
+    #[test]
+    fn generate_poseidon2_trace_with_dummy() {
+        let step_rows: Vec<Row<F>> = vec![];
+        let sponge_trace = generate_poseidon2_sponge_trace(&step_rows);
+        let trace = super::generate_poseidon2_preimage_pack_trace(&sponge_trace);
+        assert_eq!(trace.len(), MIN_TRACE_LENGTH);
+    }
+}

circuits/src/lib.rs

@@ -22,6 +22,7 @@ pub mod memory_zeroinit;
 pub mod memoryinit;
 pub mod poseidon2;
 pub mod poseidon2_output_bytes;
+pub mod poseidon2_preimage_pack;
 pub mod poseidon2_sponge;
 pub mod program;
 pub mod program_multiplicities;

circuits/src/memory/columns.rs

@@ -6,6 +6,7 @@ use plonky2::hash::hashing::PlonkyPermutation;
 use plonky2::hash::poseidon2::Poseidon2Permutation;
 use plonky2::iop::ext_target::ExtensionTarget;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
+use poseidon2::mozak_poseidon2;
 
 use crate::columns_view::{columns_view_impl, make_col_map};
 use crate::cross_table_lookup::Column;
@@ -14,7 +15,7 @@ use crate::memory_halfword::columns::HalfWordMemory;
 use crate::memory_io::columns::InputOutputMemory;
 use crate::memory_zeroinit::columns::MemoryZeroInit;
 use crate::memoryinit::columns::{MemoryInit, MemoryInitCtl};
-use crate::poseidon2_output_bytes::columns::{Poseidon2OutputBytes, BYTES_COUNT};
+use crate::poseidon2_output_bytes::columns::Poseidon2OutputBytes;
 use crate::poseidon2_sponge::columns::Poseidon2Sponge;
 use crate::rangecheck::columns::RangeCheckCtl;
 use crate::stark::mozak_stark::{MemoryTable, TableWithTypedOutput};
@@ -123,36 +124,43 @@ impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Memory<F>> {
         if (value.ops.is_permute + value.ops.is_init_permute).is_zero() {
             vec![]
         } else {
-            let rate = Poseidon2Permutation::<F>::RATE;
-            // each Field element in preimage represents a byte.
-            (0..rate)
-                .map(|i| Memory {
-                    clk: value.clk,
-                    addr: value.input_addr
-                        + F::from_canonical_u8(u8::try_from(i).expect("i > 255")),
-                    is_load: F::ONE,
-                    value: value.preimage[i],
-                    ..Default::default()
+            // each Field element in preimage represents packed data (packed bytes)
+            (0..Poseidon2Permutation::<F>::RATE)
+                .flat_map(|fe_index_inside_preimage| {
+                    let base_address = value.input_addr
+                        + mozak_poseidon2::data_capacity_fe::<F>()
+                            * F::from_canonical_usize(fe_index_inside_preimage);
+                    let unpacked = mozak_poseidon2::unpack_to_field_elements(
+                        &value.preimage[fe_index_inside_preimage],
+                    );
+
+                    (0..mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT)
+                        .map(|byte_index_inside_fe| Memory {
+                            clk: value.clk,
+                            addr: base_address + F::from_canonical_usize(byte_index_inside_fe),
+                            is_load: F::ONE,
+                            value: unpacked[byte_index_inside_fe],
+                            ..Default::default()
+                        })
+                        .collect::<Vec<_>>()
                 })
                 .collect()
         }
     }
 }
 
 impl<F: RichField> From<&Poseidon2OutputBytes<F>> for Vec<Memory<F>> {
-    fn from(value: &Poseidon2OutputBytes<F>) -> Self {
-        if value.is_executed.is_zero() {
+    fn from(output: &Poseidon2OutputBytes<F>) -> Self {
+        if output.is_executed.is_zero() {
             vec![]
         } else {
-            (0..BYTES_COUNT)
-                .map(|i| Memory {
-                    clk: value.clk,
-                    addr: value.output_addr
-                        + F::from_canonical_u8(u8::try_from(i).expect(
-                            "BYTES_COUNT of poseidon output should be representable by a u8",
-                        )),
+            (0..)
+                .zip(output.output_bytes)
+                .map(|(i, value)| Memory {
+                    clk: output.clk,
+                    addr: output.output_addr + F::from_canonical_usize(i),
                     is_store: F::ONE,
-                    value: value.output_bytes[i],
+                    value,
                     ..Default::default()
                 })
                 .collect()

circuits/src/poseidon2_preimage_pack/columns.rs

@@ -0,0 +1,96 @@
+use itertools::Itertools;
+use plonky2::hash::hash_types::RichField;
+use poseidon2::mozak_poseidon2;
+
+use crate::columns_view::{columns_view_impl, make_col_map, NumberOfColumns};
+use crate::linear_combination::Column;
+use crate::linear_combination_typed::ColumnWithTypedInput;
+use crate::memory::columns::MemoryCtl;
+use crate::poseidon2::columns::STATE_SIZE;
+use crate::poseidon2_sponge::columns::Poseidon2Sponge;
+use crate::stark::mozak_stark::{Poseidon2PreimagePackTable, TableWithTypedOutput};
+
+#[repr(C)]
+#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)]
+pub struct Poseidon2PreimagePack<F> {
+    pub clk: F,
+    pub byte_addr: F,
+    pub bytes: [F; mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT],
+    pub is_executed: F,
+}
+
+columns_view_impl!(Poseidon2PreimagePack);
+make_col_map!(PACK, Poseidon2PreimagePack);
+
+pub const NUM_POSEIDON2_PREIMAGE_PACK_COLS: usize = Poseidon2PreimagePack::<()>::NUMBER_OF_COLUMNS;
+
+impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Poseidon2PreimagePack<F>> {
+    // To make it safe for user to change constants
+    #[allow(clippy::assertions_on_constants)]
+    fn from(value: &Poseidon2Sponge<F>) -> Self {
+        if (value.ops.is_init_permute + value.ops.is_permute).is_zero() {
+            vec![]
+        } else {
+            assert!(
+                mozak_poseidon2::FIELD_ELEMENTS_RATE <= STATE_SIZE,
+                "Packing RATE (FIELD_ELEMENTS_RATE) should be less or equal than STATE_SIZE"
+            );
+            let preimage: [F; mozak_poseidon2::FIELD_ELEMENTS_RATE] = value.preimage
+                [..mozak_poseidon2::FIELD_ELEMENTS_RATE]
+                .try_into()
+                .expect("Should succeed since preimage can't be empty");
+            // For each FE of preimage we have PACK_CAP bytes
+            preimage
+                .iter()
+                .enumerate()
+                .map(|(i, fe)| Poseidon2PreimagePack {
+                    clk: value.clk,
+                    byte_addr: value.input_addr
+                        + F::from_canonical_usize(i) * mozak_poseidon2::data_capacity_fe::<F>(),
+                    bytes: mozak_poseidon2::unpack_to_field_elements(fe),
+                    is_executed: F::ONE,
+                })
+                .collect_vec()
+        }
+    }
+}
+
+columns_view_impl!(Poseidon2SpongePreimagePackCtl);
+#[repr(C)]
+#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)]
+pub struct Poseidon2SpongePreimagePackCtl<T> {
+    pub clk: T,
+    pub value: T,
+    pub byte_addr: T,
+}
+#[must_use]
+pub fn lookup_for_poseidon2_sponge() -> TableWithTypedOutput<Poseidon2SpongePreimagePackCtl<Column>>
+{
+    Poseidon2PreimagePackTable::new(
+        Poseidon2SpongePreimagePackCtl {
+            clk: PACK.clk,
+            value: ColumnWithTypedInput::reduce_with_powers(PACK.bytes, 1 << 8),
+            byte_addr: PACK.byte_addr,
+        },
+        PACK.is_executed,
+    )
+}
+
+#[must_use]
+pub fn lookup_for_input_memory() -> Vec<TableWithTypedOutput<MemoryCtl<Column>>> {
+    (0..)
+        .zip(PACK.bytes)
+        .map(|(index, value)| {
+            Poseidon2PreimagePackTable::new(
+                MemoryCtl {
+                    clk: PACK.clk,
+                    is_store: ColumnWithTypedInput::constant(0),
+                    is_load: ColumnWithTypedInput::constant(1),
+                    value,
+                    addr: PACK.byte_addr + index,
+                },
+                PACK.is_executed,
+            )
+        })
+        .collect()
+}

circuits/src/poseidon2_preimage_pack/mod.rs

@@ -0,0 +1,2 @@
+pub mod columns;
+pub mod stark;