faiss/tests/test_hnsw.cpp

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <gtest/gtest.h>

#include <cmath>
#include <cstddef>
#include <limits>
#include <random>
#include <thread>
#include <unordered_set>
#include <vector>

#include <faiss/IndexBinaryHNSW.h>
#include <faiss/IndexFlat.h>
#include <faiss/IndexHNSW.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/HNSW.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/VisitedTable.h>
#include <faiss/impl/hnsw/MinimaxHeap.h>
#include <faiss/utils/random.h>
#include <faiss/utils/utils.h>
#include <omp.h>

namespace {

struct SearchStatsEnabledGuard {
    explicit SearchStatsEnabledGuard(bool enabled)
            : previous(faiss::get_search_stats_enabled()) {
        faiss::set_search_stats_enabled(enabled);
    }
    ~SearchStatsEnabledGuard() {
        faiss::set_search_stats_enabled(previous);
    }
    bool previous;
};

} // namespace

TEST(HNSWStats, DisabledCollectionSkipsGlobalStats) {
    faiss::IndexBinaryHNSW index(8, 2);
    const uint8_t database = 0;
    index.add(1, &database);
    index.hnsw.efSearch = 1;
    SearchStatsEnabledGuard guard(false);
    faiss::hnsw_stats.reset();

    int32_t distance;
    faiss::idx_t label;
    index.search(1, &database, 1, &distance, &label);

    EXPECT_EQ(0, faiss::hnsw_stats.n1);
}

TEST(HNSWStats, EnabledCollectionAggregatesConcurrentSearchesExactly) {
    constexpr size_t num_threads = 16;
    constexpr size_t searches_per_thread = 100;

    faiss::IndexBinaryHNSW index(8, 2);
    const uint8_t database = 0;
    index.add(1, &database);
    index.hnsw.efSearch = 1;
    SearchStatsEnabledGuard guard(true);
    faiss::hnsw_stats.reset();

    std::vector<std::thread> threads;
    threads.reserve(num_threads);
    for (size_t i = 0; i < num_threads; ++i) {
        threads.emplace_back([&] {
            omp_set_num_threads(1);
            int32_t distance;
            faiss::idx_t label;
            for (size_t j = 0; j < searches_per_thread; ++j) {
                index.search(1, &database, 1, &distance, &label);
            }
        });
    }
    for (auto& thread : threads) {
        thread.join();
    }

    EXPECT_EQ(num_threads * searches_per_thread, faiss::hnsw_stats.n1);
}

int reference_pop_min(faiss::MinimaxHeap& heap, float* vmin_out) {
    assert(heap.k > 0);
    // returns min. This is an O(n) operation
    int i = heap.k - 1;
    while (i >= 0) {
        if (heap.ids[i] != -1) {
            break;
        }
        i--;
    }
    if (i == -1) {
        return -1;
    }
    int imin = i;
    float vmin = heap.dis[i];
    i--;
    while (i >= 0) {
        if (heap.ids[i] != -1 && heap.dis[i] < vmin) {
            vmin = heap.dis[i];
            imin = i;
        }
        i--;
    }
    if (vmin_out) {
        *vmin_out = vmin;
    }
    int ret = heap.ids[imin];
    heap.ids[imin] = -1;
    --heap.nvalid;

    return ret;
}

