[WIP] Update keccak-builtin

cairo_programs/keccak_builtin.cairo

@@ -0,0 +1,18 @@
+%builtins keccak
+ from starkware.cairo.common.cairo_builtins import KeccakBuiltin
+ from starkware.cairo.common.keccak_state import KeccakBuiltinState
+
+ func main{keccak_ptr: KeccakBuiltin*}() {
+     assert keccak_ptr[0].input = KeccakBuiltinState(1,2,3,4,5,6,7,8);
+     let result = keccak_ptr[0].output;
+     let keccak_ptr = keccak_ptr + KeccakBuiltin.SIZE;
+     assert result.s0 = 528644516554364142278482415480021626364691973678134577961206;
+     assert result.s1 = 768681319646568210457759892191562701823009052229295869963057;
+     assert result.s2 = 1439835513376369408063324968379272676079109225238241190228026;
+     assert result.s3 = 1150396629165612276474514703759718478742374517669870754478270;
+     assert result.s4 = 1515147102575186161827863034255579930572231617017100845406254;
+     assert result.s5 = 1412568161597072838250338588041800080889949791225997426843744;
+     assert result.s6 = 982235455376248641031519404605670648838699214888770304613539;
+     assert result.s7 = 1339947803093378278438908448344904300127577306141693325151040;
+     return ();
+}

src/hint_processor/builtin_hint_processor/keccak_utils.rs

@@ -180,10 +180,3 @@ fn left_pad(bytes_vector: &mut [u8], n_zeros: usize) -> Vec<u8> {
 
     res
 }
-
-pub(crate) fn left_pad_u64(bytes_vector: &mut [u64], n_zeros: usize) -> Vec<u64> {
-    let mut res: Vec<u64> = vec![0; n_zeros];
-    res.extend(bytes_vector.iter());
-
-    res
-}

src/main.rs

