Merge pull request #12 from IW4x/bug/rollback-menu-code

[Menus] Rollback old menu code

src/Components/Modules/AssetInterfaces/IMenuList.cpp

@@ -7,7 +7,7 @@ namespace Assets
         Utils::Memory::Allocator* allocator = builder->getAllocator();
 
         // actually gets the whole list
-        std::vector<Game::menuDef_t*> menus = Components::Menus::LoadMenu(name);
+        auto menus = Components::Menus::LoadMenu(name);
         if (menus.empty()) return;
 
         // Allocate new menu list
@@ -27,7 +27,7 @@ namespace Assets
         // Copy new menus
         for (unsigned int i = 0; i < menus.size(); ++i)
         {
-            newList->menus[i] = menus[i];
+            newList->menus[i] = menus[i].second;
         }
 
         header->menuList = newList;

src/Components/Modules/AssetInterfaces/ImenuDef_t.cpp

@@ -13,7 +13,7 @@ namespace Assets
         if (menus.size() == 0) return;
         if (menus.size() > 1) Components::Logger::Print("Menu '%s' on disk has more than one menudef in it. Only saving the first one\n", name.data());
 
-        header->menu = menus[0];
+        header->menu = menus[0].second;
     }
 
 

src/Components/Modules/Menus.cpp

@@ -2,7 +2,9 @@
 
 namespace Components
 {
-	std::unordered_map<std::string, Game::menuDef_t*> Menus::DiskMenuList;
+	std::vector<std::string> Menus::CustomMenus;
+	std::unordered_map<std::string, Game::menuDef_t*> Menus::MenuList;
+	std::unordered_map<std::string, Game::MenuList*> Menus::MenuListList;
 
 	int Menus::ReserveSourceHandle()
 	{
@@ -164,12 +166,79 @@ namespace Components
 			}
 		}
 
+		Menus::OverrideMenu(menu);
+		Menus::RemoveMenu(menu->window.name);
+		Menus::MenuList[menu->window.name] = menu;
+
 		return menu;
 	}
 
-	std::vector<Game::menuDef_t*> Menus::LoadMenu(const std::string& menu)
+	Game::MenuList* Menus::LoadCustomMenuList(const std::string& menu, Utils::Memory::Allocator* allocator)
 	{
-		std::vector<Game::menuDef_t*> menus;
+		std::vector<std::pair<bool, Game::menuDef_t*>> menus;
+		FileSystem::File menuFile(menu);
+
+		if (!menuFile.exists()) return nullptr;
+
+		Game::pc_token_t token;
+		int handle = Menus::LoadMenuSource(menu, menuFile.getBuffer());
+
+		if (Menus::IsValidSourceHandle(handle))
+		{
+			while (true)
+			{
+				ZeroMemory(&token, sizeof(token));
+
+				if (!Game::PC_ReadTokenHandle(handle, &token) || token.string[0] == '}')
+				{
+					break;
+				}
+
+				if (!_stricmp(token.string, "loadmenu"))
+				{
+					Game::PC_ReadTokenHandle(handle, &token);
+
+					Utils::Merge(&menus, Menus::LoadMenu(Utils::String::VA("ui_mp\\%s.menu", token.string)));
+				}
+
+				if (!_stricmp(token.string, "menudef"))
+				{
+					Game::menuDef_t* menudef = Menus::ParseMenu(handle);
+					if (menudef) menus.push_back({ true, menudef }); // Custom menu
+				}
+			}
+
+			Menus::FreeMenuSource(handle);
+		}
+
+		if (menus.empty()) return nullptr;
+
+		// Allocate new menu list
+		Game::MenuList* list = allocator->allocate<Game::MenuList>();
+		if (!list) return nullptr;
+
+		list->menus = allocator->allocateArray<Game::menuDef_t*>(menus.size());
+		if (!list->menus)
+		{
+			allocator->free(list);
+			return nullptr;
+		}
+
+		list->name = allocator->duplicateString(menu);
+		list->menuCount = menus.size();
+
+		// Copy new menus
+		for (unsigned int i = 0; i < menus.size(); ++i)
+		{
+			list->menus[i] = menus[i].second;
+		}
+
+		return list;
+	}
+
+	std::vector<std::pair<bool, Game::menuDef_t*>> Menus::LoadMenu(const std::string& menu)
+	{
+		std::vector<std::pair<bool, Game::menuDef_t*>> menus;
 		FileSystem::File menuFile(menu);
 
 		if (menuFile.exists())
@@ -198,7 +267,7 @@ namespace Components
 					if (!_stricmp(token.string, "menudef"))
 					{
 						Game::menuDef_t* menudef = Menus::ParseMenu(handle);
-						if (menudef) menus.push_back(menudef); // Custom menu
+						if (menudef) menus.push_back({ true, menudef }); // Custom menu
 					}
 				}
 
