use std::io::{Read, Seek, SeekFrom, Write};
use crate::error::{Error, Result};
use crate::format::btree_v2::BTreeV2Header;
use crate::format::checksum::checksum_metadata;
use crate::io::reader::HdfReader;
use super::MutableFile;
struct BTreeV2RecordLocation {
node_addr: u64,
record_offset: usize,
checksum_len: usize,
}
struct BTreeV2MutableNodeInfo {
max_nrec: usize,
cum_max_nrec: u64,
cum_max_nrec_size: usize,
}
struct BuiltBTreeV2Node {
addr: u64,
nrecords: u16,
total_records: u64,
}
impl MutableFile {
#[allow(clippy::too_many_arguments)]
pub(super) fn rewrite_btree_v2_chunk(
&mut self,
index_addr: u64,
info: &crate::hl::dataset::DatasetInfo,
chunk_coords: &[u64],
chunk_dims: &[u64],
chunk_size: u64,
chunk_addr: u64,
unfiltered_chunk_bytes: usize,
) -> Result<()> {
let scaled_coords = Self::scaled_chunk_coords(chunk_coords, chunk_dims)?;
let filtered = info
.filter_pipeline
.as_ref()
.map(|pipeline| !pipeline.filters.is_empty())
.unwrap_or(false);
let chunk_size_len = if filtered {
Self::filtered_chunk_size_len(
info.layout.version,
unfiltered_chunk_bytes,
self.superblock.sizeof_size,
)
} else {
0
};
let sa = usize::from(self.superblock.sizeof_addr);
let expected_record_size =
Self::btree_v2_chunk_record_size(sa, filtered, chunk_size_len, scaled_coords.len())?;
let mut guard = self.inner.lock();
let reader = &mut guard.reader;
let header = BTreeV2Header::read_at(reader, index_addr)?;
if usize::from(header.record_size) != expected_record_size {
return Err(Error::InvalidFormat(format!(
"v2 B-tree chunk record size {} does not match expected {expected_record_size}",
header.record_size
)));
}
if header.tree_type != 10 && header.tree_type != 11 {
return Err(Error::Unsupported(format!(
"write_chunk cannot update v2 B-tree type {} chunk indexes",
header.tree_type
)));
}
let mut records = Vec::new();
let expected_records = usize::try_from(header.total_records)
.map_err(|_| Error::InvalidFormat("v2 B-tree record count exceeds usize".into()))?;
records.reserve(expected_records.saturating_add(1));
crate::format::btree_v2::collect_all_records_into(reader, index_addr, &mut records)?;
match Self::find_btree_v2_chunk_record_index(
&records,
&scaled_coords,
filtered,
chunk_size_len,
sa,
)? {
Ok(_) => {
let record = Self::encode_btree_v2_chunk_record(
chunk_addr,
chunk_size,
&scaled_coords,
filtered,
chunk_size_len,
sa,
)?;
let location = Self::find_btree_v2_chunk_record_location(
reader,
&header,
&scaled_coords,
filtered,
chunk_size_len,
sa,
)?
.ok_or_else(|| {
Error::InvalidFormat(
"existing v2 B-tree chunk record was not found on disk".into(),
)
})?;
drop(guard);
return self.rewrite_btree_v2_record_at(&location, &record);
}
Err(index) => {
let new_record = Self::encode_btree_v2_chunk_record(
chunk_addr,
chunk_size,
&scaled_coords,
filtered,
chunk_size_len,
sa,
)?;
records.insert(index, new_record);
}
}
drop(guard);
self.rebuild_btree_v2_chunk_tree(index_addr, &header, &records)
}
fn rewrite_btree_v2_record_at(
&mut self,
location: &BTreeV2RecordLocation,
record: &[u8],
) -> Result<()> {
let record_end = location
.record_offset
.checked_add(record.len())
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
if record_end > location.checksum_len {
return Err(Error::InvalidFormat(
"v2 B-tree record extends past checksum span".into(),
));
}
let node_pos = self.physical_addr(location.node_addr, "v2 B-tree node")?;
let mut node = vec![0; location.checksum_len];
self.read_fresh_bytes_into(node_pos, &mut node)?;
node[location.record_offset..record_end].copy_from_slice(record);
let checksum = checksum_metadata(&node);
let record_pos = node_pos
.checked_add(Self::usize_to_u64(
location.record_offset,
"v2 B-tree record offset",
)?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record position overflow".into()))?;
self.write_handle.seek(SeekFrom::Start(record_pos))?;
self.write_handle.write_all(record)?;
let checksum_pos = node_pos
.checked_add(Self::usize_to_u64(
location.checksum_len,
"v2 B-tree checksum offset",
)?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree checksum position overflow".into()))?;
self.write_handle.seek(SeekFrom::Start(checksum_pos))?;
self.write_handle.write_all(&checksum.to_le_bytes())?;
Ok(())
}
pub(super) fn prune_btree_v2_chunks_outside_extent(
&mut self,
index_addr: u64,
info: &crate::hl::dataset::DatasetInfo,
new_dims: &[u64],
chunk_dims: &[u64],
) -> Result<()> {
let filtered = info
.filter_pipeline
.as_ref()
.map(|pipeline| !pipeline.filters.is_empty())
.unwrap_or(false);
let unfiltered_chunk_bytes = Self::chunk_unfiltered_byte_len(info, chunk_dims)?;
let chunk_size_len = if filtered {
Self::filtered_chunk_size_len(
info.layout.version,
unfiltered_chunk_bytes,
self.superblock.sizeof_size,
)
} else {
0
};
let sa = usize::from(self.superblock.sizeof_addr);
let expected_record_size =
Self::btree_v2_chunk_record_size(sa, filtered, chunk_size_len, chunk_dims.len())?;
let mut guard = self.inner.lock();
let reader = &mut guard.reader;
let header = BTreeV2Header::read_at(reader, index_addr)?;
if usize::from(header.record_size) != expected_record_size {
return Err(Error::InvalidFormat(format!(
"v2 B-tree chunk record size {} does not match expected {expected_record_size}",
header.record_size
)));
}
if header.tree_type != 10 && header.tree_type != 11 {
return Err(Error::Unsupported(format!(
"resize cannot prune v2 B-tree type {} chunk indexes",
header.tree_type
)));
}
let mut records = Vec::new();
crate::format::btree_v2::collect_all_records_into(reader, index_addr, &mut records)?;
let original_len = records.len();
records.retain(|record| {
Self::btree_v2_record_starts_inside_extent(
record,
new_dims,
chunk_dims,
filtered,
chunk_size_len,
sa,
)
.unwrap_or(false)
});
drop(guard);
if records.len() == original_len {
return Ok(());
}
self.rebuild_btree_v2_chunk_tree(index_addr, &header, &records)
}
fn btree_v2_record_starts_inside_extent(
record: &[u8],
dims: &[u64],
chunk_dims: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<bool> {
let scaled_coords = Self::decode_btree_v2_scaled_coords(
record,
filtered,
chunk_size_len,
sizeof_addr,
chunk_dims.len(),
)?;
for ((&scaled, &chunk), &dim) in scaled_coords.iter().zip(chunk_dims).zip(dims) {
let start = scaled
.checked_mul(chunk)
.ok_or_else(|| Error::InvalidFormat("chunk start coordinate overflow".into()))?;
if start >= dim {
return Ok(false);
}
}
Ok(true)
}
fn chunk_unfiltered_byte_len(
info: &crate::hl::dataset::DatasetInfo,
chunk_dims: &[u64],
) -> Result<usize> {
let element_size = Self::u64_to_usize(u64::from(info.datatype.size), "datatype size")?;
chunk_dims.iter().try_fold(element_size, |acc, &dim| {
let dim = Self::u64_to_usize(dim, "chunk dimension")?;
acc.checked_mul(dim)
.ok_or_else(|| Error::InvalidFormat("chunk byte size overflow".into()))
})
}
fn find_btree_v2_chunk_record_index(
records: &[Vec<u8>],
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<std::result::Result<usize, usize>> {
let mut left = 0usize;
let mut right = records.len();
while left < right {
let mid = left + (right - left) / 2;
match Self::cmp_btree_v2_scaled_coords(
&records[mid],
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
)? {
std::cmp::Ordering::Less => left = mid + 1,
std::cmp::Ordering::Equal => return Ok(Ok(mid)),
std::cmp::Ordering::Greater => right = mid,
}
}
Ok(Err(left))
}
fn rebuild_btree_v2_chunk_tree(
&mut self,
header_addr: u64,
header: &BTreeV2Header,
records: &[Vec<u8>],
) -> Result<()> {
let sa = usize::from(self.superblock.sizeof_addr);
let new_depth = self.btree_v2_required_depth(header, records.len(), sa)?;
let node_info = Self::btree_v2_mutable_node_info_for_depth(header, sa, new_depth)?;
let root = self.append_btree_v2_subtree(header, &node_info, new_depth, records)?;
self.rewrite_btree_v2_header_root(
header_addr,
header,
new_depth,
root.addr,
root.nrecords,
root.total_records,
)
}
fn append_btree_v2_subtree(
&mut self,
header: &BTreeV2Header,
node_info: &[BTreeV2MutableNodeInfo],
depth: u16,
records: &[Vec<u8>],
) -> Result<BuiltBTreeV2Node> {
if depth == 0 {
let leaf_capacity = node_info
.first()
.map(|info| info.max_nrec)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree node info is empty".into()))?;
if records.len() > leaf_capacity {
return Err(Error::InvalidFormat(
"v2 B-tree leaf records exceed node capacity".into(),
));
}
let addr = self.append_btree_v2_leaf(header, records)?;
return Ok(BuiltBTreeV2Node {
addr,
nrecords: Self::usize_to_u16(records.len(), "v2 B-tree leaf record count")?,
total_records: Self::usize_to_u64(records.len(), "v2 B-tree total record count")?,
});
}
let depth_index = usize::from(depth);
let info = node_info.get(depth_index).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal node depth is invalid".into())
})?;
let child_info = node_info
.get(depth_index - 1)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child node depth is invalid".into()))?;
let child_count = Self::btree_v2_child_count_for_records(records.len(), depth, node_info)?;
let separator_count = child_count
.checked_sub(1)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child count underflow".into()))?;
if separator_count > info.max_nrec {
return Err(Error::InvalidFormat(
"v2 B-tree internal node has too many records".into(),
));
}
let min_child_records = Self::btree_v2_min_records_for_depth(depth - 1)?;
let child_capacity = Self::u64_to_usize(
child_info
.cum_max_nrec
.min(Self::usize_to_u64(records.len(), "v2 B-tree records")?),
"v2 B-tree child capacity",
)?;
let mut record_pos = 0usize;
let mut remaining_child_records =
records.len().checked_sub(separator_count).ok_or_else(|| {
Error::InvalidFormat("invalid v2 B-tree chunk record distribution".into())
})?;
let mut children = Vec::with_capacity(child_count);
let mut separators = Vec::with_capacity(separator_count);
for child_index in 0..child_count {
let remaining_children = child_count - child_index;
let min_for_rest = remaining_children
.checked_sub(1)
.and_then(|count| count.checked_mul(min_child_records))
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree minimum record count overflow".into())
})?;
let max_for_current = child_capacity.min(
remaining_child_records
.checked_sub(min_for_rest)
.ok_or_else(|| {
Error::InvalidFormat("invalid v2 B-tree chunk record distribution".into())
})?,
);
let min_for_current = min_child_records.max(
remaining_child_records.saturating_sub(
remaining_children
.checked_sub(1)
.and_then(|count| count.checked_mul(child_capacity))
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree child capacity overflow".into())
})?,
),
);
let take = remaining_child_records
.div_ceil(remaining_children)
.clamp(min_for_current, max_for_current);
if take == 0 || take > child_capacity {
return Err(Error::InvalidFormat(
"invalid v2 B-tree chunk record distribution".into(),
));
}
let child_end = record_pos.checked_add(take).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree rebuild record offset overflow".into())
})?;
let child_records = records.get(record_pos..child_end).ok_or_else(|| {
Error::InvalidFormat("invalid v2 B-tree chunk record distribution".into())
})?;
children.push(self.append_btree_v2_subtree(
header,
node_info,
depth - 1,
child_records,
)?);
record_pos = child_end;
remaining_child_records =
remaining_child_records.checked_sub(take).ok_or_else(|| {
Error::InvalidFormat("invalid v2 B-tree chunk record distribution".into())
})?;
if child_index + 1 < child_count {
let separator = records.get(record_pos).ok_or_else(|| {
Error::InvalidFormat("invalid v2 B-tree chunk record distribution".into())
})?;
separators.push(separator.as_slice());
record_pos = record_pos.checked_add(1).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree rebuild record offset overflow".into())
})?;
}
}
if record_pos != records.len() {
return Err(Error::InvalidFormat(
"v2 B-tree chunk rebuild did not consume all records".into(),
));
}
let addr =
self.append_btree_v2_internal(header, node_info, depth, &separators, &children)?;
Ok(BuiltBTreeV2Node {
addr,
nrecords: Self::usize_to_u16(separators.len(), "v2 B-tree internal record count")?,
total_records: Self::usize_to_u64(records.len(), "v2 B-tree total record count")?,
})
}
fn encode_btree_v2_leaf(header: &BTreeV2Header, records: &[Vec<u8>]) -> Result<Vec<u8>> {
if records.len() > usize::from(u16::MAX) {
return Err(Error::Unsupported(
"v2 B-tree leaf record count exceeds u16".into(),
));
}
let record_size = usize::from(header.record_size);
let records_bytes = records
.len()
.checked_mul(record_size)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree leaf size overflow".into()))?;
let node_size = Self::btree_v2_node_size(header)?;
let compact_len = 6usize
.checked_add(records_bytes)
.and_then(|value| value.checked_add(4))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree leaf size overflow".into()))?;
if compact_len > node_size {
return Err(Error::InvalidFormat(
"v2 B-tree leaf records exceed node size".into(),
));
}
let mut leaf = Vec::with_capacity(node_size);
leaf.extend_from_slice(b"BTLF");
leaf.push(0);
leaf.push(header.tree_type);
for record in records {
if record.len() != record_size {
return Err(Error::InvalidFormat(
"v2 B-tree leaf record has wrong size".into(),
));
}
leaf.extend_from_slice(record);
}
let checksum = checksum_metadata(&leaf);
leaf.extend_from_slice(&checksum.to_le_bytes());
leaf.resize(node_size, 0);
Ok(leaf)
}
fn append_btree_v2_leaf(&mut self, header: &BTreeV2Header, records: &[Vec<u8>]) -> Result<u64> {
let leaf = Self::encode_btree_v2_leaf(header, records)?;
let addr = self.append_aligned_zeros(leaf.len(), 8)?;
self.write_handle.seek(SeekFrom::Start(addr))?;
self.write_handle.write_all(&leaf)?;
Ok(addr)
}
fn encode_btree_v2_internal(
&self,
header: &BTreeV2Header,
node_info: &[BTreeV2MutableNodeInfo],
depth: u16,
separators: &[&[u8]],
children: &[BuiltBTreeV2Node],
) -> Result<Vec<u8>> {
if children.len() != separators.len() + 1 {
return Err(Error::InvalidFormat(
"v2 B-tree internal child/record count mismatch".into(),
));
}
if separators.len() > usize::from(u16::MAX) {
return Err(Error::Unsupported(
"v2 B-tree internal record count exceeds u16".into(),
));
}
let depth_index = usize::from(depth);
if depth == 0 || depth_index >= node_info.len() {
return Err(Error::InvalidFormat(
"v2 B-tree internal node depth is invalid".into(),
));
}
if separators.len() > node_info[depth_index].max_nrec {
return Err(Error::InvalidFormat(
"v2 B-tree internal node has too many records".into(),
));
}
let child_nrecords_size = Self::bytes_needed(Self::usize_to_u64(
node_info[0].max_nrec,
"v2 B-tree leaf record capacity",
)?);
let child_all_nrecords_size = if depth > 1 {
node_info[depth_index - 1].cum_max_nrec_size
} else {
0
};
let sa = usize::from(self.superblock.sizeof_addr);
let record_size = usize::from(header.record_size);
let pointer_size = sa
.checked_add(child_nrecords_size)
.and_then(|value| value.checked_add(child_all_nrecords_size))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree pointer size overflow".into()))?;
let node_size = Self::btree_v2_node_size(header)?;
let compact_len =
6usize
.checked_add(separators.len().checked_mul(record_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal size overflow".into())
})?)
.and_then(|value| {
children
.len()
.checked_mul(pointer_size)
.and_then(|children_bytes| value.checked_add(children_bytes))
})
.and_then(|value| value.checked_add(4))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree internal size overflow".into()))?;
if compact_len > node_size {
return Err(Error::InvalidFormat(
"v2 B-tree internal records exceed node size".into(),
));
}
let mut node = Vec::with_capacity(node_size);
node.extend_from_slice(b"BTIN");
node.push(0);
node.push(header.tree_type);
for record in separators {
if record.len() != record_size {
return Err(Error::InvalidFormat(
"v2 B-tree internal separator has wrong size".into(),
));
}
node.extend_from_slice(record);
}
let mut encoded = [0u8; 8];
let child_max_node_records = node_info[depth_index - 1].max_nrec;
let child_max_total_records = node_info[depth_index - 1].cum_max_nrec;
for child in children {
if usize::from(child.nrecords) > child_max_node_records
|| child.total_records > child_max_total_records
|| child.total_records < u64::from(child.nrecords)
{
return Err(Error::InvalidFormat(
"v2 B-tree internal child record counts are inconsistent".into(),
));
}
node.extend_from_slice(Self::encode_uint_le_into(
child.addr,
&mut encoded,
sa,
"v2 B-tree child address",
)?);
node.extend_from_slice(Self::encode_uint_le_into(
u64::from(child.nrecords),
&mut encoded,
child_nrecords_size,
"v2 B-tree child record count",
)?);
if child_all_nrecords_size > 0 {
node.extend_from_slice(Self::encode_uint_le_into(
child.total_records,
&mut encoded,
child_all_nrecords_size,
"v2 B-tree child total record count",
)?);
}
}
let checksum = checksum_metadata(&node);
node.extend_from_slice(&checksum.to_le_bytes());
node.resize(node_size, 0);
Ok(node)
}
fn append_btree_v2_internal(
&mut self,
header: &BTreeV2Header,
node_info: &[BTreeV2MutableNodeInfo],
depth: u16,
separators: &[&[u8]],
children: &[BuiltBTreeV2Node],
) -> Result<u64> {
let node = self.encode_btree_v2_internal(header, node_info, depth, separators, children)?;
let addr = self.append_aligned_zeros(node.len(), 8)?;
self.write_handle.seek(SeekFrom::Start(addr))?;
self.write_handle.write_all(&node)?;
Ok(addr)
}
fn btree_v2_required_depth(
&self,
header: &BTreeV2Header,
record_count: usize,
sizeof_addr: usize,
) -> Result<u16> {
let record_count = Self::usize_to_u64(record_count, "v2 B-tree record count")?;
for depth in 0..=u16::MAX {
let node_info = Self::btree_v2_mutable_node_info_for_depth(header, sizeof_addr, depth)?;
if record_count <= node_info[usize::from(depth)].cum_max_nrec {
return Ok(depth);
}
}
Err(Error::Unsupported(
"v2 B-tree chunk index is too large to rebuild".into(),
))
}
fn btree_v2_child_count_for_records(
record_count: usize,
depth: u16,
node_info: &[BTreeV2MutableNodeInfo],
) -> Result<usize> {
if depth == 0 {
return Ok(0);
}
let depth_index = usize::from(depth);
let info = node_info.get(depth_index).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal node depth is invalid".into())
})?;
let child_capacity = node_info
.get(depth_index - 1)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child node depth is invalid".into()))?
.cum_max_nrec;
let child_capacity = Self::u64_to_usize(
child_capacity.min(Self::usize_to_u64(record_count, "v2 B-tree records")?),
"v2 B-tree child capacity",
)?;
let min_child_records = Self::btree_v2_min_records_for_depth(depth - 1)?;
let max_children = info
.max_nrec
.checked_add(1)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child count overflow".into()))?;
for child_count in 2..=max_children {
let separators = child_count - 1;
let Some(child_records) = record_count.checked_sub(separators) else {
continue;
};
let min_total = child_count
.checked_mul(min_child_records)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child count overflow".into()))?;
let max_total = child_count
.checked_mul(child_capacity)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child count overflow".into()))?;
if child_records >= min_total && child_records <= max_total {
return Ok(child_count);
}
}
Err(Error::InvalidFormat(
"invalid v2 B-tree chunk record distribution".into(),
))
}
fn btree_v2_min_records_for_depth(depth: u16) -> Result<usize> {
let mut min_records = 1usize;
for _ in 0..depth {
min_records = min_records
.checked_mul(2)
.and_then(|value| value.checked_add(1))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree minimum record overflow".into()))?;
}
Ok(min_records)
}
fn bytes_needed(mut value: u64) -> usize {
let mut bytes = 1usize;
while value > 0xff {
value >>= 8;
bytes += 1;
}
bytes
}
fn rewrite_btree_v2_header_root(
&mut self,
header_addr: u64,
_header: &BTreeV2Header,
new_depth: u16,
new_root_addr: u64,
new_root_nrecords: u16,
new_total_records: u64,
) -> Result<()> {
let sa = usize::from(self.superblock.sizeof_addr);
let ss = usize::from(self.superblock.sizeof_size);
let depth_pos = header_addr
.checked_add(u64::from(4u8 + 1 + 1 + 4 + 2))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree header offset overflow".into()))?;
let root_addr_pos = header_addr
.checked_add(u64::from(4u8 + 1 + 1 + 4 + 2 + 2 + 1 + 1))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree header offset overflow".into()))?;
let root_nrecords_pos = root_addr_pos
.checked_add(Self::usize_to_u64(sa, "v2 B-tree address width")?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree header offset overflow".into()))?;
let total_records_pos = root_nrecords_pos
.checked_add(2)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree header offset overflow".into()))?;
let checksum_pos = total_records_pos
.checked_add(Self::usize_to_u64(ss, "v2 B-tree length width")?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree header offset overflow".into()))?;
self.write_handle.seek(SeekFrom::Start(depth_pos))?;
self.write_handle.write_all(&new_depth.to_le_bytes())?;
self.write_handle.seek(SeekFrom::Start(root_addr_pos))?;
let mut encoded = [0u8; 8];
self.write_handle.write_all(Self::encode_uint_le_into(
new_root_addr,
&mut encoded,
sa,
"v2 B-tree root address",
)?)?;
self.write_handle.seek(SeekFrom::Start(root_nrecords_pos))?;
self.write_handle
.write_all(&new_root_nrecords.to_le_bytes())?;
self.write_handle.seek(SeekFrom::Start(total_records_pos))?;
self.write_uint_le(new_total_records, ss)?;
let check_len = usize::try_from(checksum_pos - header_addr).map_err(|_| {
Error::InvalidFormat("v2 B-tree header checksum span is too large".into())
})?;
let mut guard = self.inner.lock();
guard.reader.seek(header_addr)?;
let mut header_bytes = vec![0u8; check_len];
guard.reader.read_bytes_into(&mut header_bytes)?;
drop(guard);
let depth_offset = Self::checked_header_relative_offset(depth_pos, header_addr)?;
let root_addr_offset = Self::checked_header_relative_offset(root_addr_pos, header_addr)?;
let root_nrecords_offset =
Self::checked_header_relative_offset(root_nrecords_pos, header_addr)?;
let total_offset = Self::checked_header_relative_offset(total_records_pos, header_addr)?;
Self::header_window_mut(&mut header_bytes, depth_offset, 2)?
.copy_from_slice(&new_depth.to_le_bytes());
Self::header_window_mut(&mut header_bytes, root_addr_offset, sa)?.copy_from_slice(
Self::encode_uint_le_into(new_root_addr, &mut encoded, sa, "v2 B-tree root address")?,
);
Self::header_window_mut(&mut header_bytes, root_nrecords_offset, 2)?
.copy_from_slice(&new_root_nrecords.to_le_bytes());
Self::header_window_mut(&mut header_bytes, total_offset, ss)?.copy_from_slice(
Self::encode_uint_le_into(
new_total_records,
&mut encoded,
ss,
"v2 B-tree total record count",
)?,
);
let checksum = checksum_metadata(&header_bytes);
self.write_handle.seek(SeekFrom::Start(checksum_pos))?;
self.write_handle.write_all(&checksum.to_le_bytes())?;
Ok(())
}
fn btree_v2_node_size(header: &BTreeV2Header) -> Result<usize> {
usize::try_from(header.node_size)
.map_err(|_| Error::InvalidFormat("v2 B-tree node size is too large".into()))
}
fn find_btree_v2_chunk_record_location<R: Read + Seek>(
reader: &mut HdfReader<R>,
header: &BTreeV2Header,
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<Option<BTreeV2RecordLocation>> {
if header.total_records == 0 {
return Ok(None);
}
let node_info = Self::btree_v2_mutable_node_info(header, sizeof_addr)?;
Self::find_btree_v2_chunk_record_location_in_node(
reader,
header,
&node_info,
header.root_addr,
header.root_nrecords,
header.depth,
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
)
}
#[allow(clippy::too_many_arguments)]
fn find_btree_v2_chunk_record_location_in_node<R: Read + Seek>(
reader: &mut HdfReader<R>,
header: &BTreeV2Header,
node_info: &[BTreeV2MutableNodeInfo],
node_addr: u64,
nrecords: u16,
depth: u16,
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<Option<BTreeV2RecordLocation>> {
if depth == 0 {
return Self::find_btree_v2_chunk_record_location_in_leaf(
reader,
header,
node_addr,
nrecords,
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
);
}
let depth_index = usize::from(depth);
let Some(info) = node_info.get(depth_index) else {
return Err(Error::InvalidFormat(
"v2 B-tree internal node depth is invalid".into(),
));
};
if usize::from(nrecords) > info.max_nrec {
return Err(Error::InvalidFormat(
"v2 B-tree internal node has too many records".into(),
));
}
let record_size = usize::from(header.record_size);
let nrecords_usize = usize::from(nrecords);
let record_bytes = nrecords_usize.checked_mul(record_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal record bytes overflow".into())
})?;
let leaf_capacity = node_info
.first()
.map(|info| info.max_nrec)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree node info is empty".into()))?;
let child_nrec_size = Self::bytes_needed(Self::usize_to_u64(
leaf_capacity,
"v2 B-tree leaf record capacity",
)?);
let child_all_nrec_size = if depth > 1 {
node_info[depth_index - 1].cum_max_nrec_size
} else {
0
};
let child_entry_size = sizeof_addr
.checked_add(child_nrec_size)
.and_then(|value| value.checked_add(child_all_nrec_size))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child entry size overflow".into()))?;
let child_count = nrecords_usize
.checked_add(1)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child count overflow".into()))?;
let child_bytes = child_count
.checked_mul(child_entry_size)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child bytes overflow".into()))?;
let checksum_len = 6usize
.checked_add(record_bytes)
.and_then(|value| value.checked_add(child_bytes))
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal checksum span overflow".into())
})?;
let mut node = vec![0; checksum_len + 4];
reader.seek(node_addr)?;
reader.read_bytes_into(&mut node)?;
if node.get(..4) != Some(b"BTIN") {
return Err(Error::InvalidFormat(
"invalid v2 B-tree internal magic".into(),
));
}
if node.get(4) != Some(&0) || node.get(5) != Some(&header.tree_type) {
return Err(Error::InvalidFormat(
"invalid v2 B-tree internal node prefix".into(),
));
}
let records_start = 6usize;
let mut left = 0usize;
let mut right = nrecords_usize;
while left < right {
let mid = left + (right - left) / 2;
let record_offset = records_start
.checked_add(mid.checked_mul(record_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree record offset overflow".into())
})?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
let record = node
.get(record_offset..record_offset + record_size)
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal record is truncated".into())
})?;
match Self::cmp_btree_v2_scaled_coords(
record,
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
)? {
std::cmp::Ordering::Less => left = mid + 1,
std::cmp::Ordering::Equal => {
return Ok(Some(BTreeV2RecordLocation {
node_addr,
record_offset,
checksum_len,
}));
}
std::cmp::Ordering::Greater => right = mid,
}
}
let child_table_start = records_start
.checked_add(record_bytes)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child table offset overflow".into()))?;
let child_offset = child_table_start
.checked_add(left.checked_mul(child_entry_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree child entry offset overflow".into())
})?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child entry offset overflow".into()))?;
let child_addr = Self::read_le_uint(
node.get(child_offset..child_offset + sizeof_addr)
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree child address is truncated".into())
})?,
)?;
let child_nrecords_start = child_offset
.checked_add(sizeof_addr)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree child entry offset overflow".into()))?;
let child_nrecords = Self::read_le_uint(
node.get(child_nrecords_start..child_nrecords_start + child_nrec_size)
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree child record count is truncated".into())
})?,
)?;
let child_nrecords = u16::try_from(child_nrecords)
.map_err(|_| Error::InvalidFormat("v2 B-tree child record count exceeds u16".into()))?;
Self::find_btree_v2_chunk_record_location_in_node(
reader,
header,
node_info,
child_addr,
child_nrecords,
depth - 1,
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
)
}
#[allow(clippy::too_many_arguments)]
fn find_btree_v2_chunk_record_location_in_leaf<R: Read + Seek>(
reader: &mut HdfReader<R>,
header: &BTreeV2Header,
node_addr: u64,
nrecords: u16,
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<Option<BTreeV2RecordLocation>> {
let record_size = usize::from(header.record_size);
let nrecords_usize = usize::from(nrecords);
let record_bytes = nrecords_usize
.checked_mul(record_size)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree leaf record bytes overflow".into()))?;
let checksum_len = 6usize
.checked_add(record_bytes)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree leaf checksum span overflow".into()))?;
let mut node = vec![0; checksum_len + 4];
reader.seek(node_addr)?;
reader.read_bytes_into(&mut node)?;
if node.get(..4) != Some(b"BTLF") {
return Err(Error::InvalidFormat("invalid v2 B-tree leaf magic".into()));
}
if node.get(4) != Some(&0) || node.get(5) != Some(&header.tree_type) {
return Err(Error::InvalidFormat(
"invalid v2 B-tree leaf node prefix".into(),
));
}
let records_start = 6usize;
let mut left = 0usize;
let mut right = nrecords_usize;
while left < right {
let mid = left + (right - left) / 2;
let record_offset = records_start
.checked_add(mid.checked_mul(record_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree record offset overflow".into())
})?)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
let record = node
.get(record_offset..record_offset + record_size)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree leaf record is truncated".into()))?;
match Self::cmp_btree_v2_scaled_coords(
record,
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
)? {
std::cmp::Ordering::Less => left = mid + 1,
std::cmp::Ordering::Equal => {
return Ok(Some(BTreeV2RecordLocation {
node_addr,
record_offset,
checksum_len,
}));
}
std::cmp::Ordering::Greater => right = mid,
}
}
Ok(None)
}
fn btree_v2_mutable_node_info(
header: &BTreeV2Header,
sizeof_addr: usize,
) -> Result<Vec<BTreeV2MutableNodeInfo>> {
Self::btree_v2_mutable_node_info_for_depth(header, sizeof_addr, header.depth)
}
fn btree_v2_mutable_node_info_for_depth(
header: &BTreeV2Header,
sizeof_addr: usize,
depth: u16,
) -> Result<Vec<BTreeV2MutableNodeInfo>> {
let node_size = Self::btree_v2_node_size(header)?;
let record_size = usize::from(header.record_size);
if node_size <= 10 || record_size == 0 {
return Err(Error::InvalidFormat("invalid v2 B-tree node sizing".into()));
}
let leaf_max = (node_size - 10) / record_size;
if leaf_max == 0 {
return Err(Error::InvalidFormat(
"v2 B-tree leaf cannot hold any records".into(),
));
}
let leaf_max_u64 = Self::usize_to_u64(leaf_max, "v2 B-tree leaf record capacity")?;
let max_nrec_size = Self::bytes_needed(leaf_max_u64);
let depth = usize::from(depth);
let mut infos = Vec::with_capacity(depth + 1);
infos.push(BTreeV2MutableNodeInfo {
max_nrec: leaf_max,
cum_max_nrec: leaf_max_u64,
cum_max_nrec_size: 0,
});
for level in 1..=depth {
let pointer_size = sizeof_addr
.checked_add(max_nrec_size)
.and_then(|value| value.checked_add(infos[level - 1].cum_max_nrec_size))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree pointer size overflow".into()))?;
let prefix_and_pointer = 6usize.checked_add(pointer_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal node prefix overflow".into())
})?;
if node_size <= prefix_and_pointer {
return Err(Error::InvalidFormat(
"v2 B-tree internal node cannot hold records".into(),
));
}
let record_slot = record_size.checked_add(pointer_size).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree internal record slot size overflow".into())
})?;
let max_nrec = (node_size - prefix_and_pointer) / record_slot;
if max_nrec == 0 {
return Err(Error::InvalidFormat(
"v2 B-tree internal node cannot hold records".into(),
));
}
let max_nrec_u64 = Self::usize_to_u64(max_nrec, "v2 B-tree internal record capacity")?;
let cum_max_nrec = max_nrec_u64
.checked_add(1)
.and_then(|value| value.checked_mul(infos[level - 1].cum_max_nrec))
.and_then(|value| value.checked_add(max_nrec_u64))
.ok_or_else(|| {
Error::InvalidFormat("v2 B-tree cumulative record count overflow".into())
})?;
infos.push(BTreeV2MutableNodeInfo {
max_nrec,
cum_max_nrec,
cum_max_nrec_size: Self::bytes_needed(cum_max_nrec),
});
}
Ok(infos)
}
fn btree_v2_chunk_record_size(
sizeof_addr: usize,
filtered: bool,
chunk_size_len: usize,
ndims: usize,
) -> Result<usize> {
let filter_bytes = if filtered {
chunk_size_len
.checked_add(4)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record size overflow".into()))?
} else {
0
};
let coord_bytes = ndims
.checked_mul(8)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record size overflow".into()))?;
sizeof_addr
.checked_add(filter_bytes)
.and_then(|value| value.checked_add(coord_bytes))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record size overflow".into()))
}
fn encode_btree_v2_chunk_record(
addr: u64,
chunk_size: u64,
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<Vec<u8>> {
let expected_record_size = Self::btree_v2_chunk_record_size(
sizeof_addr,
filtered,
chunk_size_len,
scaled_coords.len(),
)?;
let mut record = Vec::with_capacity(expected_record_size);
Self::encode_btree_v2_chunk_record_into(
addr,
chunk_size,
scaled_coords,
filtered,
chunk_size_len,
sizeof_addr,
&mut record,
)?;
Ok(record)
}
fn encode_btree_v2_chunk_record_into(
addr: u64,
chunk_size: u64,
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
record: &mut Vec<u8>,
) -> Result<()> {
let expected_record_size = Self::btree_v2_chunk_record_size(
sizeof_addr,
filtered,
chunk_size_len,
scaled_coords.len(),
)?;
record.clear();
if record.capacity() < expected_record_size {
record.reserve_exact(expected_record_size - record.capacity());
}
let mut encoded = [0u8; 8];
record.extend_from_slice(Self::encode_uint_le_into(
addr,
&mut encoded,
sizeof_addr,
"v2 B-tree chunk address",
)?);
if filtered {
record.extend_from_slice(Self::encode_uint_le_into(
chunk_size,
&mut encoded,
chunk_size_len,
"v2 B-tree chunk size",
)?);
record.extend_from_slice(&0u32.to_le_bytes());
}
for &coord in scaled_coords {
record.extend_from_slice(&coord.to_le_bytes());
}
Ok(())
}
fn decode_btree_v2_scaled_coords(
record: &[u8],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
ndims: usize,
) -> Result<Vec<u64>> {
let mut pos = sizeof_addr;
if record.len() < pos {
return Err(Error::InvalidFormat(
"truncated v2 B-tree chunk address".into(),
));
}
if filtered {
pos = pos
.checked_add(chunk_size_len)
.and_then(|value| value.checked_add(4))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
}
let coords_end = pos
.checked_add(
ndims.checked_mul(8).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree record offset overflow".into())
})?,
)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
if record.len() < coords_end {
return Err(Error::InvalidFormat(
"truncated v2 B-tree scaled chunk coordinates".into(),
));
}
let mut coords = Vec::with_capacity(ndims);
for _ in 0..ndims {
let next = pos
.checked_add(8)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
coords.push(Self::read_le_uint(&record[pos..next])?);
pos = next;
}
Ok(coords)
}
fn cmp_btree_v2_scaled_coords(
record: &[u8],
scaled_coords: &[u64],
filtered: bool,
chunk_size_len: usize,
sizeof_addr: usize,
) -> Result<std::cmp::Ordering> {
let mut pos = sizeof_addr;
if record.len() < pos {
return Err(Error::InvalidFormat(
"truncated v2 B-tree chunk address".into(),
));
}
if filtered {
pos = pos
.checked_add(chunk_size_len)
.and_then(|value| value.checked_add(4))
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
}
let coords_end = pos
.checked_add(
scaled_coords.len().checked_mul(8).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree record offset overflow".into())
})?,
)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
if record.len() < coords_end {
return Err(Error::InvalidFormat(
"truncated v2 B-tree scaled chunk coordinates".into(),
));
}
for &scaled_coord in scaled_coords {
let next = pos
.checked_add(8)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree record offset overflow".into()))?;
match Self::read_le_uint(&record[pos..next])?.cmp(&scaled_coord) {
std::cmp::Ordering::Equal => pos = next,
ordering => return Ok(ordering),
}
}
Ok(std::cmp::Ordering::Equal)
}
fn scaled_chunk_coords(chunk_coords: &[u64], chunk_dims: &[u64]) -> Result<Vec<u64>> {
if chunk_coords.len() != chunk_dims.len() {
return Err(Error::InvalidFormat(
"chunk coordinate rank does not match chunk dimensions".into(),
));
}
chunk_coords
.iter()
.zip(chunk_dims)
.map(|(&coord, &dim)| {
if dim == 0 {
return Err(Error::InvalidFormat("chunk dimension is zero".into()));
}
if coord % dim != 0 {
return Err(Error::InvalidFormat(
"chunk coordinate is not aligned to chunk dimension".into(),
));
}
Ok(coord / dim)
})
.collect()
}
fn read_le_uint(bytes: &[u8]) -> Result<u64> {
if bytes.len() > 8 {
return Err(Error::InvalidFormat(
"little-endian integer is wider than u64".into(),
));
}
let mut value = 0u64;
for (idx, byte) in bytes.iter().enumerate() {
value |= u64::from(*byte) << (idx * 8);
}
Ok(value)
}
fn checked_header_relative_offset(pos: u64, base: u64) -> Result<usize> {
let delta = pos.checked_sub(base).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree header field precedes header base".into())
})?;
usize::try_from(delta)
.map_err(|_| Error::InvalidFormat("v2 B-tree header field offset is too large".into()))
}
fn header_window_mut(buf: &mut [u8], offset: usize, len: usize) -> Result<&mut [u8]> {
let end = offset
.checked_add(len)
.ok_or_else(|| Error::InvalidFormat("v2 B-tree header field offset overflow".into()))?;
buf.get_mut(offset..end).ok_or_else(|| {
Error::InvalidFormat("v2 B-tree header field exceeds checksum span".into())
})
}
}
#[cfg(test)]
mod tests {
use super::MutableFile;
use crate::format::btree_v2::BTreeV2Header;
fn test_header(record_size: u16) -> BTreeV2Header {
BTreeV2Header {
tree_type: 10,
node_size: 512,
record_size,
depth: 0,
split_pct: 100,
merge_pct: 40,
root_addr: 0,
root_nrecords: 0,
total_records: 0,
}
}
#[test]
fn btree_v2_chunk_record_size_rejects_coord_overflow() {
let result = MutableFile::btree_v2_chunk_record_size(8, false, 0, usize::MAX / 8 + 1);
assert!(result.is_err());
}
#[test]
fn decode_btree_v2_scaled_coords_rejects_coordinate_span_overflow() {
let result =
MutableFile::decode_btree_v2_scaled_coords(&[0; 8], false, 0, 0, usize::MAX / 8 + 1);
assert!(result.is_err());
}
#[test]
fn decode_btree_v2_scaled_coords_rejects_truncated_coordinates() {
let result = MutableFile::decode_btree_v2_scaled_coords(&[0; 7], false, 0, 0, 1);
assert!(result.is_err());
}
#[test]
fn encode_btree_v2_leaf_rejects_wrong_record_size() {
let header = test_header(8);
let result = MutableFile::encode_btree_v2_leaf(&header, &[vec![0; 7]]);
assert!(result.is_err());
}
#[test]
fn encode_btree_v2_leaf_pads_to_declared_node_size() {
let header = test_header(8);
let leaf = MutableFile::encode_btree_v2_leaf(&header, &[vec![1; 8], vec![2; 8]]).unwrap();
assert_eq!(leaf.len(), header.node_size as usize);
assert!(leaf[6 + 2 * 8 + 4..].iter().all(|byte| *byte == 0));
}
#[test]
fn encode_btree_v2_chunk_record_into_reuses_existing_buffer() {
let scaled_coords = [2, 3];
let expected = MutableFile::encode_btree_v2_chunk_record(
0x0102_0304,
0x0506,
&scaled_coords,
true,
2,
4,
)
.unwrap();
let mut record = Vec::with_capacity(expected.len());
record.extend_from_slice(&[0xff; 3]);
let capacity = record.capacity();
MutableFile::encode_btree_v2_chunk_record_into(
0x0102_0304,
0x0506,
&scaled_coords,
true,
2,
4,
&mut record,
)
.unwrap();
assert_eq!(record, expected);
assert_eq!(record.capacity(), capacity);
}
#[test]
fn btree_v2_header_window_rejects_bad_offsets() {
let err = MutableFile::checked_header_relative_offset(9, 10).unwrap_err();
assert!(
err.to_string()
.contains("v2 B-tree header field precedes header base"),
"unexpected error: {err}"
);
let mut buf = vec![0u8; 4];
let err = MutableFile::header_window_mut(&mut buf, 3, 2).unwrap_err();
assert!(
err.to_string()
.contains("v2 B-tree header field exceeds checksum span"),
"unexpected error: {err}"
);
}
}