@@ -38,7 +38,9 @@ struct Args {
 
 fn validate_layout(value: &str) -> Result<(), String> {
     match value {
-        "plain" | "small" | "dex" | "bitwise" | "perpetual_with_bitwise" | "all" => Ok(()),
+        "plain" | "small" | "dex" | "bitwise" | "perpetual_with_bitwise" | "all" | "recursive" => {
+            Ok(())
+        }
         _ => Err(format!("{value} is not a valid layout")),
     }
 }

src/vm/errors/runner_errors.rs

@@ -1,7 +1,7 @@
 use std::collections::HashSet;
 
 use super::memory_errors::MemoryError;
-use crate::types::relocatable::{MaybeRelocatable, Relocatable};
+use crate::types::relocatable::Relocatable;
 use felt::Felt;
 use thiserror::Error;
 
@@ -32,7 +32,7 @@ pub enum RunnerError {
     #[error("Given builtins are not in appropiate order")]
     DisorderedBuiltins,
     #[error("Expected integer at address {0:?} to be smaller than 2^{1}, Got {2}")]
-    IntegerBiggerThanPowerOfTwo(MaybeRelocatable, u32, Felt),
+    IntegerBiggerThanPowerOfTwo(Relocatable, u32, Felt),
     #[error("{0}")]
     EcOpSameXCoordinate(String),
     #[error("EcOpBuiltin: point {0:?} is not on the curve")]
@@ -79,6 +79,8 @@ pub enum RunnerError {
     Memory(#[from] MemoryError),
     #[error("keccak_builtin: Failed to get first input address")]
     KeccakNoFirstInput,
+    #[error("{0}: Expected integer at address {1}")]
+    BuiltinExpectedInteger(&'static str, Relocatable),
     #[error("keccak_builtin: Failed to convert input cells to u64 values")]
     KeccakInputCellsNotU64,
 }

src/vm/runners/builtin_runner/bitwise.rs

@@ -87,14 +87,14 @@ impl BitwiseBuiltinRunner {
         ) {
             if num_x.bits() > self.bitwise_builtin.total_n_bits as u64 {
                 return Err(RunnerError::IntegerBiggerThanPowerOfTwo(
-                    x_addr.into(),
+                    x_addr,
                     self.bitwise_builtin.total_n_bits,
                     num_x.clone(),
                 ));
             };
             if num_y.bits() > self.bitwise_builtin.total_n_bits as u64 {
                 return Err(RunnerError::IntegerBiggerThanPowerOfTwo(
-                    y_addr.into(),
+                    y_addr,
                     self.bitwise_builtin.total_n_bits,
                     num_y.clone(),
                 ));

src/vm/runners/builtin_runner/keccak.rs

@@ -1,4 +1,6 @@
-use crate::hint_processor::builtin_hint_processor::keccak_utils::left_pad_u64;
+use std::cell::RefCell;
+use std::collections::HashMap;
+
 use crate::math_utils::safe_div_usize;
 use crate::types::instance_definitions::keccak_instance_def::KeccakInstanceDef;
 use crate::types::relocatable::{MaybeRelocatable, Relocatable};
@@ -8,24 +10,25 @@ use crate::vm::vm_core::VirtualMachine;
 use crate::vm::vm_memory::memory::Memory;
 use crate::vm::vm_memory::memory_segments::MemorySegmentManager;
 use felt::Felt;
+use num_bigint::BigUint;
 use num_integer::div_ceil;
 use num_traits::{One, ToPrimitive};
 
 use super::KECCAK_BUILTIN_NAME;
 
-const KECCAK_ARRAY_LEN: usize = 25;
+const KECCAK_FELT_BYTE_SIZE: usize = 25; // 200 / 8
 
 #[derive(Debug, Clone)]
 pub struct KeccakBuiltinRunner {
     ratio: u32,
     pub base: usize,
     pub(crate) cells_per_instance: u32,
     pub(crate) n_input_cells: u32,
-    verified_addresses: Vec<Relocatable>,
     pub(crate) stop_ptr: Option<usize>,
     pub(crate) included: bool,
     state_rep: Vec<u32>,
     instances_per_component: u32,
+    cache: RefCell<HashMap<Relocatable, Felt>>,
 }
 
 impl KeccakBuiltinRunner {
@@ -36,10 +39,10 @@ impl KeccakBuiltinRunner {
             n_input_cells: instance_def._state_rep.len() as u32,
             cells_per_instance: instance_def.cells_per_builtin(),
             stop_ptr: None,
-            verified_addresses: Vec::new(),
             included,
             instances_per_component: instance_def._instance_per_component,
             state_rep: instance_def._state_rep.clone(),
+            cache: RefCell::new(HashMap::new()),
         }
     }
 
@@ -74,52 +77,57 @@ impl KeccakBuiltinRunner {
         if index < self.n_input_cells as usize {
             return Ok(None);
         }
-
-        let first_input_addr = address
-            .sub_usize(index)
-            .map_err(|_| RunnerError::KeccakNoFirstInput)?;
-
-        if self.verified_addresses.contains(&first_input_addr) {
-            return Ok(None);
+        if let Some(felt) = self.cache.borrow().get(&address) {
+            return Ok(Some(felt.into()));
         }
-
-        let mut input_felts_u64 = vec![];
-
-        for i in 0..self.n_input_cells {
-            let val = match memory.get(&(first_input_addr + i as usize)) {
-                Some(val) => val
-                    .as_ref()
-                    .get_int_ref()
-                    .and_then(|x| x.to_u64())
-                    .ok_or(RunnerError::KeccakInputCellsNotU64)?,
+        // index will always be less than address.offset, so we can safely unwrap here
+        let first_input_addr = address.sub_usize(index).unwrap();
+        let first_output_addr = first_input_addr + self.n_input_cells as usize;
+
+        let mut input_felts = vec![];
+
+        for i in 0..self.n_input_cells as usize {
+            let val = match memory.get(&(first_input_addr + i)) {
+                Some(value) => {
+                    let num = value
+                        .get_int_ref()
+                        .ok_or(RunnerError::BuiltinExpectedInteger(
+                            KECCAK_BUILTIN_NAME,
+                            first_input_addr + i,
+                        ))?;
+                    if num >= &(Felt::one() << self.state_rep[i]) {
+                        return Err(RunnerError::IntegerBiggerThanPowerOfTwo(
+                            first_input_addr + i,
+                            self.state_rep[i],
+                            num.clone(),
+                        ));
+                    }
+                    num.clone()
+                }
                 _ => return Ok(None),
             };
 
-            input_felts_u64.push(val)
+            input_felts.push(val)
         }
 
-        if let Some((i, bits)) = self.state_rep.iter().enumerate().next() {
-            let val = memory.get_integer(first_input_addr + i)?;
-            if val.as_ref() >= &(Felt::one() << *bits) {
-                return Err(RunnerError::IntegerBiggerThanPowerOfTwo(
-                    (first_input_addr + i).into(),
-                    *bits,
-                    val.into_owned(),
-                ));
-            }
-
-            let len = input_felts_u64.len();
-            let mut input_felts_u64 = left_pad_u64(&mut input_felts_u64, KECCAK_ARRAY_LEN - len)
-                .try_into()
-                .map_err(|_| RunnerError::SliceToArrayError)?;
-
-            keccak::f1600(&mut input_felts_u64);
-
-            return Ok(input_felts_u64
-                .get(address.offset - 1)
-                .map(|x| Felt::from(*x).into()));
+        let input_message: Vec<u8> = input_felts
+            .iter()
+            .flat_map(|x| Self::right_pad(&x.to_biguint().to_bytes_le(), KECCAK_FELT_BYTE_SIZE))
+            .collect();
+        let keccak_result = Self::keccak_f(&input_message)?;
+
+        let mut start_index = 0_usize;
+        for (i, bits) in self.state_rep.iter().enumerate() {
+            let end_index = start_index + *bits as usize / 8;
+            self.cache.borrow_mut().insert(
+                first_output_addr + i,
+                Felt::from(BigUint::from_bytes_le(
+                    &keccak_result[start_index..end_index],
+                )),
+            );
+            start_index = end_index;
         }
-        Ok(None)
+        Ok(self.cache.borrow().get(&address).map(|x| x.into()))
     }
 
     pub fn get_allocated_memory_units(&self, vm: &VirtualMachine) -> Result<usize, MemoryError> {
@@ -245,10 +253,39 @@ impl KeccakBuiltinRunner {
         // So the real number is 4 * 64 * 1024 = 262144.
         safe_div_usize(262144_usize, diluted_n_bits as usize).unwrap_or(0)
     }
+
+    fn right_pad(bytes: &[u8], final_size: usize) -> Vec<u8> {
+        let zeros: Vec<u8> = vec![0; final_size - bytes.len()];
+        let mut bytes_vector = bytes.to_vec();
+        bytes_vector.extend(zeros);
+        bytes_vector
+    }
+
+    fn keccak_f(input_message: &[u8]) -> Result<Vec<u8>, RunnerError> {
+        let bigint = BigUint::from_bytes_le(input_message);
+        let w = 64;
+        let keccak_input = (0..25)
+            .map(|i| {
+                ((&bigint >> (i * w)) & (BigUint::from(2_u8).pow(w) - BigUint::one())).to_u64()
+            })
+            .collect::<Option<Vec<u64>>>()
+            .ok_or(RunnerError::KeccakInputCellsNotU64)?;
+        // This unwrap wont fail as keccak_input is created from a (0..25) range
+        let mut keccak_input: [u64; 25] = keccak_input.try_into().unwrap();
+        keccak::f1600(&mut keccak_input);
+        Ok(keccak_input
+            .iter()
+            .enumerate()
+            .map(|(i, x)| BigUint::from(*x) << (i * w as usize))
+            .sum::<BigUint>()
+            .to_bytes_le())
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use num_traits::Num;
+
     use super::*;
     use crate::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor;
     use crate::relocatable;
@@ -529,9 +566,13 @@ mod tests {
         let result = builtin.deduce_memory_cell(Relocatable::from((0, 25)), &memory);
         assert_eq!(
             result,
-            Ok(Some(MaybeRelocatable::from(Felt::new(
-                3086936446498698982_u64
-            ))))
+            Ok(Some(MaybeRelocatable::from(
+                Felt::from_str_radix(
+                    "1006979841721999878391288827876533441431370448293338267890891",
+                    10
+                )
+                .unwrap()
+            )))
         );
     }
 
@@ -558,55 +599,51 @@ mod tests {
     }
 
     #[test]
-    fn deduce_memory_cell_offset_first_addr_error() {
-        let memory = memory![
-            ((0, 16), 43),
-            ((0, 17), 199),
-            ((0, 18), 0),
-            ((0, 19), 0),
-            ((0, 20), 0),
-            ((0, 21), 0),
-            ((0, 22), 0),
-            ((0, 23), 1),
-            ((0, 24), 0),
-            ((0, 25), 0),
-            ((0, 26), 43),
-            ((0, 27), 199),
-            ((0, 28), 0),
-            ((0, 29), 0),
-            ((0, 30), 0),
-            ((0, 31), 0),
-            ((0, 32), 0),
-            ((0, 33), 1),
-            ((0, 34), 0),
-            ((0, 35), 0)
-        ];
+    fn deduce_memory_cell_expected_integer() {
+        let memory = memory![((0, 0), (1, 2))];
 
         let mut builtin = KeccakBuiltinRunner::new(&KeccakInstanceDef::default(), true);
 
-        builtin.verified_addresses.push(Relocatable::from((0, 16)));
+        builtin.n_input_cells = 1;
+        builtin.cells_per_instance = 100;
+
+        let result = builtin.deduce_memory_cell(Relocatable::from((0, 1)), &memory);
+
+        assert_eq!(
+            result,
+            Err(RunnerError::BuiltinExpectedInteger(
+                KECCAK_BUILTIN_NAME,
+                (0, 0).into()
+            ))
+        );
+    }
+
+    #[test]
+    fn deduce_memory_cell_missing_input_cells() {
+        let memory = memory![((0, 1), (1, 2))];
+
+        let mut builtin = KeccakBuiltinRunner::new(&KeccakInstanceDef::default(), true);
+
+        builtin.n_input_cells = 1;
+        builtin.cells_per_instance = 100;
+
+        let result = builtin.deduce_memory_cell(Relocatable::from((0, 1)), &memory);
 
-        let result = builtin.deduce_memory_cell(Relocatable::from((0, 25)), &memory);
         assert_eq!(result, Ok(None));
     }
 
     #[test]
-    fn deduce_memory_cell_expected_integer() {
+    fn deduce_memory_cell_input_cell() {
         let memory = memory![((0, 0), (1, 2))];
 
         let mut builtin = KeccakBuiltinRunner::new(&KeccakInstanceDef::default(), true);
 
-        builtin.n_input_cells = 0;
+        builtin.n_input_cells = 1;
         builtin.cells_per_instance = 100;
 
-        let result = builtin.deduce_memory_cell(Relocatable::from((0, 99)), &memory);
+        let result = builtin.deduce_memory_cell(Relocatable::from((0, 0)), &memory);
 
-        assert_eq!(
-            result,
-            Err(RunnerError::Memory(MemoryError::ExpectedInteger(
-                (0, 0).into()
-            )))
-        );
+        assert_eq!(result, Ok(None));
     }
 
     #[test]
@@ -668,4 +705,19 @@ mod tests {
 
         assert_eq!(result, 16384);
     }
+
+    #[test]
+    fn right_pad() {
+        let num = [1_u8];
+        let padded_num = KeccakBuiltinRunner::right_pad(&num, 5);
+        assert_eq!(padded_num, vec![1, 0, 0, 0, 0]);
+    }
+
+    #[test]
+    fn keccak_f() {
+        let input_bytes = b"\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x05\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x06\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x07\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
+        let expected_output_bytes = b"\xf6\x98\x81\xe1\x00!\x1f.\xc4*\x8c\x0c\x7fF\xc8q8\xdf\xb9\xbe\x07H\xca7T1\xab\x16\x17\xa9\x11\xff-L\x87\xb2iY.\x96\x82x\xde\xbb\\up?uz:0\xee\x08\x1b\x15\xd6\n\xab\r\x0b\x87T:w\x0fH\xe7!f},\x08a\xe5\xbe8\x16\x13\x9a?\xad~<9\xf7\x03`\x8b\xd8\xa3F\x8aQ\xf9\n9\xcdD\xb7.X\xf7\x8e\x1f\x17\x9e \xe5i\x01rr\xdf\xaf\x99k\x9f\x8e\x84\\\xday`\xf1``\x02q+\x8e\xad\x96\xd8\xff\xff3<\xb6\x01o\xd7\xa6\x86\x9d\xea\xbc\xfb\x08\xe1\xa3\x1c\x06z\xab@\xa1\xc1\xb1xZ\x92\x96\xc0.\x01\x13g\x93\x87!\xa6\xa8z\x9c@\x0bY'\xe7\xa7Qr\xe5\xc1\xa3\xa6\x88H\xa5\xc0@9k:y\xd1Kw\xd5";
+        let output_bytes = KeccakBuiltinRunner::keccak_f(input_bytes);
+        assert_eq!(output_bytes, Ok(expected_output_bytes.to_vec()));
+    }
 }