1use crate::{bflatn_from, indexes::PageOffset, table::RowRef};
7use spacetimedb_sats::layout::{AlgebraicTypeLayout, PrimitiveType, ProductTypeElementLayout, Size, VarLenType};
8use spacetimedb_sats::{
9 algebraic_value::{ser::ValueSerializer, Packed},
10 i256,
11 sum_value::SumTag,
12 u256, AlgebraicType, AlgebraicValue, ArrayValue, ProductType, ProductValue, SumValue, F32, F64,
13};
14use std::{cell::Cell, mem};
15use thiserror::Error;
16
17#[derive(Error, Debug)]
18pub enum TypeError {
19 #[error(
20 "Attempt to read column {} of a product with only {} columns of type {:?}",
21 desired,
22 found.elements.len(),
23 found,
24 )]
25 IndexOutOfBounds { desired: usize, found: ProductType },
26 #[error("Attempt to read a column at type `{desired}`, but the column's type is {found:?}")]
27 WrongType {
28 desired: &'static str,
29 found: AlgebraicType,
30 },
31}
32
33pub unsafe trait ReadColumn: Sized {
43 fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool;
57
58 unsafe fn unchecked_read_column(row_ref: RowRef<'_>, layout: &ProductTypeElementLayout) -> Self;
138
139 fn read_column(row_ref: RowRef<'_>, idx: usize) -> Result<Self, TypeError> {
142 let layout = row_ref.row_layout().product();
143
144 let col = layout.elements.get(idx).ok_or_else(|| TypeError::IndexOutOfBounds {
147 desired: idx,
148 found: layout.product_type(),
149 })?;
150
151 if !Self::is_compatible_type(&col.ty) {
153 return Err(TypeError::WrongType {
154 desired: std::any::type_name::<Self>(),
155 found: col.ty.algebraic_type(),
156 });
157 }
158
159 Ok(unsafe {
160 Self::unchecked_read_column(row_ref, col)
167 })
168 }
169}
170
171unsafe impl ReadColumn for bool {
172 fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool {
173 matches!(ty, AlgebraicTypeLayout::Primitive(PrimitiveType::Bool))
174 }
175
176 unsafe fn unchecked_read_column(row_ref: RowRef<'_>, layout: &ProductTypeElementLayout) -> Self {
177 debug_assert!(Self::is_compatible_type(&layout.ty));
178
179 let (page, offset) = row_ref.page_and_offset();
180 let col_offset = offset + PageOffset(layout.offset);
181
182 let data = page.get_row_data(col_offset, Size(mem::size_of::<Self>() as u16));
183 let data: *const bool = data.as_ptr().cast();
184 unsafe { *data }
190 }
191}
192
193macro_rules! impl_read_column_number {
194 ($primitive_type:ident => $native_type:ty) => {
195 unsafe impl ReadColumn for $native_type {
196 fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool {
197 matches!(ty, AlgebraicTypeLayout::Primitive(PrimitiveType::$primitive_type))
198 }
199
200 unsafe fn unchecked_read_column(
201 row_ref: RowRef<'_>,
202 layout: &ProductTypeElementLayout,
203 ) -> Self {
204 debug_assert!(Self::is_compatible_type(&layout.ty));
205
206 let (page, offset) = row_ref.page_and_offset();
207 let col_offset = offset + PageOffset(layout.offset);
208
209 let data = page.get_row_data(col_offset, Size(mem::size_of::<Self>() as u16));
210 let data: Result<[u8; mem::size_of::<Self>()], _> = data.try_into();
211 let data = unsafe { data.unwrap_unchecked() };
215
216 Self::from_le_bytes(data)
217 }
218 }
219 };
220
221 ($($primitive_type:ident => $native_type:ty);* $(;)*) => {
222 $(impl_read_column_number!($primitive_type => $native_type);)*
223 };
224}
225
226impl_read_column_number! {
227 I8 => i8;
228 U8 => u8;
229 I16 => i16;
230 U16 => u16;
231 I32 => i32;
232 U32 => u32;
233 I64 => i64;
234 U64 => u64;
235 I128 => i128;
236 U128 => u128;
237 I256 => i256;
238 U256 => u256;
239 F32 => f32;
240 F64 => f64;
241}
242
243unsafe impl ReadColumn for AlgebraicValue {
244 fn is_compatible_type(_ty: &AlgebraicTypeLayout) -> bool {
245 true
246 }
247 unsafe fn unchecked_read_column(row_ref: RowRef<'_>, layout: &ProductTypeElementLayout) -> Self {
248 let curr_offset = Cell::new(layout.offset as usize);
249 let blob_store = row_ref.blob_store();
250 let (page, page_offset) = row_ref.page_and_offset();
251 let fixed_bytes = page.get_row_data(page_offset, row_ref.row_layout().size());
252
253 let res = unsafe {
258 bflatn_from::serialize_value(ValueSerializer, fixed_bytes, page, blob_store, &curr_offset, &layout.ty)
259 };
260
261 debug_assert!(res.is_ok());
262
263 unsafe { res.unwrap_unchecked() }
265 }
266}
267
268macro_rules! impl_read_column_via_av {
269 ($av_pattern:pat => $into_method:ident => $native_type:ty) => {
270 unsafe impl ReadColumn for $native_type {
271 fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool {
272 matches!(ty, $av_pattern)
273 }
274
275 unsafe fn unchecked_read_column(
276 row_ref: RowRef<'_>,
277 layout: &ProductTypeElementLayout,
278 ) -> Self {
279 debug_assert!(Self::is_compatible_type(&layout.ty));
280
281 let av = unsafe { AlgebraicValue::unchecked_read_column(row_ref, layout) };
285
286 let res = av.$into_method();
287
288 debug_assert!(res.is_ok());
289
290 unsafe { res.unwrap_unchecked() }
294 }
295 }
296 };
297
298 ($($av_pattern:pat => $into_method:ident => $native_type:ty);* $(;)*) => {
299 $(impl_read_column_via_av!($av_pattern => $into_method => $native_type);)*
300 };
301}
302
303impl_read_column_via_av! {
304 AlgebraicTypeLayout::VarLen(VarLenType::String) => into_string => Box<str>;
305 AlgebraicTypeLayout::VarLen(VarLenType::Array(_)) => into_array => ArrayValue;
306 AlgebraicTypeLayout::Sum(_) => into_sum => SumValue;
307 AlgebraicTypeLayout::Product(_) => into_product => ProductValue;
308}
309
310macro_rules! impl_read_column_via_from {
311 ($($base:ty => $target:ty);* $(;)*) => {
312 $(
313 unsafe impl ReadColumn for $target {
314 fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool {
315 <$base>::is_compatible_type(ty)
316 }
317
318 unsafe fn unchecked_read_column(row_ref: RowRef<'_>, layout: &ProductTypeElementLayout) -> Self {
319 <$target>::from(unsafe { <$base>::unchecked_read_column(row_ref, layout) })
321 }
322 }
323 )*
324 };
325}
326
327impl_read_column_via_from! {
328 u16 => spacetimedb_primitives::ColId;
329 u32 => spacetimedb_primitives::TableId;
330 u32 => spacetimedb_primitives::IndexId;
331 u32 => spacetimedb_primitives::ConstraintId;
332 u32 => spacetimedb_primitives::SequenceId;
333 u32 => spacetimedb_primitives::ScheduleId;
334 u128 => Packed<u128>;
335 i128 => Packed<i128>;
336 u256 => Box<u256>;
337 i256 => Box<i256>;
338 f32 => F32;
339 f64 => F64;
340}
341
342unsafe impl ReadColumn for SumTag {
345 fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool {
346 matches!(ty, AlgebraicTypeLayout::Sum(_))
347 }
348
349 unsafe fn unchecked_read_column(row_ref: RowRef<'_>, layout: &ProductTypeElementLayout) -> Self {
350 debug_assert!(Self::is_compatible_type(&layout.ty));
351
352 let (page, offset) = row_ref.page_and_offset();
353 let col_offset = offset + PageOffset(layout.offset);
354
355 let data = page.get_row_data(col_offset, Size(1));
356 let data: Result<[u8; 1], _> = data.try_into();
357 let [data] = unsafe { data.unwrap_unchecked() };
360
361 Self(data)
362 }
363}
364
365#[cfg(test)]
366mod test {
367 use super::*;
368 use crate::table::test::table;
369 use crate::{blob_store::HashMapBlobStore, page_pool::PagePool};
370 use proptest::{prelude::*, prop_assert_eq, proptest, test_runner::TestCaseResult};
371 use spacetimedb_sats::{product, proptest::generate_typed_row};
372
373 proptest! {
374 #![proptest_config(ProptestConfig::with_cases(if cfg!(miri) { 8 } else { 2048 }))]
375
376 #[test]
377 fn read_column_same_value((ty, val) in generate_typed_row()) {
383 let pool = PagePool::new_for_test();
384 let mut blob_store = HashMapBlobStore::default();
385 let mut table = table(ty);
386
387 let (_, row_ref) = table.insert(&pool, &mut blob_store, &val).unwrap();
388
389 for (idx, orig_col_value) in val.into_iter().enumerate() {
390 let read_col_value = row_ref.read_col::<AlgebraicValue>(idx).unwrap();
391 prop_assert_eq!(orig_col_value, read_col_value);
392 }
393 }
394
395 #[test]
396 fn read_column_wrong_type((ty, val) in generate_typed_row()) {
400 let pool = PagePool::new_for_test();
401 let mut blob_store = HashMapBlobStore::default();
402 let mut table = table(ty.clone());
403
404 let (_, row_ref) = table.insert(&pool, &mut blob_store, &val).unwrap();
405
406 for (idx, col_ty) in ty.elements.iter().enumerate() {
407 assert_wrong_type_error::<u8>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::U8)?;
408 assert_wrong_type_error::<i8>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::I8)?;
409 assert_wrong_type_error::<u16>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::U16)?;
410 assert_wrong_type_error::<i16>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::I16)?;
411 assert_wrong_type_error::<u32>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::U32)?;
412 assert_wrong_type_error::<i32>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::I32)?;
413 assert_wrong_type_error::<u64>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::U64)?;
414 assert_wrong_type_error::<i64>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::I64)?;
415 assert_wrong_type_error::<u128>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::U128)?;
416 assert_wrong_type_error::<i128>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::I128)?;
417 assert_wrong_type_error::<u256>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::U256)?;
418 assert_wrong_type_error::<i256>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::I256)?;
419 assert_wrong_type_error::<f32>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::F32)?;
420 assert_wrong_type_error::<f64>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::F64)?;
421 assert_wrong_type_error::<bool>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::Bool)?;
422 assert_wrong_type_error::<Box<str>>(row_ref, idx, &col_ty.algebraic_type, AlgebraicType::String)?;
423 }
424 }
425
426 #[test]
427 fn read_column_out_of_bounds((ty, val) in generate_typed_row()) {
431 let pool = PagePool::new_for_test();
432 let mut blob_store = HashMapBlobStore::default();
433 let mut table = table(ty.clone());
434
435 let (_, row_ref) = table.insert(&pool, &mut blob_store, &val).unwrap();
436
437 let oob = ty.elements.len();
438
439 match row_ref.read_col::<AlgebraicValue>(oob) {
440 Err(TypeError::IndexOutOfBounds { desired, found }) => {
441 prop_assert_eq!(desired, oob);
442 prop_assert_eq!(found.elements.len(), ty.elements.len());
448 for (found_col, ty_col) in found.elements.iter().zip(ty.elements.iter()) {
449 prop_assert_eq!(&found_col.algebraic_type, &ty_col.algebraic_type);
450 }
451 }
452 Err(e) => panic!("Expected TypeError::IndexOutOfBounds but found {e:?}"),
453 Ok(val) => panic!("Expected error but found Ok({val:?})"),
454 }
455 }
456 }
457
458 fn assert_wrong_type_error<Col: ReadColumn + PartialEq + std::fmt::Debug>(
463 row_ref: RowRef<'_>,
464 col_idx: usize,
465 col_ty: &AlgebraicType,
466 correct_col_ty: AlgebraicType,
467 ) -> TestCaseResult {
468 if col_ty != &correct_col_ty {
469 match row_ref.read_col::<Col>(col_idx) {
470 Err(TypeError::WrongType { desired, found }) => {
471 prop_assert_eq!(desired, std::any::type_name::<Col>());
472 prop_assert_eq!(&found, col_ty);
473 }
474 Err(e) => panic!("Expected TypeError::WrongType but found {e:?}"),
475 Ok(val) => panic!("Expected error but found Ok({val:?})"),
476 }
477 }
478 Ok(())
479 }
480
481 macro_rules! test_read_column_primitive {
486 ($name:ident { $algebraic_type:expr => $rust_type:ty = $val:expr }) => {
487 #[test]
488 fn $name() {
489 let pool = PagePool::new_for_test();
490 let mut blob_store = HashMapBlobStore::default();
491 let mut table = table(ProductType::from_iter([$algebraic_type]));
492
493 let val: $rust_type = $val;
494 let (_, row_ref) = table.insert(&pool, &mut blob_store, &product![val.clone()]).unwrap();
495
496 assert_eq!(val, row_ref.read_col::<$rust_type>(0).unwrap());
497 }
498 };
499
500
501 ($($name:ident { $algebraic_type:expr => $rust_type:ty = $val:expr };)*) => {
502 $(test_read_column_primitive! {
503 $name { $algebraic_type => $rust_type = $val }
504 })*
505 }
506 }
507
508 test_read_column_primitive! {
509 read_column_i8 { AlgebraicType::I8 => i8 = i8::MAX };
510 read_column_u8 { AlgebraicType::U8 => u8 = 0xa5 };
511 read_column_i16 { AlgebraicType::I16 => i16 = i16::MAX };
512 read_column_u16 { AlgebraicType::U16 => u16 = 0xa5a5 };
513 read_column_i32 { AlgebraicType::I32 => i32 = i32::MAX };
514 read_column_u32 { AlgebraicType::U32 => u32 = 0xa5a5a5a5 };
515 read_column_i64 { AlgebraicType::I64 => i64 = i64::MAX };
516 read_column_u64 { AlgebraicType::U64 => u64 = 0xa5a5a5a5_a5a5a5a5 };
517 read_column_i128 { AlgebraicType::I128 => i128 = i128::MAX };
518 read_column_u128 { AlgebraicType::U128 => u128 = 0xa5a5a5a5_a5a5a5a5_a5a5a5a5_a5a5a5a5 };
519 read_column_i256 { AlgebraicType::I256 => i256 = i256::MAX };
520 read_column_u256 { AlgebraicType::U256 => u256 =
521 u256::from_words(
522 0xa5a5a5a5_a5a5a5a5_a5a5a5a5_a5a5a5a5,
523 0xa5a5a5a5_a5a5a5a5_a5a5a5a5_a5a5a5a5
524 )
525 };
526
527 read_column_f32 { AlgebraicType::F32 => f32 = 1.0 };
528 read_column_f64 { AlgebraicType::F64 => f64 = 1.0 };
529
530 read_column_bool { AlgebraicType::Bool => bool = true };
531
532 read_column_empty_string { AlgebraicType::String => Box<str> = "".into() };
533
534 read_column_short_string { AlgebraicType::String => Box<str> = "short string".into() };
536
537 read_column_medium_string { AlgebraicType::String => Box<str> = "medium string.".repeat(16).into() };
539
540 read_column_long_string { AlgebraicType::String => Box<str> = "long string. ".repeat(2048).into() };
542
543 read_sum_value_plain { AlgebraicType::simple_enum(["a", "b"].into_iter()) => SumValue = SumValue::new_simple(1) };
544 read_sum_tag_plain { AlgebraicType::simple_enum(["a", "b"].into_iter()) => SumTag = SumTag(1) };
545 }
546
547 #[test]
548 fn read_sum_tag_from_sum_with_payload() {
549 let algebraic_type = AlgebraicType::sum([("a", AlgebraicType::U8), ("b", AlgebraicType::U16)]);
550
551 let pool = PagePool::new_for_test();
552 let mut blob_store = HashMapBlobStore::default();
553 let mut table = table(ProductType::from([algebraic_type]));
554
555 let val = SumValue::new(1, 42u16);
556 let (_, row_ref) = table.insert(&pool, &mut blob_store, &product![val.clone()]).unwrap();
557
558 assert_eq!(val.tag, row_ref.read_col::<SumTag>(0).unwrap().0);
559 }
560}