From 24984ea4f2a100fd4e55580c5a444692e7846d4b Mon Sep 17 00:00:00 2001 From: topjohnwu Date: Tue, 23 Nov 2021 18:08:14 -0800 Subject: [PATCH] Optimize stream for full-file writes --- native/jni/magiskboot/bootimg.cpp | 4 ++-- native/jni/magiskboot/compress.cpp | 30 +++++++++++++++-------------- native/jni/magiskboot/compress.hpp | 4 ++-- native/jni/utils/include/stream.hpp | 6 +++++- native/jni/utils/stream.cpp | 29 ++++++++++++++++++++++++---- 5 files changed, 50 insertions(+), 23 deletions(-) diff --git a/native/jni/magiskboot/bootimg.cpp b/native/jni/magiskboot/bootimg.cpp index d0b05b3e7..7aac9afeb 100644 --- a/native/jni/magiskboot/bootimg.cpp +++ b/native/jni/magiskboot/bootimg.cpp @@ -19,14 +19,14 @@ uint64_t dyn_img_hdr::j64 = 0; static void decompress(format_t type, int fd, const void *in, size_t size) { auto ptr = get_decoder(type, make_unique(fd)); - ptr->write(in, size); + ptr->write(in, size, true); } static off_t compress(format_t type, int fd, const void *in, size_t size) { auto prev = lseek(fd, 0, SEEK_CUR); { auto strm = get_encoder(type, make_unique(fd)); - strm->write(in, size); + strm->write(in, size, true); } auto now = lseek(fd, 0, SEEK_CUR); return now - prev; diff --git a/native/jni/magiskboot/compress.cpp b/native/jni/magiskboot/compress.cpp index 8007f5fa2..9845452ea 100644 --- a/native/jni/magiskboot/compress.cpp +++ b/native/jni/magiskboot/compress.cpp @@ -17,7 +17,7 @@ using namespace std; -#define bwrite filter_stream::write +#define bwrite base->write constexpr size_t CHUNK = 0x40000; constexpr size_t LZ4_UNCOMPRESSED = 0x800000; @@ -108,8 +108,7 @@ class zopfli_encoder : public chunk_out_stream { public: explicit zopfli_encoder(stream_ptr &&base) : chunk_out_stream(std::move(base), ZOPFLI_MASTER_BLOCK_SIZE), - zo{}, out(nullptr), outsize(0), crc(crc32_z(0L, Z_NULL, 0)), - in_total(0), bp(0), final(false) { + zo{}, out(nullptr), outsize(0), crc(crc32_z(0L, Z_NULL, 0)), in_total(0), bp(0) { ZopfliInitOptions(&zo); // 5 iterations is reasonable for large files @@ -130,7 +129,6 @@ public: } ~zopfli_encoder() override { - final = true; finalize(); /* CRC */ @@ -150,7 +148,7 @@ public: } protected: - bool write_chunk(const void *buf, size_t len) override { + bool write_chunk(const void *buf, size_t len, bool final) override { if (len == 0) return true; @@ -178,7 +176,6 @@ private: unsigned long crc; uint32_t in_total; unsigned char bp; - bool final; }; class bz_strm : public out_stream { @@ -465,7 +462,7 @@ public: } protected: - bool write_chunk(const void *buf, size_t len) override { + bool write_chunk(const void *buf, size_t len, bool final) override { // This is an error if (len != chunk_sz) return false; @@ -514,22 +511,27 @@ public: } protected: - bool write_chunk(const void *buf, size_t len) override { - int r = LZ4_compress_HC((const char *) buf, out_buf, len, LZ4_COMPRESSED, LZ4HC_CLEVEL_MAX); - if (r == 0) { + bool write_chunk(const void *buf, size_t len, bool final) override { + auto in = static_cast(buf); + uint32_t block_sz = LZ4_compress_HC(in, out_buf, len, LZ4_COMPRESSED, LZ4HC_CLEVEL_MAX); + if (block_sz == 0) { LOGW("LZ4HC compression failure\n"); return false; } - return bwrite(&r, sizeof(r)) && bwrite(out_buf, r); + if (bwrite(&block_sz, sizeof(block_sz)) && bwrite(out_buf, block_sz)) { + in_total += sizeof(block_sz) + block_sz; + return true; + } + return false; } private: char *out_buf; bool lg; - unsigned in_total; + uint32_t in_total; }; -stream_ptr get_encoder(format_t type, stream_ptr &&base) { +filter_strm_ptr get_encoder(format_t type, stream_ptr &&base) { switch (type) { case XZ: return make_unique(std::move(base)); @@ -551,7 +553,7 @@ stream_ptr get_encoder(format_t type, stream_ptr &&base) { } } -stream_ptr get_decoder(format_t type, stream_ptr &&base) { +filter_strm_ptr get_decoder(format_t type, stream_ptr &&base) { switch (type) { case XZ: case LZMA: diff --git a/native/jni/magiskboot/compress.hpp b/native/jni/magiskboot/compress.hpp index 4963687ef..5b9a7aef6 100644 --- a/native/jni/magiskboot/compress.hpp +++ b/native/jni/magiskboot/compress.hpp @@ -4,9 +4,9 @@ #include "format.hpp" -stream_ptr get_encoder(format_t type, stream_ptr &&base); +filter_strm_ptr get_encoder(format_t type, stream_ptr &&base); -stream_ptr get_decoder(format_t type, stream_ptr &&base); +filter_strm_ptr get_decoder(format_t type, stream_ptr &&base); void compress(const char *method, const char *infile, const char *outfile); diff --git a/native/jni/utils/include/stream.hpp b/native/jni/utils/include/stream.hpp index 9853434ef..423ea5bd4 100644 --- a/native/jni/utils/include/stream.hpp +++ b/native/jni/utils/include/stream.hpp @@ -26,6 +26,7 @@ public: ssize_t read(void *buf, size_t len) override; bool write(const void *buf, size_t len) override; + virtual bool write(const void *buf, size_t len, bool final); // Seeking while filtering does not make sense off_t seek(off_t off, int whence) final { return stream::seek(off, whence); } @@ -34,6 +35,8 @@ protected: stream_ptr base; }; +using filter_strm_ptr = std::unique_ptr; + // Buffered output stream, writing in chunks class chunk_out_stream : public filter_stream { public: @@ -48,11 +51,12 @@ public: // Reading does not make sense ssize_t read(void *buf, size_t len) final { return stream::read(buf, len); } bool write(const void *buf, size_t len) final; + bool write(const void *buf, size_t len, bool final) final; protected: // Classes inheriting this class has to call finalize() in its destructor void finalize(); - virtual bool write_chunk(const void *buf, size_t len) = 0; + virtual bool write_chunk(const void *buf, size_t len, bool final) = 0; size_t chunk_sz; diff --git a/native/jni/utils/stream.cpp b/native/jni/utils/stream.cpp index aed2813af..3dbba71ed 100644 --- a/native/jni/utils/stream.cpp +++ b/native/jni/utils/stream.cpp @@ -107,13 +107,21 @@ bool filter_stream::write(const void *buf, size_t len) { return base->write(buf, len); } -bool chunk_out_stream::write(const void *_in, size_t len) { +bool filter_stream::write(const void *buf, size_t len, bool final) { + return write(buf, len); +} + +bool chunk_out_stream::write(const void *buf, size_t len) { + return write(buf, len, false); +} + +bool chunk_out_stream::write(const void *_in, size_t len, bool final) { auto in = static_cast(_in); while (len) { if (buf_off + len >= chunk_sz) { + // Enough input for a chunk const uint8_t *src; if (buf_off) { - // Copy the rest of the chunk to internal buffer src = _buf; auto copy = chunk_sz - buf_off; memcpy(_buf + buf_off, in, copy); @@ -125,8 +133,21 @@ bool chunk_out_stream::write(const void *_in, size_t len) { in += chunk_sz; len -= chunk_sz; } - if (!write_chunk(src, chunk_sz)) + if (!write_chunk(src, chunk_sz, final && len == 0)) return false; + } else if (final) { + // Final input data, write regardless whether it is chunk sized + if (buf_off) { + memcpy(_buf + buf_off, in, len); + auto avail = buf_off + len; + buf_off = 0; + if (!write_chunk(_buf, avail, true)) + return false; + } else { + if (!write_chunk(in, len, true)) + return false; + } + break; } else { // Buffer internally if (!_buf) { @@ -142,7 +163,7 @@ bool chunk_out_stream::write(const void *_in, size_t len) { void chunk_out_stream::finalize() { if (buf_off) { - write_chunk(_buf, buf_off); + write_chunk(_buf, buf_off, true); delete[] _buf; _buf = nullptr; buf_off = 0;