106 lines
2.5 KiB
Rust
106 lines
2.5 KiB
Rust
/*
|
||
* Copyright (c) 2025 Emma Tebibyte <emma@tebibyte.media>
|
||
* Copyright (c) 2025 silty silt <silt@tebibyte.media>
|
||
* SPDX-License-Identifier: AGPL-3.0-or-later
|
||
*
|
||
* This program is free software: you can redistribute it and/or modify it
|
||
* under the terms of the GNU Affero General Public License as published by the
|
||
* Free Software Foundation, either version 3 of the License, or (at your
|
||
* option) any later version.
|
||
*
|
||
* This program is distributed in the hope that it will be useful, but WITHOUT
|
||
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
|
||
* FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License
|
||
* for more details.
|
||
*
|
||
* You should have received a copy of the GNU Affero General Public License
|
||
* along with this program. If not, see https://www.gnu.org/licenses/.
|
||
*/
|
||
|
||
use std::{
|
||
io::{ Read, Result },
|
||
mem,
|
||
};
|
||
|
||
const BUFFER_SIZE: usize = 4096;
|
||
|
||
pub struct Delimited<T: Read> {
|
||
delimiter: Vec<u8>,
|
||
buffer: Vec<u8>,
|
||
stream: T,
|
||
}
|
||
|
||
impl<T> Delimited<T> where T: Read {
|
||
pub fn new<R>(stream: T, delimiter: R) -> Self where R: AsRef<[u8]> {
|
||
Delimited {
|
||
stream,
|
||
delimiter: delimiter.as_ref().to_vec(),
|
||
buffer: Vec::with_capacity(BUFFER_SIZE),
|
||
}
|
||
}
|
||
|
||
fn find_subslice(&self) -> Option<usize> {
|
||
match self.delimiter.len() {
|
||
/* TODO: is this optimization necessary? */
|
||
1 => self.buffer.iter().position(|&b| b == self.delimiter[0]),
|
||
_ => {
|
||
self.buffer
|
||
.windows(self.delimiter.len())
|
||
.position(|w| w == self.delimiter)
|
||
},
|
||
}
|
||
}
|
||
}
|
||
|
||
impl<T> Iterator for Delimited<T> where T: Read {
|
||
type Item = Result<Vec<u8>>;
|
||
|
||
fn next(&mut self) -> Option<Self::Item> {
|
||
let mut buf = [0; BUFFER_SIZE];
|
||
|
||
loop {
|
||
if let Some(p) = self.find_subslice() {
|
||
let chunk = self.buffer.drain(..p).collect::<Vec<_>>();
|
||
|
||
let _ = self.buffer.drain(..self.delimiter.len());
|
||
|
||
return Some(Ok(chunk));
|
||
}
|
||
|
||
match self.stream.read(&mut buf) {
|
||
Ok(0) => { /* no bytes read, we’re probably done */
|
||
let _ = self.buffer.is_empty() && return None;
|
||
|
||
return Some(Ok(mem::take(&mut self.buffer)));
|
||
},
|
||
Ok(n) => {
|
||
self.buffer.extend_from_slice(&buf[..n]);
|
||
},
|
||
Err(e) => {
|
||
return Some(Err(e));
|
||
},
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use Delimited;
|
||
|
||
#[test]
|
||
fn testing() {
|
||
let d = '\u{1E}'.to_string();
|
||
let input = vec!["meow", "woof", "ribbit"];
|
||
let r = input.join(&d);
|
||
|
||
let mut output = Delimited::new(r.as_bytes(), d);
|
||
|
||
let mut i = 0;
|
||
while let Some(item) = output.next() {
|
||
assert_eq!(input[i].as_bytes(), item.unwrap());
|
||
i += 1;
|
||
}
|
||
}
|
||
}
|