Adding option to pass user data to allocator functions (#1765)

Add an option to pass user data to the allocator functions. It is common to
do this so that the host embedder can pass a struct as user data and access
that struct from the allocator, which gives the host embedder the ability to
do things such as track allocation statistics within the allocator.

Compile with `cmake -DWASM_MEM_ALLOC_WITH_USER_DATA=1` to enable
the option, and the allocator functions provided by the host embedder should
be like below (an extra argument `data` is added):
void *malloc(void *data, uint32 size) { .. }
void *realloc(void *data, uint32 size) { .. }
void free(void *data, void *ptr) { .. }

Signed-off-by: Andrew Chambers <ncham@amazon.com>

build-scripts/config_common.cmake

@@ -311,3 +311,6 @@ endif ()
 if (WAMR_BUILD_WASI_NN EQUAL 1)
   message ("     WASI-NN enabled")
 endif ()
+if (WAMR_BUILD_ALLOC_WITH_USER_DATA EQUAL 1)
+  add_definitions(-DWASM_MEM_ALLOC_WITH_USER_DATA=1)
+endif()

core/config.h

@@ -426,4 +426,8 @@
 #define WASM_ENABLE_SGX_IPFS 0
 #endif
 
+#ifndef WASM_MEM_ALLOC_WITH_USER_DATA
+#define WASM_MEM_ALLOC_WITH_USER_DATA 0
+#endif
+
 #endif /* end of _CONFIG_H_ */

core/iwasm/common/wasm_c_api.c

@@ -340,6 +340,10 @@ wasm_engine_new_internal(mem_alloc_type_t type, const MemAllocOption *opts)
             opts->allocator.free_func;
         init_args.mem_alloc_option.allocator.realloc_func =
             opts->allocator.realloc_func;
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+        init_args.mem_alloc_option.allocator.user_data =
+            opts->allocator.user_data;
+#endif
     }
     else {
         init_args.mem_alloc_option.pool.heap_buf = NULL;

core/iwasm/common/wasm_memory.c

@@ -11,16 +11,25 @@
 typedef enum Memory_Mode {
     MEMORY_MODE_UNKNOWN = 0,
     MEMORY_MODE_POOL,
-    MEMORY_MODE_ALLOCATOR
+    MEMORY_MODE_ALLOCATOR,
+    MEMORY_MODE_SYSTEM_ALLOCATOR
 } Memory_Mode;
 
 static Memory_Mode memory_mode = MEMORY_MODE_UNKNOWN;
 
 static mem_allocator_t pool_allocator = NULL;
 
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+static void *allocator_user_data = NULL;
+static void *(*malloc_func)(void *user_data, unsigned int size) = NULL;
+static void *(*realloc_func)(void *user_data, void *ptr,
+                             unsigned int size) = NULL;
+static void (*free_func)(void *user_data, void *ptr) = NULL;
+#else
 static void *(*malloc_func)(unsigned int size) = NULL;
 static void *(*realloc_func)(void *ptr, unsigned int size) = NULL;
 static void (*free_func)(void *ptr) = NULL;
+#endif
 
 static unsigned int global_pool_size;
 
@@ -39,6 +48,24 @@ wasm_memory_init_with_pool(void *mem, unsigned int bytes)
     return false;
 }
 
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+static bool
+wasm_memory_init_with_allocator(void *_user_data, void *_malloc_func,
+                                void *_realloc_func, void *_free_func)
+{
+    if (_malloc_func && _free_func && _malloc_func != _free_func) {
+        memory_mode = MEMORY_MODE_ALLOCATOR;
+        allocator_user_data = _user_data;
+        malloc_func = _malloc_func;
+        realloc_func = _realloc_func;
+        free_func = _free_func;
+        return true;
+    }
+    LOG_ERROR("Init memory with allocator (%p, %p, %p, %p) failed.\n",
+              _user_data, _malloc_func, _realloc_func, _free_func);
+    return false;
+}
+#else
 static bool
 wasm_memory_init_with_allocator(void *_malloc_func, void *_realloc_func,
                                 void *_free_func)
