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: 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}
21
22impl<F> SplitWriter<F>
23where
24    F: Fn(usize) -> String,
25{
26    pub fn create(
27        dest_dir: impl Into<PathBuf>,
28        get_file_name: F,
29        split_size: NonZeroUsize,
30    ) -> io::Result<Self> {
31        let dest_dir = dest_dir.into();
32        let first_file = File::create(dest_dir.join(get_file_name(0)))?;
33
34        Ok(Self {
35            split_size,
36            dest_dir,
37            get_file_name,
38            current_offset: 0,
39            first_file,
40            last_file: None,
41            current_i: 0,
42        })
43    }
44
45    pub fn file_count(&self) -> usize {
46        self.current_i + 1
47    }
48
49    pub fn total_size(&self) -> u64 {
50        self.split_size.get() as u64 * self.current_i as u64 + self.current_offset as u64
51    }
52
53    /// Don't do any more writes after calling this!
54    pub fn write_header(&mut self, header: &[u8]) -> io::Result<()> {
55        self.first_file.rewind()?;
56        self.first_file.write_all(header)
57    }
58}
59
60impl<F> Write for SplitWriter<F>
61where
62    F: Fn(usize) -> String,
63{
64    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
65        if buf.is_empty() {
66            return Ok(0);
67        }
68
69        if self.current_offset == self.split_size.get() {
70            self.current_i += 1;
71
72            let file_name = (self.get_file_name)(self.current_i);
73            let file_path = self.dest_dir.join(file_name);
74            let file = File::create(file_path)?;
75
76            self.last_file = Some(file);
77            self.current_offset = 0;
78        }
79
80        let current_file = match &mut self.last_file {
81            Some(last_file) => last_file,
82            None => &mut self.first_file,
83        };
84
85        let remaining = self.split_size.get() - self.current_offset;
86        let to_write = buf.len().min(remaining);
87        let written = current_file.write(&buf[..to_write])?;
88
89        self.current_offset += written;
90
91        Ok(written)
92    }
93
94    fn flush(&mut self) -> io::Result<()> {
95        if let Some(last_file) = &mut self.last_file {
96            last_file.flush()?;
97        }
98
99        self.first_file.flush()
100    }
101}
split_write/lib.rs

split_write/
lib.rs
