use std::cell::RefCell;
use std::io;
use std::io::Write;
use qubit_budget::ResourceQuantity;
use qubit_io::ByteOutput;
use super::json_output_accounting::JsonOutputAccounting;
use crate::encode::JsonEncodeError;
pub(in crate::encode) struct JsonOutputWriter<'a, W: ?Sized, R, Q>
where
Q: ResourceQuantity,
{
writer: &'a mut W,
accounting: &'a RefCell<JsonOutputAccounting<'a, R, Q>>,
has_output_budget: bool,
io_error: Option<io::Error>,
}
impl<'a, W: ?Sized, R, Q> JsonOutputWriter<'a, W, R, Q>
where
Q: ResourceQuantity,
{
#[inline]
pub(in crate::encode) fn new(writer: &'a mut W, accounting: &'a RefCell<JsonOutputAccounting<'a, R, Q>>) -> Self {
Self {
writer,
accounting,
has_output_budget: accounting.borrow().has_output_budget(),
io_error: None,
}
}
pub(in crate::encode) fn into_result(
self,
result: Result<(), serde_json::Error>,
) -> Result<(), JsonEncodeError<R, Q>> {
let violation = self.accounting.borrow_mut().take_violation();
if let Some(error) = violation {
return Err(JsonEncodeError::<R, Q>::budget(error));
}
let syntax_error = self.accounting.borrow_mut().take_syntax_error();
if let Some(error) = syntax_error {
return Err(JsonEncodeError::<R, Q>::invalid_raw_json(error));
}
if let Some(error) = self.io_error {
return Err(JsonEncodeError::<R, Q>::write(error));
}
if let Some(error) = self.accounting.borrow_mut().take_serialization_error() {
return Err(JsonEncodeError::<R, Q>::serialization(error));
}
if result.is_err() {
let error = self.accounting.borrow_mut().take_serialization_error_or_custom();
return Err(JsonEncodeError::<R, Q>::serialization(error));
}
Ok(())
}
#[must_use]
fn record_io_error(&mut self, error: io::Error) -> io::Error {
let kind = error.kind();
if self.io_error.is_none() {
self.io_error = Some(error);
}
io::Error::new(kind, "JSON output operation failed")
}
}
impl<O: ?Sized, R, Q> Write for JsonOutputWriter<'_, O, R, Q>
where
O: ByteOutput,
R: Clone,
Q: ResourceQuantity,
{
fn write(&mut self, input: &[u8]) -> io::Result<usize> {
if self.has_output_budget {
let accounting = self.accounting.borrow();
if let Err(error) = accounting.check_available(input.len()) {
drop(accounting);
self.accounting.borrow_mut().record_violation(error);
return Err(io::Error::other("JSON output budget exceeded"));
}
}
match qubit_io::Output::write(&mut *self.writer, input) {
Ok(written) => {
if written == 0 && !input.is_empty() {
let error = io::Error::new(io::ErrorKind::WriteZero, "JSON output writer accepted no bytes");
return Err(self.record_io_error(error));
}
if self.has_output_budget {
let mut accounting = self.accounting.borrow_mut();
if let Err(error) = accounting.consume(written) {
accounting.record_violation(error);
return Err(io::Error::other("JSON output budget exceeded"));
}
}
Ok(written)
}
Err(error) => {
if error.kind() == io::ErrorKind::Interrupted {
Err(error)
} else {
Err(self.record_io_error(error))
}
}
}
}
fn flush(&mut self) -> io::Result<()> {
match qubit_io::Output::flush(&mut *self.writer) {
Ok(()) => Ok(()),
Err(error) => Err(self.record_io_error(error)),
}
}
}