WIP

circuits/Cargo.toml

@@ -37,6 +37,7 @@ mozak-runner = { path = "../runner", features = ["test"] }
 proptest = "1.4"
 
 [features]
+default = ["enable_poseidon_starks"]
 enable_poseidon_starks = []
 enable_register_starks = []
 test = []

circuits/src/generation/memory.rs

@@ -159,13 +159,19 @@ pub fn generate_memory_trace<F: RichField>(
         transform_io(io_memory_public_rows),
     )
     .collect();
-
     #[cfg(feature = "enable_poseidon_starks")]
     merged_trace.extend(transform_poseidon2_sponge(poseidon2_sponge_rows));
+    let d = transform_poseidon2_sponge(poseidon2_sponge_rows).collect::<Vec<_>>();
+    println!("sponge-memory: {:?}", d);
     #[cfg(feature = "enable_poseidon_starks")]
     merged_trace.extend(transform_poseidon2_output_bytes(
         poseidon2_output_bytes_rows,
     ));
+    let d = transform_poseidon2_output_bytes(poseidon2_output_bytes_rows).collect::<Vec<_>>();
+    println!("p-output-memory: {:?}", d);
+
+    let d = transform_memory_init::<F>(memory_init_rows).collect::<Vec<_>>();
+    println!("init-memory: {:?}", d);
 
     merged_trace.sort_by_key(key);
     let mut merged_trace: Vec<_> = merged_trace

circuits/src/generation/mod.rs

@@ -13,6 +13,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 rangecheck;
@@ -59,6 +60,7 @@ use crate::generation::memoryinit::{
     generate_elf_memory_init_trace, generate_mozak_memory_init_trace,
 };
 use crate::generation::poseidon2::generate_poseidon2_trace;
+use crate::generation::poseidon2_preimage_pack::generate_poseidon2_preimage_pack_trace;
 use crate::generation::program::generate_program_rom_trace;
 use crate::stark::mozak_stark::{
     all_starks, MozakStark, PublicInputs, TableKindArray, TableKindSetBuilder,
@@ -98,6 +100,9 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     // This step generates the output only - it takes last sponge elements (with get-flag = true)
     let poseidon2_output_bytes_rows = generate_poseidon2_output_bytes_trace(&poseiden2_sponge_rows);
     #[allow(unused)]
+    let poseidon2_preimage_pack_rows =
+        generate_poseidon2_preimage_pack_trace(&poseiden2_sponge_rows);
+    #[allow(unused)]
     // This is the computation step, all the heavy logic happens here. In other words, the poseidon
     // hash
     let poseidon2_rows = generate_poseidon2_trace(&record.executed);
@@ -150,6 +155,8 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         poseidon2_sponge_stark: trace_rows_to_poly_values(poseiden2_sponge_rows),
         #[cfg(feature = "enable_poseidon_starks")]
         poseidon2_output_bytes_stark: trace_rows_to_poly_values(poseidon2_output_bytes_rows),
+        #[cfg(feature = "enable_poseidon_starks")]
+        poseidon2_preimage_pack: trace_rows_to_poly_values(poseidon2_preimage_pack_rows),
     }
     .build()
 }

circuits/src/generation/poseidon2_preimage_pack.rs