@@ -206,15 +275,171 @@ namespace Components
 			}
 		}
 
-        // store loaded menus to be freed later
+		return menus;
-        for (auto it = menus.begin(); it != menus.end(); ++it)
+	}
+
+	std::vector<std::pair<bool, Game::menuDef_t*>> Menus::LoadMenu(Game::menuDef_t* menudef)
 	{
-            Menus::DiskMenuList[(*it)->window.name] = *it;
+		std::vector<std::pair<bool, Game::menuDef_t*>> menus = Menus::LoadMenu(Utils::String::VA("ui_mp\\%s.menu", menudef->window.name));
+
+		if (menus.empty())
+		{
+			// 			// Try loading the original menu, if we can't load our custom one
+			// 			Game::menuDef_t* originalMenu = AssetHandler::FindOriginalAsset(Game::XAssetType::ASSET_TYPE_MENU, menudef->window.name).menu;
+			//
+			// 			if (originalMenu)
+			// 			{
+			// 				menus.push_back({ false, originalMenu });
+			// 			}
+			// 			else
+			// 			{
+			menus.push_back({ false, menudef }); // Native menu
+// 			}
 		}
 
 		return menus;
 	}
 
+	Game::MenuList* Menus::LoadScriptMenu(const char* menu)
+	{
+		Utils::Memory::Allocator* allocator = Utils::Memory::GetAllocator();
+
+		std::vector<std::pair<bool, Game::menuDef_t*>> menus = Menus::LoadMenu(menu);
+		if (menus.empty()) return nullptr;
+
+		// Allocate new menu list
+		Game::MenuList* newList = allocator->allocate<Game::MenuList>();
+		if (!newList) return nullptr;
+
+		newList->menus = allocator->allocateArray<Game::menuDef_t*>(menus.size());
+		if (!newList->menus)
+		{
+			allocator->free(newList);
+			return nullptr;
+		}
+
+		newList->name = allocator->duplicateString(menu);
+		newList->menuCount = menus.size();
+
+		// Copy new menus
+		for (unsigned int i = 0; i < menus.size(); ++i)
+		{
+			newList->menus[i] = menus[i].second;
+		}
+
+		Menus::RemoveMenuList(newList->name);
+		Menus::MenuListList[newList->name] = newList;
+
+		return newList;
+	}
+
+	void Menus::SafeMergeMenus(std::vector<std::pair<bool, Game::menuDef_t*>>* menus, std::vector<std::pair<bool, Game::menuDef_t*>> newMenus)
+	{
+		// Check if we overwrote a menu
+		for (auto i = menus->begin(); i != menus->end();)
+		{
+			// Try to find the native menu
+			bool found = !i->first; // Only if custom menu, try to find it
+
+			// If there is none, try to find a custom menu
+			if (!found)
+			{
+				for (auto& entry : Menus::MenuList)
+				{
+					if (i->second == entry.second)
+					{
+						found = true;
+						break;
+					}
+				}
+			}
+
+			// Remove the menu if it has been deallocated (not found)
+			if (!found)
+			{
+				i = menus->erase(i);
+				continue;
+			}
+
+			bool increment = true;
+
+			// Remove the menu if it has been loaded twice
+			for (auto& newMenu : newMenus)
+			{
+				if (i->second->window.name == std::string(newMenu.second->window.name))
+				{
+					Menus::RemoveMenu(i->second);
+
+					i = menus->erase(i);
+					increment = false;
+					break;
+				}
+			}
+
+			if (increment) ++i;
+		}
+
+		Utils::Merge(menus, newMenus);
+	}
+
+	Game::MenuList* Menus::LoadMenuList(Game::MenuList* menuList)
+	{
+		Utils::Memory::Allocator* allocator = Utils::Memory::GetAllocator();
+
+		std::vector<std::pair<bool, Game::menuDef_t*>> menus;
+
+		for (int i = 0; i < menuList->menuCount; ++i)
+		{
+			if (!menuList->menus[i]) continue;
+			Menus::SafeMergeMenus(&menus, Menus::LoadMenu(menuList->menus[i]));
+		}
+
+		// Load custom menus
+		if (menuList->name == "ui_mp/code.txt"s) // Should be menus, but code is loaded ingame
+		{
+			for (auto menu : Menus::CustomMenus)
+			{
+				bool hasMenu = false;
+				for (auto& loadedMenu : menus)
+				{
+					if (loadedMenu.second->window.name == menu)
+					{
+						hasMenu = true;
+						break;
+					}
+				}
+
+				if (!hasMenu) Menus::SafeMergeMenus(&menus, Menus::LoadMenu(menu));
+			}
+		}
+
+		// Allocate new menu list
+		Game::MenuList* newList = allocator->allocate<Game::MenuList>();
+		if (!newList) return menuList;
+
+		size_t size = menus.size();
+		newList->menus = allocator->allocateArray<Game::menuDef_t*>(size);
+		if (!newList->menus)
+		{
+			allocator->free(newList);
+			return menuList;
+		}
+
+		newList->name = allocator->duplicateString(menuList->name);
+		newList->menuCount = size;
+
+		// Copy new menus
+		for (unsigned int i = 0; i < menus.size(); ++i)
+		{
+			newList->menus[i] = menus[i].second;
+		}
+
+		Menus::RemoveMenuList(newList->name);
+		Menus::MenuListList[newList->name] = newList;
+
+		return newList;
+	}
+
 	void Menus::FreeMenuSource(int handle)
 	{
 		Utils::Memory::Allocator* allocator = Utils::Memory::GetAllocator();
@@ -258,29 +483,227 @@ namespace Components
 		Game::sourceFiles[handle] = nullptr;
 	}
 
