use std::cell::RefCell;
use std::io;
use std::io::Write;
use qubit_budget::ResourceQuantity;
use super::json_output_accounting::JsonOutputAccounting;
use crate::encode::JsonEncodeError;
pub(in crate::encode) struct JsonOutputWriter<'a, W, R, Q>
where
Q: ResourceQuantity,
{
writer: W,
accounting: &'a RefCell<JsonOutputAccounting<'a, R, Q>>,
has_output_budget: bool,
io_error: Option<io::Error>,
}
impl<'a, W, R, Q> JsonOutputWriter<'a, W, R, Q>
where
Q: ResourceQuantity,
{
#[inline]
pub(in crate::encode) fn new(writer: 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 result.is_err() {
let error = self.accounting.borrow_mut().take_serialization_error_or_custom();
return Err(JsonEncodeError::<R, Q>::serialization(error));
}
Ok(())
}
}
impl<W, R, Q> Write for JsonOutputWriter<'_, W, R, Q>
where
W: Write,
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 self.writer.write(input) {
Ok(written) => {
if written == 0 && !input.is_empty() {
let error = io::Error::new(io::ErrorKind::WriteZero, "JSON output writer accepted no bytes");
self.io_error = Some(io::Error::new(error.kind(), error.to_string()));
return Err(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) => {
self.io_error = Some(io::Error::new(error.kind(), error.to_string()));
Err(error)
}
}
}
fn flush(&mut self) -> io::Result<()> {
match self.writer.flush() {
Ok(()) => Ok(()),
Err(error) => {
self.io_error = Some(io::Error::new(error.kind(), error.to_string()));
Err(error)
}
}
}
}