#include <algorithm>
#include <assert.h>
+#ifdef HAS_ENDIAN_H
+#include <endian.h>
+#endif
#include <fcntl.h>
#include <string.h>
#include <string_view>
#include <unistd.h>
#include <zdict.h>
#include <zstd.h>
+#include <unicode/coll.h>
+#include <unicode/locid.h>
#define P4NENC_BOUND(n) ((n + 127) / 128 + (n + 32) * sizeof(uint32_t))
string zstd_compress(const string &src, ZSTD_CDict *cdict, 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);
-}
-
class PostingListBuilder {
public:
inline void add_docid(uint32_t docid);
inline void add_first_docid(uint32_t docid);
void finish();
- string encoded;
- size_t num_docids = 0;
+ vector<unsigned char> encoded;
+ size_t get_num_docids() const
+ {
+ // Updated only when we flush, so check that we're finished.
+ assert(pending_deltas.empty());
+ return num_docids;
+ }
private:
void write_header(uint32_t docid);
vector<uint32_t> pending_deltas;
+ uint32_t num_docids = 0; // Should be size_t, except the format only supports 2^32 docids per posting list anyway.
uint32_t last_docid = -1;
};
return;
}
- assert(num_docids != 0);
-
pending_deltas.push_back(docid - last_docid - 1);
last_docid = docid;
if (pending_deltas.size() == 128) {
append_block();
pending_deltas.clear();
+ num_docids += 128;
}
- ++num_docids;
}
void PostingListBuilder::add_first_docid(uint32_t docid)
// No interleaving for partial blocks.
unsigned char buf[P4NENC_BOUND(128)];
unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), pending_deltas.size(), /*interleaved=*/false, buf);
- encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
+ encoded.insert(encoded.end(), buf, end);
+
+ num_docids += pending_deltas.size();
+ pending_deltas.clear();
}
void PostingListBuilder::append_block()
unsigned char buf[P4NENC_BOUND(128)];
assert(pending_deltas.size() == 128);
unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), 128, /*interleaved=*/true, buf);
- encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
+ encoded.insert(encoded.end(), buf, end);
}
void PostingListBuilder::write_header(uint32_t docid)
{
unsigned char buf[P4NENC_BOUND(1)];
unsigned char *end = write_baseval(docid, buf);
- encoded.append(reinterpret_cast<char *>(buf), end - buf);
+ encoded.insert(encoded.end(), buf, end);
}
void DictionaryBuilder::add_file(string filename, dir_time)
std::unique_ptr<PostingListBuilder *[]> invindex;
FILE *outfp;
+ off_t outfp_pos; // Cheaper than calling ftell(outfp) all the time.
std::string current_block;
std::string tempbuf;
const size_t block_size;
std::string dir_times_compressed;
};
-
EncodingCorpus::EncodingCorpus(FILE *outfp, size_t block_size, ZSTD_CDict *cdict, bool store_dir_times)
- : invindex(new PostingListBuilder *[NUM_TRIGRAMS]), outfp(outfp), block_size(block_size), store_dir_times(store_dir_times), cdict(cdict)
+ : invindex(new PostingListBuilder *[NUM_TRIGRAMS]), outfp(outfp), outfp_pos(ftell(outfp)), block_size(block_size), store_dir_times(store_dir_times), cdict(cdict)
{
fill(invindex.get(), invindex.get() + NUM_TRIGRAMS, nullptr);
if (store_dir_times) {
uint32_t docid = num_blocks;
+ // Oh, ICU...
+ vector<uint8_t> sort_key;
+ sort_key.resize(32);
+ int32_t num_locales;
+ const icu::Locale* locales = icu::Collator::getAvailableLocales(num_locales);
+ for (int i = 0; i < num_locales; ++i) {
+ const icu::Locale &loc = locales[i];
+ if (strcmp(loc.getName(), "en_US_POSIX") == 0) {
+ continue; // Too weird.
+ }
+ UErrorCode status = U_ZERO_ERROR;
+ icu::Collator *coll = icu::Collator::createInstance(loc, status);
+ if (U_FAILURE(status)) {
+ fprintf(stderr, "ERROR: Failed to create collator\n");
+ exit(1);
+ }
+ coll->setStrength(icu::Collator::PRIMARY);
+ const char *ptr = current_block.c_str();
+ const char *end = ptr + current_block.size();
+ while (ptr < end) {
+ size_t len = strlen(ptr);
+ int32_t sortkey_len;
+ for ( ;; ) {
+ sortkey_len = coll->getSortKey(icu::UnicodeString::fromUTF8(icu::StringPiece(ptr, len)), sort_key.data(), sort_key.size());
+ if (sortkey_len < sort_key.size()) { // Note <, not <=; we need to keep a slop byte.
+ break;
+ }
+ sort_key.resize(sortkey_len * 3 / 2);
+ }
+
+ const uint8_t *keyptr = &sort_key[0];
+ const uint8_t *keyend = keyptr + sortkey_len;
+ while (keyptr < keyend - 3) {
+ // NOTE: Will read one byte past the end of the trigram, but it's OK,
+ // since we always call it from contexts where there's a terminating zero byte.
+ uint32_t trgm;
+ memcpy(&trgm, keyptr, sizeof(trgm));
+ ++keyptr;
+ trgm = le32toh(trgm);
+ add_docid(trgm & 0xffffff, docid);
+ }
+
+ ptr += len + 1;
+ }
+ }
+#if 0
// 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);
- add_docid(trgm, docid);
+ const char *end = ptr + current_block.size();
+ while (ptr < end - 3) { // Must be at least one filename left, that's at least three bytes.
+ if (ptr[0] == '\0') {
+ // This filename is zero bytes, so skip it (and the zero terminator).
+ ++ptr;
+ continue;
+ } else if (ptr[1] == '\0') {
+ // This filename is one byte, so skip it (and the zero terminator).
+ ptr += 2;
+ continue;
+ } else if (ptr[2] == '\0') {
+ // This filename is two bytes, so skip it (and the zero terminator).
+ ptr += 3;
+ continue;
+ }
+ for (;;) {
+ // NOTE: Will read one byte past the end of the trigram, but it's OK,
+ // since we always call it from contexts where there's a terminating zero byte.
+ uint32_t trgm;
+ memcpy(&trgm, ptr, sizeof(trgm));
+ ++ptr;
+ trgm = le32toh(trgm);
+ add_docid(trgm & 0xffffff, docid);
+ if (trgm <= 0xffffff) {
+ // Terminating zero byte, so we're done with this filename.
+ // Skip the remaining two bytes, and the zero terminator.
+ ptr += 3;
+ break;
}
}
- ptr += s.size() + 1;
}
+#endif
// Compress and add the filename block.
- filename_blocks.push_back(ftell(outfp));
+ filename_blocks.push_back(outfp_pos);
string compressed = zstd_compress(current_block, cdict, &tempbuf);
if (fwrite(compressed.data(), compressed.size(), 1, outfp) != 1) {
perror("fwrite()");
exit(1);
}
+ outfp_pos += compressed.size();
current_block.clear();
num_files_in_block = 0;
}
for (uint32_t trgm : all_trigrams) {
// We don't know offset yet, so set it to zero.
- Trigram to_insert{ trgm, uint32_t(corpus.get_pl_builder(trgm).num_docids), 0 };
+ Trigram to_insert{ trgm, uint32_t(corpus.get_pl_builder(trgm).get_num_docids()), 0 };
uint32_t bucket = hash_trigram(trgm, ht_size);
unsigned distance = 0;
if (path.empty()) {
path = ".";
}
+ int fd = -1;
#ifdef O_TMPFILE
- int fd = open(path.c_str(), O_WRONLY | O_TMPFILE, 0640);
- if (fd == -1) {
+ fd = open(path.c_str(), O_WRONLY | O_TMPFILE, 0640);
+ if (fd == -1 && errno != EOPNOTSUPP && errno != EISDIR) {
perror(path.c_str());
exit(1);
}
-#else
- temp_filename = string(outfile) + ".XXXXXX";
- int fd = mkstemp(&temp_filename[0]);
+#endif
if (fd == -1) {
- perror(temp_filename.c_str());
- exit(1);
- }
- if (fchmod(fd, 0640) == -1) {
- perror("fchmod");
- exit(1);
+ temp_filename = string(outfile) + ".XXXXXX";
+ fd = mkstemp(&temp_filename[0]);
+ if (fd == -1) {
+ perror(temp_filename.c_str());
+ exit(1);
+ }
+ if (fchmod(fd, 0640) == -1) {
+ perror("fchmod");
+ exit(1);
+ }
}
-#endif
if (owner != (gid_t)-1) {
if (fchown(fd, (uid_t)-1, owner) == -1) {
continue;
PostingListBuilder &pl_builder = corpus->get_pl_builder(trgm);
pl_builder.finish();
- longest_posting_list = max(longest_posting_list, pl_builder.num_docids);
- trigrams += pl_builder.num_docids;
+ longest_posting_list = max(longest_posting_list, pl_builder.get_num_docids());
+ trigrams += pl_builder.get_num_docids();
bytes_for_posting_lists += pl_builder.encoded.size();
}
size_t num_trigrams = corpus->num_trigrams();
continue;
}
- const string &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
+ const vector<unsigned char> &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
offset += encoded.size();
}
if (hashtable[i].num_docids == 0) {
continue;
}
- const string &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
+ const vector<unsigned char> &encoded = corpus->get_pl_builder(hashtable[i].trgm).encoded;
fwrite(encoded.data(), encoded.size(), 1, outfp);
}
fseek(outfp, 0, SEEK_SET);
fwrite(&hdr, sizeof(hdr), 1, outfp);
+ if (!temp_filename.empty()) {
+ if (rename(temp_filename.c_str(), outfile.c_str()) == -1) {
+ perror("rename");
+ exit(1);
+ }
+ } else {
#ifdef O_TMPFILE
- // Give the file a proper name, making it visible in the file system.
- // TODO: It would be nice to be able to do this atomically, like with rename.
- unlink(outfile.c_str());
- char procpath[256];
- snprintf(procpath, sizeof(procpath), "/proc/self/fd/%d", fileno(outfp));
- if (linkat(AT_FDCWD, procpath, AT_FDCWD, outfile.c_str(), AT_SYMLINK_FOLLOW) == -1) {
- perror("linkat");
- exit(1);
- }
-#else
- if (rename(temp_filename.c_str(), outfile.c_str()) == -1) {
- perror("rename");
- exit(1);
- }
+ // Give the file a proper name, making it visible in the file system.
+ // TODO: It would be nice to be able to do this atomically, like with rename.
+ unlink(outfile.c_str());
+ char procpath[256];
+ snprintf(procpath, sizeof(procpath), "/proc/self/fd/%d", fileno(outfp));
+ if (linkat(AT_FDCWD, procpath, AT_FDCWD, outfile.c_str(), AT_SYMLINK_FOLLOW) == -1) {
+ perror("linkat");
+ exit(1);
+ }
#endif
+ }
fclose(outfp);