sum_segment_tree/
lib.rs

1//! A fixed-capacity sum tree for weighted sampling, the structure used in
2//! prioritized experience replay.
3//!
4//! Priorities live in the leaves of a complete binary tree packed into a flat
5//! array. Every internal node stores the sum of its two children, so the total
6//! weight and "find the leaf at cumulative weight `s`" are both
7//! `O(log capacity)`. Writes are a ring buffer: once full, the oldest leaf is
8//! overwritten.
9
10/// A binary sum tree over `capacity` leaves. The capacity is rounded up to a
11/// power of two so the tree stays complete.
12#[derive(Clone, Debug)]
13pub struct SumTree {
14    // 1-indexed heap. `nodes[1]` is the root (total). Leaf `i` lives at
15    // `nodes[capacity + i]`.
16    nodes: Vec<f32>,
17    capacity: usize,
18    write: usize,
19    len: usize,
20}
21
22impl SumTree {
23    /// Create a tree whose capacity is `capacity` rounded up to a power of two
24    /// (at least one).
25    pub fn new(capacity: usize) -> Self {
26        let capacity = capacity.max(1).next_power_of_two();
27        SumTree {
28            nodes: vec![0.0; 2 * capacity],
29            capacity,
30            write: 0,
31            len: 0,
32        }
33    }
34
35    /// Number of leaves the tree can hold.
36    pub fn capacity(&self) -> usize {
37        self.capacity
38    }
39
40    /// Number of leaves currently written.
41    pub fn len(&self) -> usize {
42        self.len
43    }
44
45    pub fn is_empty(&self) -> bool {
46        self.len == 0
47    }
48
49    pub fn is_full(&self) -> bool {
50        self.len == self.capacity
51    }
52
53    /// Sum of every priority in the tree.
54    #[inline]
55    pub fn total(&self) -> f32 {
56        self.nodes[1]
57    }
58
59    /// Read the priority stored at `index`.
60    #[inline]
61    pub fn priority(&self, index: usize) -> f32 {
62        assert!(index < self.capacity, "index {index} out of bounds");
63        self.nodes[self.capacity + index]
64    }
65
66    /// Set the priority at `index` and repair the sums up to the root.
67    ///
68    /// Ancestors are adjusted by the delta rather than recomputed from both
69    /// children, halving the memory traffic per level. Over very many updates
70    /// this can accumulate floating-point drift; call [`SumTree::rebuild`] to
71    /// reset it.
72    #[inline]
73    pub fn update(&mut self, index: usize, priority: f32) {
74        assert!(index < self.capacity, "index {index} out of bounds");
75        assert!(priority >= 0.0, "priority must be non-negative");
76        let nodes = &mut self.nodes;
77        // `nodes.len()` is a power of two, so `idx & mask` is a no-op for our
78        // in-range indices while proving `idx < len` to the compiler, which
79        // drops the bounds checks without any `unsafe`.
80        let mask = nodes.len() - 1;
81        let mut i = index + self.capacity;
82        let delta = priority - nodes[i & mask];
83        nodes[i & mask] = priority;
84        while i > 1 {
85            i >>= 1;
86            nodes[i & mask] += delta;
87        }
88    }
89
90    /// Recompute every internal sum from the leaves, clearing any drift left by
91    /// repeated [`SumTree::update`] calls.
92    pub fn rebuild(&mut self) {
93        for i in (1..self.capacity).rev() {
94            self.nodes[i] = self.nodes[2 * i] + self.nodes[2 * i + 1];
95        }
96    }
97
98    /// Append a priority at the next ring position, overwriting the oldest leaf
99    /// when full. Returns the leaf index that was written.
100    #[inline]
101    pub fn push(&mut self, priority: f32) -> usize {
102        let index = self.write;
103        self.update(index, priority);
104        self.write = (self.write + 1) % self.capacity;
105        if self.len < self.capacity {
106            self.len += 1;
107        }
108        index
109    }
110
111    /// Find the leaf whose cumulative-weight interval contains `s`, returning
112    /// its index and priority. `s` is clamped to `[0, total]`. Returns `None`
113    /// when the tree is empty or its total weight is zero.
114    #[inline]
115    pub fn get(&self, s: f32) -> Option<(usize, f32)> {
116        let total = self.nodes[1];
117        if self.len == 0 || total <= 0.0 {
118            return None;
119        }
120
121        let mut s = if s.is_nan() { 0.0 } else { s.clamp(0.0, total) };
122        let capacity = self.capacity;
123        let nodes = self.nodes.as_slice();
124
125        // `nodes.len()` is a power of two, so `idx & mask` is a no-op for our
126        // in-range indices while proving `idx < len` to the compiler, which
127        // elides the per-access bounds checks without any `unsafe`.
128        let mask = nodes.len() - 1;
129        let mut i = 1;
130
131        // Branchless descent: the comparison drives the child index and the
132        // subtraction directly, which avoids a per-level mispredicted branch.
133        while i < capacity {
134            let left = 2 * i;
135            let left_sum = nodes[left & mask];
136            let go_right = (s > left_sum) as usize;
137            s -= left_sum * go_right as f32;
138            i = left + go_right;
139        }
140
141        Some((i - capacity, nodes[i & mask]))
142    }
143}
sum_segment_tree/lib.rs

sum_segment_tree/
lib.rs
