1use std::ffi::CStr;
2
3use crate::error::{check_error, Error, ErrorCode, Result};
4use crate::types::LogLevel;
5
6pub struct ConfigBuilder {
22 pub memory_limit: u64,
24 pub num_threads: u32,
26 pub enable_console_log: bool,
28 pub fts_brute_force_by_keys_ratio: Option<f32>,
30}
31
32impl ConfigBuilder {
33 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 pub fn memory_limit(mut self, bytes: u64) -> Self {
45 self.memory_limit = bytes;
46 self
47 }
48
49 pub fn num_threads(mut self, count: u32) -> Self {
51 self.num_threads = count;
52 self
53 }
54
55 pub fn enable_console_log(mut self, enable: bool) -> Self {
57 self.enable_console_log = enable;
58 self
59 }
60
61 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 pub fn build(self) -> Self {
73 self
74 }
75}
76
77impl Default for ConfigBuilder {
78 fn default() -> Self {
79 Self::new()
80 }
81}
82
83pub(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 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 unsafe { zvec_rust_sys::zvec_config_data_destroy(self.handle) };
147 }
148 }
149}
150
151pub 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#[doc(hidden)]
197pub fn shutdown() -> Result<()> {
198 check_error(unsafe { zvec_rust_sys::zvec_shutdown() })
199}
200
201pub fn is_initialized() -> bool {
203 unsafe { zvec_rust_sys::zvec_is_initialized() }
204}
205
206pub 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
217pub fn check_version(major: i32, minor: i32, patch: i32) -> bool {
219 unsafe { zvec_rust_sys::zvec_check_version(major, minor, patch) }
220}
221
222pub fn version_major() -> i32 {
224 unsafe { zvec_rust_sys::zvec_get_version_major() }
225}
226
227pub fn version_minor() -> i32 {
229 unsafe { zvec_rust_sys::zvec_get_version_minor() }
230}
231
232pub fn version_patch() -> i32 {
234 unsafe { zvec_rust_sys::zvec_get_version_patch() }
235}
236
237pub fn io_backend_type() -> u32 {
244 unsafe { zvec_rust_sys::zvec_get_io_backend_type() }
245}
246
247pub 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
261pub 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}