fix(hnsw): clamp ef to the stored point count and validate cluster lookups

docs/development.md

@@ -533,6 +533,10 @@ Use the [path conversion helper](#filesystem-paths) for file access. Preserve
 Frozen release-produced indexes and reference searches live in `tests/fixtures/persistence/`;
 keep their provenance and hashes, and do not regenerate them with the current writer.
 
+Index files are trusted input (see the usage guide): validate structure and sizes so
+that corrupt or mismatched files are reported before allocation, without treating a
+crafted file as a supported input.
+
 Do not silently reinterpret an old file. Add a magic/version discriminator, use fixed-width
 serialized fields for new formats, validate sizes before allocation, check every read, and retain a
 compatibility test fixture or an explicit rejection path. Review each affected format;

docs/docs/index/hnsw.md

@@ -62,7 +62,7 @@ see [C++ integration](../cpp_integration.md).
 | `nbits` / C++ `total_bits` | 1–9 total bits including the sign bit; Python defaults to 8. Raw-vector storage is not supported. |
 | `M`, `ef_construction` | Graph connectivity and construction search budget; Python defaults are 16 and 200. |
 | `metric`, `random_seed` | Python defaults are `"l2"` and 100; `"ip"` is also supported. |
-| `k`, `ef` | Positive result count and search window. Python requires `k <= max_elements`; omitted/zero `ef` uses `max(k, 10)`, and positive `ef` is raised to at least `k`. |
+| `k`, `ef` | Positive result count and search window. Python requires `k <= max_elements`; omitted/zero `ef` uses `max(k, 10)`, positive `ef` is raised to at least `k`, and `ef` above the stored point count is clamped to it. |
 
 `random_seed` controls graph levels and the FHT/Kac rotation (using its low 32 bits).
 With identical inputs and parameters, single-threaded construction is reproducible

docs/docs/index/usage.md

@@ -164,3 +164,9 @@ or unsupported versions. See the
 Compatibility is format-specific; preserving old readers' ability to open new
 files is not guaranteed. Consult each index's format documentation before
 downgrading the library.
+
+Index files are trusted input. Loaders validate structure and reject truncated,
+inconsistent, or unsupported files so that mistakes are reported, but they are
+not a security boundary: a crafted file can make `load` allocate its stored
+capacity or violate invariants that later `add` and `remove` calls rely on.
+Load only files written by this library.

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -73,8 +73,19 @@ class HierarchicalNSW {
         return num_stored_points_.load() - num_removed_;
     }
     [[nodiscard]] size_t num_stored_points() const { return num_stored_points_; }
+    // Cluster of a stored record by internal ID (0 through num_stored_points() - 1;
+    // these coincide with labels only for points added in order). Removed points
+    // keep their cluster.
     [[nodiscard]] PID cluster_id_of(PID internal_id) const {
-        return raw_cluster_id(internal_id) & ~kRemovedMask;
+        if (data_level0_memory_ == nullptr) {
+            throw std::logic_error(
+                "HNSW index must be constructed or loaded before cluster lookup"
+            );
+        }
+        if (internal_id >= num_stored_points_.load()) {
+            throw std::invalid_argument("HNSW internal ID is out of range");
+        }
+        return stored_cluster_id(internal_id);
     }
 
     void save(const char* filename) const;
@@ -157,6 +168,7 @@ class HierarchicalNSW {
 
     // Queries use the original dimension. Each result is (distance, external label),
     // where distance estimates squared L2 or 1 - dot according to the index metric.
+    // ef_search is raised to k and clamped to the stored point count.
     std::vector<std::vector<std::pair<float, PID>>> search(
         const float* queries,
         size_t num_queries,
@@ -288,6 +300,9 @@ class HierarchicalNSW {
     }
 
     // The stored cluster ID with the removal mark, as save and load see it.
+    [[nodiscard]] PID stored_cluster_id(PID internal_id) const {
+        return raw_cluster_id(internal_id) & ~kRemovedMask;
+    }
     [[nodiscard]] PID raw_cluster_id(PID internal_id) const {
         return *(reinterpret_cast<const PID*>(
             data_level0_memory_.get() + (internal_id * level0_record_bytes_) +

python_bindings/__init__.pyi

@@ -571,7 +571,8 @@ class HnswIndex:
         k : int
             Number of neighbors requested, between 1 and max_elements.
         ef : int, optional
-            Search breadth. 0 (default) selects max(k, 10).
+            Search breadth. 0 (default) selects max(k, 10); values above the stored
+            point count are clamped.
         num_threads : int, optional
             Worker count. Default is 1; 0 selects the available CPU count.
 

python_bindings/hnsw_bindings.cpp

@@ -376,7 +376,8 @@ queries : array_like of float, shape (n_queries, dim)
 k : int
     Number of neighbors requested, between 1 and max_elements.
 ef : int, optional
-    Search breadth. 0 (default) selects max(k, 10).
+    Search breadth. 0 (default) selects max(k, 10); values above the stored
+    point count are clamped.
 num_threads : int, optional
     Worker count. Default is 1; 0 selects the available CPU count.
 

src/index/hnsw.cpp

@@ -954,6 +954,9 @@ std::vector<std::pair<float, PID>> HierarchicalNSW::search_knn_sorted_with(
     if (ef_search == 0) {
         ef_search = kDefaultEfSearch;
     }
+    // No candidate pool can hold more distinct points than are stored, so a huge
+    // ef_search must not turn into a huge allocation.
+    ef_search = std::min(ef_search, num_stored_points_.load());
 
     SplitSingleQuery<float> query_wrapper(
         rotated_query, padded_dim_, ex_bits_, query_config_, metric_type_
@@ -989,7 +992,9 @@ std::vector<std::pair<float, PID>> HierarchicalNSW::search_knn_sorted_with(
 
     auto visited = visited_set_pool_->acquire();
     // Keep the ef_search best estimates in a bounded sorted buffer.
-    buffer::SearchBuffer<float> search_pool(std::max(ef_search, k));
+    buffer::SearchBuffer<float> search_pool(
+        std::min(std::max(ef_search, k), num_stored_points_.load())
+    );
     std::vector<AnnCandidate<float>> nearest(k);
     const size_t num_results = search_layers(
         *this,
@@ -1068,7 +1073,7 @@ void HierarchicalNSW::reconstruct_rotated(PID internal_id, float* rotated) const
     quant::rabitq_impl::unpack_sign_bits(bin.bin_code(), padded_dim_, combined.data());
 
     const auto* centroid = reinterpret_cast<const float*>(centroids_memory_.get()) +
-                           (cluster_id_of(internal_id) * padded_dim_);
+                           (stored_cluster_id(internal_id) * padded_dim_);
 
     if (ex_bits_ == 0) {
         quant::reconstruct_full_vec<float, uint16_t>(
@@ -1118,7 +1123,7 @@ float HierarchicalNSW::quantized_distance(
     }
 
     const auto* centroid = reinterpret_cast<const float*>(centroids_memory_.get()) +
-                           (cluster_id_of(target) * padded_dim_);
+                           (stored_cluster_id(target) * padded_dim_);
     const float dist_sqr = euclidean_sqr(cached.rotated.data(), centroid, padded_dim_);
     const float g_error = std::sqrt(dist_sqr);
     const float g_add = (metric_type_ == METRIC_IP)

src/index/hnsw_search_impl.hpp

@@ -56,9 +56,9 @@ inline void HierarchicalNSW::estimate_bin_direct(
     DistanceEstimate& estimate
 ) const {
     if (metric_type_ == METRIC_IP) {
-        const float query_centroid_ip = q_to_centroids[cluster_id_of(candidate_id)];
+        const float query_centroid_ip = q_to_centroids[stored_cluster_id(candidate_id)];
         const float residual_norm =
-            q_to_centroids[cluster_id_of(candidate_id) + num_clusters_];
+            q_to_centroids[stored_cluster_id(candidate_id) + num_clusters_];
         estimate = estimate_split_single_direct<Kernel>(
             bin_data(candidate_id),
             query_wrapper,
@@ -68,7 +68,7 @@ inline void HierarchicalNSW::estimate_bin_direct(
         );
     } else {
         // L2 distance
-        const float residual_norm = q_to_centroids[cluster_id_of(candidate_id)];
+        const float residual_norm = q_to_centroids[stored_cluster_id(candidate_id)];
         estimate = estimate_split_single_direct<Kernel>(
             bin_data(candidate_id),
             query_wrapper,
@@ -87,9 +87,9 @@ inline void HierarchicalNSW::estimate_full_direct(
     DistanceEstimate& estimate
 ) const {
     if (metric_type_ == METRIC_IP) {
-        const float query_centroid_ip = q_to_centroids[cluster_id_of(candidate_id)];
+        const float query_centroid_ip = q_to_centroids[stored_cluster_id(candidate_id)];
         const float residual_norm =
-            q_to_centroids[cluster_id_of(candidate_id) + num_clusters_];
+            q_to_centroids[stored_cluster_id(candidate_id) + num_clusters_];
         estimate = estimate_split_single_full_direct<Kernel>(
             bin_data(candidate_id),
             ex_data(candidate_id),
@@ -102,7 +102,7 @@ inline void HierarchicalNSW::estimate_full_direct(
         );
     } else {
         // L2 distance
-        const float residual_norm = q_to_centroids[cluster_id_of(candidate_id)];
+        const float residual_norm = q_to_centroids[stored_cluster_id(candidate_id)];
         estimate = estimate_split_single_full_direct<Kernel>(
             bin_data(candidate_id),
             ex_data(candidate_id),

tests/python/test_hnsw.py

@@ -560,3 +560,13 @@ def test_legacy_padding_round_trip_and_add(tmp_path):
     np.testing.assert_array_equal(actual[1], distances)
     loaded.add(np.ones((1, 65), dtype=np.float32))
     assert loaded.num_points == 4
+
+
+def test_huge_ef_is_clamped_to_the_stored_point_count():
+    data = np.random.default_rng(5).normal(size=(80, 64)).astype(np.float32)
+    index = HnswIndex(64, 80, nbits=4)
+    index.build(data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32))
+    exhaustive_ids, exhaustive_distances = index.search(data[:3], k=5, ef=80)
+    ids, distances = index.search(data[:3], k=5, ef=2**40)
+    np.testing.assert_array_equal(ids, exhaustive_ids)
+    np.testing.assert_array_equal(distances, exhaustive_distances)

tests/unit/rabitqlib/index/hnsw_test.cpp

@@ -1430,6 +1430,47 @@ TEST(HnswInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) {
     EXPECT_EQ(index.search(edge_queries.data(), 2, 3, 10, 1).size(), 2U);
 }
 
+TEST(HnswSearchTest, ClampsEfToStoredPointsAndChecksClusterLookups) {
+    constexpr size_t kDim = 64;
+    constexpr size_t kCount = 24;
+    constexpr size_t kClusters = 3;
+    std::mt19937 gen(906);
+    std::normal_distribution<float> normal;
+    std::vector<float> data(kCount * kDim);
+    std::vector<float> centroids(kClusters * kDim);
+    for (auto* values : {&data, &centroids}) {
+        for (auto& value : *values) {
+            value = normal(gen);
+        }
+    }
+    std::vector<PID> cluster_ids(kCount);
+    for (size_t i = 0; i < kCount; ++i) {
+        cluster_ids[i] = static_cast<PID>(i % kClusters);
+    }
+
+    HierarchicalNSW unbuilt;
+    EXPECT_THROW(unbuilt.cluster_id_of(0), std::logic_error);
+
+    HierarchicalNSW index(kCount, kDim, 4, 8, 32);
+    index.construct(
+        kClusters, centroids.data(), kCount, data.data(), cluster_ids.data(), 1
+    );
+    for (size_t i = 0; i < kCount; ++i) {
+        EXPECT_EQ(index.cluster_id_of(static_cast<PID>(i)), cluster_ids[i]);
+    }
+    EXPECT_THROW(index.cluster_id_of(static_cast<PID>(kCount)), std::invalid_argument);
+    EXPECT_THROW(index.cluster_id_of(kPidMax), std::invalid_argument);
+
+    // A search window beyond the stored point count neither allocates for it nor
+    // changes the results of an exhaustive window.
+    const auto exhaustive = index.search(data.data(), 2, 5, kCount, 1);
+    const auto huge = index.search(data.data(), 2, 5, size_t{1} << 40, 1);
+    ASSERT_EQ(huge.size(), exhaustive.size());
+    for (size_t q = 0; q < huge.size(); ++q) {
+        EXPECT_EQ(huge[q], exhaustive[q]);
+    }
+}
+
 }  // namespace
 }  // namespace rabitqlib::hnsw