armdb 0.14.2

use super::*;
use crate::scope::Progress;
use crate::{ConstTree, DbError, FixedConfig, UniqueTargets};
use crate::{WriteHook, armour::Db};
use std::sync::Arc;

#[cfg(feature = "var-collections")]
use crate::test_faults::{POISON, PoisonCodec};

fn key(a: u64, b: u64) -> K3 {
    K3 { a, b, c: 0 }
}

crate::unique_index!(
    RawBulkClaimNamespace,
    key = [u8; 8],
    name = "raw_bulk_claim_ns"
);

struct RawBulkClaimHook {
    targets: UniqueTargets,
}

impl WriteHook<[u8; 8]> for RawBulkClaimHook {
    const CLAIMS: bool = true;

    fn on_write(&self, _key: &[u8; 8], _old: Option<&[u8]>, _new: Option<&[u8]>) {}

    fn claim_targets(&self) -> Option<&UniqueTargets> {
        Some(&self.targets)
    }
}

fn seed_prefix<H: ScopeHarness>(h: &H, a: u64, count: u64) -> Scope<'_, H::Scan<'_>, 1>
where
    for<'a> H::Scan<'a>: BulkScan,
{
    for b in 1..=count {
        h.insert(key(a, b));
    }
    h.scope::<1>(&key(a, 0))
}

#[test]
fn bulk_delete_up_to_zero_is_non_observing() {
    let dir = tempfile::tempdir().unwrap();
    let h = ConstTreeHarness::open(dir.path());
    let scope = seed_prefix(&h, 1, 1);

    assert_eq!(
        scope.delete_up_to(0).unwrap(),
        Progress {
            deleted: 0,
            done: false
        }
    );
    assert_eq!(scope.count().unwrap(), 1);
}

/// The break this catches: a raw byte family deletes rows that happen not to
/// claim anything before discovering that the source has a claim target and
/// must use a transaction for every bulk mutation.
#[test]
fn bulk_raw_claim_source_refuses_before_any_delete() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let namespace = db
        .open_unique_index::<RawBulkClaimNamespace>(FixedConfig::test())
        .unwrap();
    let tree = ConstTree::<[u8; 8], 8, RawBulkClaimHook>::open_hooked(
        dir.path().join("raw"),
        Config::test(),
        RawBulkClaimHook {
            targets: UniqueTargets::new().with(&namespace),
        },
    )
    .unwrap();
    let first = 1_u64.to_be_bytes();
    let second = 2_u64.to_be_bytes();
    tree.insert(&first, &[1; 8]).unwrap();
    tree.insert(&second, &[2; 8]).unwrap();

    assert!(matches!(
        tree.scope::<0>(&[0; 8]).delete_all(),
        Err(DbError::UniqueIndexRequiresTx { .. })
    ));
    assert_eq!(tree.get(&first), Some([1; 8]));
    assert_eq!(tree.get(&second), Some([2; 8]));
}

#[test]
fn bulk_exact_chunk_needs_the_extra_probe_before_done() {
    let dir = tempfile::tempdir().unwrap();
    let h = ConstTreeHarness::open(dir.path());
    let prefix = key(2, 0);
    for b in 1..=crate::scope::bulk::CHUNK as u64 {
        h.insert(key(2, b));
    }

    let first = h
        .scope::<1>(&prefix)
        .delete_up_to(crate::scope::bulk::CHUNK)
        .unwrap();
    assert_eq!(
        first,
        Progress {
            deleted: crate::scope::bulk::CHUNK,
            done: true
        }
    );
    assert_eq!(
        h.scope::<1>(&prefix)
            .delete_up_to(crate::scope::bulk::CHUNK)
            .unwrap(),
        Progress {
            deleted: 0,
            done: true
        }
    );
}

