ic-memory 0.26.2

use crate::{
    capability::{CommittedAllocations, ValidatedAllocations},
    declaration::AllocationDeclaration,
    declaration::DeclarationSnapshot,
    ledger::{
        AllocationLedger, AllocationReservationError, AllocationRetirement,
        AllocationRetirementError, AllocationStageError, LedgerCommitError, LedgerCommitStore,
        checked_reservation_count, stage_reservation_generation, stage_retirement_generation,
        stage_validated_generation,
    },
    policy::AllocationPolicy,
    validation::{AllocationValidationError, validate_allocations},
};
use std::borrow::Cow;

///
/// AllocationBootstrap
///
/// Golden-path allocation ledger bootstrap pipeline.
///
/// This type owns allocation-governance sequencing only: recover the persisted
/// ledger, apply the owner layer's policy, validate current declarations
/// against ledger history, stage and commit the next generation, and return
/// a pending [`PendingBootstrapCommit`] after the in-memory commit store advances.
/// The persistence owner must durably write that state and explicitly confirm
/// persistence before it can obtain [`CommittedAllocations`].
///
/// `AllocationBootstrap` is for whichever layer owns a given `ic-memory`
/// ledger store. That owner may be a framework such as Canic, a library such as
/// IcyDB using `ic-memory` directly, or a standalone application canister. The
/// ownership model is not a fixed `ic-memory -> Canic -> IcyDB -> application`
/// chain.
///
/// Exactly one owner should bootstrap a given ledger store. If multiple layers
/// use `ic-memory` in the same canister, they must either compose their
/// declarations into one bootstrap owner or use distinct ledger stores and
/// allocation domains.
///
/// The owner still decides when bootstrap runs, how the ledger store is backed
/// by stable memory, and when endpoint dispatch or stable-memory handle opening
/// is allowed.
#[derive(Debug)]
pub struct AllocationBootstrap<'store> {
    store: &'store mut LedgerCommitStore,
}

impl<'store> AllocationBootstrap<'store> {
    /// Build a bootstrap pipeline over a protected ledger commit store.
    pub const fn new(store: &'store mut LedgerCommitStore) -> Self {
        Self { store }
    }

    /// Recover, validate, stage, and advance one pending allocation generation.
    pub fn validate_and_commit<P>(
        &mut self,
        snapshot: DeclarationSnapshot,
        policy: &P,
        committed_at: Option<u64>,
    ) -> Result<PendingBootstrapCommit, BootstrapError<P::Error>>
    where
        P: AllocationPolicy,
    {
        let prior = self.store.recover().map_err(BootstrapError::Ledger)?;
        self.validate_against(prior, snapshot, policy, committed_at)
    }

    /// Initialize an empty ledger store, then validate and advance a pending commit.
    ///
    /// This is the privileged genesis/import path. Normal runtime users should
    /// use [`crate::MemoryRuntime::bootstrap`] directly or the default TLS
    /// convenience bootstrap, both of which supply an empty current-format
    /// genesis ledger. A non-empty `genesis` should only be supplied by the layer
    /// that owns migration or import for this ledger store.
    ///
    /// The generic crate guarantees only that `genesis` is used when the
    /// protected physical store is empty, never when recovery sees corrupt or
    /// partially written state.
    pub fn initialize_validate_and_commit<P>(
        &mut self,
        genesis: &AllocationLedger,
        snapshot: DeclarationSnapshot,
        policy: &P,
        committed_at: Option<u64>,
    ) -> Result<PendingBootstrapCommit, BootstrapError<P::Error>>
    where
        P: AllocationPolicy,
    {
        let prior = self
            .store
            .recover_or_initialize(genesis)
            .map_err(BootstrapError::Ledger)?;
        self.validate_against(prior, snapshot, policy, committed_at)
    }