void test_popmin(int heap_size, int amount_to_put) {
    // create a heap
    faiss::MinimaxHeap mm_heap(heap_size);

    using storage_idx_t = faiss::HNSW::storage_idx_t;

    std::default_random_engine rng(123 + heap_size * amount_to_put);
    std::uniform_int_distribution<storage_idx_t> u(0, 65536);
    std::uniform_real_distribution<float> uf(0, 1);

    // generate random unique indices
    std::unordered_set<storage_idx_t> indices;
    while (static_cast<int>(indices.size()) < amount_to_put) {
        const storage_idx_t index = u(rng);
        indices.insert(index);
    }

    // put ones into the heap
    for (const auto index : indices) {
        float distance = uf(rng);
        if (distance >= 0.7f) {
            // add infinity values from time to time
            distance = std::numeric_limits<float>::infinity();
        }
        mm_heap.push(index, distance);
    }

    // clone the heap
    faiss::MinimaxHeap cloned_mm_heap = mm_heap;

    // takes ones out one by one
    while (mm_heap.size() > 0) {
        // compare heaps
        ASSERT_EQ(mm_heap.n, cloned_mm_heap.n);
        ASSERT_EQ(mm_heap.k, cloned_mm_heap.k);
        ASSERT_EQ(mm_heap.nvalid, cloned_mm_heap.nvalid);
        ASSERT_EQ(mm_heap.ids, cloned_mm_heap.ids);
        ASSERT_EQ(mm_heap.dis, cloned_mm_heap.dis);

        // use the reference pop_min for the cloned heap
        float cloned_vmin_dis = std::numeric_limits<float>::quiet_NaN();
        storage_idx_t cloned_vmin_idx =
                reference_pop_min(cloned_mm_heap, &cloned_vmin_dis);

        float vmin_dis = std::numeric_limits<float>::quiet_NaN();
        storage_idx_t vmin_idx = mm_heap.pop_min(&vmin_dis);

        // compare returns
        ASSERT_EQ(vmin_dis, cloned_vmin_dis);
        ASSERT_EQ(vmin_idx, cloned_vmin_idx);
    }

    // compare heaps again
    ASSERT_EQ(mm_heap.n, cloned_mm_heap.n);
    ASSERT_EQ(mm_heap.k, cloned_mm_heap.k);
    ASSERT_EQ(mm_heap.nvalid, cloned_mm_heap.nvalid);
    ASSERT_EQ(mm_heap.ids, cloned_mm_heap.ids);
    ASSERT_EQ(mm_heap.dis, cloned_mm_heap.dis);
}

void test_popmin_identical_distances(
        int heap_size,
        int amount_to_put,
        const float distance) {
    // create a heap
    faiss::MinimaxHeap mm_heap(heap_size);

    using storage_idx_t = faiss::HNSW::storage_idx_t;

    std::default_random_engine rng(123 + heap_size * amount_to_put);
    std::uniform_int_distribution<storage_idx_t> u(0, 65536);

    // generate random unique indices
    std::unordered_set<storage_idx_t> indices;
    while (static_cast<int>(indices.size()) < amount_to_put) {
        const storage_idx_t index = u(rng);
        indices.insert(index);
    }

    // put ones into the heap
    for (const auto index : indices) {
        mm_heap.push(index, distance);
    }

    // clone the heap
    faiss::MinimaxHeap cloned_mm_heap = mm_heap;

    // takes ones out one by one
    while (mm_heap.size() > 0) {
        // compare heaps
        ASSERT_EQ(mm_heap.n, cloned_mm_heap.n);
        ASSERT_EQ(mm_heap.k, cloned_mm_heap.k);
        ASSERT_EQ(mm_heap.nvalid, cloned_mm_heap.nvalid);
        ASSERT_EQ(mm_heap.ids, cloned_mm_heap.ids);
        ASSERT_EQ(mm_heap.dis, cloned_mm_heap.dis);

        // use the reference pop_min for the cloned heap
        float cloned_vmin_dis = std::numeric_limits<float>::quiet_NaN();
        storage_idx_t cloned_vmin_idx =
                reference_pop_min(cloned_mm_heap, &cloned_vmin_dis);

        float vmin_dis = std::numeric_limits<float>::quiet_NaN();
        storage_idx_t vmin_idx = mm_heap.pop_min(&vmin_dis);

        // compare returns
        ASSERT_EQ(vmin_dis, cloned_vmin_dis);
        ASSERT_EQ(vmin_idx, cloned_vmin_idx);
    }

    // compare heaps again
    ASSERT_EQ(mm_heap.n, cloned_mm_heap.n);
    ASSERT_EQ(mm_heap.k, cloned_mm_heap.k);
    ASSERT_EQ(mm_heap.nvalid, cloned_mm_heap.nvalid);
    ASSERT_EQ(mm_heap.ids, cloned_mm_heap.ids);
    ASSERT_EQ(mm_heap.dis, cloned_mm_heap.dis);
}

