rust 树数据结构库 indextree 使用简介

由于 rust 特殊的引用规则存在, 创建一个树形结构也变成了一件困难的事情.

本文介绍 indextree, 一个基于 arena 的树形结构库

该库的用法简单直观, 首先创建节点池 arena, 然后在上面创建新的节点, 节点由 NodeId 表示,
然后可以根据 NodeId 来索引到真正存放数据的 Node, 也可以根据 NodeId 来进行树的修改和遍历.

作者给出的使用例子如下:

use indextree::Arena;

// Create a new arena
let arena = &mut Arena::new();

// Add some new nodes to the arena
let a = arena.new_node(1);
let b = arena.new_node(2);

// Append a to b
assert!(a.append(b, arena).is_ok());
assert_eq!(b.ancestors(arena).into_iter().count(), 2);

当然实际使用中无法静态的获取树的结构, 可以使用 stack 的 push/pop 操作来动态的生成一颗树

以文档解析 token 树为例

use indextree::{Arena, NodeId};

// 树存储的数据
struct OperationType {
  ...
}

pub struct OperationTree {
  arena: Arena<OperationType>,        // 节点池 arena
  root: NodeId,                       // 根节点
  cur: NodeId,                        // 当前操作节点
  child_count: HashMap<NodeId, usize>,// 节点对应还有多少子节点未出栈
}

然后定义一个 stack 的 trait

pub trait WritableStack { // 这里也可以写成范型的
  fn push(&mut self, node_type: OperationType);
  fn pop(&mut self);
}

实现的时候维护 child_count, 并切换 cur 即可, 但需要注意文档树中有叶子节点, 为了使用方便, 叶子节点入栈时, 不更新 cur.
(也可以出栈完立刻出栈, 但使用时容易出错)

impl WritableStack for OperationTree {
  fn push(&mut self, node_type: OperationType) {
    // println!("push: {:?}", node_type);
    let has_child = node_type.has_child();
    let new_node = self.arena.new_node(node_type);
    // append, 插入一个当前节点的子节点.
    self.cur.append(new_node, &mut self.arena);
    // 更新插入节点的父亲
    self.child_count.entry(self.cur).and_modify(|e| *e += 1);
    if has_child {
      // 非节子, 替换cur
      self.cur = new_node;
      self.child_count.insert(self.cur, 0);
    }
  }

  fn pop(&mut self) {
    // println!("pop: {:?}", self.arena.get(self.cur).unwrap().get());
    if self.child_count[&self.cur] == 0 {
      // 当前节点已经没有儿子, 切换cur到父亲.
      let mut ancestors = self.cur.ancestors(&self.arena);
      ancestors.next();
      let next_cur = ancestors.next();
      if next_cur.is_some() {
        self.cur = next_cur.unwrap();
      }
    }
    self.child_count.entry(self.cur).and_modify(|e| *e -= 1);
  }
}

构建完树后, 我们要进行节点遍历使用, indextree 有非常好的 iterator 支持, 可以在当前树链上遍历、横向遍历、子节点遍历. 可以看它的节点定义, 保存了所有相关联的 NodeId.

pub struct Node<T> {
    // Keep these private (with read-only accessors) so that we can keep them
    // consistent. E.g. the parent of a node’s child is that node.
    pub(crate) parent: Option<NodeId>,
    pub(crate) previous_sibling: Option<NodeId>, 
    pub(crate) next_sibling: Option<NodeId>, 
    pub(crate) first_child: Option<NodeId>,
    pub(crate) last_child: Option<NodeId>,
    pub(crate) stamp: NodeStamp,
    /// The actual data which will be stored within the tree.
    pub(crate) data: NodeData<T>,
}

数据和访问函数的对应为:

parent: ancestors
previous_sibling: preceding_siblings
next_sibling: following_siblings
first_child: children
last_child: reversed_children

使用这些接口 再调用 iterator 相应操作就可以灵活的对树进行操作. 一个递归打印的实现如下.

pub fn dfs_fmt(&self, node_id: NodeId, f: &mut Formatter<'_>, dep: usize) -> std::fmt::Result {
  let mut ret = write!(f, "{}", "-".repeat(dep));
  if ret.is_err() {
    return ret;
  }
  ret = writeln!(f, "{:?}\n", self.get_node(node_id).get());
  if ret.is_err() {
    return ret;
  }
  let children = node_id.children(&self.arena);
  for child in children {
    let ret = self.dfs_fmt(child, f, dep + 1);
    if ret.is_err() {
      return ret;
    }
  }
  return Ok(());
}