kvdb-overlay/src/database_overlay.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/>.
 */

use std::collections::BTreeMap;

#[cfg(feature = "sled-backend")]
use sled::{Transactional, transaction::ConflictableTransactionError};

use crate::{
    Batch, Database, DatabaseOverlayState, DatabaseOverlayStateDiff, Error, Result, Tree,
    TreeOverlay, TreeOverlayIter, TreeOverlayStateDiff,
};

/// An overlay on top of an entire [`Database`] which can span multiple trees
#[derive(Clone)]
pub struct DatabaseOverlay {
    /// The [`Database`] that is being overlayed.
    db: Database,
    /// Current overlay cache state
    pub state: DatabaseOverlayState,
    /// Checkpointed cache state to revert to
    checkpoint: DatabaseOverlayState,
}

impl DatabaseOverlay {
    /// Instantiate a new [`DatabaseOverlay`] on top of a given
    /// [`Database`].
    /// Note: Provided protected trees don't have to be opened as
    /// protected, as they are setup as protected here.
    pub fn new(db: &Database, protected_tree_names: Vec<String>) -> Result<Self> {
        let initial_tree_names = db.tree_names()?;
        Ok(Self {
            db: db.clone(),
            state: DatabaseOverlayState::new(
                initial_tree_names.clone(),
                protected_tree_names.clone(),
            ),
            checkpoint: DatabaseOverlayState::new(initial_tree_names, protected_tree_names),
        })
    }

    /// Create a new [`TreeOverlay`] on top of a given `tree_name`.
    /// This function will also open a new tree inside `database`
    /// regardless of if it has existed before, so for convenience, we
    /// also provide [`DatabaseOverlay::purge_new_trees`] in case we
    /// decide we don't want to write the batches, and drop the new
    /// trees. Additionally, a boolean flag is passed to mark the
    /// oppened tree as protected, meanning that it can't be removed
    /// and its references will never be dropped.
    pub fn open_tree(
        &mut self,
        name: &str,
        #[cfg(feature = "fjall-backend")] create_options: impl FnOnce() -> fjall::KeyspaceCreateOptions,
        protected: bool,
    ) -> Result<()> {
        // Check if we have already opened this tree
        if self.state.caches.contains_key(name) {
            return Ok(());
        }

        // Open this tree in the database
        let tree = self.db.open_tree(
            name,
            #[cfg(feature = "fjall-backend")]
            create_options,
        )?;
        let mut cache = TreeOverlay::new(&tree);

        // If we are reopenning a dropped tree, grab its cache
        if let Some(diff) = self.state.dropped_trees.remove(name) {
            cache.state = (&diff).into();
        }

        // In case it hasn't existed before, we also need to track it
        // in `self.new_tree_names`.
        let name = name.to_string();
        if !self.state.initial_tree_names.contains(&name) {
            self.state.new_tree_names.push(name.clone());
        }

        self.state.caches.insert(name.clone(), cache);

        // Mark tree as protected if requested
        if protected && !self.state.protected_tree_names.contains(&name) {
            self.state.protected_tree_names.push(name);
        }

        Ok(())
    }

    /// Create a new [`TreeOverlay`] on top of a given `tree_name`.
    /// This function will also open a new tree inside `database`,
    /// using the default backend configuration, regardless of if it
    /// has existed before, so for convenience, we also provide
    /// [`DatabaseOverlay::purge_new_trees`] in case we decide we don't
    /// want to write the batches, and drop the new trees.
    /// Additionally, a boolean flag is passed to mark the oppened tree
    /// as protected, meanning that it can't be removed and its
    /// references will never be dropped.
    pub fn open_tree_default(&mut self, name: &str, protected: bool) -> Result<()> {
        #[cfg(feature = "sled-backend")]
        {
            self.open_tree(name, protected)
        }
        #[cfg(feature = "fjall-backend")]
        {
            self.open_tree(name, fjall::KeyspaceCreateOptions::default, protected)
        }
    }

