forked from bonsai/harakit
112 lines
2.9 KiB
Rust
112 lines
2.9 KiB
Rust
/*
|
|
* Copyright (c) 2025 Emma Tebibyte <emma@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::self,
|
|
};
|
|
|
|
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(stream: T, delimiter: &[u8]) -> Self {
|
|
Delimited {
|
|
stream,
|
|
delimiter: delimiter.to_vec(),
|
|
buffer: Vec::with_capacity(BUFFER_SIZE),
|
|
}
|
|
}
|
|
}
|
|
|
|
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) = find_subslice(&self.buffer, &self.delimiter) {
|
|
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) => {
|
|
if self.buffer.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
return Some(Ok(mem::take(&mut self.buffer)));
|
|
},
|
|
Ok(n) => {
|
|
let content = &buf[..n];
|
|
self.buffer.extend_from_slice(&content);
|
|
},
|
|
Err(e) => {
|
|
return Some(Err(e));
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn find_subslice(stack: &[u8], key: &[u8]) -> Option<usize> {
|
|
if key.len() == 1 {
|
|
return stack.iter().position(|&b| b == key[0]);
|
|
}
|
|
if key.len() > stack.len() {
|
|
return None;
|
|
}
|
|
for i in 0..=stack.len() - key.len() {
|
|
if &stack[i..i + key.len()] == key {
|
|
return Some(i);
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
#[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.as_bytes());
|
|
|
|
let mut i = 0;
|
|
while let Some(item) = output.next() {
|
|
assert_eq!(input[i].as_bytes(), item.unwrap());
|
|
i += 1;
|
|
}
|
|
}
|
|
}
|