icydb-core 0.264.10

use super::*;
use crate::{
    db::{
        data::DecodedDataStoreKey,
        journal::JournalSequence,
        key_taxonomy::{PrimaryKeyComponent, PrimaryKeyValue},
        positioned_overlay::{JournalOverlayPosition, PositionedOverlayRetirement},
        registry::StoreAllocationIdentity,
    },
    testing::{test_memory, test_memory_with_backing},
};
use ic_memory::ic_stable_structures::{Memory, VectorMemory};
use std::ops::Bound;

const SINGLE_MEMORY_MANAGER_BUCKET_PAGES: u64 = 128;

fn raw_key(entity: u64, id: u64) -> RawDataStoreKey {
    DecodedDataStoreKey::new(
        EntityTag::new(entity),
        &PrimaryKeyValue::Scalar(PrimaryKeyComponent::Nat64(id)),
    )
    .to_raw()
    .expect("test data key should encode")
}

fn raw_row(value: u8) -> RawRow {
    RawRow::try_new(vec![value]).expect("test raw row should be bounded")
}

fn overlay_position(sequence: u64) -> JournalOverlayPosition {
    JournalOverlayPosition::new(
        StoreAllocationIdentity::new(230, "test::data"),
        JournalSequence::new(sequence),
    )
}

fn seed_store(memory_id: u8, entries: &[(u64, u64, u8)]) -> DataStore {
    let mut store = DataStore::init_journaled(test_memory(memory_id));
    for (entity, id, row) in entries {
        store.insert_raw_for_test(raw_key(*entity, *id), raw_row(*row));
    }
    store
}

fn seed_heap_store(entries: &[(u64, u64, u8)]) -> DataStore {
    let mut store = DataStore::init_heap();
    for (entity, id, row) in entries {
        store.insert_raw_for_test(raw_key(*entity, *id), raw_row(*row));
    }
    store
}

#[test]
fn data_store_read_borrows_heap_and_live_rows_but_owns_stable_rows() {
    let heap = seed_heap_store(&[(1, 1, 11)]);
    assert!(matches!(
        heap.read(&raw_key(1, 1)),
        StoredRowRead::Borrowed(row) if row.as_bytes() == [11]
    ));

    let mut journaled = DataStore::init_journaled(test_memory(228));
    journaled
        .fold_recovered_journal_put(raw_key(1, 1), raw_row(11))
        .expect("canonical row should fold");
    assert!(matches!(
        journaled.read(&raw_key(1, 1)),
        StoredRowRead::Owned(row) if row.as_bytes() == [11]
    ));

    journaled
        .apply_recovered_journal_put(raw_key(1, 1), raw_row(22))
        .expect("live override should apply");
    assert!(matches!(
        journaled.read(&raw_key(1, 1)),
        StoredRowRead::Borrowed(row) if row.as_bytes() == [22]
    ));
}

