osobhandClaude Fable 5.1
6e84f31ed6
Dataspace and chunk dimensions are untrusted 64-bit fields, but the chunked read paths computed `num_elements() as usize * elem_size` and `chunk_dims.product() * elem_size` with plain arithmetic and fed the result to `vec![0u8; n]`. A crafted file could wrap the product (under-sizing the output buffer that chunks are then copied into) or request an allocation large enough to abort the process. - Dataspace::checked_num_elements, checked_byte_len, checked_chunk_byte_len and alloc_output (try_reserve_exact) replace the plain products and vec![0; n] at every chunked read site, plus the VDS and hyperslab paths. Overflow and allocation failure are FormatError::Overflow. - Dataspace::num_elements saturates instead of wrapping. - A zero-element dataset returns early, which also keeps the stride products in range when another dimension is huge. - parallel_read.rs: the three `c_addr + size > len` bounds checks used a raw add; they now use checked_add like the rest of the crate. Co-Authored-By: Claude Fable 5.1 <[email protected]>
2380 lines
80 KiB
Rust
2380 lines
80 KiB
Rust
//! Chunked dataset reading: B-tree v1 type 1 traversal and chunk assembly.
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
extern crate alloc;
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::{format, vec, vec::Vec};
|
|
|
|
use crate::chunk_cache::CacheAlignedBuffer;
|
|
#[cfg(feature = "std")]
|
|
use crate::chunk_cache::ChunkCache;
|
|
use crate::data_layout::DataLayout;
|
|
use crate::dataspace::Dataspace;
|
|
use crate::datatype::Datatype;
|
|
use crate::error::FormatError;
|
|
use crate::extensible_array::{ExtensibleArrayHeader, read_extensible_array_chunks};
|
|
use crate::filter_pipeline::FilterPipeline;
|
|
use crate::filters::decompress_chunk;
|
|
use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks};
|
|
#[cfg(feature = "std")]
|
|
use std::sync::Arc;
|
|
|
|
#[cfg(feature = "parallel")]
|
|
use crate::parallel_read;
|
|
|
|
#[cfg(feature = "parallel")]
|
|
use crate::lane_partition::PartitionStats;
|
|
|
|
/// Decompress all chunks into cache-line-aligned buffers, using lane-partitioned
|
|
/// parallel decompression when the `parallel` feature is enabled and the chunk
|
|
/// count exceeds the threshold.
|
|
fn decompress_all_chunks(
|
|
file_data: &[u8],
|
|
chunks: &[ChunkInfo],
|
|
pipeline: Option<&FilterPipeline>,
|
|
chunk_total_bytes: usize,
|
|
element_size: u32,
|
|
) -> Result<Vec<CacheAlignedBuffer>, FormatError> {
|
|
#[cfg(feature = "parallel")]
|
|
{
|
|
if let Some(pl) = pipeline
|
|
&& parallel_read::should_use_parallel(chunks.len())
|
|
{
|
|
// Seed from the first chunk's address and count for determinism.
|
|
let seed = chunks.first().map(|c| c.address).unwrap_or(0) ^ (chunks.len() as u64);
|
|
let (data, _stats) = parallel_read::decompress_chunks_lane_partitioned(
|
|
file_data,
|
|
chunks,
|
|
pl,
|
|
chunk_total_bytes,
|
|
element_size,
|
|
seed,
|
|
None, // auto-detect lane count
|
|
)?;
|
|
return Ok(data.into_iter().map(CacheAlignedBuffer::from_vec).collect());
|
|
}
|
|
}
|
|
|
|
// Sequential fallback — allocate into aligned buffers
|
|
let mut result = Vec::with_capacity(chunks.len());
|
|
for chunk_info in chunks {
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
|
|
|
let decompressed = if let Some(pl) = pipeline {
|
|
if chunk_info.filter_mask == 0 {
|
|
decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)?
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
result.push(CacheAlignedBuffer::from_vec(decompressed));
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
/// Decompress all chunks with lane-partitioned parallelism and return
|
|
/// per-lane diagnostics.
|
|
///
|
|
/// This is the stats-returning variant for callers who want to inspect
|
|
/// the partition balance. Only available with the `parallel` feature.
|
|
#[cfg(feature = "parallel")]
|
|
pub fn decompress_all_chunks_with_stats(
|
|
file_data: &[u8],
|
|
chunks: &[ChunkInfo],
|
|
pipeline: &FilterPipeline,
|
|
chunk_total_bytes: usize,
|
|
element_size: u32,
|
|
seed: u64,
|
|
num_lanes: Option<usize>,
|
|
) -> Result<(Vec<Vec<u8>>, PartitionStats), FormatError> {
|
|
parallel_read::decompress_chunks_lane_partitioned(
|
|
file_data,
|
|
chunks,
|
|
pipeline,
|
|
chunk_total_bytes,
|
|
element_size,
|
|
seed,
|
|
num_lanes,
|
|
)
|
|
}
|
|
|
|
/// Information about a single chunk in a chunked dataset.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ChunkInfo {
|
|
/// Size of chunk data in the file (after compression).
|
|
pub chunk_size: u32,
|
|
/// Bitmask of filters that were NOT applied (0 = all applied).
|
|
pub filter_mask: u32,
|
|
/// N-dimensional offset of this chunk in dataset space.
|
|
pub offsets: Vec<u64>,
|
|
/// File address of the chunk data.
|
|
pub address: u64,
|
|
}
|
|
|
|
/// Checks that `[offset, offset + needed)` fits within `data`, guarding the
|
|
/// addition against `usize` overflow from a crafted near-`usize::MAX` offset.
|
|
fn ensure_len(data: &[u8], offset: usize, needed: usize) -> Result<(), FormatError> {
|
|
if offset
|
|
.checked_add(needed)
|
|
.is_none_or(|end| end > data.len())
|
|
{
|
|
return Err(FormatError::UnexpectedEof {
|
|
expected: offset.saturating_add(needed),
|
|
available: data.len(),
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// `elements * elem_size` for sizes that come from the file. Dataspace and
|
|
/// chunk dimensions are untrusted 64-bit fields, so a crafted file can make
|
|
/// the plain product wrap to a small number (or to something enormous).
|
|
pub(crate) fn checked_byte_len(elements: u64, elem_size: usize) -> Result<usize, FormatError> {
|
|
usize::try_from(elements)
|
|
.ok()
|
|
.and_then(|n| n.checked_mul(elem_size))
|
|
.ok_or_else(|| {
|
|
FormatError::Overflow(format!(
|
|
"{elements} elements of {elem_size} bytes exceeds the addressable size"
|
|
))
|
|
})
|
|
}
|
|
|
|
/// Product of chunk dimensions times the element size, overflow-checked.
|
|
pub(crate) fn checked_chunk_byte_len(
|
|
chunk_dims: &[usize],
|
|
elem_size: usize,
|
|
) -> Result<usize, FormatError> {
|
|
chunk_dims
|
|
.iter()
|
|
.try_fold(elem_size, |acc, &d| acc.checked_mul(d))
|
|
.ok_or_else(|| {
|
|
FormatError::Overflow(format!(
|
|
"chunk dimensions {chunk_dims:?} x {elem_size} bytes exceeds the addressable size"
|
|
))
|
|
})
|
|
}
|
|
|
|
/// A zero-filled output buffer of `len` bytes. `vec![0; len]` aborts the
|
|
/// process when the allocation fails; a size taken from the file must surface
|
|
/// as an error instead.
|
|
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
|
|
let mut out = Vec::new();
|
|
out.try_reserve_exact(len).map_err(|_| {
|
|
FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"))
|
|
})?;
|
|
out.resize(len, 0);
|
|
Ok(out)
|
|
}
|
|
|
|
fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
|
|
let s = size as usize;
|
|
if pos.checked_add(s).is_none_or(|end| end > data.len()) {
|
|
return Err(FormatError::UnexpectedEof {
|
|
expected: pos.saturating_add(s),
|
|
available: data.len(),
|
|
});
|
|
}
|
|
let slice = &data[pos..pos + s];
|
|
Ok(match size {
|
|
2 => u16::from_le_bytes([slice[0], slice[1]]) as u64,
|
|
4 => u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) as u64,
|
|
8 => u64::from_le_bytes([
|
|
slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
|
|
]),
|
|
_ => return Err(FormatError::InvalidOffsetSize(size)),
|
|
})
|
|
}
|
|
|
|
/// Traverse B-tree v1 type 1 to collect all chunk locations.
|
|
///
|
|
/// `ndims` is the number of offset dimensions in each key, which equals
|
|
/// `chunk_dimensions.len()` from the DataLayout::Chunked message (rank+1).
|
|
pub fn collect_chunk_info(
|
|
file_data: &[u8],
|
|
btree_address: u64,
|
|
ndims: usize,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
) -> Result<Vec<ChunkInfo>, FormatError> {
|
|
collect_chunk_info_inner(file_data, btree_address, ndims, offset_size, length_size, 0)
|
|
}
|
|
|
|
/// Maximum recursion depth for chunk B-tree traversal (malformed/cyclic data
|
|
/// protection), matching `btree_v1.rs`'s `MAX_BTREE_DEPTH`.
|
|
const MAX_CHUNK_BTREE_DEPTH: usize = 64;
|
|
|
|
fn collect_chunk_info_inner(
|
|
file_data: &[u8],
|
|
btree_address: u64,
|
|
ndims: usize,
|
|
offset_size: u8,
|
|
_length_size: u8,
|
|
depth: usize,
|
|
) -> Result<Vec<ChunkInfo>, FormatError> {
|
|
if depth > MAX_CHUNK_BTREE_DEPTH {
|
|
return Err(FormatError::NestingDepthExceeded);
|
|
}
|
|
|
|
let offset = btree_address as usize;
|
|
let os = offset_size as usize;
|
|
|
|
// Parse B-tree v1 header
|
|
let header_size = 8 + os * 2;
|
|
ensure_len(file_data, offset, header_size)?;
|
|
|
|
if &file_data[offset..offset + 4] != b"TREE" {
|
|
return Err(FormatError::InvalidBTreeSignature);
|
|
}
|
|
|
|
let node_type = file_data[offset + 4];
|
|
if node_type != 1 {
|
|
return Err(FormatError::InvalidBTreeNodeType(node_type));
|
|
}
|
|
|
|
let node_level = file_data[offset + 5];
|
|
let entries_used = u16::from_le_bytes([file_data[offset + 6], file_data[offset + 7]]) as usize;
|
|
|
|
let mut pos = offset + 8 + os * 2; // skip left/right sibling
|
|
|
|
// Key size: chunk_size(4) + filter_mask(4) + ndims * offset_size
|
|
let key_size = 4 + 4 + ndims * os;
|
|
|
|
if node_level == 0 {
|
|
// Leaf node: keys and children interleaved
|
|
// key[0], child[0], key[1], child[1], ..., key[N-1], child[N-1], key[N]
|
|
let needed = entries_used * (key_size + os) + key_size;
|
|
ensure_len(file_data, pos, needed)?;
|
|
|
|
let mut chunks = Vec::with_capacity(entries_used);
|
|
for _ in 0..entries_used {
|
|
// Parse key
|
|
let chunk_size = u32::from_le_bytes([
|
|
file_data[pos],
|
|
file_data[pos + 1],
|
|
file_data[pos + 2],
|
|
file_data[pos + 3],
|
|
]);
|
|
let filter_mask = u32::from_le_bytes([
|
|
file_data[pos + 4],
|
|
file_data[pos + 5],
|
|
file_data[pos + 6],
|
|
file_data[pos + 7],
|
|
]);
|
|
let mut offsets = Vec::with_capacity(ndims);
|
|
let mut kp = pos + 8;
|
|
for _ in 0..ndims {
|
|
offsets.push(read_offset(file_data, kp, offset_size)?);
|
|
kp += os;
|
|
}
|
|
pos += key_size;
|
|
|
|
// Parse child address
|
|
let address = read_offset(file_data, pos, offset_size)?;
|
|
pos += os;
|
|
|
|
chunks.push(ChunkInfo {
|
|
chunk_size,
|
|
filter_mask,
|
|
offsets,
|
|
address,
|
|
});
|
|
}
|
|
// Skip final key
|
|
Ok(chunks)
|
|
} else {
|
|
// Internal node: recurse into children
|
|
let needed = entries_used * (key_size + os) + key_size;
|
|
ensure_len(file_data, pos, needed)?;
|
|
|
|
let mut child_addrs = Vec::with_capacity(entries_used);
|
|
for _ in 0..entries_used {
|
|
pos += key_size; // skip key
|
|
let child_addr = read_offset(file_data, pos, offset_size)?;
|
|
child_addrs.push(child_addr);
|
|
pos += os;
|
|
}
|
|
|
|
let mut all_chunks = Vec::new();
|
|
for child_addr in child_addrs {
|
|
let child_chunks = collect_chunk_info_inner(
|
|
file_data,
|
|
child_addr,
|
|
ndims,
|
|
offset_size,
|
|
_length_size,
|
|
depth + 1,
|
|
)?;
|
|
all_chunks.extend(child_chunks);
|
|
}
|
|
Ok(all_chunks)
|
|
}
|
|
}
|
|
|
|
/// Generate ChunkInfo entries for an implicit index (v4 index type 2).
|
|
///
|
|
/// Chunks are stored contiguously starting at `base_address`. No stored index;
|
|
/// addresses are computed from the chunk position.
|
|
pub fn generate_implicit_chunks(
|
|
base_address: u64,
|
|
dataset_dims: &[u64],
|
|
chunk_dimensions: &[u32],
|
|
element_size: u32,
|
|
) -> Vec<ChunkInfo> {
|
|
let rank = chunk_dimensions.len();
|
|
let chunk_byte_size: u64 =
|
|
chunk_dimensions.iter().map(|&d| d as u64).product::<u64>() * element_size as u64;
|
|
|
|
let mut num_chunks_per_dim = Vec::with_capacity(rank);
|
|
for d in 0..rank {
|
|
let ds = dataset_dims[d];
|
|
let ch = chunk_dimensions[d] as u64;
|
|
num_chunks_per_dim.push(ds.div_ceil(ch));
|
|
}
|
|
let total_chunks: u64 = num_chunks_per_dim.iter().product();
|
|
|
|
let mut chunks = Vec::with_capacity(total_chunks as usize);
|
|
for linear_idx in 0..total_chunks {
|
|
let mut offsets = vec![0u64; rank];
|
|
let mut remaining = linear_idx;
|
|
for d in (0..rank).rev() {
|
|
let nchunks = num_chunks_per_dim[d];
|
|
let chunk_idx = remaining % nchunks;
|
|
remaining /= nchunks;
|
|
offsets[d] = chunk_idx * chunk_dimensions[d] as u64;
|
|
}
|
|
|
|
chunks.push(ChunkInfo {
|
|
chunk_size: chunk_byte_size as u32,
|
|
filter_mask: 0,
|
|
offsets,
|
|
address: base_address + linear_idx * chunk_byte_size,
|
|
});
|
|
}
|
|
|
|
chunks
|
|
}
|
|
|
|
/// Read a chunked dataset, decompressing chunks as needed.
|
|
pub fn read_chunked_data(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (
|
|
chunk_dimensions,
|
|
version,
|
|
chunk_index_type,
|
|
addr_opt,
|
|
single_filtered_size,
|
|
single_filter_mask,
|
|
) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
version,
|
|
chunk_index_type,
|
|
single_chunk_filtered_size,
|
|
single_chunk_filter_mask,
|
|
} => (
|
|
chunk_dimensions,
|
|
*version,
|
|
*chunk_index_type,
|
|
*btree_address,
|
|
*single_chunk_filtered_size,
|
|
*single_chunk_filter_mask,
|
|
),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
|
|
// Both v3 and v4 include element size as last dim (rank+1)
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// Collect chunks based on version and index type
|
|
let chunks = match (version, chunk_index_type) {
|
|
(3, _) => {
|
|
let ndims = chunk_dimensions.len(); // rank+1
|
|
collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?
|
|
}
|
|
(4, Some(1)) => {
|
|
// Single chunk — one chunk covering the entire dataset
|
|
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
|
(fs as u32, single_filter_mask.unwrap_or(0))
|
|
} else {
|
|
(chunk_byte_size as u32, 0)
|
|
};
|
|
vec![ChunkInfo {
|
|
chunk_size: csize,
|
|
filter_mask: fmask,
|
|
offsets: vec![0u64; rank],
|
|
address: addr,
|
|
}]
|
|
}
|
|
(4, Some(2)) => {
|
|
// Implicit index — use spatial chunk dims only
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
generate_implicit_chunks(
|
|
addr,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
)
|
|
}
|
|
(4, Some(3)) => {
|
|
// Fixed Array — use spatial chunk dims only
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_fixed_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(4, Some(4)) => {
|
|
// Extensible Array — use spatial chunk dims only
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
ExtensibleArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_extensible_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(v, idx) => {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"unsupported chunked layout version={v}, index_type={idx:?}"
|
|
)));
|
|
}
|
|
};
|
|
|
|
// Assemble output
|
|
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
|
|
if total_bytes == 0 {
|
|
// Also keeps the stride products below in range: with a zero-sized
|
|
// dimension the total is 0 even if other dimensions are huge.
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut output = alloc_output(total_bytes)?;
|
|
|
|
let mut ds_strides = vec![1usize; rank];
|
|
for i in (0..rank.saturating_sub(1)).rev() {
|
|
ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1];
|
|
}
|
|
|
|
let mut chunk_strides = vec![1usize; rank];
|
|
for i in (0..rank.saturating_sub(1)).rev() {
|
|
chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1];
|
|
}
|
|
|
|
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
|
|
// Fast path: no filters — copy directly from file_data without intermediate alloc
|
|
if pipeline.is_none() {
|
|
for chunk_info in &chunks {
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
let chunk_data = &file_data[c_addr..c_addr + size];
|
|
|
|
if rank == 0 {
|
|
let copy_len = chunk_data.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&chunk_data[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
chunk_data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
return Ok(output);
|
|
}
|
|
|
|
// Filtered path: decompress all chunks then assemble
|
|
let decompressed_chunks = decompress_all_chunks(
|
|
file_data,
|
|
&chunks,
|
|
pipeline,
|
|
chunk_total_bytes,
|
|
elem_size as u32,
|
|
)?;
|
|
|
|
for (chunk_info, decompressed) in chunks.iter().zip(decompressed_chunks.iter()) {
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
if rank == 0 {
|
|
let copy_len = decompressed.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
decompressed,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Read a chunked dataset with caching support.
|
|
///
|
|
/// On the first call, scans the chunk index (B-tree / fixed array / etc.) once
|
|
/// and populates the cache's hash index. Subsequent calls skip the index scan
|
|
/// entirely. Decompressed chunk data is also cached with LRU eviction.
|
|
#[cfg(feature = "std")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_chunked_data_cached(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
cache: &ChunkCache,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (
|
|
chunk_dimensions,
|
|
version,
|
|
chunk_index_type,
|
|
addr_opt,
|
|
single_filtered_size,
|
|
single_filter_mask,
|
|
) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
version,
|
|
chunk_index_type,
|
|
single_chunk_filtered_size,
|
|
single_chunk_filter_mask,
|
|
} => (
|
|
chunk_dimensions,
|
|
*version,
|
|
*chunk_index_type,
|
|
*btree_address,
|
|
*single_chunk_filtered_size,
|
|
*single_chunk_filter_mask,
|
|
),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// The per-file cache is shared across datasets; bind it to this one so a
|
|
// different dataset's chunk index is never reused for this read.
|
|
cache.ensure_dataset(addr);
|
|
|
|
// Populate chunk index on first access
|
|
if !cache.has_index() {
|
|
let chunks = match (version, chunk_index_type) {
|
|
(3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?,
|
|
(4, Some(1)) => {
|
|
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
|
(fs as u32, single_filter_mask.unwrap_or(0))
|
|
} else {
|
|
(chunk_byte_size as u32, 0)
|
|
};
|
|
vec![ChunkInfo {
|
|
chunk_size: csize,
|
|
filter_mask: fmask,
|
|
offsets: vec![0u64; rank],
|
|
address: addr,
|
|
}]
|
|
}
|
|
(4, Some(2)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
generate_implicit_chunks(
|
|
addr,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
)
|
|
}
|
|
(4, Some(3)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_fixed_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(4, Some(4)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header = ExtensibleArrayHeader::parse(
|
|
file_data,
|
|
addr as usize,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
read_extensible_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(v, idx) => {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"unsupported chunked layout version={v}, index_type={idx:?}"
|
|
)));
|
|
}
|
|
};
|
|
cache.populate_index(&chunks, rank);
|
|
}
|
|
|
|
let chunks = cache.all_indexed_chunks().unwrap_or_default();
|
|
|
|
// Assemble output
|
|
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
|
|
if total_bytes == 0 {
|
|
// Also keeps the stride products below in range: with a zero-sized
|
|
// dimension the total is 0 even if other dimensions are huge.
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut output = alloc_output(total_bytes)?;
|
|
|
|
let mut ds_strides = vec![1usize; rank];
|
|
for i in (0..rank.saturating_sub(1)).rev() {
|
|
ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1];
|
|
}
|
|
|
|
let mut chunk_strides = vec![1usize; rank];
|
|
for i in (0..rank.saturating_sub(1)).rev() {
|
|
chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1];
|
|
}
|
|
|
|
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
|
|
for chunk_info in &chunks {
|
|
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
|
|
|
// Try decompressed cache first
|
|
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
|
cached
|
|
} else {
|
|
// Decompress from file
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
|
let dec = if let Some(pl) = pipeline {
|
|
if chunk_info.filter_mask == 0 {
|
|
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
cache.put_decompressed(coord, dec)
|
|
};
|
|
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
if rank == 0 {
|
|
let copy_len = decompressed.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
&decompressed,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Sweep context passed into `read_chunked_data_sweep` to enable adaptive
|
|
/// prefetching based on detected access patterns.
|
|
///
|
|
/// The caller is responsible for maintaining the `SweepContext` across
|
|
/// multiple reads on the same dataset. After each read, the context will
|
|
/// contain updated sweep detection state and any predicted next-chunk
|
|
/// coordinates.
|
|
pub struct SweepContext {
|
|
/// Sliding window of recent chunk coordinates.
|
|
pub history: Vec<Vec<u64>>,
|
|
/// Maximum window size.
|
|
pub window_size: usize,
|
|
/// Currently detected sweep direction label.
|
|
pub direction: &'static str,
|
|
/// How many chunks ahead to predict.
|
|
pub prefetch_count: usize,
|
|
/// Predicted next chunk coordinates (populated after each read).
|
|
pub predicted_next: Vec<Vec<u64>>,
|
|
}
|
|
|
|
impl SweepContext {
|
|
/// Create a new sweep context with the given window size and prefetch count.
|
|
pub fn new(window_size: usize, prefetch_count: usize) -> Self {
|
|
Self {
|
|
history: Vec::with_capacity(window_size),
|
|
window_size,
|
|
direction: "random",
|
|
prefetch_count,
|
|
predicted_next: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Create with default settings (window=12, prefetch=4).
|
|
pub fn with_defaults() -> Self {
|
|
Self::new(12, 4)
|
|
}
|
|
|
|
/// Record a chunk coordinate access and update predictions.
|
|
fn record(&mut self, coord: Vec<u64>, ndims: usize) {
|
|
if self.history.len() >= self.window_size {
|
|
self.history.remove(0);
|
|
}
|
|
self.history.push(coord);
|
|
|
|
if self.history.len() < 3 || ndims == 0 {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
// Inline sweep detection matching the algorithm in clawhdf5-io/sweep.rs
|
|
let num_deltas = self.history.len() - 1;
|
|
let mut changing = vec![0usize; ndims];
|
|
for i in 0..num_deltas {
|
|
let prev = &self.history[i];
|
|
let curr = &self.history[i + 1];
|
|
if prev.len() < ndims || curr.len() < ndims {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
for d in 0..ndims {
|
|
if curr[d] != prev[d] {
|
|
changing[d] += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
let threshold = num_deltas.div_ceil(2);
|
|
let (max_dim, max_changes) = changing.iter().enumerate().max_by_key(|(_, c)| *c).unwrap();
|
|
|
|
if *max_changes < threshold {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
let others_max = changing
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(d, _)| *d != max_dim)
|
|
.map(|(_, c)| *c)
|
|
.max()
|
|
.unwrap_or(0);
|
|
|
|
if others_max > 0 && *max_changes < others_max * 2 {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
self.direction = if max_dim == ndims - 1 {
|
|
"row_major"
|
|
} else if max_dim == 0 {
|
|
"column_major"
|
|
} else {
|
|
"slice_major"
|
|
};
|
|
|
|
// Predict next chunks
|
|
let sweep_dim = max_dim;
|
|
let mut total_step: i64 = 0;
|
|
let mut step_count: usize = 0;
|
|
for i in 1..self.history.len() {
|
|
let prev = self.history[i - 1][sweep_dim] as i64;
|
|
let curr = self.history[i][sweep_dim] as i64;
|
|
let diff = curr - prev;
|
|
if diff != 0 {
|
|
total_step += diff;
|
|
step_count += 1;
|
|
}
|
|
}
|
|
|
|
if step_count == 0 {
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
let avg_step = total_step / step_count as i64;
|
|
if avg_step == 0 {
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
let last = self.history.last().unwrap();
|
|
self.predicted_next.clear();
|
|
for i in 1..=self.prefetch_count {
|
|
let mut pred = last.clone();
|
|
let new_val = last[sweep_dim] as i64 + avg_step * i as i64;
|
|
if new_val < 0 {
|
|
break;
|
|
}
|
|
pred[sweep_dim] = new_val as u64;
|
|
self.predicted_next.push(pred);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Read a chunked dataset with caching and sweep-aware prefetching.
|
|
///
|
|
/// Extends `read_chunked_data_cached` by feeding each chunk coordinate to a
|
|
/// [`SweepContext`]. When a sweep pattern is detected, predicted next-chunk
|
|
/// coordinates are pre-populated in the cache index via `prefetch_hint`.
|
|
#[cfg(feature = "std")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_chunked_data_sweep(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
cache: &ChunkCache,
|
|
sweep: &mut SweepContext,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (
|
|
chunk_dimensions,
|
|
version,
|
|
chunk_index_type,
|
|
addr_opt,
|
|
single_filtered_size,
|
|
single_filter_mask,
|
|
) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
version,
|
|
chunk_index_type,
|
|
single_chunk_filtered_size,
|
|
single_chunk_filter_mask,
|
|
} => (
|
|
chunk_dimensions,
|
|
*version,
|
|
*chunk_index_type,
|
|
*btree_address,
|
|
*single_chunk_filtered_size,
|
|
*single_chunk_filter_mask,
|
|
),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// The per-file cache is shared across datasets; bind it to this one so a
|
|
// different dataset's chunk index is never reused for this read.
|
|
cache.ensure_dataset(addr);
|
|
|
|
// Populate chunk index on first access
|
|
if !cache.has_index() {
|
|
let chunks = match (version, chunk_index_type) {
|
|
(3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?,
|
|
(4, Some(1)) => {
|
|
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
|
(fs as u32, single_filter_mask.unwrap_or(0))
|
|
} else {
|
|
(chunk_byte_size as u32, 0)
|
|
};
|
|
vec![ChunkInfo {
|
|
chunk_size: csize,
|
|
filter_mask: fmask,
|
|
offsets: vec![0u64; rank],
|
|
address: addr,
|
|
}]
|
|
}
|
|
(4, Some(2)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
generate_implicit_chunks(
|
|
addr,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
)
|
|
}
|
|
(4, Some(3)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_fixed_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(4, Some(4)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header = ExtensibleArrayHeader::parse(
|
|
file_data,
|
|
addr as usize,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
read_extensible_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(v, idx) => {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"unsupported chunked layout version={v}, index_type={idx:?}"
|
|
)));
|
|
}
|
|
};
|
|
cache.populate_index(&chunks, rank);
|
|
}
|
|
|
|
let chunks = cache.all_indexed_chunks().unwrap_or_default();
|
|
|
|
// Assemble output
|
|
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
|
|
if total_bytes == 0 {
|
|
// Also keeps the stride products below in range: with a zero-sized
|
|
// dimension the total is 0 even if other dimensions are huge.
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut output = alloc_output(total_bytes)?;
|
|
|
|
let mut ds_strides = vec![1usize; rank];
|
|
for i in (0..rank.saturating_sub(1)).rev() {
|
|
ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1];
|
|
}
|
|
|
|
let mut chunk_strides = vec![1usize; rank];
|
|
for i in (0..rank.saturating_sub(1)).rev() {
|
|
chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1];
|
|
}
|
|
|
|
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
|
|
for chunk_info in &chunks {
|
|
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
|
|
|
// Feed coordinate to sweep detector
|
|
sweep.record(coord.clone(), rank);
|
|
|
|
// Issue prefetch hint for predicted next chunks
|
|
if !sweep.predicted_next.is_empty() {
|
|
cache.prefetch_hint(&sweep.predicted_next);
|
|
cache.set_sweep_direction(sweep.direction);
|
|
}
|
|
|
|
// Try decompressed cache first
|
|
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
|
cached
|
|
} else {
|
|
// Decompress from file
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
|
let dec = if let Some(pl) = pipeline {
|
|
if chunk_info.filter_mask == 0 {
|
|
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
cache.put_decompressed(coord, dec)
|
|
};
|
|
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
if rank == 0 {
|
|
let copy_len = decompressed.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
&decompressed,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Read chunked data using pre-computed chunk layout for fast assembly.
|
|
///
|
|
/// This path builds a `ChunkIndex` and `ChunkLayout` on first access (cached
|
|
/// in the `ChunkCache`), then uses the pre-computed row-copy plan for assembly,
|
|
/// avoiding per-element N-D coordinate math on repeated reads.
|
|
#[cfg(feature = "std")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_chunked_data_indexed(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
cache: &ChunkCache,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (
|
|
chunk_dimensions,
|
|
version,
|
|
chunk_index_type,
|
|
addr_opt,
|
|
single_filtered_size,
|
|
single_filter_mask,
|
|
) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
version,
|
|
chunk_index_type,
|
|
single_chunk_filtered_size,
|
|
single_chunk_filter_mask,
|
|
} => (
|
|
chunk_dimensions,
|
|
*version,
|
|
*chunk_index_type,
|
|
*btree_address,
|
|
*single_chunk_filtered_size,
|
|
*single_chunk_filter_mask,
|
|
),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// The per-file cache is shared across datasets; bind it to this one so a
|
|
// different dataset's chunk index is never reused for this read.
|
|
cache.ensure_dataset(addr);
|
|
|
|
// Build chunk index on first access
|
|
if !cache.has_chunk_index() {
|
|
let chunks = match (version, chunk_index_type) {
|
|
(3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?,
|
|
(4, Some(1)) => {
|
|
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
|
(fs as u32, single_filter_mask.unwrap_or(0))
|
|
} else {
|
|
(chunk_byte_size as u32, 0)
|
|
};
|
|
vec![ChunkInfo {
|
|
chunk_size: csize,
|
|
filter_mask: fmask,
|
|
offsets: vec![0u64; rank],
|
|
address: addr,
|
|
}]
|
|
}
|
|
(4, Some(2)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
generate_implicit_chunks(
|
|
addr,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
)
|
|
}
|
|
(4, Some(3)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_fixed_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(4, Some(4)) => {
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header = ExtensibleArrayHeader::parse(
|
|
file_data,
|
|
addr as usize,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
read_extensible_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(v, idx) => {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"unsupported chunked layout version={v}, index_type={idx:?}"
|
|
)));
|
|
}
|
|
};
|
|
cache.populate_chunk_index(&chunks, rank);
|
|
// Also populate the legacy index for compatibility
|
|
if !cache.has_index() {
|
|
cache.populate_index(&chunks, rank);
|
|
}
|
|
}
|
|
|
|
// Build chunk layout on first access
|
|
if !cache.has_chunk_layout() {
|
|
cache.populate_chunk_layout(&ds_dims, &chunk_dims, elem_size);
|
|
}
|
|
|
|
// Get the layout info (mappings, output size, chunk total bytes)
|
|
let (mappings_info, output_bytes, chunk_total_bytes) = cache
|
|
.with_chunk_layout(|layout| {
|
|
let info: Vec<_> = layout
|
|
.mappings
|
|
.iter()
|
|
.map(|m| (m.coord.clone(), m.file_offset, m.file_size, m.filter_mask))
|
|
.collect();
|
|
(info, layout.output_bytes, layout.chunk_total_bytes)
|
|
})
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?;
|
|
|
|
// Decompress chunks (using LRU cache where possible)
|
|
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(mappings_info.len());
|
|
for (coord, file_offset, file_size, filter_mask) in &mappings_info {
|
|
if let Some(cached) = cache.get_decompressed_aligned(coord) {
|
|
chunk_buffers.push(cached);
|
|
} else {
|
|
let c_addr = *file_offset as usize;
|
|
let size = *file_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
let raw_chunk = &file_data[c_addr..c_addr + size];
|
|
let decompressed = if let Some(pl) = pipeline {
|
|
if *filter_mask == 0 {
|
|
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
let aligned = CacheAlignedBuffer::from_vec(decompressed);
|
|
let arc = cache.put_decompressed_aligned(coord.clone(), aligned);
|
|
chunk_buffers.push(arc);
|
|
}
|
|
}
|
|
|
|
// Assemble using pre-computed layout
|
|
let mut output = vec![0u8; output_bytes];
|
|
let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|b| b.as_slice()).collect();
|
|
cache.with_chunk_layout(|layout| {
|
|
layout.assemble(&data_refs, &mut output);
|
|
});
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Copy chunk data into the output buffer at the correct N-D position.
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn copy_chunk_to_output(
|
|
chunk_data: &[u8],
|
|
output: &mut [u8],
|
|
chunk_offsets: &[usize],
|
|
chunk_dims: &[usize],
|
|
ds_dims: &[usize],
|
|
ds_strides: &[usize],
|
|
chunk_strides: &[usize],
|
|
elem_size: usize,
|
|
rank: usize,
|
|
) {
|
|
// Row-copy approach: iterate over outer dimensions, memcpy the innermost
|
|
// dimension in bulk. For 1-D data this is a single memcpy per chunk.
|
|
// For N-D data this is one memcpy per "row" (innermost dim slice).
|
|
|
|
if rank == 1 {
|
|
// Fast path for 1-D: single contiguous copy per chunk
|
|
let global_start = chunk_offsets[0];
|
|
let copy_len = chunk_dims[0].min(ds_dims[0].saturating_sub(global_start));
|
|
let (Some(src_bytes), Some(dst_start)) = (
|
|
copy_len.checked_mul(elem_size),
|
|
global_start.checked_mul(elem_size),
|
|
) else {
|
|
return;
|
|
};
|
|
if src_bytes > 0
|
|
&& dst_start
|
|
.checked_add(src_bytes)
|
|
.is_some_and(|end| end <= output.len())
|
|
&& src_bytes <= chunk_data.len()
|
|
{
|
|
output[dst_start..dst_start + src_bytes].copy_from_slice(&chunk_data[..src_bytes]);
|
|
}
|
|
return;
|
|
}
|
|
|
|
// General N-D: iterate over outer dimensions, copy innermost rows
|
|
let inner_dim = rank - 1;
|
|
let inner_chunk_len =
|
|
chunk_dims[inner_dim].min(ds_dims[inner_dim].saturating_sub(chunk_offsets[inner_dim]));
|
|
let Some(row_bytes) = inner_chunk_len.checked_mul(elem_size) else {
|
|
return;
|
|
};
|
|
|
|
if row_bytes == 0 {
|
|
return;
|
|
}
|
|
|
|
// Number of rows = product of all outer chunk dimensions
|
|
let Some(outer_count) = chunk_dims[..inner_dim]
|
|
.iter()
|
|
.try_fold(1usize, |acc, &d| acc.checked_mul(d))
|
|
else {
|
|
return;
|
|
};
|
|
|
|
// Outer strides for iterating chunk-local coordinates
|
|
let mut outer_strides = vec![1usize; inner_dim];
|
|
for i in (0..inner_dim.saturating_sub(1)).rev() {
|
|
let Some(stride) = outer_strides[i + 1].checked_mul(chunk_dims[i + 1]) else {
|
|
return;
|
|
};
|
|
outer_strides[i] = stride;
|
|
}
|
|
|
|
for outer_idx in 0..outer_count {
|
|
// Convert outer flat index to N-D chunk-local coords for dims 0..inner_dim
|
|
let mut remaining = outer_idx;
|
|
let mut ds_flat = 0usize;
|
|
let mut src_flat = 0usize;
|
|
let mut out_of_bounds = false;
|
|
|
|
for d in 0..inner_dim {
|
|
let coord_in_chunk = if inner_dim > 1 {
|
|
remaining / outer_strides[d]
|
|
} else {
|
|
remaining
|
|
};
|
|
if inner_dim > 1 {
|
|
remaining %= outer_strides[d];
|
|
}
|
|
|
|
let Some(global_coord) = chunk_offsets[d].checked_add(coord_in_chunk) else {
|
|
out_of_bounds = true;
|
|
break;
|
|
};
|
|
if global_coord >= ds_dims[d] {
|
|
out_of_bounds = true;
|
|
break;
|
|
}
|
|
let (Some(ds_term), Some(src_term)) = (
|
|
global_coord.checked_mul(ds_strides[d]),
|
|
coord_in_chunk.checked_mul(chunk_strides[d]),
|
|
) else {
|
|
out_of_bounds = true;
|
|
break;
|
|
};
|
|
let (Some(new_ds_flat), Some(new_src_flat)) =
|
|
(ds_flat.checked_add(ds_term), src_flat.checked_add(src_term))
|
|
else {
|
|
out_of_bounds = true;
|
|
break;
|
|
};
|
|
ds_flat = new_ds_flat;
|
|
src_flat = new_src_flat;
|
|
}
|
|
|
|
if out_of_bounds {
|
|
continue;
|
|
}
|
|
|
|
// Add innermost dimension offset
|
|
let Some(inner_term) = chunk_offsets[inner_dim].checked_mul(ds_strides[inner_dim]) else {
|
|
continue;
|
|
};
|
|
let Some(ds_flat) = ds_flat.checked_add(inner_term) else {
|
|
continue;
|
|
};
|
|
|
|
let (Some(src_start), Some(dst_start)) = (
|
|
src_flat.checked_mul(elem_size),
|
|
ds_flat.checked_mul(elem_size),
|
|
) else {
|
|
continue;
|
|
};
|
|
|
|
let fits = src_start
|
|
.checked_add(row_bytes)
|
|
.is_some_and(|end| end <= chunk_data.len())
|
|
&& dst_start
|
|
.checked_add(row_bytes)
|
|
.is_some_and(|end| end <= output.len());
|
|
if fits {
|
|
output[dst_start..dst_start + row_bytes]
|
|
.copy_from_slice(&chunk_data[src_start..src_start + row_bytes]);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn simple_space(dimensions: Vec<u64>) -> Dataspace {
|
|
Dataspace {
|
|
space_type: crate::dataspace::DataspaceType::Simple,
|
|
rank: dimensions.len() as u8,
|
|
dimensions,
|
|
max_dimensions: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn crafted_dimensions_are_errors_not_wraparound() {
|
|
// 2^63 * 2 wraps to 0 with a plain product; 2^40 * 2^40 wraps too.
|
|
for dims in [
|
|
vec![1u64 << 63, 2],
|
|
vec![1 << 40, 1 << 40],
|
|
vec![u64::MAX, u64::MAX],
|
|
] {
|
|
let space = simple_space(dims.clone());
|
|
assert!(
|
|
matches!(space.checked_num_elements(), Err(FormatError::Overflow(_))),
|
|
"{dims:?}"
|
|
);
|
|
// The infallible accessor saturates instead of wrapping.
|
|
assert_eq!(space.num_elements(), u64::MAX, "{dims:?}");
|
|
}
|
|
assert_eq!(simple_space(vec![3, 4]).checked_num_elements().unwrap(), 12);
|
|
// A zero-sized dimension makes the whole product 0, not an overflow.
|
|
assert_eq!(
|
|
simple_space(vec![0, 1 << 40, 1 << 40])
|
|
.checked_num_elements()
|
|
.unwrap(),
|
|
0
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn byte_length_helpers_check_overflow() {
|
|
assert_eq!(checked_byte_len(10, 8).unwrap(), 80);
|
|
assert!(matches!(
|
|
checked_byte_len(u64::MAX, 8),
|
|
Err(FormatError::Overflow(_))
|
|
));
|
|
assert_eq!(checked_chunk_byte_len(&[10, 10], 4).unwrap(), 400);
|
|
assert!(matches!(
|
|
checked_chunk_byte_len(&[usize::MAX, 2], 4),
|
|
Err(FormatError::Overflow(_))
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn unallocatable_output_is_an_error_not_an_abort() {
|
|
assert_eq!(alloc_output(16).unwrap(), vec![0u8; 16]);
|
|
assert!(matches!(
|
|
alloc_output(usize::MAX / 2),
|
|
Err(FormatError::Overflow(_))
|
|
));
|
|
}
|
|
|
|
fn write_offset(buf: &mut Vec<u8>, val: u64, size: u8) {
|
|
match size {
|
|
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
|
|
8 => buf.extend_from_slice(&val.to_le_bytes()),
|
|
_ => panic!("unsupported offset size in test"),
|
|
}
|
|
}
|
|
|
|
/// Build a B-tree v1 type 1 leaf node with given chunk infos.
|
|
fn build_chunk_btree_leaf(chunks: &[ChunkInfo], ndims: usize, offset_size: u8) -> Vec<u8> {
|
|
let _os = offset_size as usize;
|
|
let entries_used = chunks.len() as u16;
|
|
let mut buf = Vec::new();
|
|
|
|
// Header
|
|
buf.extend_from_slice(b"TREE");
|
|
buf.push(1); // node_type = 1 (raw data chunks)
|
|
buf.push(0); // node_level = 0 (leaf)
|
|
buf.extend_from_slice(&entries_used.to_le_bytes());
|
|
|
|
// Left/right sibling = undefined
|
|
let undef: u64 = if offset_size == 4 {
|
|
0xFFFFFFFF
|
|
} else {
|
|
0xFFFFFFFFFFFFFFFF
|
|
};
|
|
write_offset(&mut buf, undef, offset_size);
|
|
write_offset(&mut buf, undef, offset_size);
|
|
|
|
// Entries: key[i], child[i] pairs, then final key
|
|
for chunk in chunks {
|
|
// Key: chunk_size(4) + filter_mask(4) + ndims offsets
|
|
buf.extend_from_slice(&chunk.chunk_size.to_le_bytes());
|
|
buf.extend_from_slice(&chunk.filter_mask.to_le_bytes());
|
|
for d in 0..ndims {
|
|
let off = if d < chunk.offsets.len() {
|
|
chunk.offsets[d]
|
|
} else {
|
|
0
|
|
};
|
|
write_offset(&mut buf, off, offset_size);
|
|
}
|
|
// Child: address
|
|
write_offset(&mut buf, chunk.address, offset_size);
|
|
}
|
|
|
|
// Final key (dummy)
|
|
buf.extend_from_slice(&0u32.to_le_bytes()); // chunk_size
|
|
buf.extend_from_slice(&0u32.to_le_bytes()); // filter_mask
|
|
for _ in 0..ndims {
|
|
write_offset(&mut buf, u64::MAX, offset_size);
|
|
}
|
|
|
|
buf
|
|
}
|
|
|
|
// --- ChunkInfo collection tests ---
|
|
|
|
#[test]
|
|
fn collect_two_chunks_from_leaf() {
|
|
let ndims = 2; // rank+1 for 1D dataset
|
|
let os: u8 = 8;
|
|
|
|
let chunks = vec![
|
|
ChunkInfo {
|
|
chunk_size: 80,
|
|
filter_mask: 0,
|
|
offsets: vec![0, 0],
|
|
address: 0x1000,
|
|
},
|
|
ChunkInfo {
|
|
chunk_size: 80,
|
|
filter_mask: 0,
|
|
offsets: vec![10, 0],
|
|
address: 0x2000,
|
|
},
|
|
];
|
|
|
|
let btree = build_chunk_btree_leaf(&chunks, ndims, os);
|
|
let mut file_data = vec![0u8; 0x3000];
|
|
file_data[..btree.len()].copy_from_slice(&btree);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
|
|
assert_eq!(result.len(), 2);
|
|
assert_eq!(result[0].address, 0x1000);
|
|
assert_eq!(result[0].offsets, vec![0, 0]);
|
|
assert_eq!(result[0].chunk_size, 80);
|
|
assert_eq!(result[1].address, 0x2000);
|
|
assert_eq!(result[1].offsets, vec![10, 0]);
|
|
}
|
|
|
|
#[test]
|
|
fn collect_three_chunks() {
|
|
let ndims = 2;
|
|
let os: u8 = 8;
|
|
|
|
let chunks = vec![
|
|
ChunkInfo {
|
|
chunk_size: 40,
|
|
filter_mask: 0,
|
|
offsets: vec![0, 0],
|
|
address: 0x100,
|
|
},
|
|
ChunkInfo {
|
|
chunk_size: 40,
|
|
filter_mask: 0,
|
|
offsets: vec![5, 0],
|
|
address: 0x200,
|
|
},
|
|
ChunkInfo {
|
|
chunk_size: 40,
|
|
filter_mask: 0,
|
|
offsets: vec![10, 0],
|
|
address: 0x300,
|
|
},
|
|
];
|
|
|
|
let btree = build_chunk_btree_leaf(&chunks, ndims, os);
|
|
let mut file_data = vec![0u8; 0x1000];
|
|
file_data[..btree.len()].copy_from_slice(&btree);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
|
|
assert_eq!(result.len(), 3);
|
|
assert_eq!(result[0].address, 0x100);
|
|
assert_eq!(result[1].address, 0x200);
|
|
assert_eq!(result[2].address, 0x300);
|
|
}
|
|
|
|
#[test]
|
|
fn collect_empty_btree() {
|
|
let ndims = 2;
|
|
let os: u8 = 8;
|
|
let btree = build_chunk_btree_leaf(&[], ndims, os);
|
|
let mut file_data = vec![0u8; 0x1000];
|
|
file_data[..btree.len()].copy_from_slice(&btree);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
|
|
assert_eq!(result.len(), 0);
|
|
}
|
|
|
|
// --- Chunked read tests (synthetic) ---
|
|
|
|
use crate::dataspace::{Dataspace, DataspaceType};
|
|
use crate::datatype::{Datatype, DatatypeByteOrder};
|
|
|
|
fn make_f64_type() -> Datatype {
|
|
Datatype::FloatingPoint {
|
|
size: 8,
|
|
byte_order: DatatypeByteOrder::LittleEndian,
|
|
bit_offset: 0,
|
|
bit_precision: 64,
|
|
exponent_location: 52,
|
|
exponent_size: 11,
|
|
mantissa_location: 0,
|
|
mantissa_size: 52,
|
|
exponent_bias: 1023,
|
|
}
|
|
}
|
|
|
|
fn make_f32_type() -> Datatype {
|
|
Datatype::FloatingPoint {
|
|
size: 4,
|
|
byte_order: DatatypeByteOrder::LittleEndian,
|
|
bit_offset: 0,
|
|
bit_precision: 32,
|
|
exponent_location: 23,
|
|
exponent_size: 8,
|
|
mantissa_location: 0,
|
|
mantissa_size: 23,
|
|
exponent_bias: 127,
|
|
}
|
|
}
|
|
|
|
/// Build a synthetic file with a B-tree and chunk data for a 1D uncompressed dataset.
|
|
fn build_1d_chunked_file(
|
|
values: &[f64],
|
|
chunk_size_elems: usize,
|
|
) -> (Vec<u8>, DataLayout, Dataspace) {
|
|
let os: u8 = 8;
|
|
let elem_size = 8usize;
|
|
let ndims = 2; // rank(1) + 1
|
|
let total = values.len();
|
|
|
|
// Place chunk data starting at offset 0x2000
|
|
let mut file_data = vec![0u8; 0x10000];
|
|
let mut chunk_infos = Vec::new();
|
|
let mut data_offset = 0x2000usize;
|
|
|
|
let mut start = 0;
|
|
while start < total {
|
|
let end = (start + chunk_size_elems).min(total);
|
|
let chunk_bytes = chunk_size_elems * elem_size; // full chunk allocation
|
|
|
|
// Write chunk data (full chunk size, padding with zeros)
|
|
for (i, value) in values.iter().enumerate().take(end).skip(start) {
|
|
let byte_offset = data_offset + (i - start) * elem_size;
|
|
file_data[byte_offset..byte_offset + 8].copy_from_slice(&value.to_le_bytes());
|
|
}
|
|
|
|
chunk_infos.push(ChunkInfo {
|
|
chunk_size: chunk_bytes as u32,
|
|
filter_mask: 0,
|
|
offsets: vec![start as u64, 0],
|
|
address: data_offset as u64,
|
|
});
|
|
|
|
data_offset += chunk_bytes;
|
|
start += chunk_size_elems;
|
|
}
|
|
|
|
// Build B-tree at offset 0x100
|
|
let btree = build_chunk_btree_leaf(&chunk_infos, ndims, os);
|
|
let btree_addr = 0x100usize;
|
|
file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_size_elems as u32, elem_size as u32],
|
|
btree_address: Some(btree_addr as u64),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![total as u64],
|
|
max_dimensions: None,
|
|
};
|
|
|
|
(file_data, layout, dataspace)
|
|
}
|
|
|
|
#[test]
|
|
fn read_chunked_data_rejects_zero_dim_chunk_layout() {
|
|
// Found by fuzzing: chunk_dimensions.len() == 0 caused `ndims - 1` to
|
|
// underflow. A malformed/degenerate chunked layout must error cleanly.
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![],
|
|
btree_address: Some(0),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![10],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f64_type();
|
|
let file_data = vec![0u8; 64];
|
|
let result = read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8);
|
|
assert!(
|
|
matches!(result, Err(FormatError::ChunkedReadError(_))),
|
|
"expected a clean ChunkedReadError, got {result:?}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn copy_chunk_to_output_1d_rejects_overflowing_offset_without_panicking() {
|
|
// Found by fuzzing: `global_start * elem_size` overflowed for a
|
|
// crafted large chunk offset.
|
|
let chunk_data = vec![1u8; 16];
|
|
let mut output = vec![0u8; 16];
|
|
let chunk_offsets = [usize::MAX - 1];
|
|
let chunk_dims = [1usize];
|
|
let ds_dims = [usize::MAX];
|
|
let ds_strides = [1usize];
|
|
let chunk_strides = [1usize];
|
|
copy_chunk_to_output(
|
|
&chunk_data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
8,
|
|
1,
|
|
);
|
|
// No panic; the out-of-range write was skipped, output left untouched.
|
|
assert_eq!(output, vec![0u8; 16]);
|
|
}
|
|
|
|
#[test]
|
|
fn copy_chunk_to_output_nd_rejects_overflowing_offset_without_panicking() {
|
|
let chunk_data = vec![1u8; 16];
|
|
let mut output = vec![0u8; 16];
|
|
let chunk_offsets = [usize::MAX - 1, 0];
|
|
let chunk_dims = [1usize, 1usize];
|
|
let ds_dims = [usize::MAX, usize::MAX];
|
|
let ds_strides = [1usize, 1usize];
|
|
let chunk_strides = [1usize, 1usize];
|
|
copy_chunk_to_output(
|
|
&chunk_data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
8,
|
|
2,
|
|
);
|
|
assert_eq!(output, vec![0u8; 16]);
|
|
}
|
|
|
|
#[test]
|
|
fn read_1d_two_chunks_no_compression() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 20 * 8);
|
|
|
|
// Verify values
|
|
for i in 0..20 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn read_1d_three_chunks_partial_last() {
|
|
// 25 elements, chunk size 10 => 3 chunks, last has only 5 valid
|
|
let values: Vec<f64> = (0..25).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 25 * 8);
|
|
|
|
for i in 0..25 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "deflate")]
|
|
#[test]
|
|
fn read_1d_two_chunks_with_deflate() {
|
|
use crate::filter_pipeline::{FILTER_DEFLATE, FilterDescription, FilterPipeline};
|
|
use crate::filters::compress_chunk;
|
|
|
|
let os: u8 = 8;
|
|
let elem_size = 8usize;
|
|
let ndims = 2;
|
|
let chunk_elems = 10usize;
|
|
let total = 20usize;
|
|
|
|
let pipeline = FilterPipeline {
|
|
version: 2,
|
|
filters: vec![FilterDescription {
|
|
filter_id: FILTER_DEFLATE,
|
|
name: None,
|
|
flags: 0,
|
|
client_data: vec![6],
|
|
}],
|
|
};
|
|
|
|
let values: Vec<f64> = (0..total).map(|i| i as f64).collect();
|
|
let mut file_data = vec![0u8; 0x10000];
|
|
let mut chunk_infos = Vec::new();
|
|
let mut data_offset = 0x2000usize;
|
|
|
|
for chunk_idx in 0..2 {
|
|
let start = chunk_idx * chunk_elems;
|
|
let mut chunk_bytes = Vec::new();
|
|
for value in values.iter().skip(start).take(chunk_elems) {
|
|
chunk_bytes.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
let compressed = compress_chunk(&chunk_bytes, &pipeline, elem_size as u32).unwrap();
|
|
|
|
file_data[data_offset..data_offset + compressed.len()].copy_from_slice(&compressed);
|
|
|
|
chunk_infos.push(ChunkInfo {
|
|
chunk_size: compressed.len() as u32,
|
|
filter_mask: 0,
|
|
offsets: vec![start as u64, 0],
|
|
address: data_offset as u64,
|
|
});
|
|
|
|
data_offset += compressed.len() + 16; // some padding
|
|
}
|
|
|
|
let btree = build_chunk_btree_leaf(&chunk_infos, ndims, os);
|
|
let btree_addr = 0x100usize;
|
|
file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_elems as u32, elem_size as u32],
|
|
btree_address: Some(btree_addr as u64),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![total as u64],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f64_type();
|
|
|
|
let raw = read_chunked_data(
|
|
&file_data,
|
|
&layout,
|
|
&dataspace,
|
|
&datatype,
|
|
Some(&pipeline),
|
|
8,
|
|
8,
|
|
)
|
|
.unwrap();
|
|
|
|
for i in 0..total {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn read_2d_four_chunks() {
|
|
// 4x6 dataset with chunk size 2x3 => 4 chunks
|
|
let os: u8 = 8;
|
|
let elem_size = 4usize; // f32
|
|
let ndims = 3; // rank(2) + 1
|
|
let ds_dims = [4usize, 6];
|
|
let chunk_dims = [2usize, 3];
|
|
|
|
let values: Vec<f32> = (0..24).map(|i| i as f32).collect();
|
|
let mut file_data = vec![0u8; 0x10000];
|
|
let mut chunk_infos = Vec::new();
|
|
let mut data_offset = 0x2000usize;
|
|
|
|
// Generate chunks: (0,0), (0,3), (2,0), (2,3)
|
|
for row_start in (0..ds_dims[0]).step_by(chunk_dims[0]) {
|
|
for col_start in (0..ds_dims[1]).step_by(chunk_dims[1]) {
|
|
let mut chunk_bytes = Vec::new();
|
|
for r in 0..chunk_dims[0] {
|
|
for c in 0..chunk_dims[1] {
|
|
let gr = row_start + r;
|
|
let gc = col_start + c;
|
|
let val = if gr < ds_dims[0] && gc < ds_dims[1] {
|
|
values[gr * ds_dims[1] + gc]
|
|
} else {
|
|
0.0
|
|
};
|
|
chunk_bytes.extend_from_slice(&val.to_le_bytes());
|
|
}
|
|
}
|
|
|
|
let chunk_size = chunk_bytes.len();
|
|
file_data[data_offset..data_offset + chunk_size].copy_from_slice(&chunk_bytes);
|
|
|
|
chunk_infos.push(ChunkInfo {
|
|
chunk_size: chunk_size as u32,
|
|
filter_mask: 0,
|
|
offsets: vec![row_start as u64, col_start as u64, 0],
|
|
address: data_offset as u64,
|
|
});
|
|
|
|
data_offset += chunk_size + 8;
|
|
}
|
|
}
|
|
|
|
let btree = build_chunk_btree_leaf(&chunk_infos, ndims, os);
|
|
let btree_addr = 0x100usize;
|
|
file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_dims[0] as u32, chunk_dims[1] as u32, elem_size as u32],
|
|
btree_address: Some(btree_addr as u64),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 2,
|
|
dimensions: vec![ds_dims[0] as u64, ds_dims[1] as u64],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f32_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 24 * 4);
|
|
|
|
for i in 0..24 {
|
|
let val = f32::from_le_bytes(raw[i * 4..(i + 1) * 4].try_into().unwrap());
|
|
assert_eq!(val, i as f32, "mismatch at element {i}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn wrong_node_type_error() {
|
|
// Build a type-0 B-tree and try to collect chunk info
|
|
let mut buf = Vec::new();
|
|
buf.extend_from_slice(b"TREE");
|
|
buf.push(0); // type 0, not 1
|
|
buf.push(0);
|
|
buf.extend_from_slice(&0u16.to_le_bytes());
|
|
buf.extend_from_slice(&[0xFF; 16]); // siblings
|
|
// final key
|
|
buf.extend_from_slice(&[0u8; 24]);
|
|
|
|
let mut file_data = vec![0u8; 512];
|
|
file_data[..buf.len()].copy_from_slice(&buf);
|
|
|
|
let err = collect_chunk_info(&file_data, 0, 2, 8, 8).unwrap_err();
|
|
assert_eq!(err, FormatError::InvalidBTreeNodeType(0));
|
|
}
|
|
|
|
#[test]
|
|
fn collect_chunk_info_rejects_near_usize_max_offset() {
|
|
let file_data = vec![0u8; 64];
|
|
let result = collect_chunk_info(&file_data, u64::MAX - 4, 2, 8, 8);
|
|
assert!(
|
|
matches!(result, Err(FormatError::UnexpectedEof { .. })),
|
|
"expected a clean UnexpectedEof, got {result:?}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn collect_chunk_info_rejects_self_referencing_internal_node() {
|
|
// A type-1 internal node (level 1) whose single child address points
|
|
// back to itself: an infinite-recursion / cyclic B-tree attack.
|
|
let ndims = 2;
|
|
let os: u8 = 8;
|
|
let mut buf = Vec::new();
|
|
buf.extend_from_slice(b"TREE");
|
|
buf.push(1); // node_type = 1 (raw data chunks)
|
|
buf.push(1); // node_level = 1 (internal)
|
|
buf.extend_from_slice(&1u16.to_le_bytes()); // entries_used = 1
|
|
write_offset(&mut buf, u64::MAX, os); // left sibling undefined
|
|
write_offset(&mut buf, u64::MAX, os); // right sibling undefined
|
|
// key[0]: chunk_size(4) + filter_mask(4) + ndims offsets
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
for _ in 0..ndims {
|
|
write_offset(&mut buf, 0, os);
|
|
}
|
|
// child[0]: points back to offset 0 (this same node) — cyclic.
|
|
write_offset(&mut buf, 0, os);
|
|
// final key
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
for _ in 0..ndims {
|
|
write_offset(&mut buf, u64::MAX, os);
|
|
}
|
|
|
|
let mut file_data = vec![0u8; 256];
|
|
file_data[..buf.len()].copy_from_slice(&buf);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os);
|
|
assert!(
|
|
matches!(result, Err(FormatError::NestingDepthExceeded)),
|
|
"expected a clean NestingDepthExceeded, got {result:?}"
|
|
);
|
|
}
|
|
|
|
// --- Implicit chunk generation tests ---
|
|
|
|
#[test]
|
|
fn implicit_chunks_1d_five_chunks() {
|
|
let chunks = generate_implicit_chunks(
|
|
0x1000,
|
|
&[100],
|
|
&[20],
|
|
8, // f64
|
|
);
|
|
assert_eq!(chunks.len(), 5);
|
|
let chunk_byte_size = 20 * 8;
|
|
for (i, c) in chunks.iter().enumerate() {
|
|
assert_eq!(c.address, 0x1000 + i as u64 * chunk_byte_size as u64);
|
|
assert_eq!(c.offsets, vec![i as u64 * 20]);
|
|
assert_eq!(c.filter_mask, 0);
|
|
assert_eq!(c.chunk_size, chunk_byte_size as u32);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn implicit_chunks_2d() {
|
|
// 10x6 dataset, 4x3 chunks => ceil(10/4)=3, ceil(6/3)=2 => 6 chunks
|
|
let chunks = generate_implicit_chunks(
|
|
0x2000,
|
|
&[10, 6],
|
|
&[4, 3],
|
|
4, // f32
|
|
);
|
|
assert_eq!(chunks.len(), 6);
|
|
let chunk_byte_size = 4 * 3 * 4;
|
|
// Row-major: (0,0), (0,3), (4,0), (4,3), (8,0), (8,3)
|
|
assert_eq!(chunks[0].offsets, vec![0, 0]);
|
|
assert_eq!(chunks[1].offsets, vec![0, 3]);
|
|
assert_eq!(chunks[2].offsets, vec![4, 0]);
|
|
assert_eq!(chunks[3].offsets, vec![4, 3]);
|
|
assert_eq!(chunks[4].offsets, vec![8, 0]);
|
|
assert_eq!(chunks[5].offsets, vec![8, 3]);
|
|
for (i, c) in chunks.iter().enumerate() {
|
|
assert_eq!(c.address, 0x2000 + i as u64 * chunk_byte_size as u64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn implicit_chunks_partial_last() {
|
|
// 25 elements, chunk size 10 => 3 chunks (last partial)
|
|
let chunks = generate_implicit_chunks(0x0, &[25], &[10], 8);
|
|
assert_eq!(chunks.len(), 3);
|
|
assert_eq!(chunks[0].offsets, vec![0]);
|
|
assert_eq!(chunks[1].offsets, vec![10]);
|
|
assert_eq!(chunks[2].offsets, vec![20]);
|
|
}
|
|
|
|
// --- V4 single chunk synthetic test ---
|
|
|
|
#[test]
|
|
fn read_v4_single_chunk_synthetic() {
|
|
// Build a synthetic v4 single chunk dataset (no filters)
|
|
let values: Vec<f64> = vec![10.0, 20.0, 30.0];
|
|
let elem_size = 8usize;
|
|
let chunk_elems = 3usize;
|
|
|
|
let mut file_data = vec![0u8; 0x2000];
|
|
let data_addr = 0x1000usize;
|
|
for (i, &v) in values.iter().enumerate() {
|
|
file_data[data_addr + i * elem_size..data_addr + (i + 1) * elem_size]
|
|
.copy_from_slice(&v.to_le_bytes());
|
|
}
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_elems as u32, elem_size as u32],
|
|
btree_address: Some(data_addr as u64),
|
|
version: 4,
|
|
chunk_index_type: Some(1),
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![3],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f64_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 24);
|
|
for i in 0..3 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, values[i]);
|
|
}
|
|
}
|
|
|
|
// --- Cached read tests ---
|
|
|
|
use crate::chunk_cache::ChunkCache;
|
|
|
|
#[test]
|
|
fn cached_read_populates_index_and_returns_correct_data() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
|
|
assert!(!cache.has_index());
|
|
let raw = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert!(cache.has_index());
|
|
assert_eq!(raw.len(), 20 * 8);
|
|
for i in 0..20 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cached_read_second_call_uses_cache() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
|
|
// First read — populates index + decompressed cache
|
|
let raw1 = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert!(cache.has_index());
|
|
assert!(cache.cached_chunk_count() > 0);
|
|
|
|
// Second read — should hit the decompressed cache
|
|
let raw2 = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(raw1, raw2);
|
|
}
|
|
|
|
#[test]
|
|
fn cached_read_with_partial_last_chunk() {
|
|
let values: Vec<f64> = (0..25).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
|
|
let raw = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(raw.len(), 25 * 8);
|
|
for i in 0..25 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
// --- Sweep-aware read tests ---
|
|
|
|
#[test]
|
|
fn sweep_read_returns_correct_data() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
let mut sweep = SweepContext::with_defaults();
|
|
|
|
let raw = read_chunked_data_sweep(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(raw.len(), 20 * 8);
|
|
for i in 0..20 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_read_populates_sweep_context() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
let mut sweep = SweepContext::with_defaults();
|
|
|
|
read_chunked_data_sweep(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
|
|
)
|
|
.unwrap();
|
|
|
|
// After reading 2 chunks (offsets [0] and [10]), history should be populated
|
|
assert!(!sweep.history.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_context_unit_test() {
|
|
let mut ctx = SweepContext::with_defaults();
|
|
ctx.record(vec![0, 0], 2);
|
|
ctx.record(vec![0, 10], 2);
|
|
ctx.record(vec![0, 20], 2);
|
|
assert_eq!(ctx.direction, "row_major");
|
|
assert!(!ctx.predicted_next.is_empty());
|
|
assert_eq!(ctx.predicted_next[0], vec![0, 30]);
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_context_random() {
|
|
let mut ctx = SweepContext::with_defaults();
|
|
ctx.record(vec![0, 0], 2);
|
|
ctx.record(vec![30, 20], 2);
|
|
ctx.record(vec![10, 0], 2);
|
|
assert_eq!(ctx.direction, "random");
|
|
assert!(ctx.predicted_next.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_read_access_stats() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
let mut sweep = SweepContext::with_defaults();
|
|
|
|
read_chunked_data_sweep(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
|
|
)
|
|
.unwrap();
|
|
|
|
let stats = cache.access_stats();
|
|
// We accessed 2 chunks; the second should be sequential to the first
|
|
assert!(stats.sequential_count > 0 || stats.random_count > 0);
|
|
}
|
|
}
|