use std::{
fs::File,
io::{self, LineWriter, Write},
path::{Path, PathBuf},
str::FromStr,
sync::{Arc, Mutex, MutexGuard},
};
#[track_caller]
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<T> {
mutex.lock().unwrap_or_else(|e| e.into_inner())
}
#[derive(Debug, Clone)]
pub struct Output(OutputInner);
#[derive(Debug, Clone)]
enum OutputInner {
Stdout,
File {
path: Arc<PathBuf>,
writer: Arc<Mutex<LineWriter<File>>>,
},
}
impl Output {
pub fn stdout() -> Self {
Self(OutputInner::Stdout)
}
pub fn create(path: PathBuf) -> io::Result<Self> {
let path = Arc::new(path);
let file = File::create(&*path)?;
let writer = Arc::new(Mutex::new(LineWriter::new(file)));
Ok(Self(OutputInner::File { path, writer }))
}
pub fn is_stdout(&self) -> bool {
matches!(self.0, OutputInner::Stdout)
}
pub fn is_file(&self) -> bool {
matches!(self.0, OutputInner::File { .. })
}
pub fn path(&self) -> Option<&Path> {
match &self.0 {
OutputInner::Stdout => None,
OutputInner::File { path, .. } => Some(path),
}
}
pub fn lock(&self) -> LockedOutput<'_> {
let inner = match &self.0 {
OutputInner::Stdout => {
let writer = io::stdout().lock();
LockedOutputInner::Stdout { writer }
}
OutputInner::File { path, writer: file } => {
let writer = lock(file);
LockedOutputInner::File {
path: Arc::clone(path),
writer,
}
}
};
LockedOutput(inner)
}
}
impl FromStr for Output {
type Err = io::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if s == "-" {
return Ok(Self::stdout());
}
Self::create(PathBuf::from(s))
}
}
macro_rules! with_writer {
($inner:expr, $var:ident => $e:expr) => {
match $inner {
OutputInner::Stdout => {
let mut $var = io::stdout();
$e
}
OutputInner::File { writer, .. } => {
let mut $var = lock(writer);
$e
}
}
};
}
impl Write for Output {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
with_writer!(&self.0, writer => writer.write(buf))
}
fn flush(&mut self) -> io::Result<()> {
with_writer!(&self.0, writer => writer.flush())
}
fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
with_writer!(&self.0, writer => writer.write_vectored(bufs))
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
with_writer!(&self.0, writer => writer.write_all(buf))
}
}
#[derive(Debug)]
pub struct LockedOutput<'a>(LockedOutputInner<'a>);
impl LockedOutput<'_> {
pub fn is_stdin(&self) -> bool {
matches!(self.0, LockedOutputInner::Stdout { .. })
}
pub fn is_file(&self) -> bool {
matches!(self.0, LockedOutputInner::File { .. })
}
pub fn path(&self) -> Option<&Path> {
match &self.0 {
LockedOutputInner::Stdout { .. } => None,
LockedOutputInner::File { path, .. } => Some(path),
}
}
}
#[derive(Debug)]
enum LockedOutputInner<'a> {
Stdout {
writer: io::StdoutLock<'a>,
},
File {
path: Arc<PathBuf>,
writer: MutexGuard<'a, LineWriter<File>>,
},
}
macro_rules! with_locked_writer {
($inner:expr, $var:ident => $e:expr) => {
match $inner {
LockedOutputInner::Stdout { writer } => {
let $var = writer;
$e
}
LockedOutputInner::File { writer, .. } => {
let $var = writer;
$e
}
}
};
}
impl Write for LockedOutput<'_> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
with_locked_writer!(&mut self.0, writer => writer.write(buf))
}
fn flush(&mut self) -> io::Result<()> {
with_locked_writer!(&mut self.0, writer => writer.flush())
}
fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
with_locked_writer!(&mut self.0, writer => writer.write_vectored(bufs))
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
with_locked_writer!(&mut self.0, writer => writer.write_all(buf))
}
}