-	void Menus::FreeDiskMenu(Game::menuDef_t* menudef)
+	void Menus::FreeMenu(Game::menuDef_t* menudef)
 	{
 		Utils::Memory::Allocator* allocator = Utils::Memory::GetAllocator();
 
 		// Do i need to free expressions and strings?
 		// Or does the game take care of it?
 		// Seems like it does...
+
 		if (menudef->items)
 		{
+			// Seems like this is obsolete as well,
+			// as the game handles the memory
+
+			//for (int i = 0; i < menudef->itemCount; ++i)
+			//{
+			//	Game::Menu_FreeItemMemory(menudef->items[i]);
+			//}
+
 			allocator->free(menudef->items);
 		}
 
 		allocator->free(menudef);
 	}
 
-	void Menus::FreeEverything()
+	void Menus::FreeMenuList(Game::MenuList* menuList)
 	{
-		for (auto i = Menus::DiskMenuList.begin(); i != Menus::DiskMenuList.end(); ++i)
+		if (!menuList) return;
+		Utils::Memory::Allocator* allocator = Utils::Memory::GetAllocator();
+
+		// Keep our compiler happy
+		Game::MenuList list = { menuList->name, menuList->menuCount, menuList->menus };
+
+		if (list.name)
 		{
-			Menus::FreeDiskMenu(i->second);
+			allocator->free(list.name);
 		}
 
-		Menus::DiskMenuList.clear();
+		if (list.menus)
+		{
+			allocator->free(list.menus);
+		}
+
+		allocator->free(menuList);
+	}
+
+	void Menus::RemoveMenu(const std::string& menu)
+	{
+		auto i = Menus::MenuList.find(menu);
+		if (i != Menus::MenuList.end())
+		{
+			if (i->second) Menus::FreeMenu(i->second);
+			i = Menus::MenuList.erase(i);
+		}
+	}
+
+	void Menus::RemoveMenu(Game::menuDef_t* menudef)
+	{
+		for (auto i = Menus::MenuList.begin(); i != Menus::MenuList.end();)
+		{
+			if (i->second == menudef)
+			{
+				Menus::FreeMenu(menudef);
+				i = Menus::MenuList.erase(i);
+			}
+			else
+			{
+				++i;
+			}
+		}
+	}
+
+	void Menus::RemoveMenuList(const std::string& menuList)
+	{
+		auto i = Menus::MenuListList.find(menuList);
+		if (i != Menus::MenuListList.end())
+		{
+			if (i->second)
+			{
+				for (auto j = 0; j < i->second->menuCount; ++j)
+				{
+					Menus::RemoveMenu(i->second->menus[j]);
+				}
+
+				Menus::FreeMenuList(i->second);
+			}
+
+			i = Menus::MenuListList.erase(i);
+		}
+	}
+
+	// This is actually a really important function
+	// It checks if we have already loaded the menu we passed and replaces its instances in memory
+	// Due to deallocating the old menu, the game might crash on not being able to handle its old instance
+	// So we need to override it in our menu lists and the game's ui context
+	// EDIT: We might also remove the old instances inside RemoveMenu
+	// EDIT2: Removing old instances without having a menu to replace them with might leave a nullptr
+	// EDIT3: Wouldn't it be better to check if the new menu we're trying to load has already been loaded and not was not deallocated and return that one instead of loading a new one?
+	void Menus::OverrideMenu(Game::menuDef_t* menu)
+	{
+		if (!menu || !menu->window.name) return;
+		std::string name = menu->window.name;
+
+		// Find the old menu
+		auto i = Menus::MenuList.find(name);
+		if (i != Menus::MenuList.end())
+		{
+			// We have found it, *yay*
+			Game::menuDef_t* oldMenu = i->second;
+
+			// Replace every old instance with our new one in the ui context
+			for (int j = 0; j < Game::uiContext->menuCount; ++j)
+			{
+				if (Game::uiContext->Menus[j] == oldMenu)
+				{
+					Game::uiContext->Menus[j] = menu;
+				}
+			}
+
+			// Replace every old instance with our new one in our menu lists
+			for (auto j = Menus::MenuListList.begin(); j != Menus::MenuListList.end(); ++j)
+			{
+				Game::MenuList* list = j->second;
+
+				if (list && list->menus)
+				{
+					for (int k = 0; k < list->menuCount; ++k)
+					{
+						if (list->menus[k] == oldMenu)
+						{
+							list->menus[k] = menu;
+						}
+					}
+				}
+			}
+		}
+	}
+
+	void Menus::RemoveMenuList(Game::MenuList* menuList)
+	{
+		if (!menuList || !menuList->name) return;
+		Menus::RemoveMenuList(menuList->name);
+	}
+
+	void Menus::FreeEverything()
+	{
+		for (auto i = Menus::MenuListList.begin(); i != Menus::MenuListList.end(); ++i)
+		{
+			Menus::FreeMenuList(i->second);
+		}
+
+		Menus::MenuListList.clear();
+
+		for (auto i = Menus::MenuList.begin(); i != Menus::MenuList.end(); ++i)
+		{
+			Menus::FreeMenu(i->second);
+		}
+
+		Menus::MenuList.clear();
+	}
+
+	Game::XAssetHeader Menus::MenuFindHook(Game::XAssetType /*type*/, const std::string& filename)
+	{
+		return { Game::Menus_FindByName(Game::uiContext, filename.data()) };
+	}
+
+	Game::XAssetHeader Menus::MenuListFindHook(Game::XAssetType type, const std::string& filename)
+	{
+		Game::XAssetHeader header = { nullptr };
+
+		// Free the last menulist and ui context, as we have to rebuild it with the new menus
+		if (Menus::MenuListList.find(filename) != Menus::MenuListList.end())
+		{
+			Game::MenuList* list = Menus::MenuListList[filename];
+
+			for (int i = 0; list && list->menus && i < list->menuCount; ++i)
+			{
+				Menus::RemoveMenuFromContext(Game::uiContext, list->menus[i]);
+			}
+
+			Menus::RemoveMenuList(filename);
+		}
+
+		if (Utils::String::EndsWith(filename, ".menu"))
+		{
+			if (FileSystem::File(filename).exists())
+			{
+				header.menuList = Menus::LoadScriptMenu(filename.data());
+				if (header.menuList) return header;
+			}
+		}
+
+		Game::MenuList* menuList = Game::DB_FindXAssetHeader(type, filename.data()).menuList;
+		header.menuList = menuList;
+
+		if (menuList && reinterpret_cast<DWORD>(menuList) != 0xDDDDDDDD)
+		{
+			// Parse scriptmenus!
+			if ((menuList->menuCount > 0 && menuList->menus[0] && menuList->menus[0]->window.name == "default_menu"s))
+			{
+				if (FileSystem::File(filename).exists())
+				{
+					header.menuList = Menus::LoadScriptMenu(filename.data());
+
+					// Reset, if we didn't find scriptmenus
+					if (!header.menuList)
+					{
+						header.menuList = menuList;
+					}
+				}
+			}
+			else
+			{
+				header.menuList = Menus::LoadMenuList(menuList);
+			}
+		}
+		else
+		{
+			header.menuList = nullptr;
+		}
+
+		return header;
 	}
 
 	bool Menus::IsMenuVisible(Game::UiContext* dc, Game::menuDef_t* menu)
