split_write/
lib.rs

1// SPDX-FileCopyrightText: 2026 Manuel Quarneti <mq1@ik.me>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::{
5    fs::File,
6    io::{self, Seek, Write},
7    num::NonZeroUsize,
8    path::PathBuf,
9};
10
11#[derive(Debug)]
12pub struct SplitWriter<F> {
13    split_size: Option<NonZeroUsize>,
14    dest_dir: PathBuf,
15    get_file_name: F,
16    current_offset: usize,
17    first_file: File,
18    last_file: Option<File>,
19    current_i: usize,
20    total_size: u64,
21}
22
23impl<F> SplitWriter<F>
24where
25    F: Fn(usize) -> String,
26{
27    pub fn create(
28        dest_dir: impl Into<PathBuf>,
29        get_file_name: F,
30        split_size: Option<NonZeroUsize>,
31    ) -> io::Result<Self> {
32        let dest_dir = dest_dir.into();
33        let first_file = File::create(dest_dir.join(get_file_name(0)))?;
34
35        Ok(Self {
36            split_size,
37            dest_dir,
38            get_file_name,
39            current_offset: 0,
40            first_file,
41            last_file: None,
42            current_i: 0,
43            total_size: 0,
44        })
45    }
46
47    pub fn file_count(&self) -> usize {
48        self.current_i + 1
49    }
50
51    pub fn total_size(&self) -> u64 {
52        self.total_size
53    }
54
55    /// Don't do any more writes after calling this!
56    pub fn write_header(&mut self, header: &[u8]) -> io::Result<()> {
57        self.first_file.rewind()?;
58        self.first_file.write_all(header)
59    }
60}
61
62impl<F> Write for SplitWriter<F>
63where
64    F: Fn(usize) -> String,
65{
66    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
67        if buf.is_empty() {
68            return Ok(0);
69        }
70
71        if self
72            .split_size
73            .is_some_and(|s| s.get() == self.current_offset)
74        {
75            self.current_i += 1;
76
77            let file_name = (self.get_file_name)(self.current_i);
78            let file_path = self.dest_dir.join(file_name);
79            let file = File::create(file_path)?;
80
81            self.last_file = Some(file);
82            self.current_offset = 0;
83        }
84
85        let current_file = self.last_file.as_mut().unwrap_or(&mut self.first_file);
86
87        let written = if let Some(split_size) = self.split_size {
88            let remaining = split_size.get() - self.current_offset;
89            let to_write = buf.len().min(remaining);
90            let written = current_file.write(&buf[..to_write])?;
91            self.current_offset += written;
92            written
93        } else {
94            current_file.write(buf)?
95        };
96
97        self.total_size += written as u64;
98
99        Ok(written)
100    }
101
102    fn flush(&mut self) -> io::Result<()> {
103        if let Some(last_file) = &mut self.last_file {
104            last_file.flush()?;
105        }
106
107        self.first_file.flush()
108    }
109}
split_write/lib.rs

split_write/
lib.rs
