Unload the firmware modules we swap in, when the game unloads the library

Tekken 6 never reaches gameplay with the real mpeg.prx: it plays its intro
movie, returns to the title screen, starts loading a demo match and loads
forever. With the sceMpeg HLE it plays fine.

The game unloads its video libraries before the level load and expects the
memory back. It gets most of it - scePsmf and scePsmfPlayer do go away - but
mpeg.prx stays resident, 33KB of it, sitting in the middle of the region the
loader then asks for:

  08c64000 - 09e24000  18.6MB  taken  UserSbrk
  09e24000 - 09ed4000   720KB  free
  09ed4000 - 09edc300  33.5KB  taken  ELF/sceMpeg_library
  09edc300 - 09f44000   425KB  free
  09f44000 - 09f4c000    32KB  taken  UtilityModule/302_av_atrac3plus

0x09f44000 - 0x09e24000 is 0x120000, which is exactly the allocation that
fails. Without mpeg.prx in the way that span is one free block and the level
loads.

sceUtility notifies the per-library hooks with state 1 when a utility module is
loaded and -1 when it is unloaded. The hooks that swap in a firmware module
only ever handled the load, so nothing ever took them back out. That affects
sceMpeg, sceMp3, sceMp4 and sceAtrac alike; Tekken is just the game whose
memory budget is tight enough to notice.

The unload has to take out what we put in and nothing else, so the loaded ids
are remembered rather than looked up by name: a game like Death Jr ships its
own mpeg.prx and loads it itself, and freeing that would be freeing the game's
memory. Verified that Death Jr still decodes its 1033 frames with our loader
never touching its module.

Savestates from before this have no record of what was swapped in, so they keep
the old behaviour of leaving the modules loaded rather than risk freeing
something the game owns.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

Core/HLE/sceKernelModule.cpp

@@ -2026,6 +2026,17 @@ static PSPModule *__KernelLoadELFFromPtr(const u8 *ptr, size_t elfSize, u32 load
 	return module;
 }
 
