[C API] Implement custom allocator interface and related functions for memory management

bindings/c/include/svs/c/svs_c.h

@@ -189,6 +189,65 @@ struct svs_threadpool_interface {
     void* self;
 };
 
+/// @brief Allocators used for memory management in the SVS C API.
+enum svs_allocator_kind {
+    SVS_ALLOCATOR_KIND_DEFAULT = 0,
+    SVS_ALLOCATOR_KIND_HUGE_PAGE = 1,
+    SVS_ALLOCATOR_KIND_CUSTOM = 2
+};
+
+/// @brief Operations table for a custom allocator interface
+/// @remarks The user must ensure that the allocator implementation is thread-safe and
+/// that the provided function pointers remain valid for the lifetime of the allocator
+/// interface.
+/// @var svs_allocator_interface_ops::version
+///   Version of the allocator interface.
+/// @var svs_allocator_interface_ops::struct_size
+///   Size of the structure, used for versioning and compatibility checks.
+/// @var svs_allocator_interface_ops::allocate
+///   Function pointer to allocate memory with the specified size and alignment.
+///   @param self Pointer to the allocator instance.
+///   @param size Size of the memory to allocate in bytes.
+///   @param alignment Alignment requirement for the allocated memory in bytes.
+///   @param out_err Handle to capture any error that occurs during allocation. User code
+///   may call svs_error_set() to set the error code and message if an error occurs.
+///   @return Pointer to the allocated memory, or NULL if allocation fails.
+/// @var svs_allocator_interface_ops::deallocate
+///   Function pointer to deallocate memory previously allocated by the allocator.
+///   @param self Pointer to the allocator instance.
+///   @param ptr Pointer to the memory to deallocate.
+///   @param size Size of the memory to deallocate in bytes.
+///   @param alignment Alignment requirement for the allocated memory in bytes.
+struct svs_allocator_interface_ops {
+    uint32_t version;
+    size_t struct_size;
+    void* (*allocate)(void* self, size_t size, size_t alignment, svs_error_h out_err);
+    void (*deallocate)(void* self, void* ptr, size_t size, size_t alignment);
+};
+
+/// @brief Macro to create a user-defined allocator interface operations structure
+/// @param allocate_func Function pointer to allocate memory with the specified size and
+/// alignment
+/// @param deallocate_func Function pointer to deallocate memory previously allocated by the
+/// allocator
+#define SVS_INIT_ALLOCATOR_OPS(allocate_func, deallocate_func)     \
+    {                                                              \
+        .version = SVS_C_API_VERSION,                              \
+        .struct_size = sizeof(struct svs_allocator_interface_ops), \
+        .allocate = &allocate_func, .deallocate = &deallocate_func \
+    }
+
+/// @brief Structure representing a custom allocator interface
+/// @var svs_allocator_interface::ops
+///   Function pointers for the allocator operations.
+/// @var svs_allocator_interface::self
+///   Pointer to the user-defined allocator instance. This pointer is passed to the
+///   function pointers in @p ops when they are called.
+struct svs_allocator_interface {
+    struct svs_allocator_interface_ops* ops;
+    void* self;
+};
+
 /// @brief Operations table for a custom ID filter interface
 /// @remarks The user must ensure that the ID filter implementation is thread-safe and
 /// that the provided function pointers remain valid for the lifetime of the ID filter
@@ -424,10 +483,14 @@ typedef enum svs_algorithm_type svs_algorithm_type_t;
 typedef enum svs_data_type svs_data_type_t;
 typedef enum svs_storage_kind svs_storage_kind_t;
 typedef enum svs_threadpool_kind svs_threadpool_kind_t;
+typedef enum svs_allocator_kind svs_allocator_kind_t;
 
 typedef struct svs_threadpool_interface_ops svs_threadpool_ops_t;
 typedef struct svs_threadpool_interface svs_threadpool_t;
 typedef struct svs_threadpool_interface* svs_threadpool_i;
+typedef struct svs_allocator_interface_ops svs_allocator_ops_t;
+typedef struct svs_allocator_interface svs_allocator_t;
+typedef struct svs_allocator_interface* svs_allocator_i;
 
 typedef struct svs_id_filter_interface_ops svs_id_filter_ops_t;
 typedef struct svs_id_filter_interface svs_id_filter_t;
@@ -696,6 +759,28 @@ SVS_API bool svs_index_builder_set_threadpool_custom(
     svs_index_builder_h builder, svs_threadpool_i pool, svs_error_h out_err /*=NULL*/
 );
 
