pipecircuit 0.0.10

use super::{SecondaryIndex, SecondaryIndexInconsistency};
use crate::storage::warp_table::IndexKey;
use crate::storage::warp_table::dense_key_index::{
    FindOutcome, PrepareRemoveOutcome, PreparedRemoval,
};
use crate::storage::warp_table::secondary_index::StableSlot;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SecondaryIndexRemoveError<SecondaryKey> {
    Inconsistent {
        secondary_key: SecondaryKey,
        slot: StableSlot,
        reason: SecondaryIndexInconsistency,
    },
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[must_use]
pub enum SecondaryIndexRemoveOutcome<SecondaryKey> {
    RemovedFromBucket {
        secondary_key: SecondaryKey,
        slot: StableSlot,
        remaining_slots: u64,
    },
    RemovedBucket {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
    BucketAbsent {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
    SlotAbsent {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
}

pub type SecondaryIndexRemoveResult<SecondaryKey> =
    Result<SecondaryIndexRemoveOutcome<SecondaryKey>, SecondaryIndexRemoveError<SecondaryKey>>;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[must_use]
pub enum SecondaryIndexCommitRemoveOutcome<SecondaryKey> {
    RemovedFromBucket {
        secondary_key: SecondaryKey,
        slot: StableSlot,
        remaining_slots: u64,
    },
    RemovedBucket {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
}

#[must_use]
pub enum SecondaryIndexPrepareRemoveOutcome<'index, SecondaryKey: IndexKey> {
    Prepared(PreparedSecondaryIndexRemove<'index, SecondaryKey>),
    BucketAbsent {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
    SlotAbsent {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
}

enum RemovePlan<SecondaryKey> {
    RetainBucket { row: usize },
    RemoveBucket(PreparedRemoval<SecondaryKey>),
}

pub struct PreparedSecondaryIndexRemove<'index, SecondaryKey: IndexKey> {
    index: &'index mut SecondaryIndex<SecondaryKey>,
    plan: RemovePlan<SecondaryKey>,
    secondary_key: SecondaryKey,
    slot: StableSlot,
}

impl<SecondaryKey: IndexKey> SecondaryIndex<SecondaryKey> {
    pub fn try_prepare_remove(
        &mut self,
        secondary_key: SecondaryKey,
        slot: StableSlot,
    ) -> Result<
        SecondaryIndexPrepareRemoveOutcome<'_, SecondaryKey>,
        SecondaryIndexRemoveError<SecondaryKey>,
    > {
        self.validate_lengths(secondary_key, slot)?;
        let row = match self
            .secondary_key_index
            .find(secondary_key, &self.secondary_keys)
            .map_err(|reason| inconsistent(secondary_key, slot, reason.into()))?
        {
            FindOutcome::Present { row, .. } => row,
            FindOutcome::Absent => {
                return Ok(SecondaryIndexPrepareRemoveOutcome::BucketAbsent {
                    secondary_key,
                    slot,
                });
            }
        };
        let slot_set = self.slot_sets.get(row).ok_or_else(|| {
            inconsistent(
                secondary_key,
                slot,
                SecondaryIndexInconsistency::SlotSetUnavailable { row },
            )
        })?;
        if slot_set.is_empty() {
            return Err(inconsistent(
                secondary_key,
                slot,
                SecondaryIndexInconsistency::EmptySlotSet { row },
            ));
        }
        if !slot_set.contains(slot) {
            return Ok(SecondaryIndexPrepareRemoveOutcome::SlotAbsent {
                secondary_key,
                slot,
            });
        }
        let plan = if slot_set.len() > 1 {
            RemovePlan::RetainBucket { row }
        } else {
            RemovePlan::RemoveBucket(self.prepare_bucket_removal(secondary_key, slot, row)?)
        };
        Ok(SecondaryIndexPrepareRemoveOutcome::Prepared(
            PreparedSecondaryIndexRemove {
                index: self,
                plan,
                secondary_key,
                slot,
            },
        ))
    }

    pub fn remove(
        &mut self,
        secondary_key: SecondaryKey,
        slot: StableSlot,
    ) -> SecondaryIndexRemoveResult<SecondaryKey> {
        match self.try_prepare_remove(secondary_key, slot)? {
            SecondaryIndexPrepareRemoveOutcome::Prepared(prepared) => Ok(match prepared.commit() {
                SecondaryIndexCommitRemoveOutcome::RemovedFromBucket {
                    secondary_key,
                    slot,
                    remaining_slots,
                } => SecondaryIndexRemoveOutcome::RemovedFromBucket {
                    secondary_key,
                    slot,
                    remaining_slots,
                },
                SecondaryIndexCommitRemoveOutcome::RemovedBucket {
                    secondary_key,
                    slot,
                } => SecondaryIndexRemoveOutcome::RemovedBucket {
                    secondary_key,
                    slot,
                },
            }),
            SecondaryIndexPrepareRemoveOutcome::BucketAbsent {
                secondary_key,
                slot,
            } => Ok(SecondaryIndexRemoveOutcome::BucketAbsent {
                secondary_key,
                slot,
            }),
            SecondaryIndexPrepareRemoveOutcome::SlotAbsent {
                secondary_key,
                slot,
            } => Ok(SecondaryIndexRemoveOutcome::SlotAbsent {
                secondary_key,
                slot,
            }),
        }
    }

    fn validate_lengths(
        &self,
        secondary_key: SecondaryKey,
        slot: StableSlot,
    ) -> Result<(), SecondaryIndexRemoveError<SecondaryKey>> {
        if self.secondary_keys.len() != self.slot_sets.len() {
            return Err(inconsistent(
                secondary_key,
                slot,
                SecondaryIndexInconsistency::KeySlotSetLengthMismatch {
                    secondary_keys: self.secondary_keys.len(),
                    slot_sets: self.slot_sets.len(),
                },
            ));
        }
        Ok(())
    }

    fn prepare_bucket_removal(
        &self,
        secondary_key: SecondaryKey,
        slot: StableSlot,
        row: usize,
    ) -> Result<PreparedRemoval<SecondaryKey>, SecondaryIndexRemoveError<SecondaryKey>> {
        let last_row = self.secondary_keys.len().checked_sub(1).ok_or_else(|| {
            inconsistent(
                secondary_key,
                slot,
                SecondaryIndexInconsistency::KeyUnavailable { row },
            )
        })?;
        match self
            .secondary_key_index
            .prepare_remove(secondary_key, &self.secondary_keys, last_row)
            .map_err(|reason| inconsistent(secondary_key, slot, reason.into()))?
        {
            PrepareRemoveOutcome::Prepared(removal) if removal.row == row => Ok(removal),
            PrepareRemoveOutcome::Prepared(removal) => Err(inconsistent(
                secondary_key,
                slot,
                SecondaryIndexInconsistency::PreparedRemovalRowMismatch {
                    found_row: row,
                    prepared_row: removal.row,
                },
            )),
            PrepareRemoveOutcome::Absent => Err(inconsistent(
                secondary_key,
                slot,
                SecondaryIndexInconsistency::PreparedRemovalAbsent { row },
            )),
        }
    }
}

impl<SecondaryKey: IndexKey> PreparedSecondaryIndexRemove<'_, SecondaryKey> {
    pub const fn secondary_key(&self) -> SecondaryKey {
        self.secondary_key
    }

    pub const fn slot(&self) -> StableSlot {
        self.slot
    }

    #[allow(
        clippy::indexing_slicing,
        reason = "removal preparation proves the retained bucket row exists"
    )]
    pub fn commit(self) -> SecondaryIndexCommitRemoveOutcome<SecondaryKey> {
        match self.plan {
            RemovePlan::RetainBucket { row } => {
                let slot_set = &mut self.index.slot_sets[row];
                slot_set.commit_prepared_remove(self.slot);
                SecondaryIndexCommitRemoveOutcome::RemovedFromBucket {
                    secondary_key: self.secondary_key,
                    slot: self.slot,
                    remaining_slots: slot_set.len(),
                }
            }
            RemovePlan::RemoveBucket(removal) => {
                let row = removal.row;
                self.index
                    .secondary_key_index
                    .commit_remove(removal, &self.index.secondary_keys);
                self.index.secondary_keys.swap_remove(row);
                self.index.slot_sets.swap_remove(row);
                SecondaryIndexCommitRemoveOutcome::RemovedBucket {
                    secondary_key: self.secondary_key,
                    slot: self.slot,
                }
            }
        }
    }
}

fn inconsistent<SecondaryKey>(
    secondary_key: SecondaryKey,
    slot: StableSlot,
    reason: SecondaryIndexInconsistency,
) -> SecondaryIndexRemoveError<SecondaryKey> {
    SecondaryIndexRemoveError::Inconsistent {
        secondary_key,
        slot,
        reason,
    }
}