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

use std::collections::{BTreeMap, BTreeSet};

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

/// Struct representing [`DatabaseOverlay`] cache state.
#[derive(Debug, Clone)]
pub struct DatabaseOverlayState {
    /// Existing trees in `Database` at the time of instantiation, so
    /// we can track newly opened trees.
    pub initial_tree_names: Vec<String>,
    /// New trees that have been opened, but didn't exist in `Database`
    /// before.
    pub new_tree_names: Vec<String>,
    /// Pointers to [`TreeOverlay`] instances that have been created.
    pub caches: BTreeMap<String, TreeOverlay>,
    /// Trees that were dropped, along with their last state full diff.
    pub dropped_trees: BTreeMap<String, TreeOverlayStateDiff>,
    /// Protected trees, that we don't allow their removal, and don't
    /// drop their references if they become stale.
    pub protected_tree_names: Vec<String>,
}

impl DatabaseOverlayState {
    /// Instantiate a new [`DatabaseOverlayState`].
    pub fn new(initial_tree_names: Vec<String>, protected_tree_names: Vec<String>) -> Self {
        Self {
            initial_tree_names,
            new_tree_names: vec![],
            caches: BTreeMap::new(),
            dropped_trees: BTreeMap::new(),
            protected_tree_names,
        }
    }

    /// Aggregate all the current overlay changes into [`Batch`]
    /// instances and return vectors of [`Tree`] and their respective
    /// [`Batch`] that can be used for further operations. If there are
    /// no changes, both vectors will be empty.
    pub fn aggregate(&self) -> Result<Vec<(Tree, Batch)>> {
        let mut batches = vec![];

        for (key, cache) in self.caches.iter() {
            if self.dropped_trees.contains_key(key) {
                return Err(Error::CollectionNotFound(key.clone()));
            }

            if let Some(batch) = cache.aggregate() {
                batches.push((cache.tree.clone(), batch));
            }
        }

        Ok(batches)
    }

    /// Add provided `Database` overlay state changes to our own.
    /// If a `Tree` doesn't exists it is opened using the default
    /// backend configuration
    pub fn add_diff(&mut self, database: &Database, diff: &DatabaseOverlayStateDiff) -> Result<()> {
        self.initial_tree_names
            .retain(|x| diff.initial_tree_names.contains(x));

        for (k, (cache, drop)) in diff.caches.iter() {
            if *drop {
                assert!(!self.protected_tree_names.contains(k));
                self.new_tree_names.retain(|x| x != k);
                self.caches.remove(k);
                self.dropped_trees.insert(k.clone(), cache.clone());
                continue;
            }

            let Some(tree_overlay) = self.caches.get_mut(k) else {
                if !self.initial_tree_names.contains(k) && !self.new_tree_names.contains(k) {
                    self.new_tree_names.push(k.clone());
                }
                let mut overlay = TreeOverlay::new(&database.open_tree_default(k)?);
                overlay.add_diff(cache);
                self.caches.insert(k.clone(), overlay);
                continue;
            };

            // Add the diff to our tree overlay state
            tree_overlay.add_diff(cache);
        }

        for (k, (cache, restored)) in &diff.dropped_trees {
            // Drop the trees that are not restored
            if !restored {
                if self.dropped_trees.contains_key(k) {
                    continue;
                }
                self.new_tree_names.retain(|x| x != k);
                self.caches.remove(k);
                self.dropped_trees.insert(k.clone(), cache.clone());
                continue;
            }
            assert!(!self.protected_tree_names.contains(k));

            // Restore the tree
            self.initial_tree_names.retain(|x| x != k);
            if !self.new_tree_names.contains(k) {
                self.new_tree_names.push(k.clone());
            }

            let mut overlay = TreeOverlay::new(&database.open_tree_default(k)?);
            overlay.add_diff(cache);
            self.caches.insert(k.clone(), overlay);
        }

        Ok(())
    }

