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 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 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}