mktree/

tree.rs

//! A tree structure and display with [`egui::SelectableLabel`] at every leaves.

use crate::*;

use crossbeam_channel::Sender;
use egui::{collapsing_header::CollapsingState, RichText};
use std::collections::HashSet;

// ----------------------------------------------------------------------------
#[derive(Clone, Debug, Default)]
pub struct Node<T: Entity> {
    /// This name will be used to represent the node in the tree.
    name: String,
    /// Inner content of the node.
    content: T,
    /// Collapse state of the `CollapsingHeader`.
    open: bool,
}

impl<T: Entity> Node<T> {
    fn empty() -> Self {
        Self {
            content: T::empty(),
            open: true,
            ..Default::default()
        }
    }

    fn name(mut self, name: &str) -> Self {
        self.name = name.to_owned();
        self
    }

    fn match_by(&self, pattern: &str) -> bool {
        self.name.to_lowercase().contains(pattern)
    }
}

#[derive(Clone, Debug)]
/// Used for display tree data with [`Node`].
/// We're distinguishing the type of the nodes in order to perform further actions
/// depending on the selected node type.
pub enum TreeNode<T: Entity> {
    Root(Node<T>),
    SubTreeRoot(Node<T>),
    Leaf(Node<T>),
}

impl<T: Entity> Default for TreeNode<T> {
    /// Makes a leaf.
    fn default() -> Self {
        TreeNode::Leaf(Node::default())
    }
}

impl<T: Entity> TreeNode<T> {
    fn node_name(&self) -> &String {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => &inner.name,
        }
    }

    fn open(&self) -> bool {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => inner.open,
        }
    }

    fn open_mut(&mut self, open: bool) {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => {
                inner.open = open;
            }
        }
    }

    fn toggle_collapse_state(&mut self) {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => {
                inner.open = !inner.open;
            }
        }
    }

    fn selected(&self, selection: &HashSet<T>) -> bool {
        match self {
            // TODO: `SubTreeRoot`s won't be shown as selected for now.
            Self::Root(_) | Self::SubTreeRoot(_) => false,
            Self::Leaf(inner) => match inner.content.name() {
                // TODO: this makes assets with the same name all show as selected
                Some(name) => selection
                    .iter()
                    .filter_map(|a| a.name())
                    .collect::<HashSet<&String>>()
                    .contains(name),
                None => false,
            },
        }
    }

    fn leaf(entity: T) -> TreeNode<T> {
        TreeNode::Leaf(Node {
            name: match entity.name() {
                Some(name) => name.to_owned(),
                None => EMPTY_NODE_NAME.to_owned(),
            },
            content: entity,
            open: false,
        })
    }

    fn subtree_root(group: &str, typ: &ProjectSource) -> TreeNode<T> {
        TreeNode::SubTreeRoot(Node {
            name: group.to_owned(),
            content: T::as_group(group, typ),
            open: true,
        })
    }
}

// ----------------------------------------------------------------------------
#[derive(Clone, Debug, Default)]
/// Recursive tree struct.
pub struct MkTree<T: Entity>(pub TreeNode<T>, pub Vec<MkTree<T>>);

impl<T: Entity> MkTree<T> {
    /// Can be used to sort trees after construction.
    /// Must clone to get around the borrow checker.
    pub fn node_name(&self) -> String {
        self.0.node_name().clone()
    }

    fn open(&self) -> bool {
        self.0.open()
    }

    fn toggle_collapse_state(&mut self) {
        self.0.toggle_collapse_state()
    }

    fn node_content(&self) -> &T {
        match &self.0 {
            TreeNode::Root(inner) | TreeNode::SubTreeRoot(inner) | TreeNode::Leaf(inner) => {
                &inner.content
            }
        }
    }

    /// Recursively determines if the node should be visible, given the filter string.
    fn match_by(&self, pattern: &str) -> bool {
        match &self.0 {
            TreeNode::Leaf(inner) => inner.match_by(pattern),
            TreeNode::SubTreeRoot(_) => self
                .1
                .iter()
                .filter(|n| n.match_by(pattern))
                .next()
                .is_some(),
            TreeNode::Root(_) => true,
        }
    }

