3 files changed, 124 insertions, 68 deletions

diff --git a/src/core/file_sys/romfs.cpp b/src/core/file_sys/romfs.cpp
index ddcfe5980..fb5683a6b 100644
--- a/src/core/file_sys/romfs.cpp
+++ b/src/core/file_sys/romfs.cpp
@@ -140,7 +140,8 @@ VirtualFile CreateRomFS(VirtualDir dir, VirtualDir ext) {
         return nullptr;
 
     RomFSBuildContext ctx{dir, ext};
-    return ConcatenatedVfsFile::MakeConcatenatedFile(0, ctx.Build(), dir->GetName());
+    auto file_map = ctx.Build();
+    return ConcatenatedVfsFile::MakeConcatenatedFile(0, file_map, dir->GetName());
 }
 
 } // namespace FileSys
diff --git a/src/core/file_sys/vfs_concat.cpp b/src/core/file_sys/vfs_concat.cpp
index d23623aa0..853b893a1 100644
--- a/src/core/file_sys/vfs_concat.cpp
+++ b/src/core/file_sys/vfs_concat.cpp
@@ -10,84 +10,105 @@
 
 namespace FileSys {
 
-static bool VerifyConcatenationMapContinuity(const std::multimap<u64, VirtualFile>& map) {
-    const auto last_valid = --map.end();
-    for (auto iter = map.begin(); iter != last_valid;) {
-        const auto old = iter++;
-        if (old->first + old->second->GetSize() != iter->first) {
+ConcatenatedVfsFile::ConcatenatedVfsFile(ConcatenationMap&& concatenation_map_, std::string&& name_)
+    : concatenation_map(std::move(concatenation_map_)), name(std::move(name_)) {
+    DEBUG_ASSERT(this->VerifyContinuity());
+}
+
+bool ConcatenatedVfsFile::VerifyContinuity() const {
+    u64 last_offset = 0;
+    for (auto& entry : concatenation_map) {
+        if (entry.offset != last_offset) {
             return false;
         }
-    }
-
-    return map.begin()->first == 0;
-}
 
-ConcatenatedVfsFile::ConcatenatedVfsFile(std::vector<VirtualFile> files_, std::string name_)
-    : name(std::move(name_)) {
-    std::size_t next_offset = 0;
-    for (const auto& file : files_) {
-        files.emplace(next_offset, file);
-        next_offset += file->GetSize();
+        last_offset = entry.offset + entry.file->GetSize();
     }
-}
 
-ConcatenatedVfsFile::ConcatenatedVfsFile(std::multimap<u64, VirtualFile> files_, std::string name_)
-    : files(std::move(files_)), name(std::move(name_)) {
-    ASSERT(VerifyConcatenationMapContinuity(files));
+    return true;
 }
 
 ConcatenatedVfsFile::~ConcatenatedVfsFile() = default;
 
-VirtualFile ConcatenatedVfsFile::MakeConcatenatedFile(std::vector<VirtualFile> files,
-                                                      std::string name) {
-    if (files.empty())
+VirtualFile ConcatenatedVfsFile::MakeConcatenatedFile(const std::vector<VirtualFile>& files,
+                                                      std::string&& name) {
+    // Fold trivial cases.
+    if (files.empty()) {
         return nullptr;
-    if (files.size() == 1)
-        return files[0];
+    }
+    if (files.size() == 1) {
+        return files.front();
+    }
+
+    // Make the concatenation map from the input.
+    std::vector<ConcatenationEntry> concatenation_map;
+    concatenation_map.reserve(files.size());
+    u64 last_offset = 0;
+
+    for (auto& file : files) {
+        concatenation_map.emplace_back(ConcatenationEntry{
+            .offset = last_offset,
+            .file = file,
+        });
+
+        last_offset += file->GetSize();
+    }
 
-    return VirtualFile(new ConcatenatedVfsFile(std::move(files), std::move(name)));
+    return VirtualFile(new ConcatenatedVfsFile(std::move(concatenation_map), std::move(name)));
 }
 
 VirtualFile ConcatenatedVfsFile::MakeConcatenatedFile(u8 filler_byte,
-                                                      std::multimap<u64, VirtualFile> files,
-                                                      std::string name) {
-    if (files.empty())
+                                                      const std::multimap<u64, VirtualFile>& files,
+                                                      std::string&& name) {
+    // Fold trivial cases.
+    if (files.empty()) {
         return nullptr;
-    if (files.size() == 1)
+    }
+    if (files.size() == 1) {
         return files.begin()->second;
+    }
 
-    const auto last_valid = --files.end();
-    for (auto iter = files.begin(); iter != last_valid;) {
-        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()));
+    // Make the concatenation map from the input.
+    std::vector<ConcatenationEntry> concatenation_map;
+
+    concatenation_map.reserve(files.size());
+    u64 last_offset = 0;
+
+    // Iteration of a multimap is ordered, so offset will be strictly non-decreasing.
+    for (auto& [offset, file] : files) {
+        if (offset > last_offset) {
+            concatenation_map.emplace_back(ConcatenationEntry{
+                .offset = last_offset,
+                .file = std::make_shared<StaticVfsFile>(filler_byte, offset - last_offset),
+            });
         }
-    }
 
-    // Ensure the map starts at offset 0 (start of file), otherwise pad to fill.
-    if (files.begin()->first != 0)
-        files.emplace(0, std::make_shared<StaticVfsFile>(filler_byte, files.begin()->first));
+        concatenation_map.emplace_back(ConcatenationEntry{
+            .offset = offset,
+            .file = file,
+        });
+
+        last_offset = offset + file->GetSize();
+    }
 
-    return VirtualFile(new ConcatenatedVfsFile(std::move(files), std::move(name)));
+    return VirtualFile(new ConcatenatedVfsFile(std::move(concatenation_map), std::move(name)));
 }
 
 std::string ConcatenatedVfsFile::GetName() const {
-    if (files.empty()) {
+    if (concatenation_map.empty()) {
         return "";
     }
     if (!name.empty()) {
         return name;
     }
-    return files.begin()->second->GetName();
+    return concatenation_map.front().file->GetName();
 }
 
 std::size_t ConcatenatedVfsFile::GetSize() const {
-    if (files.empty()) {
+    if (concatenation_map.empty()) {
         return 0;
     }
-    return files.rbegin()->first + files.rbegin()->second->GetSize();
+    return concatenation_map.back().offset + concatenation_map.back().file->GetSize();
 }
 
 bool ConcatenatedVfsFile::Resize(std::size_t new_size) {
@@ -95,10 +116,10 @@ bool ConcatenatedVfsFile::Resize(std::size_t new_size) {
 }
 
 VirtualDir ConcatenatedVfsFile::GetContainingDirectory() const {
-    if (files.empty()) {
+    if (concatenation_map.empty()) {
         return nullptr;
     }
-    return files.begin()->second->GetContainingDirectory();
+    return concatenation_map.front().file->GetContainingDirectory();
 }
 
 bool ConcatenatedVfsFile::IsWritable() const {
@@ -110,25 +131,45 @@ bool ConcatenatedVfsFile::IsReadable() const {
 }
 
 std::size_t ConcatenatedVfsFile::Read(u8* data, std::size_t length, std::size_t offset) const {
-    auto entry = --files.end();
-    for (auto iter = files.begin(); iter != files.end(); ++iter) {
-        if (iter->first > offset) {
-            entry = --iter;
+    const ConcatenationEntry key{
+        .offset = offset,
+        .file = nullptr,
+    };
+
+    // Read nothing if the map is empty.
+    if (concatenation_map.empty()) {
+        return 0;
+    }
+
+    // Binary search to find the iterator to the first position we can check.
+    // It must exist, since we are not empty and are comparing unsigned integers.
+    auto it = std::prev(std::upper_bound(concatenation_map.begin(), concatenation_map.end(), key));
+    u64 cur_length = length;
+    u64 cur_offset = offset;
+
+    while (cur_length > 0 && it != concatenation_map.end()) {
+        // Check if we can read the file at this position.
+        const auto& file = it->file;
+        const u64 file_offset = it->offset;
+        const u64 file_size = file->GetSize();
+
+        if (cur_offset >= file_offset + file_size) {
+            // Entirely out of bounds read.
             break;
         }
-    }
 
-    if (entry->first + entry->second->GetSize() <= offset)
-        return 0;
+        // Read the file at this position.
+        const u64 intended_read_size = std::min<u64>(cur_length, file_size);
+        const u64 actual_read_size =
+            file->Read(data + (cur_offset - offset), intended_read_size, cur_offset - file_offset);
 
-    const auto read_in =
-        std::min<u64>(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);
+        // Update tracking.
+        cur_offset += actual_read_size;
+        cur_length -= actual_read_size;
+        it++;
     }
 
-    return entry->second->Read(data, std::min<u64>(read_in, length), offset - entry->first);
+    return cur_offset - offset;
 }
 
 std::size_t ConcatenatedVfsFile::Write(const u8* data, std::size_t length, std::size_t offset) {
diff --git a/src/core/file_sys/vfs_concat.h b/src/core/file_sys/vfs_concat.h
index 9be0261b6..6b329d545 100644
--- a/src/core/file_sys/vfs_concat.h
+++ b/src/core/file_sys/vfs_concat.h
@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <compare>
 #include <map>
 #include <memory>
 #include "core/file_sys/vfs.h"
@@ -12,19 +13,33 @@ namespace FileSys {
 // Class that wraps multiple vfs files and concatenates them, making reads seamless. Currently
 // read-only.
 class ConcatenatedVfsFile : public VfsFile {
-    explicit ConcatenatedVfsFile(std::vector<VirtualFile> files, std::string name_);
-    explicit ConcatenatedVfsFile(std::multimap<u64, VirtualFile> files, std::string name_);
+private:
+    struct ConcatenationEntry {
+        u64 offset;
+        VirtualFile file;
+
+        auto operator<=>(const ConcatenationEntry& other) const {
+            return this->offset <=> other.offset;
+        }
+    };
+    using ConcatenationMap = std::vector<ConcatenationEntry>;
+
+    explicit ConcatenatedVfsFile(std::vector<ConcatenationEntry>&& concatenation_map,
+                                 std::string&& name);
+    bool VerifyContinuity() const;
 
 public:
     ~ConcatenatedVfsFile() override;
 
     /// Wrapper function to allow for more efficient handling of files.size() == 0, 1 cases.
-    static VirtualFile MakeConcatenatedFile(std::vector<VirtualFile> files, std::string name);
+    static VirtualFile MakeConcatenatedFile(const std::vector<VirtualFile>& files,
+                                            std::string&& name);
 
     /// Convenience function that turns a map of offsets to files into a concatenated file, filling
     /// gaps with a given filler byte.
-    static VirtualFile MakeConcatenatedFile(u8 filler_byte, std::multimap<u64, VirtualFile> files,
-                                            std::string name);
+    static VirtualFile MakeConcatenatedFile(u8 filler_byte,
+                                            const std::multimap<u64, VirtualFile>& files,
+                                            std::string&& name);
 
     std::string GetName() const override;
     std::size_t GetSize() const override;
@@ -37,8 +52,7 @@ public:
     bool Rename(std::string_view new_name) override;
 
 private:
-    // Maps starting offset to file -- more efficient.
-    std::multimap<u64, VirtualFile> files;
+    ConcatenationMap concatenation_map;
     std::string name;
 };