+/// @brief Set the allocator configuration for the index builder
+/// @param builder The index builder handle
+/// @param kind The kind of allocator to use
+/// @param out_err An optional error handle to capture errors
+/// @return true on success, false on failure
+SVS_API bool svs_index_builder_set_allocator(
+    svs_index_builder_h builder, svs_allocator_kind_t kind, svs_error_h out_err /*=NULL*/
+);
+
+/// @brief Set the custom allocator for the index builder
+/// @param builder The index builder handle
+/// @param allocator The custom allocator interface
+/// @param out_err An optional error handle to capture errors
+/// @return true on success, false on failure
+/// @remarks The builder copies @p allocator and its ops table by value, so those two
+/// objects may be freed or modified after this call returns. The object referenced by
+/// @p allocator->self is not copied and must outlive the builder and every index built or
+/// loaded with it.
+SVS_API bool svs_index_builder_set_allocator_custom(
+    svs_index_builder_h builder, svs_allocator_i allocator, svs_error_h out_err /*=NULL*/
+);
+
 /// @brief Estimate the memory usage of an index based on the builder configuration and
 /// number of vectors
 /// @param builder The index builder handle

bindings/c/src/allocator.hpp

@@ -15,8 +15,15 @@
  */
 #pragma once
 
+#include "svs/c/svs_c.h"
+
+#include "error.hpp"
+
+#include <svs/core/allocator.h>
 #include <svs/core/data/simple.h>
+#include <svs/lib/float16.h>
 #include <svs/lib/memory.h>
+#include <svs/lib/meta.h>
 
 namespace svs {
 namespace c_runtime {
@@ -25,5 +32,79 @@ template <typename T, bool UseBlocked, typename Allocator = svs::lib::Allocator<
 using MaybeBlockedAlloc =
     std::conditional_t<UseBlocked, svs::data::Blocked<Allocator>, Allocator>;
 
+template <typename T> class CustomAllocator : public svs::AllocatorInterface {
+  public:
+    using value_type = T;
+
+    static void validate(const svs_allocator_i allocator) {
+        if (allocator == nullptr) {
+            throw std::invalid_argument("Custom allocator pointer cannot be null.");
+        }
+        if (allocator->ops == nullptr) {
+            throw std::invalid_argument("Custom allocator interface is not initialized.");
+        }
+        if (allocator->ops->version > svs_get_version()) {
+            throw std::invalid_argument(
+                "Custom allocator interface version is not supported."
+            );
+        }
+        if (allocator->ops->struct_size < sizeof(svs_allocator_ops_t)) {
+            throw std::invalid_argument(
+                "Incompatible custom allocator interface struct size."
+            );
+        }
+        if (allocator->ops->allocate == nullptr || allocator->ops->deallocate == nullptr) {
+            throw std::invalid_argument(
+                "Custom allocator interface has null function pointers."
+            );
+        }
+    }
+
+    CustomAllocator(const svs_allocator_ops_t& ops, void* self)
+        : ops_(ops)
+        , self_(self) {}
+
+    void* allocate(size_t n) override {
+        svs_error_desc err{SVS_ERROR_UNKNOWN, "Unknown error in custom allocator allocate"};
+
+        auto result = ops_.allocate(self_, n * sizeof(T), alignof(T), &err);
+        if (result == nullptr) {
+            throw svs::c_runtime::out_of_memory(
+                "Custom allocator failed to allocate memory: (" + std::to_string(err.code) +
+                ") " + err.message
+            );
+        }
+        return result;
+    }
+
+    void deallocate(void* p, size_t n) override {
+        ops_.deallocate(self_, p, n * sizeof(T), alignof(T));
+    }
+
+    AllocatorInterface* clone() const override {
+        return new CustomAllocator<T>(ops_, self_);
+    }
+
+    AllocatorInterface* rebind_to(DataType type) const override {
+        return svs::lib::match(
+            AllocatorInterface::rebind_types{},
+            type,
+            [this]<typename Tag>(svs::lib::Type<Tag>) -> AllocatorInterface* {
+                return new CustomAllocator<Tag>(ops_, self_);
+            }
+        );
+    }
+
+  private:
+    svs_allocator_ops_t ops_;
+    void* self_;
+};
+
+template <typename T = std::byte>
+AllocatorHandle<T> make_custom_allocator_handle(const svs_allocator_i allocator) {
+    CustomAllocator<T>::validate(allocator);
+    return AllocatorHandle<T>{
+        std::make_unique<CustomAllocator<T>>(*allocator->ops, allocator->self)};
+}
 } // namespace c_runtime
 } // namespace svs

bindings/c/src/data_builder/leanvec.hpp

