Skip to main content

zvec_rust/
config.rs

1use std::ffi::CStr;
2
3use crate::error::{check_error, Error, ErrorCode, Result};
4use crate::types::LogLevel;
5
6/// Configuration builder for initializing the zvec library.
7///
8/// This struct collects configuration values without touching the C library,
9/// making it safe to use even before the library is initialized.
10///
11/// # Example
12/// ```no_run
13/// use zvec_rust::ConfigBuilder;
14///
15/// let config = ConfigBuilder::new()
16///     .memory_limit(1024 * 1024 * 1024)
17///     .num_threads(4)
18///     .enable_console_log(true)
19///     .build();
20/// ```
21pub struct ConfigBuilder {
22    /// Memory limit in bytes (0 = use library default).
23    pub memory_limit: u64,
24    /// Number of threads for query and optimize (0 = use library default).
25    pub num_threads: u32,
26    /// Whether to enable console logging at Info level.
27    pub enable_console_log: bool,
28    /// FTS brute-force-by-keys ratio (None = use library default).
29    pub fts_brute_force_by_keys_ratio: Option<f32>,
30}
31
32impl ConfigBuilder {
33    /// Creates a new builder with default values.
34    pub fn new() -> Self {
35        ConfigBuilder {
36            memory_limit: 0,
37            num_threads: 0,
38            enable_console_log: false,
39            fts_brute_force_by_keys_ratio: None,
40        }
41    }
42
43    /// Sets the memory limit in bytes.
44    pub fn memory_limit(mut self, bytes: u64) -> Self {
45        self.memory_limit = bytes;
46        self
47    }
48
49    /// Sets the number of threads for both query and optimize.
50    pub fn num_threads(mut self, count: u32) -> Self {
51        self.num_threads = count;
52        self
53    }
54
55    /// Enables or disables console logging at Info level.
56    pub fn enable_console_log(mut self, enable: bool) -> Self {
57        self.enable_console_log = enable;
58        self
59    }
60
61    /// Sets the FTS brute-force-by-keys ratio.
62    pub fn fts_brute_force_by_keys_ratio(mut self, ratio: f32) -> Self {
63        self.fts_brute_force_by_keys_ratio = Some(ratio);
64        self
65    }
66
67    /// Finalizes the builder configuration.
68    ///
69    /// This is a no-op that returns `self` for API consistency. The builder
70    /// is a plain-data struct — no C resources are allocated here.
71    /// To apply the configuration, pass the result to [`initialize`].
72    pub fn build(self) -> Self {
73        self
74    }
75}
76
77impl Default for ConfigBuilder {
78    fn default() -> Self {
79        Self::new()
80    }
81}
82
83/// Low-level configuration handle wrapping the C API config.
84pub(crate) struct ConfigData {
85    pub(crate) handle: *mut zvec_rust_sys::zvec_config_data_t,
86}
87
88impl ConfigData {
89    pub(crate) fn new() -> Result<Self> {
90        let handle = unsafe { zvec_rust_sys::zvec_config_data_create() };
91        if handle.is_null() {
92            return Err(Error {
93                code: ErrorCode::InternalError,
94                message: "failed to create config data".into(),
95            });
96        }
97        Ok(ConfigData { handle })
98    }
99
100    pub(crate) fn set_memory_limit(&mut self, bytes: u64) -> Result<()> {
101        check_error(unsafe { zvec_rust_sys::zvec_config_data_set_memory_limit(self.handle, bytes) })
102    }
103
104    pub(crate) fn set_query_thread_count(&mut self, count: u32) -> Result<()> {
105        check_error(unsafe {
106            zvec_rust_sys::zvec_config_data_set_query_thread_count(self.handle, count)
107        })
108    }
109
110    pub(crate) fn set_optimize_thread_count(&mut self, count: u32) -> Result<()> {
111        check_error(unsafe {
112            zvec_rust_sys::zvec_config_data_set_optimize_thread_count(self.handle, count)
113        })
114    }
115
116    pub(crate) fn set_fts_brute_force_by_keys_ratio(&mut self, ratio: f32) -> Result<()> {
117        check_error(unsafe {
118            zvec_rust_sys::zvec_config_data_set_fts_brute_force_by_keys_ratio(self.handle, ratio)
119        })
120    }
121
122    pub(crate) fn set_console_log(&mut self, level: LogLevel) -> Result<()> {
123        let log_config = unsafe { zvec_rust_sys::zvec_config_log_create_console(level as u32) };
124        if log_config.is_null() {
125            return Err(Error {
126                code: ErrorCode::InternalError,
127                message: "failed to create console log config".into(),
128            });
129        }
130        // Ownership of log_config transfers to config_data on success.
131        // On failure, we must free it manually to avoid a leak.
132        let result = check_error(unsafe {
133            zvec_rust_sys::zvec_config_data_set_log_config(self.handle, log_config)
134        });
135        if result.is_err() {
136            unsafe { zvec_rust_sys::zvec_config_log_destroy(log_config) };
137        }
138        result
139    }
140}
141
142impl Drop for ConfigData {
143    fn drop(&mut self) {
144        if !self.handle.is_null() {
145            // Safety: handle was created by zvec_config_data_create
146            unsafe { zvec_rust_sys::zvec_config_data_destroy(self.handle) };
147        }
148    }
149}
150
151/// Initializes the zvec library with optional configuration.
152///
153/// Pass `None` to use default configuration, or provide a [`ConfigBuilder`]
154/// to customize memory limits, thread counts, and logging.
155///
156/// # Examples
157///
158/// ```no_run
159/// use zvec_rust::*;
160///
161/// // Default initialization
162/// initialize(None)?;
163///
164/// // With builder
165/// let config = ConfigBuilder::new()
166///     .memory_limit(1024 * 1024 * 1024)
167///     .num_threads(4)
168///     .build();
169/// initialize(Some(&config))?;
170/// # Ok::<(), zvec_rust::Error>(())
171/// ```
172pub fn initialize(config: Option<&ConfigBuilder>) -> Result<()> {
173    match config {
174        None => check_error(unsafe { zvec_rust_sys::zvec_initialize(std::ptr::null()) }),
175        Some(builder) => {
176            let mut cfg = ConfigData::new()?;
177            if builder.memory_limit > 0 {
178                cfg.set_memory_limit(builder.memory_limit)?;
179            }
180            if builder.num_threads > 0 {
181                cfg.set_query_thread_count(builder.num_threads)?;
182                cfg.set_optimize_thread_count(builder.num_threads)?;
183            }
184            if builder.enable_console_log {
185                cfg.set_console_log(LogLevel::Info)?;
186            }
187            if let Some(ratio) = builder.fts_brute_force_by_keys_ratio {
188                cfg.set_fts_brute_force_by_keys_ratio(ratio)?;
189            }
190            check_error(unsafe { zvec_rust_sys::zvec_initialize(cfg.handle as *const _) })
191        }
192    }
193}
194
195/// Shuts down the zvec library and releases all resources.
196#[doc(hidden)]
197pub fn shutdown() -> Result<()> {
198    check_error(unsafe { zvec_rust_sys::zvec_shutdown() })
199}
200
201/// Returns `true` if the library has been initialized.
202pub fn is_initialized() -> bool {
203    unsafe { zvec_rust_sys::zvec_is_initialized() }
204}
205
206/// Returns the library version string.
207pub fn version() -> String {
208    unsafe {
209        let ptr = zvec_rust_sys::zvec_get_version();
210        if ptr.is_null() {
211            return String::new();
212        }
213        CStr::from_ptr(ptr).to_string_lossy().into_owned()
214    }
215}
216
217/// Checks if the current library version meets the minimum requirements.
218pub fn check_version(major: i32, minor: i32, patch: i32) -> bool {
219    unsafe { zvec_rust_sys::zvec_check_version(major, minor, patch) }
220}
221
222/// Returns the major version number.
223pub fn version_major() -> i32 {
224    unsafe { zvec_rust_sys::zvec_get_version_major() }
225}
226
227/// Returns the minor version number.
228pub fn version_minor() -> i32 {
229    unsafe { zvec_rust_sys::zvec_get_version_minor() }
230}
231
232/// Returns the patch version number.
233pub fn version_patch() -> i32 {
234    unsafe { zvec_rust_sys::zvec_get_version_patch() }
235}
236
237/// Returns the I/O backend type code used for DiskAnn disk reads.
238///
239/// Compare against the `ZVEC_IO_BACKEND_TYPE_*` constants in [`crate::sys`]:
240/// Linux selects io_uring, then libaio, then synchronous pread; macOS ARM64
241/// uses pread; Windows uses overlapped I/O. Use [`io_backend_type_name`] for
242/// the matching name.
243pub fn io_backend_type() -> u32 {
244    unsafe { zvec_rust_sys::zvec_get_io_backend_type() }
245}
246
247/// Returns the name of an I/O backend type code.
248///
249/// One of `"io_uring"`, `"libaio"`, `"pread"`, `"windows_overlapped"`, or
250/// `"unknown"` for an unrecognized code.
251pub fn io_backend_type_name(backend_type: u32) -> String {
252    unsafe {
253        let ptr = zvec_rust_sys::zvec_get_io_backend_type_name(backend_type);
254        if ptr.is_null() {
255            return String::new();
256        }
257        CStr::from_ptr(ptr).to_string_lossy().into_owned()
258    }
259}
260
261/// Returns a human-readable description of the active I/O backend.
262///
263/// On Linux the pread description also explains why io_uring and libaio were
264/// unavailable. The returned string is a copy of a thread-local buffer owned
265/// by the zvec library.
266pub fn io_backend_description() -> String {
267    unsafe {
268        let ptr = zvec_rust_sys::zvec_get_io_backend_description();
269        if ptr.is_null() {
270            return String::new();
271        }
272        CStr::from_ptr(ptr).to_string_lossy().into_owned()
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    #[test]
281    fn config_builder_defaults() {
282        let builder = ConfigBuilder::new();
283        assert_eq!(builder.memory_limit, 0);
284        assert_eq!(builder.num_threads, 0);
285        assert!(!builder.enable_console_log);
286    }
287
288    #[test]
289    fn config_builder_chaining() {
290        let builder = ConfigBuilder::new()
291            .memory_limit(1024)
292            .num_threads(4)
293            .enable_console_log(true)
294            .build();
295        assert_eq!(builder.memory_limit, 1024);
296        assert_eq!(builder.num_threads, 4);
297        assert!(builder.enable_console_log);
298    }
299
300    #[test]
301    fn config_builder_default_trait() {
302        let builder = ConfigBuilder::default();
303        assert_eq!(builder.memory_limit, 0);
304        assert_eq!(builder.num_threads, 0);
305        assert!(!builder.enable_console_log);
306    }
307
308    #[test]
309    fn config_builder_memory_limit_setter() {
310        let builder = ConfigBuilder::new().memory_limit(2048);
311        assert_eq!(builder.memory_limit, 2048);
312    }
313
314    #[test]
315    fn config_builder_num_threads_setter() {
316        let builder = ConfigBuilder::new().num_threads(8);
317        assert_eq!(builder.num_threads, 8);
318    }
319
320    #[test]
321    fn config_builder_enable_console_log_setter() {
322        let builder = ConfigBuilder::new().enable_console_log(true);
323        assert!(builder.enable_console_log);
324    }
325
326    #[test]
327    fn config_builder_build_returns_self() {
328        let builder = ConfigBuilder::new()
329            .memory_limit(4096)
330            .num_threads(2)
331            .enable_console_log(true)
332            .build();
333        assert_eq!(builder.memory_limit, 4096);
334        assert_eq!(builder.num_threads, 2);
335        assert!(builder.enable_console_log);
336    }
337
338    #[test]
339    fn config_builder_overwrite_values() {
340        let builder = ConfigBuilder::new()
341            .memory_limit(1024)
342            .memory_limit(2048)
343            .num_threads(4)
344            .num_threads(8)
345            .enable_console_log(false)
346            .enable_console_log(true);
347        assert_eq!(builder.memory_limit, 2048);
348        assert_eq!(builder.num_threads, 8);
349        assert!(builder.enable_console_log);
350    }
351
352    #[test]
353    fn config_builder_zero_values() {
354        let builder = ConfigBuilder::new().memory_limit(0).num_threads(0);
355        assert_eq!(builder.memory_limit, 0);
356        assert_eq!(builder.num_threads, 0);
357    }
358
359    #[test]
360    fn config_builder_large_memory_limit() {
361        let builder = ConfigBuilder::new().memory_limit(u64::MAX);
362        assert_eq!(builder.memory_limit, u64::MAX);
363    }
364
365    #[test]
366    fn config_builder_fts_ratio() {
367        let builder = ConfigBuilder::new();
368        assert_eq!(builder.fts_brute_force_by_keys_ratio, None);
369
370        let builder = ConfigBuilder::new().fts_brute_force_by_keys_ratio(0.5);
371        assert_eq!(builder.fts_brute_force_by_keys_ratio, Some(0.5));
372    }
373
374    #[test]
375    fn io_backend_type_name_maps_every_code() {
376        use zvec_rust_sys::{
377            ZVEC_IO_BACKEND_TYPE_IO_URING, ZVEC_IO_BACKEND_TYPE_LIBAIO, ZVEC_IO_BACKEND_TYPE_PREAD,
378            ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED,
379        };
380
381        assert_eq!(io_backend_type_name(ZVEC_IO_BACKEND_TYPE_PREAD), "pread");
382        assert_eq!(io_backend_type_name(ZVEC_IO_BACKEND_TYPE_LIBAIO), "libaio");
383        assert_eq!(
384            io_backend_type_name(ZVEC_IO_BACKEND_TYPE_IO_URING),
385            "io_uring"
386        );
387        assert_eq!(
388            io_backend_type_name(ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED),
389            "windows_overlapped"
390        );
391        assert_eq!(io_backend_type_name(u32::MAX), "unknown");
392    }
393
394    #[test]
395    fn io_backend_introspection_reports_the_active_backend() {
396        let backend = io_backend_type();
397        let name = io_backend_type_name(backend);
398        assert!(
399            ["pread", "libaio", "io_uring", "windows_overlapped"].contains(&name.as_str()),
400            "unexpected io backend {} (code {})",
401            name,
402            backend
403        );
404        assert!(!io_backend_description().is_empty());
405    }
406}