feat: store_byte memory method (#329)

* feat: store_byte memory method

* chore: fmt fixes

* chore: fmt fixes

* address pr comments

* address pr review

* rebase

* fix: var name

crates/evm/src/instructions/memory_operations.cairo

@@ -129,8 +129,7 @@ impl MemoryOperation of MemoryOperationTrait {
         let offset = self.stack.pop_usize()?;
         let value = self.stack.pop()?;
         let value: u8 = (value.low & 0xFF).try_into().unwrap();
-        let values = array![value].span();
-        self.memory.store_n(values, offset);
+        self.memory.store_byte(value, offset);
 
         Result::Ok(())
     }

crates/evm/src/memory.cairo

@@ -11,11 +11,10 @@ use utils::constants::{
 use cmp::{max};
 use utils::{
     helpers, helpers::SpanExtensionTrait, helpers::ArrayExtensionTrait, math::Exponentiation,
-    math::WrappingExponentiation
+    math::WrappingExponentiation, math::Bitshift
 };
 use debug::PrintTrait;
 
-
 #[derive(Destruct, Default)]
 struct Memory {
     items: Felt252Dict<u128>,
@@ -26,6 +25,7 @@ trait MemoryTrait {
     fn new() -> Memory;
     fn size(ref self: Memory) -> usize;
     fn store(ref self: Memory, element: u256, offset: usize);
+    fn store_byte(ref self: Memory, value: u8, offset: usize);
     fn store_n(ref self: Memory, elements: Span<u8>, offset: usize);
     fn store_padded_segment(ref self: Memory, offset: usize, length: usize, source: Span<u8>);
     fn ensure_length(ref self: Memory, length: usize);
@@ -75,6 +75,34 @@ impl MemoryImpl of MemoryTrait {
         self.store_element(element, chunk_index, offset_in_chunk);
     }
 
+
+    /// Stores a single byte into memory at a specified offset.
+    ///
+    /// # Arguments
+    ///
+    /// * `self` - A mutable reference to the `Memory` instance to store the byte in.
+    /// * `value` - The byte value to store in memory.
+    /// * `offset` - The offset within memory to store the byte at.
+    #[inline(always)]
+    fn store_byte(ref self: Memory, value: u8, offset: usize) {
+        let new_min_bytes_len = helpers::ceil_bytes_len_to_next_32_bytes_word(offset + 1);
+        self.bytes_len = cmp::max(new_min_bytes_len, self.bytes_len);
+
+        // Get offset's memory word index and left-based offset of byte in word.
+        let (chunk_index, left_offset) = u32_safe_divmod(offset, u32_as_non_zero(16));
+
+        // As the memory words are in big-endian order, we need to convert our left-based offset
+        // to a right-based one.a
+        let right_offset = 15 - left_offset;
+        let mask: u128 = 0xFF * helpers::pow2(right_offset.into() * 8);
+
+        // First erase byte value at offset, then set the new value using bitwise ops
+        let word: u128 = self.items.get(chunk_index.into());
+        let new_word = (word & ~mask) | (value.into().shl(right_offset.into() * 8));
+        self.items.insert(chunk_index.into(), new_word);
+    }
+
+
     /// Stores a span of N bytes into memory at a specified offset.
     ///
     /// This function checks the alignment of the offset to 16-byte chunks, and handles the special case where the bytes to be

crates/evm/src/tests/test_memory.cairo

@@ -510,3 +510,82 @@ fn test_store_padded_segment_should_add_n_elements_padded_with_offset_between_tw
         'Wrong memory value'
     );
 }
+#[test]
+#[available_gas(20000000)]
+fn test_store_byte_should_store_byte_at_offset() {
+    // Given
+    let mut memory = MemoryTrait::new();
+
+    // When
+    memory.store_byte(0x01, 15);
+
+    // Then
+    assert(memory.items[0] == 0x01, 'Wrong value for word 0');
+    assert(memory.items[1] == 0x00, 'Wrong value for word 1');
+    assert(memory.bytes_len == 32, 'Wrong memory length');
+}
+#[test]
+#[available_gas(20000000)]
+fn test_store_byte_should_store_byte_at_offset_2() {
+    // Given
+    let mut memory = MemoryTrait::new();
+
+    // When
+    memory.store_byte(0xff, 14);
+
+    // Then
+    assert(memory.items[0] == 0xff00, 'Wrong value for word 0');
+    assert(memory.items[1] == 0x00, 'Wrong value for word 1');
+    assert(memory.bytes_len == 32, 'Wrong memory length');
+}
+
+#[test]
+#[available_gas(20000000)]
+fn test_store_byte_should_store_byte_at_offset_in_existing_word() {
+    // Given
+    let mut memory = MemoryTrait::new();
+    memory.items.insert(0, 0xFFFF); // Set the first word in memory to 0xFFFF;
+    memory.items.insert(1, 0xFFFF);
+
+    // When
+    memory.store_byte(0x01, 30);
+
+    // Then
+    assert(memory.items[0] == 0xFFFF, 'Wrong value for word 0');
+    assert(memory.items[1] == 0x01FF, 'Wrong value for word 1');
+    assert(memory.bytes_len == 32, 'Wrong memory length');
+}
+
+#[test]
+#[available_gas(20000000)]
+fn test_store_byte_should_store_byte_at_offset_in_new_word() {
+    // Given
+    let mut memory = MemoryTrait::new();
+
+    // When
+    memory.store_byte(0x01, 32);
+
+    // Then
+    assert(memory.items[0] == 0x0, 'Wrong value for word 0');
+    assert(memory.items[1] == 0x0, 'Wrong value for word 1');
+    assert(memory.items[2] == 0x01000000000000000000000000000000, 'Wrong value for word 2');
+    assert(memory.bytes_len == 64, 'Wrong memory length');
+}
+
+#[test]
+#[available_gas(20000000)]
+fn test_store_byte_should_store_byte_at_offset_in_new_word_with_existing_value_in_previous_word() {
+    // Given
+    let mut memory = MemoryTrait::new();
+    memory.items.insert(0, 0x0100);
+    memory.items.insert(1, 0xffffffffffffffffffffffffffffffff);
+
+    // When
+    memory.store_byte(0xAB, 17);
+
+    // Then
+    assert(memory.items[0] == 0x0100, 'Wrong value in word 0');
+    assert(memory.items[1] == 0xffABffffffffffffffffffffffffffff, 'Wrong value in word 1');
+    assert(memory.bytes_len == 32, 'Wrong memory length');
+}
+