From 497f926aab11b146fd006d95941a6377e6d0dbf6 Mon Sep 17 00:00:00 2001
From: drons <drons@list.ru>
Date: Sat, 13 Apr 2024 10:07:29 +0300
Subject: [PATCH] Declare saveIndex/loadIndex as templates to use custom
 streams

---
 hnswlib/hnswalg.h | 23 ++++++++++++++++++++---
 hnswlib/hnswlib.h |  8 ++++----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/hnswlib/hnswalg.h b/hnswlib/hnswalg.h
index f516df59..b4c8aa09 100644
--- a/hnswlib/hnswalg.h
+++ b/hnswlib/hnswalg.h
@@ -85,6 +85,16 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         loadIndex(location, s, max_elements);
     }
 
+    template<class stream_t>
+    HierarchicalNSW(
+        stream_t &stream,
+        SpaceInterface<dist_t> *s,
+        bool nmslib = false,
+        size_t max_elements = 0,
+        bool allow_replace_deleted = false)
+        : allow_replace_deleted_(allow_replace_deleted) {
+        loadIndexFromStream(stream, s, max_elements);
+    }
 
     HierarchicalNSW(
         SpaceInterface<dist_t> *s,
@@ -684,7 +694,11 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
     void saveIndex(const std::string &location) {
         std::ofstream output(location, std::ios::binary);
-
+        saveIndexToStream(output);
+    }
+    
+    template<typename stream_t>
+    void saveIndexToStream(stream_t &output) {
         writeBinaryPOD(output, offsetLevel0_);
         writeBinaryPOD(output, max_elements_);
         writeBinaryPOD(output, cur_element_count);
@@ -711,17 +725,20 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         output.close();
     }
 
-
     void loadIndex(const std::string &location, SpaceInterface<dist_t> *s, size_t max_elements_i = 0) {
         std::ifstream input(location, std::ios::binary);
 
         if (!input.is_open())
             throw std::runtime_error("Cannot open file");
+        loadIndexFromStream(input, s, max_elements_i);
+    }
 
+    template<typename stream_t>
+    void loadIndexFromStream(stream_t &input, SpaceInterface<dist_t> *s, size_t max_elements_i = 0) {
         clear();
         // get file size:
         input.seekg(0, input.end);
-        std::streampos total_filesize = input.tellg();
+        auto total_filesize = input.tellg();
         input.seekg(0, input.beg);
 
         readBinaryPOD(input, offsetLevel0_);
diff --git a/hnswlib/hnswlib.h b/hnswlib/hnswlib.h
index 7ccfbba5..f03ef85a 100644
--- a/hnswlib/hnswlib.h
+++ b/hnswlib/hnswlib.h
@@ -157,13 +157,13 @@ class pairGreater {
     }
 };
 
-template<typename T>
-static void writeBinaryPOD(std::ostream &out, const T &podRef) {
+template<typename stream_t, typename T>
+static void writeBinaryPOD(stream_t &out, const T &podRef) {
     out.write((char *) &podRef, sizeof(T));
 }
 
-template<typename T>
-static void readBinaryPOD(std::istream &in, T &podRef) {
+template<typename stream_t, typename T>
+static void readBinaryPOD(stream_t &in, T &podRef) {
     in.read((char *) &podRef, sizeof(T));
 }