Cache unit import metadata

src/wasm/worker.cpp

@@ -71,6 +71,24 @@ struct WasmAbiInfo
 	bool found = false;
 };
 
+enum class WasmUnitImportKind
+{
+	Memory,
+	Table,
+	StackPointer,
+	MemoryBase,
+	TableBase,
+	Function,
+	GotMemory,
+	GotFunction
+};
+
+struct WasmUnitImport
+{
+	WasmUnitImportKind kind;
+	String name;
+};
+
 struct WasmUnitModule
 {
 	String source_path;   // absolute .uce path
@@ -81,6 +99,8 @@ struct WasmUnitModule
 	WasmDylinkInfo dylink;
 	WasmAbiInfo abi;
 	std::optional<wasmtime::Module> module;
+	std::vector<WasmUnitImport> imports;
+	size_t got_memory_imports = 0;
 };
 
 struct WasmWorkerConfig
@@ -778,6 +798,39 @@ public:
 			return(nullptr);
 		}
 		unit->module.emplace(std::move(*compiled_or_cached));
+		for(auto import_type : unit->module->imports())
+		{
+			String mod_name(import_type.module());
+			String name(import_type.name());
+			auto extern_type = wasmtime::ExternType::from_import(import_type);
+			bool is_func_import = std::get_if<wasmtime::FuncType::Ref>(&extern_type) != 0;
+			WasmUnitImportKind kind;
+			if(mod_name == "env" && name == "memory")
+				kind = WasmUnitImportKind::Memory;
+			else if(mod_name == "env" && name == "__indirect_function_table")
+				kind = WasmUnitImportKind::Table;
+			else if(mod_name == "env" && name == "__stack_pointer")
+				kind = WasmUnitImportKind::StackPointer;
+			else if(mod_name == "env" && name == "__memory_base")
+				kind = WasmUnitImportKind::MemoryBase;
+			else if(mod_name == "env" && name == "__table_base")
+				kind = WasmUnitImportKind::TableBase;
+			else if(mod_name == "env" && is_func_import)
+				kind = WasmUnitImportKind::Function;
+			else if(mod_name == "GOT.mem")
+			{
+				kind = WasmUnitImportKind::GotMemory;
+				unit->got_memory_imports++;
+			}
+			else if(mod_name == "GOT.func")
+				kind = WasmUnitImportKind::GotFunction;
+			else
+			{
+				error = source_path + ": import policy violation: " + mod_name + "." + name;
+				return(nullptr);
+			}
+			unit->imports.push_back({ kind, std::move(name) });
+		}
 		module_cache[wasm_path] = unit;
 		return(unit);
 	}
@@ -1321,25 +1374,25 @@ private:
 		// import resolution
 		auto import_start = std::chrono::steady_clock::now();
 		std::vector<wasmtime::Extern> imports;
+		imports.reserve(mod->imports.size());
 		std::vector<std::pair<String, wasmtime::Global>> deferred_got;
