]> git.sesse.net Git - plocate/blobdiff - database-builder.cpp
Release plocate 1.1.4.
[plocate] / database-builder.cpp
index 42fc186bf3689792f6bd85519303c5134c99a4ed..da52cdb57764b6d60247c74e4318a84872503c49 100644 (file)
@@ -5,6 +5,9 @@
 
 #include <algorithm>
 #include <assert.h>
+#ifdef HAS_ENDIAN_H
+#include <endian.h>
+#endif
 #include <fcntl.h>
 #include <string.h>
 #include <string_view>
@@ -26,29 +29,18 @@ constexpr unsigned num_overflow_slots = 16;
 
 string zstd_compress(const string &src, ZSTD_CDict *cdict, string *tempbuf);
 
-static inline uint32_t read_unigram(const string_view s, size_t idx)
-{
-       if (idx < s.size()) {
-               return (unsigned char)s[idx];
-       } else {
-               return 0;
-       }
-}
-
-static inline uint32_t read_trigram(const string_view s, size_t start)
-{
-       return read_unigram(s, start) |
-               (read_unigram(s, start + 1) << 8) |
-               (read_unigram(s, start + 2) << 16);
-}
-
 class PostingListBuilder {
 public:
        inline void add_docid(uint32_t docid);
+       inline void add_first_docid(uint32_t docid);
        void finish();
 
-       string encoded;
-       size_t num_docids = 0;
+       vector<unsigned char> encoded;
+       size_t get_num_docids() const {
+               // Updated only when we flush, so check that we're finished.
+               assert(pending_deltas.empty());
+               return num_docids;
+       }
 
 private:
        void write_header(uint32_t docid);
@@ -56,7 +48,8 @@ private:
 
        vector<uint32_t> pending_deltas;
 
-       uint32_t last_block_end, last_docid = -1;
+       uint32_t num_docids = 0;  // Should be size_t, except the format only supports 2^32 docids per posting list anyway.
+       uint32_t last_docid = -1;
 };
 
 void PostingListBuilder::add_docid(uint32_t docid)
@@ -66,22 +59,20 @@ void PostingListBuilder::add_docid(uint32_t docid)
                return;
        }
 
-       if (num_docids == 0) {
-               // Very first docid.
-               write_header(docid);
-               ++num_docids;
-               last_block_end = last_docid = docid;
-               return;
-       }
-
        pending_deltas.push_back(docid - last_docid - 1);
        last_docid = docid;
        if (pending_deltas.size() == 128) {
                append_block();
                pending_deltas.clear();
-               last_block_end = docid;
+               num_docids += 128;
        }
+}
+
+void PostingListBuilder::add_first_docid(uint32_t docid)
+{
+       write_header(docid);
        ++num_docids;
+       last_docid = docid;
 }
 
 void PostingListBuilder::finish()
@@ -95,7 +86,10 @@ void PostingListBuilder::finish()
        // No interleaving for partial blocks.
        unsigned char buf[P4NENC_BOUND(128)];
        unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), pending_deltas.size(), /*interleaved=*/false, buf);
-       encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
+       encoded.insert(encoded.end(), buf, end);
+
+       num_docids += pending_deltas.size();
+       pending_deltas.clear();
 }
 
 void PostingListBuilder::append_block()
@@ -103,14 +97,14 @@ void PostingListBuilder::append_block()
        unsigned char buf[P4NENC_BOUND(128)];
        assert(pending_deltas.size() == 128);
        unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), 128, /*interleaved=*/true, buf);
-       encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
+       encoded.insert(encoded.end(), buf, end);
 }
 
 void PostingListBuilder::write_header(uint32_t docid)
 {
        unsigned char buf[P4NENC_BOUND(1)];
        unsigned char *end = write_baseval(docid, buf);
-       encoded.append(reinterpret_cast<char *>(buf), end - buf);
+       encoded.insert(encoded.end(), buf, end);
 }
 
 void DictionaryBuilder::add_file(string filename, dir_time)
@@ -191,7 +185,21 @@ public:
                return invindex[trgm] != nullptr;
        }
        size_t num_files_seen() const override { return num_files; }
