qroissant/crates/qroissant-core/src/value.rs

use bytes::Bytes;

use crate::error::CoreError;
use crate::error::CoreResult;
use crate::protocol::Attribute;
use crate::protocol::Primitive;
use crate::protocol::ValueType;

/// q atom payload.
#[derive(Clone, Debug, PartialEq)]
pub enum Atom {
    Boolean(bool),
    Guid([u8; 16]),
    Byte(u8),
    Short(i16),
    Int(i32),
    Long(i64),
    Real(f32),
    Float(f64),
    Char(u8),
    Symbol(Bytes),
    Timestamp(i64),
    Month(i32),
    Date(i32),
    Datetime(f64),
    Timespan(i64),
    Minute(i32),
    Second(i32),
    Time(i32),
}

impl Atom {
    pub fn primitive(&self) -> Primitive {
        match self {
            Self::Boolean(_) => Primitive::Boolean,
            Self::Guid(_) => Primitive::Guid,
            Self::Byte(_) => Primitive::Byte,
            Self::Short(_) => Primitive::Short,
            Self::Int(_) => Primitive::Int,
            Self::Long(_) => Primitive::Long,
            Self::Real(_) => Primitive::Real,
            Self::Float(_) => Primitive::Float,
            Self::Char(_) => Primitive::Char,
            Self::Symbol(_) => Primitive::Symbol,
            Self::Timestamp(_) => Primitive::Timestamp,
            Self::Month(_) => Primitive::Month,
            Self::Date(_) => Primitive::Date,
            Self::Datetime(_) => Primitive::Datetime,
            Self::Timespan(_) => Primitive::Timespan,
            Self::Minute(_) => Primitive::Minute,
            Self::Second(_) => Primitive::Second,
            Self::Time(_) => Primitive::Time,
        }
    }
}

/// q homogeneous vector payload.
///
/// All fixed-width numeric types store their data as raw [`Bytes`], enabling
/// zero-copy slicing from the IPC frame buffer during decode. Typed access
/// is provided via `as_*_slice()` methods using `bytemuck::cast_slice`.
#[derive(Clone, Debug, PartialEq)]
pub enum VectorData {
    Boolean(Bytes),
    Guid(Bytes),
    Byte(Bytes),
    Short(Bytes),
    Int(Bytes),
    Long(Bytes),
    Real(Bytes),
    Float(Bytes),
    Char(Bytes),
    Symbol(Vec<Bytes>),
    Timestamp(Bytes),
    Month(Bytes),
    Date(Bytes),
    Datetime(Bytes),
    Timespan(Bytes),
    Minute(Bytes),
    Second(Bytes),
    Time(Bytes),
}

impl VectorData {
    pub fn primitive(&self) -> Primitive {
        match self {
            Self::Boolean(_) => Primitive::Boolean,
            Self::Guid(_) => Primitive::Guid,
            Self::Byte(_) => Primitive::Byte,
            Self::Short(_) => Primitive::Short,
            Self::Int(_) => Primitive::Int,
            Self::Long(_) => Primitive::Long,
            Self::Real(_) => Primitive::Real,
            Self::Float(_) => Primitive::Float,
            Self::Char(_) => Primitive::Char,
            Self::Symbol(_) => Primitive::Symbol,
            Self::Timestamp(_) => Primitive::Timestamp,
            Self::Month(_) => Primitive::Month,
            Self::Date(_) => Primitive::Date,
            Self::Datetime(_) => Primitive::Datetime,
            Self::Timespan(_) => Primitive::Timespan,
            Self::Minute(_) => Primitive::Minute,
            Self::Second(_) => Primitive::Second,
            Self::Time(_) => Primitive::Time,
        }
    }

