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 let Some(split_size) = self.split_size
72            && self.current_offset == split_size.get()
73        {
74            self.current_i += 1;
75
76            let file_name = (self.get_file_name)(self.current_i);
77            let file_path = self.dest_dir.join(file_name);
78            let file = File::create(file_path)?;
79
80            self.last_file = Some(file);
81            self.current_offset = 0;
82        }
83
84        let current_file = match &mut self.last_file {
85            Some(last_file) => last_file,
86            None => &mut self.first_file,
87        };
88
89        let written = match self.split_size {
90            Some(split_size) => {
91                let remaining = split_size.get() - self.current_offset;
92                let to_write = buf.len().min(remaining);
93                let written = current_file.write(&buf[..to_write])?;
94                self.current_offset += written;
95                written
96            }
97            None => current_file.write(buf)?,
98        };
99
100        self.total_size += written as u64;
101
102        Ok(written)
103    }
104
105    fn flush(&mut self) -> io::Result<()> {
106        if let Some(last_file) = &mut self.last_file {
107            last_file.flush()?;
108        }
109
110        self.first_file.flush()
111    }
112}
split_write/lib.rs

split_write/
lib.rs
