Upgrade Lua service intercepts (#775)

Co-authored-by: Théo B. <16072534+LiquidFenrir@users.noreply.github.com>

include/services/service_manager.hpp

@@ -3,7 +3,7 @@
 #include <optional>
 #include <span>
 #include <string>
-#include <unordered_set>
+#include <unordered_map>
 
 #include "kernel_types.hpp"
 #include "logger.hpp"
@@ -94,8 +94,10 @@ class ServiceManager {
 	// For example, if we want to intercept dsp::DSP ReadPipe (Header: 0x000E00C0), the "serviceName" field would be "dsp::DSP"
 	// and the "function" field would be 0x000E00C0
 	LuaManager& lua;
-	std::unordered_set<InterceptedService> interceptedServices = {};
-	// Calling std::unordered_set<T>::size() compiles to a fairly non-trivial function call on Clang, so we store this
+
+	// Map from service intercept entries to their corresponding Lua callbacks
+	std::unordered_map<InterceptedService, int> interceptedServices = {};
+	// Calling std::unordered_map<T>::size() compiles to a non-trivial function call on Clang, so we store this
 	// separately and check it on service calls, for performance reasons
 	bool haveServiceIntercepts = false;
 
@@ -134,12 +136,23 @@ class ServiceManager {
 	Y2RService& getY2R() { return y2r; }
 	IRUserService& getIRUser() { return ir_user; }
 
-	void addServiceIntercept(const std::string& service, u32 function) {
-		interceptedServices.insert(InterceptedService(service, function));
+	void addServiceIntercept(const std::string& service, u32 function, int callbackRef) {
+		auto success = interceptedServices.try_emplace(InterceptedService(service, function), callbackRef);
+		if (!success.second) {
+			// An intercept for this service function already exists
+			// Remove the old callback and set the new one
+			lua.removeInterceptedService(service, function, success.first->second);
+			success.first->second = callbackRef;
+		}
+
 		haveServiceIntercepts = true;
 	}
 
 	void clearServiceIntercepts() {
+		for (const auto& [interceptedService, callbackRef] : interceptedServices) {
+			lua.removeInterceptedService(interceptedService.serviceName, interceptedService.function, callbackRef);
+		}
+
 		interceptedServices.clear();
 		haveServiceIntercepts = false;
 	}

include/services/service_map.hpp

@@ -0,0 +1,29 @@
+#pragma once
+#include <string>
+#include <type_traits>
+#include <utility>
+
+#include "handles.hpp"
+
+// Helpers for constructing std::maps to look up OS services.
+// We want to be able to map both service names -> services (Used for OS emulation)
+// And service handles -> services (For Lua service call intercepts)
+using ServiceMapEntry = std::pair<std::string, HorizonHandle>;
+
+// Comparator for constructing a name->handle service map
+struct ServiceMapByNameComparator {
+	// The comparators must be transparent, as our search key is different from our set key
+	// Our set key is a ServiceMapEntry, while the comparator each time is either the name or the service handle
+	using is_transparent = std::true_type;
+	bool operator()(const ServiceMapEntry& lhs, std::string_view rhs) const { return lhs.first < rhs; }
+	bool operator()(std::string_view lhs, const ServiceMapEntry& rhs) const { return lhs < rhs.first; }
+	bool operator()(const ServiceMapEntry& lhs, const ServiceMapEntry& rhs) const { return lhs.first < rhs.first; }
+};
+
+// Comparator for constructing a handle->name service map
+struct ServiceMapByHandleComparator {
+	using is_transparent = std::true_type;
+	bool operator()(const ServiceMapEntry& lhs, HorizonHandle rhs) const { return lhs.second < rhs; }
+	bool operator()(HorizonHandle lhs, const ServiceMapEntry& rhs) const { return lhs < rhs.second; }
+	bool operator()(const ServiceMapEntry& lhs, const ServiceMapEntry& rhs) const { return lhs.second < rhs.second; }
+};