Improve NAND import progress dialog

Now with cancel button and an actual progress bar. For simplicity, we do
two passes on the progress bar, one for loading the NAND into memory and
one for extracting it. The user directory is likely on an SSD, making
the extraction pass invisibly fast.

2 files changed, 38 insertions, 14 deletions

diff --git a/Source/Core/DiscIO/NANDImporter.cpp b/Source/Core/DiscIO/NANDImporter.cpp
index 54122d0626..f4c94bbb2c 100644
--- a/Source/Core/DiscIO/NANDImporter.cpp
+++ b/Source/Core/DiscIO/NANDImporter.cpp
@@ -23,8 +23,7 @@ NANDImporter::NANDImporter() : m_nand_root(File::GetUserPath(D_WIIROOT_IDX))
 }
 NANDImporter::~NANDImporter() = default;
 
-void NANDImporter::ImportNANDBin(const std::string& path_to_bin,
-                                 std::function<void()> update_callback,
+void NANDImporter::ImportNANDBin(const std::string& path_to_bin, UpdateCallback update_callback,
                                  const std::function<std::string()>& get_otp_dump_path)
 {
   m_update_callback = std::move(update_callback);
@@ -35,8 +34,13 @@ void NANDImporter::ImportNANDBin(const std::string& path_to_bin,
     return;
 
   ExportKeys();
-  ProcessEntry(0, "");
+
+  if (!ExtractFiles())
+    return;
+
   ExtractCertificates();
+
+  std::ignore = m_update_callback(Step::Extracting, m_progress_cur, m_progress_max);
 }
 
 bool NANDImporter::ReadNANDBin(const std::string& path_to_bin,
@@ -60,10 +64,8 @@ bool NANDImporter::ReadNANDBin(const std::string& path_to_bin,
 
   for (size_t i = 0; i < NAND_TOTAL_BLOCKS; i++)
   {
-    // Instead of updating on every cycle, we only update every 1000 cycles for a balance between
-    // not updating fast enough vs updating too fast
-    if (i % 1000 == 0)
-      m_update_callback();
+    if (m_update_callback(Step::Loading, int(i), int(NAND_TOTAL_BLOCKS)))
+      return false;
 
     file.ReadBytes(&m_nand[i * NAND_BLOCK_SIZE], NAND_BLOCK_SIZE);
 
@@ -121,6 +123,14 @@ bool NANDImporter::FindSuperblock()
   return true;
 }
 
+bool NANDImporter::ExtractFiles()
+{
+  constexpr u16 CLUSTER_CHAIN_END = 0xFFFB;
+  m_progress_cur = 0;
+  m_progress_max = std::ranges::count(m_superblock->fat, CLUSTER_CHAIN_END);
+  return ProcessEntry(0, "");
+}
+
 std::string NANDImporter::GetPath(const NANDFSTEntry& entry, const std::string& parent_path)
 {
   std::string name(entry.name, strnlen(entry.name, sizeof(NANDFSTEntry::name)));
@@ -131,25 +141,26 @@ std::string NANDImporter::GetPath(const NANDFSTEntry& entry, const std::string&
   return parent_path + '/' + name;
 }
 
-void NANDImporter::ProcessEntry(u16 entry_number, const std::string& parent_path)
+bool NANDImporter::ProcessEntry(u16 entry_number, const std::string& parent_path)
 {
   while (entry_number != 0xffff)
   {
     if (entry_number >= m_superblock->fst.size())
     {
       ERROR_LOG_FMT(DISCIO, "FST entry number {} out of range", entry_number);
-      return;
+      return false;
     }
 
     const NANDFSTEntry entry = m_superblock->fst[entry_number];
 
     const std::string path = GetPath(entry, parent_path);
     INFO_LOG_FMT(DISCIO, "Entry: {} Path: {}", entry, path);
-    m_update_callback();
 
     Type type = static_cast<Type>(entry.mode & 3);
     if (type == Type::File)
     {
+      if (m_update_callback(Step::Extracting, m_progress_cur++, m_progress_max))
+        return false;
       std::vector<u8> data = GetEntryData(entry);
       File::IOFile file(m_nand_root + path, "wb");
       file.WriteBytes(data.data(), data.size());
@@ -157,7 +168,8 @@ void NANDImporter::ProcessEntry(u16 entry_number, const std::string& parent_path
     else if (type == Type::Directory)
     {
       File::CreateDir(m_nand_root + path);
-      ProcessEntry(entry.sub, path);
+      if (!ProcessEntry(entry.sub, path))
+        return false;
     }
     else
     {
@@ -166,6 +178,7 @@ void NANDImporter::ProcessEntry(u16 entry_number, const std::string& parent_path
 
     entry_number = entry.sib;
   }
+  return true;
 }
 
 std::vector<u8> NANDImporter::GetEntryData(const NANDFSTEntry& entry) const
diff --git a/Source/Core/DiscIO/NANDImporter.h b/Source/Core/DiscIO/NANDImporter.h
index 164d009f5a..f22f043709 100644
--- a/Source/Core/DiscIO/NANDImporter.h
+++ b/Source/Core/DiscIO/NANDImporter.h
@@ -23,10 +23,18 @@ public:
   NANDImporter();
   ~NANDImporter();
 
+  enum class Step
+  {
+    Loading,
+    Extracting,
+  };
+  // Return true to cancel.
+  using UpdateCallback = std::function<bool(Step step, int cur, int max)>;
+
   // Extract a NAND image to the configured NAND root.
   // If the associated OTP/SEEPROM dump (keys.bin) is not included in the image,
   // get_otp_dump_path will be called to get a path to it.
-  void ImportNANDBin(const std::string& path_to_bin, std::function<void()> update_callback,
+  void ImportNANDBin(const std::string& path_to_bin, UpdateCallback update_callback,
                      const std::function<std::string()>& get_otp_dump_path);
   bool ExtractCertificates();
 
@@ -67,9 +75,10 @@ private:
   bool ReadNANDBin(const std::string& path_to_bin,
                    const std::function<std::string()>& get_otp_dump_path);
   bool FindSuperblock();
+  bool ExtractFiles();
   std::string GetPath(const NANDFSTEntry& entry, const std::string& parent_path);
   std::string FormatDebugString(const NANDFSTEntry& entry);
-  void ProcessEntry(u16 entry_number, const std::string& parent_path);
+  bool ProcessEntry(u16 entry_number, const std::string& parent_path);
   std::vector<u8> GetEntryData(const NANDFSTEntry& entry) const;
   void ExportKeys();
 
@@ -78,7 +87,9 @@ private:
   std::vector<u8> m_nand_keys;
   std::unique_ptr<Common::AES::Context> m_aes_ctx;
   std::unique_ptr<NANDSuperblock> m_superblock;
-  std::function<void()> m_update_callback;
+  UpdateCallback m_update_callback;
+  u16 m_progress_cur = 0;
+  u16 m_progress_max = 0;
 };
 }  // namespace DiscIO