weavatrix-git 0.3.1

use std::{
    fs,
    path::{Path, PathBuf},
    sync::{Mutex, OnceLock},
};

use crate::{
    HashKind, ObjectId, ObjectKind, Result, bitmap::BitmapIndex, cache::ByteCache, error::invalid,
};

#[derive(Clone)]
pub(super) struct CachedObject {
    pub(super) kind: ObjectKind,
    pub(super) data: Vec<u8>,
}

pub(crate) struct PackIndex {
    pub(super) pack_path: PathBuf,
    ids: Vec<ObjectId>,
    offsets: Vec<u64>,
    bytes: OnceLock<Vec<u8>>,
    delta_cache: Mutex<ByteCache<u64, CachedObject>>,
    bitmap: Option<BitmapIndex>,
}

impl PackIndex {
    pub(crate) fn open(path: &Path, hash: HashKind, cache_bytes: usize) -> Result<Self> {
        let data = fs::read(path)?;
        let hash_len = hash.bytes();
        if data.get(..4) != Some(&[0xff, b't', b'O', b'c']) {
            return Err(invalid("only pack index v2 is supported"));
        }
        if read_u32(&data, 4)? != 2 {
            return Err(invalid("unsupported pack index version"));
        }
        let fanout_start = 8;
        let count = usize::try_from(read_u32(&data, fanout_start + 255 * 4)?)
            .map_err(|_| invalid("pack object count overflow"))?;
        let ids_start = fanout_start + 256 * 4;
        let crc_start = ids_start
            .checked_add(count * hash_len)
            .ok_or_else(|| invalid("pack index size overflow"))?;
        let offset_start = crc_start
            .checked_add(count * 4)
            .ok_or_else(|| invalid("pack index size overflow"))?;
        let large_start = offset_start
            .checked_add(count * 4)
            .ok_or_else(|| invalid("pack index size overflow"))?;
        let trailer = hash_len * 2;
        if data.len() < large_start + trailer {
            return Err(invalid("truncated pack index"));
        }
        let mut ids = Vec::with_capacity(count);
        for index in 0..count {
            let start = ids_start + index * hash_len;
            ids.push(ObjectId::from_bytes(&data[start..start + hash_len])?);
        }
        if !ids.windows(2).all(|pair| pair[0] < pair[1]) {
            return Err(invalid("pack index identifiers are not sorted"));
        }
        let large_bytes = data.len() - large_start - trailer;
        if !large_bytes.is_multiple_of(8) {
            return Err(invalid("malformed large-offset table"));
        }
        let large_count = large_bytes / 8;
        let mut offsets = Vec::with_capacity(count);
        for index in 0..count {
            let raw = read_u32(&data, offset_start + index * 4)?;
            if raw & 0x8000_0000 == 0 {
                offsets.push(u64::from(raw));
            } else {
                let slot = usize::try_from(raw & 0x7fff_ffff).expect("u32 fits usize");
                if slot >= large_count {
                    return Err(invalid("large pack offset index is out of bounds"));
                }
                offsets.push(read_u64(&data, large_start + slot * 8)?);
            }
        }
        let pack_path = path.with_extension("pack");
        let bitmap = BitmapIndex::open(&path.with_extension("bitmap"), hash)?;
        Ok(Self {
            pack_path,
            ids,
            offsets,
            bytes: OnceLock::new(),
            delta_cache: Mutex::new(ByteCache::new(cache_bytes)),
            bitmap,
        })
    }

    pub(crate) fn offset(&self, id: ObjectId) -> Option<u64> {
        self.ids
            .binary_search(&id)
            .ok()
            .map(|index| self.offsets[index])
    }

    pub(super) fn id_at_offset(&self, offset: u64) -> Option<ObjectId> {
        self.offsets
            .iter()
            .position(|candidate| *candidate == offset)
            .map(|index| self.ids[index])
    }

    pub(super) fn object_count(&self) -> usize {
        self.ids.len()
    }

    pub(super) fn pack_bytes(&self) -> Result<&[u8]> {
        if self.bytes.get().is_none() {
            let bytes = fs::read(&self.pack_path)?;
            let _ = self.bytes.set(bytes);
        }
        Ok(self.bytes.get().expect("pack bytes initialized"))
    }

    pub(super) fn cached(&self, offset: u64) -> Option<CachedObject> {
        self.delta_cache
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .get(offset)
    }

    pub(super) fn cache(&self, offset: u64, object: CachedObject) {
        let bytes = object.data.len();
        self.delta_cache
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .insert(offset, object, bytes);
    }

    pub(crate) fn bitmap_reachable(
        &self,
        id: ObjectId,
        max_objects: usize,
    ) -> Result<Option<Vec<ObjectId>>> {
        let Some(bitmap) = &self.bitmap else {
            return Ok(None);
        };
        let Some(position) = self.ids.binary_search(&id).ok() else {
            return Ok(None);
        };
        bitmap.reachable(position, &self.bitmap_order(), max_objects)
    }

    fn bitmap_order(&self) -> Vec<ObjectId> {
        let mut positions = (0..self.ids.len()).collect::<Vec<_>>();
        positions.sort_unstable_by_key(|position| self.offsets[*position]);
        positions
            .into_iter()
            .map(|position| self.ids[position])
            .collect()
    }
}

fn read_u32(input: &[u8], offset: usize) -> Result<u32> {
    let bytes = input
        .get(offset..offset + 4)
        .ok_or_else(|| invalid("truncated pack index integer"))?;
    Ok(u32::from_be_bytes(bytes.try_into().expect("four bytes")))
}

fn read_u64(input: &[u8], offset: usize) -> Result<u64> {
    let bytes = input
        .get(offset..offset + 8)
        .ok_or_else(|| invalid("truncated pack index integer"))?;
    Ok(u64::from_be_bytes(bytes.try_into().expect("eight bytes")))
}