From 8a71fe377a1f0f7c1c1241bbf654785a0536f0dc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Capelle?= <capelle.mikael@gmail.com>
Date: Fri, 29 May 2020 23:12:11 +0200
Subject: [PATCH] Extract a bit more stuff before installation.

---
 src/installer_fomod_csharp.cpp | 68 +++++++++++++++++++---------------
 src/installer_fomod_csharp.h   | 14 ++-----
 2 files changed, 41 insertions(+), 41 deletions(-)

diff --git a/src/installer_fomod_csharp.cpp b/src/installer_fomod_csharp.cpp
index 72cf605..5f77bcc 100644
--- a/src/installer_fomod_csharp.cpp
+++ b/src/installer_fomod_csharp.cpp
@@ -87,27 +87,14 @@ std::shared_ptr<const FileTreeEntry> InstallerFomodCSharp::findInfoFile(std::sha
   return nullptr;
 }
 
-
-void InstallerFomodCSharp::appendImageFiles(std::vector<std::shared_ptr<const FileTreeEntry>>& entries, std::shared_ptr<const IFileTree> tree) const
-{
-  static std::set<QString, FileNameComparator> imageSuffixes{ "png", "jpg", "jpeg", "gif", "bmp" };
-  for (auto entry : *tree) {
-    if (entry->isDir()) {
-      appendImageFiles(entries, entry->astree());
-    }
-    else if (imageSuffixes.count(entry->suffix()) > 0) {
-      entries.push_back(entry);
-    }
-  }
-}
-
-
 bool InstallerFomodCSharp::isArchiveSupported(std::shared_ptr<const MOBase::IFileTree> tree) const {
   return findScriptFile(tree) != nullptr;
 }
 
 InstallerFomodCSharp::EInstallResult InstallerFomodCSharp::install(MOBase::GuessedValue<QString>& modName, std::shared_ptr<MOBase::IFileTree>& tree,
-  QString& version, int& modID) {
+  QString& version, int& modID) 
+{
+  static std::set<QString, FileNameComparator> imageSuffixes{ "png", "jpg", "jpeg", "gif", "bmp" };
 
   // Extract the script file:
   auto scriptFile = findScriptFile(tree);
@@ -115,15 +102,35 @@ InstallerFomodCSharp::EInstallResult InstallerFomodCSharp::install(MOBase::Guess
     return EInstallResult::RESULT_NOTATTEMPTED;
   }
 
-  std::vector<std::shared_ptr<const FileTreeEntry>> toExtract{ scriptFile };
-
   // Check if there is a info.xml:
   auto infoFile = findInfoFile(tree);
-  if (infoFile != nullptr) {
-    toExtract.push_back(infoFile);
-  }
-  appendImageFiles(toExtract, tree);
 
+  // Set containing everything to extract except the script and the info file:
+  std::set<std::shared_ptr<const FileTreeEntry>> toExtractSet{ scriptFile };
+
+  if (infoFile != nullptr) {
+    toExtractSet.insert(infoFile);
+  }
+
+  // Extract all the images:
+  tree->walk([&](const QString&, auto entry) {
+    if (entry->isFile() && imageSuffixes.count(entry->suffix()) > 0) {
+      toExtractSet.insert(entry);
+    }
+    return IFileTree::WalkReturn::CONTINUE;
+  });
+
+  // Extract everything from the fomod/ folder:
+  auto fomodFolder = findFomodDirectory(tree);
+  fomodFolder->walk([&](const QString&, auto entry) {
+    if (entry->isFile()) {
+      toExtractSet.insert(entry);
+    }
+    return IFileTree::WalkReturn::CONTINUE;
+  });
+
+  // Convert to vector:
+  std::vector toExtract(std::begin(toExtractSet), std::end(toExtractSet));
   QStringList paths(manager()->extractFiles(toExtract));
 
   // If user cancelled:
@@ -131,8 +138,14 @@ InstallerFomodCSharp::EInstallResult InstallerFomodCSharp::install(MOBase::Guess
     return EInstallResult::RESULT_CANCELED;
   }
 
+  // Create a map from entry to file path:
+  std::map<std::shared_ptr<const FileTreeEntry>, QString> entryToPath;
+  for (std::size_t i = 0; i < toExtract.size(); ++i) {
+    entryToPath[toExtract[i]] = paths[i];
+  }
+
   if (infoFile != nullptr) {
-    QFile file(paths[1]);
+    QFile file(entryToPath[infoFile]);
     if (file.open(QIODevice::ReadOnly)) {
       auto info = FomodInfoReader::readXml(file, &FomodInfoReader::parseInfo);
       if (!std::get<0>(info).isEmpty()) {
@@ -147,12 +160,6 @@ InstallerFomodCSharp::EInstallResult InstallerFomodCSharp::install(MOBase::Guess
     }
   }
 
-  // Create a map from entry to file path:
-  std::map<std::shared_ptr<const FileTreeEntry>, QString> entryToPath;
-  for (std::size_t i = 0; i < toExtract.size(); ++i) {
-    entryToPath[toExtract[i]] = paths[i];
-  }
-
   // Show the dialog:
   InstallerFomodPredialog dialog(modName, parentWidget());
   if (dialog.exec() != QDialog::Accepted) {
@@ -171,8 +178,9 @@ InstallerFomodCSharp::EInstallResult InstallerFomodCSharp::install(MOBase::Guess
   modName.update(dialog.getName(), GUESS_USER);
 
   // Run the C# script:
+  const QString scriptPath = entryToPath[scriptFile];
   CSharp::beforeInstall(this, manager(), parentWidget(), std::const_pointer_cast<IFileTree>(scriptFile->parent()->parent()), std::move(entryToPath));
-  auto result = CSharp::executeCSharpScript(paths[0]);
+  auto result = CSharp::executeCSharpScript(scriptPath);
   auto newTree = CSharp::afterInstall(result == EInstallResult::RESULT_SUCCESS);
   if (result == EInstallResult::RESULT_SUCCESS) {
     tree = newTree;
diff --git a/src/installer_fomod_csharp.h b/src/installer_fomod_csharp.h
index a3d46cf..08dfc77 100644
--- a/src/installer_fomod_csharp.h
+++ b/src/installer_fomod_csharp.h
@@ -1,6 +1,3 @@
-#ifndef INSTALLER_FOMOD_CSHARP_H
-#define INSTALLER_FOMOD_CSHARP_H
-
 /*
 Copyright (C) 2020 Holt59. All rights reserved.
 
@@ -18,6 +15,9 @@ You should have received a copy of the GNU General Public License
 along with Mod Organizer.  If not, see <http://www.gnu.org/licenses/>.
 */
 
+#ifndef INSTALLER_FOMOD_CSHARP_H
+#define INSTALLER_FOMOD_CSHARP_H
+
 #include "iplugininstallersimple.h"
 
 class InstallerFomodCSharp  : public MOBase::IPluginInstallerSimple
@@ -82,14 +82,6 @@ private:
   std::shared_ptr<const MOBase::FileTreeEntry> findScriptFile(std::shared_ptr<const MOBase::IFileTree> tree) const;
   std::shared_ptr<const MOBase::FileTreeEntry> findInfoFile(std::shared_ptr<const MOBase::IFileTree> tree) const;
 
-  /**
-   * @brief Recurse through the given tree and add all the images to the given vector.
-   *
-   * @param result Vector of entries to add the images.
-   * @param tree The tree to look files in.
-   */
-  void appendImageFiles(
-    std::vector<std::shared_ptr<const MOBase::FileTreeEntry>>& entries, std::shared_ptr<const MOBase::IFileTree> tree) const;
 };
 
 #endif
\ No newline at end of file