Refactor

src/fpindex/io/memory_directory.cpp

@@ -4,18 +4,18 @@ namespace fpindex {
 namespace io {
 
 std::shared_ptr<File> MemoryDirectory::OpenFile(const std::string &name, bool create) {
-  std::lock_guard<std::mutex> lock(mutex_);
-  auto iter = files_.find(name);
-  if (iter != files_.end()) {
-    return iter->second;
-  }
-  if (!create) {
-    return nullptr;
-  }
-  auto file = std::make_shared<MemoryFile>();
-  files_[name] = file;
-  return file;
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto iter = files_.find(name);
+    if (iter != files_.end()) {
+        return iter->second;
+    }
+    if (!create) {
+        return nullptr;
+    }
+    auto file = std::make_shared<MemoryFile>();
+    files_[name] = file;
+    return file;
 }
 
-}
-}
+}  // namespace io
+}  // namespace fpindex

src/fpindex/io/memory_directory.h

@@ -12,6 +12,7 @@ namespace io {
 class MemoryDirectory : public Directory {
  public:
     std::shared_ptr<File> OpenFile(const std::string& name, bool create = false) override;
+
  private:
     std::mutex mutex_;
     std::map<std::string, std::shared_ptr<MemoryFile>> files_;

src/fpindex/segment.cpp

@@ -1,11 +1,11 @@
 #include "fpindex/segment.h"
 
+#include <QDebug>
+
+#include "fpindex/io/file.h"
 #include "fpindex/proto/internal.pb.h"
 #include "fpindex/search_result.h"
 #include "fpindex/segment_file_format.h"
-#include "fpindex/io/file.h"
-
-#include <QDebug>
 
 namespace fpindex {
 
@@ -37,7 +37,8 @@ bool BlockBasedSegment::Search(const std::vector<uint32_t>& hashes, std::vector<
     auto prev_block_range_start = block_index.begin();
 
     for (const auto& hash : hashes) {
-        auto block_it = std::lower_bound(prev_block_range_start, block_index.end(), hash, CompareHashAgainstBlockIndexBack{});
+        auto block_it =
+            std::lower_bound(prev_block_range_start, block_index.end(), hash, CompareHashAgainstBlockIndexBack{});
         if (block_it == block_index.end()) {
             block_it = prev_block_range_start;
         }
@@ -138,10 +139,14 @@ bool Segment::Load(const std::shared_ptr<io::File>& file) {
     return true;
 }
 
-Segment::Segment(uint32_t id, std::shared_ptr<io::File> file, SegmentHeader header, size_t first_block_offset,
-                 std::vector<std::pair<uint32_t, uint32_t>>&& block_index)
-    : BlockBasedSegment(id), header_(header), file_(file), first_block_offset_(first_block_offset), block_index_(std::move(block_index)) {
-        ready_ = true;
+bool Segment::Load(std::shared_ptr<io::File> file, SegmentHeader header, size_t first_block_offset,
+                   std::vector<std::pair<uint32_t, uint32_t>>&& block_index) {
+    file_ = file;
+    header_ = header;
+    first_block_offset_ = first_block_offset;
+    block_index_ = std::move(block_index);
+    ready_ = true;
+    return true;
 }
 
 }  // namespace fpindex

src/fpindex/segment.h

@@ -26,6 +26,8 @@ class BlockBasedSegment : public BaseSegment {
     virtual bool GetBlock(size_t block_no, std::vector<std::pair<uint32_t, uint32_t>> *items) = 0;
 };
 
+class SegmentBuilder;
+
 class Segment : public BlockBasedSegment {
  public:
     Segment(uint32_t id) : BlockBasedSegment(id) {}
@@ -40,8 +42,9 @@ class Segment : public BlockBasedSegment {
     bool GetBlock(size_t block_no, std::vector<std::pair<uint32_t, uint32_t>> *items) override;
 
  private:
-    Segment(uint32_t id, std::shared_ptr<io::File> file, SegmentHeader header, size_t first_block_offset,
-            std::vector<std::pair<uint32_t, uint32_t>> &&block_index);
+    friend class fpindex::SegmentBuilder;
+    bool Load(std::shared_ptr<io::File> file, SegmentHeader header, size_t first_block_offset,
+              std::vector<std::pair<uint32_t, uint32_t>> &&block_index);
 
     std::mutex mutex_;
     std::atomic<bool> ready_{false};

src/fpindex/segment_builder.cpp

@@ -5,6 +5,7 @@
 #include <mutex>
 
 #include "fpindex/search_result.h"
+#include "fpindex/segment.h"
 #include "fpindex/segment_file_format.h"
 
 namespace fpindex {
@@ -52,7 +53,7 @@ bool SegmentBuilder::Search(const std::vector<uint32_t>& hashes, std::vector<Sea
     return true;
 }
 
-bool SegmentBuilder::Save(const std::shared_ptr<io::File> &file) {
+std::shared_ptr<Segment> SegmentBuilder::Save(const std::shared_ptr<io::File>& file) {
     std::shared_lock<std::shared_mutex> lock;
     if (!frozen_) {
         lock = std::shared_lock<std::shared_mutex>(mutex_);
@@ -65,27 +66,33 @@ bool SegmentBuilder::Save(const std::shared_ptr<io::File> &file) {
     internal::InitializeSegmentHeader(&header);
     internal::SerializeSegmentHeader(coded_stream.get(), header);
     if (coded_stream->HadError()) {
-        return false;
+        return nullptr;
     }
 
     const int block_size = header.block_size();
     const int header_size = coded_stream->ByteCount();
 
     int block_count = 0;
     auto it = data_.begin();
+    std::vector<std::pair<uint32_t, uint32_t>> block_index;
     while (it != data_.end()) {
-        it = internal::SerializeSegmentBlock(coded_stream.get(), header, it, data_.end());
+        auto next_it = internal::SerializeSegmentBlock(coded_stream.get(), header, it, data_.end());
         if (coded_stream->HadError()) {
-            return false;
+            return nullptr;
         }
+        auto first_key = it->first, last_key = std::prev(next_it)->first;
+        block_index.emplace_back(first_key, last_key);
+        it = next_it;
         block_count++;
     }
 
     if (header_size + block_count * block_size != coded_stream->ByteCount()) {
-        return false;
+        return nullptr;
     }
 
-    return true;
+    auto result = std::make_shared<Segment>(id());
+    result->Load(file, header, header_size, std::move(block_index));
+    return result;
 }
 
 }  // namespace fpindex

src/fpindex/segment_builder.h

@@ -10,6 +10,8 @@
 
 namespace fpindex {
 
+class Segment;
+
 class SegmentBuilder : public BaseSegment {
  public:
     SegmentBuilder(uint32_t id) : BaseSegment(id) {}
@@ -25,7 +27,7 @@ class SegmentBuilder : public BaseSegment {
     bool IsFrozen();
 
     // Serialize the segment data to the output stream.
-    bool Save(const std::shared_ptr<io::File> &file);
+    std::shared_ptr<Segment> Save(const std::shared_ptr<io::File>& file);
 
  private:
     std::shared_mutex mutex_;

src/fpindex/segment_builder_test.cpp

@@ -1,9 +1,10 @@
-#include "fpindex/segment.h"
 #include "fpindex/segment_builder.h"
-#include "fpindex/io/memory_file.h"
 
 #include <gtest/gtest.h>
 
+#include "fpindex/io/memory_file.h"
+#include "fpindex/segment.h"
+
 using namespace fpindex;
 using namespace testing;
 
@@ -33,18 +34,26 @@ TEST(SegmentBuilderTest, Save) {
     segment.Add(1, {1, 2, 3});
 
     auto file = std::make_shared<io::MemoryFile>();
-    segment.Save(file);
+    auto segment2 = segment.Save(file);
+
+    {
+        std::vector<uint32_t> query{1, 2, 3};
+        std::vector<SearchResult> results;
+        ASSERT_TRUE(segment2->Search(query, &results));
+        ASSERT_EQ(1, results.size());
+        ASSERT_EQ(1, results[0].id());
+        ASSERT_EQ(3, results[0].score());
+    }
 
-    Segment new_segment(0);
-    new_segment.Load(file);
+    Segment segment3(0);
+    segment3.Load(file);
 
-    std::vector<uint32_t> query{1, 2, 3};
-    std::vector<SearchResult> results;
-    ASSERT_TRUE(new_segment.Search(query, &results));
-    for (const auto& result : results) {
-        std::cout << result.id() << " " << result.score() << std::endl;
+    {
+        std::vector<uint32_t> query{1, 2, 3};
+        std::vector<SearchResult> results;
+        ASSERT_TRUE(segment3.Search(query, &results));
+        ASSERT_EQ(1, results.size());
+        ASSERT_EQ(1, results[0].id());
+        ASSERT_EQ(3, results[0].score());
     }
-    ASSERT_EQ(1, results.size());
-    ASSERT_EQ(1, results[0].id());
-    ASSERT_EQ(3, results[0].score());
 }

src/fpindex/segment_file_format.h

@@ -89,7 +89,7 @@ inline Iter SerializeSegmentBlock(io::CodedOutputStream *output, const SegmentHe
 }
 
 inline bool ParseSegmentBlockV1(io::CodedInputStream *input, const SegmentHeader &header,
-                         std::vector<std::pair<uint32_t, uint32_t>> *entries) {
+                                std::vector<std::pair<uint32_t, uint32_t>> *entries) {
     const int block_size = header.block_size();
     auto start = input->CurrentPosition();
     auto limit = input->PushLimit(block_size);
@@ -129,7 +129,7 @@ inline bool ParseSegmentBlockV1(io::CodedInputStream *input, const SegmentHeader
 }
 
 inline bool ParseSegmentBlock(io::CodedInputStream *input, const SegmentHeader &header,
-                       std::vector<std::pair<uint32_t, uint32_t>> *entries) {
+                              std::vector<std::pair<uint32_t, uint32_t>> *entries) {
     if (header.block_format() == SegmentBlockFormat::BLOCK_FORMAT_V1) {
         return ParseSegmentBlockV1(input, header, entries);
     }