WIA: Implement compression

Source/Core/DiscIO/Blob.h

@@ -25,6 +25,8 @@
 
 namespace DiscIO
 {
+enum class WIACompressionType : u32;
+
 // Increment CACHE_REVISION (GameFileCache.cpp) if the enum below is modified
 enum class BlobType
 {
@@ -174,7 +176,8 @@ bool ConvertToPlain(BlobReader* infile, const std::string& infile_path,
                     const std::string& outfile_path, CompressCB callback = nullptr,
                     void* arg = nullptr);
 bool ConvertToWIA(BlobReader* infile, const std::string& infile_path,
-                  const std::string& outfile_path, int chunk_size, CompressCB callback = nullptr,
+                  const std::string& outfile_path, WIACompressionType compression_type,
+                  int compression_level, int chunk_size, CompressCB callback = nullptr,
                   void* arg = nullptr);
 
 }  // namespace DiscIO

Source/Core/DiscIO/WIABlob.h

@@ -24,6 +24,15 @@ namespace DiscIO
 {
 class VolumeDisc;
 
+enum class WIACompressionType : u32
+{
+  None = 0,
+  Purge = 1,
+  Bzip2 = 2,
+  LZMA = 3,
+  LZMA2 = 4,
+};
+
 constexpr u32 WIA_MAGIC = 0x01414957;  // "WIA\x1" (byteswapped to little endian)
 
 class WIAFileReader : public BlobReader
@@ -56,22 +65,14 @@ public:
   };
 
   static ConversionResult ConvertToWIA(BlobReader* infile, const VolumeDisc* infile_volume,
-                                       File::IOFile* outfile, int chunk_size, CompressCB callback,
+                                       File::IOFile* outfile, WIACompressionType compression_type,
+                                       int compression_level, int chunk_size, CompressCB callback,
                                        void* arg);
 
 private:
   using SHA1 = std::array<u8, 20>;
   using WiiKey = std::array<u8, 16>;
 
-  enum class CompressionType : u32
-  {
-    None = 0,
-    Purge = 1,
-    Bzip2 = 2,
-    LZMA = 3,
-    LZMA2 = 4,
-  };
-
   // See docs/WIA.md for details about the format
 
 #pragma pack(push, 1)
@@ -253,6 +254,88 @@ private:
     bool m_error_occurred = false;
   };
 
+  class Compressor
+  {
+  public:
+    virtual ~Compressor();
+
+    // First call Start, then AddDataOnlyForPurgeHashing/Compress any number of times,
+    // then End, then GetData/GetSize any number of times.
+
+    virtual bool Start() = 0;
+    virtual bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) { return true; }
+    virtual bool Compress(const u8* data, size_t size) = 0;
+    virtual bool End() = 0;
+
+    virtual const u8* GetData() const = 0;
+    virtual size_t GetSize() const = 0;
+  };
+
+  class PurgeCompressor final : public Compressor
+  {
+  public:
+    PurgeCompressor();
+    ~PurgeCompressor();
+
+    bool Start() override;
+    bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) override;
+    bool Compress(const u8* data, size_t size) override;
+    bool End() override;
+
+    const u8* GetData() const override;
+    size_t GetSize() const override;
+
+  private:
+    std::vector<u8> m_buffer;
+    size_t m_bytes_written;
+    mbedtls_sha1_context m_sha1_context;
+  };
+
+  class Bzip2Compressor final : public Compressor
+  {
+  public:
+    Bzip2Compressor(int compression_level);
+    ~Bzip2Compressor();
+
+    bool Start() override;
+    bool Compress(const u8* data, size_t size) override;
+    bool End() override;
+
+    const u8* GetData() const override;
+    size_t GetSize() const override;
+
+  private:
+    void ExpandBuffer(size_t bytes_to_add);
+
+    bz_stream m_stream = {};
+    std::vector<u8> m_buffer;
+    int m_compression_level;
+  };
+
+  class LZMACompressor final : public Compressor
+  {
+  public:
+    LZMACompressor(bool lzma2, int compression_level, u8 compressor_data_out[7],
+                   u8* compressor_data_size_out);
+    ~LZMACompressor();
+
+    bool Start() override;
+    bool Compress(const u8* data, size_t size) override;
+    bool End() override;
+
+    const u8* GetData() const override;
+    size_t GetSize() const override;
+
+  private:
+    void ExpandBuffer(size_t bytes_to_add);
+
+    lzma_stream m_stream = LZMA_STREAM_INIT;
+    lzma_options_lzma m_options = {};
+    lzma_filter m_filters[2];
+    std::vector<u8> m_buffer;
+    bool m_initialization_failed = false;
+  };
+
   class Chunk
   {
   public:
@@ -304,6 +387,8 @@ private:
 
   static std::string VersionToString(u32 version);
 
+  static u32 LZMA2DictionarySize(u8 p);
+
   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,
@@ -317,9 +402,27 @@ private:
                                                      std::vector<PartitionEntry>* partition_entries,
                                                      std::vector<RawDataEntry>* raw_data_entries,
                                                      std::vector<DataEntry>* data_entries);
