Rust：Arena 分配器使用实践

在 Rust 开发中，常常会遇到需要生成循环依赖的数据的情况。比如控制流图，它不存在某个至高无上的所有者，所有结点都是公平存在的。那么这种情况下，如果用 Rust 自带所有权机制来管理这些结点，几乎寸步难行。你不能使用 Rc，因为不存在明确的强弱引用。实际上对于一个循环的控制流图，必然是要产生循环的强引用的。

Arena 分配的思想是，对于这种场景，产生的对象往往可以一次性全部销毁。比如我们在内存中构建一个图，然后在这个图上跑各种算法。最后一次性删掉整个图。那么就没有必要进行琐碎的内存管理，可以直接申请一块连续大内存，即 Arena。（还有一种常见方法是申请个二维数组。但是当节点数量较多时，这种方法的内存占用不可接受。）

实现 Arena 分配的方法很多。一种方法是泵分配（Bump Allocation），即申请一块连续内存，并用一个空闲指针指向开头。然后每次需要一块小内存，就指针就往后移动，划出分配出的区域。当区域全部内存用光，则开辟一个更大的区域，然后把这个区域复制过去。

bumpalo 是 Bump Allocation 在 Rust 中的一种实现。

基本使用方法如下：

1    fn it_works() {
2        use bumpalo::Bump;
3        struct Node {
4            val: i32,
5        }
6        let bump = Bump::new();
7        let node_ref = bump.alloc(Node { val: 0 });
8        node_ref.val += 1;
9    }

bump.alloc 返回一个互斥可变引用。

不过如果就是这么简单的情景，就没必要用 Arena 了。

下面介绍使用我们需要用到同作者开发的 generational-arena。

1cargo add generational-arena

代码：

 1use generational_arena::Index;
 2#[test]
 3fn it_works() {
 4    use generational_arena::Arena;
 5    struct Node {
 6        val: i32,
 7        edges: Vec<Index>,
 8    }
 9    let mut nodes :Arena<Node> = Arena::new();
10  
11    /*
12    1 --> 2 
13    |     ^ 
14    |     | 
15    \ --> 4 
16    */
17    let n1 = nodes.insert(Node {val: 1, edges: vec![]});
18    let n2 = nodes.insert(Node {val: 2, edges: vec![]});
19    let n4 = nodes.insert(Node {val: 4, edges: vec![]});
20    // connect 1->2, 1->4, 4->2
21    nodes[n1].edges.push(n2);
22    nodes[n1].edges.push(n4);
23    nodes[n4].edges.push(n2);
24   /*
25    1 --> 2 --> 3
26    |     ^     |
27    |     |     v
28    \ --> 4 --> 5
29    */
30    let n3 = nodes.insert(Node {val: 3, edges: vec![]});
31    let n5 = nodes.insert(Node {val: 5, edges: vec![]});
32    // connect 2->3, 3->5, 4->5
33    nodes[n2].edges.push(n3);
34    nodes[n3].edges.push(n5);
35    nodes[n4].edges.push(n5);
36
37}

上面的代码，我们分两阶段构建图。第一阶段构建了一个简单的图，第二阶段构建了一个更复杂的图，加入了环。

由于使用的是 Index 而不是引用，使得我们不用处理复杂的所有权关系。而且这里的索引不是使用连续数字，比起用一个简单的动态数组存个坐标的方式来说，这样的好处是，我们可以在删除节点的时候，不用移动后面的节点，从而获得更低的时间复杂度。

可以用下面的代码对图进行 DFS：

 1    // DFS
 2    let mut stack = vec![];
 3    let mut visited = vec![];
 4    stack.push(n1);
 5    while !stack.is_empty() {
 6        let node = stack.pop().unwrap();
 7        if visited.contains(&node) {
 8            continue;
 9        }
10        visited.push(node);
11        for edge in &nodes[node].edges {
12            stack.push(*edge);
13        }
14    }
15    println!("visited: {:?}", visited.iter().map(|i| nodes[*i].val).collect::<Vec<_>>());