    /// Remove provided `database` overlay state changes from our own.
    pub fn remove_diff(&mut self, diff: &DatabaseOverlayStateDiff) {
        // We have some assertions here to catch catastrophic
        // logic bugs here, as all our fields are depending on each
        // other when checking for differences.
        for (k, (cache, drop)) in diff.caches.iter() {
            // We must know the tree
            assert!(
                self.initial_tree_names.contains(k)
                    || self.new_tree_names.contains(k)
                    || self.dropped_trees.contains_key(k)
            );
            if !self.initial_tree_names.contains(k) {
                self.initial_tree_names.push(k.clone());
            }
            self.new_tree_names.retain(|x| x != k);

            // Check if tree is marked for drop
            if *drop {
                assert!(!self.protected_tree_names.contains(k));
                self.initial_tree_names.retain(|x| x != k);
                self.new_tree_names.retain(|x| x != k);
                self.caches.remove(k);
                self.dropped_trees.remove(k);
                continue;
            }

            // If the key is not in the cache, and it exists
            // in the dropped trees, update its diff
            let Some(tree_overlay) = self.caches.get_mut(k) else {
                let Some(tree_overlay) = self.dropped_trees.get_mut(k) else {
                    continue;
                };
                tree_overlay.update_values(cache);
                continue;
            };

            // If the state is unchanged, handle the stale tree
            if tree_overlay.state == cache.into() {
                // If tree is protected, we simply reset its cache
                if self.protected_tree_names.contains(k) {
                    tree_overlay.state.cache = BTreeMap::new();
                    tree_overlay.state.removed = BTreeSet::new();
                    tree_overlay.checkpoint();
                    continue;
                }

                // Drop the stale reference
                self.caches.remove(k);
                continue;
            }

            // Remove the diff from our tree overlay state
            tree_overlay.remove_diff(cache);
        }

        // Now we handle the dropped trees
        for (k, (cache, restored)) in diff.dropped_trees.iter() {
            // We must know the tree
            assert!(
                self.initial_tree_names.contains(k)
                    || self.new_tree_names.contains(k)
                    || self.dropped_trees.contains_key(k)
            );

            // Drop the trees that are not restored
            if !restored {
                assert!(!self.protected_tree_names.contains(k));
                self.initial_tree_names.retain(|x| x != k);
                self.new_tree_names.retain(|x| x != k);
                self.caches.remove(k);
                self.dropped_trees.remove(k);
                continue;
            }

            // Restore the tree
            self.initial_tree_names.retain(|x| x != k);
            if !self.new_tree_names.contains(k) {
                self.new_tree_names.push(k.clone());
            }

            // Skip if not in cache
            let Some(tree_overlay) = self.caches.get_mut(k) else {
                continue;
            };

            // If the state is unchanged, handle the stale tree
            if tree_overlay.state == cache.into() {
                // If tree is protected, we simply reset its cache
                if self.protected_tree_names.contains(k) {
                    tree_overlay.state.cache = BTreeMap::new();
                    tree_overlay.state.removed = BTreeSet::new();
                    tree_overlay.checkpoint();
                    continue;
                }

                // Drop the stale reference
                self.caches.remove(k);
                continue;
            }

            // Remove the diff from our tree overlay state
            tree_overlay.remove_diff(cache);
        }
    }
}

impl Default for DatabaseOverlayState {
    fn default() -> Self {
        Self::new(vec![], vec![])
    }
}

/// Auxilliary struct representing a [`DatabaseOverlayState`] diff log.
#[derive(Debug, Default, Clone, PartialEq)]
pub struct DatabaseOverlayStateDiff {
    /// Existing trees in `database` at the time of instantiation, so
    /// we can track newly opened trees.
    pub initial_tree_names: Vec<String>,
    /// State diff logs of all [`TreeOverlay`] instances that have been
    /// created, along with a boolean flag indicating if it should be
    /// dropped. The drop flag is always set to false, and change to
    /// true when we inverse the diff of a new tree(not in our initial
    /// tree names) and the inserts vector is empty, indicating that
    /// the tree should be dropped.
    pub caches: BTreeMap<String, (TreeOverlayStateDiff, bool)>,
    /// Trees that were dropped, along with their last state full diff,
    /// along with a boolean flag indicating if they should be
    /// restored. The restore flag is always set to false, and change
    /// to true when we inverse the diff, unless the tree is a new
    /// tree(not in our initial tree names).
    pub dropped_trees: BTreeMap<String, (TreeOverlayStateDiff, bool)>,
}

impl DatabaseOverlayStateDiff {
    /// Instantiate a new [`DatabaseOverlayStateDiff`], over the
    /// provided [`DatabaseOverlayState`].
    pub fn new(state: &DatabaseOverlayState) -> Result<Self> {
        let mut caches = BTreeMap::new();
        let mut dropped_trees = BTreeMap::new();

        for (key, cache) in state.caches.iter() {
            let diff = cache.diff(&[])?;

            // Skip if diff is empty for an existing tree
            if diff.cache.is_empty()
                && diff.removed.is_empty()
                && !state.new_tree_names.contains(key)
            {
                continue;
            }

            caches.insert(key.clone(), (diff, false));
        }

        for (key, cache) in state.dropped_trees.iter() {
            dropped_trees.insert(key.clone(), (cache.clone(), false));
        }

        Ok(Self {
            initial_tree_names: state.initial_tree_names.clone(),
            caches,
            dropped_trees,
        })
    }