#[test]
fn positioned_data_overlay_retires_exactly_across_delete_and_reinsert() {
    let key = raw_key(1, 1);
    let mut store = DataStore::init_journaled(test_memory(230));
    store
        .fold_recovered_journal_put(key.clone(), raw_row(11))
        .expect("canonical row should seed");
    store.rebuild_entity_cardinality_from_entries();

    store
        .publish_positioned_journal_entry(key.clone(), Some(raw_row(22)), overlay_position(1))
        .expect("first positioned put should publish");
    store
        .publish_positioned_journal_entry(key.clone(), None, overlay_position(2))
        .expect("later positioned delete should publish");
    let prepared = store
        .prepare_position_retirement([key.clone()], overlay_position(1))
        .expect("older batch retirement should preflight");
    assert_eq!(
        prepared.entries,
        [(key.clone(), PositionedOverlayRetirement::Superseded)]
    );
    store
        .fold_recovered_journal_put(key.clone(), raw_row(22))
        .expect("first batch should become canonical");
    store.apply_prepared_position_retirement(prepared);
    assert!(
        store.get(&key).is_none(),
        "newer tombstone must remain visible"
    );

    store
        .publish_positioned_journal_entry(key.clone(), Some(raw_row(33)), overlay_position(3))
        .expect("reinsert should supersede the tombstone");
    let prepared = store
        .prepare_position_retirement([key.clone()], overlay_position(2))
        .expect("delete retirement should preserve the reinsert");
    assert_eq!(
        prepared.entries,
        [(key.clone(), PositionedOverlayRetirement::Superseded)]
    );
    store
        .fold_recovered_journal_delete(&key)
        .expect("delete batch should become canonical");
    store.apply_prepared_position_retirement(prepared);
    assert_eq!(
        store
            .get(&key)
            .expect("reinsert must remain visible")
            .as_bytes(),
        [33],
    );

    let prepared = store
        .prepare_position_retirement([key.clone()], overlay_position(3))
        .expect("newest batch should retire exactly");
    assert_eq!(
        prepared.entries,
        [(key.clone(), PositionedOverlayRetirement::Exact)]
    );
    store
        .fold_recovered_journal_put(key.clone(), raw_row(33))
        .expect("reinsert batch should become canonical");
    store.apply_prepared_position_retirement(prepared);
    assert_eq!(
        store
            .get(&key)
            .expect("canonical reinsert should remain")
            .as_bytes(),
        [33],
    );
    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(1));
}

#[test]
fn data_store_tiny_fold_stays_within_one_memory_manager_bucket() {
    let memory = VectorMemory::default();
    let data_memory = test_memory_with_backing(23, memory.clone());
    // Attribute only the store's growth, excluding the bootstrapped ledger.
    let bootstrap_pages = memory.size();
    let mut store = DataStore::init_journaled(data_memory);

    store
        .fold_recovered_journal_put(raw_key(1, 1), raw_row(7))
        .expect("tiny folded row should persist");

    assert!(
        memory.size() - bootstrap_pages <= SINGLE_MEMORY_MANAGER_BUCKET_PAGES,
        "tiny data fold should not allocate extra MemoryManager buckets; pages={}",
        memory.size()
    );
}

fn collect_keys(store: &DataStore) -> Vec<RawDataStoreKey> {
    let mut keys = Vec::new();
    let _: Result<(), Infallible> = store.visit_entries(|key, _row| {
        keys.push(key.clone());
        Ok(StoreVisit::Continue)
    });
    keys
}

#[test]
fn data_store_visit_entries_preserves_storage_key_order() {
    let store = seed_store(221, &[(2, 1, 21), (1, 3, 13), (1, 1, 11), (1, 2, 12)]);

    let mut expected = vec![raw_key(2, 1), raw_key(1, 3), raw_key(1, 1), raw_key(1, 2)];
    expected.sort();

    assert_eq!(collect_keys(&store), expected);
}

#[test]
fn data_store_visit_range_preserves_raw_key_bounds() {
    let entries = [(1, 1, 11), (1, 2, 12), (1, 3, 13), (1, 4, 14)];
    let mut canonical = DataStore::init_journaled(test_memory(226));
    for (entity, id, value) in entries {
        canonical
            .fold_recovered_journal_put(raw_key(entity, id), raw_row(value))
            .expect("canonical seed should fold");
    }
    for store in [
        seed_heap_store(&entries),
        seed_store(222, &entries),
        canonical,
    ] {
        let lower = raw_key(1, 2);
        let upper = raw_key(1, 4);
        let bounds = (Bound::Included(&lower), Bound::Excluded(&upper));
        let mut visited = Vec::new();
        store
            .visit_range(bounds, |key, row| {
                visited.push((key.clone(), row.as_bytes()[0]));
                Ok::<_, Infallible>(StoreVisit::Continue)
            })
            .unwrap();
        assert_eq!(visited, vec![(raw_key(1, 2), 12), (raw_key(1, 3), 13)]);

        let mut keys = Vec::new();
        store
            .visit_key_range(bounds, |key| {
                keys.push(key.clone());
                Ok::<_, Infallible>(StoreVisit::Continue)
            })
            .unwrap();
        assert_eq!(keys, vec![raw_key(1, 2), raw_key(1, 3)]);
        keys.clear();
        store
            .visit_key_range_rev(bounds, |key| {
                keys.push(key.clone());
                Ok::<_, Infallible>(StoreVisit::Continue)
            })
            .unwrap();
        assert_eq!(keys, vec![raw_key(1, 3), raw_key(1, 2)]);
    }
}

