From 524da5dc54428ef892f1df7b8215b9110f339fef Mon Sep 17 00:00:00 2001
From: Eladash <elad3356p@gmail.com>
Date: Thu, 12 May 2022 19:59:33 +0300
Subject: [PATCH] LV2: Memory state post-exitspawn fixes

* Fix memory capacity if SDK version of the following executable differs from the original process'.
* Keep user memory containers, they are not freed at exitspawn!

Hw test https://github.com/elad335/myps3tests/commit/4bf60023ee3779a221da968732fc622b85266c5b
---
 rpcs3/Emu/Cell/PPUModule.cpp       | 12 +++++++-
 rpcs3/Emu/Cell/lv2/sys_process.cpp | 49 +++++++++++++++++++++++-------
 rpcs3/Emu/System.cpp               |  2 ++
 rpcs3/Emu/System.h                 |  1 +
 4 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/rpcs3/Emu/Cell/PPUModule.cpp b/rpcs3/Emu/Cell/PPUModule.cpp
index 6778a44967..fbb5f71a3a 100644
--- a/rpcs3/Emu/Cell/PPUModule.cpp
+++ b/rpcs3/Emu/Cell/PPUModule.cpp
@@ -1840,7 +1840,17 @@ bool ppu_load_exec(const ppu_exec_object& elf)
 		mem_size += 0xC000000;
 	}
 
-	g_fxo->init<lv2_memory_container>(mem_size)->used += primary_stacksize;
+	if (Emu.init_mem_containers)
+	{
+		// Refer to sys_process_exit2 for explanation
+		Emu.init_mem_containers(mem_size);
+	}
+	else
+	{
+		g_fxo->init<lv2_memory_container>(mem_size);
+	}
+
+	g_fxo->get<lv2_memory_container>().used += primary_stacksize;
 
 	ppu->cmd_push({ppu_cmd::initialize, 0});
 
diff --git a/rpcs3/Emu/Cell/lv2/sys_process.cpp b/rpcs3/Emu/Cell/lv2/sys_process.cpp
index 321df513aa..7c041e7693 100644
--- a/rpcs3/Emu/Cell/lv2/sys_process.cpp
+++ b/rpcs3/Emu/Cell/lv2/sys_process.cpp
@@ -402,25 +402,52 @@ void _sys_process_exit2(ppu_thread& ppu, s32 status, vm::ptr<sys_exit2_param> ar
 
 	// TODO: set prio, flags
 
-	std::string path = vfs::get(argv[0]);
-	std::string hdd1 = vfs::get("/dev_hdd1/");
-	std::string disc;
-
-	if (Emu.GetCat() == "DG" || Emu.GetCat() == "GD")
-		disc = vfs::get("/dev_bdvd/");
-	if (disc.empty() && !Emu.GetTitleID().empty())
-		disc = vfs::get(Emu.GetDir());
-
-	Emu.CallFromMainThread([path = std::move(path), argv = std::move(argv), envp = std::move(envp), data = std::move(data), disc = std::move(disc)
-		, hdd1 = std::move(hdd1), klic = g_fxo->get<loaded_npdrm_keys>().last_key(), old_config = Emu.GetUsedConfig()]() mutable
+	Emu.CallFromMainThread([argv = std::move(argv), envp = std::move(envp), data = std::move(data)]() mutable
 	{
 		sys_process.success("Process finished -> %s", argv[0]);
+
+		std::string disc;
+
+		if (Emu.GetCat() == "DG" || Emu.GetCat() == "GD")
+			disc = vfs::get("/dev_bdvd/");
+		if (disc.empty() && !Emu.GetTitleID().empty())
+			disc = vfs::get(Emu.GetDir());
+
+		std::string path = vfs::get(argv[0]);
+		std::string hdd1 = vfs::get("/dev_hdd1/");
+		std::string old_config = Emu.GetUsedConfig();
+
+		const u128 klic = g_fxo->get<loaded_npdrm_keys>().last_key();
+
+		using namespace id_manager;
+
+		auto func = [old_size = g_fxo->get<lv2_memory_container>().size, vec = (reader_lock{g_mutex}, g_fxo->get<id_map<lv2_memory_container>>().vec)](u32 sdk_suggested_mem) mutable
+		{
+			// Save LV2 memory containers
+			g_fxo->init<id_map<lv2_memory_container>>()->vec = std::move(vec);
+
+			// Empty the containers, accumulate their total size
+			u32 total_size = 0;
+			idm::select<lv2_memory_container>([&](u32, lv2_memory_container& ctr)
+			{
+				ctr.used = 0;
+				total_size += ctr.size;
+			});
+
+			// The default memory container capacity can only decrease after exitspawn
+			// 1. If newer SDK version suggests higher memory capacity - it is ignored
+			// 2. If newer SDK version suggests lower memory capacity - it is lowered
+			// And if 2. happens while user memory containers exist, the left space can be spent on user memory containers
+			g_fxo->init<lv2_memory_container>(std::min(old_size - total_size, sdk_suggested_mem) + total_size);
+		};
+
 		Emu.Kill(false);
 		Emu.argv = std::move(argv);
 		Emu.envp = std::move(envp);
 		Emu.data = std::move(data);
 		Emu.disc = std::move(disc);
 		Emu.hdd1 = std::move(hdd1);
+		Emu.init_mem_containers = std::move(func);
 
 		if (klic)
 		{
diff --git a/rpcs3/Emu/System.cpp b/rpcs3/Emu/System.cpp
index 3867f9a2bd..ded8379726 100644
--- a/rpcs3/Emu/System.cpp
+++ b/rpcs3/Emu/System.cpp
@@ -1806,6 +1806,7 @@ void Emulator::Kill(bool allow_autoexit)
 		disc.clear();
 		klic.clear();
 		hdd1.clear();
+		init_mem_containers = nullptr;
 		m_config_path.clear();
 		m_config_mode = cfg_mode::custom;
 		return;
@@ -1951,6 +1952,7 @@ void Emulator::Kill(bool allow_autoexit)
 	disc.clear();
 	klic.clear();
 	hdd1.clear();
+	init_mem_containers = nullptr;
 	m_config_path.clear();
 	m_config_mode = cfg_mode::custom;
 
diff --git a/rpcs3/Emu/System.h b/rpcs3/Emu/System.h
index c64f694ed4..c9eb111f09 100644
--- a/rpcs3/Emu/System.h
+++ b/rpcs3/Emu/System.h
@@ -179,6 +179,7 @@ public:
 	std::vector<u128> klic;
 	std::string disc;
 	std::string hdd1;
+	std::function<void(u32)> init_mem_containers;
 
 	u32 m_boot_source_type = 0; // CELL_GAME_GAMETYPE_SYS