prolly-map 0.2.2

use super::codec::{
    put_cid, put_varint, Reader, FORMAT_VERSION, MAX_KEY_BYTES, VECTOR_ENCODING_F32_LE,
};
use super::{ReferenceKind, TypedReference};
use crate::prolly::cid::Cid;
use crate::prolly::config::Config;
use crate::prolly::encoding::Encoding;
use crate::prolly::error::Error;
use crate::prolly::proximity::{
    DistanceMetric, HierarchyConfig, OverflowConfig, ProximityConfig, ScalarQuantizationConfig,
    VectorStorageConfig,
};
use crate::prolly::tree::Tree;

const MAGIC: &[u8; 4] = b"PRXI";
const COVERING_RADIUS_EUCLIDEAN_F64_UP: u8 = 1;
const NORMALIZATION_NONE: u8 = 0;
const NORMALIZATION_UNIT_F32_FIXED_POINT: u8 = 1;

#[derive(Clone, Debug, PartialEq)]
pub(crate) struct Descriptor {
    pub(crate) config: ProximityConfig,
    pub(crate) count: u64,
    pub(crate) directory: Tree,
    pub(crate) proximity_root: Cid,
}

impl Descriptor {
    pub(crate) fn encode(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(MAGIC);
        bytes.push(FORMAT_VERSION);
        bytes.push(0); // required flags
        bytes.push(VECTOR_ENCODING_F32_LE);
        put_varint(u64::from(self.config.dimensions), &mut bytes);
        bytes.push(self.config.metric.id());
        bytes.push(normalization_id(self.config.metric));
        put_varint(self.count, &mut bytes);
        bytes.push(self.config.hierarchy.log_chunk_size);
        bytes.extend_from_slice(&self.config.hierarchy.level_hash_seed.to_le_bytes());
        put_varint(u64::from(self.config.overflow.min_page_bytes), &mut bytes);
        put_varint(
            u64::from(self.config.overflow.target_page_bytes),
            &mut bytes,
        );
        put_varint(u64::from(self.config.overflow.max_page_bytes), &mut bytes);
        bytes.extend_from_slice(&self.config.overflow.hash_seed.to_le_bytes());
        put_varint(
            u64::from(self.config.vector_storage.inline_threshold_bytes),
            &mut bytes,
        );
        match &self.config.scalar_quantization {
            None => bytes.push(0),
            Some(config) => {
                bytes.push(1);
                put_varint(u64::from(config.group_size), &mut bytes);
            }
        }
        encode_directory(&self.directory, &mut bytes);
        put_cid(&self.proximity_root, &mut bytes);
        bytes.push(COVERING_RADIUS_EUCLIDEAN_F64_UP);
        put_cid(&configuration_fingerprint(&self.config), &mut bytes);
        put_varint(0, &mut bytes); // reserved extension bytes
        bytes
    }