+bool KernelUnloadModuleByID(SceUID moduleId) {
+	u32 error;
+	PSPModule *module = kernelObjects.Get<PSPModule>(moduleId, error);
+	if (!module) {
+		return false;
+	}
+	module->Cleanup();
+	kernelObjects.Destroy<PSPModule>(moduleId);
+	return true;
+}
+
 bool KernelModuleIsLoaded(std::string_view name) {
 	u32 error;
 	for (SceUID moduleId : loadedModules) {

Core/HLE/sceKernelModule.h

@@ -256,6 +256,12 @@ SceUID KernelLoadModule(const std::string &filename, std::string *error_string,
 // Whether a real (non-HLE-stub) module calling itself this is loaded. Lets a caller tell whether a
 // library is already provided before bringing in another copy of it.
 bool KernelModuleIsLoaded(std::string_view name);
+
+// Take back out a module we brought in ourselves (the firmware swap in sceUtility). Frees the
+// memory block, which is the point: a game that unloads its video libraries before a level load
+// expects that space back, and one 33KB module left sitting in the middle of it is enough to
+// break a contiguous allocation. Returns false if the id is not a live module.
+bool KernelUnloadModuleByID(SceUID moduleId);
 int __KernelStartModule(SceUID moduleId, u32 argsize, u32 argAddr, u32 returnValueAddr, SceKernelSMOption *smoption, bool *needsWait);
 u32 __KernelStopUnloadSelfModuleWithOrWithoutStatus(u32 exitCode, u32 argSize, u32 argp, u32 statusAddr, u32 optionAddr, bool WithStatus);
 u32 sceKernelFindModuleByUID(u32 uid);

Core/HLE/sceUtility.cpp

@@ -95,12 +95,17 @@ static const int mp4ModuleDeps[] = {0x0300, 0};
 // So the module-list check below is a guard rather than the normal path, and it costs nothing:
 // asking the list rather than remembering what we loaded means this needs no state of its own. It
 // is right after a savestate load, across games, and if a game unloads a library and asks again.
+// The firmware modules we swapped in for a library whose HLE is disabled, keyed by the utility
+// module whose load brought them in. Remembered because the unload has to take out what we put
+// in and nothing else: a game that ships its own copy loads it itself, and we must not free that.
+static std::map<int, std::vector<SceUID>> swappedFirmwareModules;
+
 struct FirmwareModule {
 	const char *path;        // in the firmware
 	const char *moduleName;  // what the module calls itself once loaded
 };
 
-static void LoadFirmwareModules(const char *library, const FirmwareModule *modules, size_t count) {
+static void LoadFirmwareModules(int utilityModule, const char *library, const FirmwareModule *modules, size_t count) {
 	for (size_t i = 0; i < count; i++) {
 		if (KernelModuleIsLoaded(modules[i].moduleName)) {
 			DEBUG_LOG(Log::sceUtility, "%s is already loaded - not loading %s on top of it",
@@ -133,13 +138,33 @@ static void LoadFirmwareModules(const char *library, const FirmwareModule *modul
 			ERROR_LOG(Log::sceUtility, "Failed to start %s (%08x)", modules[i].path, result);
 			return;
 		}
+		swappedFirmwareModules[utilityModule].push_back(id);
 		INFO_LOG(Log::sceUtility, "Loaded the real %s", modules[i].path);
 	}
 }
 
+// The other half: give the memory back when the game says it is done with the library. Reverse
+// order, since a later module may import from an earlier one.
+static void UnloadFirmwareModules(int utilityModule) {
+	auto it = swappedFirmwareModules.find(utilityModule);
+	if (it == swappedFirmwareModules.end()) {
+		return;
+	}
+	for (auto id = it->second.rbegin(); id != it->second.rend(); ++id) {
+		if (KernelUnloadModuleByID(*id)) {
+			INFO_LOG(Log::sceUtility, "Unloaded the real module %d we had swapped in", *id);
+		}
+	}
+	swappedFirmwareModules.erase(it);
+}
+
 // mpeg.prx needs sceVideocodec, sceMpegbase and sceAudiocodec from us, all of which we implement,
 // so the module itself is the only thing that has to come from somewhere real.
 static void NotifyLoadStatusMpegBase(int state, u32 loadAddr, u32 totalSize) {
+	if (state == -1) {
+		UnloadFirmwareModules(0x303);
+		return;
+	}
 	// The effective flags, not the raw setting: those also account for the compat flags, for a
 	// firmware dump that isn't there, and for the boundary a savestate restored - resolving
 	// imports one way and loading modules the other is how a game ends up with neither.
@@ -149,18 +174,22 @@ static void NotifyLoadStatusMpegBase(int state, u32 loadAddr, u32 totalSize) {
 	static const FirmwareModule modules[] = {
 		{ "flash0:/kd/mpeg.prx", "sceMpeg_library" },
 	};
-	LoadFirmwareModules("sceMpeg", modules, ARRAY_SIZE(modules));
+	LoadFirmwareModules(0x303, "sceMpeg", modules, ARRAY_SIZE(modules));
 }
 
 // libmp3.prx imports nothing but the kernel and sceAudiocodec, which we have.
 static void NotifyLoadStatusMp3(int state, u32 loadAddr, u32 totalSize) {
+	if (state == -1) {
+		UnloadFirmwareModules(0x304);
+		return;
+	}
 	if (state != 1 || !(GetEffectiveDisableHLEFlags() & DisableHLEFlags::sceMp3)) {
 		return;
 	}
 	static const FirmwareModule modules[] = {
 		{ "flash0:/kd/libmp3.prx", "sceMp3_Library" },
 	};
-	LoadFirmwareModules("sceMp3", modules, ARRAY_SIZE(modules));
+	LoadFirmwareModules(0x304, "sceMp3", modules, ARRAY_SIZE(modules));
 }
 
 static void NotifyLoadStatusAvcodec(int state, u32 loadAddr, u32 totalSize) {
@@ -171,6 +200,10 @@ static void NotifyLoadStatusAvcodec(int state, u32 loadAddr, u32 totalSize) {
 // functions from sceAudiocodec (Init and Decode) plus ordinary kernel calls, and mp4msv.prx - the
 // 41 functions libmp4 leans on - imports nothing at all.
 static void NotifyLoadStatusMp4(int state, u32 loadAddr, u32 totalSize) {
+	if (state == -1) {
+		UnloadFirmwareModules(0x308);
+		return;
+	}
 	if (state != 1) {
 		return;
 	}
@@ -194,10 +227,14 @@ static void NotifyLoadStatusMp4(int state, u32 loadAddr, u32 totalSize) {
 		{ "flash0:/kd/mp4msv.prx", "mp4msv_module" },
 		{ "flash0:/kd/libmp4.prx", "sceMp4_library" },
 	};
-	LoadFirmwareModules("sceMp4", modules, ARRAY_SIZE(modules));
+	LoadFirmwareModules(0x308, "sceMp4", modules, ARRAY_SIZE(modules));
 }
 
 static void NotifyLoadStatusAtrac(int state, u32 loadAddr, u32 totalSize) {
+	if (state == -1) {
+		UnloadFirmwareModules(0x302);
+		return;
+	}
 	if (state == 1) {
 		// The effective flags, for the same reason the loads above use them.
 		if (GetEffectiveDisableHLEFlags() & DisableHLEFlags::sceAtrac) {
@@ -208,7 +245,7 @@ static void NotifyLoadStatusAtrac(int state, u32 loadAddr, u32 totalSize) {
 			static const FirmwareModule modules[] = {
 				{ "flash0:/kd/libatrac3plus.prx", "sceATRAC3plus_Library" },
 			};
-			LoadFirmwareModules("sceAtrac", modules, ARRAY_SIZE(modules));
+			LoadFirmwareModules(0x302, "sceAtrac", modules, ARRAY_SIZE(modules));
 			return;
 		}
 
@@ -449,13 +486,14 @@ void __UtilityInit() {
 	DeactivateDialog();
 	SavedataParam::Init();
 	currentlyLoadedModules.clear();
+	swappedFirmwareModules.clear();
 	volatileUnlockEvent = CoreTiming::RegisterEvent("UtilityVolatileUnlock", UtilityVolatileUnlock);
 
 	ResetSecondsSinceLastGameSave();
 }
 
 void __UtilityDoState(PointerWrap &p) {
-	auto s = p.Section("sceUtility", 1, 6);
+	auto s = p.Section("sceUtility", 1, 7);
 	if (!s) {
 		return;
 	}
@@ -479,6 +517,14 @@ void __UtilityDoState(PointerWrap &p) {
 		}
 	}
 
+	if (s >= 7) {
+		Do(p, swappedFirmwareModules);
+	} else if (p.mode == p.MODE_READ) {
+		// An older state has no record of what we swapped in, so the unload notification will
+		// leave those modules loaded rather than risk freeing something the game owns.
+		swappedFirmwareModules.clear();
+	}
+
 	if (s >= 3) {
 		Do(p, volatileUnlockEvent);
 	} else {