use std::{
io,
sync::atomic::{AtomicBool, Ordering},
};
pub struct Iter<'a, I> {
pub inner: I,
should_interrupt: &'a AtomicBool,
}
impl<'a, I> Iter<'a, I>
where
I: Iterator,
{
pub fn new(inner: I, should_interrupt: &'a AtomicBool) -> Self {
Iter {
inner,
should_interrupt,
}
}
}
impl<I> Iterator for Iter<'_, I>
where
I: Iterator,
{
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
if self.should_interrupt.load(Ordering::Relaxed) {
return None;
}
self.inner.next()
}
}
pub struct IterWithErr<'a, I, EFN> {
pub inner: I,
make_err: Option<EFN>,
should_interrupt: &'a AtomicBool,
}
impl<'a, I, EFN, E> IterWithErr<'a, I, EFN>
where
I: Iterator,
EFN: FnOnce() -> E,
{
pub fn new(inner: I, make_err: EFN, should_interrupt: &'a AtomicBool) -> Self {
IterWithErr {
inner,
make_err: Some(make_err),
should_interrupt,
}
}
}
impl<I, EFN, E> Iterator for IterWithErr<'_, I, EFN>
where
I: Iterator,
EFN: FnOnce() -> E,
{
type Item = Result<I::Item, E>;
fn next(&mut self) -> Option<Self::Item> {
self.make_err.as_ref()?;
if self.should_interrupt.load(Ordering::Relaxed) {
return self.make_err.take().map(|f| Err(f()));
}
match self.inner.next() {
Some(next) => Some(Ok(next)),
None => {
self.make_err = None;
None
}
}
}
}
pub struct Read<'a, R> {
pub inner: R,
pub should_interrupt: &'a AtomicBool,
}
impl<R> io::Read for Read<'_, R>
where
R: io::Read,
{
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if self.should_interrupt.load(Ordering::Relaxed) {
return Err(std::io::Error::new(std::io::ErrorKind::Other, "Interrupted"));
}
self.inner.read(buf)
}
}
impl<R> io::BufRead for Read<'_, R>
where
R: io::BufRead,
{
fn fill_buf(&mut self) -> io::Result<&[u8]> {
self.inner.fill_buf()
}
fn consume(&mut self, amt: usize) {
self.inner.consume(amt);
}
}
pub struct Write<'a, W> {
pub inner: W,
pub should_interrupt: &'a AtomicBool,
}
impl<W> io::Write for Write<'_, W>
where
W: std::io::Write,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if self.should_interrupt.load(Ordering::Relaxed) {
return Err(std::io::Error::new(std::io::ErrorKind::Other, "Interrupted"));
}
self.inner.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
impl<W> io::Seek for Write<'_, W>
where
W: std::io::Seek,
{
fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
self.inner.seek(pos)
}
}