+  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);
+  static ConversionResult CompressAndWrite(File::IOFile* file, u64* bytes_written,
+                                           Compressor* compressor, const u8* data, size_t size,
+                                           size_t* size_out);
+
+  template <typename T>
+  static void PushBack(std::vector<u8>* vector, const T& x)
+  {
+    const size_t offset_in_vector = vector->size();
+    vector->resize(offset_in_vector + sizeof(T));
+    const u8* x_ptr = reinterpret_cast<const u8*>(&x);
+    std::copy(x_ptr, x_ptr + sizeof(T), vector->data() + offset_in_vector);
+  }
 
   bool m_valid;
-  CompressionType m_compression_type;
+  WIACompressionType m_compression_type;
 
   File::IOFile m_file;
   Chunk m_cached_chunk;

Source/Core/DolphinQt/ConvertDialog.cpp

@@ -26,6 +26,7 @@
 #include "Common/Logging/Log.h"
 #include "DiscIO/Blob.h"
 #include "DiscIO/ScrubbedBlob.h"
+#include "DiscIO/WIABlob.h"
 #include "DolphinQt/QtUtils/ModalMessageBox.h"
 #include "DolphinQt/QtUtils/ParallelProgressDialog.h"
 #include "UICommon/GameFile.h"
@@ -70,9 +71,17 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi
   grid_layout->addWidget(new QLabel(tr("Block Size:")), 1, 0);
   grid_layout->addWidget(m_block_size, 1, 1);
 
+  m_compression = new QComboBox;
+  grid_layout->addWidget(new QLabel(tr("Compression:")), 2, 0);
+  grid_layout->addWidget(m_compression, 2, 1);
+
+  m_compression_level = new QComboBox;
+  grid_layout->addWidget(new QLabel(tr("Compression Level:")), 3, 0);
+  grid_layout->addWidget(m_compression_level, 3, 1);
+
   m_scrub = new QCheckBox;
-  grid_layout->addWidget(new QLabel(tr("Remove Junk Data (Irreversible):")), 2, 0);
-  grid_layout->addWidget(m_scrub, 2, 1);
+  grid_layout->addWidget(new QLabel(tr("Remove Junk Data (Irreversible):")), 4, 0);
+  grid_layout->addWidget(m_scrub, 4, 1);
   m_scrub->setEnabled(
       std::none_of(m_files.begin(), m_files.end(), std::mem_fn(&UICommon::GameFile::IsDatelDisc)));
 
@@ -108,9 +117,12 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi
 
   connect(m_format, QOverload<int>::of(&QComboBox::currentIndexChanged), this,
           &ConvertDialog::OnFormatChanged);
+  connect(m_compression, QOverload<int>::of(&QComboBox::currentIndexChanged), this,
+          &ConvertDialog::OnCompressionChanged);
   connect(convert_button, &QPushButton::clicked, this, &ConvertDialog::Convert);
 
   OnFormatChanged();
+  OnCompressionChanged();
 }
 
 void ConvertDialog::AddToBlockSizeComboBox(int size)
@@ -118,6 +130,16 @@ void ConvertDialog::AddToBlockSizeComboBox(int size)
   m_block_size->addItem(QString::fromStdString(UICommon::FormatSize(size, 0)), size);
 }
 
