[AJM] Track registered codec instance lifecycle (#352)

src/SharpEmu.Libs/Audio/AjmExports.cs

@@ -10,8 +10,29 @@ namespace SharpEmu.Libs.Audio;
 
 public static class AjmExports
 {
-    private static readonly ConcurrentDictionary<uint, byte> Contexts = new();
+    private const int OrbisAjmErrorInvalidContext = unchecked((int)0x80930002);
+    private const int OrbisAjmErrorInvalidInstance = unchecked((int)0x80930003);
+    private const int OrbisAjmErrorInvalidParameter = unchecked((int)0x80930005);
+    private const int OrbisAjmErrorOutOfResources = unchecked((int)0x80930007);
+    private const int OrbisAjmErrorCodecAlreadyRegistered = unchecked((int)0x80930009);
+    private const int OrbisAjmErrorCodecNotRegistered = unchecked((int)0x8093000A);
+    private const int OrbisAjmErrorWrongRevisionFlag = unchecked((int)0x8093000B);
+    private const uint MaxCodecType = 23;
+    private const int MaxInstanceIndex = 0x2FFF;
+    private static readonly ConcurrentDictionary<uint, AjmContextState> Contexts = new();
     private static int _nextContextId;
+
+    private sealed class AjmContextState
+    {
+        public object Gate { get; } = new();
+
+        public HashSet<uint> RegisteredCodecs { get; } = new();
+
+        public Dictionary<uint, uint> InstancesBySlot { get; } = new();
+
+        public int NextInstanceIndex { get; set; }
+    }
+
     public static int AjmInitialize(CpuContext ctx)
     {
         var reserved = ctx[CpuRegister.Rdi];
@@ -29,7 +50,7 @@ public static class AjmExports
             return unchecked((int)0x806A0001);
         }
 
-        Contexts[contextId] = 0;
+        Contexts[contextId] = new AjmContextState();
         if (string.Equals(Environment.GetEnvironmentVariable("SHARPEMU_LOG_AJM"), "1", StringComparison.Ordinal))
         {
             Console.Error.WriteLine(
@@ -62,9 +83,22 @@ public static class AjmExports
         var contextId = unchecked((uint)ctx[CpuRegister.Rdi]);
         var codecType = unchecked((uint)ctx[CpuRegister.Rsi]);
         var reserved = ctx[CpuRegister.Rdx];
-        if (reserved != 0 || !Contexts.ContainsKey(contextId))
+        if (codecType >= MaxCodecType || reserved != 0)
         {
-            return unchecked((int)0x806A0001);
+            return ctx.SetReturn(OrbisAjmErrorInvalidParameter);
+        }
+
+        if (!Contexts.TryGetValue(contextId, out var state))
+        {
+            return ctx.SetReturn(OrbisAjmErrorInvalidContext);
+        }
+
+        lock (state.Gate)
+        {
+            if (!state.RegisteredCodecs.Add(codecType))
+            {
+                return ctx.SetReturn(OrbisAjmErrorCodecAlreadyRegistered);
+            }
         }
 
         if (string.Equals(Environment.GetEnvironmentVariable("SHARPEMU_LOG_AJM"), "1", StringComparison.Ordinal))
@@ -77,6 +111,97 @@ public static class AjmExports
         return 0;
     }
 
+    [SysAbiExport(
+        Nid = "AxoDrINp4J8",
+        ExportName = "sceAjmInstanceCreate",
+        Target = Generation.Gen4 | Generation.Gen5,
+        LibraryName = "libSceAjm")]
+    public static int AjmInstanceCreate(CpuContext ctx)
+    {
+        var contextId = unchecked((uint)ctx[CpuRegister.Rdi]);
+        var codecType = unchecked((uint)ctx[CpuRegister.Rsi]);
+        var flags = ctx[CpuRegister.Rdx];
+        var outputAddress = ctx[CpuRegister.Rcx];
+        if (!Contexts.TryGetValue(contextId, out var state))
+        {
+            return ctx.SetReturn(OrbisAjmErrorInvalidContext);
+        }
+
+        if (codecType >= MaxCodecType || outputAddress == 0)
+        {
+            return ctx.SetReturn(OrbisAjmErrorInvalidParameter);
+        }
+
+        if ((flags & 0x7) == 0)
+        {
+            return ctx.SetReturn(OrbisAjmErrorWrongRevisionFlag);
+        }
+
+        uint instanceId;
+        lock (state.Gate)
+        {
+            if (!state.RegisteredCodecs.Contains(codecType))
+            {
+                return ctx.SetReturn(OrbisAjmErrorCodecNotRegistered);
+            }
+
+            if (state.InstancesBySlot.Count >= MaxInstanceIndex)
+            {
+                return ctx.SetReturn(OrbisAjmErrorOutOfResources);
+            }
+
+            var nextInstanceIndex = state.NextInstanceIndex;
+            uint instanceSlot;
+            do
+            {
+                nextInstanceIndex = nextInstanceIndex % MaxInstanceIndex + 1;
+                instanceSlot = unchecked((uint)nextInstanceIndex);
+            }
+            while (state.InstancesBySlot.ContainsKey(instanceSlot));
+
+            instanceId = (codecType << 14) | instanceSlot;
+            Span<byte> value = stackalloc byte[sizeof(uint)];
+            BinaryPrimitives.WriteUInt32LittleEndian(value, instanceId);
+            if (!ctx.Memory.TryWrite(outputAddress, value))
+            {
+                return ctx.SetReturn(OrbisAjmErrorInvalidParameter);
+            }
+
+            state.NextInstanceIndex = nextInstanceIndex;
+            state.InstancesBySlot.Add(instanceSlot, instanceId);
+        }
+
+        Trace($"instance_create context={contextId} codec={codecType} flags=0x{flags:X} instance=0x{instanceId:X8}");
+        return ctx.SetReturn(0);
+    }
+
+    [SysAbiExport(
+        Nid = "RbLbuKv8zho",
+        ExportName = "sceAjmInstanceDestroy",
+        Target = Generation.Gen4 | Generation.Gen5,
+        LibraryName = "libSceAjm")]
+    public static int AjmInstanceDestroy(CpuContext ctx)
+    {
+        var contextId = unchecked((uint)ctx[CpuRegister.Rdi]);
+        var instanceId = unchecked((uint)ctx[CpuRegister.Rsi]);
+        if (!Contexts.TryGetValue(contextId, out var state))
+        {
+            return ctx.SetReturn(OrbisAjmErrorInvalidContext);
+        }
+
+        var instanceSlot = instanceId & 0x3FFF;
+        lock (state.Gate)
+        {
+            if (instanceSlot == 0 || !state.InstancesBySlot.Remove(instanceSlot))
+            {
+                return ctx.SetReturn(OrbisAjmErrorInvalidInstance);
+            }
+        }
+
+        Trace($"instance_destroy context={contextId} instance=0x{instanceId:X8}");
+        return ctx.SetReturn(0);
+    }
+
     [SysAbiExport(
         Nid = "Wi7DtlLV+KI",
         ExportName = "sceAjmModuleUnregister",
@@ -101,4 +226,18 @@ public static class AjmExports
         ctx[CpuRegister.Rax] = 0;
         return 0;
     }
+
+    internal static void ResetForTests()
+    {
+        Contexts.Clear();
+        Interlocked.Exchange(ref _nextContextId, 0);
+    }
+
+    private static void Trace(string message)
+    {
+        if (string.Equals(Environment.GetEnvironmentVariable("SHARPEMU_LOG_AJM"), "1", StringComparison.Ordinal))
+        {
+            Console.Error.WriteLine($"[LOADER][TRACE] ajm.{message}");
+        }
+    }
 }

tests/SharpEmu.Libs.Tests/Audio/AjmExportsTests.cs

@@ -0,0 +1,198 @@
+// Copyright (C) 2026 SharpEmu Emulator Project
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+using System.Buffers.Binary;
+using SharpEmu.HLE;
+using SharpEmu.Libs.Audio;
+using Xunit;
+
+namespace SharpEmu.Libs.Tests.Audio;
+
+[CollectionDefinition("AjmState", DisableParallelization = true)]
+public sealed class AjmStateCollection
+{
+    public const string Name = "AjmState";
+}
+
+[Collection(AjmStateCollection.Name)]
+public sealed class AjmExportsTests : IDisposable
+{
+    private const int InvalidContext = unchecked((int)0x80930002);
+    private const int InvalidInstance = unchecked((int)0x80930003);
+    private const int InvalidParameter = unchecked((int)0x80930005);
+    private const int CodecAlreadyRegistered = unchecked((int)0x80930009);
+    private const int CodecNotRegistered = unchecked((int)0x8093000A);
+    private const ulong MemoryBase = 0x1_0000_0000;
+    private const ulong ContextAddress = MemoryBase + 0x100;
+    private const ulong InstanceAddress = MemoryBase + 0x200;
+
+    private readonly FakeCpuMemory _memory = new(MemoryBase, 0x1000);
+    private readonly CpuContext _ctx;
+
+    public AjmExportsTests()
+    {
+        AjmExports.ResetForTests();
+        _ctx = new CpuContext(_memory, Generation.Gen5);
+    }
+
+    [Fact]
+    public void InstanceLifecycle_RegisteredCodecCreatesAndDestroysInstance()
+    {
+        var contextId = Initialize();
+
+        Assert.Equal(0, RegisterCodec(contextId, 1));
+        Assert.Equal(0, CreateInstance(contextId, 1, 0x401, InstanceAddress));
+        Assert.Equal(0x4001u, ReadUInt32(InstanceAddress));
+
+        Assert.Equal(0, DestroyInstance(contextId, 0x4001));
+        Assert.Equal(InvalidInstance, DestroyInstance(contextId, 0x4001));
+    }
+
+    [Fact]
+    public void InstanceCreate_UnregisteredCodecDoesNotWriteOutput()
+    {
+        var contextId = Initialize();
+        WriteUInt32(InstanceAddress, 0xCCCCCCCC);
+
+        Assert.Equal(CodecNotRegistered, CreateInstance(contextId, 1, 0x401, InstanceAddress));
+        Assert.Equal(0xCCCCCCCCu, ReadUInt32(InstanceAddress));
+    }
+
+    [Fact]
+    public void InstanceCreate_FaultingOutputDoesNotAdvanceInstanceId()
+    {
+        var contextId = Initialize();
+        Assert.Equal(0, RegisterCodec(contextId, 1));
+
+        Assert.Equal(InvalidParameter, CreateInstance(contextId, 1, 0x401, MemoryBase + 0x1000));
+        Assert.Equal(0, CreateInstance(contextId, 1, 0x401, InstanceAddress));
+        Assert.Equal(0x4001u, ReadUInt32(InstanceAddress));
+        Assert.Equal(0, DestroyInstance(contextId, 0x4001));
+    }
+
+    [Fact]
+    public void ModuleRegister_RejectsDuplicateAndUnknownContext()
+    {
+        var contextId = Initialize();
+
+        Assert.Equal(0, RegisterCodec(contextId, 1));
+        Assert.Equal(CodecAlreadyRegistered, RegisterCodec(contextId, 1));
+        Assert.Equal(InvalidContext, RegisterCodec(contextId + 1, 1));
+    }
+
+    [Fact]
+    public void InstanceDestroy_RejectsUnknownContextAndSlot()
+    {
+        var contextId = Initialize();
+
+        Assert.Equal(InvalidContext, DestroyInstance(contextId + 1, 1));
+        Assert.Equal(InvalidInstance, DestroyInstance(contextId, 0));
+        Assert.Equal(InvalidInstance, DestroyInstance(contextId, 1));
+    }
+
+    [Fact]
+    public void InstanceDestroy_ResolvesInstanceByMaskedSlot()
+    {
+        var contextId = Initialize();
+        Assert.Equal(0, RegisterCodec(contextId, 1));
+        Assert.Equal(0, CreateInstance(contextId, 1, 0x401, InstanceAddress));
+
+        Assert.Equal(0, DestroyInstance(contextId, 0x8001));
+        Assert.Equal(InvalidInstance, DestroyInstance(contextId, 0x4001));
+    }
+
+    [Fact]
+    public void ConcurrentInstanceCreates_ProduceUniqueLiveIds()
+    {
+        const int count = 32;
+        var contextId = Initialize();
+        Assert.Equal(0, RegisterCodec(contextId, 1));
+
+        var results = Enumerable.Range(0, count)
+            .AsParallel()
+            .Select(index =>
+            {
+                var outputAddress = MemoryBase + 0x300 + unchecked((ulong)(index * sizeof(uint)));
+                var context = new CpuContext(_memory, Generation.Gen5)
+                {
+                    [CpuRegister.Rdi] = contextId,
+                    [CpuRegister.Rsi] = 1,
+                    [CpuRegister.Rdx] = 0x401,
+                    [CpuRegister.Rcx] = outputAddress,
+                };
+                var result = AjmExports.AjmInstanceCreate(context);
+                return (result, instanceId: ReadUInt32(outputAddress));
+            })
+            .ToArray();
+
+        Assert.All(results, result => Assert.Equal(0, result.result));
+        Assert.Equal(count, results.Select(result => result.instanceId).Distinct().Count());
+        Assert.All(results, result => Assert.Equal(0, DestroyInstance(contextId, result.instanceId)));
+    }
+
+    [Fact]
+    public void InstanceLifecycleExports_RegisterForBothGenerations()
+    {
+        foreach (var generation in new[] { Generation.Gen4, Generation.Gen5 })
+        {
+            var manager = new ModuleManager();
+            manager.RegisterExports(SharpEmu.Generated.SysAbiExportRegistry.CreateExports(generation));
+
+            Assert.True(manager.TryGetExport("AxoDrINp4J8", out var create));
+            Assert.Equal("sceAjmInstanceCreate", create.Name);
+            Assert.True(manager.TryGetExport("RbLbuKv8zho", out var destroy));
+            Assert.Equal("sceAjmInstanceDestroy", destroy.Name);
+        }
+    }
+
+    public void Dispose()
+    {
+        AjmExports.ResetForTests();
+    }
+
+    private uint Initialize()
+    {
+        _ctx[CpuRegister.Rdi] = 0;
+        _ctx[CpuRegister.Rsi] = ContextAddress;
+        Assert.Equal(0, AjmExports.AjmInitialize(_ctx));
+        return ReadUInt32(ContextAddress);
+    }
+
+    private int RegisterCodec(uint contextId, uint codecType)
+    {
+        _ctx[CpuRegister.Rdi] = contextId;
+        _ctx[CpuRegister.Rsi] = codecType;
+        _ctx[CpuRegister.Rdx] = 0;
+        return AjmExports.AjmModuleRegister(_ctx);
+    }
+
+    private int CreateInstance(uint contextId, uint codecType, ulong flags, ulong outputAddress)
+    {
+        _ctx[CpuRegister.Rdi] = contextId;
+        _ctx[CpuRegister.Rsi] = codecType;
+        _ctx[CpuRegister.Rdx] = flags;
+        _ctx[CpuRegister.Rcx] = outputAddress;
+        return AjmExports.AjmInstanceCreate(_ctx);
+    }
+
+    private int DestroyInstance(uint contextId, uint instanceId)
+    {
+        _ctx[CpuRegister.Rdi] = contextId;
+        _ctx[CpuRegister.Rsi] = instanceId;
+        return AjmExports.AjmInstanceDestroy(_ctx);
+    }
+
+    private uint ReadUInt32(ulong address)
+    {
+        Span<byte> value = stackalloc byte[sizeof(uint)];
+        Assert.True(_memory.TryRead(address, value));
+        return BinaryPrimitives.ReadUInt32LittleEndian(value);
+    }
+
+    private void WriteUInt32(ulong address, uint value)
+    {
+        Span<byte> bytes = stackalloc byte[sizeof(uint)];
+        BinaryPrimitives.WriteUInt32LittleEndian(bytes, value);
+        Assert.True(_memory.TryWrite(address, bytes));
+    }
+}