@@ -293,7 +716,7 @@ namespace Components
 
 				if (originalConnect == menu) // Check if we draw the original loadscreen
 				{
-					if (Menus::DiskMenuList.find("connect") != Menus::DiskMenuList.end()) // Check if we have a custom loadscreen, to prevent drawing the original one on top
+					if (Menus::MenuList.find("connect") != Menus::MenuList.end()) // Check if we have a custom loadscreen, to prevent drawing the original one on top
 					{
 						return false;
 					}
@@ -304,110 +727,82 @@ namespace Components
 		return Game::Menu_IsVisible(dc, menu);
 	}
 
-    void Menus::AddMenuListToContext(Game::UiContext* ctx, Game::MenuList* list, int close)
+	void Menus::RemoveMenuFromContext(Game::UiContext* dc, Game::menuDef_t* menu)
 	{
-        // scriptmenu
+		// Search menu in context
-        if (std::string(list->name).find(".menu") != std::string::npos)
+		int i = 0;
+		for (; i < dc->menuCount; ++i)
 		{
-            auto menus = Menus::LoadMenu(list->name);
+			if (dc->Menus[i] == menu)
-
-            if (menus.size())
 			{
-                Logger::Print("Overriding menu '%s'\n", list->name);
+				break;
-                for (auto it = menus.begin(); it != menus.end(); ++it)
-                {
-                    if (ctx->menuCount < MAX_MENUS_IN_CONTEXT)
-                    {
-                        ctx->Menus[ctx->menuCount++] = *it;
-                    }
-
-                    if (close)
-                    {
-                        Game::Menus_CloseRequest(ctx, *it);
 			}
 		}
 
-                return; // don't add original menus
+		// Remove from stack
-            }
+		if (i < dc->menuCount)
-        }
-            
-
-        for (int i = 0; i < list->menuCount; i++)
 		{
-            Game::menuDef_t* cur = list->menus[i];
+			for (; i < dc->menuCount - 1; ++i)
-
-            if (cur->window.name == reinterpret_cast<const char*>(0xDDDDDDDD))
 			{
-                DebugBreak();
+				dc->Menus[i] = dc->Menus[i + 1];
 			}
 
-            auto menus = Menus::LoadMenu(Utils::String::VA("ui_mp/%s.menu", cur->window.name));
+			// Clear last menu
-
+			dc->Menus[--dc->menuCount] = nullptr;
-            if (menus.size())
-            {
-                Logger::Print("Overriding menu '%s'\n", cur->window.name);
-                if (menus.size() > 1) Logger::Print("Disk menu has more than one definition Using only the first one.\n");
-                cur = menus[0]; // replace menu in context with loaded one
-            }
-
-            if (ctx->menuCount < MAX_MENUS_IN_CONTEXT)
-            {
-                ctx->Menus[ctx->menuCount++] = cur;
-            }
-
-            if (close)
-            {
-                Game::Menus_CloseRequest(ctx, cur);
 		}
 	}
+
+	void Menus::Add(const std::string& menu)
+	{
+		Menus::CustomMenus.push_back(menu);
 	}
 
-    void Menus::RegisterMenuLists()
+	void Menus::RegisterCustomMenusHook()
 	{
-        Utils::Hook::Call<void()>(0x401700)(); // reset ui context
+		Game::UiContext* uiInfoArray = (Game::UiContext*)0x62E2858;
+		// Game::MenuList list;
 
-        // we can't call DB_FindXAssetHeader here because it blocks the rest of loading waiting on those 2 menulsits
+		Utils::Hook::Call<void()>(0x401700)(); // call original load functions
-        // TODO: Figure out a better way to trigger the custom menulist loading because if you skip the intro the 
+		//Game::XAssetHeader header = Game::DB_FindXAssetHeader(Game::XAssetType::ASSET_TYPE_MENULIST, "ui_mp/iw4x.txt");
-        // custom menus won't have loaded until a few seconds later. All overridden menus are already
+		//if (header.data && !(header.menuList->menuCount == 1 && !_stricmp("default_menu", header.menuList->menus[0]->window.name)))
-        // loaded so it isn't a black screen but it wont show the first time intro, credits, etc.
+		//{
-        // as soon as this loads those start to work again
+		//	// Utils::Hook::Call<void(void*, Game::MenuList*, int)>(0x401700)(uiInfoArray, header.menuList, 1); // add loaded menus
-        // if we just trigger this here it blocks the intro from showing because of the FindXAssetHeader calls
+		//	std::memcpy(&list, header.data, sizeof(Game::MenuList));
-        // that are waiting for zones to finish loading
 
-        auto loadCustomMenus = []()
+		//	for (int i = 0; i < uiInfoArray->menuCount; i++)
+		//	{
+		//		for (int j = 0; j < list.menuCount; j++)
+		//		{
+		//			if (!list.menus[j]) continue; // skip already used entries
+		//			if (!stricmp(list.menus[j]->window.name, uiInfoArray->Menus[i]->window.name))
+		//			{
+		//				uiInfoArray->Menus[i] = list.menus[j]; // overwrite UiContext pointer
+		//				list.menus[j] = nullptr; // clear entries that already exist so we don't add them later
+		//			}
+		//		}
+		//	}
+
+		//	for (int i = 0; i < list.menuCount; i++)
+		//	{
+		//		if (list.menus[i])
+		//		{
+		//			uiInfoArray->Menus[uiInfoArray->menuCount++] = list.menus[i];
+		//		}
+		//	}
+		//}
+
+		for (int i = 0; i < uiInfoArray->menuCount; i++)
 		{
-            // attempt to load iw4x menus            
+			OutputDebugStringA(Utils::String::VA("%s\n", uiInfoArray->Menus[i]->window.name));
-            Game::XAssetHeader header = Game::DB_FindXAssetHeader(Game::XAssetType::ASSET_TYPE_MENULIST, "ui_mp/iw4x.txt");
-            if (header.data && !(header.menuList->menuCount == 1 && !_stricmp("default_menu", header.menuList->menus[0]->window.name)))
-            {
-                Menus::AddMenuListToContext(Game::uiContext, header.menuList, 1);
 		}
 
-            // attempt to load mod menus
+		/*
 		header = Game::DB_FindXAssetHeader(Game::XAssetType::ASSET_TYPE_MENULIST, "ui_mp/mod.txt");
 		if (header.data && !(header.menuList->menuCount == 1 && !_stricmp("default_menu", header.menuList->menus[0]->window.name)))
 		{
-                Menus::AddMenuListToContext(Game::uiContext, header.menuList, 1);
+			Utils::Hook::Call<void(void*, Game::MenuList*, int)>(0x401700)(uiInfoArray, header.menuList, 1); // add loaded menus
 		}
-        };
+		*/
-
-        if (!FastFiles::Ready())
-        {
-            Scheduler::OnReady(loadCustomMenus, true);
-        }
-        else
-        {
-            loadCustomMenus();
-        }
-    }
-
-    void Menus::ResetContextHook(int a1)
-    {
-        // reset our lists
-        Menus::FreeEverything();
-
-        // continue with initialization
-        Utils::Hook::Call<void(int)>(0x4A57D0)(a1);
 	}
 
 	Menus::Menus()
@@ -417,17 +812,15 @@ namespace Components
 		// Ensure everything is zero'ed
 		Menus::FreeEverything();
 
+		// Intercept asset finding
+		AssetHandler::OnFind(Game::XAssetType::ASSET_TYPE_MENU, Menus::MenuFindHook);
+		AssetHandler::OnFind(Game::XAssetType::ASSET_TYPE_MENULIST, Menus::MenuListFindHook);
+
 		// Don't open connect menu
 		//Utils::Hook::Nop(0x428E48, 5);
 
 		// register custom menufiles if they exist
-         Utils::Hook(0x4A58C3, Menus::RegisterMenuLists, HOOK_CALL).install()->quick();
+		Utils::Hook(0x4A58C3, Menus::RegisterCustomMenusHook, HOOK_CALL).install()->quick();
-
-        // take control of menus in uiContext
-        Utils::Hook(0x4533C0, Menus::AddMenuListToContext, HOOK_JUMP).install()->quick();
-
-        // reset our list on UiContext reset
-        Utils::Hook(0x4B5422, Menus::ResetContextHook, HOOK_CALL).install()->quick();
 
 		// Use the connect menu open call to update server motds
 		Utils::Hook(0x428E48, []()
@@ -472,7 +865,7 @@ namespace Components
 					// Close all menus
 					Game::Menus_CloseAll(Game::uiContext);
 
-			// Free custom menus and reset uiContext list
+					// Free custom menus
 					Menus::FreeEverything();
 
 					// Only disconnect if in-game, context is updated automatically!
@@ -482,7 +875,8 @@ namespace Components
 					}
 					else
 					{
-                Menus::RegisterMenuLists(); // register custom menus
+						// Reinitialize ui context
+						Utils::Hook::Call<void()>(0x401700)();
 
 						// Reopen main menu
 						Game::Menus_OpenByName(Game::uiContext, "main_text");
@@ -497,10 +891,27 @@ namespace Components
 				{
 					Command::Execute("openmenu quickmessage");
 				});
+
+			// Define custom menus here
+			Menus::Add("ui_mp/changelog.menu");
+			Menus::Add("ui_mp/theater_menu.menu");
+			Menus::Add("ui_mp/pc_options_multi.menu");
+			Menus::Add("ui_mp/pc_options_game.menu");
+			Menus::Add("ui_mp/stats_reset.menu");
+			Menus::Add("ui_mp/stats_unlock.menu");
+			Menus::Add("ui_mp/security_increase_popmenu.menu");
+			Menus::Add("ui_mp/mod_download_popmenu.menu");
+			Menus::Add("ui_mp/popup_friends.menu");
+			Menus::Add("ui_mp/menu_first_launch.menu");
+			Menus::Add("ui_mp/startup_messages.menu");
+			Menus::Add("ui_mp/pc_store.menu");
+			Menus::Add("ui_mp/iw4x_credits.menu");
+			Menus::Add("ui_mp/resetclass.menu");
 	}
 
 	Menus::~Menus()
 	{
+		Menus::CustomMenus.clear();
 		Menus::FreeEverything();
 	}
 }

