Clean up service intercepts

CMakeLists.txt

@@ -415,7 +415,7 @@ set(HEADER_FILES include/emulator.hpp include/helpers.hpp include/termcolor.hpp
                  include/services/dsp_firmware_db.hpp include/frontend_settings.hpp include/fs/archive_twl_photo.hpp
                  include/fs/archive_twl_sound.hpp include/fs/archive_card_spi.hpp include/services/ns.hpp include/audio/audio_device.hpp
                  include/audio/audio_device_interface.hpp include/audio/libretro_audio_device.hpp include/services/ir/ir_types.hpp
-                 include/services/ir/ir_device.hpp include/services/ir/circlepad_pro.hpp
+                 include/services/ir/ir_device.hpp include/services/ir/circlepad_pro.hpp include/services/service_intercept.hpp
 )
 
 if(IOS)

include/lua_manager.hpp

@@ -63,4 +63,5 @@ class LuaManager {
 	void reset() {}
 	void signalEvent(LuaEvent e) {}
 	bool signalInterceptedService(const std::string& service, u32 function, u32 messagePointer) { return false; }
+};
 #endif

include/services/service_intercept.hpp

@@ -0,0 +1,28 @@
+#pragma once
+#include <functional>
+#include <string>
+
+#include "helpers.hpp"
+
+// We allow Lua scripts to intercept service calls and allow their own code to be ran on SyncRequests
+// 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
+struct InterceptedService {
+	std::string serviceName;  // Name of the service whose function
+	u32 function;             // Header of the function to intercept
+
+	InterceptedService(const std::string& name, u32 header) : serviceName(name), function(header) {}
+	bool operator==(const InterceptedService& other) const { return serviceName == other.serviceName && function == other.function; }
+};
+
+// Define hashing function for InterceptedService to use it with unordered_map
+namespace std {
+	template <>
+	struct hash<InterceptedService> {
+		usize operator()(const InterceptedService& s) const noexcept {
+			const usize hash1 = std::hash<std::string>{}(s.serviceName);
+			const usize hash2 = std::hash<u32>{}(s.function);
+			return hash1 ^ (hash2 << 1);
+		}
+	};
+}  // namespace std

include/services/service_manager.hpp

@@ -37,6 +37,7 @@
 #include "services/ns.hpp"
 #include "services/nwm_uds.hpp"
 #include "services/ptm.hpp"
+#include "services/service_intercept.hpp"
 #include "services/soc.hpp"
 #include "services/ssl.hpp"
 #include "services/y2r.hpp"
@@ -93,23 +94,14 @@ 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;
-	struct InterceptedService {
-		std::string serviceName;  // Name of the service whose function
-		u32 function;             // Header of the function to intercept
+	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
+	// separately and check it on service calls, for performance reasons
+	bool haveServiceIntercepts = false;
 
-		InterceptedService(const std::string& name, u32 header) : serviceName(name), function(header) {}
-		bool operator==(const InterceptedService& other) const { return serviceName == other.serviceName && function == other.function; }
-	};
-
-	struct InterceptedServiceHash {
-		usize operator()(const InterceptedService& s) const noexcept {
-			usize h1 = std::hash<std::string>{}(s.serviceName);
-			usize h2 = std::hash<u32>{}(s.function);
-			return h1 ^ (h2 << 1);
-		}
-	};
-
-	std::unordered_set<InterceptedService, InterceptedServiceHash> interceptedServices = {};
+	// Checks for whether a service call is intercepted by Lua and handles it. Returns true if Lua told us not to handle the function,
+	// or false if we should handle it as normal
+	bool checkForIntercept(u32 messagePointer, Handle handle);
 
 	// "srv:" commands
 	void enableNotification(u32 messagePointer);
@@ -142,6 +134,13 @@ 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 clearServiceIntercepts() { interceptedServices.clear(); }
+	void addServiceIntercept(const std::string& service, u32 function) {
+		interceptedServices.insert(InterceptedService(service, function));
+		haveServiceIntercepts = true;
+	}
+
+	void clearServiceIntercepts() {
+		interceptedServices.clear();
+		haveServiceIntercepts = false;
+	}
 };

src/core/services/service_manager.cpp

