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::NonZeroU64,
8    path::PathBuf,
9};
10
11#[derive(Debug)]
12pub struct SplitWriter<F> {
13    split_size: NonZeroU64,
14    dest_dir: PathBuf,
15    get_file_name: F,
16    current_pos: u64,
17    first_file: File,
18    last_file: Option<File>,
19    file_count: usize,
20}
21
22impl<F> SplitWriter<F>
23where
24    F: Fn(usize) -> String,
25{
26    pub fn try_new(
27        dest_dir: impl Into<PathBuf>,
28        get_file_name: F,
29        split_size: NonZeroU64,
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_pos: 0,
39            first_file,
40            last_file: None,
41            file_count: 1,
42        })
43    }
44
45    pub fn file_count(&self) -> usize {
46        self.file_count
47    }
48
49    pub fn total_size(&self) -> u64 {
50        self.current_pos
51    }
52
53    pub fn write_header(&mut self, header: &[u8]) -> io::Result<()> {
54        self.first_file.rewind()?;
55        self.first_file.write_all(header)
56    }
57}
58
59impl<F> Write for SplitWriter<F>
60where
61    F: Fn(usize) -> String,
62{
63    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
64        if buf.is_empty() {
65            return Ok(0);
66        }
67
68        #[allow(clippy::cast_possible_truncation)]
69        let i = (self.current_pos / self.split_size.get()) as usize;
70
71        if self.file_count <= i {
72            let idx = self.file_count;
73            let file_name = (self.get_file_name)(idx);
74            let file_path = self.dest_dir.join(file_name);
75            let file = File::create(file_path)?;
76
77            if let Some(last_file) = &mut self.last_file {
78                last_file.flush()?;
79            }
80
81            self.last_file = Some(file);
82            self.file_count += 1;
83        }
84
85        let (file, offset) = if let Some(last_file) = &mut self.last_file {
86            (last_file, self.current_pos % self.split_size.get())
87        } else {
88            (&mut self.first_file, self.current_pos)
89        };
90
91        let to_write = match usize::try_from(self.split_size.get() - offset) {
92            Ok(remaining) => buf.len().min(remaining),
93            Err(_) => buf.len(),
94        };
95
96        let n = file.write(&buf[..to_write])?;
97
98        self.current_pos += n as u64;
99
100        Ok(n)
101    }
102
103    fn flush(&mut self) -> io::Result<()> {
104        if let Some(last_file) = &mut self.last_file {
105            last_file.flush()?;
106        }
107
108        self.first_file.flush()
109    }
110}
split_write/lib.rs

split_write/
lib.rs
