Enhance the flexibility of the BinaryOutputArchive and BinaryInputArchive

parallel_hashmap/phmap_dump.h

@@ -168,22 +168,32 @@ bool parallel_hash_set<N, RefSet, Mtx_, Policy, Hash, Eq, Alloc>::phmap_load(Inp
 class BinaryOutputArchive {
 public:
     BinaryOutputArchive(const char *file_path) {
-        ofs_.open(file_path, std::ofstream::out | std::ofstream::trunc | std::ofstream::binary);
+        os_ = new std::ofstream(file_path, std::ofstream::out | std::ofstream::trunc | std::ofstream::binary);
+        need_close_ = true;
     }
 
-    ~BinaryOutputArchive() = default;
+    BinaryOutputArchive(std::ostream& os) : os_(&os), need_close_(false) {}
+
+    ~BinaryOutputArchive() {
+        if (need_close_) {
+            os_->flush();
+            delete os_;
+        }
+        need_close_ = false;
+    }
+
     BinaryOutputArchive(const BinaryOutputArchive&) = delete;
     BinaryOutputArchive& operator=(const BinaryOutputArchive&) = delete;
 
     bool saveBinary(const void *p, size_t sz) {
-        ofs_.write(reinterpret_cast<const char*>(p), (std::streamsize)sz);
+        os_->write(reinterpret_cast<const char*>(p), (std::streamsize)sz);
         return true;
     }
 
     template<typename V>
     typename std::enable_if<type_traits_internal::IsTriviallyCopyable<V>::value, bool>::type
     saveBinary(const V& v) {
-        ofs_.write(reinterpret_cast<const char *>(&v), sizeof(V));
+        os_->write(reinterpret_cast<const char *>(&v), sizeof(V));
         return true;
     }
 
@@ -194,29 +204,38 @@ class BinaryOutputArchive {
     }
 
 private:
-    std::ofstream ofs_;
+    std::ostream* os_;
+    bool need_close_;
 };
 
 
 class BinaryInputArchive {
 public:
     BinaryInputArchive(const char * file_path) {
-        ifs_.open(file_path, std::ofstream::in | std::ofstream::binary);
+        is_ = new std::ifstream(file_path, std::ifstream::in | std::ifstream::binary);
     }
+
+    BinaryInputArchive(std::istream& is) : need_close_(false), is_(&is) {}
 
-    ~BinaryInputArchive() = default;
+    ~BinaryInputArchive() {
+        if (need_close_) {
+            delete is_;
+        }
+        need_close_ = false;
+    }
+
     BinaryInputArchive(const BinaryInputArchive&) = delete;
     BinaryInputArchive& operator=(const BinaryInputArchive&) = delete;
 
     bool loadBinary(void* p, size_t sz) {
-        ifs_.read(reinterpret_cast<char*>(p),  (std::streamsize)sz);
+        is_->read(reinterpret_cast<char*>(p),  (std::streamsize)sz);
         return true;
     }
 
     template<typename V>
     typename std::enable_if<type_traits_internal::IsTriviallyCopyable<V>::value, bool>::type
     loadBinary(V* v) {
-        ifs_.read(reinterpret_cast<char *>(v), sizeof(V));
+        is_->read(reinterpret_cast<char *>(v), sizeof(V));
         return true;
     }
 
@@ -227,7 +246,8 @@ class BinaryInputArchive {
     }
 
 private:
-    std::ifstream ifs_;
+    bool need_close_;
+    std::istream* is_;
 };
 
 } // namespace phmap

tests/dump_load_test.cc

@@ -1,3 +1,7 @@
+#include <cstdint>
+#include <fstream>
+#include <parallel_hashmap/phmap.h>
+#include <sstream>
 #include <vector>
 
 #include "gtest/gtest.h"
@@ -22,6 +26,16 @@ TEST(DumpLoad, FlatHashSet_uint32) {
         EXPECT_TRUE(st2.phmap_load(ar_in));
     }
     EXPECT_TRUE(st1 == st2);
+
+    {
+        std::stringstream ss;
+        phmap::BinaryOutputArchive ar_out(ss);
+        EXPECT_TRUE(st1.phmap_dump(ar_out));
+        phmap::flat_hash_set<uint32_t> st3;
+        phmap::BinaryInputArchive ar_in(ss);
+        EXPECT_TRUE(st3.phmap_load(ar_in));
+        EXPECT_TRUE(st1 == st3);
+    }
 }
 
 TEST(DumpLoad, FlatHashMap_uint64_uint32) {
@@ -39,6 +53,16 @@ TEST(DumpLoad, FlatHashMap_uint64_uint32) {
         EXPECT_TRUE(mp2.phmap_load(ar_in));
     }
 
+    {
+        std::stringstream ss;
+        phmap::BinaryOutputArchive ar_out(ss);
+        EXPECT_TRUE(mp1.phmap_dump(ar_out));
+        phmap::flat_hash_map<uint64_t, uint32_t> mp3;
+        phmap::BinaryInputArchive ar_in(ss);
+        EXPECT_TRUE(mp3.phmap_load(ar_in));
+        EXPECT_TRUE(mp1 == mp3);
+    }
+
     EXPECT_TRUE(mp1 == mp2);
 }
 
@@ -57,6 +81,24 @@ TEST(DumpLoad, ParallelFlatHashMap_uint64_uint32) {
         EXPECT_TRUE(mp2.phmap_load(ar_in));
     }
     EXPECT_TRUE(mp1 == mp2);
+
+    // test stringstream and dump/load in the middle of the stream
+    {
+        char hello[] = "Hello";
+        std::stringstream ss;
+        ss.write(hello, 5);
+        phmap::BinaryOutputArchive ar_out(ss);
+        EXPECT_TRUE(mp1.phmap_dump(ar_out));
+        phmap::parallel_flat_hash_map<uint64_t, uint32_t> mp3;
+        phmap::BinaryInputArchive ar_in(ss);
+        char s[5];
+        ss.read(s, 5);
+        for (int i = 0; i < 5; ++i) {
+            EXPECT_EQ(hello[i], s[i]);
+        }
+        EXPECT_TRUE(mp3.phmap_load(ar_in));
+        EXPECT_TRUE(mp1 == mp3);
+    }
 }
 
 }