输出：

1visited: [1, 4, 5, 2, 3]

那么如果删除结点呢？

1    // remove node 2
2    nodes.remove(n2);
3    // connect 5->3
4    nodes[n5].edges.push(n3);
5    // DFS
6    // ...

出问题了：

1thread 'tests::it_works' panicked at 'No element at index', src/lib.rs:70:26

因为这种情况下我们需要手动管理依赖关系。删除 n2 时，与 n2 相连的边也需要删除。这时候我们就需要用到 generational-arena 的 Arena::retain 方法。

retain 的参数是一个断言。所有断言执行后为 false 的元素都会被删除。

1    // remove node 2
2    nodes.retain(|_, node| {
3        node.edges.retain(|edge| *edge != n2); // 删除所有连至 n2 的边
4        node.val != 2 // 删除 val 为 2 的结点
5    });
6    // connect 5->3
7    nodes[n5].edges.push(n3);
8    // DFS
9    // ...

运行结果

1visited: [1, 4, 5, 3]
完整代码 

 1#[cfg(test)]
 2mod tests {
 3    use generational_arena::Index;
 4
 5    #[test]
 6    fn it_works() {
 7        use generational_arena::Arena;
 8        struct Node {
 9            val: i32,
10            edges: Vec<Index>,
11        }
12        let mut nodes: Arena<Node> = Arena::new();
13
14        /*
15        1 --> 2
16        |     ^
17        |     |
18        \ --> 4
19        */
20        let n1 = nodes.insert(Node {
21            val: 1,
22            edges: vec![],
23        });
24        let n2 = nodes.insert(Node {
25            val: 2,
26            edges: vec![],
27        });
28        let n4 = nodes.insert(Node {
29            val: 4,
30            edges: vec![],
31        });
32        // connect 1->2, 1->4, 4->2
33        nodes[n1].edges.push(n2);
34        nodes[n1].edges.push(n4);
35        nodes[n4].edges.push(n2);
36        /*
37        1 --> 2 --> 3
38        |     ^ .   |
39        |     |     v
40        \ --> 4 --> 5
41        */
42        let n3 = nodes.insert(Node {
43            val: 3,
44            edges: vec![],
45        });
46        let n5 = nodes.insert(Node {
47            val: 5,
48            edges: vec![],
49        });
50        // connect 2->3, 3->5, 4->5
51        nodes[n2].edges.push(n3);
52        nodes[n3].edges.push(n5);
53        nodes[n4].edges.push(n5);
54
55        // DFS
56        // remove node 2
57        nodes.retain(|_, node| {
58            node.edges.retain(|edge| *edge != n2);
59            node.val != 2
60        });
61        // connect 5->3
62        nodes[n5].edges.push(n3);
63        // DFS
64        // ...
65        let mut stack = vec![];
66        let mut visited = vec![];
67        stack.push(n1);
68        while !stack.is_empty() {
69            let node = stack.pop().unwrap();
70            if visited.contains(&node) {
71                continue;
72            }
73            visited.push(node);
74            for edge in &nodes[node].edges {
75                stack.push(*edge);
76            }
77        }
78        println!(
79            "visited: {:?}",
80            visited.iter().map(|i| nodes[*i].val).collect::<Vec<_>>()
81        );
82    }
83}

不过，这种处理方式代价是很高昂的，因为我们需要在每次删除结点时都要遍历一遍所有结点，删除所有与其相连的边。如果我们的图非常大，这种代价是不可接受的。所以另一种方式是维持一份双向引用。

1struct Node {
2    val: i32,
3    edges: Vec<Index>,
4    back_edges: Vec<Index>,
5}

这样我们，删除结点时，只需要遍历其所有的 back_edges 进行操作。具体代码略。

最近发布

要查看全部文章，请点击右上角“归档”

随机文章