void copy_base_level_only(
        const faiss::IndexHNSWCagra& src,
        faiss::IndexHNSWCagra& dst) {
    auto n = src.ntotal;
    auto d = src.d;
    auto M = src.hnsw.nb_neighbors(0) / 2;
    auto graph_degree = src.hnsw.nb_neighbors(0);

    if (dst.storage && dst.own_fields) {
        delete dst.storage;
    }
    dst.storage = new faiss::IndexFlatL2(d);
    dst.own_fields = true;
    dst.d = d;
    dst.metric_type = src.metric_type;
    dst.is_trained = true;
    dst.keep_max_size_level0 = true;

    dst.hnsw.reset();
    dst.hnsw.assign_probas.clear();
    dst.hnsw.cum_nneighbor_per_level.clear();
    dst.hnsw.set_default_probas(M, 1.0 / std::log(M));

    dst.hnsw.prepare_level_tab(n, false);

    auto src_flat = dynamic_cast<faiss::IndexFlat*>(src.storage);
    FAISS_THROW_IF_NOT(src_flat);
    dst.storage->add(n, src_flat->get_xb());
    dst.ntotal = n;

    for (faiss::idx_t i = 0; i < n; i++) {
        size_t src_begin, src_end;
        src.hnsw.neighbor_range(i, 0, &src_begin, &src_end);

        size_t dst_begin, dst_end;
        dst.hnsw.neighbor_range(i, 0, &dst_begin, &dst_end);

        for (size_t j = 0; j < graph_degree && j < (dst_end - dst_begin); j++) {
            dst.hnsw.neighbors[dst_begin + j] =
                    src.hnsw.neighbors[src_begin + j];
        }
    }

    dst.base_level_only = true;
}

TEST(HNSW, Test_popmin) {
    std::vector<size_t> sizes = {1, 2, 3, 4, 5, 7, 9, 11, 16, 27, 32, 64, 128};
    for (const size_t size : sizes) {
        for (size_t amount = size; amount > 0; amount /= 2) {
            test_popmin(size, amount);
        }
    }
}

TEST(HNSW, Test_popmin_identical_distances) {
    std::vector<size_t> sizes = {1, 2, 3, 4, 5, 7, 9, 11, 16, 27, 32};
    for (const size_t size : sizes) {
        for (size_t amount = size; amount > 0; amount /= 2) {
            test_popmin_identical_distances(size, amount, 1.0f);
        }
    }
}

TEST(HNSW, Test_popmin_infinite_distances) {
    std::vector<size_t> sizes = {1, 2, 3, 4, 5, 7, 9, 11, 16, 27, 32};
    for (const size_t size : sizes) {
        for (size_t amount = size; amount > 0; amount /= 2) {
            test_popmin_identical_distances(
                    size, amount, std::numeric_limits<float>::infinity());
        }
    }
}

TEST(HNSW, Test_IndexHNSW_METRIC_Lp) {
    // Create an HNSW index with METRIC_Lp and metric_arg = 3
    faiss::IndexFlat storage_index(1, faiss::METRIC_Lp);
    storage_index.metric_arg = 3;
    faiss::IndexHNSW index(&storage_index, 32);

    // Add a single data point
    float data[1] = {0.0};
    index.add(1, data);

    // Prepare a query
    float query[1] = {2.0};
    float distance;
    faiss::idx_t label;

    index.search(1, query, 1, &distance, &label);

    EXPECT_NEAR(distance, 8.0, 1e-5); // Distance should be 8.0 (2^3)
    EXPECT_EQ(label, 0);              // Label should be 0
}