#[test]
fn data_store_row_preflight_can_stop_before_visiting_the_payload() {
    let store = seed_store(229, &[(1, 1, 11), (1, 2, 12), (1, 3, 13)]);
    let mut preflighted = Vec::new();
    let mut visited = Vec::new();

    let exhausted = store
        .try_visit_range_with_row_preflight(
            (
                Bound::Included(raw_key(1, 1)),
                Bound::Included(raw_key(1, 3)),
            ),
            |key| {
                preflighted.push(key.clone());
                Ok::<_, Infallible>(if key == &raw_key(1, 2) {
                    StoreVisit::Stop
                } else {
                    StoreVisit::Continue
                })
            },
            |key, row| {
                visited.push((key.clone(), row.as_bytes()[0]));
                Ok::<_, Infallible>(StoreVisit::Continue)
            },
        )
        .expect("stable-only range should use the fused visitor");

    assert_eq!(exhausted, Some(false));
    assert_eq!(preflighted, vec![raw_key(1, 1), raw_key(1, 2)]);
    assert_eq!(visited, vec![(raw_key(1, 1), 11)]);
}

#[test]
fn data_store_entity_cardinality_tracks_heap_writes() {
    let mut store = seed_heap_store(&[(1, 1, 11), (1, 2, 12), (2, 1, 21)]);

    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(2));
    assert_eq!(store.exact_entity_count(EntityTag::new(2)), Some(1));
    assert_eq!(store.exact_entity_count(EntityTag::new(3)), Some(0));

    store.insert_raw_for_test(raw_key(1, 2), raw_row(22));
    assert_eq!(
        store.exact_entity_count(EntityTag::new(1)),
        Some(2),
        "replacing a row must not change entity cardinality",
    );

    store.insert_raw_for_test(raw_key(1, 3), raw_row(13));
    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(3));

    let removed = store.remove(&raw_key(1, 1));
    assert!(removed.is_some());
    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(2));

    let removed = store.remove(&raw_key(1, 99));
    assert!(removed.is_none());
    assert_eq!(
        store.exact_entity_count(EntityTag::new(1)),
        Some(2),
        "removing a missing row must not change entity cardinality",
    );
}

#[test]
fn data_store_entity_cardinality_tracks_journaled_overlay_writes() {
    let mut store = DataStore::init_journaled(test_memory(226));
    store
        .fold_recovered_journal_put(raw_key(1, 1), raw_row(11))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 3), raw_row(13))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(2, 1), raw_row(21))
        .expect("canonical seed should fold");
    store.rebuild_entity_cardinality_from_entries();
    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(2));
    assert_eq!(store.exact_entity_count(EntityTag::new(2)), Some(1));

    store
        .apply_recovered_journal_put(raw_key(1, 0), raw_row(10))
        .expect("live put should apply");
    store
        .apply_recovered_journal_put(raw_key(1, 3), raw_row(33))
        .expect("live replacement should apply");
    store.remove(&raw_key(1, 1));
    assert_eq!(
        store.exact_entity_count(EntityTag::new(1)),
        Some(2),
        "one insert, one replacement, and one delete should leave two entity rows",
    );

    store
        .reset_journaled_live_projection()
        .expect("projection reset should succeed");
    assert_eq!(
        store.exact_entity_count(EntityTag::new(1)),
        None,
        "projection reset must not scan canonical rows to rebuild optional counts",
    );
}