    pub fn len(&self) -> usize {
        match self {
            Self::Boolean(b) | Self::Byte(b) | Self::Char(b) => b.len(),
            Self::Guid(b) => b.len() / 16,
            Self::Short(b) => b.len() / 2,
            Self::Int(b)
            | Self::Month(b)
            | Self::Date(b)
            | Self::Minute(b)
            | Self::Second(b)
            | Self::Time(b)
            | Self::Real(b) => b.len() / 4,
            Self::Long(b)
            | Self::Timestamp(b)
            | Self::Timespan(b)
            | Self::Float(b)
            | Self::Datetime(b) => b.len() / 8,
            Self::Symbol(v) => v.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the underlying raw bytes for non-Symbol variants.
    pub fn raw_bytes(&self) -> Option<&Bytes> {
        match self {
            Self::Symbol(_) => None,
            Self::Boolean(b)
            | Self::Guid(b)
            | Self::Byte(b)
            | Self::Short(b)
            | Self::Int(b)
            | Self::Long(b)
            | Self::Real(b)
            | Self::Float(b)
            | Self::Char(b)
            | Self::Timestamp(b)
            | Self::Month(b)
            | Self::Date(b)
            | Self::Datetime(b)
            | Self::Timespan(b)
            | Self::Minute(b)
            | Self::Second(b)
            | Self::Time(b) => Some(b),
        }
    }

    pub fn as_i16_slice(&self) -> &[i16] {
        match self {
            Self::Short(b) => bytemuck::cast_slice(b),
            _ => panic!("as_i16_slice called on {:?}", self.primitive()),
        }
    }

    pub fn as_i32_slice(&self) -> &[i32] {
        match self {
            Self::Int(b)
            | Self::Month(b)
            | Self::Date(b)
            | Self::Minute(b)
            | Self::Second(b)
            | Self::Time(b) => bytemuck::cast_slice(b),
            _ => panic!("as_i32_slice called on {:?}", self.primitive()),
        }
    }

    pub fn as_i64_slice(&self) -> &[i64] {
        match self {
            Self::Long(b) | Self::Timestamp(b) | Self::Timespan(b) => bytemuck::cast_slice(b),
            _ => panic!("as_i64_slice called on {:?}", self.primitive()),
        }
    }

    pub fn as_f32_slice(&self) -> &[f32] {
        match self {
            Self::Real(b) => bytemuck::cast_slice(b),
            _ => panic!("as_f32_slice called on {:?}", self.primitive()),
        }
    }

    pub fn as_f64_slice(&self) -> &[f64] {
        match self {
            Self::Float(b) | Self::Datetime(b) => bytemuck::cast_slice(b),
            _ => panic!("as_f64_slice called on {:?}", self.primitive()),
        }
    }

    // Construction helpers for tests and ingestion paths.

    pub fn from_i16s(values: &[i16]) -> Self {
        Self::Short(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_i32s(values: &[i32]) -> Self {
        Self::Int(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_i64s(values: &[i64]) -> Self {
        Self::Long(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_f32s(values: &[f32]) -> Self {
        Self::Real(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_f64s(values: &[f64]) -> Self {
        Self::Float(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_guids(values: &[[u8; 16]]) -> Self {
        let mut buf = Vec::with_capacity(values.len() * 16);
        for guid in values {
            buf.extend_from_slice(guid);
        }
        Self::Guid(Bytes::from(buf))
    }

    pub fn from_timestamps(values: &[i64]) -> Self {
        Self::Timestamp(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_months(values: &[i32]) -> Self {
        Self::Month(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_dates(values: &[i32]) -> Self {
        Self::Date(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_datetimes(values: &[f64]) -> Self {
        Self::Datetime(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_timespans(values: &[i64]) -> Self {
        Self::Timespan(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_minutes(values: &[i32]) -> Self {
        Self::Minute(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_seconds(values: &[i32]) -> Self {
        Self::Second(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }

    pub fn from_times(values: &[i32]) -> Self {
        Self::Time(Bytes::copy_from_slice(bytemuck::cast_slice(values)))
    }
}

/// q homogeneous vector with an attached q attribute.
#[derive(Clone, Debug, PartialEq)]
pub struct Vector {
    attribute: Attribute,
    data: VectorData,
}

impl Vector {
    pub fn new(attribute: Attribute, data: VectorData) -> Self {
        Self { attribute, data }
    }

    pub fn attribute(&self) -> Attribute {
        self.attribute
    }

    pub fn primitive(&self) -> Primitive {
        self.data.primitive()
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    pub fn data(&self) -> &VectorData {
        &self.data
    }
}

/// q general list.
#[derive(Clone, Debug, PartialEq)]
pub struct List {
    attribute: Attribute,
    values: Vec<Value>,
}

impl List {
    pub fn new(attribute: Attribute, values: Vec<Value>) -> Self {
        Self { attribute, values }
    }

    pub fn attribute(&self) -> Attribute {
        self.attribute
    }

    pub fn len(&self) -> usize {
        self.values.len()
    }

    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }

    pub fn values(&self) -> &[Value] {
        &self.values
    }
}

/// q dictionary.
#[derive(Clone, Debug, PartialEq)]
pub struct Dictionary {
    sorted: bool,
    keys: Box<Value>,
    values: Box<Value>,
}

impl Dictionary {
    pub fn new(sorted: bool, keys: Value, values: Value) -> Self {
        Self {
            sorted,
            keys: Box::new(keys),
            values: Box::new(values),
        }
    }

    pub fn sorted(&self) -> bool {
        self.sorted
    }

    pub fn keys(&self) -> &Value {
        &self.keys
    }

    pub fn values(&self) -> &Value {
        &self.values
    }

    pub fn len(&self) -> usize {
        self.keys.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn validate(&self) -> CoreResult<()> {
        if self.keys.len() != self.values.len() {
            return Err(CoreError::InvalidStructure(format!(
                "q dictionary key/value lengths differ: {} != {}",
                self.keys.len(),
                self.values.len()
            )));
        }

        Ok(())
    }
}

/// q table.
#[derive(Clone, Debug, PartialEq)]
pub struct Table {
    attribute: Attribute,
    column_names: Vec<Bytes>,
    columns: Vec<Value>,
}

impl Table {
    pub fn new(attribute: Attribute, column_names: Vec<Bytes>, columns: Vec<Value>) -> Self {
        Self {
            attribute,
            column_names,
            columns,
        }
    }

    pub fn attribute(&self) -> Attribute {
        self.attribute
    }

    pub fn column_names(&self) -> &[Bytes] {
        &self.column_names
    }

    pub fn columns(&self) -> &[Value] {
        &self.columns
    }

    pub fn num_columns(&self) -> usize {
        self.columns.len()
    }

    pub fn len(&self) -> usize {
        self.columns.first().map_or(0, Value::len)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn validate(&self) -> CoreResult<()> {
        if self.column_names.len() != self.columns.len() {
            return Err(CoreError::InvalidStructure(format!(
                "q table column name count {} does not match column count {}",
                self.column_names.len(),
                self.columns.len()
            )));
        }

        if let Some(expected_rows) = self.columns.first().map(Value::len) {
            for column in self.columns.iter().skip(1) {
                if column.len() != expected_rows {
                    return Err(CoreError::InvalidStructure(format!(
                        "q table column lengths differ: expected {expected_rows}, found {}",
                        column.len()
                    )));
                }
            }
        }

        Ok(())
    }
}

/// Decoded q value subset currently supported by the rewrite.
#[derive(Clone, Debug, PartialEq)]
pub enum Value {
    Atom(Atom),
    Vector(Vector),
    List(List),
    Dictionary(Dictionary),
    Table(Table),
    UnaryPrimitive { opcode: i8 },
}

impl Value {
    pub fn qtype(&self) -> ValueType {
        match self {
            Self::Atom(atom) => ValueType::atom(atom.primitive()),
            Self::Vector(vector) => ValueType::vector(vector.primitive(), vector.attribute()),
            Self::List(list) => ValueType::list(list.attribute()),
            Self::Dictionary(dictionary) => ValueType::dictionary(dictionary.sorted()),
            Self::Table(table) => ValueType::table(table.attribute()),
            Self::UnaryPrimitive { .. } => ValueType::unary_primitive(),
        }
    }

    pub fn len(&self) -> usize {
        match self {
            Self::Atom(_) | Self::UnaryPrimitive { .. } => 1,
            Self::Vector(vector) => vector.len(),
            Self::List(list) => list.len(),
            Self::Dictionary(dictionary) => dictionary.len(),
            Self::Table(table) => table.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        match self {
            Self::Atom(_) | Self::UnaryPrimitive { .. } => false,
            Self::Vector(vector) => vector.is_empty(),
            Self::List(list) => list.is_empty(),
            Self::Dictionary(dictionary) => dictionary.is_empty(),
            Self::Table(table) => table.is_empty(),
        }
    }
}