computer: fenwick + per block profitability

crates/brk_computer/src/internal/algo/expanding_percentiles.rs

@@ -1,5 +1,7 @@
 use brk_types::StoredF32;
 
+use super::fenwick::FenwickTree;
+
 /// Fast expanding percentile tracker using a Fenwick tree (Binary Indexed Tree).
 ///
 /// Values are discretized to 10 BPS (0.1%) resolution and tracked in
@@ -9,9 +11,7 @@ use brk_types::StoredF32;
 /// - 0.1% value resolution (10 BPS granularity)
 #[derive(Clone)]
 pub(crate) struct ExpandingPercentiles {
-    /// Fenwick tree storing cumulative frequency counts.
-    /// 1-indexed: tree[0] is unused, tree[1..=TREE_SIZE] hold data.
-    tree: Vec<u32>,
+    tree: FenwickTree<u32>,
     count: u32,
 }
 
@@ -24,7 +24,7 @@ const TREE_SIZE: usize = (MAX_BPS / BUCKET_BPS) as usize + 1;
 impl Default for ExpandingPercentiles {
     fn default() -> Self {
         Self {
-            tree: vec![0u32; TREE_SIZE + 1], // 1-indexed
+            tree: FenwickTree::new(TREE_SIZE),
             count: 0,
         }
     }
@@ -36,16 +36,15 @@ impl ExpandingPercentiles {
     }
 
     pub fn reset(&mut self) {
-        self.tree.iter_mut().for_each(|v| *v = 0);
+        self.tree.reset();
         self.count = 0;
     }
 
-    /// Convert f32 ratio to bucket index (1-indexed for Fenwick).
+    /// Convert f32 ratio to 0-indexed bucket.
     #[inline]
     fn to_bucket(value: f32) -> usize {
         let bps = (value as f64 * 10000.0).round() as i32;
-        let bucket = (bps / BUCKET_BPS).clamp(0, TREE_SIZE as i32 - 1);
-        bucket as usize + 1
+        (bps / BUCKET_BPS).clamp(0, TREE_SIZE as i32 - 1) as usize
     }
 
     /// Bulk-load values in O(n + N) instead of O(n log N).
@@ -57,16 +56,9 @@ impl ExpandingPercentiles {
                 continue;
             }
             self.count += 1;
-            self.tree[Self::to_bucket(v)] += 1;
-        }
-        // Convert flat frequencies to Fenwick tree in O(N)
-        for i in 1..=TREE_SIZE {
-            let parent = i + (i & i.wrapping_neg());
-            if parent <= TREE_SIZE {
-                let val = self.tree[i];
-                self.tree[parent] += val;
-            }
+            self.tree.add_raw(Self::to_bucket(v), &1);
         }
+        self.tree.build_in_place();
     }
 
     /// Add a value. O(log N).
@@ -76,28 +68,7 @@ impl ExpandingPercentiles {
             return;
         }
         self.count += 1;
-        let mut i = Self::to_bucket(value);
-        while i <= TREE_SIZE {
-            self.tree[i] += 1;
-            i += i & i.wrapping_neg();
-        }
-    }
-
-    /// Find the bucket containing the k-th element (1-indexed k).
-    /// Uses the standard Fenwick tree walk-down in O(log N).
-    #[inline]
-    fn kth(&self, mut k: u32) -> usize {
-        let mut pos = 0;
-        let mut bit = 1 << (usize::BITS - 1 - TREE_SIZE.leading_zeros());
-        while bit > 0 {
-            let next = pos + bit;
-            if next <= TREE_SIZE && self.tree[next] < k {
-                k -= self.tree[next];
-                pos = next;
-            }
-            bit >>= 1;
-        }
-        pos + 1
+        self.tree.add(Self::to_bucket(value), &1);
     }
 
     /// Compute 6 percentiles in one call. O(6 × log N).
