Make PackagesWriter streaming - you don't need to give count when creating it

1 files changed, 23 insertions, 15 deletions

diff --git a/src/packages.cc b/src/packages.cc
index 0309df1..f2c4ba6 100644
--- a/src/packages.cc
+++ b/src/packages.cc
@@ -14,26 +14,23 @@ namespace {
 
 class PackagesWriterImpl : public PackagesWriter {
 public:
-  PackagesWriterImpl(size_t count, std::ostream* out)
-    : count_(count), out_(out) {
+  PackagesWriterImpl(std::ostream* out)
+    : open_(true), count_(0), out_(out) {
     uint8_t header[8];
     memcpy(header, "TPP", 3);
     header[3] = 0x1;  // Version
-    write_u32(header + 4, count_);  // Count
+    write_u32(header + 4, 0);  // Count
     out->write(reinterpret_cast<char*>(header), 8);
-
-    if (count == 0) {
-      write_u64(header, 0);  // EOF
-      out_->write(reinterpret_cast<char*>(header), 8);
-    }
   }
 
   ~PackagesWriterImpl() {
-    assert(count_ == 0);
+    if (open_) {
+      flush();
+    }
   }
 
   void write(Package const& package, std::string const& data) override {
-    if (count_ == 0) {
+    if (!open_) {
       assert(false);
       return;
     }
@@ -52,14 +49,25 @@ public:
     out_->write(reinterpret_cast<char*>(ptr), need);
     backup.reset();
     out_->write(data.data(), data.size());
+    count_++;
+  }
 
-    if (--count_ == 0) {
-      write_u64(size, 0);  // EOF
-      out_->write(reinterpret_cast<char*>(size), 8);
+  void flush() override {
+    if (!open_) return;
+    open_ = false;
+    uint8_t footer[8];
+    write_u64(footer, 0);  // EOF
+    out_->write(reinterpret_cast<char*>(footer), 8);
+    if (count_ > 0) {
+      out_->seekp(4);
+      write_u32(footer, count_);
+      out_->write(reinterpret_cast<char*>(footer), 4);
+      count_ = 0;
     }
   }
 
 private:
+  bool open_;
   size_t count_;
   std::ostream* const out_;
 };
@@ -67,8 +75,8 @@ private:
 }  // namespace
 
 // static
-PackagesWriter* PackagesWriter::create(size_t count, std::ostream* out) {
-  return new PackagesWriterImpl(count, out);
+PackagesWriter* PackagesWriter::create(std::ostream* out) {
+  return new PackagesWriterImpl(out);
 }
 
 // static