nb-tree 0.2.0-alpha01

use replace_with::replace_with_or_abort;
use std::{hash::Hash, mem};
use tracing::{debug, error, trace, trace_span};

use crate::path::Path;

use super::{
    diff::{DiffApplyError, DiffMap, DiffNode},
    iter::depth::Traversal,
    CombinationMode, DiffTree, NodeIDX, Tree,
};

/// A tree with optional data for each node
pub type TreeBuilder<N, B> = Tree<Option<N>, B>;

impl<N, B> Tree<N, B>
where
    N: Default,
    B: Default + Eq + Hash + Clone,
{
    /// Extends the tree branches with the given `path` if it leaves the tree
    pub(super) fn extend(&mut self, path: &Path<B>, from_idx: Option<NodeIDX>) -> Vec<NodeIDX> {
        match self.get_idx(path, from_idx) {
            Ok(idx) => vec![idx],
            Err(node) => {
                let mut idxs = Vec::new();
                let (mut node_idx, path_idx) = node.unwrap_or_else(|| {
                    //tree is empty. Create a root node
                    self.insert_root(Default::default());
                    idxs.push(0);
                    (0, 0)
                });
                for branch in path.iter().skip(path_idx) {
                    node_idx = self
                        .insert_at(node_idx, branch.clone(), Default::default())
                        .1;
                    idxs.push(node_idx);
                }
                idxs
            }
        }
    }
    /// Inserts the `value` at the given `path`, extending the tree if necessary
    pub fn insert_extend(&mut self, path: &Path<B>, value: N) -> N {
        let idx = *self.extend(path, None).last().unwrap();
        mem::replace(&mut self.nodes[idx].value, value)
    }

    /// Chainable insertion with extension.
    ///
    /// Calls `insert_extend` and returns a mutable reference to self
    pub fn ix(&mut self, path: impl Into<Path<B>>, value: N) -> &mut Self {
        let ph: Path<B> = path.into();
        self.insert_extend(&ph, value);
        self
    }

    /// Apply `diff` without checking its validity beforehand
    /// The tree will grow to include changes outside it.
    /// All invalid changes will be ignored (see [DiffTree::validate()] for more information).
    pub fn force_apply_extend(&mut self, diff: DiffTree<N, B>) -> bool {
        self.force_apply_with(
            diff,
            |entry, now| {
                entry.insert_extend(now);
                true
            },
            |entry| {
                replace_with_or_abort(entry, |e| e.remove_subtree());
                true
            },
        )
    }
}

impl<N, B> Tree<N, B>
where
    N: Default + Eq,
    B: Default + Eq + Hash + Clone,
{
    pub fn combine_extend(&mut self, tree: Tree<N, B>, mode: CombinationMode) -> bool {
        self.combine_with(tree, mode, |entry, now| {
            entry.insert_extend(now);
            true
        })
    }
    /// Removes trailing default nodes
    pub fn trim_default(&mut self) {
        let mut removable = Vec::new();
        let mut start = Path::new();
        let mut path = Path::new();
        for step in self.iter_on::<&N>() {
            //TODO: prevent clone per loop
            let old_path = path.clone();
            path.apply_deref(&step);
            if start.len() >= path.len() {
                // Left the trimmable subtree
                removable.push(start);
                start = Path::new();
            }
            if **step.data() == N::default() {
                // Trimmable node
                if start.is_empty() {
                    // Trimmable subtree root
                    start = path.clone();
                }
            } else {
                // Not trimmable node
                if !start.is_empty() {
                    // Move trimmable subtree root down
                    removable.push(old_path.path_to(path.len()));
                }
            }
        }

        // Remove trimmable subtrees
        for ph in removable {
            self.remove_subtree(&ph)
                .unwrap_or_else(|_| unreachable!("Could not trim a trimmable subtree"));
        }
    }
}

impl<N, B> Tree<N, B>
where
    N: Default + Eq,
    B: Default + Eq + Hash + Clone,
{
    pub fn apply_extend(&mut self, diff: DiffTree<N, B>) -> Result<bool, DiffApplyError<B>> {
        // Check diff
        diff.is_applicable_extend(self)?;

        // Apply diff
        Ok(self.force_apply_extend(diff))
    }

    pub fn apply_diff_extend(
        &mut self,
        path: Path<B>,
        diff: DiffNode<N>,
    ) -> Result<(), DiffApplyError<B>> {
        self.apply_diff_with(path, diff, |entry, v| {
            entry.insert_extend(v);
            Ok(())
        })
    }

    pub fn apply_map_extend(&mut self, map: DiffMap<N, B>) -> bool {
        map.into_iter().fold(true, |mut res, (ph, diff)| {
            if self.apply_diff_extend(ph, diff).is_err() {
                res = false;
            }
            res
        })
    }

    /// Fails if the diff is invalid
    pub fn remove_subtree_trim(&mut self, path: &Path<B>) -> Result<(), Option<Path<B>>> {
        let mut path = path.clone();
        // Remove node and children from tree
        self.remove_subtree(&path)?;

        // Has parent? (Not root node)
        if let Some(branch) = path.pop_last() {
            // Get parent indexes
            let mut parents = self
                .get_path_idxs(&path)
                .map_err(|r| path.path_to(*r.unwrap_or((vec![0], 0)).0.last().unwrap_or(&0)))?;

            // Remove empty parent nodes
            while let Some(idx) = parents.pop() {
                if self.nodes[idx].value == Default::default()
                    && parents.last().map(|&i| self.nodes[i].children.len()) == Some(1)
                {
                    self.remove_subtree_at(*parents.last().unwrap(), path.pop_last().unwrap());
                } else {
                    break;
                }
            }

            // Remove child parent link
            if !parents.is_empty() {
                self.nodes[*parents.last().unwrap()]
                    .children
                    .remove(&branch);
            }
        }
        Ok(())
    }
}

impl<N, B> TreeBuilder<N, B>
where
    N: Default,
    B: Default + Eq + Hash + Clone,
{
    pub fn is_buildable(&self) -> bool {
        // Check for adjacency
        let mut iter = self.iter_on::<&Option<N>>();

        if let Some(Traversal::Start(Some(_))) = iter.next() {
            // First node is root and not empty
        } else {
            return false;
        };
        iter.try_fold(Path::new(), |mut ph, node| {
            // No Root node and all nodes contain data
            (ph.apply_deref(&node) && node.take().is_some()).then_some(ph)
        })
        .is_some()
    }

    pub fn build(self) -> Tree<N, B> {
        // Build Tree
        self.into_iter()
            .try_fold((Tree::new(), Path::new()), |(mut tree, mut ph), node| {
                match node {
                    Traversal::Start(data) => {
                        tree.insert_root(data.unwrap());
                    }
                    Traversal::Step { up, branch, data } => {
                        let value = data.unwrap();
                        let parent_idx = ph.last().map(|(_, idx)| *idx).unwrap_or(0);

                        // Move up
                        (0..up).for_each(|_| {
                            ph.pop_last();
                        });
                        // Move down
                        ph.push_last((branch.clone(), tree.insert_at(parent_idx, branch, value).1));
                    }
                };
                Some((tree, ph))
            })
            .map(|(tree, _)| tree)
            .unwrap()
    }
}