Skip to main content

hdf5_pure_rust/hl/mutable_file/
resize.rs

1use std::io::{Read, Seek, SeekFrom, Write};
2
3use crate::error::{Error, Result};
4use crate::format::messages::data_layout::{ChunkIndexType, LayoutClass};
5use crate::format::messages::fill_value::FILL_TIME_NEVER;
6use crate::format::object_header::{
7    self, HDR_ATTR_STORE_PHASE_CHANGE, HDR_CHUNK0_SIZE_MASK, HDR_STORE_TIMES, HDR_V2_KNOWN_FLAGS,
8};
9
10use super::MutableFile;
11
12pub(super) struct ObjectHeaderMessageLocation {
13    pub(super) msg_data_offset: u64,
14    pub(super) msg_data_len: usize,
15    pub(super) v2_checksum: Option<(u64, usize)>,
16}
17
18impl MutableFile {
19    /// Resize a chunked dataset to new dimensions.
20    ///
21    /// The new dimensions must not exceed the dataset's maximum dimensions.
22    /// Whole v2-B-tree chunks outside a shrunken extent are pruned; other
23    /// chunk data is left in place.
24    /// New regions are filled with the default fill value (zeros).
25    pub fn resize_dataset(&mut self, path: &str, new_dims: &[u64]) -> Result<()> {
26        let ds = self.dataset(path)?;
27        let info = ds.info()?;
28
29        if info.layout.layout_class != LayoutClass::Chunked {
30            return Err(Error::InvalidFormat(
31                "can only resize chunked datasets".into(),
32            ));
33        }
34
35        if new_dims.len() != info.dataspace.dims.len() {
36            return Err(Error::InvalidFormat(format!(
37                "dimension count mismatch: dataset has {} dims, new shape has {}",
38                info.dataspace.dims.len(),
39                new_dims.len()
40            )));
41        }
42
43        if let Some(ref max_dims) = info.dataspace.max_dims {
44            for (i, (&new_d, &max_d)) in new_dims.iter().zip(max_dims.iter()).enumerate() {
45                if max_d != u64::MAX && new_d > max_d {
46                    return Err(Error::InvalidFormat(format!(
47                        "dimension {i}: new size {new_d} exceeds max {max_d}"
48                    )));
49                }
50            }
51        }
52
53        let shrinks = new_dims
54            .iter()
55            .zip(&info.dataspace.dims)
56            .any(|(&new_dim, &old_dim)| new_dim < old_dim);
57        if shrinks
58            && (info.layout.version <= 3
59                || matches!(
60                    info.layout.chunk_index_type,
61                    Some(ChunkIndexType::BTreeV1)
62                        | Some(ChunkIndexType::BTreeV2)
63                        | Some(ChunkIndexType::FixedArray)
64                        | Some(ChunkIndexType::ExtensibleArray)
65                ))
66        {
67            let chunk_dims = Self::chunk_data_dims(&info)?.to_vec();
68            let index_addr = info.layout.chunk_index_addr.ok_or_else(|| {
69                Error::InvalidFormat("chunked dataset missing chunk index address".into())
70            })?;
71            self.scrub_partial_shrink_chunks(path, &ds, &info, new_dims, &chunk_dims)?;
72            if info.layout.chunk_index_type == Some(ChunkIndexType::BTreeV2) {
73                self.prune_btree_v2_chunks_outside_extent(
74                    index_addr,
75                    &info,
76                    new_dims,
77                    &chunk_dims,
78                )?;
79            }
80            let physical_eof = self.write_handle.seek(SeekFrom::End(0))?;
81            let logical_eof = physical_eof
82                .checked_sub(self.superblock.base_addr)
83                .ok_or_else(|| {
84                    Error::InvalidFormat("file EOF is before HDF5 base address".into())
85                })?;
86            self.rewrite_superblock_eof(logical_eof)?;
87            self.reopen_reader()?;
88        }
89
90        let ds_addr = ds.addr();
91        let location = self.find_message_in_oh(ds_addr, object_header::MSG_DATASPACE)?;
92
93        let mut new_space = info.dataspace.clone();
94        new_space.dims = new_dims.to_vec();
95        let new_ds_bytes = new_space.encode()?;
96
97        if new_ds_bytes.len() != location.msg_data_len {
98            return Err(Error::InvalidFormat(format!(
99                "dataspace message size changed ({} -> {}); in-place resize not possible",
100                location.msg_data_len,
101                new_ds_bytes.len()
102            )));
103        }
104
105        let msg_data_offset = self.physical_addr(location.msg_data_offset, "dataspace message")?;
106        self.write_handle.seek(SeekFrom::Start(msg_data_offset))?;
107        self.write_handle.write_all(&new_ds_bytes)?;
108        if let Some((oh_start, oh_check_len)) = location.v2_checksum {
109            self.rewrite_oh_checksum(oh_start, oh_check_len)?;
110        }
111        self.write_handle.flush()?;
112        self.reopen_reader()?;
113
114        Ok(())
115    }
116
117    fn scrub_partial_shrink_chunks(
118        &mut self,
119        path: &str,
120        ds: &crate::hl::dataset::Dataset,
121        info: &crate::hl::dataset::DatasetInfo,
122        new_dims: &[u64],
123        chunk_dims: &[u64],
124    ) -> Result<()> {
125        if new_dims.len() != info.dataspace.dims.len() || chunk_dims.len() != new_dims.len() {
126            return Ok(());
127        }
128        if chunk_dims.contains(&0) {
129            return Ok(());
130        }
131        let old_dims = &info.dataspace.dims;
132        let has_partial_boundary = old_dims.iter().zip(new_dims).zip(chunk_dims).any(
133            |((&old_dim, &new_dim), &chunk_dim)| new_dim < old_dim && new_dim % chunk_dim != 0,
134        );
135        if !has_partial_boundary {
136            return Ok(());
137        }
138
139        let element_size = Self::u64_to_usize(u64::from(info.datatype.size), "datatype size")?;
140        let old_elements = Self::checked_product_u64(old_dims, "dataset element count")?;
141        let total_bytes = Self::u64_to_usize(old_elements, "dataset element count")?
142            .checked_mul(element_size)
143            .ok_or_else(|| Error::InvalidFormat("dataset byte size overflow".into()))?;
144        let raw = ds.read_raw()?;
145        if raw.len() != total_bytes {
146            return Err(Error::InvalidFormat(format!(
147                "raw dataset read returned {} bytes, expected {total_bytes}",
148                raw.len()
149            )));
150        }
151
152        let chunk_elements = Self::checked_product_u64(chunk_dims, "chunk element count")?;
153        let chunk_bytes = Self::u64_to_usize(chunk_elements, "chunk element count")?
154            .checked_mul(element_size)
155            .ok_or_else(|| Error::InvalidFormat("chunk byte size overflow".into()))?;
156        let fill = Self::fill_value_bytes(info, element_size)?;
157        let old_strides = Self::row_major_strides(old_dims)?;
158        let chunk_strides = Self::row_major_strides(chunk_dims)?;
159        let mut starts = Self::zeroed_vec(new_dims.len(), "partial-shrink chunk starts")?;
160        self.scrub_partial_shrink_chunks_recursive(
161            path,
162            0,
163            &mut starts,
164            old_dims,
165            new_dims,
166            chunk_dims,
167            &old_strides,
168            &chunk_strides,
169            &raw,
170            chunk_bytes,
171            element_size,
172            &fill,
173        )
174    }
175
176    #[allow(clippy::too_many_arguments)]
177    fn scrub_partial_shrink_chunks_recursive(
178        &mut self,
179        path: &str,
180        dim: usize,
181        starts: &mut [u64],
182        old_dims: &[u64],
183        new_dims: &[u64],
184        chunk_dims: &[u64],
185        old_strides: &[u64],
186        chunk_strides: &[u64],
187        raw: &[u8],
188        chunk_bytes: usize,
189        element_size: usize,
190        fill: &[u8],
191    ) -> Result<()> {
192        if dim < starts.len() {
193            let mut start = 0;
194            while start < old_dims[dim] {
195                starts[dim] = start;
196                self.scrub_partial_shrink_chunks_recursive(
197                    path,
198                    dim + 1,
199                    starts,
200                    old_dims,
201                    new_dims,
202                    chunk_dims,
203                    old_strides,
204                    chunk_strides,
205                    raw,
206                    chunk_bytes,
207                    element_size,
208                    fill,
209                )?;
210                start = start.checked_add(chunk_dims[dim]).ok_or_else(|| {
211                    Error::InvalidFormat("chunk start coordinate overflow".into())
212                })?;
213            }
214            return Ok(());
215        }
216
217        if starts
218            .iter()
219            .zip(new_dims)
220            .any(|(&start, &new_dim)| start >= new_dim)
221        {
222            return Ok(());
223        }
224        let boundary = starts
225            .iter()
226            .zip(old_dims)
227            .zip(new_dims)
228            .zip(chunk_dims)
229            .any(|(((&start, &old_dim), &new_dim), &chunk_dim)| {
230                new_dim < old_dim
231                    && start < new_dim
232                    && new_dim < start.saturating_add(chunk_dim).min(old_dim)
233            });
234        if !boundary {
235            return Ok(());
236        }
237
238        let mut chunk = Self::zeroed_vec(chunk_bytes, "partial-shrink chunk buffer")?;
239        for slot in chunk.chunks_exact_mut(element_size) {
240            slot.copy_from_slice(fill);
241        }
242        let mut local = Self::zeroed_vec(starts.len(), "partial-shrink local coordinates")?;
243        Self::copy_retained_chunk_elements(
244            0,
245            starts,
246            &mut local,
247            old_dims,
248            new_dims,
249            chunk_dims,
250            old_strides,
251            chunk_strides,
252            raw,
253            &mut chunk,
254            element_size,
255        )?;
256        self.write_chunk(path, starts, &chunk)
257    }
258
259    #[allow(clippy::too_many_arguments)]
260    fn copy_retained_chunk_elements(
261        dim: usize,
262        starts: &[u64],
263        local: &mut [u64],
264        old_dims: &[u64],
265        new_dims: &[u64],
266        chunk_dims: &[u64],
267        old_strides: &[u64],
268        chunk_strides: &[u64],
269        raw: &[u8],
270        chunk: &mut [u8],
271        element_size: usize,
272    ) -> Result<()> {
273        if dim < local.len() {
274            for coord in 0..chunk_dims[dim] {
275                local[dim] = coord;
276                Self::copy_retained_chunk_elements(
277                    dim + 1,
278                    starts,
279                    local,
280                    old_dims,
281                    new_dims,
282                    chunk_dims,
283                    old_strides,
284                    chunk_strides,
285                    raw,
286                    chunk,
287                    element_size,
288                )?;
289            }
290            return Ok(());
291        }
292
293        let mut old_index = 0u64;
294        let mut chunk_index = 0u64;
295        for idx in 0..local.len() {
296            let global = starts[idx]
297                .checked_add(local[idx])
298                .ok_or_else(|| Error::InvalidFormat("chunk coordinate overflow".into()))?;
299            if global >= old_dims[idx] {
300                return Ok(());
301            }
302            if global >= new_dims[idx] {
303                return Ok(());
304            }
305            old_index =
306                old_index
307                    .checked_add(global.checked_mul(old_strides[idx]).ok_or_else(|| {
308                        Error::InvalidFormat("dataset linear index overflow".into())
309                    })?)
310                    .ok_or_else(|| Error::InvalidFormat("dataset linear index overflow".into()))?;
311            chunk_index =
312                chunk_index
313                    .checked_add(local[idx].checked_mul(chunk_strides[idx]).ok_or_else(|| {
314                        Error::InvalidFormat("chunk linear index overflow".into())
315                    })?)
316                    .ok_or_else(|| Error::InvalidFormat("chunk linear index overflow".into()))?;
317        }
318
319        let old_byte = Self::u64_to_usize(old_index, "dataset linear index")?
320            .checked_mul(element_size)
321            .ok_or_else(|| Error::InvalidFormat("dataset byte offset overflow".into()))?;
322        let chunk_byte = Self::u64_to_usize(chunk_index, "chunk linear index")?
323            .checked_mul(element_size)
324            .ok_or_else(|| Error::InvalidFormat("chunk byte offset overflow".into()))?;
325        let old_end = old_byte
326            .checked_add(element_size)
327            .ok_or_else(|| Error::InvalidFormat("dataset byte range overflow".into()))?;
328        let chunk_end = chunk_byte
329            .checked_add(element_size)
330            .ok_or_else(|| Error::InvalidFormat("chunk byte range overflow".into()))?;
331        let old_slice = raw.get(old_byte..old_end).ok_or_else(|| {
332            Error::InvalidFormat("dataset retained byte range out of bounds".into())
333        })?;
334        let chunk_slice = chunk.get_mut(chunk_byte..chunk_end).ok_or_else(|| {
335            Error::InvalidFormat("chunk retained byte range out of bounds".into())
336        })?;
337        chunk_slice.copy_from_slice(old_slice);
338        Ok(())
339    }
340
341    fn zeroed_vec<T>(len: usize, context: &str) -> Result<Vec<T>>
342    where
343        T: Default + Clone,
344    {
345        let mut out = Vec::new();
346        out.try_reserve_exact(len)
347            .map_err(|err| Error::InvalidFormat(format!("{context} allocation failed: {err}")))?;
348        out.resize(len, T::default());
349        Ok(out)
350    }
351
352    fn checked_product_u64(values: &[u64], what: &str) -> Result<u64> {
353        values.iter().try_fold(1u64, |acc, &value| {
354            acc.checked_mul(value)
355                .ok_or_else(|| Error::InvalidFormat(format!("{what} overflow")))
356        })
357    }
358
359    fn row_major_strides(dims: &[u64]) -> Result<Vec<u64>> {
360        let mut strides = vec![1; dims.len()];
361        let mut stride = 1u64;
362        for idx in (0..dims.len()).rev() {
363            strides[idx] = stride;
364            stride = stride
365                .checked_mul(dims[idx])
366                .ok_or_else(|| Error::InvalidFormat("row-major stride overflow".into()))?;
367        }
368        Ok(strides)
369    }
370
371    fn fill_value_bytes(
372        info: &crate::hl::dataset::DatasetInfo,
373        element_size: usize,
374    ) -> Result<Vec<u8>> {
375        let Some(fill) = &info.fill_value else {
376            return Self::zeroed_vec(element_size, "default fill value");
377        };
378        if fill.fill_time == FILL_TIME_NEVER {
379            return Self::zeroed_vec(element_size, "default fill value");
380        }
381        let Some(value) = fill.value.as_deref() else {
382            return Self::zeroed_vec(element_size, "default fill value");
383        };
384        if value.len() != element_size {
385            return Err(Error::Unsupported(format!(
386                "fill value size {} does not match element size {}",
387                value.len(),
388                element_size
389            )));
390        }
391        Ok(value.to_vec())
392    }
393
394    /// Find a message of the given type in an object header.
395    pub(super) fn find_message_in_oh(
396        &self,
397        oh_addr: u64,
398        target_msg_type: u16,
399    ) -> Result<ObjectHeaderMessageLocation> {
400        let mut guard = self.inner.lock();
401        let reader = &mut guard.reader;
402        reader.seek(oh_addr)?;
403
404        let mut first_bytes = [0u8; 4];
405        reader.read_bytes_into(&mut first_bytes)?;
406        if first_bytes != [b'O', b'H', b'D', b'R'] {
407            reader.seek(oh_addr)?;
408            return Self::find_message_in_v1_oh(reader, target_msg_type);
409        }
410
411        let version = reader.read_u8()?;
412        if version != 2 {
413            return Err(Error::Unsupported(
414                "resize only supported for v2 object headers".into(),
415            ));
416        }
417
418        let flags = reader.read_u8()?;
419        if flags & !HDR_V2_KNOWN_FLAGS != 0 {
420            return Err(Error::InvalidFormat(format!(
421                "object header v2 flags contain reserved bits: {flags:#04x}"
422            )));
423        }
424        if flags & HDR_STORE_TIMES != 0 {
425            reader.skip(16)?;
426        }
427        if flags & HDR_ATTR_STORE_PHASE_CHANGE != 0 {
428            reader.skip(4)?;
429        }
430
431        let chunk0_size_bytes = 1u8 << (flags & HDR_CHUNK0_SIZE_MASK);
432        let chunk0_data_size = reader.read_uint(chunk0_size_bytes)?;
433        let chunk0_data_start = reader.position()?;
434        let chunk0_data_end = chunk0_data_start
435            .checked_add(chunk0_data_size)
436            .ok_or_else(|| Error::InvalidFormat("object-header chunk range overflow".into()))?;
437        let oh_check_len = usize::try_from(chunk0_data_end - oh_addr)
438            .map_err(|_| Error::InvalidFormat("object-header checksum range overflow".into()))?;
439
440        while reader.position()? < chunk0_data_end {
441            let msg_header_pos = reader.position()?;
442            if msg_header_pos
443                .checked_add(4)
444                .is_none_or(|end| end > chunk0_data_end)
445            {
446                break;
447            }
448
449            let msg_type = u16::from(reader.read_u8()?);
450            let msg_size = usize::from(reader.read_u16()?);
451            let _msg_flags = reader.read_u8()?;
452
453            if flags & object_header::HDR_ATTR_CRT_ORDER_TRACKED != 0 {
454                reader.skip(2)?;
455            }
456
457            let msg_data_offset = reader.position()?;
458            let msg_size_u64 = Self::usize_to_u64(msg_size, "object-header message size")?;
459            if msg_data_offset
460                .checked_add(msg_size_u64)
461                .is_none_or(|end| end > chunk0_data_end)
462            {
463                return Err(Error::InvalidFormat(
464                    "object-header message payload exceeds chunk".into(),
465                ));
466            }
467
468            if msg_type == target_msg_type {
469                return Ok(ObjectHeaderMessageLocation {
470                    msg_data_offset,
471                    msg_data_len: msg_size,
472                    v2_checksum: Some((oh_addr, oh_check_len)),
473                });
474            }
475
476            reader.skip(msg_size_u64)?;
477        }
478
479        Err(Error::InvalidFormat(format!(
480            "message type {target_msg_type:#06x} not found in object header"
481        )))
482    }
483
484    fn find_message_in_v1_oh<R: Read + Seek>(
485        reader: &mut crate::io::reader::HdfReader<R>,
486        target_msg_type: u16,
487    ) -> Result<ObjectHeaderMessageLocation> {
488        let version = reader.read_u8()?;
489        if version != 1 {
490            return Err(Error::Unsupported(format!(
491                "resize only supports v1/v2 object headers, got version {version}"
492            )));
493        }
494        reader.skip(1)?;
495        let _num_messages = reader.read_u16()?;
496        let _refcount = reader.read_u32()?;
497        let chunk_data_size = u64::from(reader.read_u32()?);
498        reader.skip(4)?;
499
500        let chunk_start = reader.position()?;
501        let chunk_end = chunk_start
502            .checked_add(chunk_data_size)
503            .ok_or_else(|| Error::InvalidFormat("object header v1 chunk range overflow".into()))?;
504        let mut continuations = Vec::new();
505        if let Some(location) = Self::find_message_in_v1_message_chunk(
506            reader,
507            chunk_end,
508            target_msg_type,
509            &mut continuations,
510        )? {
511            return Ok(location);
512        }
513
514        while let Some((addr, len)) = continuations.pop() {
515            reader.seek(addr)?;
516            let chunk_end = addr.checked_add(len).ok_or_else(|| {
517                Error::InvalidFormat("object header v1 continuation range overflow".into())
518            })?;
519            if let Some(location) = Self::find_message_in_v1_message_chunk(
520                reader,
521                chunk_end,
522                target_msg_type,
523                &mut continuations,
524            )? {
525                return Ok(location);
526            }
527        }
528
529        Err(Error::InvalidFormat(format!(
530            "message type {target_msg_type:#06x} not found in object header"
531        )))
532    }
533
534    fn find_message_in_v1_message_chunk<R: Read + Seek>(
535        reader: &mut crate::io::reader::HdfReader<R>,
536        chunk_end: u64,
537        target_msg_type: u16,
538        continuations: &mut Vec<(u64, u64)>,
539    ) -> Result<Option<ObjectHeaderMessageLocation>> {
540        while reader.position()? < chunk_end {
541            let msg_header_pos = reader.position()?;
542            if msg_header_pos
543                .checked_add(8)
544                .is_none_or(|end| end > chunk_end)
545            {
546                break;
547            }
548
549            let msg_type = reader.read_u16()?;
550            let msg_size = usize::from(reader.read_u16()?);
551            let _msg_flags = reader.read_u8()?;
552            reader.skip(3)?;
553            let msg_data_offset = reader.position()?;
554            let aligned_size = Self::align_u64(
555                Self::usize_to_u64(msg_size, "object-header message size")?,
556                8,
557                "object-header message alignment",
558            )?;
559            let msg_end = msg_data_offset.checked_add(aligned_size).ok_or_else(|| {
560                Error::InvalidFormat("object-header message range overflow".into())
561            })?;
562            if msg_end > chunk_end {
563                return Err(Error::InvalidFormat(
564                    "object-header message payload exceeds chunk".into(),
565                ));
566            }
567
568            if msg_type == target_msg_type {
569                return Ok(Some(ObjectHeaderMessageLocation {
570                    msg_data_offset,
571                    msg_data_len: msg_size,
572                    v2_checksum: None,
573                }));
574            }
575
576            if msg_type == object_header::MSG_HEADER_CONTINUATION {
577                let used = u64::from(reader.sizeof_addr())
578                    .checked_add(u64::from(reader.sizeof_size()))
579                    .ok_or_else(|| {
580                        Error::InvalidFormat("object-header continuation width overflow".into())
581                    })?;
582                let msg_size_u64 = Self::usize_to_u64(msg_size, "object-header message size")?;
583                if msg_size_u64 < used {
584                    return Err(Error::InvalidFormat(
585                        "object-header continuation message is truncated".into(),
586                    ));
587                }
588                let cont_addr = reader.read_addr()?;
589                let cont_len = reader.read_length()?;
590                continuations.push((cont_addr, cont_len));
591            }
592
593            reader.seek(msg_end)?;
594        }
595        Ok(None)
596    }
597
598    fn align_u64(value: u64, align: u64, what: &str) -> Result<u64> {
599        let mask = align
600            .checked_sub(1)
601            .ok_or_else(|| Error::InvalidFormat(format!("{what} overflow")))?;
602        value
603            .checked_add(mask)
604            .map(|value| value & !mask)
605            .ok_or_else(|| Error::InvalidFormat(format!("{what} overflow")))
606    }
607}
608
609#[cfg(test)]
610mod tests {
611    use crate::format::messages::dataspace::{DataspaceMessage, DataspaceType};
612
613    #[test]
614    fn v1_dataspace_encode_preserves_message_layout() {
615        let space = DataspaceMessage {
616            version: 1,
617            space_type: DataspaceType::Simple,
618            ndims: 1,
619            dims: vec![322],
620            max_dims: Some(vec![u64::MAX]),
621        };
622        let encoded = space.encode().unwrap();
623        assert_eq!(encoded.len(), 24);
624        assert_eq!(DataspaceMessage::decode(&encoded).unwrap(), space);
625    }
626}