vfs_concat: Rewrite and fix ConcatenatedVfsFile

2 files changed, 59 insertions, 14 deletions

diff --git a/src/core/file_sys/vfs_concat.cpp b/src/core/file_sys/vfs_concat.cpp
index dc7a279a9..0c07e162e 100644
--- a/src/core/file_sys/vfs_concat.cpp
+++ b/src/core/file_sys/vfs_concat.cpp
@@ -5,10 +5,22 @@
 #include <algorithm>
 #include <utility>
 
+#include "common/assert.h"
 #include "core/file_sys/vfs_concat.h"
 
 namespace FileSys {
 
+bool VerifyConcatenationMap(std::map<u64, VirtualFile> map) {
+    for (auto iter = map.begin(); iter != --map.end();) {
+        const auto old = iter++;
+        if (old->first + old->second->GetSize() != iter->first) {
+            return false;
+        }
+    }
+
+    return map.begin()->first == 0;
+}
+
 VirtualFile ConcatenateFiles(std::vector<VirtualFile> files, std::string name) {
     if (files.empty())
         return nullptr;
@@ -27,7 +39,10 @@ ConcatenatedVfsFile::ConcatenatedVfsFile(std::vector<VirtualFile> files_, std::s
     }
 }
 
-ConcatenatedVfsFile::~ConcatenatedVfsFile() = default;
+ConcatenatedVfsFile::ConcatenatedVfsFile(std::map<u64, VirtualFile> files_, std::string name)
+    : files(std::move(files_)), name(std::move(name)) {
+    ASSERT(VerifyConcatenationMap(files));
+}
 
 std::string ConcatenatedVfsFile::GetName() const {
     if (files.empty())
@@ -62,28 +77,25 @@ bool ConcatenatedVfsFile::IsReadable() const {
 }
 
 std::size_t ConcatenatedVfsFile::Read(u8* data, std::size_t length, std::size_t offset) const {
-    auto entry = files.end();
+    std::pair<u64, VirtualFile> entry = *files.rbegin();
     for (auto iter = files.begin(); iter != files.end(); ++iter) {
         if (iter->first > offset) {
-            entry = --iter;
+            entry = *--iter;
             break;
         }
     }
 
-    // Check if the entry should be the last one. The loop above will make it end().
-    if (entry == files.end() && offset < files.rbegin()->first + files.rbegin()->second->GetSize())
-        --entry;
-
-    if (entry == files.end())
+    if (entry.first + entry.second->GetSize() <= offset)
         return 0;
 
-    const auto remaining = entry->second->GetSize() + offset - entry->first;
-    if (length > remaining) {
-        return entry->second->Read(data, remaining, offset - entry->first) +
-               Read(data + remaining, length - remaining, offset + remaining);
+    const auto read_in =
+        std::min(entry.first + entry.second->GetSize() - offset, entry.second->GetSize());
+    if (length > read_in) {
+        return entry.second->Read(data, read_in, offset - entry.first) +
+               Read(data + read_in, length - read_in, offset + read_in);
     }
 
-    return entry->second->Read(data, length, offset - entry->first);
+    return entry.second->Read(data, std::min(read_in, length), offset - entry.first);
 }
 
 std::size_t ConcatenatedVfsFile::Write(const u8* data, std::size_t length, std::size_t offset) {
@@ -93,4 +105,5 @@ std::size_t ConcatenatedVfsFile::Write(const u8* data, std::size_t length, std::
 bool ConcatenatedVfsFile::Rename(std::string_view name) {
     return false;
 }
+
 } // namespace FileSys
diff --git a/src/core/file_sys/vfs_concat.h b/src/core/file_sys/vfs_concat.h
index 717d04bdc..c65c20d15 100644
--- a/src/core/file_sys/vfs_concat.h
+++ b/src/core/file_sys/vfs_concat.h
@@ -4,22 +4,54 @@
 
 #pragma once
 
+#include <map>
 #include <memory>
 #include <string_view>
 #include <boost/container/flat_map.hpp>
 #include "core/file_sys/vfs.h"
+#include "core/file_sys/vfs_static.h"
 
 namespace FileSys {
 
+class ConcatenatedVfsFile;
+
 // Wrapper function to allow for more efficient handling of files.size() == 0, 1 cases.
 VirtualFile ConcatenateFiles(std::vector<VirtualFile> files, std::string name = "");
 
+// Convenience function that turns a map of offsets to files into a concatenated file, filling gaps
+// with template parameter.
+template <u8 filler_byte>
+VirtualFile ConcatenateFiles(std::map<u64, VirtualFile> files, std::string name = "") {
+    if (files.empty())
+        return nullptr;
+    if (files.size() == 1)
+        return files.begin()->second;
+
+    for (auto iter = files.begin(); iter != --files.end();) {
+        const auto old = iter++;
+        if (old->first + old->second->GetSize() != iter->first) {
+            files.emplace(old->first + old->second->GetSize(),
+                          std::make_shared<StaticVfsFile<filler_byte>>(iter->first - old->first -
+                                                                       old->second->GetSize()));
+        }
+    }
+
+    if (files.begin()->first != 0)
+        files.emplace(0, std::make_shared<StaticVfsFile<filler_byte>>(files.begin()->first));
+
+    return std::shared_ptr<VfsFile>(new ConcatenatedVfsFile(std::move(files), std::move(name)));
+}
+
 // Class that wraps multiple vfs files and concatenates them, making reads seamless. Currently
 // read-only.
 class ConcatenatedVfsFile : public VfsFile {
     friend VirtualFile ConcatenateFiles(std::vector<VirtualFile> files, std::string name);
 
+    template <u8 filler_byte>
+    friend VirtualFile ConcatenateFiles(std::map<u64, VirtualFile> files, std::string name);
+
     ConcatenatedVfsFile(std::vector<VirtualFile> files, std::string name);
+    ConcatenatedVfsFile(std::map<u64, VirtualFile> files, std::string name);
 
 public:
     ~ConcatenatedVfsFile() override;
@@ -36,7 +68,7 @@ public:
 
 private:
     // Maps starting offset to file -- more efficient.
-    boost::container::flat_map<u64, VirtualFile> files;
+    std::map<u64, VirtualFile> files;
     std::string name;
 };