DSPHLE/AXWii: fix wiimote audio in multiple games

Three bugs specific to older Wii games:
- The size difference between high-pass and biquad filter was not
  accounted for, causing wiimote related fields to be corrupted.
- Wiimote sample buffer pointers were advanced by 32 samples per
  millisecond instead of 6 samples. Usually hidden by the first bug.
- PB updates on Wii were being byte-swapped twice, but I've not actually
  found any Wii games that make use of PB updates.

This fixes wiimote audio in at least the following games:
- Excite Truck
- Ice Age 2: The Meltdown
- Kororinpa: Marble Mania
- Rapala Tournament Fishing
- Shrek the Third
- Super Monkey Ball: Banana Blitz
- Tiger Woods PGA Tour 07
- WarioWare: Smooth Moves (issue 11725)
- Wing Island

6 files changed, 195 insertions, 154 deletions

diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp b/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp
index 84705edd7c..426a0517d0 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AX.cpp
@@ -412,6 +412,56 @@ void AXUCode::DownloadAndMixWithVolume(u32 addr, u16 vol_main, u16 vol_auxa, u16
   }
 }
 
+// Determines if this version of the UCode has a PBLowPassFilter in its AXPB layout.
+static bool HasLpf(u32 crc)
+{
+  return crc != 0x4E8A8B21;
+}
+
+// Read a PB from MRAM/ARAM
+void AXUCode::ReadPB(Memory::MemoryManager& memory, u32 addr, AXPB& pb)
+{
+  if (HasLpf(m_crc))
+  {
+    u16* dst = (u16*)&pb;
+    memory.CopyFromEmuSwapped<u16>(dst, addr, sizeof(pb));
+  }
+  else
+  {
+    // Skip lpf field.
+
+    char* dst = (char*)&pb;
+
+    constexpr size_t lpf_off = offsetof(AXPB, lpf);
+    constexpr size_t lc_off = offsetof(AXPB, loop_counter);
+
+    memory.CopyFromEmuSwapped<u16>((u16*)dst, addr, lpf_off);
+    memset(dst + lpf_off, 0, lc_off - lpf_off);
+    memory.CopyFromEmuSwapped<u16>((u16*)(dst + lc_off), addr + lpf_off, sizeof(pb) - lc_off);
+  }
+}
+
+// Write a PB back to MRAM/ARAM
+void AXUCode::WritePB(Memory::MemoryManager& memory, u32 addr, const AXPB& pb)
+{
+  if (HasLpf(m_crc))
+  {
+    const u16* src = (const u16*)&pb;
+    memory.CopyToEmuSwapped<u16>(addr, src, sizeof(pb));
+  }
+  else
+  {
+    // We skip lpf in this layout.
+
+    const char* src = (const char*)&pb;
+    constexpr size_t lpf_off = offsetof(AXPB, lpf);
+    constexpr size_t lc_off = offsetof(AXPB, loop_counter);
+
+    memory.CopyToEmuSwapped<u16>(addr, (const u16*)src, lpf_off);
+    memory.CopyToEmuSwapped<u16>(addr + lpf_off, (const u16*)(src + lc_off), sizeof(pb) - lc_off);
+  }
+}
+
 void AXUCode::ProcessPBList(u32 pb_addr)
 {
   // Samples per millisecond. In theory DSP sampling rate can be changed from
@@ -427,10 +477,9 @@ void AXUCode::ProcessPBList(u32 pb_addr)
                           m_samples_auxA_left, m_samples_auxA_right, m_samples_auxA_surround,
                           m_samples_auxB_left, m_samples_auxB_right, m_samples_auxB_surround}};
 
-    ReadPB(memory, pb_addr, pb, m_crc);
+    ReadPB(memory, pb_addr, pb);
 
