fix: changed type defs in relocatable and added test (#127)

* fix: changed type defs in relocatable and added test

* fix: added temp_data to memory to allow negative segments

* feat: added temp_data to memory_segment_manager

* fix: test

src/vm/memory/memory.zig

@@ -34,8 +34,17 @@ pub const Memory = struct {
         std.array_hash_map.AutoContext(Relocatable),
         true,
     ),
+    // The temporary data in the memory.
+    temp_data: std.ArrayHashMap(
+        Relocatable,
+        MaybeRelocatable,
+        std.array_hash_map.AutoContext(Relocatable),
+        true,
+    ),
     // The number of segments in the memory.
     num_segments: u32,
+    // The number of temporary segments in the memory.
+    num_temp_segments: u32,
     // Validated addresses are addresses that have been validated.
     // TODO: Consider merging this with `data` and benchmarking.
     validated_addresses: std.HashMap(
@@ -69,7 +78,9 @@ pub const Memory = struct {
                 Relocatable,
                 MaybeRelocatable,
             ).init(allocator),
+            .temp_data = std.AutoArrayHashMap(Relocatable, MaybeRelocatable).init(allocator),
             .num_segments = 0,
+            .num_temp_segments = 0,
             .validated_addresses = std.AutoHashMap(
                 Relocatable,
                 bool,
@@ -86,6 +97,7 @@ pub const Memory = struct {
     pub fn deinit(self: *Self) void {
         // Clear the hash maps
         self.data.deinit();
+        self.temp_data.deinit();
         self.validated_addresses.deinit();
         // Deallocate self.
         self.allocator.destroy(self);
@@ -108,19 +120,19 @@ pub const Memory = struct {
         MemoryOutOfBounds,
     }!void {
 
-        // Check if the address is valid.
+        // Insert the value into the memory.
         if (address.segment_index < 0) {
-            return CairoVMError.InvalidMemoryAddress;
+            self.temp_data.put(address, value) catch {
+                return CairoVMError.MemoryOutOfBounds;
+            };
+        } else {
+            self.data.put(
+                address,
+                value,
+            ) catch {
+                return CairoVMError.MemoryOutOfBounds;
+            };
         }
-
-        // Insert the value into the memory.
-        self.data.put(
-            address,
-            value,
-        ) catch {
-            return CairoVMError.MemoryOutOfBounds;
-        };
-
         // TODO: Add all relevant checks.
     }
 
@@ -133,6 +145,9 @@ pub const Memory = struct {
         self: *Self,
         address: Relocatable,
     ) error{MemoryOutOfBounds}!MaybeRelocatable {
+        if (address.segment_index < 0) {
+            return self.temp_data.get(address) orelse CairoVMError.MemoryOutOfBounds;
+        }
         return self.data.get(address) orelse CairoVMError.MemoryOutOfBounds;
     }
 
@@ -190,20 +205,32 @@ test "memory set and get" {
     );
     const value_1 = relocatable.fromFelt(starknet_felt.Felt252.one());
 
+    const address_2 = Relocatable.new(
+        -1,
+        0,
+    );
+    const value_2 = relocatable.fromFelt(starknet_felt.Felt252.one());
+
     // Set a value into the memory.
     _ = try memory.set(
         address_1,
         value_1,
     );
+    _ = try memory.set(
+        address_2,
+        value_2,
+    );
 
     // Get the value from the memory.
     const maybe_value_1 = try memory.get(address_1);
+    const maybe_value_2 = try memory.get(address_2);
 
     // ************************************************************
     // *                      TEST CHECKS                         *
     // ************************************************************
     // Assert that the value is the expected value.
     try expect(maybe_value_1.eq(value_1));
+    try expect(maybe_value_2.eq(value_2));
 }
 
 test "validate existing memory for range check within bound" {

src/vm/memory/relocatable.zig

@@ -10,7 +10,7 @@ pub const Relocatable = struct {
     const Self = @This();
 
     // The index of the memory segment.
-    segment_index: u64 = 0,
+    segment_index: i64 = 0,
     // The offset in the memory segment.
     offset: u64 = 0,
 
@@ -21,7 +21,7 @@ pub const Relocatable = struct {
     // # Returns
     // A new Relocatable.
     pub fn new(
-        segment_index: u64,
+        segment_index: i64,
         offset: u64,
     ) Self {
         return .{
@@ -291,13 +291,14 @@ const expectEqual = std.testing.expectEqual;
 const expectError = std.testing.expectError;
 
 test "Relocatable: eq should return true if two Relocatable are the same." {
+    try expect(Relocatable.new(-1, 4).eq(Relocatable.new(-1, 4)));
     try expect(Relocatable.new(2, 4).eq(Relocatable.new(2, 4)));
 }
 
 test "Relocatable: eq should return false if two Relocatable are not the same." {
     const relocatable1 = Relocatable.new(2, 4);
     const relocatable2 = Relocatable.new(2, 5);
-    const relocatable3 = Relocatable.new(1, 4);
+    const relocatable3 = Relocatable.new(-1, 4);
     try expect(!relocatable1.eq(relocatable2));
     try expect(!relocatable1.eq(relocatable3));
 }

src/vm/memory/segments.zig

@@ -115,6 +115,20 @@ pub const MemorySegmentManager = struct {
         return relocatable_address;
     }
 
+    // Adds a temporary memory segment and returns the first address of the new segment.
+    pub fn addTempSegment(self: *Self) Relocatable {
+        // Increment the number of temporary segments.
+        self.memory.num_temp_segments += 1;
+
+        // Create the relocatable address for the new segment.
+        const relocatable_address = Relocatable{
+            .segment_index = -@as(i64, @intCast(self.memory.num_temp_segments)),
+            .offset = 0,
+        };
+
+        return relocatable_address;
+    }
+
     /// Retrieves the size of a memory segment by its index if available, else returns null.
     ///
     /// # Parameters
@@ -135,6 +149,15 @@ pub const MemorySegmentManager = struct {
         return self.memory.data.count();
     }
 
+    /// Retrieves the number of temporary memory segments.
+    ///
+    /// # Returns
+    ///
+    /// The number of temporary memory segments as a `usize`.
+    pub fn numTempSegments(self: *Self) usize {
+        return self.memory.temp_data.count();
+    }
+
     /// Computes and returns the effective size of memory segments.
     ///
     /// This function iterates through memory segments, calculates their effective sizes, and
@@ -213,6 +236,11 @@ test "memory segment manager" {
     // Check that the memory segment manager has one segment.
     try expect(memory_segment_manager.memory.num_segments == 1);
 
+    //Allocate a temporary memory segment.
+    const relocatable_address_2 = memory_segment_manager.addTempSegment();
+
+    try expect(memory_segment_manager.memory.num_temp_segments == 1);
+
     // Check if the relocatable address is correct.
     try expectEqual(
         Relocatable{
@@ -222,19 +250,39 @@ test "memory segment manager" {
         relocatable_address_1,
     );
 
+    try expectEqual(
+        Relocatable{
+            .segment_index = -1,
+            .offset = 0,
+        },
+        relocatable_address_2,
+    );
+
     // Allocate another memory segment.
-    const relocatable_address_2 = memory_segment_manager.addSegment();
+    const relocatable_address_3 = memory_segment_manager.addSegment();
+
+    // Allocate another temporary memory segment.
+    const relocatable_address_4 = memory_segment_manager.addTempSegment();
 
     // Check that the memory segment manager has two segments.
     try expect(memory_segment_manager.memory.num_segments == 2);
+    // Check that the memory segment manager has two temporary segments.
+    try expect(memory_segment_manager.memory.num_temp_segments == 2);
 
     // Check if the relocatable address is correct.
     try expectEqual(
         Relocatable{
             .segment_index = 1,
             .offset = 0,
         },
-        relocatable_address_2,
+        relocatable_address_3,
+    );
+    try expectEqual(
+        Relocatable{
+            .segment_index = -2,
+            .offset = 0,
+        },
+        relocatable_address_4,
     );
 }
 
@@ -254,27 +302,39 @@ test "set get integer value in segment memory" {
     // ************************************************************
     _ = memory_segment_manager.addSegment();
     _ = memory_segment_manager.addSegment();
+    _ = memory_segment_manager.addTempSegment();
+    _ = memory_segment_manager.addTempSegment();
 
-    const address = Relocatable.new(
+    const address_1 = Relocatable.new(
         0,
         0,
     );
-    const value = relocatable.fromFelt(Felt252.fromInteger(42));
+    const address_2 = Relocatable.new(
+        -1,
+        0,
+    );
+    const value_1 = relocatable.fromFelt(Felt252.fromInteger(42));
+
+    const value_2 = relocatable.fromFelt(Felt252.fromInteger(84));
 
     const wrong_address = Relocatable.new(0, 1);
 
-    _ = try memory_segment_manager.memory.set(address, value);
+    _ = try memory_segment_manager.memory.set(address_1, value_1);
+    _ = try memory_segment_manager.memory.set(address_2, value_2);
 
-    try expect(memory_segment_manager.memory.data.contains(address));
+    try expect(memory_segment_manager.memory.data.contains(address_1));
     try expect(!memory_segment_manager.memory.data.contains(wrong_address));
 
     // ************************************************************
     // *                      TEST CHECKS                         *
     // ************************************************************
-    const actual_value = try memory_segment_manager.memory.get(address);
-    const expected_value = value;
+    const actual_value_1 = try memory_segment_manager.memory.get(address_1);
+    const expected_value_1 = value_1;
+    const actual_value_2 = try memory_segment_manager.memory.get(address_2);
+    const expected_value_2 = value_2;
 
-    try expect(expected_value.eq(actual_value));
+    try expect(expected_value_1.eq(actual_value_1));
+    try expect(expected_value_2.eq(actual_value_2));
 }
 
 test "MemorySegmentManager: getSegmentUsedSize should return the size of a memory segment by its index if available" {
@@ -301,14 +361,23 @@ test "MemorySegmentManager: numSegments should return the number of segments in
     defer memory_segment_manager.deinit();
     try memory_segment_manager.memory.data.put(Relocatable.new(0, 1), .{ .felt = Felt252.fromInteger(10) });
     try memory_segment_manager.memory.data.put(Relocatable.new(1, 1), .{ .felt = Felt252.fromInteger(10) });
-    try memory_segment_manager.memory.data.put(Relocatable.new(2, 1), .{ .felt = Felt252.fromInteger(10) });
-    try memory_segment_manager.memory.data.put(Relocatable.new(3, 1), .{ .felt = Felt252.fromInteger(10) });
     try expectEqual(
-        @as(usize, 4),
+        @as(usize, 2),
         memory_segment_manager.numSegments(),
     );
 }
 
+test "MemorySegmentManager: numSegments should return the number of segments in the temporary memory" {
+    var memory_segment_manager = try MemorySegmentManager.init(std.testing.allocator);
+    defer memory_segment_manager.deinit();
+    try memory_segment_manager.memory.temp_data.put(Relocatable.new(-1, 1), .{ .felt = Felt252.fromInteger(10) });
+    try memory_segment_manager.memory.temp_data.put(Relocatable.new(-2, 1), .{ .felt = Felt252.fromInteger(10) });
+    try expectEqual(
+        @as(usize, 2),
+        memory_segment_manager.numTempSegments(),
+    );
+}
+
 test "MemorySegmentManager: computeEffectiveSize for one segment memory" {
     var memory_segment_manager = try MemorySegmentManager.init(std.testing.allocator);
     defer memory_segment_manager.deinit();