sum_segment_tree/

lib.rs

//! A fixed-capacity sum tree for weighted sampling, the structure used in
//! prioritized experience replay.
//!
//! Priorities live in the leaves of a complete binary tree packed into a flat
//! array. Every internal node stores the sum of its two children, so the total
//! weight and "find the leaf at cumulative weight `s`" are both
//! `O(log capacity)`. Writes are a ring buffer: once full, the oldest leaf is
//! overwritten.

/// A binary sum tree holding exactly `capacity` leaves.
///
/// The backing tree is padded up to a power of two internally; the extra
/// padding leaves stay at zero weight and are never returned by
/// [`SumTree::get`]. The exposed [`SumTree::capacity`] is exactly what you
/// requested.
#[derive(Clone, Debug)]
pub struct SumTree {
    // 1-indexed heap of length `2 * size`. `nodes[1]` is the root (total).
    // Leaf `i` lives at `nodes[size + i]`.
    nodes: Vec<f32>,
    // Number of leaf slots in the padded tree, a power of two >= capacity.
    size: usize,
    capacity: usize,
    write: usize,
    len: usize,
}

impl SumTree {
    /// Create a tree that holds exactly `capacity` leaves (at least one).
    pub fn new(capacity: usize) -> Self {
        let capacity = capacity.max(1);
        let size = capacity.next_power_of_two();
        SumTree {
            nodes: vec![0.0; 2 * size],
            size,
            capacity,
            write: 0,
            len: 0,
        }
    }

    /// Number of leaves the tree can hold.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Number of leaves currently written.
    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn is_full(&self) -> bool {
        self.len == self.capacity
    }

    /// Sum of every priority in the tree.
    #[inline]
    pub fn total(&self) -> f32 {
        self.nodes[1]
    }

    /// Read the priority stored at `index`.
    #[inline]
    pub fn priority(&self, index: usize) -> f32 {
        assert!(index < self.capacity, "index {index} out of bounds");
        self.nodes[self.size + index]
    }

    /// Set the priority at `index` and repair the sums up to the root.
    ///
    /// Ancestors are adjusted by the delta rather than recomputed from both
    /// children, halving the memory traffic per level. Over very many updates
    /// this can accumulate floating-point drift; call [`SumTree::rebuild`] to
    /// reset it.
    #[inline]
    pub fn update(&mut self, index: usize, priority: f32) {
        assert!(index < self.capacity, "index {index} out of bounds");
        assert!(priority >= 0.0, "priority must be non-negative");
        let nodes = &mut self.nodes;
        // `nodes.len()` is a power of two, so `idx & mask` is a no-op for our
        // in-range indices while proving `idx < len` to the compiler, which
        // drops the bounds checks without any `unsafe`.
        let mask = nodes.len() - 1;
        let mut i = index + self.size;
        let delta = priority - nodes[i & mask];
        nodes[i & mask] = priority;
        while i > 1 {
            i >>= 1;
            nodes[i & mask] += delta;
        }
    }

    /// Recompute every internal sum from the leaves, clearing any drift left by
    /// repeated [`SumTree::update`] calls.
    pub fn rebuild(&mut self) {
        for i in (1..self.size).rev() {
            self.nodes[i] = self.nodes[2 * i] + self.nodes[2 * i + 1];
        }
    }

    /// Append a priority at the next ring position, overwriting the oldest leaf
    /// when full. Returns the leaf index that was written.
    #[inline]
    pub fn push(&mut self, priority: f32) -> usize {
        let index = self.write;
        self.update(index, priority);
        self.write = (self.write + 1) % self.capacity;
        if self.len < self.capacity {
            self.len += 1;
        }
        index
    }

    /// Find the leaf whose cumulative-weight interval contains `s`, returning
    /// its index and priority. `s` is clamped to `[0, total]`. Returns `None`
    /// when the tree is empty or its total weight is zero.
    ///
    /// The returned index is always in `[0, len)`, so it is safe to use it in a
    /// side buffer sized to [`SumTree::capacity`].
    #[inline]
    pub fn get(&self, s: f32) -> Option<(usize, f32)> {
        let total = self.nodes[1];
        if self.len == 0 || total <= 0.0 {
            return None;
        }

        let mut s = if s.is_nan() { 0.0 } else { s.clamp(0.0, total) };
        let size = self.size;
        let nodes = self.nodes.as_slice();

        // `nodes.len()` is a power of two, so `idx & mask` is a no-op for our
        // in-range indices while proving `idx < len` to the compiler, which
        // elides the per-access bounds checks without any `unsafe`.
        let mask = nodes.len() - 1;
        let mut i = 1;

        // Branchless descent: the comparison drives the child index and the
        // subtraction directly, which avoids a per-level mispredicted branch.
        while i < size {
            let left = 2 * i;
            let left_sum = nodes[left & mask];
            let go_right = (s > left_sum) as usize;
            s -= left_sum * go_right as f32;
            i = left + go_right;
        }

        // Clamp to the last written slot. In exact arithmetic the descent always
        // lands on a written leaf, but this keeps the returned index in
        // `[0, len)` even if floating-point drift nudges a boundary sample past
        // the filled region.
        let leaf = (i - size).min(self.len - 1);
        Some((leaf, nodes[(size + leaf) & mask]))
    }
}