    /// Recover, policy-check, reserve, and commit one reservation generation.
    ///
    /// Declarations carry checked metadata. After recovery, batches exceeding
    /// 255 items reject before policy callbacks or historical claim checks.
    pub fn reserve_and_commit<P>(
        &mut self,
        reservations: &[AllocationDeclaration],
        policy: &P,
        committed_at: Option<u64>,
    ) -> Result<AllocationLedger, BootstrapReservationError<P::Error>>
    where
        P: AllocationPolicy,
    {
        let prior = self
            .store
            .recover()
            .map_err(BootstrapReservationError::Ledger)?;
        self.reserve_against(prior.into_ledger(), reservations, policy, committed_at)
    }

    /// Initialize an empty ledger store, then reserve and commit.
    ///
    /// This is the privileged genesis/import path for reservation staging. A
    /// non-empty `genesis` should only be supplied by the owner of migration or
    /// import for this ledger store.
    /// Recovery or initialization precedes batch validation. Batches exceeding
    /// 255 items reject before policy callbacks or historical claim checks.
    pub fn initialize_reserve_and_commit<P>(
        &mut self,
        genesis: &AllocationLedger,
        reservations: &[AllocationDeclaration],
        policy: &P,
        committed_at: Option<u64>,
    ) -> Result<AllocationLedger, BootstrapReservationError<P::Error>>
    where
        P: AllocationPolicy,
    {
        let prior = self
            .store
            .recover_or_initialize(genesis)
            .map_err(BootstrapReservationError::Ledger)?;
        self.reserve_against(prior.into_ledger(), reservations, policy, committed_at)
    }

    /// Recover, retire, and commit one explicit retirement generation.
    pub fn retire_and_commit(
        &mut self,
        retirement: &AllocationRetirement,
        committed_at: Option<u64>,
    ) -> Result<AllocationLedger, BootstrapRetirementError> {
        let prior = self
            .store
            .recover()
            .map_err(BootstrapRetirementError::Ledger)?;
        let staged =
            stage_retirement_generation(Cow::Owned(prior.into_ledger()), retirement, committed_at)
                .map_err(BootstrapRetirementError::Retirement)?;
        self.store
            .commit_generation(&staged)
            .map_err(BootstrapRetirementError::Ledger)?;
        Ok(staged)
    }

    fn reserve_against<P>(
        &mut self,
        prior: AllocationLedger,
        reservations: &[AllocationDeclaration],
        policy: &P,
        committed_at: Option<u64>,
    ) -> Result<AllocationLedger, BootstrapReservationError<P::Error>>
    where
        P: AllocationPolicy,
    {
        checked_reservation_count(reservations.len())
            .map_err(BootstrapReservationError::Reservation)?;
        for reservation in reservations {
            policy
                .validate_key(&reservation.stable_key)
                .map_err(BootstrapReservationError::Policy)?;
            policy
                .validate_reserved_slot(&reservation.stable_key, &reservation.slot)
                .map_err(BootstrapReservationError::Policy)?;
        }

        let staged = stage_reservation_generation(Cow::Owned(prior), reservations, committed_at)
            .map_err(BootstrapReservationError::Reservation)?;
        self.store
            .commit_generation(&staged)
            .map_err(BootstrapReservationError::Ledger)?;
        Ok(staged)
    }

    pub(crate) fn validate_against<P>(
        &mut self,
        prior: crate::RecoveredLedger,
        snapshot: DeclarationSnapshot,
        policy: &P,
        committed_at: Option<u64>,
    ) -> Result<PendingBootstrapCommit, BootstrapError<P::Error>>
    where
        P: AllocationPolicy,
    {
        let validated =
            validate_allocations(&prior, snapshot, policy).map_err(BootstrapError::Validation)?;
        let staged =
            stage_validated_generation(Cow::Owned(prior.into_ledger()), &validated, committed_at)
                .map_err(BootstrapError::Staging)?;
        self.store
            .commit_generation(&staged)
            .map_err(BootstrapError::Ledger)?;

        Ok(PendingBootstrapCommit {
            validated,
            ledger: staged,
        })
    }
}

