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: u64,
20}
21
22impl<F> SplitWriter<F>
23where
24    F: Fn(u64) -> 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) -> u64 {
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(u64) -> String,
62{
63    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
64        if buf.is_empty() {
65            return Ok(0);
66        }
67
68        let i = self.current_pos / self.split_size.get();
69
70        if self.file_count <= i {
71            let idx = self.file_count;
72            let file_name = (self.get_file_name)(idx);
73            let file_path = self.dest_dir.join(file_name);
74            let file = File::create(file_path)?;
75
76            if let Some(last_file) = &mut self.last_file {
77                last_file.flush()?;
78            }
79
80            self.last_file = Some(file);
81            self.file_count += 1;
82        }
83
84        let (file, offset) = match &mut self.last_file {
85            Some(last_file) => (last_file, self.current_pos % self.split_size.get()),
86            None => (&mut self.first_file, self.current_pos),
87        };
88
89        let to_write = match usize::try_from(self.split_size.get() - offset) {
90            Ok(remaining) => buf.len().min(remaining),
91            Err(_) => buf.len(),
92        };
93
94        let n = file.write(&buf[..to_write])?;
95
96        self.current_pos += n as u64;
97
98        Ok(n)
99    }
100
101    fn flush(&mut self) -> io::Result<()> {
102        if let Some(last_file) = &mut self.last_file {
103            last_file.flush()?;
104        }
105
106        self.first_file.flush()
107    }
108}
split_write/lib.rs

split_write/
lib.rs
