#![forbid(unsafe_op_in_unsafe_fn)]
extern crate alloc;
use crate::ColId;
use alloc::alloc::{alloc, dealloc, handle_alloc_error, realloc};
use core::{
alloc::Layout,
cmp::Ordering,
fmt,
hash::{Hash, Hasher},
iter,
mem::{size_of, ManuallyDrop},
ops::{Deref, DerefMut},
ptr::NonNull,
slice::{from_raw_parts, from_raw_parts_mut},
};
use either::Either;
use itertools::Itertools;
#[macro_export]
macro_rules! col_list {
($($elem:expr),* $(,)?) => {{
$crate::ColList::from([$($elem),*])
}};
}
#[repr(C)]
pub union ColList {
check: usize,
inline: ColListInline,
heap: ManuallyDrop<ColListVec>,
}
unsafe impl Sync for ColList {}
unsafe impl Send for ColList {}
impl<C: Into<ColId>> From<C> for ColList {
fn from(value: C) -> Self {
Self::new(value.into())
}
}
impl<C: Into<ColId>, const N: usize> From<[C; N]> for ColList {
fn from(cols: [C; N]) -> Self {
cols.map(|c| c.into()).into_iter().collect()
}
}
impl<C: Into<ColId>> FromIterator<C> for ColList {
fn from_iter<I: IntoIterator<Item = C>>(iter: I) -> Self {
let iter = iter.into_iter();
let (lower_bound, _) = iter.size_hint();
let mut list = Self::with_capacity(lower_bound as u16);
list.extend(iter);
list
}
}
impl<C: Into<ColId>> Extend<C> for ColList {
fn extend<T: IntoIterator<Item = C>>(&mut self, iter: T) {
let iter = iter.into_iter();
for col in iter {
self.push(col.into());
}
}
}
impl Default for ColList {
fn default() -> Self {
Self::with_capacity(0)
}
}
impl ColList {
pub fn empty() -> Self {
Self::from_inline(0)
}
pub fn new(col: ColId) -> Self {
let mut list = Self::from_inline(0);
list.push_inner(col, true);
list
}
pub fn with_capacity(cap: u16) -> Self {
if cap < Self::FIRST_HEAP_COL_U16 {
Self::from_inline(0)
} else {
Self::from_heap(ColListVec::with_capacity(cap))
}
}
fn from_inline(list: u64) -> Self {
debug_assert_eq!(list & (1 << Self::FIRST_HEAP_COL), 0);
let inline = ColListInline(list << 1 | 1);
let ret = Self { inline };
debug_assert!(ret.is_inline());
ret
}
fn from_heap(vec: ColListVec) -> Self {
let heap = ManuallyDrop::new(vec);
Self { heap }
}
pub fn as_singleton(&self) -> Option<ColId> {
let mut iter = self.iter();
match (iter.next(), iter.next()) {
(h @ Some(_), None) => h,
_ => None,
}
}
pub fn head(&self) -> Option<ColId> {
self.iter().next()
}
pub fn last(&self) -> Option<ColId> {
match self.as_inline() {
Ok(inline) => inline.last(),
Err(heap) => heap.last().copied(),
}
}
pub fn contains(&self, needle: ColId) -> bool {
match self.as_inline() {
Ok(inline) => inline.contains(needle),
Err(heap) => heap.contains(&needle),
}
}
pub fn iter(&self) -> impl '_ + Clone + Iterator<Item = ColId> {
match self.as_inline() {
Ok(inline) => Either::Left(inline.iter()),
Err(heap) => Either::Right(heap.iter().copied()),
}
}
pub fn to_u16_vec(&self) -> alloc::boxed::Box<[u16]> {
self.iter().map(u16::from).collect()
}
pub fn len(&self) -> u16 {
match self.as_inline() {
Ok(inline) => inline.len(),
Err(heap) => heap.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn push(&mut self, col: ColId) {
self.push_inner(col, self.last().is_none_or(|l| l < col));
}
fn sort_dedup(&mut self) {
if let Err(heap) = self.as_inline_mut() {
heap.sort();
let is_deduped = is_sorted_and_deduped(heap);
let wants_inline = heap.last().unwrap_or(&ColId(0)).0 < Self::FIRST_HEAP_COL_U16;
if !is_deduped || wants_inline {
*self = Self::from_iter(heap.iter().copied().dedup());
}
}
}
#[inline]
fn push_inner(&mut self, col: ColId, preserves_set_order: bool) {
let val = u16::from(col) as u64;
match (val < Self::FIRST_HEAP_COL && preserves_set_order, self.as_inline_mut()) {
(true, Ok(inline)) => inline.0 |= 1 << (val + 1),
(false, Ok(inline)) => *self = Self::from_heap(inline.heapify_and_push(col)),
(_, Err(heap)) => heap.push(col),
}
}
const FIRST_HEAP_COL: u64 = size_of::<u64>() as u64 * 8 - 1;
const FIRST_HEAP_COL_U16: u16 = Self::FIRST_HEAP_COL as u16;
#[inline]
fn as_inline(&self) -> Result<&ColListInline, &ManuallyDrop<ColListVec>> {
if self.is_inline() {
Ok(unsafe { &self.inline })
} else {
Err(unsafe { &self.heap })
}
}
#[inline]
fn as_inline_mut(&mut self) -> Result<&mut ColListInline, &mut ManuallyDrop<ColListVec>> {
if self.is_inline() {
Ok(unsafe { &mut self.inline })
} else {
Err(unsafe { &mut self.heap })
}
}
#[inline]
fn is_inline(&self) -> bool {
let addr = unsafe { self.check };
addr & 1 != 0
}
}
#[cfg(feature = "memory-usage")]
impl spacetimedb_memory_usage::MemoryUsage for ColList {
fn heap_usage(&self) -> usize {
match self.as_inline() {
Ok(_) => 0,
Err(heap) => heap.capacity() as usize,
}
}
}
impl Drop for ColList {
fn drop(&mut self) {
if let Err(heap) = self.as_inline_mut() {
unsafe { ManuallyDrop::drop(heap) };
}
}
}
impl Clone for ColList {
fn clone(&self) -> Self {
match self.as_inline() {
Ok(inline) => Self { inline: *inline },
Err(heap) => Self { heap: heap.clone() },
}
}
}
impl Eq for ColList {}
impl PartialEq for ColList {
fn eq(&self, other: &Self) -> bool {
match (self.as_inline(), other.as_inline()) {
(Ok(lhs), Ok(rhs)) => lhs == rhs,
(Err(lhs), Err(rhs)) => ***lhs == ***rhs,
_ => false,
}
}
}
impl Ord for ColList {
fn cmp(&self, other: &Self) -> Ordering {
self.iter().cmp(other.iter())
}
}
impl PartialOrd for ColList {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Hash for ColList {
fn hash<H: Hasher>(&self, state: &mut H) {
match self.as_inline() {
Ok(inline) => inline.0.hash(state),
Err(heap) => heap.hash(state),
}
}
}
impl fmt::Debug for ColList {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list().entries(self.iter()).finish()
}
}
impl From<ColSet> for ColList {
fn from(value: ColSet) -> Self {
value.0
}
}
pub enum ColOrCols<'a> {
Col(ColId),
ColList(&'a ColList),
}
impl ColOrCols<'_> {
pub fn as_singleton(&self) -> Option<ColId> {
match self {
Self::Col(col) => Some(*col),
Self::ColList(cols) => cols.as_singleton(),
}
}
pub fn iter(&self) -> impl '_ + Iterator<Item = ColId> {
match self {
Self::Col(col) => Either::Left(iter::once(*col)),
Self::ColList(cols) => Either::Right(cols.iter()),
}
}
pub fn len(&self) -> u16 {
match self {
Self::Col(_) => 1,
Self::ColList(cols) => cols.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn to_owned(self) -> ColList {
match self {
Self::Col(col) => [col].into(),
Self::ColList(list) => list.clone(),
}
}
}
impl PartialEq<ColList> for ColOrCols<'_> {
fn eq(&self, other: &ColList) -> bool {
self.iter().eq(other.iter())
}
}
impl PartialEq for ColOrCols<'_> {
fn eq(&self, other: &Self) -> bool {
self.iter().eq(other.iter())
}
}
impl Eq for ColOrCols<'_> {}
impl Ord for ColOrCols<'_> {
fn cmp(&self, other: &Self) -> Ordering {
self.iter().cmp(other.iter())
}
}
impl PartialOrd for ColOrCols<'_> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ColSet(ColList);
impl ColSet {
pub fn contains(&self, needle: ColId) -> bool {
match self.as_inline() {
Ok(inline) => inline.contains(needle),
Err(heap) => heap.binary_search(&needle).is_ok(),
}
}
}
impl<C: Into<ColId>> FromIterator<C> for ColSet {
fn from_iter<T: IntoIterator<Item = C>>(iter: T) -> Self {
Self::from(iter.into_iter().collect::<ColList>())
}
}
impl From<ColList> for ColSet {
fn from(mut list: ColList) -> Self {
list.sort_dedup();
Self(list)
}
}
impl From<&ColList> for ColSet {
fn from(value: &ColList) -> Self {
value.iter().collect()
}
}
impl From<ColOrCols<'_>> for ColSet {
fn from(value: ColOrCols<'_>) -> Self {
match value {
ColOrCols::Col(col) => ColSet(col.into()),
ColOrCols::ColList(cols) => cols.into(),
}
}
}
impl<C: Into<ColId>> From<C> for ColSet {
fn from(value: C) -> Self {
Self::from(ColList::new(value.into()))
}
}
impl Deref for ColSet {
type Target = ColList;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl fmt::Debug for ColSet {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_set().entries(self.iter()).finish()
}
}
#[derive(Clone, Copy, PartialEq)]
struct ColListInline(u64);
impl ColListInline {
fn contains(&self, needle: ColId) -> bool {
let col = needle.0;
let inline = self.undo_mark();
col < ColList::FIRST_HEAP_COL_U16 && inline & (1u64 << col) != 0
}
fn iter(&self) -> impl '_ + Clone + Iterator<Item = ColId> {
let mut value = self.undo_mark();
iter::from_fn(move || {
if value == 0 {
None
} else {
let id = ColId(value.trailing_zeros() as u16);
value &= value.wrapping_sub(1);
Some(id)
}
})
}
fn last(&self) -> Option<ColId> {
(u64::BITS - self.undo_mark().leading_zeros())
.checked_sub(1)
.map(|c| ColId(c as _))
}
fn len(&self) -> u16 {
self.undo_mark().count_ones() as u16
}
#[inline]
fn undo_mark(&self) -> u64 {
self.0 >> 1
}
fn heapify_and_push(&self, col: ColId) -> ColListVec {
let mut vec = ColListVec::with_capacity(2 * (self.len() + 1));
for col in self.iter() {
vec.push(col)
}
vec.push(col);
vec
}
}
struct ColListVec(NonNull<u16>);
impl ColListVec {
fn with_capacity(capacity: u16) -> Self {
let layout = Self::layout(capacity);
let ptr = unsafe { alloc(layout) }.cast::<u16>();
let Some(ptr_non_null) = NonNull::new(ptr) else {
handle_alloc_error(layout)
};
let mut this = Self(ptr_non_null);
unsafe {
this.set_len(0);
}
unsafe { this.set_capacity(capacity) };
this
}
fn len(&self) -> u16 {
let ptr = self.0.as_ptr();
unsafe { *ptr }
}
unsafe fn set_len(&mut self, new_len: u16) {
let ptr = self.0.as_ptr();
unsafe {
*ptr = new_len;
}
}
fn capacity(&self) -> u16 {
let ptr = self.0.as_ptr();
let capacity_ptr = unsafe { ptr.add(1) };
unsafe { *capacity_ptr }
}
unsafe fn set_capacity(&mut self, cap: u16) {
let ptr = self.0.as_ptr();
let cap_ptr = unsafe { ptr.add(1) };
unsafe {
*cap_ptr = cap;
}
}
fn push(&mut self, val: ColId) {
let len = self.len();
let cap = self.capacity();
if len == cap {
let new_cap = cap.checked_mul(2).expect("capacity overflow");
let new_layout = Self::layout(new_cap);
let old_layout = Self::layout(cap);
let old_ptr = self.0.as_ptr().cast();
let new_ptr = unsafe { realloc(old_ptr, old_layout, new_layout.size()) }.cast();
let Some(ptr_non_null) = NonNull::new(new_ptr) else {
handle_alloc_error(new_layout);
};
self.0 = ptr_non_null;
unsafe { self.set_capacity(new_cap) };
}
let base_ptr = self.0.as_ptr();
let elem_offset = 2 + len as usize;
let elem_ptr = unsafe { base_ptr.add(elem_offset) }.cast();
unsafe {
*elem_ptr = val;
}
unsafe {
self.set_len(len + 1);
}
}
fn layout(cap: u16) -> Layout {
Layout::array::<u16>(cap.checked_add(2).expect("capacity overflow") as usize).unwrap()
}
}
impl Deref for ColListVec {
type Target = [ColId];
fn deref(&self) -> &Self::Target {
let len = self.len() as usize;
let ptr = self.0.as_ptr();
let ptr = unsafe { ptr.add(2) }.cast::<ColId>();
unsafe { from_raw_parts(ptr, len) }
}
}
impl DerefMut for ColListVec {
fn deref_mut(&mut self) -> &mut Self::Target {
let len = self.len() as usize;
let ptr = self.0.as_ptr();
let ptr = unsafe { ptr.add(2) }.cast::<ColId>();
unsafe { from_raw_parts_mut(ptr, len) }
}
}
impl Drop for ColListVec {
fn drop(&mut self) {
let capacity = self.capacity();
let layout = Self::layout(capacity);
let ptr = self.0.as_ptr().cast();
unsafe { dealloc(ptr, layout) };
}
}
impl Clone for ColListVec {
fn clone(&self) -> Self {
let mut vec = ColListVec::with_capacity(self.len());
for col in self.iter().copied() {
vec.push(col);
}
vec
}
}
fn is_sorted_and_deduped(data: &[ColId]) -> bool {
match data {
[] => true,
&[mut prev, ref rest @ ..] => !rest.iter().any(|elem| {
let bad = prev >= *elem;
prev = *elem;
bad
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::collection::vec;
use proptest::prelude::*;
fn contains(list: &ColList, x: &ColId) -> bool {
list.iter().any(|y| y == *x)
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(if cfg!(miri) { 8 } else { 2048 }))]
#[test]
fn test_inline(cols in vec((0..ColList::FIRST_HEAP_COL_U16).prop_map_into(), 1..100)) {
let [head, tail @ ..] = &*cols else { unreachable!() };
let mut list = ColList::new(*head);
let mut is_inline = list.is_inline();
prop_assert!(is_inline);
prop_assert!(!list.is_empty());
prop_assert_eq!(list.len(), 1);
prop_assert_eq!(list.head(), Some(*head));
prop_assert_eq!(list.last(), Some(*head));
prop_assert_eq!(list.iter().collect::<Vec<_>>(), [*head]);
for col in tail {
is_inline &= list.last().unwrap() < *col;
list.push(*col);
prop_assert_eq!(is_inline, list.is_inline());
prop_assert!(!list.is_empty());
prop_assert_eq!(list.head(), Some(*head));
prop_assert_eq!(list.last(), Some(*col));
prop_assert_eq!(list.last(), list.iter().last());
prop_assert!(contains(&list, col));
}
prop_assert_eq!(&list.clone(), &list);
prop_assert_eq!(list.iter().collect::<Vec<_>>(), cols);
}
#[test]
fn test_heap(cols in vec((ColList::FIRST_HEAP_COL_U16..).prop_map_into(), 1..100)) {
let contains = |list: &ColList, x| list.iter().collect::<Vec<_>>().contains(x);
let head = ColId(0);
let mut list = ColList::new(head);
prop_assert!(list.is_inline());
prop_assert_eq!(list.len(), 1);
for (idx, col) in cols.iter().enumerate() {
list.push(*col);
prop_assert!(!list.is_inline());
prop_assert!(!list.is_empty());
prop_assert_eq!(list.len() as usize, idx + 2);
prop_assert_eq!(list.head(), Some(head));
prop_assert_eq!(list.last(), Some(*col));
prop_assert!(contains(&list, col));
}
prop_assert_eq!(&list.clone(), &list);
let mut cols = cols;
cols.insert(0, head);
prop_assert_eq!(list.iter().collect::<Vec<_>>(), cols);
}
#[test]
fn test_collect(cols in vec((0..100).prop_map_into(), 0..100)) {
let list = cols.iter().copied().collect::<ColList>();
prop_assert!(list.iter().eq(cols));
prop_assert_eq!(&list, &list.iter().collect::<ColList>());
}
#[test]
fn test_as_singleton(cols in vec((0..100).prop_map_into(), 0..10)) {
let list = cols.iter().copied().collect::<ColList>();
match cols.len() {
1 => {
prop_assert_eq!(list.as_singleton(), Some(cols[0]));
prop_assert_eq!(list.as_singleton(), list.head());
},
_ => prop_assert_eq!(list.as_singleton(), None),
}
}
#[test]
fn test_set_inlines(mut cols in vec((0..ColList::FIRST_HEAP_COL_U16).prop_map_into(), 1..100)) {
prop_assume!(!is_sorted_and_deduped(&cols[..]));
let list = ColList::from_iter(cols.iter().copied());
prop_assert!(!list.is_inline());
let set = ColSet::from(list);
prop_assert!(set.is_inline());
for col in cols.iter() {
prop_assert!(set.contains(*col));
}
cols.sort();
cols.dedup();
prop_assert_eq!(set.iter().collect::<Vec<_>>(), cols);
}
#[test]
fn test_set_heap(mut cols in vec((ColList::FIRST_HEAP_COL_U16..).prop_map_into(), 1..100)) {
prop_assume!(!is_sorted_and_deduped(&cols[..]));
let list = ColList::from_iter(cols.iter().copied());
prop_assert!(!list.is_inline());
let set = ColSet::from(list);
prop_assert!(!set.is_inline());
for col in cols.iter() {
prop_assert!(set.contains(*col));
}
cols.sort();
cols.dedup();
prop_assert_eq!(set.iter().collect::<Vec<_>>(), cols);
}
}
}