kvdb-overlay/tests/tree_overlay_state.rs

/* This file is part of DarkFi (https://dark.fi)
 *
 * Copyright (C) 2026-2026 Dyne.org foundation
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

//! Simulate the creation of a [`TreeOverlay`] on top of a [`Tree`]
//! instance, and perform diffs and writes to verify overlay's cache
//! diff functionality.

use kvdb_overlay::{Database, Result, TreeOverlay};

const TREE: &str = "_tree";

#[test]
fn tree_overlay_state() -> Result<()> {
    // Initialize database
    let (db, _folder) = Database::open_temp()?;

    // Initialize tree with some values and its overlay
    let tree = db.open_tree_default(TREE)?;
    tree.insert(b"key_a", b"val_a")?;
    let mut overlay = TreeOverlay::new(&tree);
    assert!(!overlay.is_empty()?);

    // Make a vector to keep track of changes
    let mut sequence = vec![];

    // Perform some changes and grab their differences
    overlay.insert(b"key_b", b"val_b")?;
    sequence.push(overlay.diff(&sequence)?);

    overlay.insert(b"key_b", b"val_bb")?;
    overlay.remove(b"key_a")?;
    sequence.push(overlay.diff(&sequence)?);

    overlay.insert(b"key_a", b"val_a")?;
    overlay.remove(b"key_b")?;
    overlay.insert(b"key_c", b"val_c")?;
    sequence.push(overlay.diff(&sequence)?);

    // Verify overlay has the correct state
    assert_eq!(overlay.state.cache.len(), 2);
    assert_eq!(
        overlay.state.cache.get(b"key_a".as_slice()),
        Some(&b"val_a".to_vec())
    );
    assert_eq!(
        overlay.state.cache.get(b"key_c".as_slice()),
        Some(&b"val_c".to_vec())
    );
    assert_eq!(overlay.state.removed.len(), 1);
    assert_eq!(
        overlay.state.removed.get(b"key_b".as_slice()),
        Some(&b"key_b".to_vec())
    );

    // Verify diffs sequence is correct
    assert_eq!(sequence.len(), 3);

    assert_eq!(sequence[0].cache.len(), 1);
    assert_eq!(
        sequence[0].cache.get(b"key_b".as_slice()),
        Some(&(None, b"val_b".to_vec()))
    );
    assert!(sequence[0].removed.is_empty());
    assert_eq!(sequence[0], sequence[0].inverse().inverse());

    assert_eq!(sequence[1].cache.len(), 1);
    assert_eq!(
        sequence[1].cache.get(b"key_b".as_slice()),
        Some(&(Some(b"val_b".to_vec()), b"val_bb".into()))
    );
    assert_eq!(sequence[1].removed.len(), 1);
    assert_eq!(
        sequence[1].removed.get(b"key_a".as_slice()),
        Some(&b"val_a".to_vec())
    );
    assert_eq!(sequence[1], sequence[1].inverse().inverse());

    assert_eq!(sequence[2].cache.len(), 2);
    assert_eq!(
        sequence[2].cache.get(b"key_a".as_slice()),
        Some(&(None, b"val_a".to_vec()))
    );
    assert_eq!(
        sequence[2].cache.get(b"key_c".as_slice()),
        Some(&(None, b"val_c".to_vec()))
    );
    assert_eq!(sequence[2].removed.len(), 1);
    assert_eq!(
        sequence[2].removed.get(b"key_b".as_slice()),
        Some(&b"val_bb".to_vec())
    );
    assert_eq!(sequence[2], sequence[2].inverse().inverse());

    // Now we are going to apply each diff and check that the database
    // has been mutated accordingly
    db.write_tree_overlays_diff(&tree, &sequence[0], false)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, Some(b"val_b".into()));
    overlay.remove_diff(&sequence[0]);

    db.write_tree_overlays_diff(&tree, &sequence[1], false)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 1);
    assert_eq!(tree.get(b"key_a")?, None);
    assert_eq!(tree.get(b"key_b")?, Some(b"val_bb".into()));
    overlay.remove_diff(&sequence[1]);

    // Since we removed the diffs, current overlay diff must be
    // the same as the last diff in the sequence
    let diff = overlay.diff(&[])?;
    assert_eq!(diff, sequence[2]);
    // Therefore we can safely use its batch
    db.write_tree_overlays_changes(&[&overlay])?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, None);
    assert_eq!(tree.get(b"key_c")?, Some(b"val_c".into()));
    overlay.remove_diff(&sequence[2]);

    // Since we removed everything, current overlay must not have
    // diffs over the tree, therefore its safe to keep using it
    let diff = overlay.diff(&[])?;
    assert!(diff.cache.is_empty());
    assert!(diff.removed.is_empty());

    // We are going to make some changes that we want to revert
    // using the corresponding diff
    overlay.insert(b"key_a", b"val_aa")?;
    overlay.insert(b"key_b", b"val_b")?;
    overlay.remove(b"key_c")?;

    // Grab the diff, apply it and verify tree state
    let diff = overlay.diff(&[])?;
    db.write_tree_overlays_diff(&tree, &diff, false)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_aa".into()));
    assert_eq!(tree.get(b"key_b")?, Some(b"val_b".into()));
    assert_eq!(tree.get(b"key_c")?, None);

    // Now we grab the diff revert batch, apply it and verity tree state
    db.write_tree_overlays_diff(&tree, &diff, true)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, None);
    assert_eq!(tree.get(b"key_c")?, Some(b"val_c".into()));

    // Now we are going to revert the diffs sequence going backwards
    // and verify tree state mutates accordingly
    db.write_tree_overlays_diff(&tree, &sequence[2], true)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 1);
    assert_eq!(tree.get(b"key_a")?, None);
    assert_eq!(tree.get(b"key_b")?, Some(b"val_bb".into()));

    db.write_tree_overlays_diff(&tree, &sequence[1], true)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, Some(b"val_b".into()));

    db.write_tree_overlays_diff(&tree, &sequence[0], true)?;
    db.flush_default_mode()?;

    // Tree has now reverted to its original state
    assert_eq!(tree.len()?, 1);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));

    Ok(())
}

#[test]
fn tree_overlay_rebuild_state() -> Result<()> {
    // Initialize database
    let (db, _folder) = Database::open_temp()?;

    // Initialize tree with some values and its overlay
    let tree = db.open_tree_default(TREE)?;
    tree.insert(b"key_a", b"val_a")?;
    let mut overlay = TreeOverlay::new(&tree);
    assert!(!overlay.is_empty()?);

    // Make two vectors to keep track of changes
    let mut sequence = vec![];
    let mut state_sequence = vec![];

    // Perform some changes and grab their differences
    overlay.insert(b"key_b", b"val_b")?;
    sequence.push(overlay.diff(&sequence)?);
    state_sequence.push(overlay.clone());

    overlay.insert(b"key_b", b"val_bb")?;
    overlay.remove(b"key_a")?;
    sequence.push(overlay.diff(&sequence)?);
    state_sequence.push(overlay.clone());

    overlay.insert(b"key_a", b"val_a")?;
    overlay.remove(b"key_b")?;
    overlay.insert(b"key_c", b"val_c")?;
    sequence.push(overlay.diff(&sequence)?);

    // Create a different overlay to rebuild
    // the previous one using the changes sequence
    let mut overlay2 = TreeOverlay::new(&tree);
    assert!(!overlay2.is_empty()?);

    // Add each diff from the sequence and verify
    // overlay has been mutated accordingly
    overlay2.add_diff(&sequence[0]);
    assert_eq!(overlay2.state.cache.len(), 1);
    assert_eq!(
        overlay2.state.cache.get(b"key_b".as_slice()),
        Some(&b"val_b".to_vec())
    );
    assert!(overlay2.state.removed.is_empty());
    assert_eq!(state_sequence[0].state, overlay2.state);

    overlay2.add_diff(&sequence[1]);
    assert_eq!(overlay2.state.cache.len(), 1);
    assert_eq!(
        overlay2.state.cache.get(b"key_b".as_slice()),
        Some(&b"val_bb".to_vec())
    );
    assert_eq!(overlay2.state.removed.len(), 1);
    assert_eq!(
        overlay2.state.removed.get(b"key_a".as_slice()),
        Some(&b"key_a".to_vec())
    );
    assert_eq!(state_sequence[1].state, overlay2.state);

    overlay2.add_diff(&sequence[2]);
    assert_eq!(overlay2.state.cache.len(), 2);
    assert_eq!(
        overlay2.state.cache.get(b"key_a".as_slice()),
        Some(&b"val_a".to_vec())
    );
    assert_eq!(
        overlay2.state.cache.get(b"key_c".as_slice()),
        Some(&b"val_c".to_vec())
    );
    assert_eq!(overlay2.state.removed.len(), 1);
    assert_eq!(
        overlay2.state.removed.get(b"key_b".as_slice()),
        Some(&b"key_b".to_vec())
    );
    assert_eq!(overlay.state, overlay2.state);

    // Now we are going to apply each diff and check that the database
    // has been mutated accordingly
    db.write_tree_overlays_diff(&tree, &sequence[0], false)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, Some(b"val_b".into()));
    overlay.remove_diff(&sequence[0]);

    db.write_tree_overlays_diff(&tree, &sequence[1], false)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 1);
    assert_eq!(tree.get(b"key_a")?, None);
    assert_eq!(tree.get(b"key_b")?, Some(b"val_bb".into()));
    overlay.remove_diff(&sequence[1]);

    // Since we removed the diffs, current overlay diff must be
    // the same as the last diff in the sequence
    let diff = overlay.diff(&[])?;
    assert_eq!(diff, sequence[2]);
    // Therefore we can safely use its batch
    db.write_tree_overlays_changes(&[&overlay])?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, None);
    assert_eq!(tree.get(b"key_c")?, Some(b"val_c".into()));
    overlay.remove_diff(&sequence[2]);

    // Since we removed everything, current overlay must not have
    // diffs over the tree, therefore its safe to keep using it
    let diff = overlay.diff(&[])?;
    assert!(diff.cache.is_empty());
    assert!(diff.removed.is_empty());

    // Now we are going to add all the inverse diffs in the overlay,
    // in reverse
    overlay.add_diff(&sequence[2].inverse());
    assert_eq!(overlay.state.cache.len(), 1);
    assert_eq!(
        overlay.state.cache.get(b"key_b".as_slice()),
        Some(&b"val_bb".to_vec())
    );
    assert_eq!(overlay.state.removed.len(), 2);
    assert_eq!(
        overlay.state.removed.get(b"key_a".as_slice()),
        Some(&b"key_a".to_vec())
    );
    assert_eq!(
        overlay.state.removed.get(b"key_c".as_slice()),
        Some(&b"key_c".to_vec())
    );

    overlay.add_diff(&sequence[1].inverse());
    assert_eq!(overlay.state.cache.len(), 2);
    assert_eq!(
        overlay.state.cache.get(b"key_a".as_slice()),
        Some(&b"val_a".to_vec())
    );
    assert_eq!(
        overlay.state.cache.get(b"key_b".as_slice()),
        Some(&b"val_b".to_vec())
    );
    assert_eq!(overlay.state.removed.len(), 1);
    assert_eq!(
        overlay.state.removed.get(b"key_c".as_slice()),
        Some(&b"key_c".to_vec())
    );

    overlay.add_diff(&sequence[0].inverse());
    assert_eq!(overlay.state.cache.len(), 1);
    assert_eq!(
        overlay.state.cache.get(b"key_a".as_slice()),
        Some(&b"val_a".to_vec())
    );
    assert_eq!(overlay.state.removed.len(), 2);
    assert_eq!(
        overlay.state.removed.get(b"key_b".as_slice()),
        Some(&b"key_b".to_vec())
    );
    assert_eq!(
        overlay.state.removed.get(b"key_c".as_slice()),
        Some(&b"key_c".to_vec())
    );

    // Now we are going to apply the overlay and verify that the
    // Tree has now reverted to its original state
    db.write_tree_overlays_changes(&[&overlay])?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 1);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));

    Ok(())
}

#[test]
fn tree_overlay_clear_state() -> Result<()> {
    // Initialize database
    let (db, _folder) = Database::open_temp()?;

    // Initialize tree with some values and its overlay
    let tree = db.open_tree_default(TREE)?;
    tree.insert(b"key_a", b"val_a")?;
    let mut overlay = TreeOverlay::new(&tree);
    assert!(!overlay.is_empty()?);

    // Make a vector to keep track of changes
    let mut sequence = vec![];

    // Perform some changes and grab their differences
    overlay.insert(b"key_b", b"val_b")?;
    sequence.push(overlay.diff(&sequence)?);

    overlay.clear()?;
    sequence.push(overlay.diff(&sequence)?);

    overlay.insert(b"key_a", b"val_a")?;
    assert!(overlay.remove(b"key_b").is_err());
    overlay.insert(b"key_c", b"val_c")?;
    sequence.push(overlay.diff(&sequence)?);

    // Verify overlay has the correct state
    assert_eq!(overlay.state.cache.len(), 2);
    assert_eq!(
        overlay.state.cache.get(b"key_a".as_slice()),
        Some(&b"val_a".into())
    );
    assert_eq!(
        overlay.state.cache.get(b"key_c".as_slice()),
        Some(&b"val_c".into())
    );
    assert!(overlay.state.removed.is_empty());

    // Verify diffs sequence is correct
    assert_eq!(sequence.len(), 3);

    assert_eq!(sequence[0].cache.len(), 1);
    assert_eq!(
        sequence[0].cache.get(b"key_b".as_slice()),
        Some(&(None, b"val_b".to_vec()))
    );
    assert!(sequence[0].removed.is_empty());
    assert_eq!(sequence[0], sequence[0].inverse().inverse());

    assert!(sequence[1].cache.is_empty());
    assert_eq!(sequence[1].removed.len(), 2);
    assert_eq!(
        sequence[1].removed.get(b"key_a".as_slice()),
        Some(&b"val_a".to_vec())
    );
    assert_eq!(
        sequence[1].removed.get(b"key_b".as_slice()),
        Some(&b"val_b".to_vec())
    );
    assert_eq!(sequence[1], sequence[1].inverse().inverse());

    assert_eq!(sequence[2].cache.len(), 2);
    assert_eq!(
        sequence[2].cache.get(b"key_a".as_slice()),
        Some(&(None, b"val_a".to_vec()))
    );
    assert_eq!(
        sequence[2].cache.get(b"key_c".as_slice()),
        Some(&(None, b"val_c".to_vec()))
    );
    assert!(sequence[2].removed.is_empty());
    assert_eq!(sequence[2], sequence[2].inverse().inverse());

    // Now we are going to apply each diff and check that the database
    // has been mutated accordingly
    db.write_tree_overlays_diff(&tree, &sequence[0], false)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, Some(b"val_b".into()));
    overlay.remove_diff(&sequence[0]);

    db.write_tree_overlays_diff(&tree, &sequence[1], false)?;
    db.flush_default_mode()?;
    assert!(tree.is_empty()?);
    overlay.remove_diff(&sequence[1]);

    // Since we removed the diffs, current overlay diff must be
    // the same as the last diff in the sequence
    let diff = overlay.diff(&[])?;
    assert_eq!(diff, sequence[2]);
    // Therefore we can safely use its batch
    db.write_tree_overlays_changes(&[&overlay])?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, None);
    assert_eq!(tree.get(b"key_c")?, Some(b"val_c".into()));
    overlay.remove_diff(&sequence[2]);

    // Since we removed everything, current overlay must not have
    // diffs over the tree, therefore its safe to keep using it
    let diff = overlay.diff(&[])?;
    assert!(diff.cache.is_empty());
    assert!(diff.removed.is_empty());

    // We are going to make some changes that we want to revert
    // using the corresponding diff
    overlay.clear()?;

    // Grab the diff, apply it and verify tree state
    let diff = overlay.diff(&[])?;
    db.write_tree_overlays_diff(&tree, &diff, false)?;
    db.flush_default_mode()?;
    assert!(tree.is_empty()?);

    // Now we grab the diff revert batch, apply it and verity tree state
    db.write_tree_overlays_diff(&tree, &diff, true)?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, None);
    assert_eq!(tree.get(b"key_c")?, Some(b"val_c".into()));

    // Now we are going to revert the diffs sequence going backwards
    // and verify tree state mutates accordingly
    db.write_tree_overlays_diff(&tree, &sequence[2], true)?;
    db.flush_default_mode()?;
    assert!(tree.is_empty()?);

    db.write_tree_overlays_diff(&tree, &sequence[1], true)?;
    db.flush_default_mode()?;
    assert_eq!(tree.len()?, 2);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));
    assert_eq!(tree.get(b"key_b")?, Some(b"val_b".into()));

    db.write_tree_overlays_diff(&tree, &sequence[0], true)?;
    db.flush_default_mode()?;

    // Tree has now reverted to its original state
    assert_eq!(tree.len()?, 1);
    assert_eq!(tree.get(b"key_a")?, Some(b"val_a".into()));

    Ok(())
}