    /// The given `T` will be used to make a leaf -- a tip of the tree -- all by itself,
    /// i.e. containing empty list of trees.
    fn single_leaf(entity: T) -> Self {
        Self(TreeNode::leaf(entity), vec![])
    }

    /// The given group name will be made into a `SubTreeRoot` directly containing the given leaves.
    pub fn leaf_group(group: &str, leaves: Vec<T>, typ: &ProjectSource) -> Self {
        Self(
            TreeNode::subtree_root(group, typ),
            leaves
                .into_iter()
                .map(|leaf| Self::single_leaf(leaf))
                .collect(),
        )
    }

    pub fn subtree_group(group: &str, subtrees: Vec<MkTree<T>>, typ: &ProjectSource) -> Self {
        Self(TreeNode::subtree_root(group, typ), subtrees)
    }

    /// A root and all other leaves at the same level.
    pub fn from_node_n_leaves(entity: T, leaves: Vec<T>) -> Self {
        Self(
            TreeNode::leaf(entity),
            leaves.into_iter().map(|c| Self::single_leaf(c)).collect(),
        )
    }

    pub fn from_node_n_subtrees(entity: T, subtrees: Vec<MkTree<T>>) -> Self {
        Self(TreeNode::leaf(entity), subtrees)
    }

    fn listen(
        &mut self,
        ui: &mut egui::Ui,
        response: egui::Response,
        selection: &mut HashSet<T>,
        sender: &Sender<TreeNodeSignal>,
    ) {
        if response.clicked() {
            if !ui.ctx().input(|i| i.modifiers).any() {
                selection.clear();
            };
            match &self.0 {
                TreeNode::Leaf(_) => {
                    // Only signal selection change in this case.
                    sender
                        .send(TreeNodeSignal::LeafClicked)
                        .expect("Channel of MkTree's selected node response has been disconnected");
                    // Saves selected `impl Entity` in `TreeContainer::selected_nodes`
                    selection.insert(self.node_content().clone());
                }
                // don't do anything when `TreeNode::SubTreeRoot` or `TreeNode::Root` gets clicked
                _ => {}
            };
            // log::info!("Selected node: {:?}", &selection);
        };
    }

    /// Recursively makes the tree UI, while collects the user selection of the leaf nodes, and sends message.
    pub fn ui(
        &mut self,
        ui: &mut egui::Ui,
        node_name: &str,
        depth: usize,
        filter: &str,
        selection: &mut HashSet<T>,
        sender: &Sender<TreeNodeSignal>,
    ) {
        let mut state = CollapsingState::load_with_default_open(
            ui.ctx(),
            ui.make_persistent_id(node_name),
            true,
        );

        state.set_open(self.open());

        let response = state
            .show_header(ui, |ui| {
                // NOTE: no semicolon, as we're returning the inner response!
                ui.selectable_label(self.0.selected(selection), node_name)
                // the `egui::Response` is returned from this `CollapsingState::show_header` method,
                // and passed to `Self::listen` to be sent via channel
                //
            })
            // then recursively draws the children content
            .body(|ui| self.children_ui(ui, &node_name, depth, filter, selection, sender));

        // toggles collapse state
        if response.0.clicked() {
            self.toggle_collapse_state();
        };

        // checks for clicks on the header UI -- selectable label
        self.listen(ui, response.1.inner, selection, sender);
    }

    fn children_ui(
        &mut self,
        ui: &mut egui::Ui,
        // this isn't used in this function body but in previous call it's used for the header
        _node_name: &str,
        depth: usize,
        filter: &str,
        selection: &mut HashSet<T>,
        sender: &Sender<TreeNodeSignal>,
    ) {
        let mut tree = std::mem::take(self);
        // modifies the children
        tree.1 = tree
            .1
            .into_iter()
            .map(|mut child| {
                let node_name = child.node_name();
                if filter.is_empty() || child.match_by(filter) {
                    child.ui(ui, &node_name, depth + 1, filter, selection, sender);
                };
                child
            })
            .collect();
        *self = tree
    }

