From 1940ac0fec81b88ea5bbb586e1e0edd35f6e98af Mon Sep 17 00:00:00 2001
From: auser1337 <154299690+auser1337@users.noreply.github.com>
Date: Sun, 8 Dec 2024 00:18:12 -0800
Subject: [PATCH] ajm: support for multiple contexts (#1690)

* ajm: support for multiple contexts

* fix sceAjmInitialize
---
 src/core/libraries/ajm/ajm.cpp | 56 ++++++++++++++++++++++++++++------
 1 file changed, 46 insertions(+), 10 deletions(-)

diff --git a/src/core/libraries/ajm/ajm.cpp b/src/core/libraries/ajm/ajm.cpp
index a3728039e..3184fa64f 100644
--- a/src/core/libraries/ajm/ajm.cpp
+++ b/src/core/libraries/ajm/ajm.cpp
@@ -19,7 +19,7 @@ constexpr int ORBIS_AJM_CHANNELMASK_QUAD = 0x0033;
 constexpr int ORBIS_AJM_CHANNELMASK_5POINT1 = 0x060F;
 constexpr int ORBIS_AJM_CHANNELMASK_7POINT1 = 0x063F;
 
-static std::unique_ptr<AjmContext> context{};
+static std::unordered_map<u32, std::unique_ptr<AjmContext>> contexts{};
 
 u32 GetChannelMask(u32 num_channels) {
     switch (num_channels) {
@@ -40,7 +40,13 @@ u32 GetChannelMask(u32 num_channels) {
 
 int PS4_SYSV_ABI sceAjmBatchCancel(const u32 context_id, const u32 batch_id) {
     LOG_INFO(Lib_Ajm, "called context_id = {} batch_id = {}", context_id, batch_id);
-    return context->BatchCancel(batch_id);
+
+    auto it = contexts.find(context_id);
+    if (it == contexts.end()) {
+        return ORBIS_AJM_ERROR_INVALID_CONTEXT;
+    }
+
+    return it->second->BatchCancel(batch_id);
 }
 
 int PS4_SYSV_ABI sceAjmBatchErrorDump() {
@@ -90,14 +96,26 @@ int PS4_SYSV_ABI sceAjmBatchStartBuffer(u32 context_id, u8* p_batch, u32 batch_s
                                         u32* out_batch_id) {
     LOG_TRACE(Lib_Ajm, "called context = {}, batch_size = {:#x}, priority = {}", context_id,
               batch_size, priority);
-    return context->BatchStartBuffer(p_batch, batch_size, priority, batch_error, out_batch_id);
+
+    auto it = contexts.find(context_id);
+    if (it == contexts.end()) {
+        return ORBIS_AJM_ERROR_INVALID_CONTEXT;
+    }
+
+    return it->second->BatchStartBuffer(p_batch, batch_size, priority, batch_error, out_batch_id);
 }
 
 int PS4_SYSV_ABI sceAjmBatchWait(const u32 context_id, const u32 batch_id, const u32 timeout,
                                  AjmBatchError* const batch_error) {
     LOG_TRACE(Lib_Ajm, "called context = {}, batch_id = {}, timeout = {}", context_id, batch_id,
               timeout);
-    return context->BatchWait(batch_id, timeout, batch_error);
+
+    auto it = contexts.find(context_id);
+    if (it == contexts.end()) {
+        return ORBIS_AJM_ERROR_INVALID_CONTEXT;
+    }
+
+    return it->second->BatchWait(batch_id, timeout, batch_error);
 }
 
 int PS4_SYSV_ABI sceAjmDecAt9ParseConfigData() {
@@ -117,12 +135,12 @@ int PS4_SYSV_ABI sceAjmFinalize() {
 
 int PS4_SYSV_ABI sceAjmInitialize(s64 reserved, u32* p_context_id) {
     LOG_INFO(Lib_Ajm, "called reserved = {}", reserved);
-    ASSERT_MSG(context == nullptr, "Multiple contexts are currently unsupported.");
     if (p_context_id == nullptr || reserved != 0) {
         return ORBIS_AJM_ERROR_INVALID_PARAMETER;
     }
-    *p_context_id = 1;
-    context = std::make_unique<AjmContext>();
+    u32 id = contexts.size() + 1;
+    *p_context_id = id;
+    contexts.emplace(id, std::make_unique<AjmContext>());
     return ORBIS_OK;
 }
 
@@ -135,12 +153,24 @@ int PS4_SYSV_ABI sceAjmInstanceCreate(u32 context_id, AjmCodecType codec_type,
                                       AjmInstanceFlags flags, u32* out_instance) {
     LOG_INFO(Lib_Ajm, "called context = {}, codec_type = {}, flags = {:#x}", context_id,
              magic_enum::enum_name(codec_type), flags.raw);
-    return context->InstanceCreate(codec_type, flags, out_instance);
+
+    auto it = contexts.find(context_id);
+    if (it == contexts.end()) {
+        return ORBIS_AJM_ERROR_INVALID_CONTEXT;
+    }
+
+    return it->second->InstanceCreate(codec_type, flags, out_instance);
 }
 
 int PS4_SYSV_ABI sceAjmInstanceDestroy(u32 context_id, u32 instance_id) {
     LOG_INFO(Lib_Ajm, "called context = {}, instance = {}", context_id, instance_id);
-    return context->InstanceDestroy(instance_id);
+
+    auto it = contexts.find(context_id);
+    if (it == contexts.end()) {
+        return ORBIS_AJM_ERROR_INVALID_CONTEXT;
+    }
+
+    return it->second->InstanceDestroy(instance_id);
 }
 
 int PS4_SYSV_ABI sceAjmInstanceExtend() {
@@ -168,7 +198,13 @@ int PS4_SYSV_ABI sceAjmModuleRegister(u32 context_id, AjmCodecType codec_type, s
     if (reserved != 0) {
         return ORBIS_AJM_ERROR_INVALID_PARAMETER;
     }
-    return context->ModuleRegister(codec_type);
+
+    auto it = contexts.find(context_id);
+    if (it == contexts.end()) {
+        return ORBIS_AJM_ERROR_INVALID_CONTEXT;
+    }
+
+    return it->second->ModuleRegister(codec_type);
 }
 
 int PS4_SYSV_ABI sceAjmModuleUnregister() {