🐛 fix memory leaks and use testing allocator in all tests

keep-starknet-strange · Oct 20, 2023 · 88e2ba9 · 88e2ba9
1 parent 989f2c4
commit 88e2ba9
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 39 deletions.
diff --git a/README.md b/README.md
@@ -73,7 +73,6 @@ zig build test --summary all
 - [ ] Differential testing against Cairo VM in Rust.
 - [ ] Memory leaks detection (i.e use tools like [valgrind](https://valgrind.org/)).
 - [ ] Check [Zig style guide](https://ziglang.org/documentation/master/#Style-Guide) and apply it.
-- [ ] Go through the code and check carefully for memory safety issues, i.e make sure we have safe deallocation of memory everywhere.
 - [ ] Create documentation.
 
 ## 📄 License

diff --git a/src/utils/log.zig b/src/utils/log.zig
@@ -36,6 +36,7 @@ pub fn logFn(
     // Capture the current time in UTC format.
     const utc_format = "YYYY-MM-DDTHH:mm:ss";
     const time_str = DateTime.now().formatAlloc(allocator, utc_format) catch unreachable;
+    defer allocator.free(time_str); // Free the memory allocated by the allocator in `formatAlloc`.
 
     // Convert the log level and scope to string using @tagName.
     const level_str = @tagName(level);

diff --git a/src/utils/time.zig b/src/utils/time.zig
@@ -504,11 +504,11 @@ pub const Duration = struct {
 const expect = std.testing.expect;
 test "time format" {
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     const utc_format = "YYYY-MM-DDTHH:mm:ss";
     const dt = DateTime.initUnix(1697696484);
     const formatted_dt = try dt.formatAlloc(allocator, utc_format);
+    defer allocator.free(formatted_dt); // Free the memory allocated by the allocator.
     const expected_formatted_dt = "2023-10-19T06:21:24";
     try expect(std.mem.eql(u8, formatted_dt, expected_formatted_dt));
 }
diff --git a/src/vm/core.zig b/src/vm/core.zig
@@ -299,8 +299,7 @@ test "update pc regular no imm" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Regular;
     instruction.op_1_addr = instructions.Op1Src.AP;
@@ -327,8 +326,7 @@ test "update pc regular with imm" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Regular;
     instruction.op_1_addr = instructions.Op1Src.Imm;
@@ -355,8 +353,7 @@ test "update pc jump with operands res null" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Jump;
     var operands = OperandsResult.default();
@@ -377,8 +374,7 @@ test "update pc jump with operands res not relocatable" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Jump;
     var operands = OperandsResult.default();
@@ -399,8 +395,7 @@ test "update pc jump with operands res relocatable" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Jump;
     var operands = OperandsResult.default();
@@ -427,8 +422,7 @@ test "update pc jump rel with operands res null" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.JumpRel;
     var operands = OperandsResult.default();
@@ -449,8 +443,7 @@ test "update pc jump rel with operands res not felt" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.JumpRel;
     var operands = OperandsResult.default();
@@ -471,8 +464,7 @@ test "update pc jump rel with operands res felt" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.JumpRel;
     var operands = OperandsResult.default();
@@ -499,8 +491,7 @@ test "update pc update jnz with operands dst zero" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Jnz;
     var operands = OperandsResult.default();
@@ -527,8 +518,7 @@ test "update pc update jnz with operands dst not zero op1 not felt" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Jnz;
     var operands = OperandsResult.default();
@@ -550,8 +540,7 @@ test "update pc update jnz with operands dst not zero op1 felt" {
     // *                 SETUP TEST CONTEXT                       *
     // ************************************************************
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
     var instruction = instructions.Instruction.default();
     instruction.pc_update = instructions.PcUpdate.Jnz;
     var operands = OperandsResult.default();

diff --git a/src/vm/memory/memory.zig b/src/vm/memory/memory.zig
@@ -13,7 +13,8 @@ pub const Memory = struct {
     // ************************************************************
     // *                        FIELDS                            *
     // ************************************************************
-
+    /// The allocator used to allocate the memory.
+    allocator: *const Allocator,
     // The data in the memory.
     data: std.HashMap(relocatable.Relocatable, relocatable.MaybeRelocatable, std.hash_map.AutoContext(relocatable.Relocatable), std.hash_map.default_max_load_percentage),
     // The number of segments in the memory.
@@ -35,6 +36,7 @@ pub const Memory = struct {
         var memory = try allocator.create(Memory);
 
         memory.* = Memory{
+            .allocator = allocator,
             .data = std.AutoHashMap(relocatable.Relocatable, relocatable.MaybeRelocatable).init(allocator.*),
             .num_segments = 0,
             .validated_addresses = std.AutoHashMap(relocatable.Relocatable, bool).init(allocator.*),
@@ -47,6 +49,8 @@ pub const Memory = struct {
         // Clear the hash maps
         self.data.deinit();
         self.validated_addresses.deinit();
+        // Deallocate self.
+        self.allocator.destroy(self);
     }
 
     // ************************************************************
@@ -88,11 +92,11 @@ pub const Memory = struct {
 
 test "memory get without value raises error" {
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
 
     // Initialize a memory instance.
     var memory = try Memory.init(&allocator);
+    defer memory.deinit();
 
     // Get a value from the memory at an address that doesn't exist.
     _ = memory.get(relocatable.Relocatable.new(0, 0)) catch |err| {
@@ -104,12 +108,11 @@ test "memory get without value raises error" {
 
 test "memory set and get" {
     // Initialize an allocator.
-    // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
 
     // Initialize a memory instance.
     var memory = try Memory.init(&allocator);
+    defer memory.deinit();
 
     const address_1 = relocatable.Relocatable.new(0, 0);
     const value_1 = relocatable.fromFelt(starknet_felt.Felt252.one());

diff --git a/src/vm/memory/segments.zig b/src/vm/memory/segments.zig
@@ -14,7 +14,8 @@ pub const MemorySegmentManager = struct {
     // ************************************************************
     // *                        FIELDS                            *
     // ************************************************************
-
+    /// The allocator used to allocate the memory.
+    allocator: *const Allocator,
     // The size of the used segments.
     segment_used_sizes: std.HashMap(u32, u32, std.hash_map.AutoContext(u32), std.hash_map.default_max_load_percentage),
     // The size of the segments.
@@ -39,6 +40,7 @@ pub const MemorySegmentManager = struct {
         var segment_manager = try allocator.create(MemorySegmentManager);
         // Initialize the values of the MemorySegmentManager struct.
         segment_manager.* = MemorySegmentManager{
+            .allocator = allocator,
             .segment_used_sizes = std.AutoHashMap(u32, u32).init(allocator.*),
             .segment_sizes = std.AutoHashMap(u32, u32).init(allocator.*),
             // Initialize the memory pointer.
@@ -57,6 +59,8 @@ pub const MemorySegmentManager = struct {
         self.public_memory_offsets.deinit();
         // Deallocate the memory.
         self.memory.deinit();
+        // Deallocate self.
+        self.allocator.destroy(self);
     }
 
     // ************************************************************
@@ -84,13 +88,13 @@ pub const MemorySegmentManager = struct {
 
 test "memory segment manager" {
     // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
 
     // Initialize a memory segment manager.
     var memory_segment_manager = try MemorySegmentManager.init(&allocator);
+    defer memory_segment_manager.deinit();
 
-    // Allocate a memory segment.
+    //Allocate a memory segment.
     const relocatable_address_1 = memory_segment_manager.addSegment();
 
     // Check that the memory segment manager has one segment.

diff --git a/src/vm/run_context.zig b/src/vm/run_context.zig
@@ -71,9 +71,12 @@ pub const RunContext = struct {
 
     /// Safe deallocation of the memory.
     pub fn deinit(self: *RunContext) void {
+        // Deallocate fields.
         self.allocator.destroy(self.pc);
         self.allocator.destroy(self.ap);
         self.allocator.destroy(self.fp);
+        // Deallocate self.
+        self.allocator.destroy(self);
     }
 
     /// Compute dst address for a given instruction.
@@ -166,15 +169,13 @@ const expectEqual = std.testing.expectEqual;
 const expectError = std.testing.expectError;
 
 test "compute_op1_addr for fp op1 addr" {
-    // Initialize an allocator.
-    // TODO: Consider using a testing allocator.
-    var allocator = std.heap.page_allocator;
+    var allocator = std.testing.allocator;
 
     const instruction =
         Instruction{ .off_0 = 1, .off_1 = 2, .off_2 = 3, .dst_reg = .FP, .op_0_reg = .AP, .op_1_addr = .FP, .res_logic = .Add, .pc_update = .Regular, .ap_update = .Regular, .fp_update = .Regular, .opcode = .NOp };
 
     const run_context = try RunContext.init_with_values(&allocator, Relocatable.new(0, 4), Relocatable.new(0, 5), Relocatable.new(0, 6));
-
+    defer run_context.deinit();
     const relocatable_addr = try run_context.compute_op_1_addr(&instruction, null);
 
     try expect(relocatable_addr.eq(Relocatable.new(0, 9)));