fn assert_runtime_bulk_operations<H: ScopeHarness>()
where
    for<'a> H::Scan<'a>: BulkScan + StrictRows,
{
    let dir = tempfile::tempdir().unwrap();
    let h = H::open(dir.path());
    let prefix = key(11, 0);
    for b in 1..=3 {
        h.insert(key(11, b));
    }

    let scope = h.scope::<1>(&prefix);
    assert_eq!(scope.count().unwrap(), 3);
    assert!(!scope.is_empty().unwrap());
    assert_eq!(scope.retain(|_, _| false).unwrap(), 3);
    assert!(h.scope::<1>(&prefix).is_empty().unwrap());

    for b in 1..=3 {
        h.insert(key(11, b));
    }
    assert_eq!(h.scope::<1>(&prefix).delete_up_to(1).unwrap().deleted, 1);
    assert_eq!(h.scope::<1>(&prefix).delete_all().unwrap(), 2);
    assert!(h.scope::<1>(&prefix).is_empty().unwrap());
}

/// The break this catches: a concrete family only keeps the bulk methods
/// nameable, while one of their public operation paths panics or is a no-op.
#[test]
fn bulk_public_operations_run_for_every_in_memory_family() {
    assert_runtime_bulk_operations::<ConstTreeHarness>();
    assert_runtime_bulk_operations::<TypedTreeHarness>();
    assert_runtime_bulk_operations::<ZeroTreeHarness>();
    #[cfg(feature = "var-collections")]
    assert_runtime_bulk_operations::<VarTreeHarness>();
    #[cfg(feature = "var-collections")]
    assert_runtime_bulk_operations::<VarTypedTreeHarness>();
    assert_runtime_bulk_operations::<ConstMapHarness>();
    assert_runtime_bulk_operations::<TypedMapHarness>();
    assert_runtime_bulk_operations::<ZeroMapHarness>();
    #[cfg(feature = "var-collections")]
    assert_runtime_bulk_operations::<VarMapHarness>();
    #[cfg(feature = "var-collections")]
    assert_runtime_bulk_operations::<VarTypedMapHarness>();
}

#[cfg(feature = "fjall")]
#[test]
fn bulk_public_operations_run_for_lsm_tree() {
    assert_runtime_bulk_operations::<LsmTreeHarness>();
}

#[test]
fn bulk_range_only_deletes_the_inner_cut() {
    let dir = tempfile::tempdir().unwrap();
    let h = ConstTreeHarness::open(dir.path());
    for b in 1..=4 {
        h.insert(key(1, b));
    }
    let prefix = key(1, 0);
    let lo = key(1, 2);
    let hi = key(1, 3);

    assert_eq!(
        h.scope::<1>(&prefix)
            .range::<2>(Bound::Included(&lo), Bound::Included(&hi))
            .delete_all()
            .unwrap(),
        2
    );
    assert_eq!(
        h.scope::<1>(&prefix)
            .keys()
            .map(|k| k.b)
            .collect::<Vec<_>>(),
        vec![1, 4]
    );
}

#[test]
fn bulk_count_and_empty_count_keys_not_rows_in_both_directions() {
    let dir = tempfile::tempdir().unwrap();
    let h = ConstMapHarness::open(dir.path());
    let prefix = key(9, 0);
    assert!(h.scope::<1>(&prefix).is_empty().unwrap());
    h.insert(key(9, 7));
    assert_eq!(h.scope::<1>(&prefix).desc().count().unwrap(), 1);
    assert!(!h.scope::<1>(&prefix).is_empty().unwrap());
    h.insert(key(10, 1));
    let empty_prefix = key(11, 0);
    assert_eq!(h.scope::<1>(&empty_prefix).count().unwrap(), 0);
    assert!(h.scope::<1>(&empty_prefix).is_empty().unwrap());

    let tree_dir = tempfile::tempdir().unwrap();
    let tree = ConstTreeHarness::open(tree_dir.path());
    for b in 1..=3 {
        tree.insert(key(8, b));
    }
    let tree_prefix = key(8, 0);
    assert_eq!(tree.scope::<1>(&tree_prefix).desc().count().unwrap(), 3);
    assert!(!tree.scope::<1>(&tree_prefix).desc().is_empty().unwrap());
}

