Refine looking up aot function with index (#3882)

* Refine looking up aot function with index

* refine the code

core/iwasm/aot/aot_runtime.c

@@ -1096,11 +1096,11 @@ aot_get_default_memory(AOTModuleInstance *module_inst)
 }
 
 AOTMemoryInstance *
-aot_get_memory_with_index(AOTModuleInstance *module_inst, uint32 index)
+aot_get_memory_with_idx(AOTModuleInstance *module_inst, uint32 mem_idx)
 {
-    if ((index >= module_inst->memory_count) || !module_inst->memories)
+    if ((mem_idx >= module_inst->memory_count) || !module_inst->memories)
         return NULL;
-    return module_inst->memories[index];
+    return module_inst->memories[mem_idx];
 }
 
 static bool
@@ -1282,20 +1282,77 @@ init_func_ptrs(AOTModuleInstance *module_inst, AOTModule *module,
     return true;
 }
 
+static int
+cmp_export_func_map(const void *a, const void *b)
+{
+    uint32 func_idx1 = ((const ExportFuncMap *)a)->func_idx;
+    uint32 func_idx2 = ((const ExportFuncMap *)b)->func_idx;
+    return func_idx1 < func_idx2 ? -1 : (func_idx1 > func_idx2 ? 1 : 0);
+}
+
+AOTFunctionInstance *
+aot_lookup_function_with_idx(AOTModuleInstance *module_inst, uint32 func_idx)
+{
+    AOTModuleInstanceExtra *extra = (AOTModuleInstanceExtra *)module_inst->e;
+    AOTFunctionInstance *export_funcs =
+        (AOTFunctionInstance *)module_inst->export_functions;
+    AOTFunctionInstance *func_inst = NULL;
+    ExportFuncMap *export_func_maps, *export_func_map, key;
+    uint64 size;
+    uint32 i;
+
+    if (module_inst->export_func_count == 0)
+        return NULL;
+
+    exception_lock(module_inst);
+
+    /* create the func_idx to export_idx maps if it hasn't been created */
+    if (!extra->export_func_maps) {
+        size = sizeof(ExportFuncMap) * (uint64)module_inst->export_func_count;
+        if (!(export_func_maps = extra->export_func_maps =
+                  runtime_malloc(size, NULL, 0))) {
+            /* allocate memory failed, lookup the export function one by one */
+            for (i = 0; i < module_inst->export_func_count; i++) {
+                if (export_funcs[i].func_index == func_idx) {
+                    func_inst = &export_funcs[i];
+                    break;
+                }
+            }
+            goto unlock_and_return;
+        }
+
+        for (i = 0; i < module_inst->export_func_count; i++) {
+            export_func_maps[i].func_idx = export_funcs[i].func_index;
+            export_func_maps[i].export_idx = i;
+        }
+
+        qsort(export_func_maps, module_inst->export_func_count,
+              sizeof(ExportFuncMap), cmp_export_func_map);
+    }
+
+    /* lookup the map to get the export_idx of the func_idx */
+    key.func_idx = func_idx;
+    export_func_map =
+        bsearch(&key, extra->export_func_maps, module_inst->export_func_count,
+                sizeof(ExportFuncMap), cmp_export_func_map);
+    if (export_func_map)
+        func_inst = &export_funcs[export_func_map->export_idx];
+
+unlock_and_return:
+    exception_unlock(module_inst);
+    return func_inst;
+}
+
 AOTFunctionInstance *
 aot_get_function_instance(AOTModuleInstance *module_inst, uint32 func_idx)
 {
     AOTModule *module = (AOTModule *)module_inst->module;
     AOTModuleInstanceExtra *extra = (AOTModuleInstanceExtra *)module_inst->e;
-    AOTFunctionInstance *export_funcs =
-        (AOTFunctionInstance *)module_inst->export_functions;
-    uint32 i;
+    AOTFunctionInstance *func_inst;
 
-    /* export functions are pre-instantiated */
-    for (i = 0; i < module_inst->export_func_count; i++) {
-        if (export_funcs[i].func_index == func_idx)
-            return &export_funcs[i];
-    }
+    /* lookup from export functions first */
+    if ((func_inst = aot_lookup_function_with_idx(module_inst, func_idx)))
+        return func_inst;
 
     exception_lock(module_inst);
 
@@ -2168,6 +2225,9 @@ aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst)
     if (module_inst->export_functions)
         wasm_runtime_free(module_inst->export_functions);
 
