trx_rs/io/
zip.rs

1use std::collections::{BTreeSet, HashMap};
2use std::fs;
3use std::io::{BufWriter, Read, Write};
4use std::path::Path;
5use std::sync::Arc;
6use zip::write::SimpleFileOptions;
7
8use super::archive_edit::{self, ArchiveOp};
9use crate::dtype::TrxScalar;
10use crate::error::{Result, TrxError};
11use crate::header::Header;
12use crate::io::filename::TrxFilename;
13use crate::mmap_backing::{vec_to_bytes, MmapBacking};
14use crate::trx_file::{DataArray, DataPerGroup, TrxFile, TrxParts};
15
16/// On-disk width for the `offsets.*` array. The TRX spec accepts both
17/// `uint32` and `uint64`; we auto-pick at write time via [`pick_for`] so
18/// every-day tractograms (≤ 4 G vertices) stay compact and only genuinely
19/// huge files pay the doubled-width cost.
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub(crate) enum OffsetsDtype {
22    U32,
23    // Reserved for the day the in-memory offset representation widens past
24    // u32. The reader already accepts `offsets.uint64`; the writer will
25    // start emitting it once `pick_for` can return this variant.
26    #[allow(dead_code)]
27    U64,
28}
29
30impl OffsetsDtype {
31    /// Pick the narrowest dtype that fits every offset in the slice.
32    pub(crate) fn pick_for(offsets: &[u32]) -> Self {
33        // The in-memory representation is `&[u32]`, so by definition every
34        // value fits in `u32`. We still keep this function as the canonical
35        // place to widen the rule if the in-memory type ever changes.
36        let _ = offsets;
37        OffsetsDtype::U32
38    }
39
40    /// Filename suffix written for this dtype (e.g. `"uint64"`).
41    pub(crate) fn suffix(self) -> &'static str {
42        match self {
43            OffsetsDtype::U32 => "uint32",
44            OffsetsDtype::U64 => "uint64",
45        }
46    }
47
48    /// Serialise an in-memory `u32` offset slice to disk bytes at this width.
49    pub(crate) fn encode(self, offsets: &[u32]) -> Vec<u8> {
50        match self {
51            OffsetsDtype::U32 => crate::mmap_backing::vec_to_bytes(offsets.to_vec()),
52            OffsetsDtype::U64 => {
53                let widened: Vec<u64> = offsets.iter().map(|&o| o as u64).collect();
54                crate::mmap_backing::vec_to_bytes(widened)
55            }
56        }
57    }
58}
59
60#[derive(Debug, Default)]
61struct TrxArchiveIndex {
62    dps: HashMap<String, String>,
63    dpv: HashMap<String, String>,
64    groups: HashMap<String, String>,
65    dpg: HashMap<String, HashMap<String, String>>,
66}
67
68fn get_entry_backing(
69    arc_mmap: &Arc<memmap2::Mmap>,
70    entry: &mut zip::read::ZipFile,
71    align_requirement: usize,
72) -> Result<MmapBacking> {
73    let size = entry.size() as usize;
74    if size == 0 {
75        return Ok(MmapBacking::Owned(Vec::new()));
76    }
77    if entry.compression() == zip::CompressionMethod::Stored {
78        let data_start = entry.data_start() as usize;
79        if data_start + size > arc_mmap.len() {
80            return Err(TrxError::Format(
81                "zip entry data exceeds archive boundaries".into(),
82            ));
83        }
84        let ptr_val = arc_mmap.as_ptr() as usize + data_start;
85        let req = align_requirement.max(1);
86        if ptr_val.is_multiple_of(req) {
87            Ok(MmapBacking::SharedSlice {
88                mmap: Arc::clone(arc_mmap),
89                offset: data_start,
90                len: size,
91            })
92        } else {
93            let slice = &arc_mmap[data_start..data_start + size];
94            let num_u64s = size.div_ceil(8);
95            let mut values = vec![0u64; num_u64s];
96            let bytes_slice = bytemuck::cast_slice_mut(&mut values);
97            bytes_slice[..size].copy_from_slice(slice);
98            Ok(MmapBacking::OwnedU64(values, size))
99        }
100    } else {
101        let num_u64s = size.div_ceil(8);
102        let mut values = vec![0u64; num_u64s];
103        let bytes_slice = bytemuck::cast_slice_mut(&mut values);
104        // We use read_exact or read_to_end into a subslice
105        Read::read_exact(entry, &mut bytes_slice[..size])?;
106        Ok(MmapBacking::OwnedU64(values, size))
107    }
108}
109
110fn load_zip_offsets(
111    arc_mmap: &Arc<memmap2::Mmap>,
112    entry: &mut zip::read::ZipFile,
113    dtype: crate::dtype::DType,
114    nb_streamlines: usize,
115    nb_vertices: usize,
116) -> Result<MmapBacking> {
117    match dtype {
118        crate::dtype::DType::UInt64 => {
119            if !entry.size().is_multiple_of(8) {
120                return Err(TrxError::Format(format!(
121                    "offsets entry size {} is not a multiple of 8 for uint64",
122                    entry.size()
123                )));
124            }
125            let mut values = vec![0u64; entry.size() as usize / 8];
126            Read::read_exact(entry, bytemuck::cast_slice_mut(&mut values))?;
127            if values.len() == nb_streamlines {
128                let mut owned: Vec<u32> = values
129                    .iter()
130                    .copied()
131                    .map(|v| {
132                        u32::try_from(v).map_err(|_| {
133                            TrxError::Format(format!("offset {v} exceeds uint32 range"))
134                        })
135                    })
136                    .collect::<Result<_>>()?;
137                owned.push(nb_vertices as u32);
138                let len = owned.len() * 4;
139                Ok(MmapBacking::OwnedU32(owned, len))
140            } else if values.len() == nb_streamlines + 1 {
141                let owned: Vec<u32> = values
142                    .iter()
143                    .copied()
144                    .map(|v| {
145                        u32::try_from(v).map_err(|_| {
146                            TrxError::Format(format!("offset {v} exceeds uint32 range"))
147                        })
148                    })
149                    .collect::<Result<_>>()?;
150                let len = owned.len() * 4;
151                Ok(MmapBacking::OwnedU32(owned, len))
152            } else {
153                Err(TrxError::Format(format!(
154                    "unexpected offset count: {} (expected {} or {})",
155                    values.len(),
156                    nb_streamlines,
157                    nb_streamlines + 1,
158                )))
159            }
160        }
161        crate::dtype::DType::UInt32 => {
162            if !entry.size().is_multiple_of(4) {
163                return Err(TrxError::Format(format!(
164                    "offsets entry size {} is not a multiple of 4 for uint32",
165                    entry.size()
166                )));
167            }
168            let num_u32 = entry.size() as usize / 4;
169            if num_u32 == nb_streamlines {
170                let mut values = vec![0u32; num_u32];
171                Read::read_exact(entry, bytemuck::cast_slice_mut(&mut values))?;
172                let mut out = values;
173                out.push(nb_vertices as u32);
174                let len = out.len() * 4;
175                Ok(MmapBacking::OwnedU32(out, len))
176            } else if num_u32 == nb_streamlines + 1 {
177                get_entry_backing(arc_mmap, entry, std::mem::align_of::<u32>())
178            } else {
179                Err(TrxError::Format(format!(
180                    "unexpected offset count: {} (expected {} or {})",
181                    num_u32,
182                    nb_streamlines,
183                    nb_streamlines + 1,
184                )))
185            }
186        }
187        other => Err(TrxError::DType(format!(
188            "offsets must be uint32 or uint64, got {other}"
189        ))),
190    }
191}
192
193/// Load a TRX file from a `.trx` zip archive.
194///
195/// Memory-maps the zip file directly from disk and parses entries in-memory
196/// without creating temporary directories.
197pub fn load_from_zip<P: TrxScalar>(path: &Path) -> Result<TrxFile<P>> {
198    let file = fs::File::open(path)?;
199    let mmap = unsafe { memmap2::Mmap::map(&file)? };
200    let arc_mmap = Arc::new(mmap);
201    let mut archive = zip::ZipArchive::new(std::io::Cursor::new(arc_mmap.as_ref()))?;
202
203    let header: Header = {
204        let mut header_entry = archive.by_name("header.json")?;
205        let mut bytes = Vec::new();
206        Read::read_to_end(&mut header_entry, &mut bytes)?;
207        serde_json::from_slice(&bytes)?
208    };
209
210    let mut positions_backing = None;
211    let mut offsets_backing = None;
212    let mut dps = HashMap::new();
213    let mut dpv = HashMap::new();
214    let mut groups = HashMap::new();
215    let mut dpg: DataPerGroup = HashMap::new();
216
217    for i in 0..archive.len() {
218        let mut entry = archive.by_index(i)?;
219        let name = entry.name().to_string();
220        if entry.is_dir() || name.ends_with('/') || name == "header.json" {
221            continue;
222        }
223
224        if let Some(rest) = name.strip_prefix("dps/") {
225            let parsed = TrxFilename::parse(rest)?;
226            let backing = get_entry_backing(&arc_mmap, &mut entry, parsed.dtype.size_of())?;
227            dps.insert(
228                parsed.name,
229                DataArray::from_backing(backing, parsed.ncols, parsed.dtype),
230            );
231        } else if let Some(rest) = name.strip_prefix("dpv/") {
232            let parsed = TrxFilename::parse(rest)?;
233            let backing = get_entry_backing(&arc_mmap, &mut entry, parsed.dtype.size_of())?;
234            dpv.insert(
235                parsed.name,
236                DataArray::from_backing(backing, parsed.ncols, parsed.dtype),
237            );
238        } else if let Some(rest) = name.strip_prefix("groups/") {
239            let parsed = TrxFilename::parse(rest)?;
240            let backing = get_entry_backing(&arc_mmap, &mut entry, parsed.dtype.size_of())?;
241            groups.insert(
242                parsed.name,
243                DataArray::from_backing(backing, parsed.ncols, parsed.dtype),
244            );
245        } else if let Some(rest) = name.strip_prefix("dpg/") {
246            if let Some((group, file_name)) = rest.split_once('/') {
247                let parsed = TrxFilename::parse(file_name)?;
248                let backing = get_entry_backing(&arc_mmap, &mut entry, parsed.dtype.size_of())?;
249                let group_map = dpg.entry(group.to_string()).or_default();
250                group_map.insert(
251                    parsed.name,
252                    DataArray::from_backing(backing, parsed.ncols, parsed.dtype),
253                );
254            }
255        } else if name.starts_with("positions.") {
256            let pos_parsed = TrxFilename::parse(&name)?;
257            if pos_parsed.dtype != P::DTYPE {
258                return Err(TrxError::DType(format!(
259                    "expected positions dtype {}, got {}",
260                    P::DTYPE,
261                    pos_parsed.dtype
262                )));
263            }
264            if pos_parsed.ncols != 3 {
265                return Err(TrxError::Format(format!(
266                    "positions must have 3 columns, got {}",
267                    pos_parsed.ncols
268                )));
269            }
270            positions_backing = Some(get_entry_backing(
271                &arc_mmap,
272                &mut entry,
273                std::mem::align_of::<P>(),
274            )?);
275        } else if name.starts_with("offsets.") {
276            let off_parsed = TrxFilename::parse(&name)?;
277            let backing = load_zip_offsets(
278                &arc_mmap,
279                &mut entry,
280                off_parsed.dtype,
281                header.nb_streamlines as usize,
282                header.nb_vertices as usize,
283            )?;
284            offsets_backing = Some(backing);
285        }
286    }
287
288    let positions_backing = match positions_backing {
289        Some(b) => b,
290        None => {
291            if header.nb_vertices == 0 {
292                MmapBacking::Owned(Vec::new())
293            } else {
294                return Err(TrxError::FileNotFound(path.join("positions")));
295            }
296        }
297    };
298
299    let offsets_backing = match offsets_backing {
300        Some(b) => b,
301        None => {
302            if header.nb_streamlines == 0 {
303                MmapBacking::Owned(vec_to_bytes(vec![0u32]))
304            } else {
305                return Err(TrxError::FileNotFound(path.join("offsets")));
306            }
307        }
308    };
309
310    Ok(TrxFile::from_parts(TrxParts {
311        header,
312        positions_backing,
313        offsets_backing,
314        dps,
315        dpv,
316        groups,
317        dpg,
318    }))
319}
320
321/// Save a `TrxFile<P>` to a `.trx` zip archive.
322///
323/// All entries are written uncompressed (Stored). DEFLATE rarely pays off on
324/// the float-dominated payload of a TRX file: compression ratios are typically
325/// <15% and write time grows substantially. Callers who want to compress the
326/// `groups/` entries (uint32 streamline-index lists, which do tend to have
327/// runs) can use [`save_to_zip_with`].
328pub fn save_to_zip<P: TrxScalar>(trx: &TrxFile<P>, path: &Path) -> Result<()> {
329    save_to_zip_with(trx, path, zip::CompressionMethod::Stored)
330}
331
332/// Save a `TrxFile<P>` to a `.trx` zip archive, applying `groups_compression`
333/// to `groups/` entries only. All other entries (header, positions, offsets,
334/// dps, dpv, dpg) are always Stored.
335pub fn save_to_zip_with<P: TrxScalar>(
336    trx: &TrxFile<P>,
337    path: &Path,
338    groups_compression: zip::CompressionMethod,
339) -> Result<()> {
340    let offsets_dtype = OffsetsDtype::pick_for(trx.offsets());
341    let file = fs::File::create(path)?;
342    let buf_writer = BufWriter::new(file);
343    let mut zip = zip::ZipWriter::new(buf_writer);
344    let stored = SimpleFileOptions::default()
345        .compression_method(zip::CompressionMethod::Stored)
346        .with_alignment(8)
347        .large_file(true);
348    let groups_opts = SimpleFileOptions::default()
349        .compression_method(groups_compression)
350        .large_file(true);
351
352    // Header
353    let header_json = trx.header().to_json()?;
354    zip.start_file("header.json", stored)?;
355    zip.write_all(header_json.as_bytes())?;
356
357    // Positions
358    let pos_filename = format!("positions.3.{}", P::DTYPE.name());
359    zip.start_file(&pos_filename, stored)?;
360    zip.write_all(trx.positions_bytes())?;
361
362    // Offsets — written at `offsets_dtype`'s width.
363    let offsets_filename = format!("offsets.{}", offsets_dtype.suffix());
364    zip.start_file(&offsets_filename, stored)?;
365    let offsets_bytes = offsets_dtype.encode(trx.offsets());
366    zip.write_all(&offsets_bytes)?;
367
368    // DPS / DPV — float-heavy, Stored.
369    write_data_map(&mut zip, "dps", trx.dps_arrays(), stored)?;
370    write_data_map(&mut zip, "dpv", trx.dpv_arrays(), stored)?;
371
372    // Groups — uint32 membership lists; honor caller's compression choice.
373    write_data_map(&mut zip, "groups", trx.group_arrays(), groups_opts)?;
374
375    // DPG — tiny per-group scalars, Stored.
376    write_dpg_map(&mut zip, "dpg", trx.dpg_arrays(), stored)?;
377
378    let mut buf_writer = zip.finish()?;
379    buf_writer.flush()?;
380    Ok(())
381}
382
383/// Append DPS arrays to a TRX zip archive, optionally overwriting existing entries.
384pub fn append_dps_to_zip(
385    path: &Path,
386    dps: &HashMap<String, DataArray>,
387    compression: zip::CompressionMethod,
388    overwrite: bool,
389) -> Result<()> {
390    let header = read_header_from_zip(path)?;
391    validate_row_count("DPS", dps, header.nb_streamlines as usize)?;
392    let index = build_archive_index(path)?;
393    let mut ops = Vec::new();
394    for (name, arr) in dps {
395        let target = data_entry_path("dps", name, arr);
396        plan_data_write(
397            &index.dps,
398            name,
399            target,
400            arr,
401            overwrite,
402            compression,
403            &mut ops,
404        )?;
405    }
406    archive_edit::apply_archive_ops(path, ops)
407}
408
409/// Append DPV arrays to a TRX zip archive, optionally overwriting existing entries.
410pub fn append_dpv_to_zip(
411    path: &Path,
412    dpv: &HashMap<String, DataArray>,
413    compression: zip::CompressionMethod,
414    overwrite: bool,
415) -> Result<()> {
416    let header = read_header_from_zip(path)?;
417    validate_row_count("DPV", dpv, header.nb_vertices as usize)?;
418    let index = build_archive_index(path)?;
419    let mut ops = Vec::new();
420    for (name, arr) in dpv {
421        let target = data_entry_path("dpv", name, arr);
422        plan_data_write(
423            &index.dpv,
424            name,
425            target,
426            arr,
427            overwrite,
428            compression,
429            &mut ops,
430        )?;
431    }
432    archive_edit::apply_archive_ops(path, ops)
433}
434
435/// Append group membership arrays to a TRX zip archive, optionally overwriting existing entries.
436pub fn append_groups_to_zip(
437    path: &Path,
438    groups: &HashMap<String, Vec<u32>>,
439    compression: zip::CompressionMethod,
440    overwrite: bool,
441) -> Result<()> {
442    let header = read_header_from_zip(path)?;
443    let index = build_archive_index(path)?;
444    let mut ops = Vec::new();
445    for (name, members) in groups {
446        validate_group_members(name, members, header.nb_streamlines as usize)?;
447        let target = format!("groups/{name}.uint32");
448        let bytes = vec_to_bytes(members.clone());
449        plan_bytes_write(
450            index.groups.get(name),
451            target,
452            bytes,
453            overwrite,
454            compression,
455            &mut ops,
456        )?;
457    }
458    archive_edit::apply_archive_ops(path, ops)
459}
460
461/// Append DPG (data-per-group) entries to a TRX zip archive, optionally overwriting existing entries.
462pub fn append_dpg_to_zip(
463    path: &Path,
464    dpg: &DataPerGroup,
465    compression: zip::CompressionMethod,
466    overwrite: bool,
467) -> Result<()> {
468    let index = build_archive_index(path)?;
469    let available_groups: BTreeSet<&str> = index.groups.keys().map(String::as_str).collect();
470    let mut ops = Vec::new();
471    for (group, arrays) in dpg {
472        if !available_groups.contains(group.as_str()) {
473            return Err(TrxError::Argument(format!(
474                "cannot add DPG entries for missing group '{group}'"
475            )));
476        }
477        let existing = index.dpg.get(group);
478        for (name, arr) in arrays {
479            let target = format!("dpg/{group}/{}", filename_for_array(name, arr));
480            let existing_path = existing.and_then(|entries| entries.get(name));
481            plan_bytes_write(
482                existing_path,
483                target,
484                arr.as_bytes().to_vec(),
485                overwrite,
486                compression,
487                &mut ops,
488            )?;
489        }
490    }
491    archive_edit::apply_archive_ops(path, ops)
492}
493
494/// Delete named DPS arrays from a TRX zip archive.
495pub fn delete_dps_from_zip(path: &Path, names: &[&str]) -> Result<()> {
496    let index = build_archive_index(path)?;
497    let mut ops = Vec::new();
498    for name in names {
499        if let Some(entry_path) = index.dps.get(*name) {
500            ops.push(ArchiveOp::Delete {
501                path: entry_path.clone(),
502            });
503        }
504    }
505    archive_edit::apply_archive_ops(path, ops)
506}
507
508/// Delete named DPV arrays from a TRX zip archive.
509pub fn delete_dpv_from_zip(path: &Path, names: &[&str]) -> Result<()> {
510    let index = build_archive_index(path)?;
511    let mut ops = Vec::new();
512    for name in names {
513        if let Some(entry_path) = index.dpv.get(*name) {
514            ops.push(ArchiveOp::Delete {
515                path: entry_path.clone(),
516            });
517        }
518    }
519    archive_edit::apply_archive_ops(path, ops)
520}
521
522/// Delete named groups (and their DPG entries) from a TRX zip archive.
523pub fn delete_groups_from_zip(path: &Path, names: &[&str]) -> Result<()> {
524    let index = build_archive_index(path)?;
525    let mut ops = Vec::new();
526    for name in names {
527        if let Some(entry_path) = index.groups.get(*name) {
528            ops.push(ArchiveOp::Delete {
529                path: entry_path.clone(),
530            });
531        }
532        ops.push(ArchiveOp::DeletePrefix {
533            prefix: format!("dpg/{name}"),
534        });
535    }
536    archive_edit::apply_archive_ops(path, ops)
537}
538
539/// Delete DPG entries for a specific group from a TRX zip archive.
540///
541/// When `names` is `None` or empty, the entire DPG prefix for the group is removed.
542/// When `names` lists specific fields, only those entries are deleted.
543pub fn delete_dpg_from_zip(path: &Path, group: &str, names: Option<&[&str]>) -> Result<()> {
544    let index = build_archive_index(path)?;
545    let mut ops = Vec::new();
546    match names {
547        None => ops.push(ArchiveOp::DeletePrefix {
548            prefix: format!("dpg/{group}"),
549        }),
550        Some([]) => ops.push(ArchiveOp::DeletePrefix {
551            prefix: format!("dpg/{group}"),
552        }),
553        Some(names) => {
554            if let Some(entries) = index.dpg.get(group) {
555                for name in names {
556                    if let Some(entry_path) = entries.get(*name) {
557                        ops.push(ArchiveOp::Delete {
558                            path: entry_path.clone(),
559                        });
560                    }
561                }
562            }
563        }
564    }
565    archive_edit::apply_archive_ops(path, ops)
566}
567
568fn read_header_from_zip(path: &Path) -> Result<Header> {
569    let bytes = archive_edit::read_archive_entry(path, "header.json")?;
570    Ok(serde_json::from_slice(&bytes)?)
571}
572
573fn build_archive_index(path: &Path) -> Result<TrxArchiveIndex> {
574    let entries = archive_edit::archive_entry_names(path)?;
575    let mut index = TrxArchiveIndex::default();
576
577    for entry in entries {
578        if let Some(rest) = entry.strip_prefix("dps/") {
579            index_entry(&mut index.dps, &entry, rest)?;
580        } else if let Some(rest) = entry.strip_prefix("dpv/") {
581            index_entry(&mut index.dpv, &entry, rest)?;
582        } else if let Some(rest) = entry.strip_prefix("groups/") {
583            index_entry(&mut index.groups, &entry, rest)?;
584        } else if let Some(rest) = entry.strip_prefix("dpg/") {
585            if let Some((group, file_name)) = rest.split_once('/') {
586                let parsed = TrxFilename::parse(file_name)?;
587                let group_entries = index.dpg.entry(group.to_string()).or_default();
588                if group_entries.insert(parsed.name, entry.clone()).is_some() {
589                    return Err(TrxError::Format(format!(
590                        "duplicate DPG entry path for group '{group}'"
591                    )));
592                }
593            }
594        }
595    }
596
597    Ok(index)
598}
599
600fn index_entry(
601    index: &mut HashMap<String, String>,
602    full_path: &str,
603    file_name: &str,
604) -> Result<()> {
605    if file_name.ends_with('/') {
606        return Ok(());
607    }
608    let parsed = TrxFilename::parse(file_name)?;
609    if index.insert(parsed.name, full_path.to_string()).is_some() {
610        return Err(TrxError::Format(format!(
611            "duplicate archive entry for '{full_path}'"
612        )));
613    }
614    Ok(())
615}
616
617fn validate_row_count(
618    kind: &str,
619    arrays: &HashMap<String, DataArray>,
620    expected_rows: usize,
621) -> Result<()> {
622    for (name, arr) in arrays {
623        if arr.nrows() != expected_rows {
624            return Err(TrxError::Format(format!(
625                "{kind} '{name}' has {} rows, expected {expected_rows}",
626                arr.nrows()
627            )));
628        }
629    }
630    Ok(())
631}
632
633fn validate_group_members(name: &str, members: &[u32], nb_streamlines: usize) -> Result<()> {
634    for &member in members {
635        if member as usize >= nb_streamlines {
636            return Err(TrxError::Format(format!(
637                "group '{name}' contains streamline index {member}, but NB_STREAMLINES is {nb_streamlines}"
638            )));
639        }
640    }
641    Ok(())
642}
643
644fn data_entry_path(prefix: &str, name: &str, arr: &DataArray) -> String {
645    format!("{prefix}/{}", filename_for_array(name, arr))
646}
647
648fn filename_for_array(name: &str, arr: &DataArray) -> String {
649    TrxFilename {
650        name: name.to_string(),
651        ncols: arr.ncols(),
652        dtype: arr.dtype(),
653    }
654    .to_filename()
655}
656
657fn plan_data_write(
658    existing: &HashMap<String, String>,
659    logical_name: &str,
660    target_path: String,
661    arr: &DataArray,
662    overwrite: bool,
663    compression: zip::CompressionMethod,
664    ops: &mut Vec<ArchiveOp>,
665) -> Result<()> {
666    plan_bytes_write(
667        existing.get(logical_name),
668        target_path,
669        arr.as_bytes().to_vec(),
670        overwrite,
671        compression,
672        ops,
673    )
674}
675
676fn plan_bytes_write(
677    existing_path: Option<&String>,
678    target_path: String,
679    bytes: Vec<u8>,
680    overwrite: bool,
681    compression: zip::CompressionMethod,
682    ops: &mut Vec<ArchiveOp>,
683) -> Result<()> {
684    match existing_path {
685        None => ops.push(ArchiveOp::Add {
686            path: target_path,
687            bytes,
688            compression,
689        }),
690        Some(_) if !overwrite => {}
691        Some(existing) if existing == &target_path => ops.push(ArchiveOp::Replace {
692            path: target_path,
693            bytes,
694            compression,
695        }),
696        Some(existing) => {
697            ops.push(ArchiveOp::Delete {
698                path: existing.clone(),
699            });
700            ops.push(ArchiveOp::Add {
701                path: target_path,
702                bytes,
703                compression,
704            });
705        }
706    }
707    Ok(())
708}
709
710fn write_data_map<W: Write + std::io::Seek>(
711    zip: &mut zip::ZipWriter<W>,
712    prefix: &str,
713    arrays: &HashMap<String, DataArray>,
714    options: SimpleFileOptions,
715) -> Result<()> {
716    for (name, arr) in arrays {
717        let filename = filename_for_array(name, arr);
718        let entry_name = format!("{prefix}/{filename}");
719        zip.start_file(&entry_name, options)?;
720        zip.write_all(arr.as_bytes())?;
721    }
722    Ok(())
723}
724
725fn write_dpg_map<W: Write + std::io::Seek>(
726    zip: &mut zip::ZipWriter<W>,
727    prefix: &str,
728    groups: &DataPerGroup,
729    options: SimpleFileOptions,
730) -> Result<()> {
731    for (group, arrays) in groups {
732        for (name, arr) in arrays {
733            let filename = filename_for_array(name, arr);
734            let entry_name = format!("{prefix}/{group}/{filename}");
735            zip.start_file(&entry_name, options)?;
736            zip.write_all(arr.as_bytes())?;
737        }
738    }
739    Ok(())
740}
741
742#[cfg(test)]
743mod tests {
744    use super::*;
745    use crate::header::Header;
746    use crate::stream::TrxStream;
747
748    #[test]
749    fn zip_round_trip_deflated_and_stored() {
750        let mut stream = TrxStream::<f32>::new(Header::identity_affine(), [100, 100, 100]);
751        stream.push_streamline(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
752        let trx = stream.finalize();
753
754        let dir = tempfile::TempDir::new().unwrap();
755        let zip_path = dir.path().join("test_deflated.trx");
756
757        save_to_zip_with(&trx, &zip_path, zip::CompressionMethod::Deflated).unwrap();
758        let loaded = load_from_zip::<f32>(&zip_path).unwrap();
759
760        assert_eq!(loaded.nb_streamlines(), 1);
761        assert_eq!(loaded.nb_vertices(), 2);
762        assert_eq!(loaded.streamline(0), trx.streamline(0));
763    }
764
765    #[test]
766    fn zip_direct_mapping_without_temp_dir() {
767        let mut stream = TrxStream::<f32>::new(Header::identity_affine(), [100, 100, 100]);
768        stream.push_streamline(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
769        let trx = stream.finalize();
770
771        let dir = tempfile::TempDir::new().unwrap();
772        let zip_path = dir.path().join("stored.trx");
773
774        save_to_zip_with(&trx, &zip_path, zip::CompressionMethod::Stored).unwrap();
775        let loaded = load_from_zip::<f32>(&zip_path).unwrap();
776
777        assert!(loaded.is_file_backed());
778        assert_eq!(loaded.nb_streamlines(), 1);
779        assert_eq!(loaded.nb_vertices(), 2);
780        assert_eq!(loaded.streamline(0), trx.streamline(0));
781    }
782
783    #[test]
784    fn zip_unaligned_offsets_size_errors() {
785        let dir = tempfile::TempDir::new().unwrap();
786        let zip_path = dir.path().join("corrupt_offsets.trx");
787
788        let file = std::fs::File::create(&zip_path).unwrap();
789        let mut zip = zip::ZipWriter::new(file);
790
791        let header = Header {
792            voxel_to_rasmm: Header::identity_affine(),
793            dimensions: [100, 100, 100],
794            nb_streamlines: 0,
795            nb_vertices: 0,
796            extra: Default::default(),
797        };
798        let json = serde_json::to_string(&header).unwrap();
799        zip.start_file("header.json", zip::write::SimpleFileOptions::default())
800            .unwrap();
801        zip.write_all(json.as_bytes()).unwrap();
802
803        // Write positions.bit32.ncols3.raw (0 vertices)
804        zip.start_file(
805            "positions.bit32.ncols3.raw",
806            zip::write::SimpleFileOptions::default(),
807        )
808        .unwrap();
809
810        // Write unaligned offsets (e.g. 5 bytes instead of multiple of 4 or 8)
811        zip.start_file(
812            "offsets.bit32.ncols1.raw",
813            zip::write::SimpleFileOptions::default(),
814        )
815        .unwrap();
816        zip.write_all(&[0u8; 5]).unwrap();
817
818        zip.finish().unwrap();
819
820        let res = load_from_zip::<f32>(&zip_path);
821        assert!(res.is_err());
822    }
823}
trx_rs/io/zip.rs

trx_rs/io/
zip.rs
