pipecircuit 0.0.11

Correct. Robust. Fast. Portable.
Documentation
use alloc::vec::Vec;

use super::{SlotLocation, StableSlot, StableSlots, StableSlotsInconsistency};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StableSlotInsertError {
    Full { maximum_rows: usize },
    Inconsistent { reason: StableSlotsInconsistency },
}

impl StableSlots {
    pub fn try_prepare_insert(
        &mut self,
    ) -> Result<PreparedStableSlotInsert<'_>, StableSlotInsertError> {
        let Self {
            maximum_rows,
            row_to_slot,
            slot_locations,
            first_free,
        } = self;
        if row_to_slot.len() >= *maximum_rows {
            return Err(StableSlotInsertError::Full {
                maximum_rows: *maximum_rows,
            });
        }
        let row =
            u32::try_from(row_to_slot.len()).map_err(|_| StableSlotInsertError::Inconsistent {
                reason: StableSlotsInconsistency::RowUnrepresentable {
                    row: row_to_slot.len(),
                },
            })?;
        let insertion = match *first_free {
            Some(slot) => prepare_reuse(slot_locations, first_free, slot)?,
            None => prepare_new(slot_locations)?,
        };
        let slot = insertion.slot();
        Ok(PreparedStableSlotInsert {
            row_to_slot,
            insertion,
            slot,
            row,
        })
    }
}

fn prepare_reuse<'slots>(
    slot_locations: &'slots mut [SlotLocation],
    first_free: &'slots mut Option<StableSlot>,
    slot: StableSlot,
) -> Result<StableSlotInsertion<'slots>, StableSlotInsertError> {
    let slot_index =
        usize::try_from(slot.value()).map_err(|_| StableSlotInsertError::Inconsistent {
            reason: StableSlotsInconsistency::SlotUnrepresentable { slot },
        })?;
    let location =
        slot_locations
            .get_mut(slot_index)
            .ok_or(StableSlotInsertError::Inconsistent {
                reason: StableSlotsInconsistency::SlotUnavailable { slot },
            })?;
    let next_free = match location {
        SlotLocation::Free(next_free) => *next_free,
        SlotLocation::Occupied(row) => {
            return Err(StableSlotInsertError::Inconsistent {
                reason: StableSlotsInconsistency::SlotOccupied { slot, row: *row },
            });
        }
    };
    Ok(StableSlotInsertion::Reused {
        location,
        first_free,
        next_free,
        slot,
    })
}

fn prepare_new(
    slot_locations: &mut Vec<SlotLocation>,
) -> Result<StableSlotInsertion<'_>, StableSlotInsertError> {
    let slot_value =
        u32::try_from(slot_locations.len()).map_err(|_| StableSlotInsertError::Inconsistent {
            reason: StableSlotsInconsistency::NewSlotUnrepresentable {
                slot: slot_locations.len(),
            },
        })?;
    Ok(StableSlotInsertion::New {
        slot_locations,
        slot: StableSlot::new(slot_value),
    })
}

enum StableSlotInsertion<'slots> {
    New {
        slot_locations: &'slots mut Vec<SlotLocation>,
        slot: StableSlot,
    },
    Reused {
        location: &'slots mut SlotLocation,
        first_free: &'slots mut Option<StableSlot>,
        next_free: Option<StableSlot>,
        slot: StableSlot,
    },
}

impl StableSlotInsertion<'_> {
    const fn slot(&self) -> StableSlot {
        match self {
            Self::New { slot, .. } | Self::Reused { slot, .. } => *slot,
        }
    }
}

pub struct PreparedStableSlotInsert<'slots> {
    row_to_slot: &'slots mut Vec<StableSlot>,
    insertion: StableSlotInsertion<'slots>,
    slot: StableSlot,
    row: u32,
}

impl PreparedStableSlotInsert<'_> {
    pub const fn slot(&self) -> StableSlot {
        self.slot
    }

    pub fn commit(self) -> StableSlot {
        match self.insertion {
            StableSlotInsertion::New { slot_locations, .. } => {
                slot_locations.push(SlotLocation::Occupied(self.row));
            }
            StableSlotInsertion::Reused {
                location,
                first_free,
                next_free,
                ..
            } => {
                *location = SlotLocation::Occupied(self.row);
                *first_free = next_free;
            }
        }
        self.row_to_slot.push(self.slot);
        self.slot
    }
}