    /// Aggregate all the 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. Provided state trees must contain all the
    /// [`Tree`] pointers the diff mutates.
    pub fn aggregate(&self, state_trees: &BTreeMap<String, Tree>) -> Result<Vec<(Tree, Batch)>> {
        let mut batches = vec![];

        for (key, (cache, drop)) in self.caches.iter() {
            if *drop {
                continue;
            }

            let Some(tree) = state_trees.get(key) else {
                return Err(Error::CollectionNotFound(key.clone()));
            };

            if let Some(batch) = cache.aggregate() {
                batches.push((tree.clone(), batch));
            }
        }

        for (key, (cache, restored)) in self.dropped_trees.iter() {
            if !restored {
                continue;
            }

            let Some(tree) = state_trees.get(key) else {
                return Err(Error::CollectionNotFound(key.clone()));
            };

            if let Some(batch) = cache.aggregate() {
                batches.push((tree.clone(), batch));
            }
        }

        Ok(batches)
    }

    /// Produces a [`DatabaseOverlayStateDiff`] containing the inverse
    /// changes from our own.
    pub fn inverse(&self) -> Self {
        let mut diff = Self {
            initial_tree_names: self.initial_tree_names.clone(),
            ..Default::default()
        };

        for (key, (cache, drop)) in self.caches.iter() {
            let inverse = cache.inverse();
            // Flip its drop flag if its a new empty tree, otherwise
            // check if its cache is empty and its a new tree.
            let drop = if inverse.cache.is_empty()
                && inverse.removed.is_empty()
                && !self.initial_tree_names.contains(key)
            {
                !drop
            } else {
                inverse.cache.is_empty() && !self.initial_tree_names.contains(key)
            };
            diff.caches.insert(key.clone(), (inverse, drop));
        }

        for (key, (cache, restored)) in self.dropped_trees.iter() {
            if !self.initial_tree_names.contains(key) {
                continue;
            }
            diff.dropped_trees
                .insert(key.clone(), (cache.clone(), !restored));
        }

        diff
    }

    /// Remove provided `database` overlay state changes from our own.
    pub fn remove_diff(&mut self, other: &Self) {
        // We have some assertions here to catch catastrophic
        // logic bugs here, as all our fields are depending on each
        // other when checking for differences.
        for initial_tree_name in &other.initial_tree_names {
            assert!(self.initial_tree_names.contains(initial_tree_name));
        }

        // First we remove each cache diff
        for (key, cache_pair) in other.caches.iter() {
            if !self.initial_tree_names.contains(key) {
                self.initial_tree_names.push(key.clone());
            }

            // If the key is not in the cache, and it exists
            // in the dropped trees, update its diff.
            let Some(tree_overlay) = self.caches.get_mut(key) else {
                let Some((tree_overlay, _)) = self.dropped_trees.get_mut(key) else {
                    continue;
                };
                tree_overlay.update_values(&cache_pair.0);
                continue;
            };

            // If the state is unchanged, handle the stale tree
            if tree_overlay == cache_pair {
                // Drop the stale reference
                self.caches.remove(key);
                continue;
            }

            // Remove the diff from our tree overlay state
            tree_overlay.0.remove_diff(&cache_pair.0);
        }

        // Now we handle the dropped trees. We must have all
        // the keys in our dropped trees keys.
        for (key, (cache, restored)) in other.dropped_trees.iter() {
            // Check if the tree was reopened
            if let Some(tree_overlay) = self.caches.get_mut(key) {
                assert!(!self.dropped_trees.contains_key(key));
                // Remove the diff from our tree overlay state
                tree_overlay.0.remove_diff(cache);
                continue;
            }
            assert!(self.dropped_trees.contains_key(key));

            // Restore tree if its flag is set to true
            if *restored {
                self.caches.insert(key.clone(), (cache.clone(), false));
            }

            // Drop the tree
            self.initial_tree_names.retain(|x| x != key);
            self.dropped_trees.remove(key);
        }
    }

    /// Auxilliary function to retrieve our newly opened trees.
    pub fn new_trees(&self) -> Vec<String> {
        let mut new_trees: Vec<String> = self.caches.keys().cloned().collect();
        new_trees.retain(|tree| !self.initial_tree_names.contains(tree));
        new_trees
    }
}