#[test]
fn data_store_entity_cardinality_ignores_fold_updates_hidden_by_live_overlay() {
    let mut store = DataStore::init_journaled(test_memory(227));
    store
        .fold_recovered_journal_put(raw_key(1, 1), raw_row(11))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 2), raw_row(12))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 3), raw_row(13))
        .expect("canonical seed should fold");
    store.rebuild_entity_cardinality_from_entries();

    store
        .apply_recovered_journal_put(raw_key(1, 2), raw_row(22))
        .expect("live replacement should apply");
    store.remove(&raw_key(1, 3));
    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(2));

    store
        .fold_recovered_journal_delete(&raw_key(1, 2))
        .expect("hidden canonical delete should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 3), raw_row(33))
        .expect("hidden canonical put should fold");

    assert_eq!(
        store.exact_entity_count(EntityTag::new(1)),
        Some(2),
        "folding canonical rows hidden by live rows or tombstones must not change visible counts",
    );
}

#[test]
fn journaled_data_store_reopens_without_materializing_entity_cardinality() {
    let memory = test_memory(230);
    let mut store = DataStore::init_journaled(memory.clone());
    store
        .fold_recovered_journal_put(raw_key(1, 1), raw_row(11))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 2), raw_row(12))
        .expect("canonical seed should fold");
    drop(store);

    let mut reopened = DataStore::init_journaled(memory);

    assert!(reopened.contains(&raw_key(1, 1)));
    assert!(reopened.contains(&raw_key(1, 2)));
    assert_eq!(
        reopened.exact_entity_count(EntityTag::new(1)),
        None,
        "startup must leave optional cardinality unavailable without scanning stable rows",
    );
    assert_eq!(
        reopened.exact_entity_cardinality_delta(EntityTag::new(1)),
        Some(0),
    );
    reopened
        .apply_recovered_journal_put(raw_key(1, 3), raw_row(13))
        .expect("post-watermark live row should apply");
    assert_eq!(
        reopened.exact_entity_cardinality_delta(EntityTag::new(1)),
        Some(1),
    );
    reopened
        .fold_recovered_journal_put(raw_key(1, 3), raw_row(13))
        .expect("the matching canonical transition should fold");
    assert_eq!(
        reopened.exact_entity_cardinality_delta(EntityTag::new(1)),
        Some(0),
        "canonical fold must consume only its exact overlay contribution",
    );
}

#[test]
fn empty_journaled_data_store_retains_exact_zero_cardinality_without_scanning() {
    let memory = test_memory(231);
    let mut store = DataStore::init_journaled(memory.clone());

    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(0));
    store
        .reset_journaled_live_projection()
        .expect("empty projection reset should succeed");
    assert_eq!(store.exact_entity_count(EntityTag::new(1)), Some(0));
    drop(store);

    let reopened = DataStore::init_journaled(memory);
    assert_eq!(reopened.exact_entity_count(EntityTag::new(1)), Some(0));
}

#[test]
fn data_store_entity_cardinality_fails_closed_for_malformed_raw_keys() {
    let mut store = seed_heap_store(&[(1, 1, 11)]);
    store.insert_raw_for_test(
        RawDataStoreKey::from_persisted_bytes(vec![1, 2, 3]),
        raw_row(99),
    );

    assert_eq!(
        store.exact_entity_count(EntityTag::new(1)),
        None,
        "malformed raw keys should disable entity-cardinality metadata",
    );
}

#[test]
fn data_store_visit_entries_can_stop_without_error() {
    let store = seed_store(224, &[(1, 1, 11), (1, 2, 12), (1, 3, 13)]);

    let mut visited = Vec::new();
    let _: Result<(), Infallible> = store.visit_entries(|key, _row| {
        visited.push(key.clone());
        Ok(if visited.len() == 2 {
            StoreVisit::Stop
        } else {
            StoreVisit::Continue
        })
    });

    assert_eq!(visited, vec![raw_key(1, 1), raw_key(1, 2)]);
}