#[test]
fn bulk_delete_up_to_handles_chunk_edges() {
    for count in [crate::scope::bulk::CHUNK - 1, crate::scope::bulk::CHUNK + 1] {
        let dir = tempfile::tempdir().unwrap();
        let h = ConstTreeHarness::open(dir.path());
        let scope = seed_prefix(&h, 3, count as u64);
        assert_eq!(
            scope.delete_up_to(crate::scope::bulk::CHUNK).unwrap(),
            Progress {
                deleted: count.min(crate::scope::bulk::CHUNK),
                done: count < crate::scope::bulk::CHUNK,
            }
        );
        assert_eq!(
            scope.count().unwrap(),
            count.saturating_sub(crate::scope::bulk::CHUNK)
        );
    }
}

#[test]
fn bulk_delete_all_resumes_in_both_directions() {
    for descending in [false, true] {
        let dir = tempfile::tempdir().unwrap();
        let h = ConstTreeHarness::open(dir.path());
        let prefix = key(4, 0);
        for b in 1..=(crate::scope::bulk::CHUNK as u64 + 1) {
            h.insert(key(4, b));
        }
        let scope = h.scope::<1>(&prefix);
        let deleted = if descending {
            scope.desc().delete_all().unwrap()
        } else {
            scope.asc().delete_all().unwrap()
        };
        assert_eq!(deleted, crate::scope::bulk::CHUNK + 1);
        assert!(h.scope::<1>(&prefix).is_empty().unwrap());
    }
}

#[test]
fn bulk_delete_all_is_idempotent_for_empty_scope() {
    let dir = tempfile::tempdir().unwrap();
    let h = ConstTreeHarness::open(dir.path());
    let prefix = key(5, 0);
    assert_eq!(h.scope::<1>(&prefix).delete_all().unwrap(), 0);
    assert_eq!(h.scope::<1>(&prefix).delete_all().unwrap(), 0);
}

#[test]
fn bulk_map_scans_a_key_in_the_final_shard() {
    let dir = tempfile::tempdir().unwrap();
    let h = ConstMapHarness::open(dir.path());
    let prefix = key(6, 0);
    let only_key = (1..)
        .map(|b| key(6, b))
        .find(|key| h.0.shard_for(key) == 2)
        .expect("a key assigned to the final map shard");
    h.insert(only_key);
    assert_eq!(h.scope::<1>(&prefix).count().unwrap(), 1);
    assert_eq!(h.scope::<1>(&prefix).delete_all().unwrap(), 1);
    assert!(h.scope::<1>(&prefix).is_empty().unwrap());
}

#[test]
fn bulk_retain_deletes_only_rejected_rows() {
    let dir = tempfile::tempdir().unwrap();
    let h = TypedTreeHarness::open(dir.path());
    let prefix = key(7, 0);
    for b in 1..=4 {
        h.insert(key(7, b));
    }

    assert_eq!(
        h.scope::<1>(&prefix)
            .retain(|_key, value| *value % 2 == 0)
            .unwrap(),
        2
    );
    assert_eq!(
        h.scope::<1>(&prefix)
            .keys()
            .map(|key| key.b)
            .collect::<Vec<_>>(),
        vec![2, 4]
    );
}

#[test]
fn bulk_retain_keeps_a_value_replaced_after_predicate() {
    let dir = tempfile::tempdir().unwrap();
    let tree = TypedTree::<K3, u64, U64Codec>::open(dir.path(), Config::test(), U64Codec).unwrap();
    let key = key(3, 1);
    tree.insert(&key, 1).unwrap();

    let removed = tree
        .scope::<0>(&K3::zeroed())
        .retain(|_, value| {
            tree.put(&key, 99).unwrap();
            *value % 2 == 0
        })
        .unwrap();

    assert_eq!(removed, 0);
    assert_eq!(tree.get(&key).map(|value| *value), Some(99));
}

/// The break this catches: a retain walk revisits keys inserted behind its
/// current cursor, so a predicate can extend the pass indefinitely.
#[test]
fn bulk_retain_leaves_a_key_inserted_behind_the_cursor() {
    let dir = tempfile::tempdir().unwrap();
    let tree = TypedTree::<K3, u64, U64Codec>::open(dir.path(), Config::test(), U64Codec).unwrap();
    let first = key(3, 1);
    let behind = key(3, 0);
    tree.insert(&first, 1).unwrap();

    assert_eq!(
        tree.scope::<1>(&key(3, 0))
            .retain(|_, _| {
                tree.insert(&behind, 2).unwrap();
                false
            })
            .unwrap(),
        1
    );
    assert_eq!(tree.get(&behind).map(|value| *value), Some(2));
}