///
/// PendingBootstrapCommit
///
/// Pending result of a successful generic allocation bootstrap commit.
///
/// The embedded [`crate::LedgerCommitStore`] has advanced, but this generic
/// layer does not own stable-memory IO. Persist the owning record first, then
/// call [`PendingBootstrapCommit::confirm_persisted`] to mint the allocation-open
/// capability.
///

#[derive(Debug, Eq, PartialEq)]
#[must_use = "persist the owning ledger record, then confirm_persisted before opening allocations"]
pub struct PendingBootstrapCommit {
    /// Staged ledger accepted by the protected generation commit.
    ledger: AllocationLedger,
    /// Validated allocation declarations awaiting persistence confirmation.
    validated: ValidatedAllocations,
}

impl PendingBootstrapCommit {
    /// Borrow the committed logical ledger for diagnostics.
    ///
    /// The persistence owner must write the owning record that contains the
    /// mutated [`crate::LedgerCommitStore`], not serialize this ledger DTO as a
    /// replacement protocol.
    #[must_use]
    pub const fn ledger(&self) -> &AllocationLedger {
        &self.ledger
    }

    /// Borrow the pre-commit validation result for diagnostics.
    #[must_use]
    pub const fn validated(&self) -> &ValidatedAllocations {
        &self.validated
    }

    /// Confirm that the owning integration durably persisted this commit.
    ///
    /// Calling this method before the stable-memory write succeeds violates the
    /// allocation protocol. The default runtime performs its stable-cell write
    /// before confirmation.
    #[must_use]
    pub fn confirm_persisted(self) -> CommittedAllocations {
        self.validated
            .confirm_persisted(self.ledger.current_generation())
    }
}

///
/// BootstrapError
///
/// Failure to recover, validate, or commit an allocation generation.
#[non_exhaustive]
#[derive(Clone, Debug, Eq, thiserror::Error, PartialEq)]
pub enum BootstrapError<P> {
    /// Ledger recovery or protected commit failed.
    #[error(transparent)]
    Ledger(LedgerCommitError),
    /// Policy or historical allocation validation failed.
    #[error(transparent)]
    Validation(AllocationValidationError<P>),
    /// Validated declarations could not be staged against the recovered ledger.
    #[error(transparent)]
    Staging(AllocationStageError),
}

///
/// BootstrapReservationError
///
/// Failure to policy-check, stage, or commit an allocation reservation.
#[non_exhaustive]
#[derive(Clone, Debug, Eq, thiserror::Error, PartialEq)]
pub enum BootstrapReservationError<P> {
    /// Ledger recovery or protected commit failed.
    #[error(transparent)]
    Ledger(LedgerCommitError),
    /// Policy adapter rejected a reservation declaration.
    #[error("allocation policy rejected a reservation")]
    Policy(P),
    /// Reservation conflicted with historical allocation facts.
    #[error(transparent)]
    Reservation(AllocationReservationError),
}