-    u32 updates_addr = HILO_TO_32(pb.updates.data);
-    u16* updates = (u16*)HLEMemory_Get_Pointer(memory, updates_addr);
+    PBUpdateData updates = LoadPBUpdates(memory, pb);
 
     for (int curr_ms = 0; curr_ms < 5; ++curr_ms)
     {
@@ -445,7 +494,7 @@ void AXUCode::ProcessPBList(u32 pb_addr)
         ptr += spms;
     }
 
-    WritePB(memory, pb_addr, pb, m_crc);
+    WritePB(memory, pb_addr, pb);
     pb_addr = HILO_TO_32(pb.next_pb);
   }
 }
diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AX.h b/Source/Core/Core/HW/DSPHLE/UCodes/AX.h
index 108a98e351..138db3aaef 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AX.h
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AX.h
@@ -30,6 +30,7 @@ class Accelerator;
 
 namespace DSP::HLE
 {
+struct AXPB;
 class DSPHLE;
 
 // We can't directly use the mixer_control field from the PB because it does
@@ -125,27 +126,6 @@ protected:
   // versions of AX.
   AXMixControl ConvertMixerControl(u32 mixer_control);
 
-  // Apply updates to a PB. Generic, used in AX GC and AX Wii.
-  template <typename PBType>
-  void ApplyUpdatesForMs(int curr_ms, PBType& pb, u16* num_updates, u16* updates)
-  {
-    auto pb_mem = Common::BitCastToArray<u16>(pb);
-
-    u32 start_idx = 0;
-    for (int i = 0; i < curr_ms; ++i)
-      start_idx += num_updates[i];
-
-    for (u32 i = start_idx; i < start_idx + num_updates[curr_ms]; ++i)
-    {
-      u16 update_off = Common::swap16(updates[2 * i]);
-      u16 update_val = Common::swap16(updates[2 * i + 1]);
-
-      pb_mem[update_off] = update_val;
-    }
-
-    pb = std::bit_cast<PBType>(pb_mem);
-  }
-
   virtual void HandleCommandList();
   void SignalWorkEnd();
 
@@ -195,6 +175,9 @@ protected:
   void DoAXState(PointerWrap& p);
 
 private:
+  void ReadPB(Memory::MemoryManager& memory, u32 addr, AXPB& pb);
+  void WritePB(Memory::MemoryManager& memory, u32 addr, const AXPB& pb);
+
   enum CmdType
   {
     CMD_SETUP = 0x00,
diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AXStructs.h b/Source/Core/Core/HW/DSPHLE/UCodes/AXStructs.h
index f819c066aa..701d692e1c 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AXStructs.h
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AXStructs.h
@@ -59,6 +59,13 @@ struct PBInitialTimeDelay
   u16 targetRight;
 };
 
+struct PBUpdate
+{
+  u16 pb_offset;
+  u16 new_value;
+};
+using PBUpdateData = std::array<PBUpdate, 32>;
+
 // Update data - read these each 1ms subframe and use them!
 // It seems that to provide higher time precisions for MIDI events, some games
 // use this thing to update the parameter blocks per 1ms sub-block (a block is 5ms).
@@ -66,7 +73,15 @@ struct PBInitialTimeDelay
 struct PBUpdates
 {
   u16 num_updates[5];
-  u16 data_hi;  // These point to main RAM. Not sure about the structure of the data.
+  u16 data_hi;  // These point to main RAM.
+  u16 data_lo;
+};
+
+// Same for Wii, where frames are only 3 ms.
+struct PBUpdatesWii
+{
+  u16 num_updates[3];
+  u16 data_hi;
   u16 data_lo;
 };
 
@@ -213,6 +228,12 @@ struct AXPB
   u16 padding[24];
 };
 
