baichun-framework-core 0.1.0

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// 树节点特征
pub trait TreeNode {
    /// 获取节点ID
    fn id(&self) -> i64;
    /// 获取父节点ID
    fn parent_id(&self) -> i64;
    /// 获取子节点
    fn children(&self) -> &Vec<Self>
    where
        Self: Sized;
    /// 获取可变子节点
    fn children_mut(&mut self) -> &mut Vec<Self>
    where
        Self: Sized;
}

/// 树形工具
pub struct TreeUtils;

impl TreeUtils {
    /// 构建树形结构
    pub fn build_tree<T>(nodes: Vec<T>) -> Vec<T>
    where
        T: TreeNode + Clone,
    {
        let mut node_map: HashMap<i64, T> = HashMap::new();
        let mut root_nodes: Vec<T> = Vec::new();

        // 构建节点映射
        for node in nodes.iter() {
            node_map.insert(node.id(), node.clone());
        }

        // 构建树形结构
        for node in nodes {
            let parent_id = node.parent_id();
            if parent_id == 0 {
                root_nodes.push(node);
            } else if let Some(parent) = node_map.get_mut(&parent_id) {
                parent.children_mut().push(node);
            }
        }

        root_nodes
    }

    /// 获取所有子节点ID
    pub fn get_child_ids<T>(nodes: &[T], parent_id: i64) -> Vec<i64>
    where
        T: TreeNode,
    {
        let mut child_ids = Vec::new();
        for node in nodes {
            if node.parent_id() == parent_id {
                child_ids.push(node.id());
                child_ids.extend(Self::get_child_ids(nodes, node.id()));
            }
        }
        child_ids
    }

    /// 获取所有父节点ID
    pub fn get_parent_ids<T>(nodes: &[T], child_id: i64) -> Vec<i64>
    where
        T: TreeNode,
    {
        let mut parent_ids = Vec::new();
        if let Some(node) = nodes.iter().find(|n| n.id() == child_id) {
            let parent_id = node.parent_id();
            if parent_id != 0 {
                parent_ids.push(parent_id);
                parent_ids.extend(Self::get_parent_ids(nodes, parent_id));
            }
        }
        parent_ids
    }
}

/// 通用树节点
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommonTreeNode<T> {
    /// 节点ID
    pub id: i64,
    /// 父节点ID
    pub parent_id: i64,
    /// 节点数据
    pub data: T,
    /// 子节点
    pub children: Vec<CommonTreeNode<T>>,
}

impl<T> TreeNode for CommonTreeNode<T> {
    fn id(&self) -> i64 {
        self.id
    }

    fn parent_id(&self) -> i64 {
        self.parent_id
    }

    fn children(&self) -> &Vec<Self> {
        &self.children
    }

    fn children_mut(&mut self) -> &mut Vec<Self> {
        &mut self.children
    }
}