From f2e833b5c474c8f5e42d413c8994174bc6fccd16 Mon Sep 17 00:00:00 2001
From: Pokechu22 <Pokechu022@gmail.com>
Date: Wed, 22 Jun 2022 13:20:00 -0700
Subject: [PATCH] DSPHLE: Eliminate global state in AX uCode

This also increases accuracy as to when specific mail is allowed, and correctly handles masking of the 0xCDD1 mails.
---
 Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp | 105 +++++++++++++----------
 Source/Core/Core/HW/DSPHLE/UCodes/AX.h   |   9 ++
 2 files changed, 71 insertions(+), 43 deletions(-)

diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp b/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp
index c53d7a8e8e..7a7d9e5d39 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp
@@ -653,53 +653,72 @@ void AXUCode::SendAUXAndMix(u32 main_auxa_up, u32 auxb_s_up, u32 main_l_dl, u32
 
 void AXUCode::HandleMail(u32 mail)
 {
-  // Indicates if the next message is a command list address.
-  static bool next_is_cmdlist = false;
-  static u16 cmdlist_size = 0;
-
-  bool set_next_is_cmdlist = false;
-
-  if (next_is_cmdlist)
+  if (m_upload_setup_in_progress)
   {
-    CopyCmdList(mail, cmdlist_size);
+    PrepareBootUCode(mail);
+    return;
+  }
+
+  switch (m_mail_state)
+  {
+  case MailState::WaitingForCmdListSize:
+    if ((mail & MAIL_CMDLIST_MASK) == MAIL_CMDLIST)
+    {
+      // A command list address is going to be sent next.
+      m_cmdlist_size = static_cast<u16>(mail & ~MAIL_CMDLIST_MASK);
+      m_mail_state = MailState::WaitingForCmdListAddress;
+    }
+    else
+    {
+      ERROR_LOG_FMT(DSPHLE, "Unknown mail sent to AX::HandleMail; expected command list: {:08x}",
+                    mail);
+    }
+    break;
+
+  case MailState::WaitingForCmdListAddress:
+    CopyCmdList(mail, m_cmdlist_size);
     HandleCommandList();
     m_cmdlist_size = 0;
     SignalWorkEnd();
-  }
-  else if (m_upload_setup_in_progress)
-  {
-    PrepareBootUCode(mail);
-  }
-  else if (mail == MAIL_RESUME)
-  {
-    // Acknowledge the resume request
-    m_mail_handler.PushMail(DSP_RESUME, true);
-  }
-  else if (mail == MAIL_NEW_UCODE)
-  {
-    m_upload_setup_in_progress = true;
-  }
-  else if (mail == MAIL_RESET)
-  {
-    m_dsphle->SetUCode(UCODE_ROM);
-  }
-  else if (mail == MAIL_CONTINUE)
-  {
-    // We don't have to do anything here - the CPU does not wait for a ACK
-    // and sends a cmdlist mail just after.
-  }
-  else if ((mail & MAIL_CMDLIST_MASK) == MAIL_CMDLIST)
-  {
-    // A command list address is going to be sent next.
-    set_next_is_cmdlist = true;
-    cmdlist_size = (u16)(mail & ~MAIL_CMDLIST_MASK);
-  }
-  else
-  {
-    ERROR_LOG_FMT(DSPHLE, "Unknown mail sent to AX::HandleMail: {:08x}", mail);
-  }
+    m_mail_state = MailState::WaitingForNextTask;
+    break;
 
-  next_is_cmdlist = set_next_is_cmdlist;
+  case MailState::WaitingForNextTask:
+    if ((mail & TASK_MAIL_MASK) != TASK_MAIL_TO_DSP)
+    {
+      WARN_LOG_FMT(DSPHLE, "Rendering task without prefix CDD1: {:08x}", mail);
+      mail = TASK_MAIL_TO_DSP | (mail & ~TASK_MAIL_MASK);
+      // The actual uCode does not check for the CDD1 prefix.
+    }
+
+    switch (mail)
+    {
+    case MAIL_RESUME:
+      // Acknowledge the resume request
+      m_mail_handler.PushMail(DSP_RESUME, true);
+      m_mail_state = MailState::WaitingForCmdListSize;
+      break;
+
+    case MAIL_NEW_UCODE:
+      m_upload_setup_in_progress = true;
+      break;
+
+    case MAIL_RESET:
+      m_dsphle->SetUCode(UCODE_ROM);
+      break;
+
+    case MAIL_CONTINUE:
+      // We don't have to do anything here - the CPU does not wait for a ACK
+      // and sends a cmdlist mail just after.
+      m_mail_state = MailState::WaitingForCmdListSize;
+      break;
+
+    default:
+      WARN_LOG_FMT(DSPHLE, "Unknown task mail: {:08x}", mail);
+      break;
+    }
+    break;
+  }
 }
 
 void AXUCode::CopyCmdList(u32 addr, u16 size)
@@ -712,7 +731,6 @@ void AXUCode::CopyCmdList(u32 addr, u16 size)
 
   for (u32 i = 0; i < size; ++i, addr += 2)
     m_cmdlist[i] = HLEMemory_Read_U16(addr);
-  m_cmdlist_size = size;
 }
 
 void AXUCode::Update()
@@ -728,6 +746,7 @@ void AXUCode::DoAXState(PointerWrap& p)
 {
   p.Do(m_cmdlist);
   p.Do(m_cmdlist_size);
+  p.Do(m_mail_state);
 
   p.Do(m_samples_main_left);
   p.Do(m_samples_main_right);
diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AX.h b/Source/Core/Core/HW/DSPHLE/UCodes/AX.h
index c51cdb34b9..d63659a995 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AX.h
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AX.h
@@ -200,5 +200,14 @@ private:
     CMD_COMPRESSOR = 0x12,
     CMD_SEND_AUX_AND_MIX = 0x13,
   };
+
+  enum class MailState
+  {
+    WaitingForCmdListSize,
+    WaitingForCmdListAddress,
+    WaitingForNextTask,
+  };
+
+  MailState m_mail_state = MailState::WaitingForCmdListSize;
 };
 }  // namespace DSP::HLE