pipecircuit 0.0.22

use super::{SecondaryIndex, SecondaryIndexInconsistency};
use crate::storage::warp_table::IndexKey;
use crate::storage::warp_table::dense_key_index::{
    CommitOutcome, PrepareInsertOutcome, PreparedCommit, PreparedVacancy,
};
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.prepare_primary_insert(secondary_key, slot, vacancy, row)?;
                let slot_set = SlotSet::one(slot);
                SecondaryIndexInsertion::VacantBucket {
                    index: self,
                    insertion,
                    slot_set,
                }
            }
        };
        Ok(PreparedSecondaryIndexInsert {
            insertion,
            secondary_key,
            slot,
        })
    }

    fn prepare_primary_insert(
        &self,
        secondary_key: SecondaryKey,
        slot: StableSlot,
        vacancy: PreparedVacancy,
        row: usize,
    ) -> Result<PreparedCommit, SecondaryIndexInsertError<SecondaryKey>> {
        self.secondary_key_index
            .prepare_commit(vacancy, row)
            .map_err(|reason| SecondaryIndexInsertError::Inconsistent {
                secondary_key,
                slot,
                reason: reason.into(),
            })
    }
}

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,
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn reserved_primary_row_preserves_the_secondary_key_slot_and_vacancy()
    -> Result<(), &'static str> {
        let mut index = SecondaryIndex::<u32>::try_new(1).map_err(|_| "index unavailable")?;
        let PrepareInsertOutcome::Vacant(vacancy) = index
            .secondary_key_index
            .prepare_insert(7, &index.secondary_keys)
            .map_err(|_| "vacancy unavailable")?
        else {
            return Err("fresh index was occupied");
        };
        let slot = StableSlot::new(11);
        let row = usize::try_from(u32::MAX).map_err(|_| "row boundary unavailable")?;
        assert_eq!(
            index.prepare_primary_insert(7, slot, vacancy, row),
            Err(SecondaryIndexInsertError::Inconsistent {
                secondary_key: 7,
                slot,
                reason: SecondaryIndexInconsistency::InsertRowReserved {
                    bucket: vacancy.bucket(),
                    row,
                },
            })
        );
        assert!(index.is_empty());
        assert_eq!(
            index
                .try_prepare_insert(7, slot)
                .map_err(|_| "valid insertion refused")?
                .commit(),
            SecondaryIndexInsertOutcome::CreatedBucket {
                secondary_key: 7,
                slot
            }
        );
        Ok(())
    }
}