mod any;
mod fixed;
mod varlen;
use std::str;
use bigdecimal::BigDecimal;
use num_bigint::BigInt;
use reifydb_value::{
reifydb_assertions,
util::bitvec::BitVec,
value::{
container::{blob::BlobContainer, number::NumberContainer, utf8::Utf8Container},
datetime::DateTime,
decimal::Decimal,
frame::{column::FrameColumn, data::FrameColumnData, frame::Frame},
int::Int,
row_number::RowNumber,
uint::Uint,
value_type::ValueType,
},
};
use crate::{
error::DecodeError,
frame::{
encoding::dict::{decode_dict_blob, decode_dict_table_bytes, decode_dict_utf8, read_index},
format::{
COL_FLAG_HAS_NONES, COLUMN_DESCRIPTOR_SIZE, Encoding, FRAME_HEADER_SIZE, MESSAGE_HEADER_SIZE,
META_HAS_CREATED_AT, META_HAS_ROW_NUMBERS, META_HAS_UPDATED_AT, RBCF_MAGIC, RBCF_VERSION,
dict_index_width_from_flags,
},
},
tag::TypeTag,
};
pub fn decode_frames(data: &[u8]) -> Result<Vec<Frame>, DecodeError> {
let mut pos = 0;
check_len(data, pos, MESSAGE_HEADER_SIZE)?;
let magic = read_u32(data, pos);
pos += 4;
if magic != RBCF_MAGIC {
return Err(DecodeError::InvalidMagic(magic));
}
let version = read_u16(data, pos);
pos += 2;
if version != RBCF_VERSION {
return Err(DecodeError::UnsupportedVersion(version));
}
let _flags = read_u16(data, pos);
pos += 2;
let frame_count = read_u32(data, pos) as usize;
pos += 4;
let _total_size = read_u32(data, pos) as usize;
pos += 4;
let mut frames = Vec::with_capacity(frame_count);
for _ in 0..frame_count {
let (frame, new_pos) = decode_frame(data, pos)?;
frames.push(frame);
pos = new_pos;
}
reifydb_assertions! {
assert!(
pos == _total_size,
"the RBCF message header declared a total size that disagrees with the bytes consumed while decoding, so the message is truncated or carries trailing bytes and a peer would mis-frame the following message (declared={} consumed={})",
_total_size,
pos
);
}
Ok(frames)
}
struct FrameHeader {
row_count: usize,
column_count: usize,
meta_flags: u8,
}
fn decode_frame(data: &[u8], start: usize) -> Result<(Frame, usize), DecodeError> {
let (header, pos) = read_frame_header(data, start)?;
let (row_numbers, pos) = read_row_numbers(data, pos, header.row_count, header.meta_flags)?;
let (created_at, pos) =
read_datetime_array(data, pos, header.row_count, header.meta_flags, META_HAS_CREATED_AT)?;
let (updated_at, pos) =
read_datetime_array(data, pos, header.row_count, header.meta_flags, META_HAS_UPDATED_AT)?;
let (columns, pos) = read_frame_columns(data, pos, header.column_count)?;
Ok((
Frame {
row_numbers,
created_at,
updated_at,
columns,
},
pos,
))
}
#[inline]
fn read_frame_header(data: &[u8], start: usize) -> Result<(FrameHeader, usize), DecodeError> {
let mut pos = start;
check_len(data, pos, FRAME_HEADER_SIZE)?;
let row_count = read_u32(data, pos) as usize;
pos += 4;
let column_count = read_u16(data, pos) as usize;
pos += 2;
let meta_flags = data[pos];
pos += 1;
let _reserved = data[pos];
pos += 1;
let _frame_size = read_u32(data, pos);
pos += 4;
Ok((
FrameHeader {
row_count,
column_count,
meta_flags,
},
pos,
))
}
#[inline]
fn read_row_numbers(
data: &[u8],
mut pos: usize,
row_count: usize,
meta_flags: u8,
) -> Result<(Vec<RowNumber>, usize), DecodeError> {
if meta_flags & META_HAS_ROW_NUMBERS == 0 {
return Ok((Vec::new(), pos));
}
check_len(data, pos, row_count * 8)?;
let mut row_numbers = Vec::with_capacity(row_count);
for _ in 0..row_count {
let v = read_u64(data, pos);
pos += 8;
row_numbers.push(RowNumber::new(v));
}
Ok((row_numbers, pos))
}
#[inline]
fn read_datetime_array(
data: &[u8],
mut pos: usize,
row_count: usize,
meta_flags: u8,
flag: u8,
) -> Result<(Vec<DateTime>, usize), DecodeError> {
if meta_flags & flag == 0 {
return Ok((Vec::new(), pos));
}
check_len(data, pos, row_count * 8)?;
let mut values = Vec::with_capacity(row_count);
for _ in 0..row_count {
let v = read_u64(data, pos);
pos += 8;
values.push(DateTime::from_nanos(v));
}
Ok((values, pos))
}
#[inline]
fn read_frame_columns(
data: &[u8],
mut pos: usize,
column_count: usize,
) -> Result<(Vec<FrameColumn>, usize), DecodeError> {
let mut columns = Vec::with_capacity(column_count);
for _ in 0..column_count {
let (col, new_pos) = decode_column(data, pos)?;
columns.push(col);
pos = new_pos;
}
Ok((columns, pos))
}
fn decode_column(data: &[u8], start: usize) -> Result<(FrameColumn, usize), DecodeError> {
let mut pos = start;
check_len(data, pos, COLUMN_DESCRIPTOR_SIZE)?;
let type_code = data[pos];
pos += 1;
let encoding_byte = data[pos];
pos += 1;
let flags = data[pos];
pos += 1;
let _reserved = data[pos];
pos += 1;
let name_len = read_u16(data, pos) as usize;
pos += 2;
let _reserved2 = read_u16(data, pos);
pos += 2;
let row_count = read_u32(data, pos) as usize;
pos += 4;
let nones_len = read_u32(data, pos) as usize;
pos += 4;
let data_len = read_u32(data, pos) as usize;
pos += 4;
let offsets_len = read_u32(data, pos) as usize;
pos += 4;
let extra_len = read_u32(data, pos) as usize;
pos += 4;
let encoding = Encoding::from_u8(encoding_byte).ok_or(DecodeError::UnknownEncoding(encoding_byte))?;
let has_nones = flags & COL_FLAG_HAS_NONES != 0;
check_len(data, pos, name_len)?;
let name = str::from_utf8(&data[pos..pos + name_len])
.map_err(|e| DecodeError::InvalidData(format!("invalid column name: {}", e)))?
.to_string();
pos += name_len;
let name_pad = (4 - (name_len % 4)) % 4;
pos += name_pad;
let result = (|| -> Result<(FrameColumnData, usize), DecodeError> {
let mut pos = pos;
let nones = if has_nones && nones_len > 0 {
check_len(data, pos, nones_len)?;
let bv = decode_bitvec(&data[pos..pos + nones_len], row_count);
pos += nones_len;
Some(bv)
} else {
pos += nones_len;
None
};
check_len(data, pos, data_len)?;
let data_bytes = &data[pos..pos + data_len];
pos += data_len;
check_len(data, pos, offsets_len)?;
let offsets_bytes = &data[pos..pos + offsets_len];
pos += offsets_len;
check_len(data, pos, extra_len)?;
let extra_bytes = &data[pos..pos + extra_len];
pos += extra_len;
let col_data = decode_column_dispatch(
type_code,
encoding,
flags,
row_count,
data_bytes,
offsets_bytes,
extra_bytes,
)?;
let col_data = if let Some(bitvec) = nones {
FrameColumnData::Option {
inner: Box::new(col_data),
bitvec,
}
} else {
col_data
};
Ok((col_data, pos))
})()
.map_err(|e| DecodeError::ColumnDecodeFailed {
column_name: name.clone(),
row_index: None,
source: Box::new(e),
})?;
let (col_data, pos) = result;
Ok((
FrameColumn {
name,
data: col_data,
},
pos,
))
}
pub(crate) fn column_type_from_code(type_code: u8) -> Result<ValueType, DecodeError> {
let tag = TypeTag::from_byte(type_code)?;
if tag.depth() != 0 {
return Err(DecodeError::InvalidData(format!(
"column type code 0x{type_code:02X} carries option depth"
)));
}
tag.to_type()
}
fn decode_column_dispatch(
type_code: u8,
encoding: Encoding,
flags: u8,
row_count: usize,
data: &[u8],
offsets: &[u8],
extra: &[u8],
) -> Result<FrameColumnData, DecodeError> {
let ty = column_type_from_code(type_code)?;
match encoding {
Encoding::Plain | Encoding::BitPack => {
if ty == ValueType::Any {
return any::decode_any_column(row_count, data);
}
if let Some(result) = fixed::decode_fixed_plain(type_code, row_count, data) {
return result;
}
if let Some(result) = varlen::decode_varlen_plain(type_code, row_count, data, offsets) {
return result;
}
Err(DecodeError::UnsupportedType(format!("{:?}", ty)))
}
Encoding::Dict => match ty {
ValueType::Utf8 => {
let index_width = dict_index_width_from_flags(flags);
let strings = decode_dict_utf8(data, extra, row_count, index_width)?;
Ok(FrameColumnData::Utf8(Utf8Container::new(strings)))
}
ValueType::Blob => {
let index_width = dict_index_width_from_flags(flags);
let blobs = decode_dict_blob(data, extra, row_count, index_width)?;
Ok(FrameColumnData::Blob(BlobContainer::new(blobs)))
}
ValueType::Int => {
let index_width = dict_index_width_from_flags(flags);
let dict_entries = decode_dict_table_bytes(extra)?;
let mut values = Vec::with_capacity(row_count);
for i in 0..row_count {
let idx = read_index(data, i, index_width) as usize;
if idx >= dict_entries.len() {
return Err(DecodeError::InvalidData(format!(
"dict index {} out of range (dict has {} entries)",
idx,
dict_entries.len()
)));
}
let big = BigInt::from_signed_bytes_le(&dict_entries[idx]);
values.push(Int(big));
}
Ok(FrameColumnData::Int(NumberContainer::new(values)))
}
ValueType::Uint => {
let index_width = dict_index_width_from_flags(flags);
let dict_entries = decode_dict_table_bytes(extra)?;
let mut values = Vec::with_capacity(row_count);
for i in 0..row_count {
let idx = read_index(data, i, index_width) as usize;
if idx >= dict_entries.len() {
return Err(DecodeError::InvalidData(format!(
"dict index {} out of range (dict has {} entries)",
idx,
dict_entries.len()
)));
}
let big = BigInt::from_signed_bytes_le(&dict_entries[idx]);
values.push(Uint(big));
}
Ok(FrameColumnData::Uint(NumberContainer::new(values)))
}
ValueType::Decimal => {
let index_width = dict_index_width_from_flags(flags);
let dict_entries = decode_dict_table_bytes(extra)?;
let mut values = Vec::with_capacity(row_count);
for i in 0..row_count {
let idx = read_index(data, i, index_width) as usize;
if idx >= dict_entries.len() {
return Err(DecodeError::InvalidData(format!(
"dict index {} out of range (dict has {} entries)",
idx,
dict_entries.len()
)));
}
let s = str::from_utf8(&dict_entries[idx]).map_err(|e| {
DecodeError::InvalidData(format!("invalid decimal string: {}", e))
})?;
let dec: BigDecimal = s.parse().map_err(|e| {
DecodeError::InvalidData(format!("invalid decimal: {}", e))
})?;
values.push(Decimal::new(dec));
}
Ok(FrameColumnData::Decimal(NumberContainer::new(values)))
}
_ => Err(DecodeError::InvalidData(format!("Dict encoding not supported for type {:?}", ty))),
},
Encoding::Rle => match ty {
ValueType::Int | ValueType::Uint | ValueType::Decimal => {
varlen::decode_rle_varlen_column(type_code, row_count, data)
}
_ => fixed::decode_rle_column(type_code, row_count, data),
},
Encoding::Delta => fixed::decode_delta_column(type_code, row_count, data),
Encoding::DeltaRle => fixed::decode_delta_rle_column(type_code, row_count, data),
}
}
fn decode_bitvec(data: &[u8], len: usize) -> BitVec {
BitVec::from_raw(data.to_vec(), len)
}
#[inline]
fn read_u16(data: &[u8], pos: usize) -> u16 {
u16::from_le_bytes([data[pos], data[pos + 1]])
}
#[inline]
fn read_u32(data: &[u8], pos: usize) -> u32 {
u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]])
}
#[inline]
fn read_u64(data: &[u8], pos: usize) -> u64 {
u64::from_le_bytes([
data[pos],
data[pos + 1],
data[pos + 2],
data[pos + 3],
data[pos + 4],
data[pos + 5],
data[pos + 6],
data[pos + 7],
])
}
fn check_len(data: &[u8], pos: usize, needed: usize) -> Result<(), DecodeError> {
if pos + needed > data.len() {
Err(DecodeError::UnexpectedEof {
expected: needed,
available: data.len().saturating_sub(pos),
})
} else {
Ok(())
}
}