+struct PBHighPassFilter
+{
+  u16 on;
+  u16 unk[3];
+};
+
 struct PBBiquadFilter
 {
   u16 on;
@@ -252,13 +273,18 @@ struct AXPBWii
   PBMixerWii mixer;
   PBInitialTimeDelay initial_time_delay;
   PBDpopWii dpop;
+  PBUpdatesWii updates;  // Not present in all versions of the struct.
   PBVolumeEnvelope vol_env;
   PBAudioAddr audio_addr;
   PBADPCMInfo adpcm;
   PBSampleRateConverter src;
   PBADPCMLoopInfo adpcm_loop_info;
   PBLowPassFilter lpf;
-  PBBiquadFilter biquad;
+  union
+  {
+    PBHighPassFilter hpf;
+    PBBiquadFilter biquad;
+  };
 
   // WIIMOTE :D
   u16 remote;
@@ -269,7 +295,7 @@ struct AXPBWii
   PBSampleRateConverterWM remote_src;
   PBInfImpulseResponseWM remote_iir;
 
-  u16 pad[12];  // align us, captain! (32B)
+  u16 pad[2];  // align us, captain! (32B)
 };
 
 // TODO: All these enums have changed a lot for Wii
diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AXVoice.h b/Source/Core/Core/HW/DSPHLE/UCodes/AXVoice.h
index 677e6a7989..ee85f8fb4b 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AXVoice.h
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AXVoice.h
@@ -12,6 +12,7 @@
 #endif
 
 #include <algorithm>
+#include <bit>
 #include <functional>
 #include <memory>
 
@@ -47,6 +48,34 @@ namespace
 // Useful macro to convert xxx_hi + xxx_lo to xxx for 32 bits.
 #define HILO_TO_32(name) ((u32(name##_hi) << 16) | name##_lo)
 
+PBUpdateData LoadPBUpdates(Memory::MemoryManager& memory, const PB_TYPE& pb)
+{
+  PBUpdateData updates;
+  u32 updates_addr = HILO_TO_32(pb.updates.data);
+  memory.CopyFromEmuSwapped((u16*)updates.data(), updates_addr, sizeof(updates));
+  return updates;
+}
+
+// Apply updates to a PB.
+void ApplyUpdatesForMs(int curr_ms, PB_TYPE& pb, u16* num_updates, const PBUpdateData& updates)
+{
+  auto pb_mem = Common::BitCastToArray<u16>(pb);
+
+  u32 start_idx = 0;
+  for (int i = 0; i < curr_ms; ++i)
+    start_idx += num_updates[i];
+
+  for (u32 i = start_idx; i < start_idx + num_updates[curr_ms]; ++i)
+  {
+    u16 update_off = updates[i].pb_offset;
+    u16 update_val = updates[i].new_value;
+
+    pb_mem[update_off] = update_val;
+  }
+
+  pb = std::bit_cast<PB_TYPE>(pb_mem);
+}
+
 // Used to pass a large amount of buffers to the mixing function.
 union AXBuffers
 {
@@ -83,69 +112,14 @@ union AXBuffers
 #ifdef AX_GC
   int* ptrs[9];
 #else
-  int* ptrs[20];
+  struct
+  {
+    int* regular_ptrs[12];
+    int* wiimote_ptrs[8];
+  };
 #endif
 };
 
