use std::sync::Arc;
use vibesql_storage::Row;
use vibesql_types::SqlValue;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct RowReference {
table_id: u16,
row_index: u32,
}
impl RowReference {
#[inline]
pub const fn new(table_id: u16, row_index: u32) -> Self {
Self { table_id, row_index }
}
#[inline]
pub const fn table_id(&self) -> u16 {
self.table_id
}
#[inline]
pub const fn row_index(&self) -> u32 {
self.row_index
}
#[inline]
pub fn range(table_id: u16, start: u32, end: u32) -> Vec<Self> {
(start..end).map(|idx| Self::new(table_id, idx)).collect()
}
#[inline]
pub fn from_selection(table_id: u16, indices: &[u32]) -> Vec<Self> {
indices.iter().map(|&idx| Self::new(table_id, idx)).collect()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct JoinedRowRef {
pub left: RowReference,
pub right: Option<RowReference>,
}
impl JoinedRowRef {
#[inline]
pub const fn matched(left: RowReference, right: RowReference) -> Self {
Self { left, right: Some(right) }
}
#[inline]
pub const fn left_only(left: RowReference) -> Self {
Self { left, right: None }
}
#[inline]
pub const fn is_matched(&self) -> bool {
self.right.is_some()
}
}
pub trait RowResolver {
fn resolve(&self, reference: &RowReference) -> Option<&Row>;
fn resolve_column(&self, reference: &RowReference, column_idx: usize) -> Option<&SqlValue>;
fn resolve_batch(&self, references: &[RowReference]) -> Vec<Option<&Row>> {
references.iter().map(|r| self.resolve(r)).collect()
}
}
pub struct VecRowResolver<'a> {
table_id: u16,
rows: &'a [Row],
}
impl<'a> VecRowResolver<'a> {
pub fn new(table_id: u16, rows: &'a [Row]) -> Self {
Self { table_id, rows }
}
}
impl<'a> RowResolver for VecRowResolver<'a> {
fn resolve(&self, reference: &RowReference) -> Option<&Row> {
if reference.table_id == self.table_id {
self.rows.get(reference.row_index as usize)
} else {
None
}
}
fn resolve_column(&self, reference: &RowReference, column_idx: usize) -> Option<&SqlValue> {
self.resolve(reference).and_then(|row| row.get(column_idx))
}
}
pub struct MultiTableResolver<'a> {
tables: Vec<(u16, &'a [Row])>,
}
impl<'a> MultiTableResolver<'a> {
pub fn new() -> Self {
Self { tables: Vec::new() }
}
pub fn add_table(&mut self, table_id: u16, rows: &'a [Row]) {
self.tables.push((table_id, rows));
}
pub fn from_tables(tables: Vec<(u16, &'a [Row])>) -> Self {
Self { tables }
}
}
impl<'a> Default for MultiTableResolver<'a> {
fn default() -> Self {
Self::new()
}
}
impl<'a> RowResolver for MultiTableResolver<'a> {
fn resolve(&self, reference: &RowReference) -> Option<&Row> {
for (table_id, rows) in &self.tables {
if *table_id == reference.table_id {
return rows.get(reference.row_index as usize);
}
}
None
}
fn resolve_column(&self, reference: &RowReference, column_idx: usize) -> Option<&SqlValue> {
self.resolve(reference).and_then(|row| row.get(column_idx))
}
}
pub struct OwnedRowSource {
table_id: u16,
rows: Arc<Vec<Row>>,
}
impl OwnedRowSource {
pub fn new(table_id: u16, rows: Vec<Row>) -> Self {
Self { table_id, rows: Arc::new(rows) }
}
#[inline]
pub fn table_id(&self) -> u16 {
self.table_id
}
#[inline]
pub fn row_count(&self) -> usize {
self.rows.len()
}
#[inline]
pub fn rows(&self) -> &[Row] {
&self.rows
}
#[inline]
pub fn reference(&self, row_index: u32) -> RowReference {
RowReference::new(self.table_id, row_index)
}
#[inline]
pub fn resolve(&self, reference: &RowReference) -> Option<&Row> {
if reference.table_id == self.table_id {
self.rows.get(reference.row_index as usize)
} else {
None
}
}
pub fn share(&self) -> Self {
Self { table_id: self.table_id, rows: Arc::clone(&self.rows) }
}
}
impl Clone for OwnedRowSource {
fn clone(&self) -> Self {
self.share()
}
}
#[cfg(test)]
mod row_ref_tests {
use super::*;
#[test]
fn test_row_reference_creation() {
let r = RowReference::new(1, 42);
assert_eq!(r.table_id(), 1);
assert_eq!(r.row_index(), 42);
}
#[test]
fn test_row_reference_range() {
let refs = RowReference::range(0, 10, 15);
assert_eq!(refs.len(), 5);
assert_eq!(refs[0].row_index(), 10);
assert_eq!(refs[4].row_index(), 14);
}
#[test]
fn test_vec_resolver() {
let rows = vec![
Row::new(vec![SqlValue::Integer(1)]),
Row::new(vec![SqlValue::Integer(2)]),
Row::new(vec![SqlValue::Integer(3)]),
];
let resolver = VecRowResolver::new(0, &rows);
let ref1 = RowReference::new(0, 1);
let resolved = resolver.resolve(&ref1).unwrap();
assert_eq!(resolved.get(0), Some(&SqlValue::Integer(2)));
let ref_wrong = RowReference::new(1, 1);
assert!(resolver.resolve(&ref_wrong).is_none());
}
#[test]
fn test_multi_table_resolver() {
let table0 = [Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A"))])];
let table1 = [Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B"))])];
let resolver = MultiTableResolver::from_tables(vec![(0, &table0[..]), (1, &table1[..])]);
let ref0 = RowReference::new(0, 0);
let ref1 = RowReference::new(1, 0);
assert_eq!(
resolver.resolve_column(&ref0, 0),
Some(&SqlValue::Varchar(arcstr::ArcStr::from("A")))
);
assert_eq!(
resolver.resolve_column(&ref1, 0),
Some(&SqlValue::Varchar(arcstr::ArcStr::from("B")))
);
}
#[test]
fn test_joined_row_ref() {
let left = RowReference::new(0, 10);
let right = RowReference::new(1, 20);
let matched = JoinedRowRef::matched(left, right);
assert!(matched.is_matched());
assert_eq!(matched.left.row_index(), 10);
assert_eq!(matched.right.unwrap().row_index(), 20);
let left_only = JoinedRowRef::left_only(left);
assert!(!left_only.is_matched());
}
#[test]
fn test_owned_row_source() {
let rows =
vec![Row::new(vec![SqlValue::Integer(100)]), Row::new(vec![SqlValue::Integer(200)])];
let source = OwnedRowSource::new(0, rows);
assert_eq!(source.row_count(), 2);
let r = source.reference(1);
assert_eq!(source.resolve(&r).unwrap().get(0), Some(&SqlValue::Integer(200)));
let shared = source.share();
assert_eq!(shared.row_count(), 2);
}
}