DirectoryBlob: Prepare more for supporting multiple partitions

2 files changed, 31 insertions, 11 deletions

diff --git a/Source/Core/DiscIO/DirectoryBlob.cpp b/Source/Core/DiscIO/DirectoryBlob.cpp
index e8db47d29a..0b16c4da7b 100644
--- a/Source/Core/DiscIO/DirectoryBlob.cpp
+++ b/Source/Core/DiscIO/DirectoryBlob.cpp
@@ -10,6 +10,7 @@
 #include <cstddef>
 #include <cstring>
 #include <locale>
+#include <map>
 #include <memory>
 #include <set>
 #include <string>
@@ -156,13 +157,21 @@ std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(const std::stri
 }
 
 DirectoryBlobReader::DirectoryBlobReader(const std::string& root_directory)
-    : m_root_directory(root_directory), m_game_partition(root_directory, {})
+    : m_root_directory(root_directory)
 {
-  m_is_wii = m_game_partition.IsWii();
+  DirectoryBlobPartition game_partition(root_directory, {});
+  m_is_wii = game_partition.IsWii();
 
-  if (m_is_wii)
+  if (!m_is_wii)
+  {
+    m_gamecube_pseudopartition = std::move(game_partition);
+  }
+  else
   {
-    SetNonpartitionDiscHeader(m_game_partition.GetHeader());
+    SetNonpartitionDiscHeader(game_partition.GetHeader());
+
+    m_partitions.emplace(GAME_PARTITION_ADDRESS, std::move(game_partition));
+
     SetPartitionTable();
     SetWiiRegionData();
     SetTMDAndTicket();
@@ -205,8 +214,9 @@ bool DirectoryBlobReader::Read(u64 offset, u64 length, u8* buffer)
 {
   // TODO: We don't handle raw access to the encrypted area of Wii discs correctly.
 
-  return ReadInternal(offset, length, buffer,
-                      m_is_wii ? m_nonpartition_contents : m_game_partition.GetContents());
+  const std::set<DiscContent>& contents =
+      m_is_wii ? m_nonpartition_contents : m_gamecube_pseudopartition.GetContents();
+  return ReadInternal(offset, length, buffer, contents);
 }
 
 bool DirectoryBlobReader::SupportsReadWiiDecrypted() const
@@ -216,10 +226,14 @@ bool DirectoryBlobReader::SupportsReadWiiDecrypted() const
 
 bool DirectoryBlobReader::ReadWiiDecrypted(u64 offset, u64 size, u8* buffer, u64 partition_offset)
 {
-  if (!m_is_wii || partition_offset != GAME_PARTITION_ADDRESS)
+  if (!m_is_wii)
+    return false;
+
+  auto it = m_partitions.find(partition_offset);
+  if (it == m_partitions.end())
     return false;
 
-  return ReadInternal(offset, size, buffer, m_game_partition.GetContents());
+  return ReadInternal(offset, size, buffer, it->second.GetContents());
 }
 
 BlobType DirectoryBlobReader::GetBlobType() const
diff --git a/Source/Core/DiscIO/DirectoryBlob.h b/Source/Core/DiscIO/DirectoryBlob.h
index 647ac11380..23048b64b0 100644
--- a/Source/Core/DiscIO/DirectoryBlob.h
+++ b/Source/Core/DiscIO/DirectoryBlob.h
@@ -5,6 +5,7 @@
 #pragma once
 
 #include <cstddef>
+#include <map>
 #include <memory>
 #include <optional>
 #include <set>
@@ -55,6 +56,7 @@ private:
 class DirectoryBlobPartition
 {
 public:
+  DirectoryBlobPartition() = default;
   DirectoryBlobPartition(const std::string& root_directory, std::optional<bool> is_wii);
 
   // We do not allow copying, because it might mess up the pointers inside DiscContents
@@ -90,9 +92,9 @@ private:
   std::vector<u8> m_fst_data;
 
   std::string m_root_directory;
-  bool m_is_wii;
+  bool m_is_wii = false;
   // GameCube has no shift, Wii has 2 bit shift
-  u32 m_address_shift;
+  u32 m_address_shift = 0;
 };
 
 class DirectoryBlobReader : public BlobReader
@@ -126,8 +128,12 @@ private:
 
   std::string m_root_directory;
 
-  DirectoryBlobPartition m_game_partition;
+  // For GameCube:
+  DirectoryBlobPartition m_gamecube_pseudopartition;
+
+  // For Wii:
   std::set<DiscContent> m_nonpartition_contents;
+  std::map<u64, DirectoryBlobPartition> m_partitions;
 
   bool m_is_wii;