use std::collections::{BTreeMap, HashSet};
use std::io::{Read, Seek, SeekFrom};
use crate::block::BlockDescriptor;
use crate::catalog::BLOCK_SIZE;
use crate::error::VssError;
use crate::store::{StoreBlockHeader, STORE_BLOCK_HEADER_LEN};
use crate::VssVolume;
const DIFF_AREA_RECORD_TYPE: u32 = 0x0003;
const BITMAP_RECORD_TYPE: u32 = 0x0006;
const BLOCK_DESCRIPTOR_LEN: usize = 32;
const SUB_BLOCK_COUNT: usize = 32;
const SUB_BLOCK_SIZE: usize = BLOCK_SIZE / SUB_BLOCK_COUNT;
const MAX_STORE_BLOCKS: usize = 1 << 20;
pub struct Snapshot<'v, R> {
reader: &'v mut R,
volume_size: u64,
block_map: BTreeMap<u64, Vec<BlockDescriptor>>,
bitmap: Vec<u8>,
}
impl<R: Read + Seek> VssVolume<R> {
pub fn snapshot(&mut self, index: usize) -> Result<Snapshot<'_, R>, VssError> {
let (store_header_offset, store_bitmap_offset) = {
let count = self.stores.len();
let d = self
.stores
.get(index)
.ok_or(VssError::StoreIndexOutOfRange { index, count })?;
let header = d
.store_header_offset
.ok_or(VssError::StoreInfoUnavailable { index })?;
let bitmap = d
.store_bitmap_offset
.ok_or(VssError::StoreInfoUnavailable { index })?;
(header, bitmap)
};
let diff_start = store_header_offset.saturating_add(BLOCK_SIZE as u64);
let block_map = build_block_map(&mut self.reader, diff_start, self.volume_size)?;
let bitmap = build_bitmap(&mut self.reader, store_bitmap_offset, self.volume_size)?;
Ok(Snapshot {
reader: &mut self.reader,
volume_size: self.volume_size,
block_map,
bitmap,
})
}
}
impl<R: Read + Seek> Snapshot<'_, R> {
#[must_use]
fn is_unallocated(&self, block_number: u64) -> bool {
let byte = block_number / 8;
let bit = (block_number % 8) as u32;
usize::try_from(byte)
.ok()
.and_then(|i| self.bitmap.get(i))
.is_some_and(|b| b & (1u8 << bit) != 0)
}
pub fn read_block(&mut self, offset: u64) -> Result<[u8; BLOCK_SIZE], VssError> {
let block_number = offset / BLOCK_SIZE as u64;
let base = block_number.saturating_mul(BLOCK_SIZE as u64);
let mut out = [0u8; BLOCK_SIZE];
match self.block_map.get(&base) {
Some(descriptors) if !descriptors.is_empty() => {
let last_plain = descriptors.iter().rfind(|d| {
!d.flags.is_forwarder() && !d.flags.is_overlay() && !d.flags.is_not_used()
});
let base_source = last_plain.map_or(base, |d| d.store_offset);
read_block_at(self.reader, base_source, self.volume_size, &mut out)?;
for d in descriptors
.iter()
.filter(|d| d.flags.is_overlay() && !d.flags.is_not_used())
{
let mut overlay = [0u8; BLOCK_SIZE];
if read_block_at(self.reader, d.store_offset, self.volume_size, &mut overlay)? {
apply_overlay(&mut out, &overlay, d.allocation_bitmap);
}
}
Ok(out)
}
_ => {
if self.is_unallocated(block_number) {
Ok(out) } else {
read_block_at(self.reader, base, self.volume_size, &mut out)?;
Ok(out)
}
}
}
}
pub fn read_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<(), VssError> {
let mut written = 0usize;
let mut cursor = offset;
while written < buf.len() {
let base = (cursor / BLOCK_SIZE as u64).saturating_mul(BLOCK_SIZE as u64);
let within = (cursor - base) as usize;
let block = self.read_block(base)?;
let n = (BLOCK_SIZE - within).min(buf.len() - written);
buf[written..written + n].copy_from_slice(&block[within..within + n]);
written += n;
cursor = cursor.saturating_add(n as u64);
}
Ok(())
}
}
fn apply_overlay(out: &mut [u8; BLOCK_SIZE], overlay: &[u8; BLOCK_SIZE], allocation_bitmap: u32) {
for i in 0..SUB_BLOCK_COUNT {
if allocation_bitmap & (1u32 << i) != 0 {
let start = i * SUB_BLOCK_SIZE;
let end = start + SUB_BLOCK_SIZE;
out[start..end].copy_from_slice(&overlay[start..end]);
}
}
}
fn read_block_at<R: Read + Seek>(
reader: &mut R,
offset: u64,
volume_size: u64,
out: &mut [u8; BLOCK_SIZE],
) -> Result<bool, VssError> {
match offset.checked_add(BLOCK_SIZE as u64) {
Some(end) if end <= volume_size => {
reader.seek(SeekFrom::Start(offset))?;
reader.read_exact(out)?;
Ok(true)
}
_ => {
*out = [0u8; BLOCK_SIZE];
Ok(false)
}
}
}
fn build_block_map<R: Read + Seek>(
reader: &mut R,
first: u64,
volume_size: u64,
) -> Result<BTreeMap<u64, Vec<BlockDescriptor>>, VssError> {
let mut map: BTreeMap<u64, Vec<BlockDescriptor>> = BTreeMap::new();
let mut visited: HashSet<u64> = HashSet::new();
let mut next = first;
let mut blocks = 0usize;
while next != 0 && blocks < MAX_STORE_BLOCKS {
if !visited.insert(next) {
break; }
match next.checked_add(BLOCK_SIZE as u64) {
Some(end) if end <= volume_size => {}
_ => break, }
reader.seek(SeekFrom::Start(next))?;
let mut block = vec![0u8; BLOCK_SIZE];
reader.read_exact(&mut block)?;
let header = StoreBlockHeader::parse(&block);
if header.record_type != DIFF_AREA_RECORD_TYPE {
break;
}
let mut off = STORE_BLOCK_HEADER_LEN;
while off + BLOCK_DESCRIPTOR_LEN <= BLOCK_SIZE {
let record = &block[off..off + BLOCK_DESCRIPTOR_LEN];
if record.iter().all(|&x| x == 0) {
break; }
let descriptor = BlockDescriptor::parse(record);
map.entry(descriptor.original_offset)
.or_default()
.push(descriptor);
off += BLOCK_DESCRIPTOR_LEN;
}
blocks += 1;
next = header.next_offset;
}
Ok(map)
}
fn build_bitmap<R: Read + Seek>(
reader: &mut R,
first: u64,
volume_size: u64,
) -> Result<Vec<u8>, VssError> {
let mut pieces: Vec<(u64, Vec<u8>)> = Vec::new();
let mut visited: HashSet<u64> = HashSet::new();
let mut next = first;
let mut blocks = 0usize;
while next != 0 && blocks < MAX_STORE_BLOCKS {
if !visited.insert(next) {
break; }
match next.checked_add(BLOCK_SIZE as u64) {
Some(end) if end <= volume_size => {}
_ => break, }
reader.seek(SeekFrom::Start(next))?;
let mut block = vec![0u8; BLOCK_SIZE];
reader.read_exact(&mut block)?;
let header = StoreBlockHeader::parse(&block);
if header.record_type != BITMAP_RECORD_TYPE {
break;
}
pieces.push((
header.relative_offset,
block[STORE_BLOCK_HEADER_LEN..BLOCK_SIZE].to_vec(),
));
blocks += 1;
next = header.next_offset;
}
pieces.sort_by_key(|(rel, _)| *rel);
let mut bitmap = Vec::with_capacity(
pieces
.len()
.saturating_mul(BLOCK_SIZE - STORE_BLOCK_HEADER_LEN),
);
for (_, payload) in pieces {
bitmap.extend_from_slice(&payload);
}
Ok(bitmap)
}