X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=plocate-build.cpp;h=028a1ee06c620d8274ae707ba1905041ac288a17;hb=fd6198891d6fd9642effc0843fef6f23b991af3e;hp=dc2fe3c0960d3db7487727df1b9354fd6000d6b0;hpb=c48d3b9a8af50db3f5ccb08caffc267bcd0aada9;p=plocate diff --git a/plocate-build.cpp b/plocate-build.cpp index dc2fe3c..028a1ee 100644 --- a/plocate-build.cpp +++ b/plocate-build.cpp @@ -1,59 +1,38 @@ +#include "database-builder.h" +#include "db.h" +#include "dprintf.h" + +#include +#include #include +#include +#include +#include +#include +#include +#include +#include #include +#include #include -#include -#include #include -#include -#include -#include -#include -#include -#include -#include -#include #include -#include - -#define VINT_IN -#include "vint.h" -#include "vp4.h" - -#define P4NENC_BOUND(n) ((n+127)/128+(n+32)*sizeof(uint32_t)) -#define dprintf(...) -//#define dprintf(...) fprintf(stderr, __VA_ARGS__); +#include +#include using namespace std; using namespace std::chrono; -string zstd_compress(const string &src, 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); -} +bool use_debug = false; -enum -{ - DBE_NORMAL = 0, /* A non-directory file */ - DBE_DIRECTORY = 1, /* A directory */ - DBE_END = 2 /* End of directory contents; contains no name */ +enum { + DBE_NORMAL = 0, /* A non-directory file */ + DBE_DIRECTORY = 1, /* A directory */ + DBE_END = 2 /* End of directory contents; contains no name */ }; // From mlocate. -struct db_header -{ +struct db_header { uint8_t magic[8]; uint32_t conf_size; uint8_t version; @@ -62,324 +41,214 @@ struct db_header }; // From mlocate. -struct db_directory -{ +struct db_directory { uint64_t time_sec; uint32_t time_nsec; uint8_t pad[4]; }; -class PostingListBuilder { -public: - void add_docid(uint32_t docid); - void finish(); - - string encoded; - size_t num_docids = 0; - -private: - void write_header(uint32_t docid); - void append_block(); - - uint32_t pending_docids[128]; - unsigned num_pending_docids = 0; - - uint32_t last_block_end; -}; - -void PostingListBuilder::add_docid(uint32_t docid) +string read_cstr(FILE *fp) { - // Deduplicate against the last inserted value, if any. - if (num_pending_docids == 0) { - if (encoded.empty()) { - // Very first docid. - write_header(docid); - ++num_docids; - last_block_end = docid; - return; - } else if (docid == last_block_end) { - return; + string ret; + for (;;) { + int ch = getc(fp); + if (ch == -1) { + perror("getc"); + exit(1); } - } else { - if (docid == pending_docids[num_pending_docids - 1]) { - return; + if (ch == 0) { + return ret; } + ret.push_back(ch); } - - pending_docids[num_pending_docids++] = docid; - if (num_pending_docids == 128) { - append_block(); - num_pending_docids = 0; - last_block_end = docid; - } - ++num_docids; -} - -void PostingListBuilder::finish() -{ - if (num_pending_docids == 0) { - return; - } - - assert(!encoded.empty()); // write_header() should already have run. - - // No interleaving for partial blocks. - unsigned char buf[P4NENC_BOUND(128)]; - unsigned char *end = p4d1enc32(pending_docids, num_pending_docids, buf, last_block_end); - encoded.append(reinterpret_cast(buf), reinterpret_cast(end)); -} - -void PostingListBuilder::append_block() -{ - unsigned char buf[P4NENC_BOUND(128)]; - assert(num_pending_docids == 128); - unsigned char *end = p4d1enc128v32(pending_docids, 128, buf, last_block_end); - encoded.append(reinterpret_cast(buf), reinterpret_cast(end)); } -void PostingListBuilder::write_header(uint32_t docid) +void handle_directory(FILE *fp, DatabaseReceiver *receiver) { - unsigned char buf[P4NENC_BOUND(1)]; - size_t bytes = p4nd1enc128v32(&docid, 1, buf); - encoded.append(reinterpret_cast(buf), bytes); -} - -class Corpus { -public: - Corpus(size_t block_size) : block_size(block_size) {} - void add_file(string filename); - void flush_block(); - - vector filename_blocks; - unordered_map invindex; - size_t num_files = 0, num_files_in_block = 0, num_blocks = 0; - -private: - string current_block; - string tempbuf; - const size_t block_size; -}; - -void Corpus::add_file(string filename) -{ - ++num_files; - if (!current_block.empty()) { - current_block.push_back('\0'); - } - current_block += filename; - if (++num_files_in_block == block_size) { - flush_block(); - } -} - -void Corpus::flush_block() -{ - if (current_block.empty()) { - return; - } - - uint32_t docid = num_blocks; - - // 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); - invindex[trgm].add_docid(docid); - } + db_directory dummy; + if (fread(&dummy, sizeof(dummy), 1, fp) != 1) { + if (feof(fp)) { + return; + } else { + perror("fread"); } - ptr += s.size() + 1; } - // Compress and add the filename block. - filename_blocks.push_back(zstd_compress(current_block, &tempbuf)); - - current_block.clear(); - num_files_in_block = 0; - ++num_blocks; -} - -const char *handle_directory(const char *ptr, Corpus *corpus) -{ - ptr += sizeof(db_directory); - - string dir_path = ptr; - ptr += dir_path.size() + 1; + string dir_path = read_cstr(fp); if (dir_path == "/") { dir_path = ""; } - - for ( ;; ) { - uint8_t type = *ptr++; + + for (;;) { + int type = getc(fp); if (type == DBE_NORMAL) { - string filename = ptr; - corpus->add_file(dir_path + "/" + filename); - ptr += filename.size() + 1; + string filename = read_cstr(fp); + receiver->add_file(dir_path + "/" + filename, unknown_dir_time); } else if (type == DBE_DIRECTORY) { - string dirname = ptr; - corpus->add_file(dir_path + "/" + dirname); - ptr += dirname.size() + 1; + string dirname = read_cstr(fp); + receiver->add_file(dir_path + "/" + dirname, unknown_dir_time); } else { - return ptr; + return; // Probably end. } } } -void read_mlocate(const char *filename, Corpus *corpus) +void read_plaintext(FILE *fp, DatabaseReceiver *receiver) { - int fd = open(filename, O_RDONLY); - if (fd == -1) { - perror(filename); - exit(1); - } - off_t len = lseek(fd, 0, SEEK_END); - if (len == -1) { - perror("lseek"); - exit(1); - } - const char *data = (char *)mmap(nullptr, len, PROT_READ, MAP_SHARED, fd, /*offset=*/0); - if (data == MAP_FAILED) { - perror("mmap"); + if (fseek(fp, 0, SEEK_SET) != 0) { + perror("fseek"); exit(1); } - const db_header *hdr = (const db_header *)data; - - // TODO: Care about the base path. - string path = data + sizeof(db_header); - uint64_t offset = sizeof(db_header) + path.size() + 1 + ntohl(hdr->conf_size); - - const char *ptr = data + offset; - while (ptr < data + len) { - ptr = handle_directory(ptr, corpus); - } - - munmap((void *)data, len); - close(fd); -} - -string zstd_compress(const string &src, string *tempbuf) -{ - size_t max_size = ZSTD_compressBound(src.size()); - if (tempbuf->size() < max_size) { - tempbuf->resize(max_size); - } - size_t size = ZSTD_compress(&(*tempbuf)[0], max_size, src.data(), src.size(), /*level=*/6); - return string(tempbuf->data(), size); -} - -void do_build(const char *infile, const char *outfile, int block_size) -{ - //steady_clock::time_point start = steady_clock::now(); - - Corpus corpus(block_size); - - read_mlocate(infile, &corpus); - if (false) { // To read a plain text file. - FILE *fp = fopen(infile, "r"); - while (!feof(fp)) { - char buf[1024]; - if (fgets(buf, 1024, fp) == nullptr || feof(fp)) { + while (!feof(fp)) { + char buf[1024]; + if (fgets(buf, sizeof(buf), fp) == nullptr) { + break; + } + string s(buf); + assert(!s.empty()); + while (s.back() != '\n' && !feof(fp)) { + // The string was longer than the buffer, so read again. + if (fgets(buf, sizeof(buf), fp) == nullptr) { break; } - string s(buf); - if (s.back() == '\n') s.pop_back(); - corpus.add_file(move(s)); + s += buf; } - fclose(fp); + if (!s.empty() && s.back() == '\n') + s.pop_back(); + receiver->add_file(move(s), unknown_dir_time); } - corpus.flush_block(); - dprintf("Read %zu files from %s\n", corpus.num_files, infile); +} - size_t trigrams = 0, longest_posting_list = 0; - size_t bytes_used = 0; - for (auto &[trigram, pl_builder] : corpus.invindex) { - pl_builder.finish(); - longest_posting_list = max(longest_posting_list, pl_builder.num_docids); - trigrams += pl_builder.num_docids; - bytes_used += pl_builder.encoded.size(); +void read_mlocate(FILE *fp, DatabaseReceiver *receiver) +{ + if (fseek(fp, 0, SEEK_SET) != 0) { + perror("fseek"); + exit(1); } - dprintf("%zu files, %zu different trigrams, %zu entries, avg len %.2f, longest %zu\n", - corpus.num_files, corpus.invindex.size(), trigrams, double(trigrams) / corpus.invindex.size(), longest_posting_list); - dprintf("%zu bytes used for posting lists (%.2f bits/entry)\n", bytes_used, 8 * bytes_used / double(trigrams)); - - //steady_clock::time_point end = steady_clock::now(); - dprintf("Building posting lists took %.1f ms.\n\n", 1e3 * duration(end - start).count()); - vector all_trigrams; - for (auto &[trigram, pl_builder] : corpus.invindex) { - all_trigrams.push_back(trigram); + db_header hdr; + if (fread(&hdr, sizeof(hdr), 1, fp) != 1) { + perror("short read"); + exit(1); } - sort(all_trigrams.begin(), all_trigrams.end()); - umask(0027); - FILE *outfp = fopen(outfile, "wb"); - - // Write the header. - uint64_t num_trigrams = corpus.invindex.size(); // 64 to get alignment. - fwrite(&num_trigrams, sizeof(num_trigrams), 1, outfp); + // TODO: Care about the base path. + string path = read_cstr(fp); - // Find out where the posting lists will end. - uint64_t offset = sizeof(uint64_t) * 2 + 16 * all_trigrams.size(); - uint64_t filename_index_offset = offset; - for (auto &[trigram, pl_builder] : corpus.invindex) { - filename_index_offset += pl_builder.encoded.size(); + if (fseek(fp, ntohl(hdr.conf_size), SEEK_CUR) != 0) { + perror("skip conf block"); + exit(1); } - fwrite(&filename_index_offset, sizeof(filename_index_offset), 1, outfp); - size_t bytes_for_trigrams = 0, bytes_for_posting_lists = 0, bytes_for_filename_index = 0, bytes_for_filenames = 0; - for (uint32_t trgm : all_trigrams) { - // 16 bytes: trgm (4), number of docids (4), byte offset (8) - fwrite(&trgm, sizeof(trgm), 1, outfp); // One byte wasted, but OK. - const PostingListBuilder &pl_builder = corpus.invindex[trgm]; - uint32_t num_docids = pl_builder.num_docids; - fwrite(&num_docids, sizeof(num_docids), 1, outfp); - fwrite(&offset, sizeof(offset), 1, outfp); - offset += pl_builder.encoded.size(); - bytes_for_trigrams += 16; + while (!feof(fp)) { + handle_directory(fp, receiver); } +} - // Write the actual posting lists. - for (uint32_t trgm : all_trigrams) { - const string &encoded = corpus.invindex[trgm].encoded; - fwrite(encoded.data(), encoded.size(), 1, outfp); - bytes_for_posting_lists += encoded.size(); +void do_build(const char *infile, const char *outfile, int block_size, bool plaintext) +{ + FILE *infp = fopen(infile, "rb"); + if (infp == nullptr) { + perror(infile); + exit(1); } - // Stick an empty block at the end as sentinel. - corpus.filename_blocks.push_back(""); - - // Write the offsets to the filenames. - offset = filename_index_offset + corpus.filename_blocks.size() * sizeof(offset); - for (const string &filename : corpus.filename_blocks) { - fwrite(&offset, sizeof(offset), 1, outfp); - offset += filename.size(); - bytes_for_filename_index += sizeof(offset); - bytes_for_filenames += filename.size(); + // Train the dictionary by sampling real blocks. + // The documentation for ZDICT_trainFromBuffer() claims that a reasonable + // dictionary size is ~100 kB, but 1 kB seems to actually compress better for us, + // and decompress just as fast. + DictionaryBuilder builder(/*blocks_to_keep=*/1000, block_size); + if (plaintext) { + read_plaintext(infp, &builder); + } else { + read_mlocate(infp, &builder); } - - // Write the actual filenames. - for (const string &filename : corpus.filename_blocks) { - fwrite(filename.data(), filename.size(), 1, outfp); + string dictionary = builder.train(1024); + + DatabaseBuilder db(outfile, /*owner=*/-1, block_size, dictionary, /*check_visibility=*/true); + DatabaseReceiver *corpus = db.start_corpus(/*store_dir_times=*/false); + if (plaintext) { + read_plaintext(infp, corpus); + } else { + read_mlocate(infp, corpus); } + fclose(infp); - fclose(outfp); + dprintf("Read %zu files from %s\n", corpus->num_files_seen(), infile); + db.finish_corpus(); +} - //size_t total_bytes = (bytes_for_trigrams + bytes_for_posting_lists + bytes_for_filename_index + bytes_for_filenames); +void usage() +{ + printf( + "Usage: plocate-build MLOCATE_DB PLOCATE_DB\n" + "\n" + "Generate plocate index from mlocate.db, typically /var/lib/mlocate/mlocate.db.\n" + "Normally, the destination should be /var/lib/mlocate/plocate.db.\n" + "\n" + " -b, --block-size SIZE number of filenames to store in each block (default 32)\n" + " -p, --plaintext input is a plaintext file, not an mlocate database\n" + " --help print this help\n" + " --version print version information\n"); +} - dprintf("Block size: %7d files\n", block_size); - dprintf("Trigrams: %'7.1f MB\n", bytes_for_trigrams / 1048576.0); - dprintf("Posting lists: %'7.1f MB\n", bytes_for_posting_lists / 1048576.0); - dprintf("Filename index: %'7.1f MB\n", bytes_for_filename_index / 1048576.0); - dprintf("Filenames: %'7.1f MB\n", bytes_for_filenames / 1048576.0); - dprintf("Total: %'7.1f MB\n", total_bytes / 1048576.0); - dprintf("\n"); +void version() +{ + printf("plocate-build %s\n", PACKAGE_VERSION); + printf("Copyright 2020 Steinar H. Gunderson\n"); + printf("License GPLv2+: GNU GPL version 2 or later .\n"); + printf("This is free software: you are free to change and redistribute it.\n"); + printf("There is NO WARRANTY, to the extent permitted by law.\n"); } int main(int argc, char **argv) { - do_build(argv[1], argv[2], 32); + static const struct option long_options[] = { + { "block-size", required_argument, 0, 'b' }, + { "plaintext", no_argument, 0, 'p' }, + { "help", no_argument, 0, 'h' }, + { "version", no_argument, 0, 'V' }, + { "debug", no_argument, 0, 'D' }, // Not documented. + { 0, 0, 0, 0 } + }; + + int block_size = 32; + bool plaintext = false; + + setlocale(LC_ALL, ""); + for (;;) { + int option_index = 0; + int c = getopt_long(argc, argv, "b:hpVD", long_options, &option_index); + if (c == -1) { + break; + } + switch (c) { + case 'b': + block_size = atoi(optarg); + break; + case 'p': + plaintext = true; + break; + case 'h': + usage(); + exit(0); + case 'V': + version(); + exit(0); + case 'D': + use_debug = true; + break; + default: + exit(1); + } + } + + if (argc - optind != 2) { + usage(); + exit(1); + } + + do_build(argv[optind], argv[optind + 1], block_size, plaintext); exit(EXIT_SUCCESS); }