X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=plocate-build.cpp;h=c42bd86113cd2eae8ba53cba7ca9625e573d9465;hb=cbecd483f57c465b2ad6d3867c760c2e5b5e79aa;hp=3f1c0c7492d1c683f715fd5dd6e08b7858382a07;hpb=8502ed727dcd5652fcec6aa68d934cd2af116052;p=plocate

diff --git a/plocate-build.cpp b/plocate-build.cpp
index 3f1c0c7..c42bd86 100644
--- a/plocate-build.cpp
+++ b/plocate-build.cpp
@@ -1,4 +1,5 @@
 #include "db.h"
+#include "dprintf.h"
 #include "turbopfor-encode.h"
 
 #include <algorithm>
@@ -22,8 +23,6 @@
 #include <zstd.h>
 
 #define P4NENC_BOUND(n) ((n + 127) / 128 + (n + 32) * sizeof(uint32_t))
-#define dprintf(...)
-//#define dprintf(...) fprintf(stderr, __VA_ARGS__);
 
 #define NUM_TRIGRAMS 16777216
 
@@ -33,6 +32,7 @@ using namespace std::chrono;
 string zstd_compress(const string &src, ZSTD_CDict *cdict, string *tempbuf);
 
 constexpr unsigned num_overflow_slots = 16;
+bool use_debug = false;
 
 static inline uint32_t read_unigram(const string_view s, size_t idx)
 {
@@ -222,7 +222,10 @@ 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());
-	dprintf(stderr, "Sampled %zu bytes in %zu blocks, built a dictionary of size %zu\n", dictionary_buf.size(), lengths.size(), ret);
+	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);
 
 	sampled_blocks.clear();
@@ -261,6 +264,7 @@ public:
 		}
 		return *invindex[trgm];
 	}
+	size_t num_trigrams() const;
 
 private:
 	unique_ptr<PostingListBuilder *[]> invindex;
@@ -317,6 +321,17 @@ void Corpus::flush_block()
 	++num_blocks;
 }
 
+size_t Corpus::num_trigrams() const
+{
+	size_t num = 0;
+	for (unsigned trgm = 0; trgm < NUM_TRIGRAMS; ++trgm) {
+		if (invindex[trgm] != nullptr) {
+			++num;
+		}
+	}
+	return num;
+}
+
 string read_cstr(FILE *fp)
 {
 	string ret;
@@ -363,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)
+{
+	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)
 {
-	FILE *fp = fopen(filename, "rb");
-	if (fp == nullptr) {
-		perror(filename);
+	if (fseek(fp, 0, SEEK_SET) != 0) {
+		perror("fseek");
 		exit(1);
 	}
 
@@ -382,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)
@@ -462,12 +502,22 @@ 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");
+	if (outfp == nullptr) {
+		perror(outfile);
+		exit(1);
+	}
 
 	// Write the header.
 	Header hdr;
@@ -487,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);
 
@@ -496,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();
@@ -538,8 +584,9 @@ void do_build(const char *infile, const char *outfile, int block_size)
 		trigrams += pl_builder.num_docids;
 		bytes_for_posting_lists += pl_builder.encoded.size();
 	}
+	size_t num_trigrams = corpus.num_trigrams();
 	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);
+	        corpus.num_files, num_trigrams, trigrams, double(trigrams) / num_trigrams, longest_posting_list);
 	dprintf("%zu bytes used for posting lists (%.2f bits/entry)\n", bytes_for_posting_lists, 8 * bytes_for_posting_lists / double(trigrams));
 
 	dprintf("Building posting lists took %.1f ms.\n\n", 1e3 * duration<float>(steady_clock::now() - start).count());
@@ -599,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);
@@ -620,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");
 }
@@ -637,17 +685,20 @@ 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.
 		{ 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:hV", long_options, &option_index);
+		int c = getopt_long(argc, argv, "b:hpVD", long_options, &option_index);
 		if (c == -1) {
 			break;
 		}
@@ -655,12 +706,18 @@ int main(int argc, char **argv)
 		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);
 		}
@@ -671,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);
 }