@@ -0,0 +1,67 @@
+use plonky2::hash::hash_types::RichField;
+
+use super::MIN_TRACE_LENGTH;
+use crate::poseidon2_preimage_pack::columns::Poseidon2PreimagePack;
+use crate::poseidon2_sponge::columns::Poseidon2Sponge;
+
+fn pad_trace<F: RichField>(
+    mut trace: Vec<Poseidon2PreimagePack<F>>,
+) -> Vec<Poseidon2PreimagePack<F>> {
+    trace.resize(
+        trace.len().next_power_of_two().max(MIN_TRACE_LENGTH),
+        Poseidon2PreimagePack::default(),
+    );
+    trace
+}
+
+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();
+    let trace = pad_trace(trace);
+    log::trace!("trace {:?}", trace);
+    trace
+}
+
+#[cfg(test)]
+mod tests {
+    use mozak_runner::vm::Row;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+
+    use crate::generation::poseidon2_sponge::generate_poseidon2_sponge_trace;
+    use crate::generation::MIN_TRACE_LENGTH;
+    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_output_bytes_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/generation/poseidon2_sponge.rs

@@ -2,7 +2,7 @@ use itertools::Itertools;
 use mozak_runner::vm::Row;
 use plonky2::hash::hash_types::RichField;
 use plonky2::hash::hashing::PlonkyPermutation;
-use plonky2::hash::poseidon2::Poseidon2Permutation;
+use plonky2::hash::poseidon2::{Poseidon2Permutation, WIDTH};
 
 use crate::generation::MIN_TRACE_LENGTH;
 use crate::poseidon2_sponge::columns::{Ops, Poseidon2Sponge};
@@ -40,6 +40,11 @@ fn unroll_sponge_data<F: RichField>(row: &Row<F>) -> Vec<Poseidon2Sponge<F>> {
             .sponge_data
             .get(i as usize)
             .expect("unroll_count not consistent with number of permutations");
+        let padded_fes = Poseidon2Permutation::<F>::RATE - sponge_datum.padded_count;
+        let mut is_padded = [F::ZERO; WIDTH];
+        for i in 0..padded_fes {
+            is_padded[Poseidon2Permutation::<F>::RATE - i] = F::ONE;
+        }
         unroll.push(Poseidon2Sponge {
             clk: F::from_canonical_u64(row.state.clk),
             ops,
@@ -49,6 +54,7 @@ fn unroll_sponge_data<F: RichField>(row: &Row<F>) -> Vec<Poseidon2Sponge<F>> {
             preimage: sponge_datum.preimage,
             output: sponge_datum.output,
             gen_output: sponge_datum.gen_output,
+            is_padded,
         });
         input_addr += rate_size;
         input_len -= rate_size;
@@ -76,7 +82,6 @@ pub fn generate_poseidon2_sponge_trace<F: RichField>(
 #[cfg(test)]
 mod test {
     use mozak_runner::poseidon2::MozakPoseidon2;
-    use plonky2::field::goldilocks_field::GoldilocksField;
     use plonky2::field::types::Field;
     use plonky2::hash::hashing::PlonkyPermutation;
     use plonky2::hash::poseidon2::Poseidon2Permutation;
@@ -92,7 +97,7 @@ mod test {
     #[test]
     fn generate_poseidon2_sponge_trace() {
         let data = "😇 Mozak is knowledge arguments based technology".to_string();
-        let data_len_in_bytes = MozakPoseidon2::<GoldilocksField>::padding(data.as_bytes()).len();
+        let data_len_in_bytes = MozakPoseidon2::padding(data.as_bytes()).len();
         let input_start_addr = 1024;
         let output_start_addr = 2048;
         let (_program, record) = create_poseidon2_test(&[Poseidon2Test {
@@ -105,9 +110,8 @@ mod test {
         let trace = super::generate_poseidon2_sponge_trace(&step_rows);
 
         let rate_size = Poseidon2Permutation::<F>::RATE;
-        let sponge_count = (data_len_in_bytes
-            / MozakPoseidon2::<GoldilocksField>::DATA_CAPACITY_PER_FIELD_ELEMENT)
-            / rate_size;
+        let sponge_count =
+            (data_len_in_bytes / MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT) / rate_size;
         for (i, value) in trace.iter().enumerate().take(sponge_count) {
             assert_eq!(
                 value.input_addr,

circuits/src/lib.rs

@@ -21,6 +21,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 rangecheck;

circuits/src/memory/columns.rs

@@ -1,5 +1,6 @@
 use core::ops::Add;
 
+use mozak_runner::poseidon2::MozakPoseidon2;
 use plonky2::field::extension::Extendable;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
@@ -112,21 +113,35 @@ impl<F: RichField> From<&FullWordMemory<F>> for Vec<Memory<F>> {
 
 impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Memory<F>> {
     fn from(value: &Poseidon2Sponge<F>) -> Self {
-        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()
+        if (value.ops.is_permute + value.ops.is_init_permute).is_one() {
+            // each Field element in preimage represents packed data (packed bytes)
+            (0..Poseidon2Permutation::<F>::RATE).filter(|i| value.is_padded[*i].is_one())
+                .flat_map(|i| {
+                    // base-address is an input-address + RATE * index-inside-preimage
+                    let base_address = value.input_addr
+                        + F::from_canonical_u64(
+                            u64::try_from(MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT).expect(
+                                "MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT should be cast-able to u64",
+                            ),
+                        ) * F::from_canonical_u8(u8::try_from(i).expect("i > 255"));
+                    // Throw away leading byte since "be"
+                    let packed = &value.preimage[i].clone().to_canonical_u64().to_be_bytes()
+                        [MozakPoseidon2::BYTES_PER_FIELD_ELEMENT
+                            - MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT..];
+                    (0..MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT)
+                        .map(|j| Memory {
+                            clk: value.clk,
+                            addr: base_address
+                                + F::from_canonical_u8(u8::try_from(j).expect("j > 255")),
+                            is_load: F::ONE,
+                            value: F::from_canonical_u8(packed[j]),
+                            ..Default::default()
+                        })
+                        .collect::<Vec<_>>()
                 })
                 .collect()
+        } else {
+            vec![]
         }
     }
 }

circuits/src/poseidon2_preimage_pack/columns.rs

@@ -0,0 +1,110 @@
+use itertools::Itertools;
+use mozak_runner::poseidon2::MozakPoseidon2;
+use plonky2::field::types::Field;
+use plonky2::hash::hash_types::RichField;
+
+use crate::columns_view::{columns_view_impl, make_col_map, NumberOfColumns};
+use crate::linear_combination::Column;
+use crate::poseidon2::columns::STATE_SIZE;
+use crate::poseidon2_sponge::columns::Poseidon2Sponge;
+
+// FIXME: MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT;
+pub const BYTES_COUNT: usize = MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT;
+
+#[repr(C)]
+#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)]
+pub struct Poseidon2PreimagePack<F> {
+    pub clk: F,
+    pub addr: F,
+    pub bytes: [F; MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT],
+    pub is_executed: F,
+}
+
+columns_view_impl!(Poseidon2PreimagePack);
+make_col_map!(Poseidon2PreimagePack);
+
+pub const NUM_POSEIDON2_PREIMAGE_PACK_COLS: usize = Poseidon2PreimagePack::<()>::NUMBER_OF_COLUMNS;
+
+impl<F: RichField> From<&Poseidon2Sponge<F>> for Vec<Poseidon2PreimagePack<F>> {
+    fn from(value: &Poseidon2Sponge<F>) -> Self {
+        if (value.ops.is_init_permute + value.ops.is_permute).is_one() {
+            assert!(
+                MozakPoseidon2::FIELD_ELEMENTS_RATE < STATE_SIZE,
+                "Packing RATE should be less than STATE_SIZE"
+            );
+            // FIXME: 8 should be MozakPoseidon2::FE_RATE
+            let preimage: [F; MozakPoseidon2::FIELD_ELEMENTS_RATE] = value.preimage
+                [..MozakPoseidon2::FIELD_ELEMENTS_RATE]
+                .try_into()
+                .expect("Should succeed since preimage can't be empty");
+            let mut base_address = value.input_addr;
+            let mut index = 0;
+            // For each FE of preimage we have BYTES_COUNT bytes
+            let result = preimage
+                .iter()
+                .map(|fe| {
+                    // Note: assumed `to_be_bytes`, otherwise another side of the array should be
+                    // taken
+                    let bytes: Vec<_> = fe.clone().to_canonical_u64().to_be_bytes()
+                        [MozakPoseidon2::BYTES_PER_FIELD_ELEMENT
+                            - MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT..]
+                        .into_iter()
+                        .map(|e| F::from_canonical_u8(*e))
+                        .collect();
+                    let addr = base_address;
+                    let i = index;
+                    index+=1;
+                    base_address = base_address + F::from_canonical_u64( u64::try_from(MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT).expect("Cast from usize to u64 for MozakPoseidon2::BYTES_PER_FIELD_ELEMENT should succeed"));
+                    Poseidon2PreimagePack {
+                        clk: value.clk,
+                        addr,
+                        bytes: <[F; MozakPoseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT]>::try_from(
+                            bytes,
+                        )
+                        .unwrap(),
+                        is_executed: F::ONE - value.is_padded[i],
+                    }
+                })
+                .collect_vec();
+            println!("poseidon-value: {:?}", value);
+            println!("preimage-result: {:?}", result);
+            return result;
+        }
+        vec![]
+    }
+}
+
+#[must_use]
+pub fn data_for_poseidon2_sponge<F: Field>() -> Vec<Column<F>> {
+    let data = col_map().map(Column::from);
+    vec![
+        data.clk,
+        data.addr,
+        Column::<F>::reduce_with_powers(&data.bytes, F::from_canonical_u16(1 << 8)),
+    ]
+}
+
+#[must_use]
+pub fn filter_for_poseidon2_sponge<F: Field>() -> Column<F> {
+    col_map().map(Column::from).is_executed
+}
+
+#[must_use]
+pub fn data_for_input_memory<F: Field>(index: u8) -> Vec<Column<F>> {
+    assert!(
+        usize::try_from(index).unwrap() < BYTES_COUNT,
+        "poseidon2-preimage data_for_input_memory: index can be 0..{:?}",
+        BYTES_COUNT
+    );
+    let data = col_map().map(Column::from);
+    vec![
+        data.clk,
+        Column::constant(F::ZERO),               // is_store
+        Column::constant(F::ONE),                // is_load
+        data.bytes[index as usize].clone(),      // value
+        data.addr + F::from_canonical_u8(index), // address
+    ]
+}
+
+#[must_use]
+pub fn filter_for_input_memory<F: Field>() -> Column<F> { col_map().map(Column::from).is_executed }

circuits/src/poseidon2_preimage_pack/mod.rs

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