Escape unprintable characters when outputting filenames to a terminal.

[plocate] / plocate-build.cpp

diff --git a/plocate-build.cpp b/plocate-build.cpp
index 2093ab2010578109eb4dd02412a9a037edc7ea19..c42bd86113cd2eae8ba53cba7ca9625e573d9465 100644 (file)
--- a/plocate-build.cpp
+++ b/plocate-build.cpp
@@ -222,6 +222,9 @@ string DictionaryBuilder::train(size_t buf_size)
        string buf;
        buf.resize(buf_size);
        size_t ret = ZDICT_trainFromBuffer(&buf[0], buf_size, dictionary_buf.data(), lengths.data(), lengths.size());
+       if (ret == size_t(-1)) {
+               return "";
+       }
        dprintf("Sampled %zu bytes in %zu blocks, built a dictionary of size %zu\n", dictionary_buf.size(), lengths.size(), ret);
        buf.resize(ret);
 
@@ -375,11 +378,37 @@ void handle_directory(FILE *fp, DatabaseReceiver *receiver)
        }
 }
 
-void read_mlocate(const char *filename, DatabaseReceiver *receiver)
+void read_plaintext(FILE *fp, DatabaseReceiver *receiver)
 {
-       FILE *fp = fopen(filename, "rb");
-       if (fp == nullptr) {
-               perror(filename);
+       if (fseek(fp, 0, SEEK_SET) != 0) {
+               perror("fseek");
+               exit(1);
+       }
+
+       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;
+                       }
+                       s += buf;
+               }
+               if (!s.empty() && s.back() == '\n')
+                       s.pop_back();
+               receiver->add_file(move(s));
+       }
+}
+
+void read_mlocate(FILE *fp, DatabaseReceiver *receiver)
+{
+       if (fseek(fp, 0, SEEK_SET) != 0) {
+               perror("fseek");
                exit(1);
        }
 
@@ -394,7 +423,6 @@ void read_mlocate(const char *filename, DatabaseReceiver *receiver)
        while (!feof(fp)) {
                handle_directory(fp, receiver);
        }
-       fclose(fp);
 }
 
 string zstd_compress(const string &src, ZSTD_CDict *cdict, string *tempbuf)
@@ -474,9 +502,15 @@ unique_ptr<Trigram[]> create_hashtable(Corpus &corpus, const vector<uint32_t> &a
        return ht;
 }
 
-void do_build(const char *infile, const char *outfile, int block_size)
+void do_build(const char *infile, const char *outfile, int block_size, bool plaintext)
 {
-       steady_clock::time_point start __attribute__((unused)) = steady_clock::now();
+       steady_clock::time_point start = steady_clock::now();
+
+       FILE *infp = fopen(infile, "rb");
+       if (infp == nullptr) {
+               perror(infile);
+               exit(1);
+       }
 
        umask(0027);
        FILE *outfp = fopen(outfile, "wb");
@@ -503,7 +537,11 @@ void do_build(const char *infile, const char *outfile, int block_size)
        // 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);
-       read_mlocate(infile, &builder);
+       if (plaintext) {
+               read_plaintext(infp, &builder);
+       } else {
+               read_mlocate(infp, &builder);
+       }
        string dictionary = builder.train(1024);
        ZSTD_CDict *cdict = ZSTD_createCDict(dictionary.data(), dictionary.size(), /*level=*/6);
 
@@ -512,21 +550,13 @@ void do_build(const char *infile, const char *outfile, int block_size)
        hdr.zstd_dictionary_length_bytes = dictionary.size();
 
        Corpus corpus(outfp, block_size, cdict);
-       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)) {
-                               break;
-                       }
-                       string s(buf);
-                       if (s.back() == '\n')
-                               s.pop_back();
-                       corpus.add_file(move(s));
-               }
-               fclose(fp);
+       if (plaintext) {
+               read_plaintext(infp, &corpus);
+       } else {
+               read_mlocate(infp, &corpus);
        }
+       fclose(infp);
+
        corpus.flush_block();
        dprintf("Read %zu files from %s\n", corpus.num_files, infile);
        hdr.num_docids = corpus.filename_blocks.size();
@@ -616,7 +646,7 @@ void do_build(const char *infile, const char *outfile, int block_size)
        fwrite(&hdr, sizeof(hdr), 1, outfp);
        fclose(outfp);
 
-       size_t total_bytes __attribute__((unused)) = (bytes_for_hashtable + bytes_for_posting_lists + bytes_for_filename_index + bytes_for_filenames);
+       size_t total_bytes = (bytes_for_hashtable + bytes_for_posting_lists + bytes_for_filename_index + bytes_for_filenames);
 
        dprintf("Block size:     %7d files\n", block_size);
        dprintf("Dictionary:     %'7.1f MB\n", dictionary.size() / 1048576.0);
@@ -637,6 +667,7 @@ void usage()
                "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");
 }
@@ -654,6 +685,7 @@ int main(int argc, char **argv)
 {
        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.
@@ -661,11 +693,12 @@ int main(int argc, char **argv)
        };
 
        int block_size = 32;
+       bool plaintext = false;
 
        setlocale(LC_ALL, "");
        for (;;) {
                int option_index = 0;
-               int c = getopt_long(argc, argv, "b:hVD", long_options, &option_index);
+               int c = getopt_long(argc, argv, "b:hpVD", long_options, &option_index);
                if (c == -1) {
                        break;
                }
@@ -673,6 +706,9 @@ int main(int argc, char **argv)
                case 'b':
                        block_size = atoi(optarg);
                        break;
+               case 'p':
+                       plaintext = true;
+                       break;
                case 'h':
                        usage();
                        exit(0);
@@ -692,6 +728,6 @@ int main(int argc, char **argv)
                exit(1);
        }
 
-       do_build(argv[optind], argv[optind + 1], block_size);
+       do_build(argv[optind], argv[optind + 1], block_size, plaintext);
        exit(EXIT_SUCCESS);
 }