Merge pull request #2480 from o1-labs/zkvm/syscalls/fix-preimagekey-l…

…ookup-revisit

o1vm/src/keccak/interpreter.rs

@@ -482,9 +482,12 @@ where
         }
     }
 
-    /// When in Squeeze mode, writes a Lookup containing the 31byte output of the hash (excludes the MSB)
+    /// When in Squeeze mode, writes a Lookup containing the 31byte output of
+    /// the hash (excludes the MSB)
     /// - if is_squeeze, adds 1 lookup
     /// - otherwise, adds 0 lookups
+    /// NOTE: this is excluding the MSB (which is then substituted with the
+    ///       file descriptor).
     fn lookup_syscall_hash(&mut self, step: Steps) {
         let bytes31 = (1..32).fold(Self::zero(), |acc, i| {
             acc * Self::two_pow(8) + self.sponge_byte(i)

o1vm/src/mips/column.rs

@@ -31,6 +31,8 @@ pub(crate) const MIPS_LENGTH_BYTES_OFF: usize = 89;
 pub(crate) const MIPS_HAS_N_BYTES_OFF: usize = 93;
 /// The maximum size of a chunk (4 bytes)
 pub(crate) const MIPS_CHUNK_BYTES_LEN: usize = 4;
+/// The location of the preimage key as a field element of 248bits
+pub(crate) const MIPS_PREIMAGE_KEY: usize = 97;
 
 /// The number of columns used for relation witness in the MIPS circuit
 pub const N_MIPS_REL_COLS: usize = SCRATCH_SIZE + 2;

o1vm/src/mips/constraints.rs

@@ -5,10 +5,9 @@ use crate::{
             ColumnAlias as MIPSColumn, MIPS_BYTE_COUNTER_OFF, MIPS_CHUNK_BYTES_LEN,
             MIPS_END_OF_PREIMAGE_OFF, MIPS_HASH_COUNTER_OFF, MIPS_HAS_N_BYTES_OFF,
             MIPS_LENGTH_BYTES_OFF, MIPS_NUM_BYTES_READ_OFF, MIPS_PREIMAGE_BYTES_OFF,
-            MIPS_PREIMAGE_CHUNK_OFF,
+            MIPS_PREIMAGE_CHUNK_OFF, MIPS_PREIMAGE_KEY,
         },
         interpreter::InterpreterEnv,
-        registers::REGISTER_PREIMAGE_KEY_START,
     },
     E,
 };
@@ -406,6 +405,9 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         // preimage in this instruction
         let this_chunk = self.variable(Self::Position::ScratchState(MIPS_PREIMAGE_CHUNK_OFF));
 
+        // The preimage key composed of 248 bits
+        let preimage_key = self.variable(Self::Position::ScratchState(MIPS_PREIMAGE_KEY));
+
         // The (at most) 4 bytes that are being processed from the preimage
         let bytes: [_; MIPS_CHUNK_BYTES_LEN] = array::from_fn(|i| {
             self.variable(Self::Position::ScratchState(MIPS_PREIMAGE_BYTES_OFF + i))
@@ -541,6 +543,11 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         // Byte checks with lookups: both preimage and length bytes are checked
         // TODO: think of a way to merge these together to perform 4 lookups
         // instead of 8 per row
+        // FIXME: understand if length bytes can ever be read together with
+        // preimage bytes. If not, then we can merge the lookups and just run
+        // 4 lookups per row for the byte checks. AKA: does the oracle always
+        // read the length bytes first and then the preimage bytes, with no
+        // overlapping?
         for byte in bytes.iter() {
             self.add_lookup(Lookup::read_one(
                 LookupTableIDs::ByteLookup,
@@ -575,14 +582,6 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         }
 
         // COMMUNICATION CHANNEL: Read hash output
-        // FIXME: check if the most significant byte is zero or 0x02
-        //        so we know what exactly needs to be passed to the lookup
-        let preimage_key = (0..8).fold(Expr::from(0), |acc, i| {
-            acc * Expr::from(2u64.pow(32))
-                + self.variable(Self::Position::ScratchState(
-                    REGISTER_PREIMAGE_KEY_START + i,
-                ))
-        });
         // If no more bytes left to be read, then the end of the preimage is
         // true.
         // TODO: keep track of counter to diminish the number of bytes at

o1vm/src/mips/witness.rs

@@ -10,7 +10,7 @@ use crate::{
             ColumnAlias as Column, MIPS_BYTE_COUNTER_OFF, MIPS_CHUNK_BYTES_LEN,
             MIPS_END_OF_PREIMAGE_OFF, MIPS_HASH_COUNTER_OFF, MIPS_HAS_N_BYTES_OFF,
             MIPS_LENGTH_BYTES_OFF, MIPS_NUM_BYTES_READ_OFF, MIPS_PREIMAGE_BYTES_OFF,
-            MIPS_PREIMAGE_CHUNK_OFF,
+            MIPS_PREIMAGE_CHUNK_OFF, MIPS_PREIMAGE_KEY,
         },
         interpreter::{
             self, ITypeInstruction, Instruction, InterpreterEnv, JTypeInstruction, RTypeInstruction,
@@ -21,6 +21,7 @@ use crate::{
 };
 use ark_ff::Field;
 use core::panic;
+use kimchi::o1_utils::Two;
 use log::{debug, info};
 use std::{
     array,
@@ -44,7 +45,9 @@ pub const NUM_INSTRUCTION_LOOKUP_TERMS: usize = 5;
 pub const NUM_LOOKUP_TERMS: usize =
     NUM_GLOBAL_LOOKUP_TERMS + NUM_DECODING_LOOKUP_TERMS + NUM_INSTRUCTION_LOOKUP_TERMS;
 // TODO: Delete and use a vector instead
-pub const SCRATCH_SIZE: usize = 97; // MIPS + hash_counter + byte_counter + eof + num_bytes_read + chunk + bytes + length + has_n_bytes
+// MIPS + hash_counter + byte_counter + eof + num_bytes_read + chunk + bytes
+// + length + has_n_bytes + chunk_bytes + preimage
+pub const SCRATCH_SIZE: usize = 98;
 
 #[derive(Clone, Default)]
 pub struct SyscallEnv {
@@ -757,6 +760,13 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreI
         if self.preimage_bytes_read == preimage_len as u64 {
             self.write_column(Column::ScratchState(MIPS_END_OF_PREIMAGE_OFF), 1);
 
+            // Store preimage key in the witness excluding the MSB as 248 bits
+            // so it can be used for the communication channel between Keccak
+            let bytes31 = (1..32).fold(Fp::zero(), |acc, i| {
+                acc * Fp::two_pow(8) + Fp::from(self.preimage_key.unwrap()[i])
+            });
+            self.write_field_column(Self::Position::ScratchState(MIPS_PREIMAGE_KEY), bytes31);
+
             debug!("Preimage has been read entirely, triggering Keccak process");
             self.keccak_env = Some(KeccakEnv::<Fp>::new(
                 self.hash_counter,