+    if (extra->export_func_maps)
+        wasm_runtime_free(extra->export_func_maps);
+
 #if WASM_ENABLE_MULTI_MEMORY != 0
     if (module_inst->export_memories)
         wasm_runtime_free(module_inst->export_memories);

core/iwasm/aot/aot_runtime.h

@@ -109,6 +109,13 @@ typedef struct AOTFunctionInstance {
     } u;
 } AOTFunctionInstance;
 
+/* Map of a function index to the element ith in
+   the export functions array */
+typedef struct ExportFuncMap {
+    uint32 func_idx;
+    uint32 export_idx;
+} ExportFuncMap;
+
 typedef struct AOTModuleInstanceExtra {
     DefPointer(const uint32 *, stack_sizes);
     /*
@@ -120,6 +127,13 @@ typedef struct AOTModuleInstanceExtra {
     MemBound shared_heap_start_off;
 
     WASMModuleInstanceExtraCommon common;
+
+    /**
+     * maps of func indexes to export func indexes, which
+     * is sorted by func index for a quick lookup and is
+     * created only when first time used.
+     */
+    ExportFuncMap *export_func_maps;
     AOTFunctionInstance **functions;
     uint32 function_count;
 #if WASM_ENABLE_MULTI_MODULE != 0
@@ -556,6 +570,13 @@ aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst);
 AOTFunctionInstance *
 aot_lookup_function(const AOTModuleInstance *module_inst, const char *name);
 
+/**
+ * Lookup an exported function in the AOT module instance with
+ * the function index.
+ */
+AOTFunctionInstance *
+aot_lookup_function_with_idx(AOTModuleInstance *module_inst, uint32 func_idx);
+
 AOTMemoryInstance *
 aot_lookup_memory(AOTModuleInstance *module_inst, char const *name);
 
@@ -563,7 +584,7 @@ AOTMemoryInstance *
 aot_get_default_memory(AOTModuleInstance *module_inst);
 
 AOTMemoryInstance *
-aot_get_memory_with_index(AOTModuleInstance *module_inst, uint32 index);
+aot_get_memory_with_idx(AOTModuleInstance *module_inst, uint32 mem_idx);
 
 /**
  * Get a function in the AOT module instance.
@@ -574,7 +595,7 @@ aot_get_memory_with_index(AOTModuleInstance *module_inst, uint32 index);
  * @return the function instance found
  */
 AOTFunctionInstance *
-aot_get_function_instance(AOTModuleInstance *module_inst, uint32_t func_idx);
+aot_get_function_instance(AOTModuleInstance *module_inst, uint32 func_idx);
 
 /**
  * Call the given AOT function of a AOT module instance with

core/iwasm/common/wasm_c_api.c

@@ -3383,21 +3383,8 @@ wasm_func_call(const wasm_func_t *func, const wasm_val_vec_t *params,
         if (!(func_comm_rt = func->func_comm_rt)) {
             AOTModuleInstance *inst_aot =
                 (AOTModuleInstance *)func->inst_comm_rt;
-            AOTModule *module_aot = (AOTModule *)inst_aot->module;
-            uint32 export_i = 0, export_func_j = 0;
-
-            for (; export_i < module_aot->export_count; ++export_i) {
-                AOTExport *export = module_aot->exports + export_i;
-                if (export->kind == EXPORT_KIND_FUNC) {
-                    if (export->index == func->func_idx_rt) {
-                        func_comm_rt =
-                            aot_lookup_function(inst_aot, export->name);
-                        ((wasm_func_t *)func)->func_comm_rt = func_comm_rt;
-                        break;
-                    }
-                    export_func_j++;
-                }
-            }
+            func_comm_rt = ((wasm_func_t *)func)->func_comm_rt =
+                aot_lookup_function_with_idx(inst_aot, func->func_idx_rt);
         }
 #endif
     }

core/iwasm/common/wasm_memory.c

@@ -1579,8 +1579,7 @@ wasm_runtime_get_memory(WASMModuleInstanceCommon *module_inst, uint32 index)
 
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT)
-        return aot_get_memory_with_index((AOTModuleInstance *)module_inst,
-                                         index);
+        return aot_get_memory_with_idx((AOTModuleInstance *)module_inst, index);
 #endif
 
     return NULL;