///
/// BootstrapRetirementError
///
/// Failure to stage or commit an explicit allocation retirement.
#[non_exhaustive]
#[derive(Clone, Debug, Eq, thiserror::Error, PartialEq)]
pub enum BootstrapRetirementError {
    /// Ledger recovery or protected commit failed.
    #[error(transparent)]
    Ledger(LedgerCommitError),
    /// Retirement conflicted with historical allocation facts.
    #[error(transparent)]
    Retirement(AllocationRetirementError),
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        declaration::AllocationDeclaration,
        ledger::{AllocationHistory, AllocationLedger, AllocationState},
        schema::SchemaMetadata,
        slot::MemoryManagerSlot,
    };

    #[derive(Debug, Eq, PartialEq)]
    struct TestPolicy;

    impl AllocationPolicy for TestPolicy {
        type Error = &'static str;

        fn validate_key(&self, _key: &crate::StableKey) -> Result<(), Self::Error> {
            Ok(())
        }

        fn validate_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            Ok(())
        }

        fn validate_reserved_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            Ok(())
        }
    }

    #[derive(Debug, Eq, PartialEq)]
    struct RejectReservedPolicy;

    impl AllocationPolicy for RejectReservedPolicy {
        type Error = &'static str;

        fn validate_key(&self, _key: &crate::StableKey) -> Result<(), Self::Error> {
            Ok(())
        }

        fn validate_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            Ok(())
        }

        fn validate_reserved_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            Err("reserved slot rejected")
        }
    }

    #[derive(Debug, Eq, PartialEq)]
    struct RejectActivePolicy;

    impl AllocationPolicy for RejectActivePolicy {
        type Error = &'static str;

        fn validate_key(&self, _key: &crate::StableKey) -> Result<(), Self::Error> {
            Ok(())
        }

        fn validate_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            Err("active slot rejected")
        }

        fn validate_reserved_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            Ok(())
        }
    }

    struct PolicyMustNotRun;

    impl AllocationPolicy for PolicyMustNotRun {
        type Error = &'static str;

        fn validate_key(&self, _key: &crate::StableKey) -> Result<(), Self::Error> {
            panic!("policy received an invalid reservation")
        }

        fn validate_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            panic!("policy received an invalid reservation")
        }

        fn validate_reserved_slot(
            &self,
            _key: &crate::StableKey,
            _slot: &MemoryManagerSlot,
        ) -> Result<(), Self::Error> {
            panic!("policy received an invalid reservation")
        }
    }

    fn ledger() -> AllocationLedger {
        AllocationLedger {
            current_generation: 0,
            allocation_history: AllocationHistory::default(),
        }
    }

    fn declaration() -> AllocationDeclaration {
        AllocationDeclaration::new(
            "app.users.v1",
            MemoryManagerSlot::new(100).expect("usable slot"),
            None,
            SchemaMetadata::default(),
        )
        .expect("declaration")
    }

    #[test]
    fn validate_and_commit_publishes_committed_generation() {
        let mut store = LedgerCommitStore::default();
        store.commit(&ledger()).expect("initial ledger");
        let snapshot = DeclarationSnapshot::new(vec![declaration()]).expect("snapshot");

        let commit = AllocationBootstrap::new(&mut store)
            .validate_and_commit(snapshot, &TestPolicy, Some(42))
            .expect("bootstrap commit");

        assert_eq!(commit.ledger().current_generation, 1);
        assert_eq!(commit.ledger().allocation_history.records().len(), 1);
        assert_eq!(commit.ledger().allocation_history.generations().len(), 1);
        assert_eq!(store.recover().unwrap().ledger(), commit.ledger());
        assert_eq!(commit.confirm_persisted().generation(), 1);
    }

    #[test]
    fn initialize_validate_and_commit_seeds_empty_ledger_store() {
        let mut store = LedgerCommitStore::default();
        let snapshot = DeclarationSnapshot::new(vec![declaration()]).expect("snapshot");

        let commit = AllocationBootstrap::new(&mut store)
            .initialize_validate_and_commit(&ledger(), snapshot, &TestPolicy, Some(42))
            .expect("bootstrap commit");

        assert_eq!(commit.ledger().current_generation, 1);
        assert_eq!(commit.ledger().allocation_history.records().len(), 1);
        assert_eq!(commit.confirm_persisted().generation(), 1);
    }

    #[test]
    fn initialize_validate_and_commit_fails_closed_on_corrupt_store() {
        let mut store = LedgerCommitStore::default();
        store
            .write_corrupt_inactive_ledger(&ledger())
            .expect("corrupt ledger");
        let snapshot = DeclarationSnapshot::new(vec![declaration()]).expect("snapshot");

        let err = AllocationBootstrap::new(&mut store)
            .initialize_validate_and_commit(&ledger(), snapshot, &TestPolicy, Some(42))
            .expect_err("corrupt state");

        assert!(matches!(err, BootstrapError::Ledger(_)));
    }

    #[test]
    fn reserve_and_commit_policy_checks_and_commits_reservation() {
        let mut store = LedgerCommitStore::default();
        store.commit(&ledger()).expect("initial ledger");
        let reservation = declaration();

        let committed = AllocationBootstrap::new(&mut store)
            .reserve_and_commit(&[reservation], &TestPolicy, Some(42))
            .expect("reservation commit");

        assert_eq!(committed.current_generation, 1);
        assert_eq!(committed.allocation_history.records().len(), 1);
        assert_eq!(
            committed.allocation_history.records()[0].state(),
            AllocationState::Reserved
        );
        assert_eq!(store.recover().unwrap().ledger(), &committed);

        // The first reservation changes local staging state before the second
        // tries to move an existing key. Neither borrowed staging nor bootstrap
        // may expose that partial batch or advance the protected store.
        let reservations = [
            AllocationDeclaration::memory_manager_unlabeled("app.future.v1", 101).unwrap(),
            AllocationDeclaration::memory_manager_unlabeled("app.users.v1", 102).unwrap(),
        ];
        let before = store.clone();
        let expected = committed
            .stage_reservation_generation(&reservations, None)
            .unwrap_err();
        assert!(matches!(
            expected,
            AllocationReservationError::StableKeySlotConflict { .. }
        ));
        assert_eq!(committed, *store.recover().unwrap().ledger());
        let error = AllocationBootstrap::new(&mut store)
            .reserve_and_commit(&reservations, &TestPolicy, None)
            .unwrap_err();
        assert_eq!(error, BootstrapReservationError::Reservation(expected));
        assert_eq!(store, before);
    }

    #[test]
    fn initialize_reserve_and_commit_seeds_empty_store() {
        let mut store = LedgerCommitStore::default();
        let reservation = declaration();

        let committed = AllocationBootstrap::new(&mut store)
            .initialize_reserve_and_commit(&ledger(), &[reservation], &TestPolicy, Some(42))
            .expect("reservation commit");

        assert_eq!(committed.current_generation, 1);
        assert_eq!(
            committed.allocation_history.records()[0].state(),
            AllocationState::Reserved
        );
    }

    #[test]
    fn reserve_and_commit_rejects_policy_failure_before_commit() {
        let mut store = LedgerCommitStore::default();
        store.commit(&ledger()).expect("initial ledger");
        let reservation = declaration();

        let err = AllocationBootstrap::new(&mut store)
            .reserve_and_commit(&[reservation], &RejectReservedPolicy, Some(42))
            .expect_err("policy failure");
        let recovered = store.recover().expect("recovered");

        assert!(matches!(err, BootstrapReservationError::Policy(_)));
        assert_eq!(recovered.current_generation(), 0);
        assert_eq!(recovered.ledger().allocation_history().records(), []);
    }

    #[test]
    fn reservation_pipeline_accepts_empty_and_full_slot_domain_batches() {
        for count in [0_u8, 255] {
            let reservations = (0..count)
                .map(|id| {
                    AllocationDeclaration::memory_manager_unlabeled(
                        format!("app.future{id}.v1"),
                        id,
                    )
                    .expect("reservation")
                })
                .collect::<Vec<_>>();
            let mut store = LedgerCommitStore::default();
            store.commit(&ledger()).expect("initial ledger");

            let committed = AllocationBootstrap::new(&mut store)
                .reserve_and_commit(&reservations, &TestPolicy, Some(42))
                .expect("bounded batch commits");

            assert_eq!(committed.current_generation(), 1);
            assert_eq!(
                committed.allocation_history().records().len(),
                usize::from(count)
            );
            assert_eq!(
                committed.allocation_history().generations()[0].declaration_count(),
                u32::from(count)
            );
            assert_eq!(store.recover().unwrap().ledger(), &committed);
        }
    }

    #[test]
    fn oversized_reservations_reject_before_policy_and_preserve_existing_store() {
        let reservations = vec![declaration(); 256];
        for initialize in [false, true] {
            let mut store = LedgerCommitStore::default();
            store.commit(&ledger()).expect("initial ledger");
            let before = store.clone();
            let mut bootstrap = AllocationBootstrap::new(&mut store);
            let error = if initialize {
                bootstrap.initialize_reserve_and_commit(
                    &ledger(),
                    &reservations,
                    &PolicyMustNotRun,
                    None,
                )
            } else {
                bootstrap.reserve_and_commit(&reservations, &PolicyMustNotRun, None)
            }
            .expect_err("oversized batch must fail before policy");

            assert_eq!(
                error,
                BootstrapReservationError::Reservation(
                    AllocationReservationError::TooManyReservations { count: 256 }
                )
            );
            assert_eq!(store, before);
        }
    }

    #[test]
    fn oversized_initial_reservations_preserve_genesis_before_policy() {
        let reservations = vec![declaration(); 256];
        let mut store = LedgerCommitStore::default();
        let mut expected = LedgerCommitStore::default();
        expected.commit(&ledger()).expect("expected genesis");

        let error = AllocationBootstrap::new(&mut store)
            .initialize_reserve_and_commit(&ledger(), &reservations, &PolicyMustNotRun, None)
            .expect_err("size rejection precedes policy checks");

        assert_eq!(
            error,
            BootstrapReservationError::Reservation(
                AllocationReservationError::TooManyReservations { count: 256 }
            )
        );
        assert_eq!(store, expected);
    }

    #[test]
    fn reservation_policy_alone_does_not_activate_reserved_allocation() {
        let mut store = LedgerCommitStore::default();
        store.commit(&ledger()).expect("initial ledger");
        let reservation = declaration();
        AllocationBootstrap::new(&mut store)
            .reserve_and_commit(&[reservation], &TestPolicy, Some(42))
            .expect("reservation commit");
        let snapshot = DeclarationSnapshot::new(vec![declaration()]).expect("snapshot");

        let err = AllocationBootstrap::new(&mut store)
            .validate_and_commit(snapshot, &RejectActivePolicy, Some(43))
            .expect_err("active validation must run");
        let recovered = store.recover().expect("recovered");

        assert!(matches!(
            err,
            BootstrapError::Validation(AllocationValidationError::Policy("active slot rejected"))
        ));
        assert_eq!(
            recovered.ledger().allocation_history().records()[0].state(),
            AllocationState::Reserved
        );
    }

    #[test]
    fn retire_and_commit_tombstones_through_protected_commit() {
        let mut store = LedgerCommitStore::default();
        store.commit(&ledger()).expect("initial ledger");
        let snapshot = DeclarationSnapshot::new(vec![declaration()]).expect("snapshot");
        let _pending = AllocationBootstrap::new(&mut store)
            .validate_and_commit(snapshot, &TestPolicy, Some(42))
            .expect("active commit");
        let retirement = AllocationRetirement::new(
            "app.users.v1",
            MemoryManagerSlot::new(100).expect("usable slot"),
        )
        .expect("retirement");

        let committed = AllocationBootstrap::new(&mut store)
            .retire_and_commit(&retirement, Some(43))
            .expect("retirement commit");

        assert_eq!(committed.current_generation, 2);
        assert_eq!(
            committed.allocation_history.records()[0].state(),
            AllocationState::Retired { generation: 2 }
        );
        assert_eq!(store.recover().unwrap().ledger(), &committed);
    }

    #[test]
    fn retire_and_commit_rejects_unknown_key_before_commit() {
        let mut store = LedgerCommitStore::default();
        store.commit(&ledger()).expect("initial ledger");
        let retirement = AllocationRetirement::new(
            "app.users.v1",
            MemoryManagerSlot::new(100).expect("usable slot"),
        )
        .expect("retirement");

        let err = AllocationBootstrap::new(&mut store)
            .retire_and_commit(&retirement, Some(43))
            .expect_err("unknown key");
        let recovered = store.recover().expect("recovered");

        assert!(matches!(err, BootstrapRetirementError::Retirement(_)));
        assert_eq!(recovered.current_generation(), 0);
        assert_eq!(recovered.ledger().allocation_history().records(), []);
    }
}