-// Determines if this version of the UCode has a PBLowPassFilter in its AXPB layout.
-bool HasLpf(u32 crc)
-{
-  switch (crc)
-  {
-  case 0x4E8A8B21:
-    return false;
-  default:
-    return true;
-  }
-}
-
-// Read a PB from MRAM/ARAM
-void ReadPB(Memory::MemoryManager& memory, u32 addr, PB_TYPE& pb, u32 crc)
-{
-  if (HasLpf(crc))
-  {
-    u16* dst = (u16*)&pb;
-    memory.CopyFromEmuSwapped<u16>(dst, addr, sizeof(pb));
-  }
-  else
-  {
-    // The below is a terrible hack in order to support two different AXPB layouts.
-    // We skip lpf in this layout.
-
-    char* dst = (char*)&pb;
-
-    constexpr size_t lpf_off = offsetof(AXPB, lpf);
-    constexpr size_t lc_off = offsetof(AXPB, loop_counter);
-
-    memory.CopyFromEmuSwapped<u16>((u16*)dst, addr, lpf_off);
-    memset(dst + lpf_off, 0, lc_off - lpf_off);
-    memory.CopyFromEmuSwapped<u16>((u16*)(dst + lc_off), addr + lpf_off, sizeof(pb) - lc_off);
-  }
-}
-
-// Write a PB back to MRAM/ARAM
-void WritePB(Memory::MemoryManager& memory, u32 addr, const PB_TYPE& pb, u32 crc)
-{
-  if (HasLpf(crc))
-  {
-    const u16* src = (const u16*)&pb;
-    memory.CopyToEmuSwapped<u16>(addr, src, sizeof(pb));
-  }
-  else
-  {
-    // The below is a terrible hack in order to support two different AXPB layouts.
-    // We skip lpf in this layout.
-
-    const char* src = (const char*)&pb;
-
-    constexpr size_t lpf_off = offsetof(AXPB, lpf);
-    constexpr size_t lc_off = offsetof(AXPB, loop_counter);
-
-    memory.CopyToEmuSwapped<u16>(addr, (const u16*)src, lpf_off);
-    memory.CopyToEmuSwapped<u16>(addr + lpf_off, (const u16*)(src + lc_off), sizeof(pb) - lc_off);
-  }
-}
-
 // Simulated accelerator state.
 class HLEAccelerator final : public Accelerator
 {
diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.cpp b/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.cpp
index e41f13a5ef..356136bb52 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.cpp
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.cpp
@@ -360,64 +360,75 @@ void AXWiiUCode::GenerateVolumeRamp(u16* output, u16 vol1, u16 vol2, size_t nval
   }
 }
 
-bool AXWiiUCode::ExtractUpdatesFields(AXPBWii& pb, u16* num_updates, u16* updates,
-                                      u32* updates_addr)
+void AXWiiUCode::ReadPB(Memory::MemoryManager& memory, u32 addr, AXPBWii& pb)
 {
-  auto pb_mem = Common::BitCastToArray<u16>(pb);
-
-  if (!m_old_axwii)
-    return false;
-
-  // Copy the num_updates field.
-  memcpy(num_updates, &pb_mem[41], 6);
-
-  // Get the address of the updates data
-  u16 addr_hi = pb_mem[44];
-  u16 addr_lo = pb_mem[45];
-  u32 addr = HILO_TO_32(addr);
-  auto& memory = m_dsphle->GetSystem().GetMemory();
-  u16* ptr = (u16*)HLEMemory_Get_Pointer(memory, addr);
-
-  *updates_addr = addr;
-
-  // Copy the updates data and change the offset to match a PB without
-  // updates data.
-  u32 updates_count = num_updates[0] + num_updates[1] + num_updates[2];
-  for (u32 i = 0; i < updates_count; ++i)
+  // The Wii PB memory layout changed twice.
+  // For HLE, we use the largest struct version.
+  char* dst = (char*)&pb;
+  constexpr size_t updates_begin = offsetof(AXPBWii, updates);
+  constexpr size_t updates_end = offsetof(AXPBWii, updates) + sizeof(PBUpdatesWii);
+  constexpr size_t gap_begin = offsetof(AXPBWii, hpf) + sizeof(PBHighPassFilter);
+  constexpr size_t gap_end = offsetof(AXPBWii, biquad) + sizeof(PBBiquadFilter);
+  switch (m_crc)
   {
-    u16 update_off = Common::swap16(ptr[2 * i]);
-    u16 update_val = Common::swap16(ptr[2 * i + 1]);
-
-    if (update_off > 45)
-      update_off -= 5;
-
-    updates[2 * i] = update_off;
-    updates[2 * i + 1] = update_val;
+  case 0x7699af32:
+  case 0xfa450138:
+    // The hpf field is a bit smaller than the biquad. Skip the difference.
+    memory.CopyFromEmuSwapped<u16>((u16*)dst, addr, gap_begin);
+    memset(dst + gap_begin, 0, gap_end - gap_begin);
+    memory.CopyFromEmuSwapped<u16>((u16*)(dst + gap_end), addr + gap_begin, sizeof(pb) - gap_end);
+    break;
+  case 0xd9c4bf34:
+  case 0xadbc06bd:
+    // Skip updates field and skip gap after hpf.
+    memory.CopyFromEmuSwapped<u16>((u16*)dst, addr, updates_begin);
+    memset(dst + updates_begin, 0, sizeof(PBUpdatesWii));
+    memory.CopyFromEmuSwapped<u16>((u16*)(dst + updates_end), addr + updates_begin,
+                                   gap_begin - updates_end);
+    memset(dst + gap_begin, 0, gap_end - gap_begin);
+    memory.CopyFromEmuSwapped<u16>((u16*)(dst + gap_end), addr + gap_begin, sizeof(pb) - gap_end);
+    break;
+  case 0x347112ba:
+  case 0x4cc52064:
+    // Just skip updates field.
+    memory.CopyFromEmuSwapped<u16>((u16*)dst, addr, updates_begin);
+    memset(dst + updates_begin, 0, sizeof(PBUpdatesWii));
+    memory.CopyFromEmuSwapped<u16>((u16*)(dst + updates_end), addr + updates_begin,
+                                   sizeof(pb) - updates_end);
+    break;
   }
-
-  // Remove the updates data from the PB
-  memmove(&pb_mem[41], &pb_mem[46], sizeof(pb) - 2 * 46);
-
-  pb = std::bit_cast<AXPBWii>(pb_mem);
-
-  return true;
 }
 
