pipecircuit 0.0.9

use super::{SecondaryIndex, SecondaryIndexInconsistency};
use crate::storage::warp_table::IndexKey;
use crate::storage::warp_table::dense_key_index::{
    CommitOutcome, PrepareInsertOutcome, PreparedCommit,
};
use crate::storage::warp_table::secondary_index::{
    PreparedSlotSetInsert, SlotSet, SlotSetInsertError, SlotSetInsertOutcome, StableSlot,
};

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

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[must_use]
pub enum SecondaryIndexInsertOutcome<SecondaryKey> {
    CreatedBucket {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
    AddedToBucket {
        secondary_key: SecondaryKey,
        slot: StableSlot,
    },
}

enum SecondaryIndexInsertion<'index, SecondaryKey: IndexKey> {
    ExistingBucket(PreparedSlotSetInsert<'index>),
    VacantBucket {
        index: &'index mut SecondaryIndex<SecondaryKey>,
        insertion: PreparedCommit,
        slot_set: SlotSet,
    },
}

pub struct PreparedSecondaryIndexInsert<'index, SecondaryKey: IndexKey> {
    insertion: SecondaryIndexInsertion<'index, SecondaryKey>,
    secondary_key: SecondaryKey,
    slot: StableSlot,
}

impl<SecondaryKey: IndexKey> SecondaryIndex<SecondaryKey> {
    pub fn try_prepare_insert(
        &mut self,
        secondary_key: SecondaryKey,
        slot: StableSlot,
    ) -> Result<
        PreparedSecondaryIndexInsert<'_, SecondaryKey>,
        SecondaryIndexInsertError<SecondaryKey>,
    > {
        if self.secondary_keys.len() != self.slot_sets.len() {
            return Err(SecondaryIndexInsertError::Inconsistent {
                secondary_key,
                slot,
                reason: SecondaryIndexInconsistency::KeySlotSetLengthMismatch {
                    secondary_keys: self.secondary_keys.len(),
                    slot_sets: self.slot_sets.len(),
                },
            });
        }
        let prepared = self
            .secondary_key_index
            .prepare_insert(secondary_key, &self.secondary_keys)
            .map_err(|reason| SecondaryIndexInsertError::Inconsistent {
                secondary_key,
                slot,
                reason: reason.into(),
            })?;
        let insertion = match prepared {
            PrepareInsertOutcome::Present { row } => {
                let slot_set =
                    self.slot_sets
                        .get_mut(row)
                        .ok_or(SecondaryIndexInsertError::Inconsistent {
                            secondary_key,
                            slot,
                            reason: SecondaryIndexInconsistency::SlotSetUnavailable { row },
                        })?;
                let prepared =
                    slot_set
                        .try_prepare_insert(slot)
                        .map_err(|reason| match reason {
                            SlotSetInsertError::AlreadyPresent { slot } => {
                                SecondaryIndexInsertError::AlreadyPresent {
                                    secondary_key,
                                    slot,
                                }
                            }
                        })?;
                SecondaryIndexInsertion::ExistingBucket(prepared)
            }
            PrepareInsertOutcome::Vacant(vacancy) => {
                if self.secondary_keys.len() >= self.maximum_distinct_keys {
                    return Err(SecondaryIndexInsertError::Full {
                        secondary_key,
                        slot,
                        maximum_distinct_keys: self.maximum_distinct_keys,
                    });
                }
                let row = self.secondary_keys.len();
                let insertion = self
                    .secondary_key_index
                    .prepare_commit(vacancy, row)
                    .map_err(|reason| SecondaryIndexInsertError::Inconsistent {
                        secondary_key,
                        slot,
                        reason: reason.into(),
                    })?;
                let slot_set = SlotSet::one(slot);
                SecondaryIndexInsertion::VacantBucket {
                    index: self,
                    insertion,
                    slot_set,
                }
            }
        };
        Ok(PreparedSecondaryIndexInsert {
            insertion,
            secondary_key,
            slot,
        })
    }
}

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

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

    pub fn commit(self) -> SecondaryIndexInsertOutcome<SecondaryKey> {
        match self.insertion {
            SecondaryIndexInsertion::ExistingBucket(prepared) => {
                let SlotSetInsertOutcome::Inserted { slot } = prepared.commit();
                SecondaryIndexInsertOutcome::AddedToBucket {
                    secondary_key: self.secondary_key,
                    slot,
                }
            }
            SecondaryIndexInsertion::VacantBucket {
                index,
                insertion,
                slot_set,
            } => {
                let CommitOutcome::Committed { .. } =
                    index.secondary_key_index.commit_prepared(insertion);
                index.secondary_keys.push(self.secondary_key);
                index.slot_sets.push(slot_set);
                SecondaryIndexInsertOutcome::CreatedBucket {
                    secondary_key: self.secondary_key,
                    slot: self.slot,
                }
            }
        }
    }
}