diff --git a/src/core/crypto/ncz_decompressor.cpp b/src/core/crypto/ncz_decompressor.cpp index 5f56e55375..a183f7ec1e 100644 --- a/src/core/crypto/ncz_decompressor.cpp +++ b/src/core/crypto/ncz_decompressor.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -120,9 +121,13 @@ std::optional, std::optional>> Parse return std::make_pair(std::move(sections), std::move(bh)); } -// Sequential reader over independently-compressed blocks, one block cached (like nsz_rs). +// Sequential reader over independently-compressed blocks. Keeps a small LRU cache of recently +// decoded blocks so random reads (e.g. update NCA with thousands of sections) don't re-decompress +// the same block over and over. class BlockReader { public: + static constexpr std::size_t BlockCacheSize = 8; + BlockReader(FileSys::VirtualFile f, const BlockHeader& bh, u64 payload_start) : source{std::move(f)}, block_size{1ULL << bh.exponent}, sizes{bh.sizes}, decompressed_size{bh.decompressed_size} { @@ -139,11 +144,13 @@ public: while (done < out.size() && position < decompressed_size) { const u64 id = position / block_size; if (id >= sizes.size()) break; - if (cache_id != id && !Load(id)) break; + const std::vector* block = GetBlock(id); + if (block == nullptr) break; const u64 in = position - id * block_size; + if (in >= block->size()) break; const std::size_t n = static_cast( - std::min(out.size() - done, cache.size() - in)); - std::memcpy(out.data() + done, cache.data() + in, n); + std::min(out.size() - done, block->size() - in)); + std::memcpy(out.data() + done, block->data() + in, n); done += n; position += n; } @@ -151,23 +158,44 @@ public: } private: - bool Load(u64 id) { + const std::vector* GetBlock(u64 id) { + // LRU hit: move to front. + for (auto it = cache.begin(); it != cache.end(); ++it) { + if (it->first == id) { + cache.splice(cache.begin(), cache, it); + return &cache.front().second; + } + } + + // Miss: decode the block. + std::vector buffer; + if (cache.size() >= BlockCacheSize) { + buffer = std::move(cache.back().second); + cache.pop_back(); + } + if (!DecodeBlock(id, buffer)) { + return nullptr; + } + cache.emplace_front(id, std::move(buffer)); + return &cache.front().second; + } + + bool DecodeBlock(u64 id, std::vector& buffer) { u64 want = block_size; if (id + 1 == sizes.size()) { const u64 rem = decompressed_size % block_size; if (rem > 0) want = rem; // avoid the 0-size bug from the original nsz } - try { cache.resize(static_cast(want)); } + try { buffer.resize(static_cast(want)); } catch (const std::bad_alloc&) { return false; } const u32 csize = sizes[id]; if (csize < want) { const auto comp = source->ReadBytes(csize, offsets[id]); if (comp.size() != csize || - Common::Compression::DecompressDataZSTD(comp, cache) != want) return false; - } else if (source->Read(cache.data(), cache.size(), offsets[id]) != cache.size()) { + Common::Compression::DecompressDataZSTD(comp, buffer) != want) return false; + } else if (source->Read(buffer.data(), buffer.size(), offsets[id]) != buffer.size()) { return false; } - cache_id = id; return true; } @@ -176,8 +204,7 @@ private: std::vector offsets; std::vector sizes; u64 decompressed_size; - u64 cache_id = ~0ULL; - std::vector cache; + std::list>> cache; // LRU: front = most recent }; class BlockNczFile final : public FileSys::VfsFile {