TEST(HNSW, Test_IndexHNSWCagra_BaseLevelOnly_RangeSearch) {
    int d = 8;
    int nb = 100;
    int nq = 5;
    int M = 4;

    std::vector<float> xb(nb * d);
    std::vector<float> xq(nq * d);
    faiss::float_rand(xb.data(), xb.size(), 1234);
    faiss::float_rand(xq.data(), xq.size(), 4321);

    faiss::IndexHNSWCagra index(d, M, faiss::METRIC_L2);
    index.add(nb, xb.data());
    index.base_level_only = true;
    index.num_base_level_search_entrypoints = 8;

    faiss::IndexHNSWCagra dst_index;
    copy_base_level_only(index, dst_index);
    dst_index.num_base_level_search_entrypoints = 8;

    faiss::RangeSearchResult res(nq);
    float radius = 1e9f;
    dst_index.range_search(nq, xq.data(), radius, &res);

    for (int i = 0; i < nq; i++) {
        auto count = res.lims[i + 1] - res.lims[i];
        EXPECT_GT(count, 0);
    }
}

class HNSWTest : public testing::Test {
   protected:
    HNSWTest() {
        xb = std::make_unique<std::vector<float>>(d * nb);
        xb->reserve(d * nb);
        faiss::float_rand(xb->data(), d * nb, 12345);
        index = std::make_unique<faiss::IndexHNSWFlat>(d, M);
        index->add(nb, xb->data());
        xq = std::unique_ptr<std::vector<float>>(
                new std::vector<float>(d * nq));
        xq->reserve(d * nq);
        faiss::float_rand(xq->data(), d * nq, 12345);
        dis = std::unique_ptr<faiss::DistanceComputer>(
                index->storage->get_distance_computer());
        dis->set_query(xq->data() + 0 * index->d);
    }

    const int d = 64;
    const int nb = 2000;
    const int M = 4;
    const int nq = 10;
    const int k = 10;
    std::unique_ptr<std::vector<float>> xb;
    std::unique_ptr<std::vector<float>> xq;
    std::unique_ptr<faiss::DistanceComputer> dis;
    std::unique_ptr<faiss::IndexHNSWFlat> index;
};

/** Do a BFS on the candidates list */
int reference_search_from_candidates(
        const faiss::HNSW& hnsw,
        faiss::DistanceComputer& qdis,
        faiss::ResultHandler& res,
        faiss::MinimaxHeap& candidates,
        faiss::VisitedTable& vt,
        faiss::HNSWStats& stats,
        int level,
        int nres_in,
        const faiss::SearchParametersHNSW* params) {
    int nres = nres_in;
    int ndis = 0;

    // can be overridden by search params
    bool do_dis_check = params ? params->check_relative_distance
                               : hnsw.check_relative_distance;
    int efSearch = params ? params->efSearch : hnsw.efSearch;
    const faiss::IDSelector* sel = params ? params->sel : nullptr;

    faiss::HNSW::C::T threshold = res.threshold;
    for (int i = 0; i < candidates.size(); i++) {
        faiss::idx_t v1 = candidates.ids[i];
        float d = candidates.dis[i];
        FAISS_ASSERT(v1 >= 0);
        if (!sel || sel->is_member(v1)) {
            if (d < threshold) {
                if (res.add_result(d, v1)) {
                    threshold = res.threshold;
                }
            }
        }
        vt.set(v1);
    }

    int nstep = 0;

    while (candidates.size() > 0) {
        float d0 = 0;
        int v0 = candidates.pop_min(&d0);

        if (do_dis_check) {
            // tricky stopping condition: there are more that ef
            // distances that are processed already that are smaller
            // than d0

            int n_dis_below = candidates.count_below(d0);
            if (n_dis_below >= efSearch) {
                break;
            }
        }

        size_t begin, end;
        hnsw.neighbor_range(v0, level, &begin, &end);

        // a reference version
        for (size_t j = begin; j < end; j++) {
            int v1 = hnsw.neighbors[j];
            if (v1 < 0) {
                break;
            }
            if (vt.get(v1)) {
                continue;
            }
            vt.set(v1);
            ndis++;
            float d = qdis(v1);
            if (!sel || sel->is_member(v1)) {
                if (d < threshold) {
                    if (res.add_result(d, v1)) {
                        threshold = res.threshold;
                        nres += 1;
                    }
                }
            }

            candidates.push(v1, d);
        }

        nstep++;
        if (!do_dis_check && nstep > efSearch) {
            break;
        }
    }

    if (level == 0) {
        stats.n1++;
        if (candidates.size() == 0) {
            stats.n2++;
        }
        stats.ndis += ndis;
        stats.nhops += nstep;
    }

    return nres;
}

