WIA: Implement multithreaded compression

Source/Core/DiscIO/WIABlob.cpp

@@ -11,6 +11,7 @@
 #include <limits>
 #include <map>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <utility>
 
@@ -30,6 +31,7 @@
 
 #include "DiscIO/Blob.h"
 #include "DiscIO/DiscExtractor.h"
+#include "DiscIO/MultithreadedCompressor.h"
 #include "DiscIO/Volume.h"
 #include "DiscIO/VolumeWii.h"
 #include "DiscIO/WiiEncryptionCache.h"
@@ -909,27 +911,36 @@ WIAFileReader::LZMACompressor::LZMACompressor(bool lzma2, int compression_level,
 
   if (!lzma2)
   {
+    if (compressor_data_size_out)
       *compressor_data_size_out = 5;
 
+    if (compressor_data_out)
+    {
       ASSERT(m_options.lc < 9);
       ASSERT(m_options.lp < 5);
       ASSERT(m_options.pb < 5);
-    compressor_data_out[0] = static_cast<u8>((m_options.pb * 5 + m_options.lp) * 9 + m_options.lc);
+      compressor_data_out[0] =
+          static_cast<u8>((m_options.pb * 5 + m_options.lp) * 9 + m_options.lc);
 
       // The dictionary size is stored as a 32-bit little endian unsigned integer
       static_assert(sizeof(m_options.dict_size) == sizeof(u32));
       std::memcpy(compressor_data_out + 1, &m_options.dict_size, sizeof(u32));
     }
+  }
   else
   {
+    if (compressor_data_size_out)
       *compressor_data_size_out = 1;
 
+    if (compressor_data_out)
+    {
       u8 encoded_dict_size = 0;
       while (encoded_dict_size < 40 && m_options.dict_size > LZMA2DictionarySize(encoded_dict_size))
         ++encoded_dict_size;
 
       compressor_data_out[0] = encoded_dict_size;
     }
+  }
 
   m_filters[0].id = lzma2 ? LZMA_FILTER_LZMA2 : LZMA_FILTER_LZMA1;
   m_filters[0].options = &m_options;
@@ -1246,7 +1257,7 @@ WIAFileReader::PartitionDataEntry WIAFileReader::CreatePartitionDataEntry(
                             Common::swap32(group_index), Common::swap32(groups)};
 }
 