@@ -140,7 +140,8 @@ struct lib::
 };
 
 template <bool UseBlocked, typename F> void for_leanvec_specializations(F&& f) {
-    using byte_alloc = svs::c_runtime::MaybeBlockedAlloc<std::byte, UseBlocked>;
+    using byte_alloc = svs::c_runtime::
+        MaybeBlockedAlloc<std::byte, UseBlocked, AllocatorHandle<std::byte>>;
 
 #define X(P, S, D) f.template operator()<LeanVecDataBuilder<P, S, byte_alloc>, D>();
 #define XX(P, S) X(P, S, DistanceL2) X(P, S, DistanceIP) X(P, S, DistanceCosineSimilarity)

bindings/c/src/data_builder/lvq.hpp

@@ -153,7 +153,8 @@ struct lib::DispatchConverter<
 };
 
 template <bool UseBlocked, typename F> void for_lvq_specializations(F&& f) {
-    using byte_alloc = svs::c_runtime::MaybeBlockedAlloc<std::byte, UseBlocked>;
+    using byte_alloc = svs::c_runtime::
+        MaybeBlockedAlloc<std::byte, UseBlocked, AllocatorHandle<std::byte>>;
 #define X(P, S, D) f.template operator()<LVQDataBuilder<P, S, byte_alloc>, D>();
 #define XX(P, S) X(P, S, DistanceL2) X(P, S, DistanceIP) X(P, S, DistanceCosineSimilarity)
     // Pattern:

bindings/c/src/data_builder/simple.hpp

@@ -91,8 +91,10 @@ struct lib::DispatchConverter<const c_runtime::Storage*, SimpleDataBuilder<T, Al
 };
 
 template <bool UseBlocked, typename F> void for_simple_specializations(F&& f) {
-    using float_alloc = svs::c_runtime::MaybeBlockedAlloc<float, UseBlocked>;
-    using float16_alloc = svs::c_runtime::MaybeBlockedAlloc<svs::Float16, UseBlocked>;
+    using float_alloc =
+        svs::c_runtime::MaybeBlockedAlloc<float, UseBlocked, AllocatorHandle<float>>;
+    using float16_alloc = svs::c_runtime::
+        MaybeBlockedAlloc<svs::Float16, UseBlocked, AllocatorHandle<svs::Float16>>;
 #define X(T, A, D) f.template operator()<SimpleDataBuilder<T, A>, D>();
 #define XX(T, A) X(T, A, DistanceL2) X(T, A, DistanceIP) X(T, A, DistanceCosineSimilarity)
     XX(float, float_alloc)

bindings/c/src/data_builder/sq.hpp

@@ -93,8 +93,10 @@ struct lib::DispatchConverter<const c_runtime::Storage*, SQDataBuilder<T, Alloc>
 };
 
 template <bool UseBlocked, typename F> void for_sq_specializations(F&& f) {
-    using int8_alloc = svs::c_runtime::MaybeBlockedAlloc<int8_t, UseBlocked>;
-    using uint8_alloc = svs::c_runtime::MaybeBlockedAlloc<uint8_t, UseBlocked>;
+    using int8_alloc =
+        svs::c_runtime::MaybeBlockedAlloc<int8_t, UseBlocked, AllocatorHandle<int8_t>>;
+    using uint8_alloc =
+        svs::c_runtime::MaybeBlockedAlloc<uint8_t, UseBlocked, AllocatorHandle<uint8_t>>;
 #define X(T, A, D) f.template operator()<SQDataBuilder<T, A>, D>();
 #define XX(T, A) X(T, A, DistanceL2) X(T, A, DistanceIP) X(T, A, DistanceCosineSimilarity)
     XX(uint8_t, uint8_alloc)

bindings/c/src/dispatcher_dynamic_vamana.cpp

@@ -45,21 +45,29 @@ svs::DynamicVamana build_dynamic_vamana_index(
     DataBuilder builder,
     Distance D,
     svs::threads::ThreadPoolHandle pool,
+    const AllocatorHandle<std::byte>& allocator_handle,
     size_t blocksize_bytes
 ) {
     svs::data::BlockingParameters block_params;
     if (blocksize_bytes != 0) {
         block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes);
     }
     using allocator_type = typename DataBuilder::allocator_type;
-    auto allocator = allocator_type{block_params};
-    auto data = builder.build(std::move(src_data.first), pool, allocator);
+    using value_type = typename allocator_type::value_type;
+
+    auto data_allocator_handle = svs::lib::rebind_allocator<value_type>(allocator_handle);
+    auto data_allocator = allocator_type{block_params, data_allocator_handle};
+    auto data = builder.build(std::move(src_data.first), pool, data_allocator);
+
+    auto graph_allocator_handle = svs::lib::rebind_allocator<uint32_t>(allocator_handle);
+    auto graph_allocator = svs::data::Blocked{block_params, graph_allocator_handle};
     return svs::DynamicVamana::build<float>(
         build_params,
         std::move(data),
         std::move(src_data.second),
         std::move(D),
-        std::move(pool)
+        std::move(pool),
+        graph_allocator
     );
 }
 