+void ConvertDialog::AddToCompressionComboBox(const QString& name, DiscIO::WIACompressionType type)
+{
+  m_compression->addItem(name, static_cast<int>(type));
+}
+
+void ConvertDialog::AddToCompressionLevelComboBox(int level)
+{
+  m_compression_level->addItem(QString::number(level), level);
+}
+
 void ConvertDialog::OnFormatChanged()
 {
   // Because DVD timings are emulated as if we can't read less than an entire ECC block at once
@@ -131,6 +153,9 @@ void ConvertDialog::OnFormatChanged()
   const DiscIO::BlobType format = static_cast<DiscIO::BlobType>(m_format->currentData().toInt());
 
   m_block_size->clear();
+  m_compression->clear();
+
+  // Populate m_block_size
   switch (format)
   {
   case DiscIO::BlobType::GCZ:
@@ -181,7 +206,57 @@ void ConvertDialog::OnFormatChanged()
     break;
   }
 
+  // Populate m_compression
+  switch (format)
+  {
+  case DiscIO::BlobType::GCZ:
+    m_compression->setEnabled(true);
+    AddToCompressionComboBox(QStringLiteral("Deflate"), DiscIO::WIACompressionType::None);
+    break;
+  case DiscIO::BlobType::WIA:
+  {
+    m_compression->setEnabled(true);
+
+    // i18n: %1 is the name of a compression method (e.g. LZMA)
+    const QString slow = tr("%1 (slow)");
+
+    AddToCompressionComboBox(tr("No Compression"), DiscIO::WIACompressionType::None);
+    AddToCompressionComboBox(QStringLiteral("Purge"), DiscIO::WIACompressionType::Purge);
+    AddToCompressionComboBox(slow.arg(QStringLiteral("bzip2")), DiscIO::WIACompressionType::Bzip2);
+    AddToCompressionComboBox(slow.arg(QStringLiteral("LZMA")), DiscIO::WIACompressionType::LZMA);
+    AddToCompressionComboBox(slow.arg(QStringLiteral("LZMA2")), DiscIO::WIACompressionType::LZMA2);
+
+    break;
+  }
+  default:
+    m_compression->setEnabled(false);
+    break;
+  }
+
   m_block_size->setEnabled(m_block_size->count() > 1);
+  m_compression->setEnabled(m_compression->count() > 1);
+}
+
+void ConvertDialog::OnCompressionChanged()
+{
+  m_compression_level->clear();
+
+  switch (static_cast<DiscIO::WIACompressionType>(m_compression->currentData().toInt()))
+  {
+  case DiscIO::WIACompressionType::Bzip2:
+  case DiscIO::WIACompressionType::LZMA:
+  case DiscIO::WIACompressionType::LZMA2:
+    for (int i = 1; i <= 9; ++i)
+      AddToCompressionLevelComboBox(i);
+
+    m_compression_level->setCurrentIndex(4);
+
+    break;
+  default:
+    break;
+  }
+
+  m_compression_level->setEnabled(m_compression_level->count() > 1);
 }
 
 bool ConvertDialog::ShowAreYouSureDialog(const QString& text)
@@ -200,6 +275,9 @@ void ConvertDialog::Convert()
 {
   const DiscIO::BlobType format = static_cast<DiscIO::BlobType>(m_format->currentData().toInt());
   const int block_size = m_block_size->currentData().toInt();
+  const DiscIO::WIACompressionType compression =
+      static_cast<DiscIO::WIACompressionType>(m_compression->currentData().toInt());
+  const int compression_level = m_compression_level->currentData().toInt();
   const bool scrub = m_scrub->isChecked();
 
   if (scrub && format == DiscIO::BlobType::PLAIN)
@@ -369,9 +447,9 @@ void ConvertDialog::Convert()
       else if (format == DiscIO::BlobType::WIA)
       {
         good = std::async(std::launch::async, [&] {
-          const bool good =
-              DiscIO::ConvertToWIA(blob_reader.get(), original_path, dst_path.toStdString(),
-                                   block_size, &CompressCB, &progress_dialog);
+          const bool good = DiscIO::ConvertToWIA(
+              blob_reader.get(), original_path, dst_path.toStdString(), compression,
+              compression_level, block_size, &CompressCB, &progress_dialog);
           progress_dialog.Reset();
           return good;
         });

Source/Core/DolphinQt/ConvertDialog.h

@@ -14,6 +14,11 @@
 class QCheckBox;
 class QComboBox;
 
+namespace DiscIO
+{
+enum class WIACompressionType : u32;
+}
+
 namespace UICommon
 {
 class GameFile;
@@ -29,15 +34,20 @@ public:
 
 private slots:
   void OnFormatChanged();
+  void OnCompressionChanged();
   void Convert();
 
 private:
   void AddToBlockSizeComboBox(int size);
+  void AddToCompressionComboBox(const QString& name, DiscIO::WIACompressionType type);
+  void AddToCompressionLevelComboBox(int level);
 
   bool ShowAreYouSureDialog(const QString& text);
 
   QComboBox* m_format;
   QComboBox* m_block_size;
+  QComboBox* m_compression;
+  QComboBox* m_compression_level;
   QCheckBox* m_scrub;
   QList<std::shared_ptr<const UICommon::GameFile>> m_files;
 };