src/Components/Modules/Menus.hpp

@@ -1,7 +1,6 @@
 #pragma once
 
 #define MAX_SOURCEFILES	64
-#define MAX_MENUS_IN_CONTEXT 640
 #undef LoadMenu
 
 namespace Components
@@ -13,31 +12,51 @@ namespace Components
 		~Menus();
 
 		static void FreeEverything();
-        static void RegisterMenuLists();
 
-        // used to load assets for zonebuilder
+		static void Add(const std::string& menu);
-        static std::vector<Game::menuDef_t*> LoadMenu(const std::string& file);
+
+		static Game::MenuList* Menus::LoadCustomMenuList(const std::string& menu, Utils::Memory::Allocator* allocator);
+		static std::vector<std::pair<bool, Game::menuDef_t*>> LoadMenu(Game::menuDef_t* menudef);
+		static std::vector<std::pair<bool, Game::menuDef_t*>> LoadMenu(const std::string& file);
 		
 	private:
-		static std::unordered_map<std::string, Game::menuDef_t*> DiskMenuList;
+		static std::unordered_map<std::string, Game::menuDef_t*> MenuList;
+		static std::unordered_map<std::string, Game::MenuList*> MenuListList;
+		static std::vector<std::string> CustomMenus;
+
+		static Game::XAssetHeader MenuFindHook(Game::XAssetType type, const std::string& filename);
+		static Game::XAssetHeader MenuListFindHook(Game::XAssetType type, const std::string& filename);
+
+		static Game::MenuList* LoadMenuList(Game::MenuList* menuList);
+		static Game::MenuList* LoadScriptMenu(const char* menu);
+
+		static void SafeMergeMenus(std::vector<std::pair<bool, Game::menuDef_t*>>* menus, std::vector<std::pair<bool, Game::menuDef_t*>> newMenus);
 
