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

   1 #include "database-builder.h"
   2 #include "db.h"
   3 #include "dprintf.h"
   4
   5 #include <algorithm>
   6 #include <arpa/inet.h>
   7 #include <assert.h>
   8 #include <chrono>
   9 #include <getopt.h>
  10 #include <iosfwd>
  11 #include <locale.h>
  12 #include <math.h>
  13 #include <memory>
  14 #include <random>
  15 #include <stdint.h>
  16 #include <stdio.h>
  17 #include <stdlib.h>
  18 #include <string.h>
  19 #include <string>
  20 #include <sys/stat.h>
  21 #include <utility>
  22 #include <vector>
  23
  24 using namespace std;
  25 using namespace std::chrono;
  26
  27 bool use_debug = false;
  28
  29 enum {
  30         DBE_NORMAL = 0, /* A non-directory file */
  31         DBE_DIRECTORY = 1, /* A directory */
  32         DBE_END = 2 /* End of directory contents; contains no name */
  33 };
  34
  35 // From mlocate.
  36 struct db_header {
  37         uint8_t magic[8];
  38         uint32_t conf_size;
  39         uint8_t version;
  40         uint8_t check_visibility;
  41         uint8_t pad[2];
  42 };
  43
  44 // From mlocate.
  45 struct db_directory {
  46         uint64_t time_sec;
  47         uint32_t time_nsec;
  48         uint8_t pad[4];
  49 };
  50
  51 string read_cstr(FILE *fp)
  52 {
  53         string ret;
  54         for (;;) {
  55                 int ch = getc(fp);
  56                 if (ch == -1) {
  57                         perror("getc");
  58                         exit(1);
  59                 }
  60                 if (ch == 0) {
  61                         return ret;
  62                 }
  63                 ret.push_back(ch);
  64         }
  65 }
  66
  67 void handle_directory(FILE *fp, DatabaseReceiver *receiver)
  68 {
  69         db_directory dummy;
  70         if (fread(&dummy, sizeof(dummy), 1, fp) != 1) {
  71                 if (feof(fp)) {
  72                         return;
  73                 } else {
  74                         perror("fread");
  75                 }
  76         }
  77
  78         string dir_path = read_cstr(fp);
  79         if (dir_path == "/") {
  80                 dir_path = "";
  81         }
  82
  83         for (;;) {
  84                 int type = getc(fp);
  85                 if (type == DBE_NORMAL) {
  86                         string filename = read_cstr(fp);
  87                         receiver->add_file(dir_path + "/" + filename, unknown_dir_time);
  88                 } else if (type == DBE_DIRECTORY) {
  89                         string dirname = read_cstr(fp);
  90                         receiver->add_file(dir_path + "/" + dirname, unknown_dir_time);
  91                 } else {
  92                         return;  // Probably end.
  93                 }
  94         }
  95 }
  96
  97 void read_plaintext(FILE *fp, DatabaseReceiver *receiver)
  98 {
  99         if (fseek(fp, 0, SEEK_SET) != 0) {
 100                 perror("fseek");
 101                 exit(1);
 102         }
 103
 104         while (!feof(fp)) {
 105                 char buf[1024];
 106                 if (fgets(buf, sizeof(buf), fp) == nullptr) {
 107                         break;
 108                 }
 109                 string s(buf);
 110                 assert(!s.empty());
 111                 while (s.back() != '\n' && !feof(fp)) {
 112                         // The string was longer than the buffer, so read again.
 113                         if (fgets(buf, sizeof(buf), fp) == nullptr) {
 114                                 break;
 115                         }
 116                         s += buf;
 117                 }
 118                 if (!s.empty() && s.back() == '\n')
 119                         s.pop_back();
 120                 receiver->add_file(move(s), unknown_dir_time);
 121         }
 122 }
 123
 124 void read_mlocate(FILE *fp, DatabaseReceiver *receiver)
 125 {
 126         if (fseek(fp, 0, SEEK_SET) != 0) {
 127                 perror("fseek");
 128                 exit(1);
 129         }
 130
 131         db_header hdr;
 132         if (fread(&hdr, sizeof(hdr), 1, fp) != 1) {
 133                 perror("short read");
 134                 exit(1);
 135         }
 136
 137         // TODO: Care about the base path.
 138         string path = read_cstr(fp);
 139
 140         if (fseek(fp, ntohl(hdr.conf_size), SEEK_CUR) != 0) {
 141                 perror("skip conf block");
 142                 exit(1);
 143         }
 144
 145         while (!feof(fp)) {
 146                 handle_directory(fp, receiver);
 147         }
 148 }
 149
 150 void do_build(const char *infile, const char *outfile, int block_size, bool plaintext)
 151 {
 152         FILE *infp = fopen(infile, "rb");
 153         if (infp == nullptr) {
 154                 perror(infile);
 155                 exit(1);
 156         }
 157
 158         // Train the dictionary by sampling real blocks.
 159         // The documentation for ZDICT_trainFromBuffer() claims that a reasonable
 160         // dictionary size is ~100 kB, but 1 kB seems to actually compress better for us,
 161         // and decompress just as fast.
 162         DictionaryBuilder builder(/*blocks_to_keep=*/1000, block_size);
 163         if (plaintext) {
 164                 read_plaintext(infp, &builder);
 165         } else {
 166                 read_mlocate(infp, &builder);
 167         }
 168         string dictionary = builder.train(1024);
 169
 170         DatabaseBuilder db(outfile, /*owner=*/-1, block_size, dictionary, /*check_visibility=*/true);
 171         DatabaseReceiver *corpus = db.start_corpus(/*store_dir_times=*/false);
 172         if (plaintext) {
 173                 read_plaintext(infp, corpus);
 174         } else {
 175                 read_mlocate(infp, corpus);
 176         }
 177         fclose(infp);
 178
 179         dprintf("Read %zu files from %s\n", corpus->num_files_seen(), infile);
 180         db.finish_corpus();
 181 }
 182
 183 void usage()
 184 {
 185         printf(
 186                 "Usage: plocate-build MLOCATE_DB PLOCATE_DB\n"
 187                 "\n"
 188                 "Generate plocate index from mlocate.db, typically /var/lib/mlocate/mlocate.db.\n"
 189                 "Normally, the destination should be /var/lib/mlocate/plocate.db.\n"
 190                 "\n"
 191                 "  -b, --block-size SIZE  number of filenames to store in each block (default 32)\n"
 192                 "  -p, --plaintext        input is a plaintext file, not an mlocate database\n"
 193                 "      --help             print this help\n"
 194                 "      --version          print version information\n");
 195 }
 196
 197 void version()
 198 {
 199         printf("plocate-build %s\n", PACKAGE_VERSION);
 200         printf("Copyright 2020 Steinar H. Gunderson\n");
 201         printf("License GPLv2+: GNU GPL version 2 or later <https://gnu.org/licenses/gpl.html>.\n");
 202         printf("This is free software: you are free to change and redistribute it.\n");
 203         printf("There is NO WARRANTY, to the extent permitted by law.\n");
 204 }
 205
 206 int main(int argc, char **argv)
 207 {
 208         static const struct option long_options[] = {
 209                 { "block-size", required_argument, 0, 'b' },
 210                 { "plaintext", no_argument, 0, 'p' },
 211                 { "help", no_argument, 0, 'h' },
 212                 { "version", no_argument, 0, 'V' },
 213                 { "debug", no_argument, 0, 'D' },  // Not documented.
 214                 { 0, 0, 0, 0 }
 215         };
 216
 217         int block_size = 32;
 218         bool plaintext = false;
 219
 220         setlocale(LC_ALL, "");
 221         for (;;) {
 222                 int option_index = 0;
 223                 int c = getopt_long(argc, argv, "b:hpVD", long_options, &option_index);
 224                 if (c == -1) {
 225                         break;
 226                 }
 227                 switch (c) {
 228                 case 'b':
 229                         block_size = atoi(optarg);
 230                         break;
 231                 case 'p':
 232                         plaintext = true;
 233                         break;
 234                 case 'h':
 235                         usage();
 236                         exit(0);
 237                 case 'V':
 238                         version();
 239                         exit(0);
 240                 case 'D':
 241                         use_debug = true;
 242                         break;
 243                 default:
 244                         exit(1);
 245                 }
 246         }
 247
 248         if (argc - optind != 2) {
 249                 usage();
 250                 exit(1);
 251         }
 252
 253         do_build(argv[optind], argv[optind + 1], block_size, plaintext);
 254         exit(EXIT_SUCCESS);
 255 }