use super::{
blob_store::BlobHash,
indexes::{Byte, Bytes, PageOffset},
};
use crate::{static_assert_align, static_assert_size};
use core::iter;
use core::marker::PhantomData;
use core::mem::{self};
use spacetimedb_sats::layout::{Size, VAR_LEN_REF_LAYOUT};
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
#[repr(C)]
pub struct VarLenRef {
pub length_in_bytes: u16,
pub first_granule: PageOffset,
}
static_assert_size!(VarLenRef, 4);
static_assert_align!(VarLenRef, 2);
const _: () = assert!(VAR_LEN_REF_LAYOUT.size as usize == mem::size_of::<VarLenRef>());
const _: () = assert!(VAR_LEN_REF_LAYOUT.align as usize == mem::align_of::<VarLenRef>());
impl VarLenRef {
#[inline]
pub const fn is_large_blob(self) -> bool {
self.length_in_bytes == Self::LARGE_BLOB_SENTINEL
}
pub const LARGE_BLOB_SENTINEL: u16 = u16::MAX;
#[inline]
pub const fn large_blob(first_granule: PageOffset) -> Self {
Self {
length_in_bytes: Self::LARGE_BLOB_SENTINEL,
first_granule,
}
}
#[inline]
pub const fn granules_used(&self) -> usize {
VarLenGranule::bytes_to_granules(self.length_in_bytes as usize).0
}
#[inline]
pub const fn is_null(self) -> bool {
self.first_granule.is_var_len_null()
}
pub const NULL: Self = Self {
length_in_bytes: 0,
first_granule: PageOffset::VAR_LEN_NULL,
};
}
const _BLOB_SENTINEL_MORE_THAN_MAX_OBJ_SIZE: () =
assert!(VarLenGranule::OBJECT_SIZE_BLOB_THRESHOLD < VarLenRef::LARGE_BLOB_SENTINEL as usize);
const _GRANULES_USED_FOR_BLOB_IS_CONSISTENT: () = {
let vlr = VarLenRef::large_blob(PageOffset::VAR_LEN_NULL);
assert!(vlr.is_large_blob() == (vlr.granules_used() == 1));
};
pub fn is_granule_offset_aligned(offset: PageOffset) -> bool {
offset.0 == offset.0 & VarLenGranuleHeader::NEXT_BITMASK
}
#[derive(Copy, Clone)]
pub struct VarLenGranuleHeader(u16);
impl VarLenGranuleHeader {
const SIZE: usize = mem::size_of::<Self>();
const LEN_BITS: u16 = 6;
const LEN_BITMASK: u16 = (1 << Self::LEN_BITS) - 1;
#[allow(clippy::assertions_on_constants)]
const _ASSERT_LEN_BITMASK_FITS_ALL_POSSIBLE_GRANULE_LENGTHS: () =
assert!(VarLenGranule::DATA_SIZE <= Self::LEN_BITMASK as usize);
const NEXT_BITMASK: u16 = !Self::LEN_BITMASK;
fn with_len(self, len: u8) -> Self {
let mut new = self;
new.0 &= !Self::LEN_BITMASK;
let capped_len = (len as u16) & Self::LEN_BITMASK;
debug_assert_eq!(
capped_len, len as u16,
"Len {len} overflows the length of a `VarLenGranule`",
);
new.0 |= capped_len;
debug_assert_eq!(self.next(), new.next(), "`set_len` has modified `next`");
debug_assert_eq!(
new.len() as u16,
capped_len,
"`set_len` has not inserted the correct `len`: expected {:x}, found {:x}",
capped_len,
new.len()
);
new
}
fn with_next(self, PageOffset(next): PageOffset) -> Self {
let mut new = self;
new.0 &= !Self::NEXT_BITMASK;
let aligned_next = next & Self::NEXT_BITMASK;
debug_assert_eq!(aligned_next, next, "Next {next:x} is unaligned");
new.0 |= aligned_next;
debug_assert_eq!(self.len(), new.len(), "`set_next` has modified `len`");
debug_assert_eq!(
new.next().0,
aligned_next,
"`set_next` has not inserted the correct `next`"
);
new
}
pub fn new(len: u8, next: PageOffset) -> Self {
Self(0).with_len(len).with_next(next)
}
const fn len(&self) -> u8 {
(self.0 & Self::LEN_BITMASK) as u8
}
pub const fn next(&self) -> PageOffset {
PageOffset(self.0 & Self::NEXT_BITMASK)
}
}
#[repr(C)] #[repr(align(64))] pub struct VarLenGranule {
pub header: VarLenGranuleHeader,
pub data: [Byte; Self::DATA_SIZE],
}
impl VarLenGranule {
pub const SIZE: Size = Size(64);
pub const DATA_SIZE: usize = Self::SIZE.len() - VarLenGranuleHeader::SIZE;
pub const OBJECT_MAX_GRANULES_BEFORE_BLOB: usize = 16;
pub const OBJECT_SIZE_BLOB_THRESHOLD: usize = Self::DATA_SIZE * Self::OBJECT_MAX_GRANULES_BEFORE_BLOB;
pub const fn space_to_granules(available_len: Size) -> usize {
available_len.len() / Self::SIZE.len()
}
pub const fn bytes_to_granules(len_in_bytes: usize) -> (usize, bool) {
if len_in_bytes > VarLenGranule::OBJECT_SIZE_BLOB_THRESHOLD {
(1, true)
} else {
(len_in_bytes.div_ceil(Self::DATA_SIZE), false)
}
}
pub fn chunks(bytes: &[u8]) -> impl DoubleEndedIterator<Item = &[u8]> {
bytes.chunks(Self::DATA_SIZE)
}
pub fn data(&self) -> &[u8] {
let len = self.header.len() as usize;
&self.data[0..len]
}
pub fn blob_hash(&self) -> BlobHash {
self.data().try_into().unwrap()
}
}
#[allow(clippy::assertions_on_constants)]
const _VLG_CAN_STORE_BLOB_HASH: () = assert!(VarLenGranule::DATA_SIZE >= BlobHash::SIZE);
pub unsafe trait VarLenMembers {
type Iter<'this, 'row>: Iterator<Item = &'row VarLenRef>
where
Self: 'this;
type IterMut<'this, 'row>: Iterator<Item = &'row mut VarLenRef>
where
Self: 'this;
unsafe fn visit_var_len_mut<'this, 'row>(&'this self, row: &'row mut Bytes) -> Self::IterMut<'this, 'row>;
unsafe fn visit_var_len<'this, 'row>(&'this self, row: &'row Bytes) -> Self::Iter<'this, 'row>;
}
#[derive(Copy, Clone)]
pub struct AlignedVarLenOffsets<'a>(&'a [u16]);
impl<'a> AlignedVarLenOffsets<'a> {
pub const fn from_offsets(offsets: &'a [u16]) -> Self {
Self(offsets)
}
}
unsafe impl VarLenMembers for AlignedVarLenOffsets<'_> {
type Iter<'this, 'row>
= AlignedVarLenOffsetsIter<'this, 'row>
where
Self: 'this;
type IterMut<'this, 'row>
= AlignedVarLenOffsetsIterMut<'this, 'row>
where
Self: 'this;
unsafe fn visit_var_len<'this, 'row>(&'this self, row: &'row Bytes) -> Self::Iter<'this, 'row> {
AlignedVarLenOffsetsIter {
offsets: self,
_row_lifetime: PhantomData,
row: row.as_ptr(),
next_offset_idx: 0,
}
}
unsafe fn visit_var_len_mut<'this, 'row>(&'this self, row: &'row mut Bytes) -> Self::IterMut<'this, 'row> {
AlignedVarLenOffsetsIterMut {
offsets: self,
_row_lifetime: PhantomData,
row: row.as_mut_ptr(),
next_offset_idx: 0,
}
}
}
pub struct AlignedVarLenOffsetsIter<'offsets, 'row> {
offsets: &'offsets AlignedVarLenOffsets<'offsets>,
_row_lifetime: PhantomData<&'row Bytes>,
row: *const Byte,
next_offset_idx: usize,
}
impl<'row> Iterator for AlignedVarLenOffsetsIter<'_, 'row> {
type Item = &'row VarLenRef;
fn next(&mut self) -> Option<Self::Item> {
if self.next_offset_idx >= self.offsets.0.len() {
None
} else {
let curr_offset_idx = self.next_offset_idx;
self.next_offset_idx += 1;
let elt_ptr: *const VarLenRef =
unsafe { self.row.add(curr_offset_idx * mem::align_of::<VarLenRef>()).cast() };
Some(unsafe { &*elt_ptr })
}
}
}
pub struct AlignedVarLenOffsetsIterMut<'offsets, 'row> {
offsets: &'offsets AlignedVarLenOffsets<'offsets>,
_row_lifetime: PhantomData<&'row mut Bytes>,
row: *mut Byte,
next_offset_idx: usize,
}
impl<'row> Iterator for AlignedVarLenOffsetsIterMut<'_, 'row> {
type Item = &'row mut VarLenRef;
fn next(&mut self) -> Option<Self::Item> {
if self.next_offset_idx >= self.offsets.0.len() {
None
} else {
let curr_offset_idx = self.next_offset_idx;
self.next_offset_idx += 1;
let elt_ptr: *mut VarLenRef =
unsafe { self.row.add(curr_offset_idx * mem::align_of::<VarLenRef>()).cast() };
Some(unsafe { &mut *elt_ptr })
}
}
}
#[derive(Copy, Clone)]
pub struct NullVarLenVisitor;
unsafe impl VarLenMembers for NullVarLenVisitor {
type Iter<'this, 'row> = iter::Empty<&'row VarLenRef>;
type IterMut<'this, 'row> = iter::Empty<&'row mut VarLenRef>;
unsafe fn visit_var_len<'this, 'row>(&'this self, _row: &'row Bytes) -> Self::Iter<'this, 'row> {
iter::empty()
}
unsafe fn visit_var_len_mut<'this, 'row>(&'this self, _row: &'row mut Bytes) -> Self::IterMut<'this, 'row> {
iter::empty()
}
}
#[cfg(test)]
mod test {
use super::*;
use proptest::prelude::*;
fn generate_var_len_offset() -> impl Strategy<Value = PageOffset> {
(0u16..(1 << 10)).prop_map(|unaligned| PageOffset(unaligned * VarLenGranule::SIZE.0))
}
fn generate_len() -> impl Strategy<Value = u8> {
0..(VarLenGranule::DATA_SIZE as u8)
}
proptest! {
#[test]
fn granule_header_bitbashing(len in generate_len(), next in generate_var_len_offset(), len2 in generate_len(), next2 in generate_var_len_offset()) {
let header = VarLenGranuleHeader::new(len, next);
prop_assert_eq!(len, header.len());
prop_assert_eq!(next, header.next());
let header_new_len = header.with_len(len2);
prop_assert_eq!(len2, header_new_len.len());
prop_assert_eq!(next, header_new_len.next());
let header_new_next = header.with_next(next2);
prop_assert_eq!(len, header_new_next.len());
prop_assert_eq!(next2, header_new_next.next());
prop_assert_eq!(header_new_len.with_next(next2).0, header_new_next.with_len(len2).0);
}
}
}