git.sesse.net Git - plocate/blob - plocate-build.cpp

   1 #include "db.h"
   2 #include "turbopfor-encode.h"
   3
   4 #include <algorithm>
   5 #include <assert.h>
   6 #include <chrono>
   7 #include <iosfwd>
   8 #include <math.h>
   9 #include <memory>
  10 #include <stdint.h>
  11 #include <stdio.h>
  12 #include <stdlib.h>
  13 #include <string.h>
  14 #include <string>
  15 #include <string_view>
  16 #include <sys/stat.h>
  17 #include <utility>
  18 #include <vector>
  19 #include <zstd.h>
  20
  21 #define P4NENC_BOUND(n) ((n + 127) / 128 + (n + 32) * sizeof(uint32_t))
  22 #define dprintf(...)
  23 //#define dprintf(...) fprintf(stderr, __VA_ARGS__);
  24
  25 #define NUM_TRIGRAMS 16777216
  26
  27 using namespace std;
  28 using namespace std::chrono;
  29
  30 string zstd_compress(const string &src, string *tempbuf);
  31
  32 constexpr unsigned num_overflow_slots = 16;
  33
  34 static inline uint32_t read_unigram(const string_view s, size_t idx)
  35 {
  36         if (idx < s.size()) {
  37                 return (unsigned char)s[idx];
  38         } else {
  39                 return 0;
  40         }
  41 }
  42
  43 static inline uint32_t read_trigram(const string_view s, size_t start)
  44 {
  45         return read_unigram(s, start) |
  46                 (read_unigram(s, start + 1) << 8) |
  47                 (read_unigram(s, start + 2) << 16);
  48 }
  49
  50 enum {
  51         DBE_NORMAL = 0, /* A non-directory file */
  52         DBE_DIRECTORY = 1, /* A directory */
  53         DBE_END = 2 /* End of directory contents; contains no name */
  54 };
  55
  56 // From mlocate.
  57 struct db_header {
  58         uint8_t magic[8];
  59         uint32_t conf_size;
  60         uint8_t version;
  61         uint8_t check_visibility;
  62         uint8_t pad[2];
  63 };
  64
  65 // From mlocate.
  66 struct db_directory {
  67         uint64_t time_sec;
  68         uint32_t time_nsec;
  69         uint8_t pad[4];
  70 };
  71
  72 class PostingListBuilder {
  73 public:
  74         inline void add_docid(uint32_t docid);
  75         void finish();
  76
  77         string encoded;
  78         size_t num_docids = 0;
  79
  80 private:
  81         void write_header(uint32_t docid);
  82         void append_block();
  83
  84         vector<uint32_t> pending_deltas;
  85
  86         uint32_t last_block_end, last_docid = -1;
  87 };
  88
  89 void PostingListBuilder::add_docid(uint32_t docid)
  90 {
  91         // Deduplicate against the last inserted value, if any.
  92         if (docid == last_docid) {
  93                 return;
  94         }
  95
  96         if (num_docids == 0) {
  97                 // Very first docid.
  98                 write_header(docid);
  99                 ++num_docids;
 100                 last_block_end = last_docid = docid;
 101                 return;
 102         }
 103
 104         pending_deltas.push_back(docid - last_docid - 1);
 105         last_docid = docid;
 106         if (pending_deltas.size() == 128) {
 107                 append_block();
 108                 pending_deltas.clear();
 109                 last_block_end = docid;
 110         }
 111         ++num_docids;
 112 }
 113
 114 void PostingListBuilder::finish()
 115 {
 116         if (pending_deltas.empty()) {
 117                 return;
 118         }
 119
 120         assert(!encoded.empty());  // write_header() should already have run.
 121
 122         // No interleaving for partial blocks.
 123         unsigned char buf[P4NENC_BOUND(128)];
 124         unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), pending_deltas.size(), /*interleaved=*/false, buf);
 125         encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
 126 }
 127
 128 void PostingListBuilder::append_block()
 129 {
 130         unsigned char buf[P4NENC_BOUND(128)];
 131         assert(pending_deltas.size() == 128);
 132         unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), 128, /*interleaved=*/true, buf);
 133         encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
 134 }
 135
 136 void PostingListBuilder::write_header(uint32_t docid)
 137 {
 138         unsigned char buf[P4NENC_BOUND(1)];
 139         unsigned char *end = write_baseval(docid, buf);
 140         encoded.append(reinterpret_cast<char *>(buf), end - buf);
 141 }
 142
 143 class Corpus {
 144 public:
 145         Corpus(FILE *outfp, size_t block_size)
 146                 : invindex(new PostingListBuilder *[NUM_TRIGRAMS]), outfp(outfp), block_size(block_size)
 147         {
 148                 fill(invindex.get(), invindex.get() + NUM_TRIGRAMS, nullptr);
 149         }
 150         ~Corpus()
 151         {
 152                 for (unsigned i = 0; i < NUM_TRIGRAMS; ++i) {
 153                         delete invindex[i];
 154                 }
 155         }
 156
 157         void add_file(string filename);
 158         void flush_block();
 159
 160         vector<uint64_t> filename_blocks;
 161         size_t num_files = 0, num_files_in_block = 0, num_blocks = 0;
 162         bool seen_trigram(uint32_t trgm)
 163         {
 164                 return invindex[trgm] != nullptr;
 165         }
 166         PostingListBuilder &get_pl_builder(uint32_t trgm)
 167         {
 168                 if (invindex[trgm] == nullptr) {
 169                         invindex[trgm] = new PostingListBuilder;
 170                 }
 171                 return *invindex[trgm];
 172         }
 173
 174 private:
 175         unique_ptr<PostingListBuilder *[]> invindex;
 176         FILE *outfp;
 177         string current_block;
 178         string tempbuf;
 179         const size_t block_size;
 180 };
 181
 182 void Corpus::add_file(string filename)
 183 {
 184         ++num_files;
 185         if (!current_block.empty()) {
 186                 current_block.push_back('\0');
 187         }
 188         current_block += filename;
 189         if (++num_files_in_block == block_size) {
 190                 flush_block();
 191         }
 192 }
 193
 194 void Corpus::flush_block()
 195 {
 196         if (current_block.empty()) {
 197                 return;
 198         }
 199
 200         uint32_t docid = num_blocks;
 201
 202         // Create trigrams.
 203         const char *ptr = current_block.c_str();
 204         while (ptr < current_block.c_str() + current_block.size()) {
 205                 string_view s(ptr);
 206                 if (s.size() >= 3) {
 207                         for (size_t j = 0; j < s.size() - 2; ++j) {
 208                                 uint32_t trgm = read_trigram(s, j);
 209                                 get_pl_builder(trgm).add_docid(docid);
 210                         }
 211                 }
 212                 ptr += s.size() + 1;
 213         }
 214
 215         // Compress and add the filename block.
 216         filename_blocks.push_back(ftell(outfp));
 217         string compressed = zstd_compress(current_block, &tempbuf);
 218         if (fwrite(compressed.data(), compressed.size(), 1, outfp) != 1) {
 219                 perror("fwrite()");
 220                 exit(1);
 221         }
 222
 223         current_block.clear();
 224         num_files_in_block = 0;
 225         ++num_blocks;
 226 }
 227
 228 string read_cstr(FILE *fp)
 229 {
 230         string ret;
 231         for (;;) {
 232                 int ch = getc(fp);
 233                 if (ch == -1) {
 234                         perror("getc");
 235                         exit(1);
 236                 }
 237                 if (ch == 0) {
 238                         return ret;
 239                 }
 240                 ret.push_back(ch);
 241         }
 242 }
 243
 244 void handle_directory(FILE *fp, Corpus *corpus)
 245 {
 246         db_directory dummy;
 247         if (fread(&dummy, sizeof(dummy), 1, fp) != 1) {
 248                 if (feof(fp)) {
 249                         return;
 250                 } else {
 251                         perror("fread");
 252                 }
 253         }
 254
 255         string dir_path = read_cstr(fp);
 256         if (dir_path == "/") {
 257                 dir_path = "";
 258         }
 259
 260         for (;;) {
 261                 int type = getc(fp);
 262                 if (type == DBE_NORMAL) {
 263                         string filename = read_cstr(fp);
 264                         corpus->add_file(dir_path + "/" + filename);
 265                 } else if (type == DBE_DIRECTORY) {
 266                         string dirname = read_cstr(fp);
 267                         corpus->add_file(dir_path + "/" + dirname);
 268                 } else {
 269                         return;  // Probably end.
 270                 }
 271         }
 272 }
 273
 274 void read_mlocate(const char *filename, Corpus *corpus)
 275 {
 276         FILE *fp = fopen(filename, "rb");
 277         if (fp == nullptr) {
 278                 perror(filename);
 279                 exit(1);
 280         }
 281
 282         db_header hdr;
 283         if (fread(&hdr, sizeof(hdr), 1, fp) != 1) {
 284                 perror("short read");
 285                 exit(1);
 286         }
 287
 288         // TODO: Care about the base path.
 289         string path = read_cstr(fp);
 290         while (!feof(fp)) {
 291                 handle_directory(fp, corpus);
 292         }
 293         fclose(fp);
 294 }
 295
 296 string zstd_compress(const string &src, string *tempbuf)
 297 {
 298         size_t max_size = ZSTD_compressBound(src.size());
 299         if (tempbuf->size() < max_size) {
 300                 tempbuf->resize(max_size);
 301         }
 302         size_t size = ZSTD_compress(&(*tempbuf)[0], max_size, src.data(), src.size(), /*level=*/6);
 303         return string(tempbuf->data(), size);
 304 }
 305
 306 bool is_prime(uint32_t x)
 307 {
 308         if ((x % 2) == 0 || (x % 3) == 0) {
 309                 return false;
 310         }
 311         uint32_t limit = ceil(sqrt(x));
 312         for (uint32_t factor = 5; factor <= limit; ++factor) {
 313                 if ((x % factor) == 0) {
 314                         return false;
 315                 }
 316         }
 317         return true;
 318 }
 319
 320 uint32_t next_prime(uint32_t x)
 321 {
 322         if ((x % 2) == 0) {
 323                 ++x;
 324         }
 325         while (!is_prime(x)) {
 326                 x += 2;
 327         }
 328         return x;
 329 }
 330
 331 unique_ptr<Trigram[]> create_hashtable(Corpus &corpus, const vector<uint32_t> &all_trigrams, uint32_t ht_size, uint32_t num_overflow_slots)
 332 {
 333         unique_ptr<Trigram[]> ht(new Trigram[ht_size + num_overflow_slots + 1]);  // 1 for the sentinel element at the end.
 334         for (unsigned i = 0; i < ht_size + num_overflow_slots + 1; ++i) {
 335                 ht[i].trgm = uint32_t(-1);
 336                 ht[i].num_docids = 0;
 337                 ht[i].offset = 0;
 338         }
 339         for (uint32_t trgm : all_trigrams) {
 340                 // We don't know offset yet, so set it to zero.
 341                 Trigram to_insert{ trgm, uint32_t(corpus.get_pl_builder(trgm).num_docids), 0 };
 342
 343                 uint32_t bucket = hash_trigram(trgm, ht_size);
 344                 unsigned distance = 0;
 345                 while (ht[bucket].num_docids != 0) {
 346                         // Robin Hood hashing; reduces the longest distance by a lot.
 347                         unsigned other_distance = bucket - hash_trigram(ht[bucket].trgm, ht_size);
 348                         if (distance > other_distance) {
 349                                 swap(to_insert, ht[bucket]);
 350                                 distance = other_distance;
 351                         }
 352
 353                         ++bucket, ++distance;
 354                         if (distance > num_overflow_slots) {
 355                                 return nullptr;
 356                         }
 357                 }
 358                 ht[bucket] = to_insert;
 359         }
 360         return ht;
 361 }
 362
 363 void do_build(const char *infile, const char *outfile, int block_size)
 364 {
 365         steady_clock::time_point start __attribute__((unused)) = steady_clock::now();
 366
 367         umask(0027);
 368         FILE *outfp = fopen(outfile, "wb");
 369
 370         // Write the header.
 371         Header hdr;
 372         memcpy(hdr.magic, "\0plocate", 8);
 373         hdr.version = -1;  // Mark as broken.
 374         hdr.hashtable_size = 0;  // Not known yet.
 375         hdr.extra_ht_slots = num_overflow_slots;
 376         hdr.num_docids = 0;
 377         hdr.hash_table_offset_bytes = -1;  // We don't know these offsets yet.
 378         hdr.filename_index_offset_bytes = -1;
 379         fwrite(&hdr, sizeof(hdr), 1, outfp);
 380
 381         Corpus corpus(outfp, block_size);
 382
 383         read_mlocate(infile, &corpus);
 384         if (false) {  // To read a plain text file.
 385                 FILE *fp = fopen(infile, "r");
 386                 while (!feof(fp)) {
 387                         char buf[1024];
 388                         if (fgets(buf, 1024, fp) == nullptr || feof(fp)) {
 389                                 break;
 390                         }
 391                         string s(buf);
 392                         if (s.back() == '\n')
 393                                 s.pop_back();
 394                         corpus.add_file(move(s));
 395                 }
 396                 fclose(fp);
 397         }
 398         corpus.flush_block();
 399         dprintf("Read %zu files from %s\n", corpus.num_files, infile);
 400         hdr.num_docids = corpus.filename_blocks.size();
 401
 402         // Stick an empty block at the end as sentinel.
 403         corpus.filename_blocks.push_back(ftell(outfp));
 404         const size_t bytes_for_filenames = corpus.filename_blocks.back() - corpus.filename_blocks.front();
 405
 406         // Write the offsets to the filenames.
 407         hdr.filename_index_offset_bytes = ftell(outfp);
 408         const size_t bytes_for_filename_index = corpus.filename_blocks.size() * sizeof(uint64_t);
 409         fwrite(corpus.filename_blocks.data(), corpus.filename_blocks.size(), sizeof(uint64_t), outfp);
 410         corpus.filename_blocks.clear();
 411         corpus.filename_blocks.shrink_to_fit();
 412
 413         // Finish up encoding the posting lists.
 414         size_t trigrams = 0, longest_posting_list = 0;
 415         size_t bytes_for_posting_lists = 0;
 416         for (unsigned trgm = 0; trgm < NUM_TRIGRAMS; ++trgm) {
 417                 if (!corpus.seen_trigram(trgm))
 418                         continue;
 419                 PostingListBuilder &pl_builder = corpus.get_pl_builder(trgm);
 420                 pl_builder.finish();
 421                 longest_posting_list = max(longest_posting_list, pl_builder.num_docids);
 422                 trigrams += pl_builder.num_docids;
 423                 bytes_for_posting_lists += pl_builder.encoded.size();
 424         }
 425         dprintf("%zu files, %zu different trigrams, %zu entries, avg len %.2f, longest %zu\n",
 426                 corpus.num_files, corpus.invindex.size(), trigrams, double(trigrams) / corpus.invindex.size(), longest_posting_list);
 427         dprintf("%zu bytes used for posting lists (%.2f bits/entry)\n", bytes_for_posting_lists, 8 * bytes_for_posting_lists / double(trigrams));
 428
 429         dprintf("Building posting lists took %.1f ms.\n\n", 1e3 * duration<float>(steady_clock::now() - start).count());
 430
 431         // Find the used trigrams.
 432         vector<uint32_t> all_trigrams;
 433         for (unsigned trgm = 0; trgm < NUM_TRIGRAMS; ++trgm) {
 434                 if (corpus.seen_trigram(trgm)) {
 435                         all_trigrams.push_back(trgm);
 436                 }
 437         }
 438
 439         // Create the hash table.
 440         unique_ptr<Trigram[]> hashtable;
 441         uint32_t ht_size = next_prime(all_trigrams.size());
 442         for (;;) {
 443                 hashtable = create_hashtable(corpus, all_trigrams, ht_size, num_overflow_slots);
 444                 if (hashtable == nullptr) {
 445                         dprintf("Failed creating hash table of size %u, increasing by 5%% and trying again.\n", ht_size);
 446                         ht_size = next_prime(ht_size * 1.05);
 447                 } else {
 448                         dprintf("Created hash table of size %u.\n\n", ht_size);
 449                         break;
 450                 }
 451         }
 452
 453         // Find the offsets for each posting list.
 454         size_t bytes_for_hashtable = (ht_size + num_overflow_slots + 1) * sizeof(Trigram);
 455         uint64_t offset = ftell(outfp) + bytes_for_hashtable;
 456         for (unsigned i = 0; i < ht_size + num_overflow_slots + 1; ++i) {
 457                 hashtable[i].offset = offset;  // Needs to be there even for empty slots.
 458                 if (hashtable[i].num_docids == 0) {
 459                         continue;
 460                 }
 461
 462                 const string &encoded = corpus.get_pl_builder(hashtable[i].trgm).encoded;
 463                 offset += encoded.size();
 464         }
 465
 466         // Write the hash table.
 467         hdr.hash_table_offset_bytes = ftell(outfp);
 468         hdr.hashtable_size = ht_size;
 469         fwrite(hashtable.get(), ht_size + num_overflow_slots + 1, sizeof(Trigram), outfp);
 470
 471         // Write the actual posting lists.
 472         for (unsigned i = 0; i < ht_size + num_overflow_slots + 1; ++i) {
 473                 if (hashtable[i].num_docids == 0) {
 474                         continue;
 475                 }
 476                 const string &encoded = corpus.get_pl_builder(hashtable[i].trgm).encoded;
 477                 fwrite(encoded.data(), encoded.size(), 1, outfp);
 478         }
 479
 480         // Rewind, and write the updated header.
 481         hdr.version = 0;
 482         fseek(outfp, 0, SEEK_SET);
 483         fwrite(&hdr, sizeof(hdr), 1, outfp);
 484         fclose(outfp);
 485
 486         size_t total_bytes __attribute__((unused)) = (bytes_for_hashtable + bytes_for_posting_lists + bytes_for_filename_index + bytes_for_filenames);
 487
 488         dprintf("Block size:     %7d files\n", block_size);
 489         dprintf("Hash table:     %'7.1f MB\n", bytes_for_hashtable / 1048576.0);
 490         dprintf("Posting lists:  %'7.1f MB\n", bytes_for_posting_lists / 1048576.0);
 491         dprintf("Filename index: %'7.1f MB\n", bytes_for_filename_index / 1048576.0);
 492         dprintf("Filenames:      %'7.1f MB\n", bytes_for_filenames / 1048576.0);
 493         dprintf("Total:          %'7.1f MB\n", total_bytes / 1048576.0);
 494         dprintf("\n");
 495 }
 496
 497 int main(int argc, char **argv)
 498 {
 499         do_build(argv[1], argv[2], 32);
 500         exit(EXIT_SUCCESS);
 501 }