From a06a3de65b1c9da1e3aaf873ee19aa4333d2692f Mon Sep 17 00:00:00 2001
From: Sam Williams <sam@arweave.org>
Date: Wed, 12 Feb 2025 18:36:50 -0500
Subject: [PATCH] chore: merge Oleg's mem fixes

---
 native/hb_beamr/hb_beamr.c | 3 ++-
 native/hb_beamr/hb_wasm.c  | 9 +++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/native/hb_beamr/hb_beamr.c b/native/hb_beamr/hb_beamr.c
index 78c25d60..15530131 100644
--- a/native/hb_beamr/hb_beamr.c
+++ b/native/hb_beamr/hb_beamr.c
@@ -110,7 +110,8 @@ static void wasm_driver_output(ErlDrvData raw, char *buff, ErlDrvSizeT bufflen)
         long size_l = (long)size;
         ei_decode_binary(buff, &index, wasm_binary, &size_l);
         ei_get_type(buff, &index, &type, &mode_size);
-        mode = driver_alloc(mode_size);
+        // the init message size + '\0' character
+        mode = driver_alloc(mode_size + 1);
         ei_decode_atom(buff, &index, mode);
         LoadWasmReq* mod_bin = driver_alloc(sizeof(LoadWasmReq));
         mod_bin->proc = proc;
diff --git a/native/hb_beamr/hb_wasm.c b/native/hb_beamr/hb_wasm.c
index 348777ca..41463ef8 100644
--- a/native/hb_beamr/hb_wasm.c
+++ b/native/hb_beamr/hb_wasm.c
@@ -158,8 +158,11 @@ void wasm_initialize_runtime(void* raw) {
     wasm_module_exports(proc->module, &exports);
 
     // Create Erlang lists for imports
-    ErlDrvTermData* init_msg = driver_alloc(sizeof(ErlDrvTermData) * (2 + (13 * imports.size) + (11 * exports.size)));
+    int init_msg_size = sizeof(ErlDrvTermData) * (2 + 3 + 5 + (13 * imports.size) + (11 * exports.size));
+    ErlDrvTermData* init_msg = driver_alloc(init_msg_size);
     int msg_i = 0;
+
+    // 2 in the init_msg_size
     init_msg[msg_i++] = ERL_DRV_ATOM;
     init_msg[msg_i++] = atom_execution_result;
 
@@ -179,7 +182,7 @@ void wasm_initialize_runtime(void* raw) {
             // TODO: Handle other types of imports?
             continue;
         }
-
+        // 13 items in the each import message
         init_msg[msg_i++] = ERL_DRV_ATOM;
         init_msg[msg_i++] = driver_mk_atom((char*)wasm_externtype_to_kind_string(type));
         init_msg[msg_i++] = ERL_DRV_STRING;
@@ -255,6 +258,7 @@ void wasm_initialize_runtime(void* raw) {
         get_function_sig(type, type_str);
         DRV_DEBUG("Export: %s [%s] -> %s", name->data, kind_str, type_str);
 
+        // 10 elements for each exported function
         init_msg[msg_i++] = ERL_DRV_ATOM;
         init_msg[msg_i++] = driver_mk_atom(kind_str);
         init_msg[msg_i++] = ERL_DRV_STRING;
@@ -267,6 +271,7 @@ void wasm_initialize_runtime(void* raw) {
         init_msg[msg_i++] = 3;
     }
 
+    // 5 closing elements
     init_msg[msg_i++] = ERL_DRV_NIL;
     init_msg[msg_i++] = ERL_DRV_LIST;
     init_msg[msg_i++] = (exports.size) + 1;