    /// Drop a tree from the overlay.
    pub fn drop_tree(&mut self, name: &str) -> Result<()> {
        // Check if tree is protected
        let name = name.to_string();
        if self.state.protected_tree_names.contains(&name) {
            return Err(Error::ProtectedTreeDrop(name));
        }

        // Check if already removed
        if self.state.dropped_trees.contains_key(&name) {
            return Err(Error::CollectionNotFound(name));
        }

        // Check if its a new tree we created
        if self.state.new_tree_names.contains(&name) {
            self.state.new_tree_names.retain(|x| x != &name);
            let tree = match self.get_cache(&name) {
                Ok(cache) => &cache.tree,
                _ => &self.db.open_tree_default(&name)?,
            };
            let diff = TreeOverlayStateDiff::new_dropped(tree);
            self.state.caches.remove(&name);
            self.state.dropped_trees.insert(name, diff);

            return Ok(());
        }

        // Check if tree existed in the database
        if !self.state.initial_tree_names.contains(&name) {
            return Err(Error::CollectionNotFound(name));
        }

        let tree = match self.get_cache(&name) {
            Ok(cache) => &cache.tree,
            _ => &self.db.open_tree_default(&name)?,
        };
        let diff = TreeOverlayStateDiff::new_dropped(tree);
        self.state.caches.remove(&name);
        self.state.dropped_trees.insert(name, diff);

        Ok(())
    }

    /// Drop newly created trees from the database. This is a
    /// convenience function that should be used when we decide that we
    /// don't want to apply any cache changes, and we want to revert
    /// back to the initial state.
    pub fn purge_new_trees(&self) -> Result<()> {
        for i in &self.state.new_tree_names {
            self.db.drop_tree(i)?;
        }

        Ok(())
    }

    /// Fetch the cache for a given tree.
    fn get_cache(&self, name: &str) -> Result<&TreeOverlay> {
        let name = name.to_string();
        if self.state.dropped_trees.contains_key(&name) {
            return Err(Error::CollectionNotFound(name));
        }

        if let Some(v) = self.state.caches.get(&name) {
            return Ok(v);
        }

        Err(Error::CollectionNotFound(name))
    }

    /// Fetch a mutable reference to the cache for a given tree.
    fn get_cache_mut(&mut self, name: &str) -> Result<&mut TreeOverlay> {
        let name = name.to_string();
        if self.state.dropped_trees.contains_key(&name) {
            return Err(Error::CollectionNotFound(name));
        }

        if let Some(v) = self.state.caches.get_mut(&name) {
            return Ok(v);
        }

        Err(Error::CollectionNotFound(name))
    }

    /// Fetch all our caches current [`Tree`] pointers.
    pub fn get_state_trees(&self) -> BTreeMap<String, Tree> {
        // Grab our state tree pointers
        let mut state_trees = BTreeMap::new();
        for (name, cache) in self.state.caches.iter() {
            state_trees.insert(name.clone(), cache.tree.clone());
        }

        state_trees
    }

    /// Returns `true` if the overlay contains a value for a specified
    /// key in the specified tree cache.
    pub fn contains_key(&self, name: &str, key: &[u8]) -> Result<bool> {
        let cache = self.get_cache(name)?;
        cache.contains_key(key)
    }

    /// Retrieve a value from the overlay if it exists in the specified
    /// tree cache.
    pub fn get(&self, name: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let cache = self.get_cache(name)?;
        cache.get(key)
    }

    /// Returns `true` if specified tree cache is empty.
    pub fn is_empty(&self, name: &str) -> Result<bool> {
        let cache = self.get_cache(name)?;
        cache.is_empty()
    }