faiss::HNSWStats reference_greedy_update_nearest(
        const faiss::HNSW& hnsw,
        faiss::DistanceComputer& qdis,
        int level,
        faiss::HNSW::storage_idx_t& nearest,
        float& d_nearest) {
    faiss::HNSWStats stats;

    for (;;) {
        faiss::HNSW::storage_idx_t prev_nearest = nearest;

        size_t begin, end;
        hnsw.neighbor_range(nearest, level, &begin, &end);

        size_t ndis = 0;

        for (size_t i = begin; i < end; i++) {
            faiss::HNSW::storage_idx_t v = hnsw.neighbors[i];
            if (v < 0) {
                break;
            }
            ndis += 1;
            float dis = qdis(v);
            if (dis < d_nearest) {
                nearest = v;
                d_nearest = dis;
            }
        }
        // update stats
        stats.ndis += ndis;
        stats.nhops += 1;

        if (nearest == prev_nearest) {
            return stats;
        }
    }
}

std::priority_queue<faiss::HNSW::Node> reference_search_from_candidate_unbounded(
        const faiss::HNSW& hnsw,
        const faiss::HNSW::Node& node,
        faiss::DistanceComputer& qdis,
        int ef,
        faiss::VisitedTable* vt,
        faiss::HNSWStats& stats) {
    int ndis = 0;
    std::priority_queue<faiss::HNSW::Node> top_candidates;
    std::priority_queue<
            faiss::HNSW::Node,
            std::vector<faiss::HNSW::Node>,
            std::greater<faiss::HNSW::Node>>
            candidates;

    top_candidates.push(node);
    candidates.push(node);

    vt->set(node.second);

    while (!candidates.empty()) {
        float d0;
        faiss::HNSW::storage_idx_t v0;
        std::tie(d0, v0) = candidates.top();

        if (d0 > top_candidates.top().first) {
            break;
        }

        candidates.pop();

        size_t begin, end;
        hnsw.neighbor_range(v0, 0, &begin, &end);

        for (size_t j = begin; j < end; ++j) {
            int v1 = hnsw.neighbors[j];

            if (v1 < 0) {
                break;
            }
            if (vt->get(v1)) {
                continue;
            }

            vt->set(v1);

            float d1 = qdis(v1);
            ++ndis;

            if (top_candidates.top().first > d1 ||
                static_cast<int>(top_candidates.size()) < ef) {
                candidates.emplace(d1, v1);
                top_candidates.emplace(d1, v1);

                if (static_cast<int>(top_candidates.size()) > ef) {
                    top_candidates.pop();
                }
            }
        }

        stats.nhops += 1;
    }

    ++stats.n1;
    if (candidates.size() == 0) {
        ++stats.n2;
    }
    stats.ndis += ndis;

    return top_candidates;
}

TEST_F(HNSWTest, TEST_search_from_candidate_unbounded) {
    omp_set_num_threads(1);
    auto nearest = index->hnsw.entry_point;
    float d_nearest = (*dis)(nearest);
    auto node = faiss::HNSW::Node(d_nearest, nearest);
    std::unique_ptr<faiss::VisitedTable> vt =
            faiss::VisitedTable::create(index->ntotal);
    faiss::HNSWStats stats;

    // actual version
    auto top_candidates = faiss::hnsw_detail::search_from_candidate_unbounded(
            index->hnsw, node, *dis, k, vt.get(), stats);

    auto reference_nearest = index->hnsw.entry_point;
    float reference_d_nearest = (*dis)(nearest);
    auto reference_node =
            faiss::HNSW::Node(reference_d_nearest, reference_nearest);
    std::unique_ptr<faiss::VisitedTable> reference_vt =
            faiss::VisitedTable::create(index->ntotal);
    faiss::HNSWStats reference_stats;

    // reference version
    auto reference_top_candidates = reference_search_from_candidate_unbounded(
            index->hnsw,
            reference_node,
            *dis,
            k,
            reference_vt.get(),
            reference_stats);
    EXPECT_EQ(stats.ndis, reference_stats.ndis);
    EXPECT_EQ(stats.nhops, reference_stats.nhops);
    EXPECT_EQ(stats.n1, reference_stats.n1);
    EXPECT_EQ(stats.n2, reference_stats.n2);
    EXPECT_EQ(top_candidates.size(), reference_top_candidates.size());
}

