use alloc::format;
use alloc::string::String;
use alloc::string::ToString;
use alloc::vec;
use alloc::vec::Vec;
use core::fmt::{self, Write as _};
use deser_core::ext::Number;
use deser_core::ser::SerializeRef;
use deser_core::ser::{self, EventSink, SerializeDriver};
use deser_core::{Atom, BytesFormat, Error, ErrorKind, Event, Serialize, State};
use crate::parser::{Dialect, load_u32, load_u64};
#[inline(always)]
fn load_u128(bytes: &[u8], pos: usize) -> u128 {
u128::from_ne_bytes(*bytes[pos..].first_chunk().unwrap())
}
use crate::{Escape, Nulls, QuoteStyle, Terminator};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SerializerConfig {
delimiter: u8,
quote: Option<u8>,
double_quote: bool,
escape: Escape,
terminator: Terminator,
quote_style: QuoteStyle,
headers: bool,
columns: Option<&'static [&'static str]>,
nulls: Nulls,
flexible: bool,
escape_formulas: bool,
context: deser_core::Context,
}
impl Default for SerializerConfig {
fn default() -> SerializerConfig {
SerializerConfig::new()
}
}
impl SerializerConfig {
pub const fn new() -> SerializerConfig {
SerializerConfig {
delimiter: b',',
quote: Some(b'"'),
double_quote: true,
escape: Escape::None,
terminator: Terminator::Newline,
quote_style: QuoteStyle::Necessary,
headers: true,
columns: None,
nulls: Nulls::None,
flexible: false,
escape_formulas: false,
context: deser_core::Context::new(),
}
}
pub const fn builder() -> SerializerConfigBuilder {
SerializerConfigBuilder::new()
}
pub const fn into_builder(self) -> SerializerConfigBuilder {
SerializerConfigBuilder { value: self }
}
pub fn set_context(&mut self, context: deser_core::Context) {
self.context = context;
}
pub fn context(&self) -> &deser_core::Context {
&self.context
}
#[inline]
fn apply_context(&self, driver: &mut SerializeDriver<'_>) {
if !self.context.is_empty() {
driver.set_default_context(self.context.clone());
}
}
pub const fn tsv() -> SerializerConfig {
let mut config = SerializerConfig::new();
config.set_delimiter(b'\t');
config.set_quote(None);
config.set_escape(Escape::Backslash);
config.set_nulls(Nulls::Text("\\N"));
config
}
pub const fn set_delimiter(&mut self, delimiter: u8) {
self.delimiter = delimiter;
}
pub const fn set_quote(&mut self, quote: Option<u8>) {
self.quote = quote;
}
pub const fn set_double_quote(&mut self, yes: bool) {
self.double_quote = yes;
}
pub const fn set_escape(&mut self, escape: Escape) {
self.escape = escape;
}
pub const fn set_terminator(&mut self, terminator: Terminator) {
self.terminator = terminator;
}
pub const fn set_quote_style(&mut self, style: QuoteStyle) {
self.quote_style = style;
}
pub const fn set_headers(&mut self, yes: bool) {
self.headers = yes;
}
pub const fn set_columns(&mut self, names: &'static [&'static str]) {
self.columns = Some(names);
}
pub const fn set_nulls(&mut self, nulls: Nulls) {
self.nulls = nulls;
}
pub const fn set_flexible(&mut self, yes: bool) {
self.flexible = yes;
}
pub const fn set_escape_formulas(&mut self, yes: bool) {
self.escape_formulas = yes;
}
pub fn to_string<T: Serialize + ?Sized>(&self, value: &T) -> Result<String, Error> {
self.to_string_ref(SerializeRef::new(&value))
}
pub fn to_string_with<F, T: Serialize + ?Sized>(
&self,
value: &T,
setup: F,
) -> Result<String, Error>
where
F: FnOnce(&mut SerializeDriver<'_>),
{
let mut driver = SerializeDriver::new(&value);
setup(&mut driver);
self.apply_context(&mut driver);
let mut out = Vec::new();
self.write_whole(&mut WriterState::default(), &mut driver, true, &mut out)?;
Ok(into_string(out))
}
fn to_string_ref(&self, value: SerializeRef<'_>) -> Result<String, Error> {
let mut driver = SerializeDriver::from_ref(value);
self.apply_context(&mut driver);
let mut out = Vec::new();
self.write_whole(&mut WriterState::default(), &mut driver, true, &mut out)?;
Ok(into_string(out))
}
pub(crate) fn write(
&self,
state: &mut WriterState,
driver: &mut SerializeDriver<'_>,
document: bool,
out: &mut Vec<u8>,
limit: usize,
) -> Result<bool, Error> {
let drive: DriveFn = if limit == usize::MAX {
drive_whole
} else {
drive_partial
};
self.write_with(state, driver, document, out, limit, drive)
}
pub(crate) fn write_whole(
&self,
state: &mut WriterState,
driver: &mut SerializeDriver<'_>,
document: bool,
out: &mut Vec<u8>,
) -> Result<(), Error> {
self.write_with(state, driver, document, out, usize::MAX, drive_whole)
.map(|_| ())
}
fn write_with(
&self,
state: &mut WriterState,
driver: &mut SerializeDriver<'_>,
document: bool,
out: &mut Vec<u8>,
limit: usize,
drive: DriveFn,
) -> Result<bool, Error> {
let dialect = match state.dialect {
Some(ref dialect) => dialect,
None => state.dialect.insert(Dialect::new(
self.delimiter,
self.quote,
self.double_quote,
self.escape,
self.terminator,
None,
)?),
};
let mut writer = RecordWriter {
encoder: FieldEncoder {
config: self,
dialect,
bytes: BytesFormat::of(driver.state()),
plain: matches!(self.quote_style, QuoteStyle::Necessary | QuoteStyle::Never)
&& self.nulls == Nulls::None
&& !self.escape_formulas,
},
names: state.names.take(),
len: state.len,
document,
direct: false,
is_map: false,
fields: 0,
record_start: 0,
field_ends: core::mem::take(&mut state.buffers.field_ends),
record: core::mem::take(&mut state.buffers.record),
scratch: Scratch::new(core::mem::take(&mut state.buffers.scratch)),
open: false,
limit,
out,
};
let had_names = writer.names.is_some();
let rv = drive(driver, &mut writer);
if rv.is_ok() || had_names {
state.names = writer.names;
}
if rv.is_ok() {
state.len = writer.len;
}
state.buffers = Buffers {
field_ends: writer.field_ends,
record: writer.record,
scratch: writer.scratch.bytes,
};
rv
}
}
#[derive(Debug, Clone)]
#[must_use]
pub struct SerializerConfigBuilder {
value: SerializerConfig,
}
impl SerializerConfigBuilder {
pub const fn new() -> SerializerConfigBuilder {
SerializerConfigBuilder {
value: SerializerConfig::new(),
}
}
pub const fn delimiter(mut self, delimiter: u8) -> SerializerConfigBuilder {
self.value.set_delimiter(delimiter);
self
}
pub const fn quote(mut self, quote: Option<u8>) -> SerializerConfigBuilder {
self.value.set_quote(quote);
self
}
pub const fn double_quote(mut self, yes: bool) -> SerializerConfigBuilder {
self.value.set_double_quote(yes);
self
}
pub const fn escape(mut self, escape: Escape) -> SerializerConfigBuilder {
self.value.set_escape(escape);
self
}
pub const fn terminator(mut self, terminator: Terminator) -> SerializerConfigBuilder {
self.value.set_terminator(terminator);
self
}
pub const fn quote_style(mut self, style: QuoteStyle) -> SerializerConfigBuilder {
self.value.set_quote_style(style);
self
}
pub const fn headers(mut self, yes: bool) -> SerializerConfigBuilder {
self.value.set_headers(yes);
self
}
pub const fn columns(mut self, names: &'static [&'static str]) -> SerializerConfigBuilder {
self.value.set_columns(names);
self
}
pub const fn nulls(mut self, nulls: Nulls) -> SerializerConfigBuilder {
self.value.set_nulls(nulls);
self
}
pub const fn flexible(mut self, yes: bool) -> SerializerConfigBuilder {
self.value.set_flexible(yes);
self
}
pub const fn escape_formulas(mut self, yes: bool) -> SerializerConfigBuilder {
self.value.set_escape_formulas(yes);
self
}
pub fn context(mut self, context: deser_core::Context) -> SerializerConfigBuilder {
self.value.set_context(context);
self
}
pub const fn build(self) -> SerializerConfig {
let value = unsafe { core::ptr::read(&self.value) };
core::mem::forget(self);
value
}
}
impl Default for SerializerConfigBuilder {
fn default() -> SerializerConfigBuilder {
SerializerConfigBuilder::new()
}
}
type DriveFn = fn(&mut SerializeDriver<'_>, &mut RecordWriter<'_>) -> Result<bool, Error>;
fn drive_whole(
driver: &mut SerializeDriver<'_>,
writer: &mut RecordWriter<'_>,
) -> Result<bool, Error> {
driver
.drive(|event, state| writer.event(event, state))
.map(|()| true)
}
fn drive_partial(
driver: &mut SerializeDriver<'_>,
writer: &mut RecordWriter<'_>,
) -> Result<bool, Error> {
driver.drive_until(writer)
}
#[derive(Clone, Default)]
struct Buffers {
field_ends: Vec<usize>,
record: Record,
scratch: Vec<u8>,
}
impl core::fmt::Debug for Buffers {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Buffers").finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Default)]
pub(crate) struct WriterState {
names: Option<Vec<String>>,
len: Option<usize>,
dialect: Option<Dialect>,
buffers: Buffers,
}
impl WriterState {
fn with_headers(names: Vec<String>) -> WriterState {
WriterState {
len: Some(names.len()),
names: Some(names),
dialect: None,
buffers: Buffers::default(),
}
}
}
#[derive(Debug, Clone)]
pub struct Serializer {
config: SerializerConfig,
state: WriterState,
out: Vec<u8>,
document: bool,
in_progress: bool,
}
impl Default for Serializer {
fn default() -> Serializer {
Serializer::new()
}
}
impl Serializer {
pub fn new() -> Serializer {
Serializer::with_config(SerializerConfig::new())
}
pub fn with_config(config: SerializerConfig) -> Serializer {
Serializer::with_state(config, WriterState::default(), false)
}
pub fn with_headers<I, S>(config: SerializerConfig, names: I) -> Serializer
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
let names = names.into_iter().map(Into::into).collect();
Serializer::with_state(config, WriterState::with_headers(names), false)
}
pub fn document(config: SerializerConfig) -> Serializer {
Serializer::with_state(config, WriterState::default(), true)
}
fn with_state(config: SerializerConfig, state: WriterState, document: bool) -> Serializer {
Serializer {
config,
state,
out: Vec::new(),
document,
in_progress: false,
}
}
pub fn config(&self) -> &SerializerConfig {
&self.config
}
pub fn headers(&self) -> Option<&[String]> {
self.state.names.as_deref()
}
pub fn serialize<T: Serialize + ?Sized>(&mut self, value: &T) -> Result<(), Error> {
ser::Serializer::serialize(self, value)
}
pub fn serialize_with<F, T: Serialize + ?Sized>(
&mut self,
value: &T,
setup: F,
) -> Result<(), Error>
where
F: FnOnce(&mut SerializeDriver<'_>),
{
ser::Serializer::serialize_with(self, value, setup)
}
pub fn as_str(&self) -> &str {
unsafe { core::str::from_utf8_unchecked(&self.out) }
}
pub fn finish(self) -> String {
into_string(self.out)
}
}
impl ser::Serializer for Serializer {
fn drive(&mut self, driver: &mut SerializeDriver<'_>) -> Result<(), Error> {
if !self.config.context.is_empty() {
driver.set_default_context(self.config.context.clone());
}
if self.in_progress {
return Err(Error::in_progress());
}
let len = self.out.len();
match self
.config
.write_whole(&mut self.state, driver, self.document, &mut self.out)
{
Ok(()) => Ok(()),
Err(err) => {
self.out.truncate(len);
Err(err)
}
}
}
}
impl ser::StreamSerializer for Serializer {
fn output(&self) -> &[u8] {
&self.out
}
fn clear_output(&mut self) {
self.out.clear();
}
fn supports_partial(&self) -> bool {
self.document
}
fn drive_partial(
&mut self,
driver: &mut SerializeDriver<'_>,
limit: usize,
) -> Result<bool, Error> {
if !self.config.context.is_empty() {
driver.set_default_context(self.config.context.clone());
}
if !self.document || (limit == usize::MAX && !self.in_progress) {
ser::Serializer::drive(self, driver)?;
return Ok(true);
}
let len = self.out.len();
match self
.config
.write(&mut self.state, driver, true, &mut self.out, limit)
{
Ok(done) => {
self.in_progress = !done;
Ok(done)
}
Err(err) => {
self.out.truncate(len);
Err(err)
}
}
}
fn in_progress(&self) -> bool {
self.in_progress
}
}
#[cfg(feature = "io")]
impl SerializerConfig {
pub fn writer<W: std::io::Write>(&self, writer: W) -> deser_core::io::Writer<W, Serializer> {
deser_core::io::Writer::new(writer, Serializer::with_config(self.clone()))
}
pub fn to_writer<W: std::io::Write, T: Serialize + ?Sized>(
&self,
writer: W,
value: &T,
) -> Result<(), Error> {
deser_core::io::to_writer(writer, Serializer::document(self.clone()), value)
}
}
#[cfg(feature = "io")]
pub fn to_writer<W: std::io::Write, T: Serialize + ?Sized>(
writer: W,
value: &T,
) -> Result<(), Error> {
SerializerConfig::new().to_writer(writer, value)
}
fn into_string(out: Vec<u8>) -> String {
debug_assert!(core::str::from_utf8(&out).is_ok());
unsafe { String::from_utf8_unchecked(out) }
}
pub fn to_string<T: Serialize + ?Sized>(value: &T) -> Result<String, Error> {
SerializerConfig::new().to_string(value)
}
#[derive(Clone, Default)]
struct Record {
keys: Vec<u8>,
key_ends: Vec<usize>,
fields: Vec<u8>,
field_ends: Vec<usize>,
}
impl Record {
fn clear(&mut self) {
self.keys.clear();
self.key_ends.clear();
self.fields.clear();
self.field_ends.clear();
}
fn key(&self, index: usize) -> &[u8] {
let start = if index == 0 {
0
} else {
self.key_ends[index - 1]
};
&self.keys[start..self.key_ends[index]]
}
fn field(&self, index: usize) -> &[u8] {
let start = if index == 0 {
0
} else {
self.field_ends[index - 1]
};
&self.fields[start..self.field_ends[index]]
}
}
struct RecordWriter<'a> {
encoder: FieldEncoder<'a>,
names: Option<Vec<String>>,
len: Option<usize>,
document: bool,
direct: bool,
is_map: bool,
fields: usize,
record_start: usize,
field_ends: Vec<usize>,
record: Record,
scratch: Scratch,
open: bool,
limit: usize,
out: &'a mut Vec<u8>,
}
impl EventSink for RecordWriter<'_> {
#[inline]
fn event(
&mut self,
event: Event<'_>,
_value: SerializeRef<'_>,
state: &mut State,
) -> Result<(), Error> {
RecordWriter::event(self, event, state)
}
#[inline]
fn pause(&mut self) -> bool {
!self.open && self.out.len() >= self.limit
}
}
impl RecordWriter<'_> {
#[inline]
fn event(&mut self, event: Event<'_>, state: &State) -> Result<(), Error> {
if let Event::Atom(ref atom) = event
&& state.depth() == usize::from(self.document) + 1
{
return self.atom(atom, state.is_map_key());
}
self.structure(event, state)
}
fn structure(&mut self, event: Event<'_>, state: &State) -> Result<(), Error> {
let depth = match event {
Event::MapStart(_) | Event::SeqStart(_) => state.depth().saturating_sub(1),
_ => state.depth(),
};
let record_depth = usize::from(self.document);
match event {
Event::SeqStart(_) | Event::SeqEnd if self.document && depth == 0 => {}
_ if depth < record_depth => {
return Err(Error::new(
ErrorKind::UnsupportedType,
"CSV documents are sequences of records",
));
}
Event::MapStart(_) | Event::SeqStart(_) if depth == record_depth => {
if self.names.is_none()
&& let Some(columns) = self.encoder.config.columns
{
let names: Vec<String> = columns.iter().map(|name| name.to_string()).collect();
if self.encoder.config.headers {
self.write_names(&names)?;
}
self.names = Some(names);
}
self.is_map = matches!(event, Event::MapStart(_));
self.open = true;
self.direct = !self.is_map || self.names.is_some();
self.fields = 0;
self.record_start = self.out.len();
self.field_ends.clear();
self.record.clear();
}
Event::MapEnd | Event::SeqEnd if depth == record_depth => {
self.open = false;
self.finish_record()?
}
_ if depth == record_depth => {
return Err(Error::new(
ErrorKind::UnsupportedType,
"CSV records must be maps or sequences",
));
}
Event::MapStart(_) | Event::SeqStart(_) | Event::MapEnd | Event::SeqEnd => {
return Err(Error::new(
ErrorKind::UnsupportedType,
"CSV fields cannot hold maps or sequences",
));
}
Event::Atom(ref atom) => return self.atom(atom, state.is_map_key()),
}
Ok(())
}
#[inline(always)]
fn atom(&mut self, atom: &Atom<'_>, is_key: bool) -> Result<(), Error> {
if is_key {
return self.key(atom);
}
let text = self.encoder.text(atom, &mut self.scratch)?;
if self.direct {
if self.fields > 0 {
self.out.push(self.encoder.dialect.delimiter);
}
self.encoder.encode(text, self.out)?;
self.field_ends.push(self.out.len());
} else {
self.encoder.encode(text, &mut self.record.fields)?;
self.record.field_ends.push(self.record.fields.len());
}
self.fields += 1;
Ok(())
}
#[inline(always)]
fn key(&mut self, atom: &Atom<'_>) -> Result<(), Error> {
if self.direct
&& let Atom::Str(key) = atom
&& let Some(names) = &self.names
&& let Some(name) = names.get(self.fields)
&& same_key(name.as_bytes(), key.as_bytes())
{
return Ok(());
}
self.other_key(atom)
}
#[inline(never)]
fn other_key(&mut self, atom: &Atom<'_>) -> Result<(), Error> {
let key = match atom {
Atom::Null | Atom::Bytes(_) => None,
atom => self.encoder.text(atom, &mut self.scratch)?,
};
let key = key.ok_or_else(unsupported_key)?;
if !self.direct {
self.record.keys.extend_from_slice(key.bytes);
} else {
let names = self.names.as_deref().unwrap_or_default();
if let Some(name) = names.get(self.fields)
&& same_key(name.as_bytes(), key.bytes)
{
return Ok(());
}
let key = key.bytes.to_vec();
self.collect();
self.record.keys.extend_from_slice(&key);
}
self.record.key_ends.push(self.record.keys.len());
Ok(())
}
fn collect(&mut self) {
let names = self.names.as_deref().unwrap_or_default();
let mut start = self.record_start;
for (index, &end) in self.field_ends.iter().enumerate() {
if index > 0 {
start += 1;
}
self.record.fields.extend_from_slice(&self.out[start..end]);
self.record.field_ends.push(self.record.fields.len());
self.record.keys.extend_from_slice(names[index].as_bytes());
self.record.key_ends.push(self.record.keys.len());
start = end;
}
self.out.truncate(self.record_start);
self.direct = false;
}
fn finish_record(&mut self) -> Result<(), Error> {
let start = self.record_start;
if !self.is_map {
self.check_len(self.fields)?;
return self.terminate(self.fields, start);
}
if self.direct {
let len = self.names.as_ref().map_or(0, Vec::len);
for column in self.fields..len {
if column > 0 {
self.out.push(self.encoder.dialect.delimiter);
}
self.encoder.encode(None, self.out)?;
}
self.check_len(len)?;
return self.terminate(len, start);
}
let count = self.record.field_ends.len();
if self.names.is_none() {
let names: Vec<String> = (0..count)
.map(|index| String::from_utf8_lossy(self.record.key(index)).into_owned())
.collect();
if self.encoder.config.headers {
self.write_names(&names)?;
}
self.names = Some(names);
}
let len = self.names.as_ref().map_or(0, Vec::len);
self.check_len(len)?;
let names = self.names.as_ref().unwrap();
let start = self.out.len();
let mut order = vec![None; len];
for index in 0..count {
let key = String::from_utf8_lossy(self.record.key(index));
match names.iter().position(|name| *name == key) {
Some(column) if order[column].is_none() => order[column] = Some(index),
Some(_) => {
return Err(Error::new(
ErrorKind::DuplicateKey,
format!("duplicate field `{}`", key),
));
}
None => {
return Err(Error::new(
ErrorKind::UnknownField,
format!("field `{}` is not a column", key),
));
}
}
}
for (column, index) in order.into_iter().enumerate() {
if column > 0 {
self.out.push(self.encoder.dialect.delimiter);
}
match index {
Some(index) => self.out.extend_from_slice(self.record.field(index)),
None => self.encoder.encode(None, self.out)?,
}
}
self.terminate(len, start)
}
fn check_len(&mut self, len: usize) -> Result<(), Error> {
match self.len {
Some(expected) if expected != len && !self.encoder.config.flexible => Err(Error::new(
ErrorKind::WrongLength,
format!("record has {} fields, expected {}", len, expected),
)),
Some(_) => Ok(()),
None => {
self.len = Some(len);
Ok(())
}
}
}
fn terminate(&mut self, len: usize, start: usize) -> Result<(), Error> {
if len == 1 && self.out.len() == start {
match self.encoder.dialect.quote {
Some(quote) => self.out.extend_from_slice(&[quote, quote]),
None => {
return Err(Error::new(
ErrorKind::InvalidValue,
"a record with a single empty field needs quotes",
));
}
}
}
match self.encoder.config.terminator {
Terminator::Newline => self.out.push(b'\n'),
Terminator::CrLf => self.out.extend_from_slice(b"\r\n"),
Terminator::Byte(byte) => self.out.push(byte),
}
Ok(())
}
fn write_names(&mut self, names: &[String]) -> Result<(), Error> {
let start = self.out.len();
for (index, name) in names.iter().enumerate() {
if index > 0 {
self.out.push(self.encoder.dialect.delimiter);
}
let text = Text {
bytes: name.as_bytes(),
numeric: false,
};
self.encoder.encode(Some(text), self.out)?;
}
self.terminate(names.len(), start)
}
}
#[derive(Clone, Copy)]
struct Text<'a> {
bytes: &'a [u8],
numeric: bool,
}
#[derive(Clone, Copy)]
struct FieldEncoder<'a> {
config: &'a SerializerConfig,
dialect: &'a Dialect,
bytes: BytesFormat,
plain: bool,
}
impl FieldEncoder<'_> {
#[inline(always)]
fn encode(&self, text: Option<Text<'_>>, out: &mut Vec<u8>) -> Result<(), Error> {
if self.plain
&& let Some(Text { bytes, .. }) = text
&& !self.dialect.has_special(bytes)
{
push_bytes(out, bytes);
return Ok(());
}
self.encode_special(text, out)
}
#[inline(never)]
fn encode_special(&self, text: Option<Text<'_>>, out: &mut Vec<u8>) -> Result<(), Error> {
let config = self.config;
let Some(Text { bytes, numeric }) = text else {
if let Nulls::Text(null) = config.nulls {
out.extend_from_slice(null.as_bytes());
}
return Ok(());
};
let reads_as_null = match config.nulls {
Nulls::None => false,
Nulls::Empty => bytes.is_empty(),
Nulls::Text(null) => bytes == null.as_bytes(),
};
let formula = config.escape_formulas
&& !numeric
&& matches!(
bytes.first(),
Some(b'=' | b'+' | b'-' | b'@' | b'\t' | b'\r')
);
let has_special = self.dialect.has_special(bytes);
let quote_style = match config.quote_style {
QuoteStyle::Always => true,
QuoteStyle::NonNumeric => !numeric,
QuoteStyle::Necessary | QuoteStyle::Never => false,
};
if !quote_style && !has_special && !reads_as_null && !formula {
out.extend_from_slice(bytes);
return Ok(());
}
let prefix: &[u8] = if formula { b"'" } else { b"" };
let escape = config.escape.byte();
let prefix_is_special = formula && self.dialect.is_special(b'\'');
let quoted = quote_style
|| (formula && self.dialect.quote.is_some() && config.quote_style != QuoteStyle::Never)
|| (escape.is_none() && (has_special || prefix_is_special || reads_as_null))
|| (reads_as_null && bytes.is_empty());
let text = prefix.iter().chain(bytes).copied();
if !quoted {
for (index, byte) in text.enumerate() {
if self.dialect.is_special(byte) || (index == 0 && reads_as_null) {
out.push(escape.unwrap_or(b'\\'));
out.push(self.escaped(byte));
} else {
out.push(byte);
}
}
return Ok(());
}
let quote = match self.dialect.quote {
Some(quote) if config.quote_style != QuoteStyle::Never => quote,
_ => {
return Err(Error::new(
ErrorKind::InvalidValue,
format!(
"field {:?} needs to be quoted",
String::from_utf8_lossy(bytes)
),
));
}
};
out.push(quote);
if formula {
self.push_quoted(b'\'', quote, out)?;
}
let mut bytes = bytes;
loop {
let run = bytes
.iter()
.position(|&b| b == quote || Some(b) == escape)
.unwrap_or(bytes.len());
out.extend_from_slice(&bytes[..run]);
let Some((&byte, rest)) = bytes[run..].split_first() else {
break;
};
bytes = rest;
self.push_quoted(byte, quote, out)?;
}
out.push(quote);
Ok(())
}
fn push_quoted(&self, byte: u8, quote: u8, out: &mut Vec<u8>) -> Result<(), Error> {
let escape = self.config.escape.byte();
if byte == quote && self.config.double_quote {
out.extend_from_slice(&[quote, quote]);
} else if byte == quote || Some(byte) == escape {
match escape {
Some(escape) => out.extend_from_slice(&[escape, self.escaped(byte)]),
None => {
return Err(Error::new(
ErrorKind::InvalidValue,
"quotes in quoted fields need to be doubled or escaped",
));
}
}
} else {
out.push(byte);
}
Ok(())
}
fn escaped(&self, byte: u8) -> u8 {
match (self.config.escape, byte) {
(Escape::Backslash, b'\t') => b't',
(Escape::Backslash, b'\n') => b'n',
(Escape::Backslash, b'\r') => b'r',
(_, byte) => byte,
}
}
#[inline(always)]
fn text<'a>(
&self,
atom: &'a Atom<'_>,
scratch: &'a mut Scratch,
) -> Result<Option<Text<'a>>, Error> {
let (bytes, numeric) = match *atom {
Atom::Str(ref value) | Atom::Lexical(ref value) => (value.as_bytes(), false),
Atom::Null => return Ok(None),
Atom::Bool(value) => (if value { &b"true"[..] } else { b"false" }, false),
Atom::U64(value) => (scratch.int.format(value).as_bytes(), true),
Atom::I64(value) => (scratch.int.format(value).as_bytes(), true),
Atom::F64(value) => (scratch.float(value), true),
Atom::F32(value) => (scratch.float(value), true),
_ => return self.other_text(atom, &mut scratch.bytes),
};
Ok(Some(Text { bytes, numeric }))
}
#[inline(never)]
fn other_text<'a>(
&self,
atom: &'a Atom<'_>,
scratch: &'a mut Vec<u8>,
) -> Result<Option<Text<'a>>, Error> {
scratch.clear();
let numeric = match *atom {
Atom::Char(value) => {
scratch.extend_from_slice(value.encode_utf8(&mut [0; 4]).as_bytes());
false
}
Atom::Bytes(ref bytes) => {
let format = bytes.fallback.copied().unwrap_or(self.bytes);
let text = format
.encode(bytes)
.or_else(|| BytesFormat::BASE64.encode(bytes))
.unwrap_or_default();
scratch.extend_from_slice(text.as_bytes());
false
}
Atom::Ext(ref ext) => {
if let Some(number) = ext.downcast_value_ref::<Number>() {
scratch.extend_from_slice(number.as_str().as_bytes());
} else if let Some(value) = ext.downcast_ref::<u128>() {
let _ = write!(ByteWriter(scratch), "{}", value);
} else if let Some(value) = ext.downcast_ref::<i128>() {
let _ = write!(ByteWriter(scratch), "{}", value);
} else {
return match ext.fallback() {
Atom::Ext(_) => Err(Error::new(
ErrorKind::UnsupportedType,
format!("CSV does not support {}", ext.name()),
)),
fallback => {
let mut inner = Scratch::new(Vec::new());
let numeric = match self.text(&fallback, &mut inner)? {
Some(text) => {
let numeric = text.numeric;
scratch.extend_from_slice(text.bytes);
numeric
}
None => return Ok(None),
};
Ok(Some(Text {
bytes: scratch,
numeric,
}))
}
};
}
true
}
Atom::Implicit(ref value) => {
let mut inner = Scratch::new(Vec::new());
return Ok(match self.text(&value.value().to_atom(), &mut inner)? {
Some(text) => {
let numeric = text.numeric;
scratch.extend_from_slice(text.bytes);
Some(Text {
bytes: scratch,
numeric,
})
}
None => None,
});
}
_ => {
return Err(Error::new(
ErrorKind::UnsupportedType,
format!("CSV does not support {}", atom.name()),
));
}
};
Ok(Some(Text {
bytes: scratch,
numeric,
}))
}
}
struct ByteWriter<'a>(&'a mut Vec<u8>);
impl fmt::Write for ByteWriter<'_> {
fn write_str(&mut self, s: &str) -> fmt::Result {
self.0.extend_from_slice(s.as_bytes());
Ok(())
}
}
struct Scratch {
bytes: Vec<u8>,
int: itoa::Buffer,
float: zmij::Buffer,
}
impl Scratch {
fn new(bytes: Vec<u8>) -> Scratch {
Scratch {
bytes,
int: itoa::Buffer::new(),
float: zmij::Buffer::new(),
}
}
#[inline]
fn float<F: zmij::Float>(&mut self, value: F) -> &[u8] {
self.float.format(value).as_bytes()
}
}
#[inline(always)]
fn push_bytes(out: &mut Vec<u8>, bytes: &[u8]) {
let len = bytes.len();
if len > 32 {
out.extend_from_slice(bytes);
return;
}
out.reserve(32);
let start = out.len();
unsafe {
let dst = out.as_mut_ptr().add(start);
match len {
0 => {}
1..=3 => {
*dst = bytes[0];
*dst.add(len / 2) = bytes[len / 2];
*dst.add(len - 1) = bytes[len - 1];
}
4..=8 => {
dst.cast::<u32>().write_unaligned(load_u32(bytes, 0));
dst.add(len - 4)
.cast::<u32>()
.write_unaligned(load_u32(bytes, len - 4));
}
9..=16 => {
dst.cast::<u64>().write_unaligned(load_u64(bytes, 0));
dst.add(len - 8)
.cast::<u64>()
.write_unaligned(load_u64(bytes, len - 8));
}
_ => {
dst.cast::<u128>().write_unaligned(load_u128(bytes, 0));
dst.add(len - 16)
.cast::<u128>()
.write_unaligned(load_u128(bytes, len - 16));
}
}
out.set_len(start + len);
}
}
#[inline(always)]
fn same_key(name: &[u8], key: &[u8]) -> bool {
let len = name.len();
if len != key.len() {
return false;
}
match len {
0 => true,
1..=3 => {
name[0] == key[0] && name[len / 2] == key[len / 2] && name[len - 1] == key[len - 1]
}
4..=8 => {
load_u32(name, 0) == load_u32(key, 0)
&& load_u32(name, len - 4) == load_u32(key, len - 4)
}
9..=16 => {
load_u64(name, 0) == load_u64(key, 0)
&& load_u64(name, len - 8) == load_u64(key, len - 8)
}
_ => name == key,
}
}
#[cold]
fn unsupported_key() -> Error {
Error::new(
ErrorKind::UnsupportedType,
"the keys of records must be strings, numbers or booleans",
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_push_bytes() {
let text: Vec<u8> = (b'a'..=b'z').chain(b'A'..=b'Z').collect();
for len in 0..=text.len() {
for prefix in [0, 1, 7] {
let mut out = vec![b'-'; prefix];
push_bytes(&mut out, &text[..len]);
assert_eq!(&out[prefix..], &text[..len]);
assert_eq!(out.len(), prefix + len);
}
}
}
#[test]
fn test_same_key() {
let text: Vec<u8> = (b'a'..=b'z').chain(b'A'..=b'Z').collect();
for len in 0..=text.len() {
let name = &text[..len];
assert!(same_key(name, &text[..len]));
if len > 0 {
assert!(!same_key(name, &text[..len - 1]));
assert!(!same_key(&text[..len - 1], name));
}
for pos in 0..len {
let mut key = name.to_vec();
key[pos] ^= 1;
assert!(!same_key(name, &key), "{:?}", key);
}
}
}
}