pipecircuit 0.0.25

Correct. Robust. Fast. Portable. Pick four.
Documentation
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;