TEST_F(HNSWTest, TEST_greedy_update_nearest) {
    omp_set_num_threads(1);

    auto nearest = index->hnsw.entry_point;
    float d_nearest = (*dis)(nearest);
    auto reference_nearest = index->hnsw.entry_point;
    float reference_d_nearest = (*dis)(reference_nearest);

    // actual version
    auto stats = faiss::hnsw_detail::greedy_update_nearest(
            index->hnsw, *dis, 0, nearest, d_nearest);

    // reference version
    auto reference_stats = reference_greedy_update_nearest(
            index->hnsw, *dis, 0, reference_nearest, reference_d_nearest);
    EXPECT_EQ(stats.ndis, reference_stats.ndis);
    EXPECT_EQ(stats.nhops, reference_stats.nhops);
    EXPECT_EQ(stats.n1, reference_stats.n1);
    EXPECT_EQ(stats.n2, reference_stats.n2);
    EXPECT_NEAR(d_nearest, reference_d_nearest, 0.01);
    EXPECT_EQ(nearest, reference_nearest);
}

TEST_F(HNSWTest, TEST_search_from_candidates) {
    omp_set_num_threads(1);

    std::vector<faiss::idx_t> I(k * nq);
    std::vector<float> D(k * nq);
    std::vector<faiss::idx_t> reference_I(k * nq);
    std::vector<float> reference_D(k * nq);
    using RH = faiss::HeapBlockResultHandler<faiss::HNSW::C>;

    std::unique_ptr<faiss::VisitedTable> vt =
            faiss::VisitedTable::create(index->ntotal);
    std::unique_ptr<faiss::VisitedTable> reference_vt =
            faiss::VisitedTable::create(index->ntotal);
    int num_candidates = 10;
    faiss::MinimaxHeap candidates(num_candidates);
    faiss::MinimaxHeap reference_candidates(num_candidates);

    for (int i = 0; i < num_candidates; i++) {
        vt->set(i);
        reference_vt->set(i);
        candidates.push(i, (*dis)(i));
        reference_candidates.push(i, (*dis)(i));
    }

    faiss::HNSWStats stats;
    RH bres(nq, D.data(), I.data(), k);
    faiss::HeapBlockResultHandler<faiss::HNSW::C>::SingleResultHandler res(
            bres);

    res.begin(0);
    faiss::hnsw_detail::search_from_candidates(
            index->hnsw, *dis, res, candidates, *vt, stats, 0, 0, nullptr);
    res.end();

    faiss::HNSWStats reference_stats;
    RH reference_bres(nq, reference_D.data(), reference_I.data(), k);
    faiss::HeapBlockResultHandler<faiss::HNSW::C>::SingleResultHandler
            reference_res(reference_bres);
    reference_res.begin(0);
    reference_search_from_candidates(
            index->hnsw,
            *dis,
            reference_res,
            reference_candidates,
            *reference_vt,
            reference_stats,
            0,
            0,
            nullptr);
    reference_res.end();
    for (int i = 0; i < nq; i++) {
        for (int j = 0; j < k; j++) {
            EXPECT_NEAR(I[i * k + j], reference_I[i * k + j], 0.1);
            EXPECT_NEAR(D[i * k + j], reference_D[i * k + j], 0.1);
        }
    }
    EXPECT_EQ(reference_stats.ndis, stats.ndis);
    EXPECT_EQ(reference_stats.nhops, stats.nhops);
    EXPECT_EQ(reference_stats.n1, stats.n1);
    EXPECT_EQ(reference_stats.n2, stats.n2);
}

