use std::borrow::Cow;
use rudb_common::{Error, Result};
use crate::dialect::Dialect;
pub fn record(
bytes: &[u8],
from: usize,
dialect: Dialect,
eof: bool,
out: &mut Vec<String>,
) -> Result<Option<usize>> {
if from >= bytes.len() {
return Ok(None);
}
let quote = dialect.quote_byte();
let escape = dialect.escape_byte();
let mut at = from;
let mut count = 0;
let mut field = Vec::new();
loop {
field.clear();
if bytes.get(at) == Some("e) {
at += 1;
loop {
let Some(&byte) = bytes.get(at) else { return Ok(None) };
if byte == escape && escape != quote {
let Some(&next) = bytes.get(at + 1) else { return Ok(None) };
field.push(next);
at += 2;
continue;
}
if byte == escape && bytes.get(at + 1) == Some("e) {
field.push(quote);
at += 2;
continue;
}
if byte == quote {
at += 1;
break;
}
field.push(byte);
at += 1;
}
match bytes.get(at) {
None if !eof => return Ok(None),
None => {}
Some(&byte) if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' => {}
Some(&byte) => return Err(after_quote(byte)),
}
} else {
while let Some(&byte) = bytes.get(at) {
if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' {
break;
}
field.push(byte);
at += 1;
}
if at >= bytes.len() && !eof {
return Ok(None);
}
}
place(out, count, &field);
count += 1;
match bytes.get(at) {
Some(&byte) if byte == dialect.delimiter => at += 1,
Some(b'\r') => {
at += 1;
if bytes.get(at) == Some(&b'\n') {
at += 1;
} else if at >= bytes.len() && !eof {
return Ok(None);
}
break;
}
Some(b'\n') => {
at += 1;
break;
}
Some(_) => unreachable!("a field stops at a delimiter, a line ending or the end"),
None => break,
}
}
out.truncate(count);
Ok(Some(at))
}
fn place(out: &mut Vec<String>, at: usize, field: &[u8]) {
let text = String::from_utf8_lossy(field);
match out.get_mut(at) {
Some(held) => {
held.clear();
held.push_str(&text);
}
None => out.push(text.into_owned()),
}
}
fn after_quote(byte: u8) -> Error {
Error::io(format!(
"a quoted value is followed by '{}' rather than by a delimiter or the end of the line",
byte as char
))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Span {
start: u32,
end: u32,
}
const ESCAPED: u32 = 1 << 31;
impl Span {
pub const MOST: usize = (ESCAPED - 1) as usize;
#[must_use]
pub fn new(start: usize, end: usize, escaped: bool) -> Self {
debug_assert!(start <= end && end <= Self::MOST);
#[allow(clippy::cast_possible_truncation)]
let (start, end) = (start as u32, end as u32);
Self { start, end: end | if escaped { ESCAPED } else { 0 } }
}
#[must_use]
pub const fn start(self) -> usize {
self.start as usize
}
#[must_use]
pub const fn end(self) -> usize {
(self.end & !ESCAPED) as usize
}
#[must_use]
pub const fn escaped(self) -> bool {
self.end & ESCAPED != 0
}
#[must_use]
pub const fn is_empty(self) -> bool {
self.start() == self.end()
}
#[must_use]
pub const fn len(self) -> usize {
self.end() - self.start()
}
#[must_use]
pub fn raw(self, bytes: &[u8]) -> &[u8] {
&bytes[self.start()..self.end()]
}
#[must_use]
pub fn text(self, bytes: &[u8], dialect: Dialect) -> Cow<'_, str> {
let raw = self.raw(bytes);
if !self.escaped() {
return String::from_utf8_lossy(raw);
}
Cow::Owned(String::from_utf8_lossy(&unescape(raw, dialect)).into_owned())
}
}
fn unescape(raw: &[u8], dialect: Dialect) -> Vec<u8> {
let quote = dialect.quote_byte();
let escape = dialect.escape_byte();
let mut out = Vec::with_capacity(raw.len());
let mut at = 0;
while let Some(&byte) = raw.get(at) {
if byte == escape
&& escape != quote
&& let Some(&next) = raw.get(at + 1)
{
out.push(next);
at += 2;
continue;
}
if byte == escape && raw.get(at + 1) == Some("e) {
out.push(quote);
at += 2;
} else {
out.push(byte);
at += 1;
}
}
out
}
#[derive(Debug, Clone, Default)]
pub struct Records {
spans: Vec<Span>,
ends: Vec<usize>,
}
impl Records {
pub fn clear(&mut self) {
self.spans.clear();
self.ends.clear();
}
#[must_use]
pub fn len(&self) -> usize {
self.ends.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.ends.is_empty()
}
#[must_use]
pub fn fields(&self, at: usize) -> &[Span] {
let start = if at == 0 { 0 } else { self.ends[at - 1] };
&self.spans[start..self.ends[at]]
}
#[must_use]
pub fn field(&self, row: usize, column: usize) -> Option<Span> {
let start = if row == 0 { 0 } else { self.ends[row - 1] };
let at = start + column;
if at < self.ends[row] { Some(self.spans[at]) } else { None }
}
pub fn shift(&mut self, by: usize) {
if by == 0 {
return;
}
#[allow(clippy::cast_possible_truncation)]
let by = by as u32;
for span in &mut self.spans {
span.start -= by;
span.end -= by;
}
}
}
pub fn records(
bytes: &[u8],
from: usize,
dialect: Dialect,
eof: bool,
limit: usize,
out: &mut Records,
) -> Result<usize> {
let quote = dialect.quote_byte();
let delimiter = dialect.delimiter;
let structural = |byte: u8| byte == quote || byte == b'\n' || byte == b'\r';
if dialect.escape_byte() == quote && !structural(delimiter) {
blocks(bytes, from, dialect, eof, limit, out)
} else {
let mut at = from;
while out.len() < limit {
let Some(next) = spans(bytes, at, dialect, eof, &mut out.spans)? else { break };
out.ends.push(out.spans.len());
at = next;
}
Ok(at)
}
}
fn blocks(
bytes: &[u8],
from: usize,
dialect: Dialect,
eof: bool,
limit: usize,
out: &mut Records,
) -> Result<usize> {
let delimiter = dialect.delimiter;
let quote = dialect.quote_byte();
let len = bytes.len();
let mut record = from;
loop {
let mut mark = out.spans.len();
let mut field = record;
let mut quoted = false;
let mut inside = 0u64;
let mut block = record;
'blocks: while block < len && out.len() < limit {
let end = len.min(block + 64);
let (quotes, ends, lines) = masks(&bytes[block..end], delimiter, quote);
let prefix = prefix_xor(quotes) ^ inside;
inside = 0u64.wrapping_sub(prefix >> 63);
let mut structural = ends & !prefix;
while structural != 0 {
let bit = structural.trailing_zeros() as usize;
structural &= structural - 1;
let at = block + bit;
if at < field {
continue;
}
let low = field.saturating_sub(block);
if quoted || quotes & below(bit) & !below(low) != 0 {
let Some(escaped) = enclosed(&bytes[field..at], quote) else {
break 'blocks;
};
out.spans.push(Span::new(field + 1, at - 1, escaped));
quoted = false;
} else {
out.spans.push(Span::new(field, at, false));
}
if lines >> bit & 1 == 0 {
field = at + 1;
continue;
}
let mut next = at + 1;
if bytes[at] == b'\r' {
match bytes.get(next) {
Some(b'\n') => next += 1,
Some(_) => {}
None if eof => {}
None => {
out.spans.truncate(mark);
return Ok(record);
}
}
}
out.ends.push(out.spans.len());
mark = out.spans.len();
record = next;
field = next;
if out.len() == limit {
return Ok(record);
}
}
if field < end && quotes >> field.saturating_sub(block) != 0 {
quoted = true;
}
block = end;
}
if out.len() >= limit || record >= len {
return Ok(record);
}
out.spans.truncate(mark);
let Some(next) = spans(bytes, record, dialect, eof, &mut out.spans)? else {
return Ok(record);
};
out.ends.push(out.spans.len());
record = next;
}
}
#[inline]
fn masks(block: &[u8], delimiter: u8, quote: u8) -> (u64, u64, u64) {
let needles = [quote, delimiter, b'\n', b'\r'];
let [quotes, delimiters, newlines, returns] = if let Ok(full) = block.try_into() {
rudb_vector::bytes::masks(full, needles)
} else {
let mut padded = [0u8; 64];
padded[..block.len()].copy_from_slice(block);
let live = below(block.len());
rudb_vector::bytes::masks(&padded, needles).map(|mask| mask & live)
};
let lines = newlines | returns;
(quotes, delimiters | lines, lines)
}
const fn prefix_xor(mut bits: u64) -> u64 {
bits ^= bits << 1;
bits ^= bits << 2;
bits ^= bits << 4;
bits ^= bits << 8;
bits ^= bits << 16;
bits ^= bits << 32;
bits
}
const fn below(count: usize) -> u64 {
(1u64 << count) - 1
}
fn enclosed(field: &[u8], quote: u8) -> Option<bool> {
let [first, inner @ .., last] = field else { return None };
if *first != quote || *last != quote {
return None;
}
let mut escaped = false;
let mut at = 0;
while let Some(found) = inner[at..].iter().position(|&byte| byte == quote) {
if inner.get(at + found + 1) != Some("e) {
return None;
}
escaped = true;
at += found + 2;
}
Some(escaped)
}
fn spans(
bytes: &[u8],
from: usize,
dialect: Dialect,
eof: bool,
out: &mut Vec<Span>,
) -> Result<Option<usize>> {
if from >= bytes.len() {
return Ok(None);
}
let mark = out.len();
let quote = dialect.quote_byte();
let escape = dialect.escape_byte();
let mut at = from;
loop {
let span;
if bytes.get(at) == Some("e) {
at += 1;
let start = at;
let mut escaped = false;
loop {
let Some(&byte) = bytes.get(at) else {
out.truncate(mark);
return Ok(None);
};
if byte == escape && escape != quote {
if at + 1 >= bytes.len() {
out.truncate(mark);
return Ok(None);
}
escaped = true;
at += 2;
continue;
}
if byte == escape && bytes.get(at + 1) == Some("e) {
escaped = true;
at += 2;
continue;
}
if byte == quote {
break;
}
at += 1;
}
span = Span::new(start, at, escaped);
at += 1;
match bytes.get(at) {
None if !eof => {
out.truncate(mark);
return Ok(None);
}
None => {}
Some(&byte) if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' => {}
Some(&byte) => {
out.truncate(mark);
return Err(after_quote(byte));
}
}
} else {
let start = at;
while let Some(&byte) = bytes.get(at) {
if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' {
break;
}
at += 1;
}
if at >= bytes.len() && !eof {
out.truncate(mark);
return Ok(None);
}
span = Span::new(start, at, false);
}
out.push(span);
match bytes.get(at) {
Some(&byte) if byte == dialect.delimiter => at += 1,
Some(b'\r') => {
at += 1;
if bytes.get(at) == Some(&b'\n') {
at += 1;
} else if at >= bytes.len() && !eof {
out.truncate(mark);
return Ok(None);
}
break;
}
Some(b'\n') => {
at += 1;
break;
}
Some(_) => unreachable!("a field stops at a delimiter, a line ending or the end"),
None => break,
}
}
Ok(Some(at))
}
#[cfg(test)]
mod tests {
use super::*;
fn split(bytes: &[u8], dialect: Dialect) -> Vec<Vec<String>> {
let mut rows = Vec::new();
let mut fields = Vec::new();
let mut at = 0;
while at < bytes.len() {
let next = record(bytes, at, dialect, true, &mut fields)
.expect("splits")
.expect("a whole record");
rows.push(fields.clone());
at = next;
}
rows
}
fn comma() -> Dialect {
Dialect { delimiter: b',', quote: Some(b'"'), escape: Some(b'"'), header: true }
}
#[test]
fn a_line_of_fields_is_the_fields_of_that_line() {
assert_eq!(split(b"a,b,c\n1,2,3\n", comma()), [["a", "b", "c"], ["1", "2", "3"]]);
}
#[test]
fn the_last_line_does_not_need_a_newline_on_it() {
assert_eq!(split(b"a,b\n1,2", comma()), [["a", "b"], ["1", "2"]]);
}
#[test]
fn all_three_line_endings_end_a_line() {
assert_eq!(split(b"a\r\nb\rc\n", comma()), [["a"], ["b"], ["c"]]);
}
#[test]
fn a_quoted_field_may_hold_the_delimiter_and_a_newline() {
assert_eq!(split(b"1,\"x,y\"\n", comma()), [["1", "x,y"]]);
assert_eq!(split(b"1,\"x\ny\"\n", comma()), [["1", "x\ny"]]);
}
#[test]
fn a_doubled_quote_inside_a_quoted_field_is_one_quote() {
assert_eq!(split(b"1,\"say \"\"hi\"\"\"\n", comma()), [["1", "say \"hi\""]]);
}
#[test]
fn an_empty_field_is_an_empty_string_here_and_becomes_a_null_above() {
assert_eq!(split(b"1,,3\n", comma()), [["1", "", "3"]]);
assert_eq!(split(b"1,\"\",3\n", comma()), [["1", "", "3"]]);
}
#[test]
fn a_trailing_delimiter_makes_a_last_empty_field() {
assert_eq!(split(b"1|x|\n", Dialect { delimiter: b'|', ..comma() }), [["1", "x", ""]]);
}
#[test]
fn a_quote_in_the_middle_of_a_bare_field_is_just_a_character() {
assert_eq!(split(b"1,he said \"hi\"\n", comma()), [["1", "he said \"hi\""]]);
}
#[test]
fn a_record_that_the_buffer_does_not_hold_all_of_is_not_a_record_yet() {
let mut fields = Vec::new();
assert_eq!(record(b"a,b", 0, comma(), false, &mut fields).unwrap(), None);
assert_eq!(record(b"a,\"b", 0, comma(), true, &mut fields).unwrap(), None);
assert_eq!(record(b"a,b\n", 0, comma(), false, &mut fields).unwrap(), Some(4));
}
#[test]
fn rubbish_after_a_closing_quote_is_an_error_rather_than_a_guess() {
let mut fields = Vec::new();
let error = record(b"\"x\"y,2\n", 0, comma(), true, &mut fields).unwrap_err();
assert!(error.message().contains("quoted value"), "{error}");
}
#[test]
fn utf8_survives_being_read_one_byte_at_a_time() {
assert_eq!(split("a,héllo\n".as_bytes(), comma()), [["a", "héllo"]]);
}
#[test]
fn a_backslash_escape_inside_quotes_keeps_whatever_byte_follows_it() {
let backslash = Dialect { escape: Some(b'\\'), ..comma() };
assert_eq!(
split(b"1,\"O\\\"Brien, Pat\"\n2,\"a\\\\b\",\"x\\ny\"\n3,c\\d\n", backslash),
vec![vec!["1", "O\"Brien, Pat"], vec!["2", "a\\b", "xny"], vec!["3", "c\\d"]]
);
let mut fields = Vec::new();
assert!(record(b"\"a\"\"b\"\n", 0, backslash, true, &mut fields).is_err());
}
}
#[cfg(test)]
mod agree {
use super::*;
type Outcome = (Vec<Vec<String>>, usize, Option<String>);
fn by_record(bytes: &[u8], from: usize, dialect: Dialect, eof: bool) -> Outcome {
let mut rows = Vec::new();
let mut fields = Vec::new();
let mut at = from;
loop {
match record(bytes, at, dialect, eof, &mut fields) {
Ok(Some(next)) => {
rows.push(fields.clone());
at = next;
}
Ok(None) => return (rows, at, None),
Err(error) => return (rows, at, Some(error.to_string())),
}
}
}
fn by_chunk(bytes: &[u8], from: usize, dialect: Dialect, eof: bool, limit: usize) -> Outcome {
let mut rows = Vec::new();
let mut at = from;
let mut out = Records::default();
loop {
out.clear();
let result = records(bytes, at, dialect, eof, limit, &mut out);
for row in 0..out.len() {
rows.push(
out.fields(row)
.iter()
.map(|span| span.text(bytes, dialect).into_owned())
.collect(),
);
}
match result {
Ok(next) => {
at = next;
if out.len() < limit {
return (rows, at, None);
}
}
Err(error) => {
let mut from = at;
let mut fields = Vec::new();
for _ in 0..out.len() {
from = record(bytes, from, dialect, eof, &mut fields).unwrap().unwrap();
}
return (rows, from, Some(error.to_string()));
}
}
}
}
fn check(bytes: &[u8], dialect: Dialect, cuts: &[usize]) {
for eof in [true, false] {
let expected = by_record(bytes, 0, dialect, eof);
for limit in [1, 2, 3, 7, 8192] {
assert_eq!(
by_chunk(bytes, 0, dialect, eof, limit),
expected,
"{:?} under {dialect:?}, eof {eof}, {limit} at a time",
String::from_utf8_lossy(bytes),
);
}
}
let whole = by_record(bytes, 0, dialect, true);
for &cut in cuts {
let cut = cut.min(bytes.len());
let (mut rows, at, error) = by_chunk(&bytes[..cut], 0, dialect, false, 8192);
if error.is_some() {
assert_eq!(error, whole.2, "an error in the front is the error in the whole");
continue;
}
let (rest, end, error) = by_chunk(bytes, at, dialect, true, 8192);
rows.extend(rest);
assert_eq!(
(rows, end, error),
whole,
"{:?} under {dialect:?} cut at {cut}",
String::from_utf8_lossy(bytes),
);
}
}
fn dialects() -> [Dialect; 5] {
let comma =
Dialect { delimiter: b',', quote: Some(b'"'), escape: Some(b'"'), header: false };
[
comma,
Dialect { quote: None, escape: None, ..comma },
Dialect { delimiter: b'|', ..comma },
Dialect { delimiter: b'\t', quote: Some(b'\''), escape: Some(b'\''), ..comma },
Dialect { escape: Some(b'\\'), ..comma },
]
}
#[test]
fn the_tricky_ones_split_the_same_both_ways() {
let long = "x".repeat(61);
let cases: Vec<Vec<u8>> = vec![
b"".to_vec(),
b"\n".to_vec(),
b"\n\n".to_vec(),
b"a".to_vec(),
b"a,b,c\n1,2,3\n".to_vec(),
b"a,b\n1,2".to_vec(),
b"a\r\nb\rc\n".to_vec(),
b"a\r".to_vec(),
b"a\r\r\n\n".to_vec(),
b"1,\"x,y\"\n2,\"x\ny\"\n".to_vec(),
b"1,\"say \"\"hi\"\"\"\n".to_vec(),
b"1,,3\n1,\"\",3\n".to_vec(),
b"1|x|\n1,x,\n".to_vec(),
b"1,he said \"hi\"\n2,x\n".to_vec(),
b"1,he said \"hi\n2,x\n".to_vec(),
b"\"x\"y,2\n".to_vec(),
b"\"x\"\"\n".to_vec(),
b"\"x\"\"".to_vec(),
b"a,\"b".to_vec(),
b"a,\"b\"".to_vec(),
b"\"".to_vec(),
b"\"\"\"\"\n".to_vec(),
b"\"a\"\"\",b\n".to_vec(),
b"\"a\\\"b\",c\n\"a\\\\\"\n".to_vec(),
b"\"a\\\\b\",\"x\\ny\",\"\\\n\"\n\"a\\".to_vec(),
b"'a,b'\t'c''d'\n".to_vec(),
b"x\xffy,\"\xfe\"\"\"\n".to_vec(),
"h\u{e9}llo,w\u{f6}rld\n".as_bytes().to_vec(),
format!("{long},\"a\nb\",c\n{long}\r\n\"{long}\"\"{long}\",d\n").into_bytes(),
format!("{long}ab\r\n{long}abc\r\n").into_bytes(),
format!("\"{long}\"\"\",\"\n\"\n{long},x\"\n").into_bytes(),
];
for bytes in &cases {
let cuts: Vec<usize> = (0..=bytes.len()).collect();
for dialect in dialects() {
check(bytes, dialect, &cuts);
}
}
}
struct Rng(u64);
impl Rng {
fn next(&mut self) -> u64 {
self.0 ^= self.0 << 13;
self.0 ^= self.0 >> 7;
self.0 ^= self.0 << 17;
self.0
}
fn below(&mut self, n: usize) -> usize {
(self.next() % n as u64) as usize
}
}
fn generate(rng: &mut Rng) -> Vec<u8> {
const PIECES: [&[u8]; 20] = [
b"a",
b"1",
b"-2.5",
b"xyz",
b",",
b"|",
b"\t",
b"\"",
b"\"\"",
b"'",
b"\\",
b"\n",
b"\r",
b"\r\n",
b" ",
b"\xc3\xa9",
b"\xff",
b"2020-01-02",
b"",
b"0123456789abcdef",
];
let mut out = Vec::new();
let rows = rng.below(40);
for _ in 0..rows {
let fields = 1 + rng.below(5);
for field in 0..fields {
if field > 0 {
out.push(b",,,|\t"[rng.below(5)]);
}
let mut body = Vec::new();
for _ in 0..rng.below(6) {
let piece = PIECES[rng.below(PIECES.len())];
if rng.below(10) == 0 {
body.extend(std::iter::repeat_n(b'q', rng.below(90)));
}
body.extend_from_slice(piece);
}
match rng.below(4) {
0 => {
out.push(b'"');
for &byte in &body {
if byte == b'"' {
out.push(b'"');
}
out.push(byte);
}
out.push(b'"');
}
1 => {
out.push(b'"');
out.extend_from_slice(&body);
out.push(b'"');
}
_ => out.extend_from_slice(&body),
}
}
out.extend_from_slice([&b"\n"[..], b"\r\n", b"\r"][rng.below(3)]);
}
if rng.below(3) == 0 {
out.truncate(out.len().saturating_sub(1 + rng.below(3)));
}
out
}
#[test]
fn thousands_of_generated_files_split_the_same_both_ways() {
let mut rng = Rng(0x9e37_79b9_7f4a_7c15);
for _ in 0..3000 {
let bytes = generate(&mut rng);
let cuts: Vec<usize> = (0..4).map(|_| rng.below(bytes.len() + 1)).collect();
for dialect in dialects() {
check(&bytes, dialect, &cuts);
}
}
}
}