-		for(auto import_type : module.imports())
+		deferred_got.reserve(mod->got_memory_imports);
+		wasmtime::GlobalType base_global_type(wasmtime::ValType::i32(), false);
+		wasmtime::GlobalType got_global_type(wasmtime::ValType::i32(), true);
+		for(const auto& import : mod->imports)
 		{
-			String mod_name(import_type.module());
+			const String& name = import.name;
-			String name(import_type.name());
+			if(import.kind == WasmUnitImportKind::Memory)
-			auto extern_type = wasmtime::ExternType::from_import(import_type);
-			bool is_func_import = std::get_if<wasmtime::FuncType::Ref>(&extern_type) != 0;
-
-			if(mod_name == "env" && name == "memory")
 			{
 				imports.push_back(*memory);
 				continue;
 			}
-			if(mod_name == "env" && name == "__indirect_function_table")
+			if(import.kind == WasmUnitImportKind::Table)
 			{
 				imports.push_back(*table);
 				continue;
 			}
-			if(mod_name == "env" && name == "__stack_pointer")
+			if(import.kind == WasmUnitImportKind::StackPointer)
 			{
 				auto sp = core->get(cx, "__stack_pointer");
 				if(!sp || !std::get_if<wasmtime::Global>(&*sp))
@@ -1347,17 +1400,16 @@ private:
 				imports.push_back(std::get<wasmtime::Global>(*sp));
 				continue;
 			}
-			if(mod_name == "env" && (name == "__memory_base" || name == "__table_base"))
+			if(import.kind == WasmUnitImportKind::MemoryBase || import.kind == WasmUnitImportKind::TableBase)
 			{
-				int32_t value = name == "__memory_base" ? (int32_t)memory_base : (int32_t)table_base;
+				int32_t value = import.kind == WasmUnitImportKind::MemoryBase ? (int32_t)memory_base : (int32_t)table_base;
-				wasmtime::GlobalType global_type(wasmtime::ValType::i32(), false);
+				auto global = wasmtime::Global::create(cx, base_global_type, wasmtime::Val(value));
-				auto global = wasmtime::Global::create(cx, global_type, wasmtime::Val(value));
 				if(!global)
 					return("global create failed: " + String(global.err().message()));
 				imports.push_back(global.ok());
 				continue;
 			}
-			if(mod_name == "env" && is_func_import)
+			if(import.kind == WasmUnitImportKind::Function)
 			{
 				auto resolve_start = std::chrono::steady_clock::now();
 				auto func = resolve_func(name);
@@ -1369,9 +1421,8 @@ private:
 				imports.push_back(*func);
 				continue;
 			}
-			if(mod_name == "GOT.mem")
+			if(import.kind == WasmUnitImportKind::GotMemory)
 			{
-				wasmtime::GlobalType global_type(wasmtime::ValType::i32(), true);
 				u32 address = 0;
 				auto resolve_start = std::chrono::steady_clock::now();
 				bool resolved = resolve_data(name, address);
@@ -1380,7 +1431,7 @@ private:
 					std::chrono::steady_clock::now() - resolve_start).count();
 				if(resolved)
 				{
-					auto global = wasmtime::Global::create(cx, global_type, wasmtime::Val((int32_t)address));
+					auto global = wasmtime::Global::create(cx, got_global_type, wasmtime::Val((int32_t)address));
 					if(!global)
 						return("global create failed: " + String(global.err().message()));
 					imports.push_back(global.ok());
@@ -1389,7 +1440,7 @@ private:
 				{
 					// provisional; self-resolved from the unit's own export
 					// (plus its memory base) after instantiation
-					auto global = wasmtime::Global::create(cx, global_type, wasmtime::Val((int32_t)0));
+					auto global = wasmtime::Global::create(cx, got_global_type, wasmtime::Val((int32_t)0));
 					if(!global)
 						return("global create failed: " + String(global.err().message()));
 					deferred_got.push_back({ name, global.ok() });
@@ -1397,7 +1448,7 @@ private:
 				}
 				continue;
 			}
-			if(mod_name == "GOT.func")
+			if(import.kind == WasmUnitImportKind::GotFunction)
 			{
 				auto resolve_start = std::chrono::steady_clock::now();
 				auto func = resolve_func(name);
@@ -1410,14 +1461,12 @@ private:
 				error = place_funcref(*func, slot);
 				if(error != "")
 					return(error);
-				wasmtime::GlobalType global_type(wasmtime::ValType::i32(), true);
+				auto global = wasmtime::Global::create(cx, got_global_type, wasmtime::Val((int32_t)slot));
-				auto global = wasmtime::Global::create(cx, global_type, wasmtime::Val((int32_t)slot));
 				if(!global)
 					return("global create failed: " + String(global.err().message()));
 				imports.push_back(global.ok());
 				continue;
 			}
-			return(source_path + ": import policy violation: " + mod_name + "." + name);
 		}
 		unit_import_total_us += (u64)std::chrono::duration_cast<std::chrono::microseconds>(
 			std::chrono::steady_clock::now() - import_start).count();