use iqdb_types::DistanceMetric;
pub use iqdb_cache::{CacheConfig, EvictionPolicy};
pub use iqdb_hnsw::HnswConfig;
pub use iqdb_ivf::IvfConfig;
pub use iqdb_persist::{Compression, FsyncPolicy};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum IndexKind {
#[default]
Flat,
Hnsw(HnswConfig),
Ivf(IvfConfig),
}
impl IndexKind {
#[must_use]
pub(crate) const fn tag(&self) -> u8 {
match self {
Self::Flat => 0,
Self::Hnsw(_) => 1,
Self::Ivf(_) => 2,
}
}
}
#[derive(Debug, Clone, Default)]
pub(crate) struct CoreConfig {
pub(crate) index: IndexKind,
pub(crate) cache: Option<CacheConfig>,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct Durability {
pub(crate) fsync: FsyncPolicy,
pub(crate) compression: Compression,
}
impl Default for Durability {
fn default() -> Self {
Self {
fsync: FsyncPolicy::Always,
compression: Compression::None,
}
}
}
#[derive(Debug, Clone)]
pub struct IqdbConfig {
dim: usize,
metric: DistanceMetric,
core: CoreConfig,
durability: Durability,
}
impl IqdbConfig {
#[must_use]
pub fn new(dim: usize, metric: DistanceMetric) -> Self {
Self {
dim,
metric,
core: CoreConfig::default(),
durability: Durability::default(),
}
}
#[must_use]
pub fn index(mut self, kind: IndexKind) -> Self {
self.core.index = kind;
self
}
#[must_use]
pub fn cache(mut self, cache: CacheConfig) -> Self {
self.core.cache = Some(cache);
self
}
#[must_use]
pub fn fsync(mut self, policy: FsyncPolicy) -> Self {
self.durability.fsync = policy;
self
}
#[must_use]
pub fn compression(mut self, compression: Compression) -> Self {
self.durability.compression = compression;
self
}
#[must_use]
pub fn dim(&self) -> usize {
self.dim
}
#[must_use]
pub fn metric(&self) -> DistanceMetric {
self.metric
}
#[must_use]
pub fn index_kind(&self) -> IndexKind {
self.core.index
}
#[must_use]
pub fn is_cached(&self) -> bool {
self.core.cache.is_some()
}
pub(crate) fn into_parts(self) -> (usize, DistanceMetric, CoreConfig, Durability) {
(self.dim, self.metric, self.core, self.durability)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_index_kind_is_flat() {
assert_eq!(IndexKind::default(), IndexKind::Flat);
assert_eq!(IndexKind::Flat.tag(), 0);
}
#[test]
fn index_kind_tags_are_stable() {
assert_eq!(IndexKind::Hnsw(HnswConfig::default()).tag(), 1);
assert_eq!(IndexKind::Ivf(IvfConfig::default()).tag(), 2);
}
#[test]
fn fluent_builder_threads_choices() {
let cfg = IqdbConfig::new(8, DistanceMetric::Cosine)
.index(IndexKind::Ivf(IvfConfig::default().with_n_probes(4)))
.cache(CacheConfig::new().capacity(512));
assert_eq!(cfg.dim(), 8);
assert_eq!(cfg.metric(), DistanceMetric::Cosine);
assert!(cfg.is_cached());
assert!(matches!(cfg.index_kind(), IndexKind::Ivf(_)));
}
#[test]
fn defaults_are_flat_uncached_and_safely_durable() {
let cfg = IqdbConfig::new(4, DistanceMetric::Euclidean);
assert!(!cfg.is_cached());
assert_eq!(cfg.index_kind(), IndexKind::Flat);
let (dim, metric, core, durability) = cfg.into_parts();
assert_eq!(dim, 4);
assert_eq!(metric, DistanceMetric::Euclidean);
assert!(core.cache.is_none());
assert_eq!(durability.fsync, FsyncPolicy::Always);
assert_eq!(durability.compression, Compression::None);
}
#[test]
fn durability_knobs_thread_through() {
let cfg = IqdbConfig::new(4, DistanceMetric::Cosine)
.fsync(FsyncPolicy::Never)
.compression(Compression::Lz4);
let (_, _, _, durability) = cfg.into_parts();
assert_eq!(durability.fsync, FsyncPolicy::Never);
assert_eq!(durability.compression, Compression::Lz4);
}
}