@@ -109,7 +80,9 @@ impl ExpandingPercentiles {
         }
         for (i, &q) in qs.iter().enumerate() {
             let k = ((q * self.count as f64).ceil() as u32).clamp(1, self.count);
-            out[i] = (self.kth(k) as u32 - 1) * BUCKET_BPS as u32;
+            // kth with 0-indexed k: k-1; result is 0-indexed bucket
+            let bucket = self.tree.kth(k - 1, |n| *n);
+            out[i] = bucket as u32 * BUCKET_BPS as u32;
         }
     }
 }

crates/brk_computer/src/internal/algo/fenwick.rs

@@ -0,0 +1,182 @@
+/// Trait for types that can be stored in a Fenwick tree.
+pub(crate) trait FenwickNode: Clone + Copy + Default {
+    fn add_assign(&mut self, other: &Self);
+}
+
+impl FenwickNode for u32 {
+    #[inline(always)]
+    fn add_assign(&mut self, other: &Self) {
+        *self += other;
+    }
+}
+
+/// Generic Fenwick tree (Binary Indexed Tree) over arbitrary node types.
+///
+/// Uses 0-indexed buckets externally; 1-indexed internally.
+/// Provides O(log N) point-update, prefix-sum, and kth walk-down.
+#[derive(Clone)]
+pub(crate) struct FenwickTree<N: FenwickNode> {
+    /// 1-indexed tree array. Position 0 is unused.
+    tree: Vec<N>,
+    size: usize,
+}
+
+impl<N: FenwickNode> FenwickTree<N> {
+    pub fn new(size: usize) -> Self {
+        Self {
+            tree: vec![N::default(); size + 1],
+            size,
+        }
+    }
+
+    #[inline]
+    pub fn size(&self) -> usize {
+        self.size
+    }
+
+    pub fn reset(&mut self) {
+        self.tree.fill(N::default());
+    }
+
+    /// Point-update: add `delta` to the node at `bucket` (0-indexed).
+    #[inline]
+    pub fn add(&mut self, bucket: usize, delta: &N) {
+        let mut i = bucket + 1;
+        while i <= self.size {
+            self.tree[i].add_assign(delta);
+            i += i & i.wrapping_neg();
+        }
+    }
+
+    /// Prefix sum of buckets [0, bucket] inclusive (0-indexed).
+    pub fn prefix_sum(&self, bucket: usize) -> N {
+        let mut result = N::default();
+        let mut i = bucket + 1;
+        while i > 0 {
+            result.add_assign(&self.tree[i]);
+            i -= i & i.wrapping_neg();
+        }
+        result
+    }
+
+    /// Find the 0-indexed bucket containing the k-th element (0-indexed k).
+    ///
+    /// `field_fn` extracts the relevant count field from a node.
+    /// The value type `V` must support comparison and subtraction
+    /// (works with `u32`, `i64`, `i128`).
+    #[inline]
+    pub fn kth<V, F>(&self, k: V, field_fn: F) -> usize
+    where
+        V: Copy + PartialOrd + std::ops::SubAssign,
+        F: Fn(&N) -> V,
+    {
+        debug_assert!(self.size > 0);
+        let mut pos = 0usize;
+        let mut remaining = k;
+        let mut bit = 1usize << (usize::BITS - 1 - self.size.leading_zeros() as u32);
+        while bit > 0 {
+            let next = pos + bit;
+            if next <= self.size {
+                let val = field_fn(&self.tree[next]);
+                if remaining >= val {
+                    remaining -= val;
+                    pos = next;
+                }
+            }
+            bit >>= 1;
+        }
+        pos // 0-indexed bucket
+    }
+
+    /// Write a raw frequency delta at a bucket. Does NOT maintain the Fenwick invariant.
+    /// Call [`build_in_place`] after all raw writes.
+    #[inline]
+    pub fn add_raw(&mut self, bucket: usize, delta: &N) {
+        self.tree[bucket + 1].add_assign(delta);
+    }
+
+    /// Convert raw frequencies (written via [`add_raw`]) into a valid Fenwick tree. O(size).
+    pub fn build_in_place(&mut self) {
+        for i in 1..=self.size {
+            let parent = i + (i & i.wrapping_neg());
+            if parent <= self.size {
+                let child = self.tree[i];
+                self.tree[parent].add_assign(&child);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn basic_add_and_prefix_sum() {
+        let mut tree = FenwickTree::<u32>::new(10);
+        tree.add(0, &3);
+        tree.add(1, &2);
+        tree.add(5, &7);
+
+        assert_eq!(tree.prefix_sum(0), 3);
+        assert_eq!(tree.prefix_sum(1), 5);
+        assert_eq!(tree.prefix_sum(4), 5);
+        assert_eq!(tree.prefix_sum(5), 12);
+        assert_eq!(tree.prefix_sum(9), 12);
+    }
+
+    #[test]
+    fn kth_walk_down() {
+        let mut tree = FenwickTree::<u32>::new(5);
+        // freq: [3, 2, 0, 5, 1]
+        tree.add(0, &3);
+        tree.add(1, &2);
+        tree.add(3, &5);
+        tree.add(4, &1);
+
+        // kth(0) = first element → bucket 0
+        assert_eq!(tree.kth(0u32, |n| *n), 0);
+        // kth(2) = 3rd element → bucket 0 (last of bucket 0)
+        assert_eq!(tree.kth(2u32, |n| *n), 0);
+        // kth(3) = 4th element → bucket 1
+        assert_eq!(tree.kth(3u32, |n| *n), 1);
+        // kth(4) = 5th element → bucket 1
+        assert_eq!(tree.kth(4u32, |n| *n), 1);
+        // kth(5) = 6th element → bucket 3 (bucket 2 is empty)
+        assert_eq!(tree.kth(5u32, |n| *n), 3);
+        // kth(10) = 11th element → bucket 4
+        assert_eq!(tree.kth(10u32, |n| *n), 4);
+    }
+
+    #[test]
+    fn build_in_place_matches_add() {
+        let mut tree_add = FenwickTree::<u32>::new(8);
+        tree_add.add(0, &5);
+        tree_add.add(2, &3);
+        tree_add.add(5, &7);
+        tree_add.add(7, &1);
+
+        let mut tree_bulk = FenwickTree::<u32>::new(8);
+        tree_bulk.add_raw(0, &5);
+        tree_bulk.add_raw(2, &3);
+        tree_bulk.add_raw(5, &7);
+        tree_bulk.add_raw(7, &1);
+        tree_bulk.build_in_place();
+
+        for i in 0..8 {
+            assert_eq!(
+                tree_add.prefix_sum(i),
+                tree_bulk.prefix_sum(i),
+                "mismatch at bucket {i}"
+            );
+        }
+    }
+
+    #[test]
+    fn reset_clears_all() {
+        let mut tree = FenwickTree::<u32>::new(10);
+        tree.add(3, &42);
+        tree.reset();
+        assert_eq!(tree.prefix_sum(9), 0);
+    }
+}

crates/brk_computer/src/internal/algo/mod.rs

@@ -1,12 +1,15 @@
 mod aggregation;
 mod drawdown;
+mod expanding_percentiles;
+mod fenwick;
 mod sliding_distribution;
 mod sliding_median;
-pub(crate) mod sliding_window;
-mod expanding_percentiles;
+mod sliding_window;
 
 pub(crate) use aggregation::*;
 pub(crate) use drawdown::*;
+pub(crate) use expanding_percentiles::*;
+pub(crate) use fenwick::*;
 pub(crate) use sliding_distribution::*;
 pub(crate) use sliding_median::*;
-pub(crate) use expanding_percentiles::*;
+pub(crate) use sliding_window::*;

crates/brk_computer/src/internal/mod.rs

@@ -1,5 +1,5 @@
 mod aggregate;
-pub(crate) mod algo;
+mod algo;
 mod containers;
 mod db_utils;
 mod derived;
@@ -8,7 +8,7 @@ mod from_tx;
 mod indexes;
 mod rolling;
 mod traits;
-pub mod transform;
+mod transform;
 mod value;
 
 pub(crate) use aggregate::*;