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