iridium-db 0.3.0

const STRUCTURED_VECTOR_MAGIC: &[u8; 4] = b"IRV1";
const FLAG_NORMALIZED: u8 = 0b0000_0001;
const HEADER_LEN: usize = 18;
const QUANTIZED_I8_SCALE_LEN: usize = 4;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum VectorEncoding {
    F32 = 0,
    QuantizedI8 = 1,
}

impl VectorEncoding {
    pub fn from_u8(value: u8) -> Option<Self> {
        match value {
            0 => Some(Self::F32),
            1 => Some(Self::QuantizedI8),
            _ => None,
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Self::F32 => "f32",
            Self::QuantizedI8 => "quantized_i8",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum VectorMetric {
    Cosine = 1,
    Euclidean = 2,
}

impl VectorMetric {
    pub fn from_u8(value: u8) -> Option<Self> {
        match value {
            1 => Some(Self::Cosine),
            2 => Some(Self::Euclidean),
            _ => None,
        }
    }

    pub fn from_function(function: &str) -> Option<Self> {
        match function {
            "vector.cosine" => Some(Self::Cosine),
            "vector.euclidean" => Some(Self::Euclidean),
            _ => None,
        }
    }

    pub fn as_function_name(&self) -> &'static str {
        match self {
            Self::Cosine => "vector.cosine",
            Self::Euclidean => "vector.euclidean",
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Cosine => "cosine",
            Self::Euclidean => "euclidean",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VectorSpaceDescriptor {
    pub space_id: u32,
    pub dimension: u16,
    pub metric: VectorMetric,
    pub encoding: VectorEncoding,
    pub normalized: bool,
}

#[derive(Debug, Clone, PartialEq)]
pub struct StructuredVector {
    pub descriptor: VectorSpaceDescriptor,
    pub values: Vec<f32>,
    pub norm: f32,
}

#[derive(Debug, Clone, PartialEq)]
pub enum DecodedVectorPayload {
    Structured(StructuredVector),
    LegacyF32(Vec<f32>),
}

pub fn encode_vector_payload_f32(
    space_id: u32,
    metric: VectorMetric,
    values: &[f32],
    normalized: bool,
) -> Vec<u8> {
    let dimension = u16::try_from(values.len()).unwrap_or(u16::MAX);
    let norm = vector_norm(values);
    let mut out = Vec::with_capacity(HEADER_LEN + values.len() * 4);
    out.extend_from_slice(STRUCTURED_VECTOR_MAGIC);
    out.extend_from_slice(&space_id.to_le_bytes());
    out.extend_from_slice(&dimension.to_le_bytes());
    out.push(VectorEncoding::F32 as u8);
    out.push(metric as u8);
    out.push(if normalized { FLAG_NORMALIZED } else { 0 });
    out.push(0);
    out.extend_from_slice(&norm.to_le_bytes());
    for value in values {
        out.extend_from_slice(&value.to_le_bytes());
    }
    out
}

pub fn encode_vector_payload_quantized_i8(
    space_id: u32,
    metric: VectorMetric,
    values: &[f32],
    normalized: bool,
) -> Result<Vec<u8>, String> {
    if values.is_empty() {
        return Err("quantized_i8 vector payload must not be empty".to_string());
    }
    let dimension = u16::try_from(values.len())
        .map_err(|_| "quantized_i8 vector dimension exceeds u16".to_string())?;
    let max_abs = values
        .iter()
        .map(|value| value.abs())
        .fold(0.0_f32, f32::max);
    let scale = if max_abs <= f32::EPSILON {
        1.0_f32
    } else {
        max_abs / 127.0_f32
    };
    let mut out = Vec::with_capacity(HEADER_LEN + QUANTIZED_I8_SCALE_LEN + values.len());
    out.extend_from_slice(STRUCTURED_VECTOR_MAGIC);
    out.extend_from_slice(&space_id.to_le_bytes());
    out.extend_from_slice(&dimension.to_le_bytes());
    out.push(VectorEncoding::QuantizedI8 as u8);
    out.push(metric as u8);
    out.push(if normalized { FLAG_NORMALIZED } else { 0 });
    out.push(0);
    out.extend_from_slice(&vector_norm(values).to_le_bytes());
    out.extend_from_slice(&scale.to_le_bytes());
    for value in values {
        let quantized = (*value / scale).round().clamp(-127.0, 127.0) as i8;
        out.push(quantized as u8);
    }
    Ok(out)
}

pub fn encode_legacy_f32_payload(values: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(values.len() * 4);
    for value in values {
        out.extend_from_slice(&value.to_le_bytes());
    }
    out
}

pub fn decode_vector_payload(payload: &[u8]) -> Result<DecodedVectorPayload, String> {
    if payload.starts_with(STRUCTURED_VECTOR_MAGIC) {
        return decode_structured_vector_payload(payload);
    }
    decode_legacy_f32_payload(payload)
        .map(DecodedVectorPayload::LegacyF32)
        .ok_or_else(|| "legacy vector payload is not valid packed f32 bytes".to_string())
}

pub fn parse_inline_vector_param(param: &str) -> Option<Vec<f32>> {
    let idx = param.find(':')?;
    let suffix = &param[idx + 1..];
    if suffix.is_empty() {
        return None;
    }
    let mut out = Vec::new();
    for value in suffix.split(':') {
        out.push(value.parse::<f32>().ok()?);
    }
    if out.is_empty() {
        None
    } else {
        Some(out)
    }
}

pub fn vector_norm(values: &[f32]) -> f32 {
    let sum = values
        .iter()
        .map(|value| {
            let value = *value as f64;
            value * value
        })
        .sum::<f64>();
    sum.sqrt() as f32
}

fn decode_structured_vector_payload(payload: &[u8]) -> Result<DecodedVectorPayload, String> {
    if payload.len() < HEADER_LEN {
        return Err("structured vector payload truncated".to_string());
    }
    let space_id = u32::from_le_bytes(payload[4..8].try_into().unwrap());
    let dimension = u16::from_le_bytes(payload[8..10].try_into().unwrap());
    let encoding = VectorEncoding::from_u8(payload[10])
        .ok_or_else(|| format!("unsupported vector encoding {}", payload[10]))?;
    let metric = VectorMetric::from_u8(payload[11])
        .ok_or_else(|| format!("unsupported vector metric {}", payload[11]))?;
    let normalized = payload[12] & FLAG_NORMALIZED != 0;
    let norm = f32::from_le_bytes(payload[14..18].try_into().unwrap());
    let body = &payload[HEADER_LEN..];

    let values = match encoding {
        VectorEncoding::F32 => decode_legacy_f32_payload(body).ok_or_else(|| {
            "structured f32 vector payload body is not valid packed f32 bytes".to_string()
        })?,
        VectorEncoding::QuantizedI8 => decode_quantized_i8_payload(body, dimension as usize)?,
    };

    if values.len() != dimension as usize {
        return Err(format!(
            "vector payload dimension mismatch: header {} values {}",
            dimension,
            values.len()
        ));
    }

    Ok(DecodedVectorPayload::Structured(StructuredVector {
        descriptor: VectorSpaceDescriptor {
            space_id,
            dimension,
            metric,
            encoding,
            normalized,
        },
        values,
        norm,
    }))
}

fn decode_quantized_i8_payload(payload: &[u8], dimension: usize) -> Result<Vec<f32>, String> {
    if payload.len() != QUANTIZED_I8_SCALE_LEN + dimension {
        return Err(format!(
            "quantized_i8 payload length mismatch: expected {}, got {}",
            QUANTIZED_I8_SCALE_LEN + dimension,
            payload.len()
        ));
    }
    let scale = f32::from_le_bytes(payload[0..4].try_into().unwrap());
    if !scale.is_finite() || scale <= 0.0 {
        return Err("quantized_i8 payload has invalid scale".to_string());
    }
    Ok(payload[4..]
        .iter()
        .map(|value| (*value as i8) as f32 * scale)
        .collect())
}

fn decode_legacy_f32_payload(payload: &[u8]) -> Option<Vec<f32>> {
    if payload.is_empty() || !payload.len().is_multiple_of(4) {
        return None;
    }
    Some(
        payload
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
            .collect(),
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn structured_vector_round_trip() {
        let payload = encode_vector_payload_f32(7, VectorMetric::Cosine, &[1.0, -2.0, 3.0], false);
        let decoded = decode_vector_payload(&payload).unwrap();
        let DecodedVectorPayload::Structured(vector) = decoded else {
            panic!("expected structured payload");
        };
        assert_eq!(vector.descriptor.space_id, 7);
        assert_eq!(vector.descriptor.dimension, 3);
        assert_eq!(vector.descriptor.metric, VectorMetric::Cosine);
        assert_eq!(vector.descriptor.encoding, VectorEncoding::F32);
        assert!(!vector.descriptor.normalized);
        assert_eq!(vector.values, vec![1.0, -2.0, 3.0]);
    }

    #[test]
    fn legacy_f32_round_trip() {
        let payload = encode_legacy_f32_payload(&[0.5, 1.5]);
        let decoded = decode_vector_payload(&payload).unwrap();
        let DecodedVectorPayload::LegacyF32(values) = decoded else {
            panic!("expected legacy payload");
        };
        assert_eq!(values, vec![0.5, 1.5]);
    }

    #[test]
    fn quantized_i8_round_trip() {
        let payload =
            encode_vector_payload_quantized_i8(11, VectorMetric::Cosine, &[1.0, -0.5, 0.25], false)
                .unwrap();
        let decoded = decode_vector_payload(&payload).unwrap();
        let DecodedVectorPayload::Structured(vector) = decoded else {
            panic!("expected structured quantized payload");
        };
        assert_eq!(vector.descriptor.space_id, 11);
        assert_eq!(vector.descriptor.dimension, 3);
        assert_eq!(vector.descriptor.encoding, VectorEncoding::QuantizedI8);
        assert_eq!(vector.descriptor.metric, VectorMetric::Cosine);
        assert_eq!(vector.values.len(), 3);
        assert!((vector.values[0] - 1.0).abs() < 0.02);
        assert!((vector.values[1] + 0.5).abs() < 0.02);
        assert!((vector.values[2] - 0.25).abs() < 0.02);
    }

    #[test]
    fn parse_inline_vector_param_extracts_values() {
        assert_eq!(
            parse_inline_vector_param("$q:1:0:-2"),
            Some(vec![1.0, 0.0, -2.0])
        );
        assert_eq!(parse_inline_vector_param("$q"), None);
    }
}