-        // Loading
-        static int ReserveSourceHandle();
-        static bool IsValidSourceHandle(int handle);
-        static Game::menuDef_t* ParseMenu(int handle);
 		static Game::script_t* LoadMenuScript(const std::string& name, const std::string& buffer);
 		static int LoadMenuSource(const std::string& name, const std::string& buffer);
 
-        // Freeing
+		static int ReserveSourceHandle();
-		static void FreeMenuSource(int handle);
+		static bool IsValidSourceHandle(int handle);
-		static void FreeDiskMenu(Game::menuDef_t* menudef);
+
+		static Game::menuDef_t* ParseMenu(int handle);
+
+		static void FreeMenuSource(int handle);
+
+		static void FreeMenuList(Game::MenuList* menuList);
+		static void FreeMenu(Game::menuDef_t* menudef);
+
+		static void RemoveMenu(const std::string& menu);
+		static void RemoveMenu(Game::menuDef_t* menudef);
+		static void RemoveMenuList(const std::string& menuList);
+		static void RemoveMenuList(Game::MenuList* menuList);
+
+		static void OverrideMenu(Game::menuDef_t* menu);
 
-        // Etc.
 		static bool IsMenuVisible(Game::UiContext* dc, Game::menuDef_t* menu);
 
-        // Manage menus in uiContext
+		static void RemoveMenuFromContext(Game::UiContext* dc, Game::menuDef_t* menu);
-        static void AddMenuListToContext(Game::UiContext* ctx, Game::MenuList* list, int close);
+
-        static void ResetContextHook(int a1);
+		static void RegisterCustomMenusHook();
 
 		// Ugly!
 		static int KeywordHash(char* key);