use std::{
fs::File,
io::{self, BufRead, BufReader, Read},
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 Input(InputInner);
#[derive(Debug, Clone)]
enum InputInner {
Stdin,
File {
path: Arc<PathBuf>,
reader: Arc<Mutex<BufReader<File>>>,
},
}
impl Input {
pub fn stdin() -> Self {
Self(InputInner::Stdin)
}
pub fn open(path: PathBuf) -> io::Result<Self> {
let path = Arc::new(path);
let file = File::open(&*path)?;
let reader = Arc::new(Mutex::new(BufReader::new(file)));
Ok(Self(InputInner::File { path, reader }))
}
pub fn is_stdin(&self) -> bool {
matches!(self.0, InputInner::Stdin)
}
pub fn is_file(&self) -> bool {
matches!(self.0, InputInner::File { .. })
}
pub fn path(&self) -> Option<&Path> {
match &self.0 {
InputInner::Stdin => None,
InputInner::File { path, .. } => Some(path),
}
}
pub fn lock(&self) -> LockedInput<'_> {
let inner = match &self.0 {
InputInner::Stdin => {
let reader = io::stdin().lock();
LockedInputInner::Stdin { reader }
}
InputInner::File { path, reader: file } => {
let reader = lock(file);
LockedInputInner::File {
path: Arc::clone(path),
reader,
}
}
};
LockedInput(inner)
}
}
impl FromStr for Input {
type Err = io::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if s == "-" {
return Ok(Self::stdin());
}
Self::open(PathBuf::from(s))
}
}
macro_rules! with_reader {
($inner:expr, $var:ident => $e:expr) => {
match $inner {
InputInner::Stdin => {
let mut $var = io::stdin();
$e
}
InputInner::File { reader, .. } => {
let mut $var = lock(reader);
$e
}
}
};
}
impl Read for Input {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
with_reader!(&self.0, r => r.read(buf))
}
fn read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> {
with_reader!(&self.0, r => r.read_vectored(bufs))
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
with_reader!(&self.0, r => r.read_to_end(buf))
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
with_reader!(&self.0, r => r.read_to_string(buf))
}
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
with_reader!(&self.0, r => r.read_exact(buf))
}
}
#[derive(Debug)]
pub struct LockedInput<'a>(LockedInputInner<'a>);
impl LockedInput<'_> {
pub fn is_stdin(&self) -> bool {
matches!(self.0, LockedInputInner::Stdin { .. })
}
pub fn is_file(&self) -> bool {
matches!(self.0, LockedInputInner::File { .. })
}
pub fn path(&self) -> Option<&Path> {
match &self.0 {
LockedInputInner::Stdin { .. } => None,
LockedInputInner::File { path, .. } => Some(path),
}
}
}
#[derive(Debug)]
enum LockedInputInner<'a> {
Stdin {
reader: io::StdinLock<'a>,
},
File {
path: Arc<PathBuf>,
reader: MutexGuard<'a, BufReader<File>>,
},
}
macro_rules! with_locked_reader {
($inner:expr, $var:ident => $e:expr) => {
match $inner {
LockedInputInner::Stdin { reader } => {
let $var = reader;
$e
}
LockedInputInner::File { reader, .. } => {
let $var = reader;
$e
}
}
};
}
impl Read for LockedInput<'_> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
with_locked_reader!(&mut self.0, r => r.read(buf))
}
fn read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> {
with_locked_reader!(&mut self.0, r => r.read_vectored(bufs))
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
with_locked_reader!(&mut self.0, r => r.read_to_end(buf))
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
with_locked_reader!(&mut self.0, r => r.read_to_string(buf))
}
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
with_locked_reader!(&mut self.0, r => r.read_exact(buf))
}
}
impl BufRead for LockedInput<'_> {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
with_locked_reader!(&mut self.0, r => r.fill_buf())
}
fn consume(&mut self, amt: usize) {
with_locked_reader!(&mut self.0, r => r.consume(amt))
}
}