-void AXWiiUCode::ReinjectUpdatesFields(AXPBWii& pb, u16* num_updates, u32 updates_addr)
+void AXWiiUCode::WritePB(Memory::MemoryManager& memory, u32 addr, const AXPBWii& pb)
 {
-  auto pb_mem = Common::BitCastToArray<u16>(pb);
-
-  // Make some space
-  memmove(&pb_mem[46], &pb_mem[41], sizeof(pb) - 2 * 46);
-
-  // Reinsert previous values
-  pb_mem[41] = num_updates[0];
-  pb_mem[42] = num_updates[1];
-  pb_mem[43] = num_updates[2];
-  pb_mem[44] = updates_addr >> 16;
-  pb_mem[45] = updates_addr & 0xFFFF;
-
-  pb = std::bit_cast<AXPBWii>(pb_mem);
+  const char* src = (const char*)&pb;
+  constexpr size_t updates_begin = offsetof(AXPBWii, updates);
+  constexpr size_t updates_end = offsetof(AXPBWii, updates) + sizeof(PBUpdatesWii);
+  constexpr size_t gap_begin = offsetof(AXPBWii, hpf) + sizeof(PBHighPassFilter);
+  constexpr size_t gap_end = offsetof(AXPBWii, biquad) + sizeof(PBBiquadFilter);
+  switch (m_crc)
+  {
+  case 0x7699af32:
+  case 0xfa450138:
+    memory.CopyToEmuSwapped<u16>(addr, (const u16*)src, gap_begin);
+    memory.CopyToEmuSwapped<u16>(addr + gap_begin, (const u16*)(src + gap_end),
+                                 sizeof(pb) - gap_end);
+    break;
+  case 0xd9c4bf34:
+  case 0xadbc06bd:
+    memory.CopyToEmuSwapped<u16>(addr, (const u16*)src, updates_begin);
+    memory.CopyToEmuSwapped<u16>(addr + updates_begin, (const u16*)(src + updates_end),
+                                 sizeof(PBUpdatesWii));
+    memory.CopyToEmuSwapped<u16>(addr + gap_begin, (const u16*)(src + gap_end),
+                                 sizeof(pb) - gap_end);
+    break;
+  case 0x347112ba:
+  case 0x4cc52064:
+    memory.CopyToEmuSwapped<u16>(addr, (const u16*)src, updates_begin);
+    memory.CopyToEmuSwapped<u16>(addr + updates_begin, (const u16*)(src + updates_end),
+                                 sizeof(pb) - updates_end);
+    break;
+  }
 }
 
 void AXWiiUCode::ProcessPBList(u32 pb_addr)