-WIAFileReader::ConversionResult WIAFileReader::SetUpDataEntriesForWriting(
+ConversionResultCode WIAFileReader::SetUpDataEntriesForWriting(
     const VolumeDisc* volume, int chunk_size, u64 iso_size, u32* total_groups,
     std::vector<PartitionEntry>* partition_entries, std::vector<RawDataEntry>* raw_data_entries,
     std::vector<DataEntry>* data_entries)
@@ -1299,7 +1310,7 @@ WIAFileReader::ConversionResult WIAFileReader::SetUpDataEntriesForWriting(
         volume->ReadSwappedAndShifted(partition.offset + 0x2bc, PARTITION_NONE);
 
     if (!data_offset || !data_size)
-      return ConversionResult::ReadFailed;
+      return ConversionResultCode::ReadFailed;
 
     const u64 data_start = partition.offset + *data_offset;
     const u64 data_end = data_start + *data_size;
@@ -1326,11 +1337,11 @@ WIAFileReader::ConversionResult WIAFileReader::SetUpDataEntriesForWriting(
     const std::optional<u64> fst_size = GetFSTSize(*volume, partition);
 
     if (!fst_offset || !fst_size)
-      return ConversionResult::ReadFailed;
+      return ConversionResultCode::ReadFailed;
 
     const IOS::ES::TicketReader& ticket = volume->GetTicket(partition);
     if (!ticket.IsValid())
-      return ConversionResult::ReadFailed;
+      return ConversionResultCode::ReadFailed;
 
     add_raw_data_entry(last_partition_end_offset, partition.offset - last_partition_end_offset);
 
@@ -1359,126 +1370,7 @@ WIAFileReader::ConversionResult WIAFileReader::SetUpDataEntriesForWriting(
 
   add_raw_data_entry(last_partition_end_offset, iso_size - last_partition_end_offset);
 
-  return ConversionResult::Success;
-}
-
-bool WIAFileReader::TryReuseGroup(std::vector<GroupEntry>* group_entries, size_t* groups_written,
-                                  std::map<ReuseID, GroupEntry>* reusable_groups,
-                                  std::optional<ReuseID> reuse_id)
-{
-  if (!reuse_id)
-    return false;
-
-  const auto it = reusable_groups->find(*reuse_id);
-  if (it != reusable_groups->end())
-  {
-    (*group_entries)[*groups_written] = it->second;
-    ++*groups_written;
-  }
-
-  return it != reusable_groups->end();
-}
-
-WIAFileReader::ConversionResult WIAFileReader::CompressAndWriteGroup(
-    File::IOFile* file, u64* bytes_written, std::vector<GroupEntry>* group_entries,
-    size_t* groups_written, Compressor* compressor, bool compressed_exception_lists,
-    const std::vector<u8>& exception_lists, const std::vector<u8>& main_data,
-    std::map<ReuseID, GroupEntry>* reusable_groups, std::optional<ReuseID> reuse_id)
-{
-  const auto all_zero = [](const std::vector<u8>& data) {
-    return std::all_of(data.begin(), data.end(), [](u8 x) { return x == 0; });
-  };
-
-  if (all_zero(exception_lists) && all_zero(main_data))
-  {
-    (*group_entries)[*groups_written] = GroupEntry{0, 0};
-    ++*groups_written;
-    return ConversionResult::Success;
-  }
-
-  if (TryReuseGroup(group_entries, groups_written, reusable_groups, reuse_id))
-    return ConversionResult::Success;
-
-  const u64 data_offset = *bytes_written;
-
-  if (compressor)
-  {
-    if (!compressor->Start())
-      return ConversionResult::InternalError;
-  }
-
-  if (!exception_lists.empty())
-  {
-    if (compressed_exception_lists && compressor)
-    {
-      if (!compressor->Compress(exception_lists.data(), exception_lists.size()))
-        return ConversionResult::InternalError;
-    }
-    else
-    {
-      *bytes_written += exception_lists.size();
-      if (!file->WriteArray(exception_lists.data(), exception_lists.size()))
-        return ConversionResult::WriteFailed;
-
-      const u64 offset_of_padding = *bytes_written;
-      if (!compressed_exception_lists)
-      {
-        if (!PadTo4(file, bytes_written))
-          return ConversionResult::WriteFailed;
-      }
-      const u64 padding_written = *bytes_written - offset_of_padding;
-
-      // Some extra stuff we have to do because Purge for some reason is supposed to hash
-      // the exception lists and the following padding but not actually compress them...
-      if (compressor)
-      {
-        if (!compressor->AddPrecedingDataOnlyForPurgeHashing(exception_lists.data(),
-                                                             exception_lists.size()))
-        {
-          return ConversionResult::InternalError;
-        }
-
-        constexpr u32 ZEROES = 0;
-        if (!compressor->AddPrecedingDataOnlyForPurgeHashing(reinterpret_cast<u8*>(ZEROES),
-                                                             padding_written))
-        {
-          return ConversionResult::InternalError;
-        }
-      }
-    }
-  }
-
-  if (compressor)
-  {
-    if (!compressor->Compress(main_data.data(), main_data.size()))
-      return ConversionResult::InternalError;
-    if (!compressor->End())
-      return ConversionResult::InternalError;
-  }
-
-  const u8* data = compressor ? compressor->GetData() : main_data.data();
-  const size_t size = compressor ? compressor->GetSize() : main_data.size();
-
-  *bytes_written += size;
-  if (!file->WriteArray(data, size))
-    return ConversionResult::WriteFailed;
-
-  if (*bytes_written >> 2 > std::numeric_limits<u32>::max())
-    return ConversionResult::InternalError;
-
-  ASSERT((data_offset & 3) == 0);
-  GroupEntry& group_entry = (*group_entries)[*groups_written];
-  group_entry.data_offset = Common::swap32(static_cast<u32>(data_offset >> 2));
-  group_entry.data_size = Common::swap32(static_cast<u32>(*bytes_written - data_offset));
-  ++*groups_written;
-
-  if (reuse_id)
-    reusable_groups->emplace(*reuse_id, group_entry);
-
-  if (!PadTo4(file, bytes_written))
-    return ConversionResult::WriteFailed;
-
-  return ConversionResult::Success;
+  return ConversionResultCode::Success;
 }
 
 std::optional<std::vector<u8>> WIAFileReader::Compress(Compressor* compressor, const u8* data,
@@ -1496,6 +1388,326 @@ std::optional<std::vector<u8>> WIAFileReader::Compress(Compressor* compressor, c
   return std::vector<u8>(data, data + size);
 }
 
+void WIAFileReader::SetUpCompressor(std::unique_ptr<Compressor>* compressor,
+                                    WIACompressionType compression_type, int compression_level,
+                                    WIAHeader2* header_2)
+{
+  switch (compression_type)
+  {
+  case WIACompressionType::None:
+    *compressor = nullptr;
+    break;
+  case WIACompressionType::Purge:
+    *compressor = std::make_unique<PurgeCompressor>();
+    break;
+  case WIACompressionType::Bzip2:
+    *compressor = std::make_unique<Bzip2Compressor>(compression_level);
+    break;
+  case WIACompressionType::LZMA:
+  case WIACompressionType::LZMA2:
+  {
+    u8* compressor_data = nullptr;
+    u8* compressor_data_size = nullptr;
+
+    if (header_2)
+    {
+      compressor_data = header_2->compressor_data;
+      compressor_data_size = &header_2->compressor_data_size;
+    }
+
+    const bool lzma2 = compression_type == WIACompressionType::LZMA2;
+    *compressor = std::make_unique<LZMACompressor>(lzma2, compression_level, compressor_data,
+                                                   compressor_data_size);
+    break;
+  }
+  }
+}
+
+ConversionResult<WIAFileReader::OutputParameters>
+WIAFileReader::ProcessAndCompress(CompressThreadState* state, CompressParameters parameters,
+                                  const std::vector<PartitionEntry>& partition_entries,
+                                  const std::vector<DataEntry>& data_entries,
+                                  std::map<ReuseID, GroupEntry>* reusable_groups,
+                                  std::mutex* reusable_groups_mutex, u64 exception_lists_per_chunk,
+                                  bool compressed_exception_lists)
+{
+  const auto all_are = [](const std::vector<u8>& data, u8 x) {
+    return std::all_of(data.begin(), data.end(), [x](u8 y) { return x == y; });
+  };
+
+  const auto all_zero = [&all_are](const std::vector<u8>& data) { return all_are(data, 0); };
+
+  const auto all_same = [&all_are](const std::vector<u8>& data) {
+    return all_are(data, data.front());
+  };
+
+  const auto reuse_id_exists = [reusable_groups,
+                                reusable_groups_mutex](const std::optional<ReuseID>& reuse_id) {
+    if (!reuse_id)
+      return false;
+
+    std::lock_guard guard(*reusable_groups_mutex);
+    const auto it = reusable_groups->find(*reuse_id);
+    return it != reusable_groups->end();
+  };
+
+  std::optional<ReuseID> reuse_id;
+
+  state->exceptions_buffer.clear();
+
+  if (!parameters.data_entry->is_partition)
+  {
+    if (all_same(parameters.data))
+      reuse_id = ReuseID{nullptr, parameters.data.size(), false, parameters.data.front()};
+  }
+  else
+  {
+    const PartitionEntry& partition_entry = partition_entries[parameters.data_entry->index];
+
+    mbedtls_aes_context aes_context;
+    mbedtls_aes_setkey_dec(&aes_context, partition_entry.partition_key.data(), 128);
+
+    const u64 groups = Common::AlignUp(parameters.data.size(), VolumeWii::GROUP_TOTAL_SIZE) /
+                       VolumeWii::GROUP_TOTAL_SIZE;
+
+    ASSERT(parameters.data.size() % VolumeWii::BLOCK_TOTAL_SIZE == 0);
+    const u64 blocks = parameters.data.size() / VolumeWii::BLOCK_TOTAL_SIZE;
+    const u64 bytes_to_write = blocks * VolumeWii::BLOCK_DATA_SIZE;
+
+    const auto create_reuse_id = [&partition_entry, bytes_to_write](u8 value, bool decrypted) {
+      return ReuseID{&partition_entry.partition_key, bytes_to_write, decrypted, value};
+    };
+
+    // Set this group as reusable if the encrypted data is all_same
+    if (all_same(parameters.data))
+      reuse_id = create_reuse_id(parameters.data.front(), false);
+
+    if (reuse_id_exists(reuse_id))
+      return OutputParameters{{}, {}, reuse_id, parameters.bytes_read, parameters.group_index};
+
+    std::vector<std::vector<HashExceptionEntry>> exception_lists(exception_lists_per_chunk);
+
+    for (u64 i = 0; i < groups; ++i)
+    {
+      const u64 offset_of_group = i * VolumeWii::GROUP_TOTAL_SIZE;
+      const u64 write_offset_of_group = i * VolumeWii::GROUP_DATA_SIZE;
+
+      const u64 blocks_in_this_group =
+          std::min<u64>(VolumeWii::BLOCKS_PER_GROUP, blocks - i * VolumeWii::BLOCKS_PER_GROUP);
+
+      for (u32 j = 0; j < VolumeWii::BLOCKS_PER_GROUP; ++j)
+      {
+        if (j < blocks_in_this_group)
+        {
+          const u64 offset_of_block = offset_of_group + j * VolumeWii::BLOCK_TOTAL_SIZE;
+          VolumeWii::DecryptBlockData(parameters.data.data() + offset_of_block,
+                                      state->decryption_buffer[j].data(), &aes_context);
+        }
+        else
+        {
+          state->decryption_buffer[j].fill(0);
+        }
+      }
+
+      VolumeWii::HashGroup(state->decryption_buffer.data(), state->hash_buffer.data());
+
+      for (u64 j = 0; j < blocks_in_this_group; ++j)
+      {
+        const u64 offset_of_block = offset_of_group + j * VolumeWii::BLOCK_TOTAL_SIZE;
+        const u64 hash_offset_of_block = j * VolumeWii::BLOCK_HEADER_SIZE;
+
+        VolumeWii::HashBlock hashes;
+        VolumeWii::DecryptBlockHashes(parameters.data.data() + offset_of_block, &hashes,
+                                      &aes_context);
+
+        const auto compare_hash = [&](size_t offset_in_block) {
+          ASSERT(offset_in_block + sizeof(SHA1) <= VolumeWii::BLOCK_HEADER_SIZE);
+
+          const u8* desired_hash = reinterpret_cast<u8*>(&hashes) + offset_in_block;
+          const u8* computed_hash = reinterpret_cast<u8*>(&state->hash_buffer[j]) + offset_in_block;
+
+          if (!std::equal(desired_hash, desired_hash + sizeof(SHA1), computed_hash))
+          {
+            const u64 hash_offset = hash_offset_of_block + offset_in_block;
+            ASSERT(hash_offset <= std::numeric_limits<u16>::max());
+
+            HashExceptionEntry& exception = exception_lists[i].emplace_back();
+            exception.offset = static_cast<u16>(Common::swap16(hash_offset));
+            std::memcpy(exception.hash.data(), desired_hash, sizeof(SHA1));
+          }
+        };
+
+        const auto compare_hashes = [&compare_hash](size_t offset, size_t size) {
+          for (size_t l = 0; l < size; l += sizeof(SHA1))
+            // The std::min is to ensure that we don't go beyond the end of HashBlock with
+            // padding_2, which is 32 bytes long (not divisible by sizeof(SHA1), which is 20).
+            compare_hash(offset + std::min(l, size - sizeof(SHA1)));
+        };
+
+        using HashBlock = VolumeWii::HashBlock;
+        compare_hashes(offsetof(HashBlock, h0), sizeof(HashBlock::h0));
+        compare_hashes(offsetof(HashBlock, padding_0), sizeof(HashBlock::padding_0));
+        compare_hashes(offsetof(HashBlock, h1), sizeof(HashBlock::h1));
+        compare_hashes(offsetof(HashBlock, padding_1), sizeof(HashBlock::padding_1));
+        compare_hashes(offsetof(HashBlock, h2), sizeof(HashBlock::h2));
+        compare_hashes(offsetof(HashBlock, padding_2), sizeof(HashBlock::padding_2));
+      }
+
+      for (u64 j = 0; j < blocks_in_this_group; ++j)
+      {
+        const u64 write_offset_of_block = write_offset_of_group + j * VolumeWii::BLOCK_DATA_SIZE;
+        std::memcpy(parameters.data.data() + write_offset_of_block,
+                    state->decryption_buffer[j].data(), VolumeWii::BLOCK_DATA_SIZE);
+      }
+    }
+
+    bool have_exceptions = false;
+
+    for (const std::vector<HashExceptionEntry>& exception_list : exception_lists)
+    {
+      const u16 exceptions = Common::swap16(static_cast<u16>(exception_list.size()));
+      PushBack(&state->exceptions_buffer, exceptions);
+      for (const HashExceptionEntry& exception : exception_list)
+        PushBack(&state->exceptions_buffer, exception);
+      if (!exception_list.empty())
+        have_exceptions = true;
+    }
+
+    parameters.data.resize(bytes_to_write);
+
+    // Set this group as reusable if it lacks exceptions and the decrypted data is all_same
+    if (!reuse_id && !have_exceptions && all_same(parameters.data))
+      reuse_id = create_reuse_id(parameters.data.front(), true);
+  }
+
+  // Special case - a compressed size of zero is treated by WIA as meaning the data is all zeroes
+  if (all_zero(state->exceptions_buffer) && all_zero(parameters.data))
+    return OutputParameters{{}, {}, reuse_id, parameters.bytes_read, parameters.group_index};
+
+  if (reuse_id_exists(reuse_id))
+    return OutputParameters{{}, {}, reuse_id, parameters.bytes_read, parameters.group_index};
+
+  if (state->compressor)
+  {
+    if (!state->compressor->Start())
+      return ConversionResultCode::InternalError;
+  }
+
+  if (!state->exceptions_buffer.empty())
+  {
+    if (compressed_exception_lists && state->compressor)
+    {
+      if (!state->compressor->Compress(state->exceptions_buffer.data(),
+                                       state->exceptions_buffer.size()))
+      {
+        return ConversionResultCode::InternalError;
+      }
+
+      state->exceptions_buffer.clear();
+    }
+    else
+    {
+      if (!compressed_exception_lists)
+      {
+        while (state->exceptions_buffer.size() % 4 != 0)
+          state->exceptions_buffer.push_back(0);
+      }
+
+      if (state->compressor)
+      {
+        if (!state->compressor->AddPrecedingDataOnlyForPurgeHashing(
+                state->exceptions_buffer.data(), state->exceptions_buffer.size()))
+        {
+          return ConversionResultCode::InternalError;
+        }
+      }
+    }
+  }
+
+  if (state->compressor)
+  {
+    if (!state->compressor->Compress(parameters.data.data(), parameters.data.size()))
+      return ConversionResultCode::InternalError;
+    if (!state->compressor->End())
+      return ConversionResultCode::InternalError;
+  }
+
+  if (state->compressor)
+  {
+    const u8* data = state->compressor->GetData();
+    const size_t size = state->compressor->GetSize();
+
+    parameters.data.resize(size);
+    std::copy(data, data + size, parameters.data.data());
+  }
+
+  return OutputParameters{state->exceptions_buffer, std::move(parameters.data), reuse_id,
+                          parameters.bytes_read, parameters.group_index};
+}
+
+ConversionResultCode WIAFileReader::Output(const OutputParameters& parameters,
+                                           File::IOFile* outfile,
+                                           std::map<ReuseID, GroupEntry>* reusable_groups,
+                                           std::mutex* reusable_groups_mutex,
+                                           GroupEntry* group_entry, u64* bytes_written)
+{
+  if (parameters.reuse_id)
+  {
+    std::lock_guard guard(*reusable_groups_mutex);
+    const auto it = reusable_groups->find(*parameters.reuse_id);
+    if (it != reusable_groups->end())
+    {
+      *group_entry = it->second;
+      return ConversionResultCode::Success;
+    }
+  }
+
+  const size_t data_size = parameters.exception_lists.size() + parameters.main_data.size();
+
+  if (*bytes_written >> 2 > std::numeric_limits<u32>::max())
+    return ConversionResultCode::InternalError;
+
+  ASSERT((*bytes_written & 3) == 0);
+  group_entry->data_offset = Common::swap32(static_cast<u32>(*bytes_written >> 2));
+  group_entry->data_size = Common::swap32(static_cast<u32>(data_size));
+
+  if (!outfile->WriteArray(parameters.exception_lists.data(), parameters.exception_lists.size()))
+    return ConversionResultCode::WriteFailed;
+  if (!outfile->WriteArray(parameters.main_data.data(), parameters.main_data.size()))
+    return ConversionResultCode::WriteFailed;
+
+  *bytes_written += data_size;
+
+  if (parameters.reuse_id)
+  {
+    std::lock_guard guard(*reusable_groups_mutex);
+    reusable_groups->emplace(*parameters.reuse_id, *group_entry);
+  }
+
+  if (!PadTo4(outfile, bytes_written))
+    return ConversionResultCode::WriteFailed;
+
+  return ConversionResultCode::Success;
+}
+
+ConversionResultCode WIAFileReader::RunCallback(size_t groups_written, u64 bytes_read,
+                                                u64 bytes_written, u32 total_groups, u64 iso_size,
+                                                CompressCB callback, void* arg)
+{
+  int ratio = 0;
+  if (bytes_read != 0)
+    ratio = static_cast<int>(100 * bytes_written / bytes_read);
+
+  const std::string text =
+      StringFromFormat(Common::GetStringT("%i of %i blocks. Compression ratio %i%%").c_str(),
+                       groups_written, total_groups, ratio);
+
+  const float completion = static_cast<float>(bytes_read) / iso_size;
+
+  return callback(text, completion, arg) ? ConversionResultCode::Success :
+                                           ConversionResultCode::Canceled;
+}
+
 bool WIAFileReader::WriteHeader(File::IOFile* file, const u8* data, size_t size, u64 upper_bound,
                                 u64* bytes_written, u64* offset_out)
 {
@@ -1516,7 +1728,7 @@ bool WIAFileReader::WriteHeader(File::IOFile* file, const u8* data, size_t size,
   return PadTo4(file, bytes_written);
 }
 
-WIAFileReader::ConversionResult
+ConversionResultCode
 WIAFileReader::ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
                             File::IOFile* outfile, WIACompressionType compression_type,
                             int compression_level, int chunk_size, CompressCB callback, void* arg)
@@ -1530,63 +1742,22 @@ WIAFileReader::ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
 
   u64 bytes_read = 0;
   u64 bytes_written = 0;
-  size_t groups_written = 0;
+  size_t groups_processed = 0;
 
   WIAHeader1 header_1{};
   WIAHeader2 header_2{};
 
-  std::unique_ptr<Compressor> compressor;
-  switch (compression_type)
-  {
-  case WIACompressionType::None:
-    compressor = nullptr;
-    break;
-  case WIACompressionType::Purge:
-    compressor = std::make_unique<PurgeCompressor>();
-    break;
-  case WIACompressionType::Bzip2:
-    compressor = std::make_unique<Bzip2Compressor>(compression_level);
-    break;
-  case WIACompressionType::LZMA:
-    compressor = std::make_unique<LZMACompressor>(
-        false, compression_level, header_2.compressor_data, &header_2.compressor_data_size);
-    break;
-  case WIACompressionType::LZMA2:
-    compressor = std::make_unique<LZMACompressor>(true, compression_level, header_2.compressor_data,
-                                                  &header_2.compressor_data_size);
-    break;
-  }
-
   std::vector<PartitionEntry> partition_entries;
   std::vector<RawDataEntry> raw_data_entries;
   std::vector<GroupEntry> group_entries;
 
-  const auto run_callback = [&] {
-    int ratio = 0;
-    if (bytes_read != 0)
-      ratio = static_cast<int>(100 * bytes_written / bytes_read);
-
-    const std::string temp =
-        StringFromFormat(Common::GetStringT("%i of %i blocks. Compression ratio %i%%").c_str(),
-                         groups_written, group_entries.size(), ratio);
-
-    float completion = 0.0f;
-    if (group_entries.size() != 0)
-      completion = static_cast<float>(groups_written) / group_entries.size();
-
-    return callback(temp, completion, arg);
-  };
-
-  if (!run_callback())
-    return ConversionResult::Canceled;
-
   u32 total_groups;
   std::vector<DataEntry> data_entries;
 
-  const ConversionResult set_up_data_entries_result =
+  const ConversionResultCode set_up_data_entries_result =
       SetUpDataEntriesForWriting(infile_volume, chunk_size, iso_size, &total_groups,
                                  &partition_entries, &raw_data_entries, &data_entries);
-  if (set_up_data_entries_result != ConversionResult::Success)
+  if (set_up_data_entries_result != ConversionResultCode::Success)
     return set_up_data_entries_result;
 
   group_entries.resize(total_groups);
@@ -1610,251 +1781,128 @@ WIAFileReader::ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
   bytes_written = headers_size_upper_bound;
 
   if (!infile->Read(0, header_2.disc_header.size(), header_2.disc_header.data()))
-    return ConversionResult::ReadFailed;
+    return ConversionResultCode::ReadFailed;
   // We intentially do not increment bytes_read here, since these bytes will be read again
 
-  const auto all_same = [](const std::vector<u8>& data) {
-    const u8 first_byte = data.front();
-    return std::all_of(data.begin(), data.end(), [first_byte](u8 x) { return x == first_byte; });
+  std::map<ReuseID, GroupEntry> reusable_groups;
+  std::mutex reusable_groups_mutex;
+
+  const auto set_up_compress_thread_state = [&](CompressThreadState* state) {
+    SetUpCompressor(&state->compressor, compression_type, compression_level, nullptr);
+    return ConversionResultCode::Success;
   };
 
-  using WiiBlockData = std::array<u8, VolumeWii::BLOCK_DATA_SIZE>;
+  const auto process_and_compress = [&](CompressThreadState* state, CompressParameters parameters) {
+    return ProcessAndCompress(state, std::move(parameters), partition_entries, data_entries,
+                              &reusable_groups, &reusable_groups_mutex, exception_lists_per_chunk,
+                              compressed_exception_lists);
+  };
 
-  std::vector<u8> exceptions_buffer;
-  std::vector<WiiBlockData> decryption_buffer;
-  std::vector<VolumeWii::HashBlock> hash_buffer;
+  const auto output = [&](OutputParameters parameters) {
+    const ConversionResultCode result =
+        Output(parameters, outfile, &reusable_groups, &reusable_groups_mutex,
+               &group_entries[parameters.group_index], &bytes_written);
 
-  std::map<ReuseID, GroupEntry> reusable_groups;
+    if (result != ConversionResultCode::Success)
+      return result;
 
-  if (!partition_entries.empty())
-  {
-    decryption_buffer.resize(VolumeWii::BLOCKS_PER_GROUP);
-    hash_buffer.resize(VolumeWii::BLOCKS_PER_GROUP);
-  }
+    return RunCallback(parameters.group_index + 1, parameters.bytes_read, bytes_written,
+                       total_groups, iso_size, callback, arg);
+  };
+
+  MultithreadedCompressor<CompressThreadState, CompressParameters, OutputParameters> mt_compressor(
+      set_up_compress_thread_state, process_and_compress, output);
 
   for (const DataEntry& data_entry : data_entries)
   {
+    u32 first_group;
+    u32 last_group;
+
+    u64 data_offset;
+    u64 data_size;
+
     if (data_entry.is_partition)
     {
       const PartitionEntry& partition_entry = partition_entries[data_entry.index];
       const PartitionDataEntry& partition_data_entry =
           partition_entry.data_entries[data_entry.partition_data_index];
 
-      const u32 first_group = Common::swap32(partition_data_entry.group_index);
-      const u32 last_group = first_group + Common::swap32(partition_data_entry.number_of_groups);
+      first_group = Common::swap32(partition_data_entry.group_index);
+      last_group = first_group + Common::swap32(partition_data_entry.number_of_groups);
 
-      const u64 data_offset =
-          Common::swap32(partition_data_entry.first_sector) * VolumeWii::BLOCK_TOTAL_SIZE;
-      const u64 data_size =
+      data_offset = Common::swap32(partition_data_entry.first_sector) * VolumeWii::BLOCK_TOTAL_SIZE;
+      data_size =
           Common::swap32(partition_data_entry.number_of_sectors) * VolumeWii::BLOCK_TOTAL_SIZE;
-
-      ASSERT(groups_written == first_group);
-      ASSERT(bytes_read == data_offset);
-
-      mbedtls_aes_context aes_context;
-      mbedtls_aes_setkey_dec(&aes_context, partition_entry.partition_key.data(), 128);
-
-      for (u32 i = first_group; i < last_group; ++i)
-      {
-        const u64 bytes_to_read = std::min<u64>(chunk_size, data_offset + data_size - bytes_read);
-
-        const u64 groups = Common::AlignUp(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE) /
-                           VolumeWii::GROUP_TOTAL_SIZE;
-
-        ASSERT(bytes_to_read % VolumeWii::BLOCK_TOTAL_SIZE == 0);
-        const u64 blocks = bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE;
-        const u64 bytes_to_write = blocks * VolumeWii::BLOCK_DATA_SIZE;
-
-        buffer.resize(bytes_to_read);
-        if (!infile->Read(bytes_read, bytes_to_read, buffer.data()))
-          return ConversionResult::ReadFailed;
-        bytes_read += bytes_to_read;
-
-        const auto create_reuse_id = [&partition_entry, bytes_to_write](u8 value, bool decrypted) {
-          return ReuseID{&partition_entry.partition_key, bytes_to_write, decrypted, value};
-        };
-
-        std::optional<ReuseID> reuse_id;
-
-        // Set this group as reusable if the encrypted data is all_same
-        if (all_same(buffer))
-          reuse_id = create_reuse_id(buffer.front(), false);
-
-        if (!TryReuseGroup(&group_entries, &groups_written, &reusable_groups, reuse_id))
-        {
-          std::vector<std::vector<HashExceptionEntry>> exception_lists(exception_lists_per_chunk);
-
-          for (u64 j = 0; j < groups; ++j)
-          {
-            const u64 offset_of_group = j * VolumeWii::GROUP_TOTAL_SIZE;
-            const u64 write_offset_of_group = j * VolumeWii::GROUP_DATA_SIZE;
-
-            const u64 blocks_in_this_group = std::min<u64>(
-                VolumeWii::BLOCKS_PER_GROUP, blocks - j * VolumeWii::BLOCKS_PER_GROUP);
-
-            for (u32 k = 0; k < VolumeWii::BLOCKS_PER_GROUP; ++k)
-            {
-              if (k < blocks_in_this_group)
-              {
-                const u64 offset_of_block = offset_of_group + k * VolumeWii::BLOCK_TOTAL_SIZE;
-                VolumeWii::DecryptBlockData(buffer.data() + offset_of_block,
-                                            decryption_buffer[k].data(), &aes_context);
-              }
-              else
-              {
-                decryption_buffer[k].fill(0);
-              }
-            }
-
-            VolumeWii::HashGroup(decryption_buffer.data(), hash_buffer.data());
-
-            for (u64 k = 0; k < blocks_in_this_group; ++k)
-            {
-              const u64 offset_of_block = offset_of_group + k * VolumeWii::BLOCK_TOTAL_SIZE;
-              const u64 hash_offset_of_block = k * VolumeWii::BLOCK_HEADER_SIZE;
-
-              VolumeWii::HashBlock hashes;
-              VolumeWii::DecryptBlockHashes(buffer.data() + offset_of_block, &hashes, &aes_context);
-
-              const auto compare_hash = [&](size_t offset_in_block) {
-                ASSERT(offset_in_block + sizeof(SHA1) <= VolumeWii::BLOCK_HEADER_SIZE);
-
-                const u8* desired_hash = reinterpret_cast<u8*>(&hashes) + offset_in_block;
-                const u8* computed_hash = reinterpret_cast<u8*>(&hash_buffer[k]) + offset_in_block;
-
-                if (!std::equal(desired_hash, desired_hash + sizeof(SHA1), computed_hash))
-                {
-                  const u64 hash_offset = hash_offset_of_block + offset_in_block;
-                  ASSERT(hash_offset <= std::numeric_limits<u16>::max());
-
-                  HashExceptionEntry& exception = exception_lists[j].emplace_back();
-                  exception.offset = static_cast<u16>(Common::swap16(hash_offset));
-                  std::memcpy(exception.hash.data(), desired_hash, sizeof(SHA1));
-                }
-              };
-
-              const auto compare_hashes = [&compare_hash](size_t offset, size_t size) {
-                for (size_t l = 0; l < size; l += sizeof(SHA1))
-                  // The std::min is to ensure that we don't go beyond the end of HashBlock with
-                  // padding_2, which is 32 bytes long (not divisible by sizeof(SHA1), which is 20).
-                  compare_hash(offset + std::min(l, size - sizeof(SHA1)));
-              };
-
-              using HashBlock = VolumeWii::HashBlock;
-              compare_hashes(offsetof(HashBlock, h0), sizeof(HashBlock::h0));
-              compare_hashes(offsetof(HashBlock, padding_0), sizeof(HashBlock::padding_0));
-              compare_hashes(offsetof(HashBlock, h1), sizeof(HashBlock::h1));
-              compare_hashes(offsetof(HashBlock, padding_1), sizeof(HashBlock::padding_1));
-              compare_hashes(offsetof(HashBlock, h2), sizeof(HashBlock::h2));
-              compare_hashes(offsetof(HashBlock, padding_2), sizeof(HashBlock::padding_2));
-            }
-
-            for (u64 k = 0; k < blocks_in_this_group; ++k)
-            {
-              std::memcpy(buffer.data() + write_offset_of_group + k * VolumeWii::BLOCK_DATA_SIZE,
-                          decryption_buffer[k].data(), VolumeWii::BLOCK_DATA_SIZE);
-            }
-          }
-
-          bool have_exceptions = false;
-
-          exceptions_buffer.clear();
-          for (const std::vector<HashExceptionEntry>& exception_list : exception_lists)
-          {
-            const u16 exceptions = Common::swap16(static_cast<u16>(exception_list.size()));
-            PushBack(&exceptions_buffer, exceptions);
-            for (const HashExceptionEntry& exception : exception_list)
-              PushBack(&exceptions_buffer, exception);
-            if (!exception_list.empty())
-              have_exceptions = true;
-          }
-
-          buffer.resize(bytes_to_write);
-
-          // Set this group as reusable if it lacks exceptions and the decrypted data is all_same
-          if (!reuse_id && !have_exceptions && all_same(buffer))
-            reuse_id = create_reuse_id(buffer.front(), true);
-
-          const ConversionResult write_result = CompressAndWriteGroup(
-              outfile, &bytes_written, &group_entries, &groups_written, compressor.get(),
-              compressed_exception_lists, exceptions_buffer, buffer, &reusable_groups, reuse_id);
-
-          if (write_result != ConversionResult::Success)
-            return write_result;
-        }
-
-        if (!run_callback())
-          return ConversionResult::Canceled;
-      }
     }
     else
     {
       const RawDataEntry& raw_data_entry = raw_data_entries[data_entry.index];
 
-      const u32 first_group = Common::swap32(raw_data_entry.group_index);
-      const u32 last_group = first_group + Common::swap32(raw_data_entry.number_of_groups);
+      first_group = Common::swap32(raw_data_entry.group_index);
+      last_group = first_group + Common::swap32(raw_data_entry.number_of_groups);
 
-      u64 data_offset = Common::swap64(raw_data_entry.data_offset);
-      u64 data_size = Common::swap64(raw_data_entry.data_size);
+      data_offset = Common::swap64(raw_data_entry.data_offset);
+      data_size = Common::swap64(raw_data_entry.data_size);
 
       const u64 skipped_data = data_offset % VolumeWii::BLOCK_TOTAL_SIZE;
       data_offset -= skipped_data;
       data_size += skipped_data;
+    }
 
-      ASSERT(groups_written == first_group);
+    ASSERT(groups_processed == first_group);
     ASSERT(bytes_read == data_offset);
 
-      exceptions_buffer.clear();
-
     for (u32 i = first_group; i < last_group; ++i)
     {
+      const ConversionResultCode status = mt_compressor.GetStatus();
+      if (status != ConversionResultCode::Success)
+        return status;
+
       const u64 bytes_to_read = std::min<u64>(chunk_size, data_offset + data_size - bytes_read);
 
       buffer.resize(bytes_to_read);
-
       if (!infile->Read(bytes_read, bytes_to_read, buffer.data()))
-          return ConversionResult::ReadFailed;
+        return ConversionResultCode::ReadFailed;
       bytes_read += bytes_to_read;
 
-        std::optional<ReuseID> reuse_id;
-        if (all_same(buffer))
-          reuse_id = ReuseID{nullptr, bytes_to_read, false, buffer.front()};
+      mt_compressor.CompressAndWrite(
+          CompressParameters{buffer, &data_entry, bytes_read, groups_processed});
 
-        const ConversionResult write_result = CompressAndWriteGroup(
-            outfile, &bytes_written, &group_entries, &groups_written, compressor.get(),
-            compressed_exception_lists, exceptions_buffer, buffer, &reusable_groups, reuse_id);
-
-        if (write_result != ConversionResult::Success)
-          return write_result;
-
-        if (!run_callback())
-          return ConversionResult::Canceled;
-      }
+      ++groups_processed;
     }
   }
 
-  ASSERT(groups_written == total_groups);
+  ASSERT(groups_processed == total_groups);
   ASSERT(bytes_read == iso_size);
 
+  mt_compressor.Shutdown();
+
+  const ConversionResultCode status = mt_compressor.GetStatus();
+  if (status != ConversionResultCode::Success)
+    return status;
+
+  std::unique_ptr<Compressor> compressor;
+  SetUpCompressor(&compressor, compression_type, compression_level, &header_2);
+
   const std::optional<std::vector<u8>> compressed_raw_data_entries = Compress(
       compressor.get(), reinterpret_cast<u8*>(raw_data_entries.data()), raw_data_entries_size);
   if (!compressed_raw_data_entries)
-    return ConversionResult::InternalError;
+    return ConversionResultCode::InternalError;
 
   const std::optional<std::vector<u8>> compressed_group_entries =
       Compress(compressor.get(), reinterpret_cast<u8*>(group_entries.data()), group_entries_size);
   if (!compressed_group_entries)
-    return ConversionResult::InternalError;
+    return ConversionResultCode::InternalError;
 
   bytes_written = sizeof(WIAHeader1) + sizeof(WIAHeader2);
   if (!outfile->Seek(sizeof(WIAHeader1) + sizeof(WIAHeader2), SEEK_SET))
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
 
   u64 partition_entries_offset;
   if (!WriteHeader(outfile, reinterpret_cast<u8*>(partition_entries.data()), partition_entries_size,
                    headers_size_upper_bound, &bytes_written, &partition_entries_offset))
   {
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
   }
 
   u64 raw_data_entries_offset;
@@ -1862,14 +1910,14 @@ WIAFileReader::ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
                    compressed_raw_data_entries->size(), headers_size_upper_bound, &bytes_written,
                    &raw_data_entries_offset))
   {
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
   }
 
   u64 group_entries_offset;
   if (!WriteHeader(outfile, compressed_group_entries->data(), compressed_group_entries->size(),
                    headers_size_upper_bound, &bytes_written, &group_entries_offset))
   {
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
   }
 
   u32 disc_type = 0;
@@ -1916,14 +1964,14 @@ WIAFileReader::ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
                    header_1.header_1_hash.data());
 
   if (!outfile->Seek(0, SEEK_SET))
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
 
   if (!outfile->WriteArray(&header_1, 1))
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
   if (!outfile->WriteArray(&header_2, 1))
-    return ConversionResult::WriteFailed;
+    return ConversionResultCode::WriteFailed;
 
-  return ConversionResult::Success;
+  return ConversionResultCode::Success;
 }
 
 bool ConvertToWIA(BlobReader* infile, const std::string& infile_path,
@@ -1942,28 +1990,28 @@ bool ConvertToWIA(BlobReader* infile, const std::string& infile_path,
 
   std::unique_ptr<VolumeDisc> infile_volume = CreateDisc(infile_path);
 
-  const WIAFileReader::ConversionResult result =
+  const ConversionResultCode result =
       WIAFileReader::ConvertToWIA(infile, infile_volume.get(), &outfile, compression_type,
                                   compression_level, chunk_size, callback, arg);
 
-  if (result == WIAFileReader::ConversionResult::ReadFailed)
+  if (result == ConversionResultCode::ReadFailed)
     PanicAlertT("Failed to read from the input file \"%s\".", infile_path.c_str());
 
-  if (result == WIAFileReader::ConversionResult::WriteFailed)
+  if (result == ConversionResultCode::WriteFailed)
   {
     PanicAlertT("Failed to write the output file \"%s\".\n"
                 "Check that you have enough space available on the target drive.",
                 outfile_path.c_str());
   }
 
-  if (result != WIAFileReader::ConversionResult::Success)
+  if (result != ConversionResultCode::Success)
   {
     // Remove the incomplete output file
     outfile.Close();
     File::Delete(outfile_path);
   }
 
-  return result == WIAFileReader::ConversionResult::Success;
+  return result == ConversionResultCode::Success;
 }
 
 }  // namespace DiscIO