#[test]
fn bulk_retain_allows_the_predicate_to_read_the_same_collection() {
    let dir = tempfile::tempdir().unwrap();
    let tree = TypedTree::<K3, u64, U64Codec>::open(dir.path(), Config::test(), U64Codec).unwrap();
    let key = key(4, 1);
    tree.insert(&key, 7).unwrap();

    let removed = tree
        .scope::<0>(&K3::zeroed())
        .retain(|scanned_key, _| {
            assert!(tree.get(scanned_key).is_some());
            *scanned_key != key
        })
        .unwrap();

    assert_eq!(removed, 1);
    assert!(!tree.contains(&key));
}

#[test]
fn bulk_retain_does_not_count_a_concurrently_removed_key() {
    use std::sync::{Arc, Barrier};

    let dir = tempfile::tempdir().unwrap();
    let tree = Arc::new(
        TypedTree::<K3, u64, U64Codec>::open(dir.path(), Config::test(), U64Codec).unwrap(),
    );
    let key = key(5, 1);
    tree.insert(&key, 1).unwrap();
    let barrier = Arc::new(Barrier::new(2));

    std::thread::scope(|threads| {
        let delete_tree = Arc::clone(&tree);
        let delete_barrier = Arc::clone(&barrier);
        threads.spawn(move || {
            delete_barrier.wait();
            delete_tree.delete(&key).unwrap();
            delete_barrier.wait();
        });

        let removed = tree
            .scope::<0>(&K3::zeroed())
            .retain(|_, _| {
                barrier.wait();
                barrier.wait();
                false
            })
            .unwrap();
        assert_eq!(removed, 0);
    });

    assert!(!tree.contains(&key));
}

#[cfg(feature = "fjall")]
#[test]
fn bulk_lsm_scope_operations_are_strict_and_mutate_normally() {
    let dir = tempfile::tempdir().unwrap();
    let h = LsmTreeHarness::open(dir.path());
    let prefix = key(8, 0);
    for b in 1..=4 {
        h.insert(key(8, b));
    }

    assert_eq!(h.scope::<1>(&prefix).count().unwrap(), 4);
    assert!(!h.scope::<1>(&prefix).is_empty().unwrap());
    assert_eq!(
        h.scope::<1>(&prefix)
            .retain(|_, value| value.v % 2 == 0)
            .unwrap(),
        2
    );
    assert_eq!(h.scope::<1>(&prefix).delete_all().unwrap(), 2);
    assert!(h.scope::<1>(&prefix).is_empty().unwrap());
}

#[cfg(feature = "fjall")]
#[test]
fn bulk_lsm_key_guard_error_aborts_count_and_delete() {
    let dir = tempfile::tempdir().unwrap();
    let h = LsmTreeHarness::open(dir.path());
    let prefix = key(9, 0);
    h.insert(key(9, 1));

    h.force_next_key_guard_error();
    assert!(h.scope::<1>(&prefix).count().is_err());

    h.force_next_key_guard_error();
    assert!(h.scope::<1>(&prefix).delete_all().is_err());
    assert_eq!(h.scope::<1>(&prefix).count().unwrap(), 1);
}

#[cfg(feature = "var-collections")]
#[test]
fn bulk_retain_stops_on_an_undecodable_var_typed_row() {
    let dir = tempfile::tempdir().unwrap();
    let tree: VarTypedTree<[u8; 8], u64, PoisonCodec> =
        VarTypedTree::open(dir.path(), Config::test(), PoisonCodec).unwrap();
    tree.put(&1_u64.to_be_bytes(), &1).unwrap();
    tree.put(&2_u64.to_be_bytes(), &POISON).unwrap();

    assert_eq!(tree.scope::<0>(&[0; 8]).count().unwrap(), 2);

    assert!(matches!(
        tree.scope::<0>(&[0; 8]).asc().retain(|_, _| false),
        Err(DbError::CorruptedEntry { .. })
    ));
    assert!(!tree.contains(&1_u64.to_be_bytes()));
    assert!(tree.contains(&2_u64.to_be_bytes()));
}