TEST_F(HNSWTest, TEST_search_neighbors_to_add) {
    omp_set_num_threads(1);

    std::unique_ptr<faiss::VisitedTable> vt =
            faiss::VisitedTable::create(index->ntotal);
    std::unique_ptr<faiss::VisitedTable> reference_vt =
            faiss::VisitedTable::create(index->ntotal);

    std::priority_queue<faiss::HNSW::NodeDistCloser> link_targets;
    std::priority_queue<faiss::HNSW::NodeDistCloser> reference_link_targets;

    faiss::hnsw_detail::search_neighbors_to_add(
            index->hnsw,
            *dis,
            link_targets,
            index->hnsw.entry_point,
            (*dis)(index->hnsw.entry_point),
            index->hnsw.max_level,
            *vt,
            false);

    faiss::hnsw_detail::search_neighbors_to_add(
            index->hnsw,
            *dis,
            reference_link_targets,
            index->hnsw.entry_point,
            (*dis)(index->hnsw.entry_point),
            index->hnsw.max_level,
            *reference_vt,
            true);

    EXPECT_EQ(link_targets.size(), reference_link_targets.size());
    while (!link_targets.empty()) {
        auto val = link_targets.top();
        auto reference_val = reference_link_targets.top();
        EXPECT_EQ(val.d, reference_val.d);
        EXPECT_EQ(val.id, reference_val.id);
        link_targets.pop();
        reference_link_targets.pop();
    }
}

TEST_F(HNSWTest, TEST_nb_neighbors_bound) {
    omp_set_num_threads(1);
    EXPECT_EQ(index->hnsw.nb_neighbors(0), 8);
    EXPECT_EQ(index->hnsw.nb_neighbors(1), 4);
    EXPECT_EQ(index->hnsw.nb_neighbors(2), 4);
    EXPECT_EQ(index->hnsw.nb_neighbors(3), 4);
    // picking a large number to trigger an exception based on checking bounds
    EXPECT_THROW(index->hnsw.nb_neighbors(100), faiss::FaissException);
}

TEST_F(HNSWTest, TEST_search_level_0) {
    omp_set_num_threads(1);
    std::vector<faiss::idx_t> I(k * nq);
    std::vector<float> D(k * nq);

    using RH = faiss::HeapBlockResultHandler<faiss::HNSW::C>;
    RH bres1(nq, D.data(), I.data(), k);
    faiss::HeapBlockResultHandler<faiss::HNSW::C>::SingleResultHandler res1(
            bres1);
    RH bres2(nq, D.data(), I.data(), k);
    faiss::HeapBlockResultHandler<faiss::HNSW::C>::SingleResultHandler res2(
            bres2);

    faiss::HNSWStats stats1, stats2;
    std::unique_ptr<faiss::VisitedTable> vt1 =
            faiss::VisitedTable::create(index->ntotal);
    std::unique_ptr<faiss::VisitedTable> vt2 =
            faiss::VisitedTable::create(index->ntotal);
    auto nprobe = 5;
    const faiss::HNSW::storage_idx_t values[] = {1, 2, 3, 4, 5};
    const faiss::HNSW::storage_idx_t* nearest_i = values;
    const float distances[] = {0.1, 0.2, 0.3, 0.4, 0.5};
    const float* nearest_d = distances;

    // search_type == 1
    res1.begin(0);
    index->hnsw.search_level_0(
            *dis, res1, nprobe, nearest_i, nearest_d, 1, stats1, *vt1, nullptr);
    res1.end();

    // search_type == 2
    res2.begin(0);
    index->hnsw.search_level_0(
            *dis, res2, nprobe, nearest_i, nearest_d, 2, stats2, *vt2, nullptr);
    res2.end();

    // search_type 1 calls search_from_candidates in a loop nprobe times.
    // search_type 2 pushes the candidates and just calls search_from_candidates
    // once, so those stats will be much less.
    EXPECT_GT(stats1.ndis, stats2.ndis);
    EXPECT_GT(stats1.nhops, stats2.nhops);
    EXPECT_GT(stats1.n1, stats2.n1);
    EXPECT_GT(stats1.n2, stats2.n2);
}