    /// Recursively counts (by accumulation) the number of leaves.
    fn leaf_len(&self, count: &mut usize) -> usize {
        if let TreeNode::Leaf(_) = self.0 {
            *count += 1;
        };
        self.1.iter().for_each(|t| {
            t.leaf_len(count);
        });
        *count
    }

    fn collapse_all(&mut self, open: bool) {
        match self.0 {
            TreeNode::Root(_) => {
                // do not expand nor close
            }
            _ => {
                self.0.open_mut(open);
            }
        }
        self.1.iter_mut().for_each(|t| t.collapse_all(open));
    }
}

#[derive(Debug, Clone, Default)]
/// Container for the asset tree, where first field holds the currently selected asset (via primary mouse click).
pub struct TreeContainer<T: Entity> {
    /// The final tree with `tree.0` of `TreeNode::Root` type.
    tree: MkTree<T>,

    root_nice_name: String,

    total_leaves: usize,

    filter: String,

    /// This contains only non-empty ProductionAsset, but such item might not have `asset_name`, i.e.
    /// having only `category.main_type` value.
    selected_nodes: HashSet<T>,

    /// Whether next toggle action should expand or collapse the tree recursively.
    batch_collapse: bool,
}

impl<T: Entity> TreeContainer<T> {
    pub fn uninitialized_root() -> Self {
        Self {
            tree: MkTree(
                TreeNode::Root(Node::empty().name(TREE_ROOT_UNINITIALIZED_NAME)),
                vec![],
            ),
            ..Default::default()
        }
    }

    pub fn subtrees_mut(&mut self, subtrees: Vec<MkTree<T>>, root_name: &str, show_count: bool) {
        self.tree = MkTree(TreeNode::Root(Node::empty().name(root_name)), subtrees);
        self.count_leaves();
        self.make_root_nice_name(root_name, show_count);
    }

    pub fn selected_nodes(&self) -> &HashSet<T> {
        &self.selected_nodes
    }

    pub fn selected_nodes_mut(&mut self, selected_nodes: HashSet<T>) {
        self.selected_nodes = selected_nodes;
    }

    // TODO: provide select all leaves?
    pub fn clear_selection(&mut self) {
        self.selected_nodes = HashSet::new();
    }

    pub fn filter_ui(&mut self, width: f32, ui: &mut egui::Ui) {
        ui.horizontal(|ui| {
            ui.label("Filter by Name:");
            ui.add(egui::TextEdit::singleline(&mut self.filter).desired_width(width));
            // forces lowercase conversion
            self.filter = self.filter.to_lowercase();
            if ui.button("ｘ").clicked() {
                self.filter.clear();
            }
        });
    }

    pub fn batch_collapse_ui(&mut self, ui: &mut egui::Ui) {
        let text = RichText::new(if self.batch_collapse {
            "⏷ Expand Tree"
        } else {
            "➖ Collapse Tree"
        });
        if ui.button(text).clicked() {
            self.tree.collapse_all(self.batch_collapse);
            self.batch_collapse = !self.batch_collapse;
        };
    }

    pub fn tree_ui(&mut self, ui: &mut egui::Ui, sender: &Sender<TreeNodeSignal>) {
        self.tree.ui(
            ui,
            &self.root_nice_name,
            0,
            &self.filter,
            &mut self.selected_nodes,
            sender,
        )
    }

    /// Counts all the leaves.
    fn count_leaves(&mut self) {
        // must reset the count first
        self.total_leaves = 0;
        self.tree.leaf_len(&mut self.total_leaves);
    }

    /// Appends after the root name of the tree the count of leaves.
    fn make_root_nice_name(&mut self, root_name: &str, show_count: bool) {
        if show_count {
            self.root_nice_name = format!("{}: total {}", root_name, self.total_leaves);
        } else {
            self.root_nice_name = root_name.to_owned();
        }
    }
}