From d64a3ab6ecfbf83e69a2d31cf0b3fb4f6a6e3368 Mon Sep 17 00:00:00 2001
From: Benbuck Nason <bnason@netflix.com>
Date: Fri, 13 Sep 2024 19:53:01 -0700
Subject: [PATCH] Fix aot multi export memory support (#3791)

---
 core/iwasm/aot/aot_runtime.c | 48 ++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/core/iwasm/aot/aot_runtime.c b/core/iwasm/aot/aot_runtime.c
index dc6cba3f..13664ca0 100644
--- a/core/iwasm/aot/aot_runtime.c
+++ b/core/iwasm/aot/aot_runtime.c
@@ -1413,6 +1413,36 @@ create_export_funcs(AOTModuleInstance *module_inst, AOTModule *module,
     return true;
 }
 
+#if WASM_ENABLE_MULTI_MEMORY != 0
+static WASMExportMemInstance *
+export_memories_instantiate(const AOTModule *module,
+                            AOTModuleInstance *module_inst,
+                            uint32 export_mem_count, char *error_buf,
+                            uint32 error_buf_size)
+{
+    WASMExportMemInstance *export_memories, *export_memory;
+    AOTExport *export = module->exports;
+    uint32 i;
+    uint64 total_size =
+        sizeof(WASMExportMemInstance) * (uint64)export_mem_count;
+
+    if (!(export_memory = export_memories =
+              runtime_malloc(total_size, error_buf, error_buf_size))) {
+        return NULL;
+    }
+
+    for (i = 0; i < module->export_count; i++, export ++)
+        if (export->kind == EXPORT_KIND_MEMORY) {
+            export_memory->name = export->name;
+            export_memory->memory = module_inst->memories[export->index];
+            export_memory++;
+        }
+
+    bh_assert((uint32)(export_memory - export_memories) == export_mem_count);
+    return export_memories;
+}
+#endif /* end of if WASM_ENABLE_MULTI_MEMORY != 0 */
+
 static bool
 create_exports(AOTModuleInstance *module_inst, AOTModule *module,
                char *error_buf, uint32 error_buf_size)
@@ -1439,6 +1469,19 @@ create_exports(AOTModuleInstance *module_inst, AOTModule *module,
         }
     }
 
+#if WASM_ENABLE_MULTI_MEMORY == 0
+    bh_assert(module_inst->export_memory_count <= 1);
+#else
+    if (module_inst->export_memory_count) {
+        module_inst->export_memories = export_memories_instantiate(
+            module, module_inst, module_inst->export_memory_count, error_buf,
+            error_buf_size);
+        if (!module_inst->export_memories) {
+            return false;
+        }
+    }
+#endif
+
     return create_export_funcs(module_inst, module, error_buf, error_buf_size);
 }
 
@@ -2082,6 +2125,11 @@ aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst)
     if (module_inst->export_functions)
         wasm_runtime_free(module_inst->export_functions);
 
+#if WASM_ENABLE_MULTI_MEMORY != 0
+    if (module_inst->export_memories)
+        wasm_runtime_free(module_inst->export_memories);
+#endif
+
     if (extra->functions) {
         uint32 func_idx;
         for (func_idx = 0; func_idx < extra->function_count; ++func_idx) {