1 files changed, 6 insertions, 5 deletions
diff --git a/src/matrix.zig b/src/matrix.zig
index eaf9685..e566bc4 100644
--- a/src/matrix.zig
+++ b/src/matrix.zig
@@ -1,5 +1,6 @@
 const std = @import("std");
 const ArrayList = std.ArrayList;
+const ArenaAllocator = std.heap.ArenaAllocator;
 
 const matrix = struct {
     h: usize,
@@ -7,11 +8,11 @@ const matrix = struct {
     x: ArrayList(f32),
 };
 
-pub fn new(h: usize, w: usize, x: []f32) matrix {
+pub fn new(arena: *ArenaAllocator, h: usize, w: usize, x: []f32) matrix {
     var m = matrix{
         .h = h,
         .w = w,
-        .x = ArrayList(f32).init(std.heap.page_allocator),
+        .x = ArrayList(f32).init(&arena.allocator),
     };
 
     var i: usize = 0;
@@ -37,8 +38,8 @@ pub fn dot(a: []f32, b: []f32, len: usize, step: usize) f32 {
     return x;
 }
 
-pub fn multiply(a: matrix, b: matrix) matrix {
-    var x = ArrayList(f32).init(std.heap.page_allocator);
+pub fn multiply(arena: *ArenaAllocator, a: matrix, b: matrix) matrix {
+    var x = ArrayList(f32).init(&arena.allocator);
     var i: usize = 0;
     var k: usize = 0;
 
@@ -50,5 +51,5 @@ pub fn multiply(a: matrix, b: matrix) matrix {
         k += a.w;
     }
 
-    return new(a.h, b.w, x.items);
+    return new(arena, a.h, b.w, x.items);
 }