std::vector<uint64_t> v(s, 0);

for (size_t i = 0; i < s; ++i) {

v[i] = rand();

}

return v;

std::vector<float> v(s, 0);

for (size_t i = 0; i < s; ++i) {

v[i] = rand();

}

return v;

const int nsubcodes = 97;

const int minbits = 1;

const int maxbits = 24;

const std::vector<uint64_t> values = random_vector(nsubcodes);

for (int nbits = minbits; nbits <= maxbits; ++nbits) {

std::cerr << "nbits = " << nbits << std::endl;

const uint64_t mask = (1ull << nbits) - 1;

std::unique_ptr<uint8_t[]> codes(

new uint8_t[(nsubcodes * maxbits + 7) / 8]);

// NOTE(hoss): Necessary scope to ensure trailing bits are flushed to

// mem.

{

faiss::PQEncoderGeneric encoder(codes.get(), nbits);

for (const auto& v : values) {

encoder.encode(v & mask);

}

}

faiss::PQDecoderGeneric decoder(codes.get(), nbits);

for (int i = 0; i < nsubcodes; ++i) {

uint64_t v = decoder.decode();

EXPECT_EQ(values[i] & mask, v);

}

}

const int nsubcodes = 100;

const std::vector<uint64_t> values = random_vector(nsubcodes);

const uint64_t mask = 0xFF;

std::unique_ptr<uint8_t[]> codes(new uint8_t[nsubcodes]);

faiss::PQEncoder8 encoder(codes.get(), 8);

for (const auto& v : values) {

encoder.encode(v & mask);

}

faiss::PQDecoder8 decoder(codes.get(), 8);

for (int i = 0; i < nsubcodes; ++i) {

uint64_t v = decoder.decode();

EXPECT_EQ(values[i] & mask, v);

}

const int nsubcodes = 100;

const std::vector<uint64_t> values = random_vector(nsubcodes);

const uint64_t mask = 0xFFFF;

std::unique_ptr<uint8_t[]> codes(new uint8_t[2 * nsubcodes]);

faiss::PQEncoder16 encoder(codes.get(), 16);

for (const auto& v : values) {

encoder.encode(v & mask);

}

faiss::PQDecoder16 decoder(codes.get(), 16);

for (int i = 0; i < nsubcodes; ++i) {

uint64_t v = decoder.decode();

EXPECT_EQ(values[i] & mask, v);

}

int d = 20, ntotal = 1000, M = 5, nbits = 4;

const std::vector<float> ds = random_vector_float(ntotal * d);

faiss::IndexPQFastScan index(d, M, nbits);

index.train(ntotal, ds.data());

index.add(ntotal, ds.data());

for (int j = 0; j < 10; j++) {

int vector_id = rand() % ntotal;

std::vector<uint8_t> old(ntotal * M);

std::vector<uint8_t> code(M);

for (int i = 0; i < ntotal; i++) {

for (int sq = 0; sq < M; sq++) {

old[i * M + sq] = faiss::pq4_get_packed_element(

index.codes.data(), index.bbs, M, i, sq);

}

}

for (int sq = 0; sq < M; sq++) {

faiss::pq4_set_packed_element(

index.codes.data(),

((old[vector_id * M + sq] + 3) % 16),

index.bbs,

M,

vector_id,

sq);

}

for (int i = 0; i < ntotal; i++) {

for (int sq = 0; sq < M; sq++) {

uint8_t newcode = faiss::pq4_get_packed_element(

index.codes.data(), index.bbs, M, i, sq);

uint8_t oldcode = old[i * M + sq];

if (i == vector_id) {

EXPECT_EQ(newcode, (oldcode + 3) % 16);

} else {

EXPECT_EQ(newcode, oldcode);

}

}

}

}