Source/Core/DiscIO/WIABlob.h

@@ -8,6 +8,7 @@
 #include <limits>
 #include <map>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <utility>
 
@@ -19,6 +20,7 @@
 #include "Common/File.h"
 #include "Common/Swap.h"
 #include "DiscIO/Blob.h"
+#include "DiscIO/MultithreadedCompressor.h"
 #include "DiscIO/WiiEncryptionCache.h"
 
 namespace DiscIO
@@ -56,19 +58,11 @@ public:
   bool SupportsReadWiiDecrypted() const override;
   bool ReadWiiDecrypted(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset) override;
 
-  enum class ConversionResult
-  {
-    Success,
-    Canceled,
-    ReadFailed,
-    WriteFailed,
-    InternalError,
-  };
-
-  static ConversionResult ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
-                                       File::IOFile* outfile, WIACompressionType compression_type,
-                                       int compression_level, int chunk_size, CompressCB callback,
-                                       void* arg);
+  static ConversionResultCode ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
+                                           File::IOFile* outfile,
+                                           WIACompressionType compression_type,
+                                           int compression_level, int chunk_size,
+                                           CompressCB callback, void* arg);
 
 private:
   using SHA1 = std::array<u8, 20>;
@@ -417,6 +411,38 @@ private:
     u8 value;
   };
 
