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}
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: Option<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        match self.split_size {
51            Some(split_size) => {
52                split_size.get() as u64 * self.current_i as u64 + self.current_offset as u64
53            }
54            None => self.current_offset as u64,
55        }
56    }
57
58    /// Don't do any more writes after calling this!
59    pub fn write_header(&mut self, header: &[u8]) -> io::Result<()> {
60        self.first_file.rewind()?;
61        self.first_file.write_all(header)
62    }
63}
64
65impl<F> Write for SplitWriter<F>
66where
67    F: Fn(usize) -> String,
68{
69    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
70        if buf.is_empty() {
71            return Ok(0);
72        }
73
74        if let Some(split_size) = self.split_size
75            && self.current_offset == split_size.get()
76        {
77            self.current_i += 1;
78
79            let file_name = (self.get_file_name)(self.current_i);
80            let file_path = self.dest_dir.join(file_name);
81            let file = File::create(file_path)?;
82
83            self.last_file = Some(file);
84            self.current_offset = 0;
85        }
86
87        let current_file = match &mut self.last_file {
88            Some(last_file) => last_file,
89            None => &mut self.first_file,
90        };
91
92        let written = match self.split_size {
93            Some(split_size) => {
94                let remaining = split_size.get() - self.current_offset;
95                let to_write = buf.len().min(remaining);
96                current_file.write(&buf[..to_write])?
97            }
98            None => current_file.write(buf)?,
99        };
100
101        self.current_offset += written;
102
103        Ok(written)
104    }
105
106    fn flush(&mut self) -> io::Result<()> {
107        if let Some(last_file) = &mut self.last_file {
108            last_file.flush()?;
109        }
110
111        self.first_file.flush()
112    }
113}
split_write/lib.rs

split_write/
lib.rs
