From e5cef03d24ab91d4334890c75a1034e983d3124c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Tue, 26 Nov 2024 11:23:23 +0100
Subject: [PATCH] fix/program-size4: fix padding in edge cases

---
 ceno_zkvm/src/instructions.rs        | 2 +-
 ceno_zkvm/src/tables/mod.rs          | 3 +--
 ceno_zkvm/src/tables/ram/ram_impl.rs | 4 ++--
 ceno_zkvm/src/witness.rs             | 5 ++++-
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/ceno_zkvm/src/instructions.rs b/ceno_zkvm/src/instructions.rs
index e87675dd6..b10fced17 100644
--- a/ceno_zkvm/src/instructions.rs
+++ b/ceno_zkvm/src/instructions.rs
@@ -73,7 +73,7 @@ pub trait Instruction<E: ExtensionField> {
             })
             .collect::<Result<(), ZKVMError>>()?;
 
-        let num_padding_instances = raw_witin.num_padding_instances();
+        let num_padding_instances = raw_witin.num_padding_instances(None);
         if num_padding_instances > 0 {
             // Fill the padding based on strategy
 
diff --git a/ceno_zkvm/src/tables/mod.rs b/ceno_zkvm/src/tables/mod.rs
index f498b868e..c6c90d61d 100644
--- a/ceno_zkvm/src/tables/mod.rs
+++ b/ceno_zkvm/src/tables/mod.rs
@@ -57,8 +57,7 @@ pub fn padding_zero<F: SmallField>(
     num_cols: usize,
     num_instances: Option<usize>,
 ) {
-    // Fill the padding with zeros, if any.
-    let num_padding_instances = table.num_padding_instances();
+    let num_padding_instances = table.num_padding_instances(num_instances);
     if num_padding_instances > 0 {
         let nthreads =
             std::env::var("RAYON_NUM_THREADS").map_or(8, |s| s.parse::<usize>().unwrap_or(8));
diff --git a/ceno_zkvm/src/tables/ram/ram_impl.rs b/ceno_zkvm/src/tables/ram/ram_impl.rs
index 2f16c6bff..39c7e6d27 100644
--- a/ceno_zkvm/src/tables/ram/ram_impl.rs
+++ b/ceno_zkvm/src/tables/ram/ram_impl.rs
@@ -117,7 +117,7 @@ impl<NVRAM: NonVolatileTable + Send + Sync + Clone> NonVolatileTableConfig<NVRAM
         );
 
         let mut init_table = RowMajorMatrix::<F>::new(NVRAM::len(&self.params), num_fixed);
-        assert_eq!(init_table.num_padding_instances(), 0);
+        assert_eq!(init_table.num_padding_instances(None), 0);
 
         init_table
             .par_iter_mut()
@@ -247,7 +247,7 @@ impl<NVRAM: NonVolatileTable + Send + Sync + Clone> PubIOTableConfig<NVRAM> {
         assert!(NVRAM::len(&self.params).is_power_of_two());
 
         let mut init_table = RowMajorMatrix::<F>::new(NVRAM::len(&self.params), num_fixed);
-        assert_eq!(init_table.num_padding_instances(), 0);
+        assert_eq!(init_table.num_padding_instances(None), 0);
 
         init_table
             .par_iter_mut()
diff --git a/ceno_zkvm/src/witness.rs b/ceno_zkvm/src/witness.rs
index e85360aee..b7021b243 100644
--- a/ceno_zkvm/src/witness.rs
+++ b/ceno_zkvm/src/witness.rs
@@ -62,8 +62,11 @@ impl<T: Sized + Sync + Clone + Send + Copy> RowMajorMatrix<T> {
         self.values.len() / self.num_col - self.num_padding_rows
     }
 
-    pub fn num_padding_instances(&self) -> usize {
+    pub fn num_padding_instances(&self, num_instances: Option<usize>) -> usize {
         self.num_padding_rows
+            + num_instances
+                .map(|ni| self.num_instances() - ni)
+                .unwrap_or(0)
     }
 
     pub fn iter_rows(&self) -> Chunks<MaybeUninit<T>> {