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