    pub(crate) fn decode(bytes: &[u8]) -> Result<Self, Error> {
        let mut reader = Reader::new(bytes, "descriptor");
        reader.exact(MAGIC)?;
        reader.version()?;
        if reader.u8()? != 0 {
            return Err(reader.invalid("unknown required flags"));
        }
        if reader.u8()? != VECTOR_ENCODING_F32_LE {
            return Err(reader.invalid("unsupported vector encoding"));
        }
        let dimensions =
            u32::try_from(reader.varint()?).map_err(|_| reader.invalid("dimensions exceed u32"))?;
        let metric = DistanceMetric::from_id(reader.u8()?)?;
        if reader.u8()? != normalization_id(metric) {
            return Err(reader.invalid("metric normalization policy mismatch"));
        }
        let count = reader.varint()?;
        let hierarchy = HierarchyConfig {
            log_chunk_size: reader.u8()?,
            level_hash_seed: reader.u64_le()?,
        };
        let overflow = OverflowConfig {
            min_page_bytes: read_u32(&mut reader, "minimum page bytes")?,
            target_page_bytes: read_u32(&mut reader, "target page bytes")?,
            max_page_bytes: read_u32(&mut reader, "maximum page bytes")?,
            hash_seed: reader.u64_le()?,
        };
        let vector_storage = VectorStorageConfig {
            inline_threshold_bytes: read_u32(&mut reader, "inline vector threshold")?,
        };
        let scalar_quantization = match reader.u8()? {
            0 => None,
            1 => Some(ScalarQuantizationConfig {
                group_size: read_u32(&mut reader, "scalar quantization group size")?,
            }),
            _ => return Err(reader.invalid("invalid scalar quantization tag")),
        };
        let directory = decode_directory(&mut reader)?;
        let proximity_root = reader.cid()?;
        if reader.u8()? != COVERING_RADIUS_EUCLIDEAN_F64_UP {
            return Err(reader.invalid("unsupported covering-bound encoding"));
        }
        let stored_fingerprint = reader.cid()?;
        if reader.varint()? != 0 {
            return Err(reader.invalid("reserved extensions are not supported"));
        }
        reader.finish()?;

        let config = ProximityConfig {
            dimensions,
            metric,
            hierarchy,
            overflow,
            vector_storage,
            scalar_quantization,
        };
        config.validate()?;
        let expected_fingerprint = configuration_fingerprint(&config);
        if stored_fingerprint != expected_fingerprint {
            return Err(Error::InvalidProximityObject {
                kind: "descriptor",
                reason: "configuration fingerprint mismatch".to_owned(),
            });
        }
        validate_directory_config(&directory.config)?;
        Ok(Self {
            config,
            count,
            directory,
            proximity_root,
        })
    }

    #[allow(dead_code)] // Used by typed graph traversal in the integration slice.
    pub(crate) fn references(bytes: &[u8]) -> Result<Vec<TypedReference>, Error> {
        let descriptor = Self::decode(bytes)?;
        let mut references = Vec::with_capacity(2);
        if let Some(root) = descriptor.directory.root {
            references.push(TypedReference {
                kind: ReferenceKind::OrderedNode,
                cid: root,
            });
        }
        references.push(TypedReference {
            kind: ReferenceKind::ProximityNode,
            cid: descriptor.proximity_root,
        });
        Ok(references)
    }
}

fn normalization_id(metric: DistanceMetric) -> u8 {
    if metric == DistanceMetric::Cosine {
        NORMALIZATION_UNIT_F32_FIXED_POINT
    } else {
        NORMALIZATION_NONE
    }
}

fn read_u32(reader: &mut Reader<'_>, field: &str) -> Result<u32, Error> {
    u32::try_from(reader.varint()?).map_err(|_| reader.invalid(format!("{field} exceeds u32")))
}

fn encode_directory(tree: &Tree, out: &mut Vec<u8>) {
    match &tree.root {
        Some(root) => {
            out.push(1);
            put_cid(root, out);
        }
        None => out.push(0),
    }
    put_varint(tree.config.min_chunk_size as u64, out);
    put_varint(tree.config.max_chunk_size as u64, out);
    put_varint(u64::from(tree.config.chunking_factor), out);
    out.extend_from_slice(&tree.config.hash_seed.to_le_bytes());
    match &tree.config.encoding {
        Encoding::Raw => out.push(0),
        Encoding::Cbor => out.push(1),
        Encoding::Json => out.push(2),
        Encoding::Custom(name) => {
            out.push(3);
            put_varint(name.len() as u64, out);
            out.extend_from_slice(name.as_bytes());
        }
    }
    put_option(tree.config.node_cache_max_nodes, out);
    put_option(tree.config.node_cache_max_bytes, out);
}