+  struct CompressThreadState
+  {
+    using WiiBlockData = std::array<u8, VolumeWii::BLOCK_DATA_SIZE>;
+
+    std::unique_ptr<Compressor> compressor;
+
+    std::vector<WiiBlockData> decryption_buffer =
+        std::vector<WiiBlockData>(VolumeWii::BLOCKS_PER_GROUP);
+
+    std::vector<VolumeWii::HashBlock> hash_buffer =
+        std::vector<VolumeWii::HashBlock>(VolumeWii::BLOCKS_PER_GROUP);
+
+    std::vector<u8> exceptions_buffer;
+  };
+
+  struct CompressParameters
+  {
+    std::vector<u8> data;
+    const DataEntry* data_entry;
+    u64 bytes_read;
+    size_t group_index;
+  };
+
+  struct OutputParameters
+  {
+    std::vector<u8> exception_lists;
+    std::vector<u8> main_data;
+    std::optional<ReuseID> reuse_id;
+    u64 bytes_read;
+    size_t group_index;
+  };
+
   static bool PadTo4(File::IOFile* file, u64* bytes_written);
   static void AddRawDataEntry(u64 offset, u64 size, int chunk_size, u32* total_groups,
                               std::vector<RawDataEntry>* raw_data_entries,
@@ -425,24 +451,34 @@ private:
   CreatePartitionDataEntry(u64 offset, u64 size, u32 index, int chunk_size, u32* total_groups,
                            const std::vector<PartitionEntry>& partition_entries,
                            std::vector<DataEntry>* data_entries);
-  static ConversionResult SetUpDataEntriesForWriting(const VolumeDisc* volume, int chunk_size,
-                                                     u64 iso_size, u32* total_groups,
-                                                     std::vector<PartitionEntry>* partition_entries,
+  static ConversionResultCode
+  SetUpDataEntriesForWriting(const VolumeDisc* volume, int chunk_size, u64 iso_size,
+                             u32* total_groups, std::vector<PartitionEntry>* partition_entries,
                              std::vector<RawDataEntry>* raw_data_entries,
                              std::vector<DataEntry>* data_entries);
-  static bool TryReuseGroup(std::vector<GroupEntry>* group_entries, size_t* groups_written,
-                            std::map<ReuseID, GroupEntry>* reusable_groups,
-                            std::optional<ReuseID> reuse_id);
-  static ConversionResult CompressAndWriteGroup(
-      File::IOFile* file, u64* bytes_written, std::vector<GroupEntry>* group_entries,
-      size_t* groups_written, Compressor* compressor, bool compressed_exception_lists,
-      const std::vector<u8>& exception_lists, const std::vector<u8>& main_data,
-      std::map<ReuseID, GroupEntry>* reusable_groups, std::optional<ReuseID> reuse_id);
   static std::optional<std::vector<u8>> Compress(Compressor* compressor, const u8* data,
                                                  size_t size);
   static bool WriteHeader(File::IOFile* file, const u8* data, size_t size, u64 upper_bound,
                           u64* bytes_written, u64* offset_out);
 
+  static void SetUpCompressor(std::unique_ptr<Compressor>* compressor,
+                              WIACompressionType compression_type, int compression_level,
+                              WIAHeader2* header_2);
+  static ConversionResult<OutputParameters>
+  ProcessAndCompress(CompressThreadState* state, CompressParameters parameters,
+                     const std::vector<PartitionEntry>& partition_entries,
+                     const std::vector<DataEntry>& data_entries,
+                     std::map<ReuseID, GroupEntry>* reusable_groups,
+                     std::mutex* reusable_groups_mutex, u64 exception_lists_per_chunk,
+                     bool compressed_exception_lists);
+  static ConversionResultCode Output(const OutputParameters& parameters, File::IOFile* outfile,
+                                     std::map<ReuseID, GroupEntry>* reusable_groups,
+                                     std::mutex* reusable_groups_mutex, GroupEntry* group_entry,
+                                     u64* bytes_written);
+  static ConversionResultCode RunCallback(size_t groups_written, u64 bytes_read, u64 bytes_written,
+                                          u32 total_groups, u64 iso_size, CompressCB callback,
+                                          void* arg);
+
   template <typename T>
   static void PushBack(std::vector<u8>* vector, const T& x)
   {