@@ -216,25 +216,9 @@ void ServiceManager::publishToSubscriber(u32 messagePointer) {
 }
 
 void ServiceManager::sendCommandToService(u32 messagePointer, Handle handle) {
-	if (interceptedServices.size() != 0) [[unlikely]] {
-		// Check if there's a Lua handler for this function and call it
-		u32 function = mem.read32(messagePointer);
-
-		for (auto [serviceName, serviceHandle] : serviceMap) {
-			if (serviceHandle == handle) {
-				auto intercept = InterceptedService(std::string(serviceName), function);
-				if (interceptedServices.contains(intercept)) {
-					printf("Call to intercepted service\n");
-
-					// If the Lua handler returns true, it means the service is handled entirely
-					// From Lua, and we shouldn't do anything else here.
-					if (lua.signalInterceptedService(intercept.serviceName, function, messagePointer)) {
-						return;
-					}
-				}
-
-				break;
-			}
+	if (haveServiceIntercepts) [[unlikely]] {
+		if (checkForIntercept(messagePointer, handle)) [[unlikely]] {
+			return;
 		}
 	}
 
@@ -281,4 +265,25 @@ void ServiceManager::sendCommandToService(u32 messagePointer, Handle handle) {
 		case KernelHandles::Y2R: y2r.handleSyncRequest(messagePointer); break;
 		default: Helpers::panic("Sent IPC message to unknown service %08X\n Command: %08X", handle, mem.read32(messagePointer));
 	}
+}
+
+bool ServiceManager::checkForIntercept(u32 messagePointer, Handle handle) {
+	// Check if there's a Lua handler for this function and call it
+	const u32 function = mem.read32(messagePointer);
+
+	for (auto [serviceName, serviceHandle] : serviceMap) {
+		if (serviceHandle == handle) {
+			auto intercept = InterceptedService(std::string(serviceName), function);
+			if (interceptedServices.contains(intercept)) {
+				// If the Lua handler returns true, it means the service is handled entirely
+				// From Lua, and we shouldn't do anything else here.
+				return lua.signalInterceptedService(intercept.serviceName, function, messagePointer);
+			}
+
+			break;
+		}
+	}
+
+	// Lua did not intercept the service, so emulate it normally
+	return false;
 }

src/lua.cpp

@@ -1,5 +1,4 @@
 #ifdef PANDA3DS_ENABLE_LUA
-#include <fmt/format.h>
 #include <teakra/disassembler.h>
 
 #include <array>
@@ -119,8 +118,8 @@ bool LuaManager::signalInterceptedService(const std::string& service, u32 functi
 
 	if (status != LUA_OK) {
 		const char* err = lua_tostring(L, -1);
-		fprintf(stderr, "Lua error in interceptService: %s\n", err ? err : "(unknown error)");
-		lua_pop(L, 1);  // Remove error message from stack
+		fprintf(stderr, "Lua: Error in interceptService: %s\n", err);
+		lua_pop(L, 1);  // Pop error message from stack
 		return false;   // Have the C++ handle the service call
 	}
 
@@ -236,7 +235,7 @@ static int loadROMThunk(lua_State* L) {
 	return 1;
 }
 
-static int interceptServiceThunk(lua_State* L) {
+static int addServiceInterceptThunk(lua_State* L) {
 	// Service name argument is invalid, report that loading failed and exit
 	if (lua_type(L, 1) != LUA_TSTRING) {
 		lua_pushboolean(L, 0);
@@ -256,11 +255,14 @@ static int interceptServiceThunk(lua_State* L) {
 	const u32 function = (u32)lua_tointeger(L, 2);
 	const auto serviceName = std::string(str, nameLength);
 	LuaManager::g_emulator->getServiceManager().addServiceIntercept(serviceName, function);
-
-	fmt::print("Intercepting call to {} (Function id: {:08X})\n", serviceName, function);
 	return 2;
 }
 
+static int clearServiceInterceptsThunk(lua_State* L) {
+	LuaManager::g_emulator->getServiceManager().clearServiceIntercepts();
+	return 0;
+}
+
 static int getButtonsThunk(lua_State* L) {
 	auto buttons = LuaManager::g_emulator->getServiceManager().getHID().getOldButtons();
 	lua_pushinteger(L, static_cast<lua_Integer>(buttons));
@@ -347,7 +349,8 @@ static constexpr luaL_Reg functions[] = {
 	{ "__getButton", getButtonThunk },
 	{ "__disassembleARM", disassembleARMThunk },
 	{ "__disassembleTeak", disassembleTeakThunk },
-	{"__interceptService", interceptServiceThunk},
+	{"__addServiceIntercept", addServiceInterceptThunk },
+	{"__clearServiceIntercepts", clearServiceInterceptsThunk },
 	{ nullptr, nullptr },
 };
 // clang-format on
@@ -388,7 +391,8 @@ void LuaManager::initializeThunks() {
 
 		disassembleARM = function(pc, instruction) return GLOBALS.__disassembleARM(pc, instruction) end,
 		disassembleTeak = function(opcode, exp) return GLOBALS.__disassembleTeak(opcode, exp or 0) end,
-		interceptService = function(service, func) return GLOBALS.__interceptService(service, func) end,
+		addServiceIntercept = function(service, func) return GLOBALS.__addServiceIntercept(service, func) end,
+		clearServiceIntercepts = function() return GLOBALS.__clearServiceIntercepts() end,
 
 		Frame = __Frame,
 		ButtonA = __ButtonA,
@@ -397,6 +401,8 @@ void LuaManager::initializeThunks() {
 		ButtonY = __ButtonY,
 		ButtonL = __ButtonL,
 		ButtonR = __ButtonR,
+		ButtonZL = __ButtonZL,
+		ButtonZR = __ButtonZR,
 		ButtonUp = __ButtonUp,
 		ButtonDown = __ButtonDown,
 		ButtonLeft = __ButtonLeft,