678 lines
22 KiB
Rust
678 lines
22 KiB
Rust
//! Chunk B-tree index cache and pre-computed layout for fast repeated reads.
|
|
//!
|
|
//! [`ChunkIndex`] scans the chunk index structure (B-tree v1, fixed array, etc.)
|
|
//! once and builds a flat `HashMap<ChunkCoord, ChunkInfo>` for O(1) lookups.
|
|
//!
|
|
//! [`ChunkLayout`] pre-computes contiguous row-copy operations for each chunk,
|
|
//! turning the N-D → flat coordinate math into a series of `memcpy` calls.
|
|
//! This eliminates per-element coordinate arithmetic during output assembly.
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
extern crate alloc;
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::{vec, vec::Vec};
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::collections::BTreeMap;
|
|
#[cfg(feature = "std")]
|
|
use std::collections::HashMap;
|
|
|
|
use crate::chunk_cache::ChunkCoord;
|
|
use crate::chunked_read::ChunkInfo;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ChunkIndex — flat HashMap of coord → ChunkInfo
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A cached chunk index that maps chunk coordinates to file location metadata.
|
|
///
|
|
/// Built once from a B-tree (or other index structure) traversal, then provides
|
|
/// O(1) lookups for any chunk coordinate.
|
|
pub struct ChunkIndex {
|
|
#[cfg(feature = "std")]
|
|
map: HashMap<ChunkCoord, ChunkInfo>,
|
|
#[cfg(not(feature = "std"))]
|
|
map: BTreeMap<ChunkCoord, ChunkInfo>,
|
|
rank: usize,
|
|
}
|
|
|
|
impl ChunkIndex {
|
|
/// Build a chunk index from a pre-collected list of `ChunkInfo`.
|
|
///
|
|
/// `rank` is the spatial dimensionality — offsets are truncated to the first
|
|
/// `rank` elements (B-tree v1 stores rank+1 offsets with element size as the
|
|
/// last dimension).
|
|
pub fn build(chunks: &[ChunkInfo], rank: usize) -> Self {
|
|
#[cfg(feature = "std")]
|
|
let mut map = HashMap::with_capacity(chunks.len());
|
|
#[cfg(not(feature = "std"))]
|
|
let mut map = BTreeMap::new();
|
|
|
|
for ci in chunks {
|
|
let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
|
|
map.insert(coord, ci.clone());
|
|
}
|
|
Self { map, rank }
|
|
}
|
|
|
|
/// O(1) lookup of a chunk by its spatial coordinate.
|
|
#[inline]
|
|
pub fn lookup(&self, coords: &[u64]) -> Option<&ChunkInfo> {
|
|
self.map.get(coords)
|
|
}
|
|
|
|
/// Number of chunks in the index.
|
|
#[inline]
|
|
pub fn num_chunks(&self) -> usize {
|
|
self.map.len()
|
|
}
|
|
|
|
/// The spatial rank this index was built for.
|
|
#[inline]
|
|
pub fn rank(&self) -> usize {
|
|
self.rank
|
|
}
|
|
|
|
/// Iterate over all indexed chunks.
|
|
pub fn iter(&self) -> impl Iterator<Item = (&ChunkCoord, &ChunkInfo)> {
|
|
self.map.iter()
|
|
}
|
|
|
|
/// Return all chunk infos as a Vec (order unspecified).
|
|
pub fn all_chunks(&self) -> Vec<ChunkInfo> {
|
|
self.map.values().cloned().collect()
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ChunkLayout — pre-computed row-copy operations
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A single contiguous memory copy within the output assembly.
|
|
#[derive(Debug, Clone)]
|
|
pub struct RowCopy {
|
|
/// Byte offset within the decompressed chunk buffer.
|
|
pub src_offset: usize,
|
|
/// Byte offset within the output buffer.
|
|
pub dst_offset: usize,
|
|
/// Number of bytes to copy.
|
|
pub len: usize,
|
|
}
|
|
|
|
/// Pre-computed mapping for one chunk: its metadata plus the row-copy plan.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ChunkMapping {
|
|
/// Spatial coordinate of this chunk.
|
|
pub coord: ChunkCoord,
|
|
/// File byte address of the compressed chunk.
|
|
pub file_offset: u64,
|
|
/// Size of the compressed chunk in the file.
|
|
pub file_size: u32,
|
|
/// Filter mask (0 = all filters applied).
|
|
pub filter_mask: u32,
|
|
/// Pre-computed row-copy operations for assembling this chunk into output.
|
|
pub copies: Vec<RowCopy>,
|
|
}
|
|
|
|
/// Pre-computed layout for assembling all chunks of a dataset into a flat
|
|
/// output buffer.
|
|
///
|
|
/// Built once from the chunk index + dataset/chunk dimensions, then reused
|
|
/// on every read. Turns the per-element N-D → flat coordinate math into a
|
|
/// sequence of `memcpy` calls.
|
|
pub struct ChunkLayout {
|
|
pub mappings: Vec<ChunkMapping>,
|
|
/// Total bytes of the output buffer.
|
|
pub output_bytes: usize,
|
|
/// Element size in bytes.
|
|
pub elem_size: usize,
|
|
/// Chunk total bytes (uncompressed).
|
|
pub chunk_total_bytes: usize,
|
|
}
|
|
|
|
impl ChunkLayout {
|
|
/// Build the chunk layout from a chunk index and dataset geometry.
|
|
///
|
|
/// - `index`: the chunk index (coord → ChunkInfo)
|
|
/// - `ds_dims`: dataset dimensions (e.g. [100, 200])
|
|
/// - `chunk_dims`: chunk dimensions (e.g. [10, 20])
|
|
/// - `elem_size`: size of a single element in bytes
|
|
pub fn build(
|
|
index: &ChunkIndex,
|
|
ds_dims: &[usize],
|
|
chunk_dims: &[usize],
|
|
elem_size: usize,
|
|
) -> Self {
|
|
let rank = ds_dims.len();
|
|
|
|
// Compute dataset strides (row-major)
|
|
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];
|
|
}
|
|
|
|
// Compute chunk strides (row-major within a chunk)
|
|
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_elements: usize = chunk_dims.iter().product();
|
|
let chunk_total_bytes = chunk_total_elements * elem_size;
|
|
let total_elements: usize = ds_dims.iter().product();
|
|
let output_bytes = total_elements * elem_size;
|
|
|
|
let mut mappings = Vec::with_capacity(index.num_chunks());
|
|
|
|
for (_coord, ci) in index.iter() {
|
|
let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
|
|
let chunk_offsets: Vec<usize> = coord.iter().map(|&o| o as usize).collect();
|
|
|
|
let copies = if rank == 0 {
|
|
// Scalar dataset — single copy
|
|
let len = chunk_total_bytes.min(output_bytes);
|
|
vec![RowCopy {
|
|
src_offset: 0,
|
|
dst_offset: 0,
|
|
len,
|
|
}]
|
|
} else {
|
|
compute_row_copies(
|
|
&chunk_offsets,
|
|
chunk_dims,
|
|
ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
)
|
|
};
|
|
|
|
mappings.push(ChunkMapping {
|
|
coord,
|
|
file_offset: ci.address,
|
|
file_size: ci.chunk_size,
|
|
filter_mask: ci.filter_mask,
|
|
copies,
|
|
});
|
|
}
|
|
|
|
Self {
|
|
mappings,
|
|
output_bytes,
|
|
elem_size,
|
|
chunk_total_bytes,
|
|
}
|
|
}
|
|
|
|
/// Assemble output from pre-decompressed chunk buffers using the
|
|
/// pre-computed row-copy plan.
|
|
///
|
|
/// `chunk_data` must be in the same order as `self.mappings`.
|
|
pub fn assemble(&self, chunk_data: &[&[u8]], output: &mut [u8]) {
|
|
debug_assert_eq!(chunk_data.len(), self.mappings.len());
|
|
for (mapping, data) in self.mappings.iter().zip(chunk_data.iter()) {
|
|
for copy in &mapping.copies {
|
|
let src_end = copy.src_offset + copy.len;
|
|
let dst_end = copy.dst_offset + copy.len;
|
|
debug_assert!(
|
|
src_end <= data.len() && dst_end <= output.len(),
|
|
"chunk copy out of bounds: src_end={src_end} data_len={} dst_end={dst_end} out_len={}",
|
|
data.len(),
|
|
output.len(),
|
|
);
|
|
if src_end <= data.len() && dst_end <= output.len() {
|
|
output[copy.dst_offset..dst_end]
|
|
.copy_from_slice(&data[copy.src_offset..src_end]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compute contiguous row-copy operations for a single chunk.
|
|
///
|
|
/// The innermost dimension is always contiguous in both chunk and output
|
|
/// memory (row-major layout). We iterate over the outer dimensions to find
|
|
/// the valid range of the innermost dim, then emit one `RowCopy` per
|
|
/// contiguous run.
|
|
fn compute_row_copies(
|
|
chunk_offsets: &[usize],
|
|
chunk_dims: &[usize],
|
|
ds_dims: &[usize],
|
|
ds_strides: &[usize],
|
|
chunk_strides: &[usize],
|
|
elem_size: usize,
|
|
rank: usize,
|
|
) -> Vec<RowCopy> {
|
|
// Number of "rows" = product of all dimensions except the innermost
|
|
let inner_dim = rank - 1;
|
|
let inner_chunk_len = chunk_dims[inner_dim];
|
|
let inner_ds_len = ds_dims[inner_dim];
|
|
let inner_chunk_offset = chunk_offsets[inner_dim];
|
|
|
|
// How many elements of the innermost dim are valid (not out-of-bounds)?
|
|
let inner_valid = if inner_chunk_offset >= inner_ds_len {
|
|
0
|
|
} else {
|
|
inner_chunk_len.min(inner_ds_len - inner_chunk_offset)
|
|
};
|
|
|
|
if inner_valid == 0 {
|
|
return Vec::new();
|
|
}
|
|
|
|
// Number of outer "rows" (product of dims 0..rank-1)
|
|
let outer_dims = &chunk_dims[..inner_dim];
|
|
let outer_count: usize = outer_dims.iter().product();
|
|
|
|
// Compute outer-only strides (row-major within the outer dimensions only).
|
|
// These are different from chunk_strides which include the innermost dim.
|
|
let mut outer_strides = vec![1usize; inner_dim];
|
|
for i in (0..inner_dim.saturating_sub(1)).rev() {
|
|
outer_strides[i] = outer_strides[i + 1] * outer_dims[i + 1];
|
|
}
|
|
|
|
let mut copies = Vec::with_capacity(outer_count);
|
|
|
|
for outer_flat in 0..outer_count {
|
|
// Convert outer flat index to N-D coordinates in dims 0..rank-1
|
|
let mut remaining = outer_flat;
|
|
let mut ds_flat_base = 0usize;
|
|
let mut chunk_flat_base = 0usize;
|
|
let mut out_of_bounds = false;
|
|
|
|
for d in 0..inner_dim {
|
|
let coord_in_chunk = remaining / outer_strides[d];
|
|
remaining %= outer_strides[d];
|
|
|
|
let global_coord = chunk_offsets[d] + coord_in_chunk;
|
|
if global_coord >= ds_dims[d] {
|
|
out_of_bounds = true;
|
|
break;
|
|
}
|
|
ds_flat_base += global_coord * ds_strides[d];
|
|
chunk_flat_base += coord_in_chunk * chunk_strides[d];
|
|
}
|
|
|
|
if out_of_bounds {
|
|
continue;
|
|
}
|
|
|
|
// Add the innermost dimension contribution
|
|
ds_flat_base += inner_chunk_offset * ds_strides[inner_dim];
|
|
// chunk_flat_base already aligned to start of inner row
|
|
|
|
let src_offset = chunk_flat_base * elem_size;
|
|
let dst_offset = ds_flat_base * elem_size;
|
|
let len = inner_valid * elem_size;
|
|
|
|
copies.push(RowCopy {
|
|
src_offset,
|
|
dst_offset,
|
|
len,
|
|
});
|
|
}
|
|
|
|
copies
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::chunked_read::ChunkInfo;
|
|
|
|
fn make_chunk(offsets: Vec<u64>, address: u64, size: u32) -> ChunkInfo {
|
|
ChunkInfo {
|
|
chunk_size: size,
|
|
filter_mask: 0,
|
|
offsets,
|
|
address,
|
|
}
|
|
}
|
|
|
|
// --- ChunkIndex tests ---
|
|
|
|
#[test]
|
|
fn index_build_and_lookup() {
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 0, 4], 0x1000, 80),
|
|
make_chunk(vec![10, 0, 4], 0x2000, 80),
|
|
make_chunk(vec![10, 5, 4], 0x3000, 80),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 2);
|
|
assert_eq!(index.num_chunks(), 3);
|
|
assert_eq!(index.rank(), 2);
|
|
|
|
let c0 = index.lookup(&[0, 0]).unwrap();
|
|
assert_eq!(c0.address, 0x1000);
|
|
|
|
let c1 = index.lookup(&[10, 0]).unwrap();
|
|
assert_eq!(c1.address, 0x2000);
|
|
|
|
let c2 = index.lookup(&[10, 5]).unwrap();
|
|
assert_eq!(c2.address, 0x3000);
|
|
|
|
assert!(index.lookup(&[5, 0]).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn index_empty() {
|
|
let index = ChunkIndex::build(&[], 2);
|
|
assert_eq!(index.num_chunks(), 0);
|
|
assert!(index.lookup(&[0, 0]).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn index_1d() {
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 8], 0x100, 40),
|
|
make_chunk(vec![5, 8], 0x200, 40),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 1);
|
|
assert_eq!(index.num_chunks(), 2);
|
|
|
|
assert_eq!(index.lookup(&[0]).unwrap().address, 0x100);
|
|
assert_eq!(index.lookup(&[5]).unwrap().address, 0x200);
|
|
}
|
|
|
|
#[test]
|
|
fn index_3d() {
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 0, 0, 4], 0x100, 100),
|
|
make_chunk(vec![0, 0, 5, 4], 0x200, 100),
|
|
make_chunk(vec![2, 3, 0, 4], 0x300, 100),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 3);
|
|
assert_eq!(index.num_chunks(), 3);
|
|
|
|
assert_eq!(index.lookup(&[0, 0, 0]).unwrap().address, 0x100);
|
|
assert_eq!(index.lookup(&[0, 0, 5]).unwrap().address, 0x200);
|
|
assert_eq!(index.lookup(&[2, 3, 0]).unwrap().address, 0x300);
|
|
}
|
|
|
|
// --- ChunkLayout tests ---
|
|
|
|
#[test]
|
|
fn layout_1d_two_chunks() {
|
|
// 20-element 1D dataset, chunk size 10, f64
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 8], 0x1000, 80),
|
|
make_chunk(vec![10, 8], 0x2000, 80),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 1);
|
|
let layout = ChunkLayout::build(&index, &[20], &[10], 8);
|
|
|
|
assert_eq!(layout.output_bytes, 160); // 20 * 8
|
|
assert_eq!(layout.mappings.len(), 2);
|
|
|
|
// Each 1D chunk should have exactly 1 row-copy
|
|
for mapping in &layout.mappings {
|
|
assert_eq!(mapping.copies.len(), 1);
|
|
assert_eq!(mapping.copies[0].len, 80); // 10 * 8
|
|
}
|
|
|
|
// Verify assembly
|
|
let chunk0: Vec<u8> = (0..10u64).flat_map(|i| (i as f64).to_le_bytes()).collect();
|
|
let chunk1: Vec<u8> = (10..20u64).flat_map(|i| (i as f64).to_le_bytes()).collect();
|
|
|
|
let mut output = vec![0u8; 160];
|
|
// Find which mapping goes where
|
|
let data_refs: Vec<&[u8]> = layout
|
|
.mappings
|
|
.iter()
|
|
.map(|m| {
|
|
if m.coord == vec![0] {
|
|
chunk0.as_slice()
|
|
} else {
|
|
chunk1.as_slice()
|
|
}
|
|
})
|
|
.collect();
|
|
layout.assemble(&data_refs, &mut output);
|
|
|
|
for i in 0..20u64 {
|
|
let val = f64::from_le_bytes(
|
|
output[i as usize * 8..(i as usize + 1) * 8]
|
|
.try_into()
|
|
.unwrap(),
|
|
);
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn layout_2d_four_chunks() {
|
|
// 4x6 dataset, chunk 2x3, f32 (4 bytes)
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 0, 4], 0x100, 24),
|
|
make_chunk(vec![0, 3, 4], 0x200, 24),
|
|
make_chunk(vec![2, 0, 4], 0x300, 24),
|
|
make_chunk(vec![2, 3, 4], 0x400, 24),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 2);
|
|
let layout = ChunkLayout::build(&index, &[4, 6], &[2, 3], 4);
|
|
|
|
assert_eq!(layout.output_bytes, 96); // 24 * 4
|
|
assert_eq!(layout.mappings.len(), 4);
|
|
|
|
// Each 2x3 chunk should have 2 row-copies (one per row)
|
|
for mapping in &layout.mappings {
|
|
assert_eq!(mapping.copies.len(), 2);
|
|
// Each copy should be 3 elements * 4 bytes = 12 bytes
|
|
for copy in &mapping.copies {
|
|
assert_eq!(copy.len, 12);
|
|
}
|
|
}
|
|
|
|
// Build chunk data and verify assembly
|
|
let mut chunk_data_map: Vec<(Vec<u64>, Vec<u8>)> = Vec::new();
|
|
for row_start in [0usize, 2] {
|
|
for col_start in [0usize, 3] {
|
|
let mut data = Vec::new();
|
|
for r in 0..2 {
|
|
for c in 0..3 {
|
|
let val = ((row_start + r) * 6 + (col_start + c)) as f32;
|
|
data.extend_from_slice(&val.to_le_bytes());
|
|
}
|
|
}
|
|
chunk_data_map.push((vec![row_start as u64, col_start as u64], data));
|
|
}
|
|
}
|
|
|
|
let data_refs: Vec<&[u8]> = layout
|
|
.mappings
|
|
.iter()
|
|
.map(|m| {
|
|
chunk_data_map
|
|
.iter()
|
|
.find(|(c, _)| c == &m.coord)
|
|
.unwrap()
|
|
.1
|
|
.as_slice()
|
|
})
|
|
.collect();
|
|
|
|
let mut output = vec![0u8; 96];
|
|
layout.assemble(&data_refs, &mut output);
|
|
|
|
for i in 0..24 {
|
|
let val = f32::from_le_bytes(output[i * 4..(i + 1) * 4].try_into().unwrap());
|
|
assert_eq!(val, i as f32, "mismatch at element {i}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn layout_partial_last_chunk() {
|
|
// 25-element 1D, chunk 10 => 3 chunks, last has 5 valid
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 8], 0x100, 80),
|
|
make_chunk(vec![10, 8], 0x200, 80),
|
|
make_chunk(vec![20, 8], 0x300, 80),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 1);
|
|
let layout = ChunkLayout::build(&index, &[25], &[10], 8);
|
|
|
|
assert_eq!(layout.output_bytes, 200); // 25 * 8
|
|
|
|
// Find the chunk at offset 20 — should only copy 5 elements
|
|
let partial = layout
|
|
.mappings
|
|
.iter()
|
|
.find(|m| m.coord == vec![20])
|
|
.unwrap();
|
|
assert_eq!(partial.copies.len(), 1);
|
|
assert_eq!(partial.copies[0].len, 40); // 5 * 8 bytes
|
|
}
|
|
|
|
#[test]
|
|
fn layout_2d_partial_boundary() {
|
|
// 5x5 dataset, chunk 3x3 => 4 chunks, some partial
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 0, 4], 0x100, 36),
|
|
make_chunk(vec![0, 3, 4], 0x200, 36),
|
|
make_chunk(vec![3, 0, 4], 0x300, 36),
|
|
make_chunk(vec![3, 3, 4], 0x400, 36),
|
|
];
|
|
let index = ChunkIndex::build(&chunks, 2);
|
|
let layout = ChunkLayout::build(&index, &[5, 5], &[3, 3], 4);
|
|
|
|
assert_eq!(layout.output_bytes, 100); // 25 * 4
|
|
|
|
// Chunk at (0,0): 3x3 fully valid → 3 rows, 3 elems each
|
|
let c00 = layout
|
|
.mappings
|
|
.iter()
|
|
.find(|m| m.coord == vec![0, 0])
|
|
.unwrap();
|
|
assert_eq!(c00.copies.len(), 3);
|
|
for copy in &c00.copies {
|
|
assert_eq!(copy.len, 12); // 3 * 4
|
|
}
|
|
|
|
// Chunk at (0,3): 3x2 valid (cols 3,4) → 3 rows, 2 elems each
|
|
let c03 = layout
|
|
.mappings
|
|
.iter()
|
|
.find(|m| m.coord == vec![0, 3])
|
|
.unwrap();
|
|
assert_eq!(c03.copies.len(), 3);
|
|
for copy in &c03.copies {
|
|
assert_eq!(copy.len, 8); // 2 * 4
|
|
}
|
|
|
|
// Chunk at (3,0): 2x3 valid (rows 3,4) → 2 rows, 3 elems each
|
|
let c30 = layout
|
|
.mappings
|
|
.iter()
|
|
.find(|m| m.coord == vec![3, 0])
|
|
.unwrap();
|
|
assert_eq!(c30.copies.len(), 2);
|
|
|
|
// Chunk at (3,3): 2x2 valid → 2 rows, 2 elems each
|
|
let c33 = layout
|
|
.mappings
|
|
.iter()
|
|
.find(|m| m.coord == vec![3, 3])
|
|
.unwrap();
|
|
assert_eq!(c33.copies.len(), 2);
|
|
for copy in &c33.copies {
|
|
assert_eq!(copy.len, 8); // 2 * 4
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn layout_3d_basic() {
|
|
// 4x4x4 dataset, chunk 2x2x2, f32 (4 bytes) => 8 chunks
|
|
let mut chunks = Vec::new();
|
|
let elem_size = 4u32;
|
|
let chunk_bytes = 2 * 2 * 2 * elem_size;
|
|
let mut addr = 0x1000u64;
|
|
for x in [0u64, 2] {
|
|
for y in [0u64, 2] {
|
|
for z in [0u64, 2] {
|
|
chunks.push(make_chunk(
|
|
vec![x, y, z, elem_size as u64],
|
|
addr,
|
|
chunk_bytes,
|
|
));
|
|
addr += chunk_bytes as u64;
|
|
}
|
|
}
|
|
}
|
|
|
|
let index = ChunkIndex::build(&chunks, 3);
|
|
let layout = ChunkLayout::build(&index, &[4, 4, 4], &[2, 2, 2], 4);
|
|
|
|
assert_eq!(layout.output_bytes, 256); // 64 * 4
|
|
assert_eq!(layout.mappings.len(), 8);
|
|
|
|
// Each 2x2x2 chunk: outer dims are 2x2=4 rows, each row copies 2 elements
|
|
for mapping in &layout.mappings {
|
|
assert_eq!(mapping.copies.len(), 4); // 2*2 outer rows
|
|
for copy in &mapping.copies {
|
|
assert_eq!(copy.len, 8); // 2 * 4 bytes
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn assemble_matches_element_by_element() {
|
|
// Verify that layout assembly produces identical output to the
|
|
// original copy_chunk_to_output approach (2D case)
|
|
let ds_dims = [6usize, 8];
|
|
let chunk_dims = [3usize, 4];
|
|
let elem_size = 8usize; // f64
|
|
|
|
let mut chunks = Vec::new();
|
|
let mut addr = 0x1000u64;
|
|
for r in (0..6).step_by(3) {
|
|
for c in (0..8).step_by(4) {
|
|
chunks.push(make_chunk(
|
|
vec![r as u64, c as u64, elem_size as u64],
|
|
addr,
|
|
(3 * 4 * elem_size) as u32,
|
|
));
|
|
addr += (3 * 4 * elem_size) as u64;
|
|
}
|
|
}
|
|
|
|
let index = ChunkIndex::build(&chunks, 2);
|
|
let layout = ChunkLayout::build(&index, &ds_dims, &chunk_dims, elem_size);
|
|
|
|
// Build chunk data
|
|
let mut chunk_buffers: Vec<(Vec<u64>, Vec<u8>)> = Vec::new();
|
|
for mapping in &layout.mappings {
|
|
let row_start = mapping.coord[0] as usize;
|
|
let col_start = mapping.coord[1] as usize;
|
|
let mut buf = vec![0u8; 3 * 4 * elem_size];
|
|
for r in 0..3 {
|
|
for c in 0..4 {
|
|
let val = ((row_start + r) * 8 + (col_start + c)) as f64;
|
|
let offset = (r * 4 + c) * elem_size;
|
|
buf[offset..offset + 8].copy_from_slice(&val.to_le_bytes());
|
|
}
|
|
}
|
|
chunk_buffers.push((mapping.coord.clone(), buf));
|
|
}
|
|
|
|
let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|(_, b)| b.as_slice()).collect();
|
|
let mut output = vec![0u8; 6 * 8 * elem_size];
|
|
layout.assemble(&data_refs, &mut output);
|
|
|
|
// Verify every element
|
|
for r in 0..6 {
|
|
for c in 0..8 {
|
|
let idx = r * 8 + c;
|
|
let val = f64::from_le_bytes(output[idx * 8..(idx + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, idx as f64, "mismatch at ({r}, {c})");
|
|
}
|
|
}
|
|
}
|
|
}
|