-       PostingListBuilder &get_pl_builder(uint32_t trgm);
+       PostingListBuilder &get_pl_builder(uint32_t trgm)
+       {
+               return *invindex[trgm];
+       }
+
+       void add_docid(uint32_t trgm, uint32_t docid)
+       {
+               if (invindex[trgm] == nullptr) {
+                       invindex[trgm] = new PostingListBuilder;
+                       invindex[trgm]->add_first_docid(docid);
+               } else {
+                       invindex[trgm]->add_docid(docid);
+               }
+       }
+
        size_t num_trigrams() const;
        std::string get_compressed_dir_times();
 
@@ -229,14 +237,6 @@ EncodingCorpus::~EncodingCorpus()
        }
 }
 
-PostingListBuilder &EncodingCorpus::get_pl_builder(uint32_t trgm)
-{
-       if (invindex[trgm] == nullptr) {
-               invindex[trgm] = new PostingListBuilder;
-       }
-       return *invindex[trgm];
-}
-
 void EncodingCorpus::add_file(string filename, dir_time dt)
 {
        ++num_files;
@@ -303,15 +303,36 @@ void EncodingCorpus::flush_block()
 
        // Create trigrams.
        const char *ptr = current_block.c_str();
-       while (ptr < current_block.c_str() + current_block.size()) {
-               string_view s(ptr);
-               if (s.size() >= 3) {
-                       for (size_t j = 0; j < s.size() - 2; ++j) {
-                               uint32_t trgm = read_trigram(s, j);
-                               get_pl_builder(trgm).add_docid(docid);
+       const char *end = ptr + current_block.size();
+       while (ptr < end - 3) {  // Must be at least one filename left, that's at least three bytes.
+               if (ptr[0] == '\0') {
+                       // This filename is zero bytes, so skip it (and the zero terminator).
+                       ++ptr;
+                       continue;
+               } else if (ptr[1] == '\0') {
+                       // This filename is one byte, so skip it (and the zero terminator).
+                       ptr += 2;
+                       continue;
+               } else if (ptr[2] == '\0') {
+                       // This filename is two bytes, so skip it (and the zero terminator).
+                       ptr += 3;
+                       continue;
+               }
+               for ( ;; ) {
+                       // NOTE: Will read one byte past the end of the trigram, but it's OK,
+                       // since we always call it from contexts where there's a terminating zero byte.
+                       uint32_t trgm;
+                       memcpy(&trgm, ptr, sizeof(trgm));
+                       ++ptr;
+                       trgm = le32toh(trgm);
+                       add_docid(trgm & 0xffffff, docid);
+                       if (trgm <= 0xffffff) {
+                               // Terminating zero byte, so we're done with this filename.
+                               // Skip the remaining two bytes, and the zero terminator.
+                               ptr += 3;
+                               break;
                        }
                }
-               ptr += s.size() + 1;
        }
 
        // Compress and add the filename block.
@@ -432,7 +453,7 @@ unique_ptr<Trigram[]> create_hashtable(EncodingCorpus &corpus, const vector<uint
        }
        for (uint32_t trgm : all_trigrams) {
                // We don't know offset yet, so set it to zero.
-               Trigram to_insert{ trgm, uint32_t(corpus.get_pl_builder(trgm).num_docids), 0 };
+               Trigram to_insert{ trgm, uint32_t(corpus.get_pl_builder(trgm).get_num_docids()), 0 };
 
                uint32_t bucket = hash_trigram(trgm, ht_size);
                unsigned distance = 0;
@@ -568,8 +589,8 @@ void DatabaseBuilder::finish_corpus()
                        continue;
                PostingListBuilder &pl_builder = corpus->get_pl_builder(trgm);
                pl_builder.finish();
-               longest_posting_list = max(longest_posting_list, pl_builder.num_docids);
-               trigrams += pl_builder.num_docids;
+               longest_posting_list = max(longest_posting_list, pl_builder.get_num_docids());
+               trigrams += pl_builder.get_num_docids();
                bytes_for_posting_lists += pl_builder.encoded.size();
        }
        size_t num_trigrams = corpus->num_trigrams();
@@ -610,7 +631,7 @@ void DatabaseBuilder::finish_corpus()
                        continue;
                }
 
-               const string &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
+               const vector<unsigned char> &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
                offset += encoded.size();
        }
 
@@ -624,7 +645,7 @@ void DatabaseBuilder::finish_corpus()
                if (hashtable[i].num_docids == 0) {
                        continue;
                }
-               const string &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
+               const vector<unsigned char> &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
                fwrite(encoded.data(), encoded.size(), 1, outfp);
        }