@@ -70,18 +78,25 @@ svs::DynamicVamana load_dynamic_vamana_index(
     DataLoader loader,
     Distance D,
     svs::threads::ThreadPoolHandle pool,
+    const AllocatorHandle<std::byte>& allocator_handle,
     size_t blocksize_bytes
 ) {
     svs::data::BlockingParameters block_params;
     if (blocksize_bytes != 0) {
         block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes);
     }
     using allocator_type = typename DataLoader::allocator_type;
-    auto allocator = allocator_type{block_params};
+    using value_type = typename allocator_type::value_type;
+    auto data_allocator_handle = svs::lib::rebind_allocator<value_type>(allocator_handle);
+    auto allocator = allocator_type{block_params, data_allocator_handle};
     auto data = loader.load(directory / "data", allocator);
+
+    auto graph_allocator_handle = svs::lib::rebind_allocator<uint32_t>(allocator_handle);
+    auto graph_allocator = svs::data::Blocked{block_params, graph_allocator_handle};
+
     return svs::DynamicVamana::assemble<float>(
         directory / "config",
-        svs::GraphLoader{directory / "graph"},
+        svs::GraphLoader{directory / "graph", graph_allocator},
         std::move(data),
         std::move(D),
         std::move(pool)
@@ -118,6 +133,7 @@ using BuildDynamicIndexDispatcher = svs::lib::Dispatcher<
     const Storage*,
     svs::DistanceType,
     svs::threads::ThreadPoolHandle,
+    const AllocatorHandle<std::byte>&,
     size_t>;
 
 const BuildDynamicIndexDispatcher& build_dynamic_vamana_index_dispatcher() {
@@ -136,6 +152,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_build(
     const Storage* storage,
     svs::DistanceType distance_type,
     svs::threads::ThreadPoolHandle pool,
+    const AllocatorHandle<std::byte>& allocator_handle,
     size_t blocksize_bytes
 ) {
     return build_dynamic_vamana_index_dispatcher().invoke(
@@ -144,6 +161,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_build(
         storage,
         distance_type,
         std::move(pool),
+        allocator_handle,
         blocksize_bytes
     );
 }
@@ -154,6 +172,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load(
     const Storage* storage,
     svs::DistanceType distance_type,
     svs::threads::ThreadPoolHandle pool,
+    const AllocatorHandle<std::byte>& allocator_handle,
     size_t blocksize_bytes
 ) {
     return build_dynamic_vamana_index_dispatcher().invoke(
@@ -162,6 +181,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load(
         storage,
         distance_type,
         std::move(pool),
+        allocator_handle,
         blocksize_bytes
     );
 }
@@ -175,24 +195,26 @@ svs::index::vamana::MemoryBreakdown dispatch_dynamic_vamana_memory_estimate(
     size_t blocksize_bytes
 ) {
     svs::index::vamana::MemoryBreakdown breakdown{};
+
+    // Graph size
     // Graph: SimpleBlockedData<uint32_t> with num_vectors rows and (max_degree + 1)
     // cols; the +1 slot stores the per-node neighbor count.
     using index_type = uint32_t;
+    using graph_allocator_type = svs::data::Blocked<svs::lib::Allocator<index_type>>;
+
     const size_t max_degree = build_params.graph_max_degree;
 
-    // TODO Fix/refactor DynamicVamana index builder to use proper allocator type and
-    // blocking parameters for graph, so that the memory estimate can be accurate for
-    // blocked data. For now, we use the default blocking parameters.
-    // There is MutableVamanaIndex deduction guides for index building defined in
-    // dynamic_index.h which set SimpleBlockedGraph as default graph type.
-    using graph_type = graphs::SimpleBlockedGraph<index_type>;
-    using graph_data_type = typename graph_type::data_type;
-    using allocator_type = graph_data_type::allocator_type;
-    using graph_builder_type = svs::SimpleDataBuilder<index_type, allocator_type>;
+    svs::data::BlockingParameters block_params;
+    if (blocksize_bytes != 0) {
+        block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes);
+    }
+    auto graph_allocator = graph_allocator_type{block_params};
+    auto graph_data_builder = svs::SimpleDataBuilder<index_type, graph_allocator_type>{};
 
     breakdown.graph_bytes =
-        graph_builder_type{}.estimate_size(num_vectors, (max_degree + 1));
+        graph_data_builder.estimate_size(num_vectors, (max_degree + 1), graph_allocator);
 
+    // Data size
     breakdown.data_bytes =
         estimate_data_size_blocked(storage, num_vectors, dimension, blocksize_bytes);