@@ -54,24 +81,38 @@ wasm_memory_init_with_allocator(void *_malloc_func, void *_realloc_func,
               _realloc_func, _free_func);
     return false;
 }
+#endif
 
 bool
 wasm_runtime_memory_init(mem_alloc_type_t mem_alloc_type,
                          const MemAllocOption *alloc_option)
 {
-    if (mem_alloc_type == Alloc_With_Pool)
+    if (mem_alloc_type == Alloc_With_Pool) {
         return wasm_memory_init_with_pool(alloc_option->pool.heap_buf,
                                           alloc_option->pool.heap_size);
-    else if (mem_alloc_type == Alloc_With_Allocator)
+    }
+    else if (mem_alloc_type == Alloc_With_Allocator) {
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+        return wasm_memory_init_with_allocator(
+            alloc_option->allocator.user_data,
+            alloc_option->allocator.malloc_func,
+            alloc_option->allocator.realloc_func,
+            alloc_option->allocator.free_func);
+#else
         return wasm_memory_init_with_allocator(
             alloc_option->allocator.malloc_func,
             alloc_option->allocator.realloc_func,
             alloc_option->allocator.free_func);
-    else if (mem_alloc_type == Alloc_With_System_Allocator)
-        return wasm_memory_init_with_allocator(os_malloc, os_realloc, os_free);
-    else
+#endif
+    }
+    else if (mem_alloc_type == Alloc_With_System_Allocator) {
+        memory_mode = MEMORY_MODE_SYSTEM_ALLOCATOR;
+        return true;
+    }
+    else {
         return false;
     }
+}
 
 void
 wasm_runtime_memory_destroy()
@@ -110,8 +151,15 @@ wasm_runtime_malloc_internal(unsigned int size)
     else if (memory_mode == MEMORY_MODE_POOL) {
         return mem_allocator_malloc(pool_allocator, size);
     }
-    else {
+    else if (memory_mode == MEMORY_MODE_ALLOCATOR) {
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+        return malloc_func(allocator_user_data, size);
+#else
         return malloc_func(size);
+#endif
+    }
+    else {
+        return os_malloc(size);
     }
 }
 
@@ -126,12 +174,19 @@ wasm_runtime_realloc_internal(void *ptr, unsigned int size)
     else if (memory_mode == MEMORY_MODE_POOL) {
         return mem_allocator_realloc(pool_allocator, ptr, size);
     }
-    else {
+    else if (memory_mode == MEMORY_MODE_ALLOCATOR) {
         if (realloc_func)
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+            return realloc_func(allocator_user_data, ptr, size);
+#else
             return realloc_func(ptr, size);
+#endif
         else
             return NULL;
     }
+    else {
+        return os_realloc(ptr, size);
+    }
 }
 
 static inline void
@@ -152,8 +207,15 @@ wasm_runtime_free_internal(void *ptr)
     else if (memory_mode == MEMORY_MODE_POOL) {
         mem_allocator_free(pool_allocator, ptr);
     }
-    else {
+    else if (memory_mode == MEMORY_MODE_ALLOCATOR) {
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+        free_func(allocator_user_data, ptr);
+#else
         free_func(ptr);
+#endif
+    }
+    else {
+        os_free(ptr);
     }
 }
 

core/iwasm/include/wasm_c_api.h

@@ -186,6 +186,9 @@ typedef union MemAllocOption {
         void *malloc_func;
         void *realloc_func;
         void *free_func;
+        /* allocator user data, only used when
+          WASM_MEM_ALLOC_WITH_USER_DATA is defined */
+        void *user_data;
     } allocator;
 } MemAllocOption;
 #endif

core/iwasm/include/wasm_export.h

@@ -117,6 +117,9 @@ typedef union MemAllocOption {
         void *malloc_func;
         void *realloc_func;
         void *free_func;
+        /* allocator user data, only used when
+            WASM_MEM_ALLOC_WITH_USER_DATA is defined */
+        void *user_data;
     } allocator;
 } MemAllocOption;
 #endif