    /// Returns last value from the overlay if the specified tree cache
    /// is not empty.
    pub fn last(&self, name: &str) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
        let cache = self.get_cache(name)?;
        cache.last()
    }

    /// Insert a key to a new value in the specified tree cache,
    /// returning the last value if it was set.
    pub fn insert(&mut self, name: &str, key: &[u8], value: &[u8]) -> Result<Option<Vec<u8>>> {
        let cache = self.get_cache_mut(name)?;
        cache.insert(key, value)
    }

    /// Delete a value in the specified tree cache, returning the old
    /// value if it existed.
    pub fn remove(&mut self, name: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let cache = self.get_cache_mut(name)?;
        cache.remove(key)
    }

    /// Removes all values from the specified tree cache and marks all
    /// its tree records as removed.
    pub fn clear(&mut self, name: &str) -> Result<()> {
        let cache = self.get_cache_mut(name)?;
        cache.clear()
    }

    /// Aggregate all the current overlay changes into [`Batch`]
    /// instances and return a vector of `Tree` and their respective
    /// `Batch` that can be used for further operations. If there are
    /// no changes, vector will be empty.
    fn aggregate(&self) -> Result<Vec<(Tree, Batch)>> {
        self.state.aggregate()
    }

    /// Ensure all new trees that have been opened exist in the
    /// database by reopening them with the default backend
    /// configuration, atomically apply all batches on all trees as a
    /// transaction, and drop dropped trees from the database. This
    /// function **does not** perform a db flush. This should be done
    /// externally, since then there is a choice to perform either
    /// blocking or async IO. After execution is successful, caller
    /// should *NOT* use the overlay again.
    pub fn apply(&mut self) -> Result<()> {
        // Ensure new trees exist
        let new_tree_names = self.state.new_tree_names.clone();
        for tree_names in &new_tree_names {
            let tree = self.db.open_tree_default(tree_names)?;
            // Update cache tree pointer, it must exist
            let cache = self.get_cache_mut(tree_names)?;
            cache.tree = tree;
        }

        // Drop removed trees
        for tree in self.state.dropped_trees.keys() {
            self.db.drop_tree(tree)?;
        }

        // Aggregate batches
        let batches = self.aggregate()?;
        if batches.is_empty() {
            return Ok(());
        }

        #[cfg(feature = "sled-backend")]
        {
            // Grab all referenced trees
            let trees: Vec<&sled::Tree> = batches.iter().map(|(t, _)| t.tree()).collect();

            // Perform an atomic transaction over all the collected trees and
            // apply the batches.
            if let Err(e) = trees.transaction(|trees| {
                for (i, tree) in trees.iter().enumerate() {
                    // Build and apply its batch
                    let mut sled_batch = sled::Batch::default();
                    for (key, value) in &batches[i].1.writes {
                        match value {
                            Some(v) => sled_batch.insert(key.as_slice(), v.as_slice()),
                            None => sled_batch.remove(key.as_slice()),
                        }
                    }
                    tree.apply_batch(&sled_batch)?;
                }

                Ok::<(), ConflictableTransactionError<sled::Error>>(())
            }) {
                return Err(Error::Transaction(e.to_string()));
            };
        }
        #[cfg(feature = "fjall-backend")]
        {
            // Grab a batch over the whole database
            let mut fjall_batch = self.db.fjall_batch();

            // Aggregate the overlay changes into the batch
            for (tree, batch) in batches {
                for (key, value) in batch.writes {
                    match value {
                        Some(v) => fjall_batch.insert(tree.tree(), key, v),
                        None => fjall_batch.remove(tree.tree(), key),
                    }
                }
            }

            // Apply the batch
            fjall_batch.commit()?;
        }

        Ok(())
    }

    /// Checkpoint current cache state so we can revert to it, if
    /// needed.
    pub fn checkpoint(&mut self) {
        self.checkpoint = self.state.clone();
    }

    /// Revert to current cache state checkpoint. This function will
    /// not drop new trees from the `db`, so caller should handle it.
    pub fn revert_to_checkpoint(&mut self) {
        self.state = self.checkpoint.clone();
    }

    /// Calculate differences from provided overlay state changes
    /// sequence. This can be used when we want to keep track of
    /// consecutive individual changes performed over the current
    /// overlay state. If the sequence is empty, current state
    /// is returned as the diff.
    pub fn diff(&self, sequence: &[DatabaseOverlayStateDiff]) -> Result<DatabaseOverlayStateDiff> {
        // Grab current state
        let mut current = DatabaseOverlayStateDiff::new(&self.state)?;

        // Remove provided diffs sequence
        for diff in sequence {
            current.remove_diff(diff);
        }

        Ok(current)
    }

    /// Add provided `db` overlay state changes from our own.
    pub fn add_diff(&mut self, diff: &DatabaseOverlayStateDiff) -> Result<()> {
        self.state.add_diff(&self.db, diff)
    }

    /// Remove provided `db` overlay state changes from our own.
    pub fn remove_diff(&mut self, diff: &DatabaseOverlayStateDiff) {
        self.state.remove_diff(diff)
    }

    /// For a provided `DatabaseOverlayStateDiff`, ensure all trees
    /// exist in the database by reopening them with the default
    /// backend configuration, atomically apply all batches on all
    /// trees as a transaction, and drop dropped trees from the
    /// database. After that, remove the state changes from our own.
    /// This is will also mutate the initial trees, based on what was
    /// oppened and/or dropped. This function **does not** perform a db
    /// flush. This should be done externally, since then there is a
    /// choice to perform either blocking or async IO.
    pub fn apply_diff(&mut self, diff: &DatabaseOverlayStateDiff) -> Result<()> {
        // We assert that the diff doesn't try to drop any of our
        // protected trees.
        for name in diff.dropped_trees.keys() {
            if self.state.protected_tree_names.contains(name) {
                return Err(Error::ProtectedTreeDrop(name.clone()));
            }
        }
        for (name, (_, drop)) in diff.caches.iter() {
            if *drop && self.state.protected_tree_names.contains(name) {
                return Err(Error::ProtectedTreeDrop(name.clone()));
            }
        }

        // Grab current state trees
        let mut state_trees = self.get_state_trees();

        // Ensure diff trees exist
        for (name, (_, drop)) in diff.caches.iter() {
            // Check if its an unknown tree
            if !self.state.initial_tree_names.contains(name)
                && !self.state.new_tree_names.contains(name)
            {
                self.state.new_tree_names.push(name.clone());
            }

            // Check if it should be dropped
            if *drop {
                self.db.drop_tree(name)?;
                continue;
            }

            if !state_trees.contains_key(name) {
                let tree = self.db.open_tree_default(name)?;
                state_trees.insert(name.clone(), tree);
            }
        }

        // Drop removed trees and ensure restored trees exist
        for (name, (_, restored)) in diff.dropped_trees.iter() {
            if !restored {
                state_trees.remove(name);
                self.db.drop_tree(name)?;
                continue;
            }

            // Check if its an unknown tree
            if !self.state.initial_tree_names.contains(name)
                && !self.state.new_tree_names.contains(name)
            {
                self.state.new_tree_names.push(name.clone());
            }

            if !state_trees.contains_key(name) {
                let tree = self.db.open_tree_default(name)?;
                state_trees.insert(name.clone(), tree);
            }
        }

        // Aggregate batches
        let batches = diff.aggregate(&state_trees)?;
        if batches.is_empty() {
            self.remove_diff(diff);
            return Ok(());
        }

        #[cfg(feature = "sled-backend")]
        {
            // Grab all referenced trees
            let trees: Vec<&sled::Tree> = batches.iter().map(|(t, _)| t.tree()).collect();

            // Perform an atomic transaction over all the collected trees and
            // apply the batches.
            if let Err(e) = trees.transaction(|trees| {
                for (i, tree) in trees.iter().enumerate() {
                    // Build and apply its batch
                    let mut sled_batch = sled::Batch::default();
                    for (key, value) in &batches[i].1.writes {
                        match value {
                            Some(v) => sled_batch.insert(key.as_slice(), v.as_slice()),
                            None => sled_batch.remove(key.as_slice()),
                        }
                    }
                    tree.apply_batch(&sled_batch)?;
                }

                Ok::<(), ConflictableTransactionError<sled::Error>>(())
            }) {
                return Err(Error::Transaction(e.to_string()));
            };
        }
        #[cfg(feature = "fjall-backend")]
        {
            // Grab a batch over the whole database
            let mut fjall_batch = self.db.fjall_batch();

            // Aggregate the overlay changes into the batch
            for (tree, batch) in batches {
                for (key, value) in batch.writes {
                    match value {
                        Some(v) => fjall_batch.insert(tree.tree(), key, v),
                        None => fjall_batch.remove(tree.tree(), key),
                    }
                }
            }

            // Apply the batch
            fjall_batch.commit()?;
        }

        // Remove changes from our current state
        self.remove_diff(diff);

        Ok(())
    }

    /// Retrieve an immutable itterator from the overlay if the
    /// specified tree cache exists.
    pub fn iter(&self, name: &str) -> Result<TreeOverlayIter<'_>> {
        let cache = self.get_cache(name)?;
        Ok(cache.iter())
    }
}