fn decode_directory(reader: &mut Reader<'_>) -> Result<Tree, Error> {
    let root = match reader.u8()? {
        0 => None,
        1 => Some(reader.cid()?),
        _ => return Err(reader.invalid("invalid directory root tag")),
    };
    let min_chunk_size = reader.usize()?;
    let max_chunk_size = reader.usize()?;
    let chunking_factor = read_u32(reader, "chunking factor")?;
    let hash_seed = reader.u64_le()?;
    let encoding = match reader.u8()? {
        0 => Encoding::Raw,
        1 => Encoding::Cbor,
        2 => Encoding::Json,
        3 => {
            let name = reader.bytes(MAX_KEY_BYTES)?;
            Encoding::Custom(
                String::from_utf8(name)
                    .map_err(|_| reader.invalid("custom encoding name is not UTF-8"))?,
            )
        }
        _ => return Err(reader.invalid("unknown directory encoding")),
    };
    Ok(Tree {
        root,
        config: Config {
            min_chunk_size,
            max_chunk_size,
            chunking_factor,
            hash_seed,
            encoding,
            node_cache_max_nodes: get_option(reader)?,
            node_cache_max_bytes: get_option(reader)?,
        },
    })
}

fn put_option(value: Option<usize>, out: &mut Vec<u8>) {
    match value {
        Some(value) => {
            out.push(1);
            put_varint(value as u64, out);
        }
        None => out.push(0),
    }
}

fn get_option(reader: &mut Reader<'_>) -> Result<Option<usize>, Error> {
    match reader.u8()? {
        0 => Ok(None),
        1 => Ok(Some(reader.usize()?)),
        _ => Err(reader.invalid("invalid option tag")),
    }
}

fn validate_directory_config(config: &Config) -> Result<(), Error> {
    if config.min_chunk_size == 0
        || config.max_chunk_size < config.min_chunk_size
        || config.chunking_factor == 0
    {
        return Err(Error::InvalidProximityObject {
            kind: "descriptor",
            reason: "invalid ordered-directory configuration".to_owned(),
        });
    }
    Ok(())
}

fn configuration_fingerprint(config: &ProximityConfig) -> Cid {
    let mut bytes = Vec::new();
    bytes.extend_from_slice(b"PCFG");
    bytes.push(FORMAT_VERSION);
    put_varint(u64::from(config.dimensions), &mut bytes);
    bytes.push(config.metric.id());
    bytes.push(normalization_id(config.metric));
    bytes.push(config.hierarchy.log_chunk_size);
    bytes.extend_from_slice(&config.hierarchy.level_hash_seed.to_le_bytes());
    put_varint(u64::from(config.overflow.min_page_bytes), &mut bytes);
    put_varint(u64::from(config.overflow.target_page_bytes), &mut bytes);
    put_varint(u64::from(config.overflow.max_page_bytes), &mut bytes);
    bytes.extend_from_slice(&config.overflow.hash_seed.to_le_bytes());
    put_varint(
        u64::from(config.vector_storage.inline_threshold_bytes),
        &mut bytes,
    );
    match &config.scalar_quantization {
        None => bytes.push(0),
        Some(config) => {
            bytes.push(1);
            put_varint(u64::from(config.group_size), &mut bytes);
        }
    }
    Cid::from_bytes(&bytes)
}

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

    #[test]
    fn descriptor_round_trip_preserves_every_shape_field() {
        let mut config = ProximityConfig::new(4);
        config.metric = DistanceMetric::Cosine;
        config.hierarchy.level_hash_seed = 99;
        config.overflow.hash_seed = 17;
        config.scalar_quantization = Some(ScalarQuantizationConfig { group_size: 2 });
        let descriptor = Descriptor {
            config,
            count: 12,
            directory: Tree::new(Config::default()),
            proximity_root: Cid::from_bytes(b"root"),
        };
        let bytes = descriptor.encode();
        assert_eq!(&bytes[..5], b"PRXI\x02");
        assert_eq!(Descriptor::decode(&bytes).unwrap(), descriptor);
        assert_eq!(Descriptor::references(&bytes).unwrap().len(), 1);

        let mut bad_flags = bytes.clone();
        bad_flags[5] = 1;
        assert!(Descriptor::decode(&bad_flags).is_err());

        let mut bad_fingerprint = bytes.clone();
        let fingerprint_last = bad_fingerprint.len() - 2;
        bad_fingerprint[fingerprint_last] ^= 1;
        assert!(Descriptor::decode(&bad_fingerprint).is_err());

        let mut trailing = bytes;
        trailing.push(0);
        assert!(Descriptor::decode(&trailing).is_err());
    }
}