#[test]
fn heap_data_store_preserves_order_bounds_and_early_stop() {
    let store = seed_heap_store(&[(2, 1, 21), (1, 3, 13), (1, 1, 11), (1, 2, 12)]);

    let mut expected = vec![raw_key(2, 1), raw_key(1, 3), raw_key(1, 1), raw_key(1, 2)];
    expected.sort();
    assert_eq!(collect_keys(&store), expected);

    let mut ranged = Vec::new();
    let _: Result<(), Infallible> = store.visit_range(
        (
            Bound::Included(raw_key(1, 1)),
            Bound::Excluded(raw_key(1, 3)),
        ),
        |key, row| {
            ranged.push((key.clone(), row.as_bytes()[0]));
            Ok(StoreVisit::Continue)
        },
    );
    assert_eq!(ranged, vec![(raw_key(1, 1), 11), (raw_key(1, 2), 12)]);

    let mut ranged_keys = Vec::new();
    let _: Result<(), Infallible> = store.visit_key_range(
        (
            Bound::Included(raw_key(1, 1)),
            Bound::Excluded(raw_key(1, 3)),
        ),
        |key| {
            ranged_keys.push(key.clone());
            Ok(StoreVisit::Continue)
        },
    );
    assert_eq!(ranged_keys, vec![raw_key(1, 1), raw_key(1, 2)]);

    let mut stopped = Vec::new();
    let _: Result<(), Infallible> = store.visit_entries(|key, _| {
        stopped.push(key.clone());
        Ok(if stopped.len() == 2 {
            StoreVisit::Stop
        } else {
            StoreVisit::Continue
        })
    });
    assert_eq!(stopped, vec![raw_key(1, 1), raw_key(1, 2)]);
}

#[test]
fn journaled_mixed_data_range_traversal_streams_without_snapshot() {
    let mut store = DataStore::init_journaled(test_memory(225));
    store
        .fold_recovered_journal_put(raw_key(1, 1), raw_row(11))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 3), raw_row(13))
        .expect("canonical seed should fold");
    store
        .fold_recovered_journal_put(raw_key(1, 5), raw_row(15))
        .expect("canonical seed should fold");
    store
        .apply_recovered_journal_put(raw_key(1, 0), raw_row(10))
        .expect("live put should apply");
    store
        .apply_recovered_journal_put(raw_key(1, 4), raw_row(14))
        .expect("live put should apply");
    store
        .apply_recovered_journal_put(raw_key(1, 5), raw_row(55))
        .expect("live override should apply");
    store.remove(&raw_key(1, 1));

    let mut asc = Vec::new();
    let _: Result<(), Infallible> = store.visit_range(
        (
            Bound::Included(raw_key(1, 0)),
            Bound::Included(raw_key(1, 5)),
        ),
        |key, row| {
            asc.push((key.clone(), row.as_bytes()[0]));
            Ok(if asc.len() == 2 {
                StoreVisit::Stop
            } else {
                StoreVisit::Continue
            })
        },
    );
    assert_eq!(asc, vec![(raw_key(1, 0), 10), (raw_key(1, 3), 13)]);

    let mut desc_keys = Vec::new();
    let _: Result<(), Infallible> = store.visit_key_range_rev(
        (
            Bound::Included(raw_key(1, 0)),
            Bound::Included(raw_key(1, 5)),
        ),
        |key| {
            desc_keys.push(key.clone());
            Ok(StoreVisit::Continue)
        },
    );
    assert_eq!(
        desc_keys,
        vec![raw_key(1, 5), raw_key(1, 4), raw_key(1, 3), raw_key(1, 0),],
        "key-only traversal must merge live overrides and tombstones without reading row values",
    );

    let checkpoint = raw_key(1, 1);
    let mut canonical_rows = Vec::new();
    store
        .visit_canonical_entries_after(Some(&checkpoint), |key, row| {
            canonical_rows.push((key.clone(), row.as_bytes()[0]));
            Ok(false)
        })
        .unwrap();
    assert_eq!(
        canonical_rows,
        vec![(raw_key(1, 3), 13), (raw_key(1, 5), 15)]
    );
}