pipecircuit 0.0.24

use super::{
    MembershipLookup, WarpTableMembershipRow, WorkspaceSelection, WorkspaceSelectionError,
};
use crate::storage::{
    IndexKey,
    query_workspace::{
        QueryChunk, QueryChunkStorage, QueryScratchStorage, QueryWorkspace, WorkspaceMembership,
    },
    warp_table::{
        WarpTable,
        secondary_index::{SecondaryIndex, StableSlots},
    },
};

mod allocated;
pub use allocated::AllocatedBulkUpdateError;
pub(in crate::storage::warp_table) use allocated::AllocatedBulkUpdateWorkspace;

pub struct BulkUpdateWorkspace<Chunks, Scratch, Pending, Group, Slots> {
    query: QueryWorkspace<Chunks, Scratch>,
    pending: Pending,
    group: Group,
    slots: Slots,
}

impl<Chunks, Scratch, Pending, Group, Slots>
    BulkUpdateWorkspace<Chunks, Scratch, Pending, Group, Slots>
{
    pub const fn new(
        query: QueryWorkspace<Chunks, Scratch>,
        pending: Pending,
        group: Group,
        slots: Slots,
    ) -> Self {
        Self {
            query,
            pending,
            group,
            slots,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BulkUpdateError {
    Selection(WorkspaceSelectionError),
    InsufficientStorage { required: usize, available: usize },
}
impl From<WorkspaceSelectionError> for BulkUpdateError {
    fn from(reason: WorkspaceSelectionError) -> Self {
        Self::Selection(reason)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[must_use]
pub enum BulkSetOutcome {
    Unchanged { selected: u64 },
    Updated { selected: u64, changed: u64 },
}

#[doc(hidden)]
pub struct BulkSelection<'update> {
    pub(in crate::storage::warp_table) selected: &'update WorkspaceMembership<'update>,
    pub(in crate::storage::warp_table) pending: &'update mut QueryChunk,
    pub(in crate::storage::warp_table) group: &'update mut QueryChunk,
    pub(in crate::storage::warp_table) slots: &'update mut [u16],
}
impl BulkSelection<'_> {
    pub fn set_secondary<Key: IndexKey>(
        &mut self,
        index: &mut SecondaryIndex<Key>,
        column: &mut [Key],
        stable_slots: &StableSlots,
        target: Key,
    ) -> Result<BulkSetOutcome, BulkUpdateError> {
        index.set_column_selection(column, stable_slots, target, self)
    }
}

pub trait WarpTableBulkRow: WarpTableMembershipRow {
    type BulkQuery<'table, 'workspace, Brand: 'table>
    where
        Self: 'table;
    fn bulk_query<'table, 'workspace, Brand: 'table>(
        query: BulkQuerySource<'table, 'workspace, Self, Brand>,
    ) -> Self::BulkQuery<'table, 'workspace, Brand>
    where
        Self: 'table;
    type BulkMut<'table>
    where
        Self: 'table;
    fn bulk_mut<'table>(
        columns: &'table mut Self::Columns,
        indexes: &'table mut Self::Indexes,
        selection: BulkSelection<'table>,
    ) -> Self::BulkMut<'table>;
}

#[doc(hidden)]
pub struct BulkQuerySource<'table, 'workspace, Row: WarpTableMembershipRow, Brand> {
    table: &'table WarpTable<Row, Brand>,
    membership: core::cell::Cell<Option<WorkspaceMembership<'workspace>>>,
}

impl<'table, 'workspace, Row: WarpTableMembershipRow, Brand: 'table>
    BulkQuerySource<'table, 'workspace, Row, Brand>
{
    pub fn find<Key>(
        &self,
        find: impl FnOnce(
            Row::MembershipQuery<'table, Brand>,
        ) -> MembershipLookup<'table, Row, Brand, Key>,
    ) -> Result<WorkspaceSelection<'table, 'workspace, Row, Brand>, WorkspaceSelectionError> {
        let membership = self
            .membership
            .take()
            .ok_or(WorkspaceSelectionError::SelectionAlreadyStarted)?;
        find(self.table.query_memberships()).in_workspace(membership)
    }
}

pub struct BulkQuery<'table, 'workspace, Row: WarpTableBulkRow + 'table, Brand: 'table> {
    query: Row::BulkQuery<'table, 'workspace, Brand>,
    lifetimes: core::marker::PhantomData<(&'table Row, &'workspace mut ())>,
}

impl<'table, 'workspace, Row: WarpTableBulkRow, Brand> core::ops::Deref
    for BulkQuery<'table, 'workspace, Row, Brand>
{
    type Target = Row::BulkQuery<'table, 'workspace, Brand>;
    fn deref(&self) -> &Self::Target {
        &self.query
    }
}

impl<Row: WarpTableBulkRow, Brand> WarpTable<Row, Brand> {
    pub fn bulk_update_in<Chunks, Scratch, Pending, Group, Slots, Output>(
        &mut self,
        workspace: &mut BulkUpdateWorkspace<Chunks, Scratch, Pending, Group, Slots>,
        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, BulkUpdateError>
    where
        Chunks: QueryChunkStorage,
        Scratch: QueryScratchStorage,
        Pending: QueryScratchStorage,
        Group: QueryScratchStorage,
        Slots: AsMut<[u16]>,
    {
        let selected = select(BulkQuery {
            query: Row::bulk_query(BulkQuerySource {
                table: self,
                membership: core::cell::Cell::new(Some(
                    workspace
                        .query
                        .membership()
                        .map_err(WorkspaceSelectionError::from)?,
                )),
            }),
            lifetimes: core::marker::PhantomData,
        })?;
        if !core::ptr::eq(self, selected.table) {
            return Err(WorkspaceSelectionError::DifferentTables.into());
        }
        let membership = selected.into_workspace();
        update(Row::bulk_mut(
            &mut self.columns,
            &mut self.indexes,
            BulkSelection {
                selected: &membership,
                pending: workspace.pending.as_mut(),
                group: workspace.group.as_mut(),
                slots: workspace.slots.as_mut(),
            },
        ))
    }
}

#[cfg(test)]
mod tests;