fix: validate advanced encoder preconditions and bound HNSW prefetch look-ahead

docs/docs/rabitq/quantizer.md

@@ -171,6 +171,12 @@ int main() {
 }
 ```
 
+The advanced encoders validate their preconditions and throw
+`std::invalid_argument` otherwise: `dim` must be at least 2 (the error factor
+divides by `dim - 1`), the code element type must hold `bits` (nine bits need
+`uint16_t`), inputs must be finite, and `SplitSingleQuery` needs a padded
+dimension that is a positive multiple of 32.
+
 #### Querying
 
 Score a query using the encoded vector and its factors.

include/rabitqlib/core/detail/rabitq_impl.hpp

@@ -7,6 +7,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <functional>
+#include <limits>
 #include <stdexcept>
 #include <type_traits>
 #include <vector>
@@ -50,6 +51,11 @@ inline CodeFactorTerms<T> code_factor_terms(
     MetricType metric_type,
     const char* invalid_metric_message
 ) {
+    // The error factor divides by dim - 1; the stable API and the indexes already
+    // enforce this, the advanced encoders reach here directly.
+    if (dim < 2) {
+        throw std::invalid_argument("RaBitQ factors require dim >= 2");
+    }
     // A zero residual needs no centered code or inner-product divisions.
     if (l2_sqr == 0) {
         if (metric_type == METRIC_L2) {
@@ -863,6 +869,13 @@ static inline void rabitq_full_impl_with_scratch(
     double t_const,
     Scratch<T>* scratch
 ) {
+    // One sign bit plus at most eight magnitude bits (xy_bits::kMaxCombinedBits).
+    if (total_bits < 1 || total_bits > 9 ||
+        total_bits > static_cast<size_t>(std::numeric_limits<TP>::digits)) {
+        throw std::invalid_argument(
+            "total_bits must be in [1, 9] and fit the code element type"
+        );
+    }
     const size_t ex_bits = total_bits - 1;
     if (ex_bits > 0) {
         RowMajorArray<T> owned_residual;

include/rabitqlib/core/query.hpp

@@ -66,6 +66,12 @@ class SplitSingleQuery {
                 "and base_bits + ex_bits <= 9"
             );
         }
+        // The query-bit transposition kernels load whole 32-coordinate blocks.
+        if (padded_dim == 0 || padded_dim % 32 != 0) {
+            throw std::invalid_argument(
+                "SplitSingleQuery requires padded_dim to be a positive multiple of 32"
+            );
+        }
 
         validate_metric_type(metric_type);
 

include/rabitqlib/utils/space.hpp

@@ -19,7 +19,7 @@
 
 namespace rabitqlib {
 
-inline uint64_t load_unaligned_u64(const void* data) {
+static inline uint64_t load_unaligned_u64(const void* data) {
     uint64_t value = 0;
     std::memcpy(&value, data, sizeof(value));
     return value;
@@ -302,11 +302,13 @@ void new_transpose_bin_512(
 float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim);
 float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim);
 
-inline float mask_ip_x0_q_old(const float* query, const uint64_t* data, size_t padded_dim) {
+static inline float mask_ip_x0_q_old(
+    const float* query, const uint64_t* data, size_t padded_dim
+) {
     return mask_ip_x0_q(query, data, padded_dim);
 }
 
-inline float ip_x0_q(
+static inline float ip_x0_q(
     const uint64_t* data,
     const uint64_t* query,
     float delta,
@@ -337,7 +339,7 @@ static inline uint32_t ip_bin_bin(const uint64_t* q, const uint64_t* d, size_t p
     return ret;
 }
 
-inline uint32_t ip_byte_bin(
+static inline uint32_t ip_byte_bin(
     const uint64_t* q, const uint64_t* d, size_t padded_dim, size_t b_query
 ) {
     uint32_t ret = 0;
@@ -349,7 +351,7 @@ inline uint32_t ip_byte_bin(
     return ret;
 }
 
-inline size_t popcount(const uint64_t* __restrict__ d, size_t length) {
+static inline size_t popcount(const uint64_t* __restrict__ d, size_t length) {
     size_t ret = 0;
     const size_t words = length / 64;
     const auto* data = reinterpret_cast<const uint8_t*>(d);

src/index/hnsw.cpp

@@ -620,8 +620,10 @@ maxheap<std::pair<float, PID>> HierarchicalNSW::search_layer(
         const size_t num_neighbors = link_count(links);
         const PID* neighbors = links + 1;
 
-        prefetch(*neighbors);
-        prefetch(*(neighbors + 1));
+        // Slots beyond the link count hold stale or unwritten IDs.
+        for (size_t j = 0; j < std::min<size_t>(2, num_neighbors); ++j) {
+            prefetch(neighbors[j]);
+        }
 
         for (size_t j = 0; j < num_neighbors; ++j) {
             PID candidate_id = neighbors[j];

src/index/hnsw_search_impl.hpp

@@ -256,7 +256,8 @@ inline void HierarchicalNSW::search_base_layer_direct(
         const PID* links = level0_links(current_node_id);
         const size_t num_neighbors = link_count(links);
 
-        for (size_t p = 0; p < kPrefetchLookahead; ++p) {
+        // Slots beyond the link count hold stale IDs; never form addresses from them.
+        for (size_t p = 0; p < std::min(kPrefetchLookahead, num_neighbors); ++p) {
             memory::mem_prefetch_l1(bin_data(*(links + 1 + p)), prefetch_size);
         }
         // Iterate over neighbors. (List starts at index 1.)
@@ -299,9 +300,13 @@ inline void HierarchicalNSW::search_base_layer_direct(
                 search_pool.insert(candidate_id, candidate_estimate.distance);
             }
 
-            memory::mem_prefetch_l2(
-                reinterpret_cast<const char*>(level0_links(search_pool.next_id())), 2
-            );
+            // A full pool parks the displaced entry, possibly with its checked bit
+            // set, in the spare slot that next_id() would read.
+            if (search_pool.has_next()) {
+                memory::mem_prefetch_l2(
+                    reinterpret_cast<const char*>(level0_links(search_pool.next_id())), 2
+                );
+            }
         }
     }
 }

tests/unit/rabitqlib/quantization/rabitq_test.cpp

@@ -6,6 +6,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <limits>
+#include <stdexcept>
 #include <string>
 #include <utility>
 #include <vector>
@@ -675,6 +676,83 @@ TEST(RabitqSignConventionTest, ZeroResidualDoesNotInflateReconstructionNorm) {
                             << std::to_string(tan_sq);
 }
 
+TEST(RabitqPreconditionTest, EncodersRejectSingleDimensionInput) {
+    const std::array<float, 1> data{0.5F};
+    const std::array<float, 1> centroid{0.0F};
+    std::array<uint16_t, 1> code{};
+    float f_add = 0;
+    float f_rescale = 0;
+    float f_error = 0;
+    // code_factor_terms divides by dim - 1; every advanced encoder reaches it.
+    EXPECT_THROW(
+        quantize_full_single(
+            data.data(), centroid.data(), 1, 4, code.data(), f_add, f_rescale, f_error
+        ),
+        std::invalid_argument
+    );
+    EXPECT_THROW(
+        quantize_full_single(data.data(), 1, 1, code.data(), f_add, f_rescale, f_error),
+        std::invalid_argument
+    );
+}
+
+TEST(RabitqPreconditionTest, FullCodeRejectsBitWidthsTheElementTypeCannotHold) {
+    constexpr size_t kDim = 64;
+    std::array<float, kDim> data;
+    data.fill(0.25F);
+    const std::array<float, kDim> centroid{};
+    std::array<uint8_t, kDim> narrow{};
+    std::array<uint16_t, kDim> wide{};
+    float f_add = 0;
+    float f_rescale = 0;
+    float f_error = 0;
+    // Nine bits (sign plus eight magnitude bits) need more than eight bits per code.
+    EXPECT_THROW(
+        quantize_full_single(
+            data.data(), centroid.data(), kDim, 9, narrow.data(), f_add, f_rescale, f_error
+        ),
+        std::invalid_argument
+    );
+    EXPECT_NO_THROW(quantize_full_single(
+        data.data(), centroid.data(), kDim, 8, narrow.data(), f_add, f_rescale, f_error
+    ));
+    EXPECT_NO_THROW(quantize_full_single(
+        data.data(), centroid.data(), kDim, 9, wide.data(), f_add, f_rescale, f_error
+    ));
+    for (size_t bits : {0U, 10U}) {
+        EXPECT_THROW(
+            quantize_full_single(
+                data.data(),
+                centroid.data(),
+                kDim,
+                bits,
+                wide.data(),
+                f_add,
+                f_rescale,
+                f_error
+            ),
+            std::invalid_argument
+        );
+    }
+}
+
+TEST(RabitqPreconditionTest, SplitSingleQueryRequiresPaddedDimensionMultipleOf32) {
+    std::vector<float> query(96, 0.5F);
+    for (size_t padded_dim : {0U, 16U, 48U, 95U}) {
+        SCOPED_TRACE(padded_dim);
+        EXPECT_THROW(
+            (SplitSingleQuery<float>(query.data(), padded_dim, 3, RabitqConfig{})),
+            std::invalid_argument
+        );
+    }
+    for (size_t padded_dim : {32U, 64U, 96U}) {
+        SCOPED_TRACE(padded_dim);
+        EXPECT_NO_THROW(
+            (SplitSingleQuery<float>(query.data(), padded_dim, 3, RabitqConfig{}))
+        );
+    }
+}
+
 }  // namespace
 }  // namespace rabitqlib::quant