@@ -439,25 +450,25 @@ void AXWiiUCode::ProcessPBList(u32 pb_addr)
                           m_samples_aux1,      m_samples_wm2,        m_samples_aux2,
                           m_samples_wm3,       m_samples_aux3}};
 
-    ReadPB(memory, pb_addr, pb, m_crc);
+    ReadPB(memory, pb_addr, pb);
 
-    u16 num_updates[3];
-    u16 updates[1024];
-    u32 updates_addr;
-    if (ExtractUpdatesFields(pb, num_updates, updates, &updates_addr))
+    if (m_old_axwii &&
+        (pb.updates.num_updates[0] | pb.updates.num_updates[1] | pb.updates.num_updates[2]))
     {
+      PBUpdateData updates = LoadPBUpdates(memory, pb);
       for (int curr_ms = 0; curr_ms < 3; ++curr_ms)
       {
-        ApplyUpdatesForMs(curr_ms, pb, num_updates, updates);
+        ApplyUpdatesForMs(curr_ms, pb, pb.updates.num_updates, updates);
         ProcessVoice(static_cast<HLEAccelerator*>(m_accelerator.get()), pb, buffers, spms,
                      ConvertMixerControl(HILO_TO_32(pb.mixer_control)),
                      m_coeffs_checksum ? m_coeffs.data() : nullptr, m_new_filter);
 
         // Forward the buffers
-        for (auto& ptr : buffers.ptrs)
+        for (auto& ptr : buffers.regular_ptrs)
           ptr += spms;
+        for (auto& ptr : buffers.wiimote_ptrs)
+          ptr += 6;
       }
-      ReinjectUpdatesFields(pb, num_updates, updates_addr);
     }
     else
     {
@@ -466,7 +477,7 @@ void AXWiiUCode::ProcessPBList(u32 pb_addr)
                    m_coeffs_checksum ? m_coeffs.data() : nullptr, m_new_filter);
     }
 
-    WritePB(memory, pb_addr, pb, m_crc);
+    WritePB(memory, pb_addr, pb);
     pb_addr = HILO_TO_32(pb.next_pb);
   }
 }
diff --git a/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.h b/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.h
index 85d3300384..a281734090 100644
--- a/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.h
+++ b/Source/Core/Core/HW/DSPHLE/UCodes/AXWii.h
@@ -46,11 +46,6 @@ protected:
   u16 m_last_main_volume = 0;
   u16 m_last_aux_volumes[3]{};
 
-  // If needed, extract the updates related fields from a PB. We need to
-  // reinject them afterwards so that the correct PB typs is written to RAM.
-  bool ExtractUpdatesFields(AXPBWii& pb, u16* num_updates, u16* updates, u32* updates_addr);
-  void ReinjectUpdatesFields(AXPBWii& pb, u16* num_updates, u32 updates_addr);
-
   // Convert a mixer_control bitfield to our internal representation for that
   // value. Required because that bitfield has a different meaning in some
   // versions of AX.
@@ -73,6 +68,9 @@ protected:
   void OutputWMSamples(u32* addresses);  // 4 addresses
 
 private:
+  void ReadPB(Memory::MemoryManager& memory, u32 addr, AXPBWii& pb);
+  void WritePB(Memory::MemoryManager& memory, u32 addr, const AXPBWii& pb);
+
   enum CmdType
   {
     CMD_SETUP = 0x00,