TEST(VisitedTableReuse, ReusesGrowsAndResets) {
    // get_reusable() hands back the same thread-local object across calls
    // (no per-call allocation) ...
    faiss::VisitedTable& t1 = faiss::VisitedTable::get_reusable(100);
    faiss::VisitedTable& t2 = faiss::VisitedTable::get_reusable(100);
    EXPECT_EQ(&t1, &t2);

    // ... and hands it back clean: a mark set before re-acquiring must not
    // survive, since a reused table may carry stamps from a prior search.
    t2.set(42);
    EXPECT_TRUE(t2.get(42));
    faiss::VisitedTable& t3 = faiss::VisitedTable::get_reusable(100);
    EXPECT_EQ(&t2, &t3);
    EXPECT_FALSE(t3.get(42));

    // Growing to a larger size keeps the same object and exposes the new range.
    faiss::VisitedTable& t4 = faiss::VisitedTable::get_reusable(1000);
    EXPECT_EQ(&t3, &t4);
    EXPECT_FALSE(t4.get(900));
    EXPECT_TRUE(t4.set(900));
    EXPECT_TRUE(t4.get(900));
}

// Restores the process-global OpenMP thread count on scope exit, so a test that
// changes it via omp_set_num_threads does not leak the setting into later
// tests.
struct ScopedOmpThreads {
    int saved = omp_get_max_threads();
    ~ScopedOmpThreads() {
        omp_set_num_threads(saved);
    }
};

TEST_F(HNSWTest, TEST_search_advances_one_generation_per_query) {
    ScopedOmpThreads omp_guard;
    omp_set_num_threads(1);
    index->hnsw.use_visited_hashset = false;

    auto& reusable = dynamic_cast<faiss::VisitedTableVector&>(
            faiss::VisitedTable::get_reusable(index->ntotal, false));
    std::fill(reusable.visited.begin(), reusable.visited.end(), 0);
    reusable.visno = 10;

    std::vector<faiss::idx_t> labels(k);
    std::vector<float> distances(k);
    index->search(1, xq->data(), k, distances.data(), labels.data());

    // Acquiring the reusable table advances 10 -> 11 to clear stale state.
    // HNSW::search then advances 11 -> 12 after the query. The outer wrapper
    // must not perform a third advance.
    EXPECT_EQ(12, reusable.visno);
}

TEST_F(HNSWTest, TEST_search_reuse_correctness) {
    ScopedOmpThreads omp_guard;

    std::vector<faiss::idx_t> I1(k * nq), I2(k * nq);
    std::vector<float> D1(k * nq), D2(k * nq);

    // Two back-to-back searches must be identical: the reused visited table
    // must not carry state across search() calls.
    omp_set_num_threads(1);
    index->search(nq, xq->data(), k, D1.data(), I1.data());
    index->search(nq, xq->data(), k, D2.data(), I2.data());
    EXPECT_EQ(I1, I2);
    EXPECT_EQ(D1, D2);

    // The reused versioned-array path must match a hash-set reference.
    std::vector<faiss::idx_t> Iref(k * nq);
    std::vector<float> Dref(k * nq);
    index->hnsw.use_visited_hashset = true;
    index->search(nq, xq->data(), k, Dref.data(), Iref.data());
    index->hnsw.use_visited_hashset = false;
    std::vector<faiss::idx_t> Iarr(k * nq);
    std::vector<float> Darr(k * nq);
    index->search(nq, xq->data(), k, Darr.data(), Iarr.data());
    EXPECT_EQ(Iref, Iarr);
    EXPECT_EQ(Dref, Darr);

    // Multi-threaded search must match single-threaded: each thread reuses its
    // own thread-local table.
    index->hnsw.use_visited_hashset = std::nullopt;
    omp_set_num_threads(4);
    std::vector<faiss::idx_t> Imt(k * nq);
    std::vector<float> Dmt(k * nq);
    index->search(nq, xq->data(), k, Dmt.data(), Imt.data());
    EXPECT_EQ(I1, Imt);
    EXPECT_EQ(D1, Dmt);
}