use alloc::vec::Vec;
use crate::error::Error;
const U64_BYTES: usize = core::mem::size_of::<u64>();
const F32_BYTES: usize = core::mem::size_of::<f32>();
const U32_BYTES: usize = core::mem::size_of::<u32>();
const UUID_BYTES: usize = core::mem::size_of::<u128>();
const USIZE_FIELDS: usize = 14;
const F32_FIELDS: usize = 10;
const U32_FIELDS: usize = 3;
const F32S_AT: usize = USIZE_FIELDS * U64_BYTES;
const U32S_AT: usize = F32S_AT + F32_FIELDS * F32_BYTES;
const DB_UUID_AT: usize = U32S_AT + U32_FIELDS * U32_BYTES;
pub const RESERVED_AT: usize = DB_UUID_AT + UUID_BYTES;
const RESERVED_LEN: usize = 8;
pub const ENCODED_LEN: usize = RESERVED_AT + RESERVED_LEN;
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct Config {
pub dim: usize,
pub max_bytes: usize,
pub max_text: usize,
pub max_blob: usize,
pub shards_facts: usize,
pub shards_entities: usize,
pub shards_edges: usize,
pub shards_temporal: usize,
pub shards_postings: usize,
pub bm25_k1: f32,
pub bm25_b: f32,
pub rrf_k: u32,
pub w_bm25: f32,
pub w_vec: f32,
pub w_graph: f32,
pub w_time: f32,
pub w_recency: f32,
pub half_life_days: u32,
pub graph_depth: u32,
pub graph_decay: f32,
pub similar_cos: f32,
pub similar_jaccard: f32,
pub hnsw_m: usize,
pub hnsw_m0: usize,
pub hnsw_ef_construction: usize,
pub hnsw_ef_search: usize,
pub flat_to_hnsw: usize,
pub db_uuid: u128,
}
impl Default for Config {
fn default() -> Self {
Self {
dim: 0,
max_bytes: 2 * 1024 * 1024 * 1024,
max_text: 4096,
max_blob: 64 * 1024,
shards_facts: 1024,
shards_entities: 256,
shards_edges: 512,
shards_temporal: 512,
shards_postings: 2048,
bm25_k1: 1.2,
bm25_b: 0.75,
rrf_k: 60,
w_bm25: 1.0,
w_vec: 1.0,
w_graph: 1.0,
w_time: 1.0,
w_recency: 0.25,
half_life_days: 180,
graph_depth: 2,
graph_decay: 0.5,
similar_cos: 0.85,
similar_jaccard: 0.5,
hnsw_m: 16,
hnsw_m0: 32,
hnsw_ef_construction: 200,
hnsw_ef_search: 64,
flat_to_hnsw: 24_000,
db_uuid: 0,
}
}
}
fn check_weight(v: f32, what: &'static str) -> Result<(), Error> {
if v.is_finite() && v >= 0.0 {
Ok(())
} else {
Err(Error::ConfigMismatch(what))
}
}
fn check_unit(v: f32, what: &'static str) -> Result<(), Error> {
if v.is_finite() && (0.0..=1.0).contains(&v) {
Ok(())
} else {
Err(Error::ConfigMismatch(what))
}
}
impl Config {
pub fn validate(&self) -> Result<(), Error> {
if self.dim > 4096 {
return Err(Error::ConfigMismatch("dim must be <= 4096"));
}
for (shards, what) in [
(self.shards_facts, "shards_facts must be a power of two"),
(
self.shards_entities,
"shards_entities must be a power of two",
),
(self.shards_edges, "shards_edges must be a power of two"),
(
self.shards_temporal,
"shards_temporal must be a power of two",
),
(
self.shards_postings,
"shards_postings must be a power of two",
),
] {
if !shards.is_power_of_two() {
return Err(Error::ConfigMismatch(what));
}
}
if self.max_text == 0 || self.max_text > self.max_blob {
return Err(Error::ConfigMismatch("max_text must be in 1..=max_blob"));
}
if self.max_blob > self.max_bytes {
return Err(Error::ConfigMismatch("max_blob must be <= max_bytes"));
}
if !(self.bm25_k1.is_finite() && self.bm25_k1 > 0.0) {
return Err(Error::ConfigMismatch("bm25_k1 must be positive"));
}
check_unit(self.bm25_b, "bm25_b must be in [0, 1]")?;
if self.rrf_k == 0 {
return Err(Error::ConfigMismatch("rrf_k must be >= 1"));
}
check_weight(self.w_bm25, "w_bm25 must be finite and >= 0")?;
check_weight(self.w_vec, "w_vec must be finite and >= 0")?;
check_weight(self.w_graph, "w_graph must be finite and >= 0")?;
check_weight(self.w_time, "w_time must be finite and >= 0")?;
check_weight(self.w_recency, "w_recency must be finite and >= 0")?;
if self.half_life_days == 0 {
return Err(Error::ConfigMismatch("half_life_days must be >= 1"));
}
if self.graph_depth > 4 {
return Err(Error::ConfigMismatch("graph_depth must be <= 4"));
}
if !(self.graph_decay.is_finite() && self.graph_decay > 0.0 && self.graph_decay <= 1.0) {
return Err(Error::ConfigMismatch("graph_decay must be in (0, 1]"));
}
check_unit(self.similar_cos, "similar_cos must be in [0, 1]")?;
check_unit(self.similar_jaccard, "similar_jaccard must be in [0, 1]")?;
if self.hnsw_m < 2 {
return Err(Error::ConfigMismatch("hnsw_m must be >= 2"));
}
if self.hnsw_m0 < self.hnsw_m {
return Err(Error::ConfigMismatch("hnsw_m0 must be >= hnsw_m"));
}
if self.hnsw_ef_construction < self.hnsw_m {
return Err(Error::ConfigMismatch(
"hnsw_ef_construction must be >= hnsw_m",
));
}
if self.hnsw_ef_search == 0 {
return Err(Error::ConfigMismatch("hnsw_ef_search must be >= 1"));
}
if self.flat_to_hnsw == 0 {
return Err(Error::ConfigMismatch("flat_to_hnsw must be >= 1"));
}
Ok(())
}
pub fn encode(&self, out: &mut Vec<u8>) {
out.reserve(ENCODED_LEN);
for v in [
self.dim,
self.max_bytes,
self.max_text,
self.max_blob,
self.shards_facts,
self.shards_entities,
self.shards_edges,
self.shards_temporal,
self.shards_postings,
self.hnsw_m,
self.hnsw_m0,
self.hnsw_ef_construction,
self.hnsw_ef_search,
self.flat_to_hnsw,
] {
out.extend_from_slice(&(v as u64).to_le_bytes());
}
for v in [
self.bm25_k1,
self.bm25_b,
self.w_bm25,
self.w_vec,
self.w_graph,
self.w_time,
self.w_recency,
self.graph_decay,
self.similar_cos,
self.similar_jaccard,
] {
out.extend_from_slice(&v.to_le_bytes());
}
for v in [self.rrf_k, self.half_life_days, self.graph_depth] {
out.extend_from_slice(&v.to_le_bytes());
}
out.extend_from_slice(&self.db_uuid.to_le_bytes());
out.extend_from_slice(&[0u8; RESERVED_LEN]);
}
pub fn decode(bytes: &[u8]) -> Result<Self, Error> {
if bytes.len() != ENCODED_LEN {
return Err(Error::Corrupt("config block length mismatch"));
}
let mut at = 0usize;
let mut take_usize = || -> Result<usize, Error> {
let v = u64::from_le_bytes(bytes[at..at + U64_BYTES].try_into().unwrap());
at += U64_BYTES;
usize::try_from(v)
.map_err(|_| Error::ConfigMismatch("database requires a 64-bit address space"))
};
let dim = take_usize()?;
let max_bytes = take_usize()?;
let max_text = take_usize()?;
let max_blob = take_usize()?;
let shards_facts = take_usize()?;
let shards_entities = take_usize()?;
let shards_edges = take_usize()?;
let shards_temporal = take_usize()?;
let shards_postings = take_usize()?;
let hnsw_m = take_usize()?;
let hnsw_m0 = take_usize()?;
let hnsw_ef_construction = take_usize()?;
let hnsw_ef_search = take_usize()?;
let flat_to_hnsw = take_usize()?;
let mut at = F32S_AT;
let mut take_f32 = || {
let v = f32::from_le_bytes(bytes[at..at + F32_BYTES].try_into().unwrap());
at += F32_BYTES;
v
};
let bm25_k1 = take_f32();
let bm25_b = take_f32();
let w_bm25 = take_f32();
let w_vec = take_f32();
let w_graph = take_f32();
let w_time = take_f32();
let w_recency = take_f32();
let graph_decay = take_f32();
let similar_cos = take_f32();
let similar_jaccard = take_f32();
let mut at = U32S_AT;
let mut take_u32 = || {
let v = u32::from_le_bytes(bytes[at..at + U32_BYTES].try_into().unwrap());
at += U32_BYTES;
v
};
let rrf_k = take_u32();
let half_life_days = take_u32();
let graph_depth = take_u32();
let db_uuid = u128::from_le_bytes(bytes[DB_UUID_AT..RESERVED_AT].try_into().unwrap());
if bytes[RESERVED_AT..ENCODED_LEN] != [0u8; RESERVED_LEN] {
return Err(Error::Corrupt("reserved config bytes must be zero"));
}
let cfg = Self {
dim,
max_bytes,
max_text,
max_blob,
shards_facts,
shards_entities,
shards_edges,
shards_temporal,
shards_postings,
bm25_k1,
bm25_b,
rrf_k,
w_bm25,
w_vec,
w_graph,
w_time,
w_recency,
half_life_days,
graph_depth,
graph_decay,
similar_cos,
similar_jaccard,
hnsw_m,
hnsw_m0,
hnsw_ef_construction,
hnsw_ef_search,
flat_to_hnsw,
db_uuid,
};
cfg.validate()?;
Ok(cfg)
}
}