use super::*;
use alloc::{boxed::Box, vec::Vec};
#[derive(Debug)]
pub enum AllocatedBulkUpdateError {
Reservation(alloc::collections::TryReserveError),
Update(BulkUpdateError),
}
impl From<alloc::collections::TryReserveError> for AllocatedBulkUpdateError {
fn from(reason: alloc::collections::TryReserveError) -> Self {
Self::Reservation(reason)
}
}
impl From<BulkUpdateError> for AllocatedBulkUpdateError {
fn from(reason: BulkUpdateError) -> Self {
Self::Update(reason)
}
}
pub(in crate::storage::warp_table) type AllocatedBulkUpdateWorkspace =
BulkUpdateWorkspace<Vec<QueryChunk>, QueryChunk, QueryChunk, QueryChunk, Vec<u16>>;
impl AllocatedBulkUpdateWorkspace {
#[cold]
#[inline(never)]
fn try_new(capacity: usize) -> Result<Box<Self>, AllocatedBulkUpdateError> {
let count = crate::storage::query_workspace::required_query_chunks(capacity)
.map_err(|reason| BulkUpdateError::Selection(reason.into()))?;
let mut chunks = Vec::new();
chunks.try_reserve_exact(count)?;
chunks.resize_with(count, || QueryChunk::EMPTY);
let mut slots = Vec::new();
slots.try_reserve_exact(capacity.min(65536))?;
slots.resize(capacity.min(65536), 0);
let query = QueryWorkspace::new(capacity, chunks, QueryChunk::EMPTY)
.map_err(|refusal| BulkUpdateError::Selection(refusal.reason.into()))?;
Ok(Box::new(Self::new(
query,
QueryChunk::EMPTY,
QueryChunk::EMPTY,
slots,
)))
}
}
struct BulkWorkspaceLease<'table, Row: WarpTableBulkRow, Brand> {
table: &'table mut WarpTable<Row, Brand>,
workspace: Option<Box<AllocatedBulkUpdateWorkspace>>,
}
impl<Row: WarpTableBulkRow, Brand> Drop for BulkWorkspaceLease<'_, Row, Brand> {
fn drop(&mut self) {
core::mem::swap(&mut self.table.bulk_workspace, &mut self.workspace);
}
}
fn reserve_workspace(
storage: &mut Option<Box<AllocatedBulkUpdateWorkspace>>,
capacity: usize,
) -> Result<&mut AllocatedBulkUpdateWorkspace, AllocatedBulkUpdateError> {
match storage {
Some(workspace) => Ok(workspace),
None => Ok(storage.insert(AllocatedBulkUpdateWorkspace::try_new(capacity)?)),
}
}
impl<Row: WarpTableBulkRow, Brand> WarpTable<Row, Brand> {
pub fn reserve_bulk_updates(&mut self) -> Result<(), AllocatedBulkUpdateError> {
reserve_workspace(&mut self.bulk_workspace, self.maximum_rows)?;
Ok(())
}
pub fn bulk_update<Output>(
&mut self,
select: impl for<'table, 'workspace> FnOnce(
BulkQuery<'table, 'workspace, Row, Brand>,
) -> Result<
WorkspaceSelection<'table, 'workspace, Row, Brand>,
WorkspaceSelectionError,
>,
update: impl FnOnce(Row::BulkMut<'_>) -> Result<Output, BulkUpdateError>,
) -> Result<Output, AllocatedBulkUpdateError> {
let mut lease = BulkWorkspaceLease {
workspace: core::mem::take(&mut self.bulk_workspace),
table: self,
};
let workspace = reserve_workspace(&mut lease.workspace, lease.table.maximum_rows)?;
lease
.table
.bulk_update_in(workspace, select, update)
.map_err(AllocatedBulkUpdateError::from)
}
}
#[cfg(test)]
mod tests;