use std::io;
use std::marker::PhantomData;
use crate::json::traits::{Write, WriteArray, WriteAs, WriteKeyAs, WriteObject};
use crate::num::atof::is_number;
use crate::num::dtoa::{MAX_FLOAT_BYTES, write_f32, write_f64};
use crate::num::itoa::{MAX_INT_DIGITS, write_u64};
use crate::options::{Options, Standard};
use crate::swar::{escape_mask, first_match, load_u64, needs_escape};
pub const DEFAULT_SINK_BUFFER: usize = 8 * 1024;
pub const fn quoted_key<const N: usize>(key: &str) -> [u8; N] {
let bytes = key.as_bytes();
assert!(
N == bytes.len() + 3,
"structio: `quoted_key` was given a length that is not the key's plus its punctuation"
);
let mut out = [0u8; N];
out[0] = b'"';
let mut i = 0;
while i < bytes.len() {
out[1 + i] = bytes[i];
i += 1;
}
out[N - 2] = b'"';
out[N - 1] = b':';
out
}
pub struct Writer<'a, O: Options = Standard> {
buf: Vec<u8>,
limit: usize,
threshold: usize,
sink: Option<Sink<'a>>,
text_from: usize,
depth: usize,
options: PhantomData<fn() -> O>,
}
struct Sink<'a> {
out: &'a mut dyn io::Write,
err: Option<io::Error>,
finished: bool,
}
impl Drop for Sink<'_> {
fn drop(&mut self) {
debug_assert!(
self.finished,
"structio: a sink Writer was dropped without `finish`, truncating its output"
);
}
}
impl<O: Options> Default for Writer<'static, O> {
fn default() -> Self {
Self::new()
}
}
impl<O: Options> Writer<'static, O> {
#[inline]
pub fn new() -> Self {
Writer::appending(Vec::new())
}
#[inline]
pub fn with_capacity(n: usize) -> Self {
Writer::appending(Vec::with_capacity(n))
}
#[inline]
pub fn from_vec(mut buf: Vec<u8>) -> Self {
buf.clear();
Writer::appending(buf)
}
#[inline]
pub fn appending(buf: Vec<u8>) -> Self {
Writer {
limit: buf.capacity(),
threshold: usize::MAX,
text_from: buf.len(),
buf,
sink: None,
depth: 0,
options: PhantomData,
}
}
}
pub(crate) fn append_in_place<O: Options>(
out: &mut Vec<u8>,
make: impl FnOnce(Vec<u8>) -> Writer<'static, O>,
write: impl FnOnce(&mut Writer<'static, O>),
) {
struct Handback<'o, O: Options> {
w: Writer<'static, O>,
out: &'o mut Vec<u8>,
keep: usize,
}
impl<O: Options> Drop for Handback<'_, O> {
fn drop(&mut self) {
let mut buf = core::mem::take(&mut self.w.buf);
buf.truncate(self.keep);
*self.out = buf;
}
}
let keep = out.len();
let mut h = Handback {
w: make(core::mem::take(out)),
out,
keep,
};
write(&mut h.w);
h.keep = usize::MAX;
}
impl<'a, O: Options> Writer<'a, O> {
#[inline]
pub fn to_sink(out: &'a mut dyn io::Write) -> Self {
Self::to_sink_with_capacity(out, DEFAULT_SINK_BUFFER)
}
pub fn to_sink_with_capacity(out: &'a mut dyn io::Write, capacity: usize) -> Self {
let capacity = capacity.max(1);
let buf = Vec::with_capacity(capacity);
Writer {
limit: buf.capacity().min(capacity),
threshold: capacity,
text_from: 0,
buf,
sink: Some(Sink {
out,
err: None,
finished: false,
}),
depth: 0,
options: PhantomData,
}
}
#[inline]
pub fn as_bytes(&self) -> &[u8] {
&self.buf
}
#[inline]
pub fn into_string(self) -> String {
assert!(
core::str::from_utf8(&self.buf[..self.text_from]).is_ok(),
"structio: `Writer::appending` was given bytes that are not UTF-8, so `into_string` \
cannot hand back a `String`; use `into_vec`"
);
debug_assert!(core::str::from_utf8(&self.buf).is_ok());
unsafe { String::from_utf8_unchecked(self.buf) }
}
#[inline]
pub fn into_vec(self) -> Vec<u8> {
self.buf
}
#[inline]
pub fn len(&self) -> usize {
self.buf.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
pub fn finish(mut self) -> io::Result<()> {
let Some(sink) = self.sink.as_mut() else {
return Ok(());
};
sink.finished = true;
if sink.err.is_none()
&& !self.buf.is_empty()
&& let Err(e) = sink.out.write_all(&self.buf)
{
sink.err = Some(e);
}
self.buf.clear();
match sink.err.take() {
Some(e) => Err(e),
None => Ok(()),
}
}
#[inline(always)]
fn room(&mut self, n: usize) {
if self.buf.len() + n > self.limit {
self.spill(n);
}
}
#[cold]
fn spill(&mut self, n: usize) {
self.drain();
if self.buf.len() + n > self.buf.capacity() {
self.buf.reserve(n);
}
self.relimit();
}
#[inline]
fn relimit(&mut self) {
self.limit = self.buf.capacity().min(self.threshold);
}
fn drain(&mut self) {
let n = self.buf.len();
if n < 2 || self.sink.is_none() {
return;
}
let mut cut = n - 1;
while cut > 0 && (self.buf[cut] & 0xC0) == 0x80 {
cut -= 1;
}
if cut == 0 {
return;
}
let sink = self.sink.as_mut().expect("checked above");
if sink.err.is_none()
&& let Err(e) = sink.out.write_all(&self.buf[..cut])
{
sink.err = Some(e);
}
self.buf.copy_within(cut.., 0);
self.buf.truncate(n - cut);
}
#[inline(always)]
pub fn push(&mut self, b: u8) {
assert!(
b.is_ascii(),
"structio: Writer::push requires an ASCII byte"
);
self.room(1);
let len = self.buf.len();
unsafe {
self.buf.as_mut_ptr().add(len).write(b);
self.buf.set_len(len + 1);
}
}
#[inline(always)]
pub fn raw(&mut self, s: &str) {
self.append(s.as_bytes());
}
#[inline(always)]
pub(crate) fn raw_bytes(&mut self, b: &[u8]) {
self.append(b);
}
#[inline(always)]
fn append(&mut self, bytes: &[u8]) {
self.room(bytes.len());
let len = self.buf.len();
unsafe {
core::ptr::copy_nonoverlapping(
bytes.as_ptr(),
self.buf.as_mut_ptr().add(len),
bytes.len(),
);
self.buf.set_len(len + bytes.len());
}
}
#[inline(always)]
fn overwrite_last(&mut self, b: u8) {
let n = self.buf.len();
debug_assert!(n > 0);
self.buf[n - 1] = b;
}
#[inline]
pub(crate) fn line(&mut self) {
if O::PRETTY {
self.push(b'\n');
self.spaces(self.depth * O::INDENT);
}
}
fn spaces(&mut self, mut n: usize) {
const RUN: &[u8; 32] = &[b' '; 32];
while n > 0 {
let take = n.min(RUN.len());
self.append(&RUN[..take]);
n -= take;
}
}
#[inline]
pub(crate) fn item(&mut self) {
if O::NEW_LINES_IN_ARRAYS {
self.line();
} else if O::PRETTY && self.buf.last() == Some(&b',') {
self.push(b' ');
}
}
#[inline]
pub(crate) fn colon(&mut self) {
self.push(b':');
if O::PRETTY {
self.push(b' ');
}
}
#[inline(always)]
fn breaks_lines(bracket: u8) -> bool {
O::PRETTY && (O::NEW_LINES_IN_ARRAYS || matches!(bracket, b'{' | b'}'))
}
#[inline]
pub(crate) fn open(&mut self, bracket: u8) {
self.push(bracket);
if Self::breaks_lines(bracket) {
self.depth += 1;
}
}
#[inline]
pub(crate) fn close(&mut self, bracket: u8) {
if Self::breaks_lines(bracket) {
self.depth -= 1;
}
self.close_over_comma(bracket);
}
#[inline(always)]
pub(crate) fn key(&mut self, prefix: &str) {
self.line();
self.raw(prefix);
if O::PRETTY {
self.push(b' ');
}
}
#[inline(always)]
pub(crate) fn append_fixed<const MAX: usize>(&mut self, src: &[u8; MAX], n: usize) {
assert!(n <= MAX);
self.room(MAX);
let len = self.buf.len();
unsafe {
core::ptr::copy_nonoverlapping(src.as_ptr(), self.buf.as_mut_ptr().add(len), MAX);
self.buf.set_len(len + n);
}
}
#[inline]
pub fn write_object<T: WriteObject>(&mut self, value: &T) {
self.open(b'{');
value.write_fields(self);
self.close(b'}');
}
#[inline(always)]
pub fn member<T: Write + ?Sized>(&mut self, prefix: &str, value: &T) {
if O::SKIP_NULL && value.is_null() {
return;
}
self.key(prefix);
value.write(self);
self.push(b',');
}
#[inline(always)]
pub fn member_with<A: WriteAs<T>, T: ?Sized>(&mut self, prefix: &str, value: &T) {
if O::SKIP_NULL && A::is_null(value) {
return;
}
self.key(prefix);
A::write(value, self);
self.push(b',');
}
#[inline]
pub fn write_tagged<T: Write + ?Sized>(&mut self, prefix: &str, value: &T) {
self.open(b'{');
self.key(prefix);
value.write(self);
self.push(b',');
self.close(b'}');
}
#[inline]
pub fn write_internally_tagged<T: WriteObject + ?Sized>(
&mut self,
prefix: &str,
name: &str,
value: &T,
) {
self.open(b'{');
self.key(prefix);
self.write_str(name);
self.push(b',');
value.write_fields(self);
self.close(b'}');
}
#[inline]
pub fn write_array<T: WriteArray>(&mut self, value: &T) {
self.open(b'[');
value.write_elements(self);
self.close(b']');
}
#[inline(always)]
pub fn element<T: Write + ?Sized>(&mut self, value: &T) {
self.item();
value.write(self);
self.push(b',');
}
#[inline]
pub fn write_seq<'i, T, I>(&mut self, items: I)
where
T: Write + 'i,
I: IntoIterator<Item = &'i T>,
{
self.open(b'[');
for value in items {
self.item();
value.write(self);
self.push(b',');
}
self.close(b']');
}
#[inline]
pub fn write_seq_with<'i, A, T, I>(&mut self, items: I)
where
A: WriteAs<T>,
T: 'i + ?Sized,
I: IntoIterator<Item = &'i T>,
{
self.open(b'[');
for value in items {
self.item();
A::write(value, self);
self.push(b',');
}
self.close(b']');
}
#[inline]
pub fn write_keyed<'i, K, V, I>(&mut self, entries: I)
where
K: crate::json::impls::ToJsonKey + 'i,
V: Write + 'i,
I: IntoIterator<Item = (&'i K, &'i V)>,
{
self.open(b'{');
for (k, v) in entries {
self.line();
k.write_key(self);
self.colon();
v.write(self);
self.push(b',');
}
self.close(b'}');
}
#[inline]
pub fn write_keyed_with<'i, KA, VA, K, V, I>(&mut self, entries: I)
where
KA: WriteKeyAs<K>,
VA: WriteAs<V>,
K: 'i + ?Sized,
V: 'i + ?Sized,
I: IntoIterator<Item = (&'i K, &'i V)>,
{
self.open(b'{');
for (k, v) in entries {
self.line();
KA::write_key(k, self);
self.colon();
VA::write(v, self);
self.push(b',');
}
self.close(b'}');
}
#[inline(always)]
fn close_over_comma(&mut self, close: u8) {
debug_assert!(close.is_ascii());
if Self::breaks_lines(close) {
if self.buf.last() == Some(&b',') {
self.buf.pop();
self.line();
}
self.push(close);
} else if self.buf.last() == Some(&b',') {
self.overwrite_last(close);
} else {
self.push(close);
}
}
#[inline(always)]
pub fn write_bool(&mut self, v: bool) {
if v {
self.append_fixed(b"true\0\0\0\0", 4);
} else {
self.append_fixed(b"false\0\0\0", 5);
}
}
#[inline(always)]
pub fn write_null(&mut self) {
self.append_fixed(b"null\0\0\0\0", 4);
}
#[inline]
pub fn write_u64(&mut self, v: u64) {
let mut tmp = [0u8; MAX_INT_DIGITS];
let n = write_u64(v, &mut tmp);
self.append_fixed(&tmp, n);
}
#[inline]
pub fn write_i64(&mut self, v: i64) {
if v < 0 {
self.push(b'-');
self.write_u64((v as u64).wrapping_neg());
} else {
self.write_u64(v as u64);
}
}
pub fn write_u128(&mut self, v: u128) {
if let Ok(small) = u64::try_from(v) {
return self.write_u64(small);
}
let mut tmp = [0u8; 39];
let mut p = tmp.len();
let mut v = v;
while v > 0 {
p -= 1;
tmp[p] = b'0' + (v % 10) as u8;
v /= 10;
}
self.append(&tmp[p..]);
}
pub fn write_i128_raw(&mut self, v: i128) {
if v < 0 {
self.push(b'-');
self.write_u128((v as u128).wrapping_neg());
} else {
self.write_u128(v as u128);
}
}
#[inline]
pub fn write_f64(&mut self, v: f64) {
let mut tmp = [0u8; MAX_FLOAT_BYTES];
match write_f64(v, &mut tmp) {
Some(n) => self.append_fixed(&tmp, n),
None => self.write_null(),
}
}
#[inline]
pub fn write_f32(&mut self, v: f32) {
let mut tmp = [0u8; MAX_FLOAT_BYTES];
match write_f32(v, &mut tmp) {
Some(n) => self.append_fixed(&tmp, n),
None => self.write_null(),
}
}
#[inline]
pub fn write_number_str(&mut self, s: &str) {
debug_assert!(
is_number(s),
"structio: Writer::write_number_str requires a JSON number literal"
);
self.raw(s);
}
pub fn write_str(&mut self, s: &str) {
let bytes = s.as_bytes();
let n = bytes.len();
self.room(n + 2 + 8);
self.push(b'"');
let mut i = 0;
let mut o = self.buf.len();
while i < n {
if i + 8 <= n {
let word = unsafe { load_u64(bytes, i) };
unsafe { store_u64(&mut self.buf, o, word) };
let m = escape_mask(word);
if m == 0 {
i += 8;
o += 8;
continue;
}
let k = first_match(m);
i += k;
o += k;
} else if !needs_escape(bytes[i]) {
unsafe { store_u8(&mut self.buf, o, bytes[i]) };
i += 1;
o += 1;
continue;
}
let c = bytes[i];
unsafe { self.buf.set_len(o) };
self.write_escape(c);
i += 1;
self.room(n - i + 8);
o = self.buf.len();
}
unsafe { self.buf.set_len(o) };
self.push(b'"');
}
#[cold]
fn write_escape(&mut self, c: u8) {
match c {
b'"' => self.raw_bytes(b"\\\""),
b'\\' => self.raw_bytes(b"\\\\"),
0x08 => self.raw_bytes(b"\\b"),
0x0C => self.raw_bytes(b"\\f"),
b'\n' => self.raw_bytes(b"\\n"),
b'\r' => self.raw_bytes(b"\\r"),
b'\t' => self.raw_bytes(b"\\t"),
_ => {
const HEX: &[u8; 16] = b"0123456789abcdef";
self.raw_bytes(b"\\u00");
self.push(HEX[(c >> 4) as usize]);
self.push(HEX[(c & 0xF) as usize]);
}
}
}
}
#[inline(always)]
unsafe fn store_u64(buf: &mut Vec<u8>, at: usize, word: u64) {
debug_assert!(at + 8 <= buf.capacity());
unsafe { (buf.as_mut_ptr().add(at) as *mut u64).write_unaligned(word.to_le()) };
}
#[inline(always)]
unsafe fn store_u8(buf: &mut Vec<u8>, at: usize, b: u8) {
debug_assert!(at < buf.capacity());
unsafe { buf.as_mut_ptr().add(at).write(b) };
}