count_tracker.rs

use slotmap::{SlotMap, SecondaryMap, new_key_type, Key};
use std::collections::HashSet;

new_key_type! {
    struct BucketKey;
}

new_key_type! {
    struct IndexKey;
}

#[derive(Clone, Copy, Debug)]
struct CountBucket {
    count: i64,
    prev: BucketKey,
    next: BucketKey,
}


#[derive(Clone, Debug)]
pub struct CountTracker {
    counts: Vec<i64>,
    in_bucket: Vec<(BucketKey, IndexKey)>,
    buckets: SlotMap<BucketKey, CountBucket>,
    indices: SecondaryMap<BucketKey, SlotMap<IndexKey, usize>>,
    zero_bucket: BucketKey,
    min_bucket: BucketKey,
    max_bucket: BucketKey,
}


impl CountTracker {
    pub fn new(n: usize) -> Self {
        let mut buckets = SlotMap::with_key();
        let zero_bucket = buckets.insert(CountBucket {
            count: 0,
            prev: BucketKey::null(),
            next: BucketKey::null(),
        });

        let mut indices = SecondaryMap::new();
        let mut zero_bucket_indices = SlotMap::with_key();
        let in_bucket = (0..n).map(|i| (zero_bucket, zero_bucket_indices.insert(i))).collect();
        indices.insert(zero_bucket, zero_bucket_indices);

        Self {
            counts: vec![0; n],
            in_bucket,
            buckets,
            indices,
            zero_bucket,
            min_bucket: zero_bucket,
            max_bucket: zero_bucket,
        }
    }

    fn set_prev(&mut self, to_update: BucketKey, new_prev: BucketKey) {
        if let Some(update_bucket) = self.buckets.get_mut(to_update) {
            update_bucket.prev = new_prev;
        } else {
            self.max_bucket = new_prev;
        }
    }

    fn set_next(&mut self, to_update: BucketKey, new_next: BucketKey) {
        if let Some(update_bucket) = self.buckets.get_mut(to_update) {
            update_bucket.next = new_next;
        } else {
            self.min_bucket = new_next;
        }
    }

    fn get_inc_bucket(&mut self, bucket_key: BucketKey) -> BucketKey {
        let CountBucket { count, next, .. } = self.buckets[bucket_key];
        if let Some(next_bucket) = self.buckets.get(next) {
            if next_bucket.count == count + 1 {
                return next;
            }
        }

        // Create new bucket and insert.
        let new_bucket = CountBucket {
            count: count + 1,
            prev: bucket_key,
            next: next,
        };

        let new_bucket_key = self.buckets.insert(new_bucket);
        self.set_prev(next, new_bucket_key);
        self.set_next(bucket_key, new_bucket_key);
        self.indices.insert(new_bucket_key, SlotMap::with_key());
        new_bucket_key
    }

    fn get_dec_bucket(&mut self, bucket_key: BucketKey) -> BucketKey {
        let CountBucket { count, prev, .. } = self.buckets[bucket_key];
        if let Some(prev_bucket) = self.buckets.get(prev) {
            if prev_bucket.count == count - 1 {
                return prev;
            }
        }

        // Create new bucket and insert.
        let new_bucket = CountBucket {
            count: count - 1,
            prev: prev,
            next: bucket_key,
        };

        let new_bucket_key = self.buckets.insert(new_bucket);
        self.set_prev(bucket_key, new_bucket_key);
        self.set_next(prev, new_bucket_key);
        self.indices.insert(new_bucket_key, SlotMap::with_key());
        new_bucket_key
    }

    fn maybe_collapse_bucket(&mut self, bucket_key: BucketKey) {
        if self.indices[bucket_key].len() > 0 || bucket_key == self.zero_bucket {
            return;
        }

        let CountBucket { prev, next, .. } = self.buckets[bucket_key];
        self.set_prev(next, prev);
        self.set_next(prev, next);
        self.buckets.remove(bucket_key);
        self.indices.remove(bucket_key);
    }

    pub fn inc(&mut self, i: usize) {
        let (old_bucket_key, old_index_key) = self.in_bucket[i];
        let new_bucket_key = self.get_inc_bucket(old_bucket_key);
        self.indices[old_bucket_key].remove(old_index_key);
        let new_index_key = self.indices[new_bucket_key].insert(i);
        self.maybe_collapse_bucket(old_bucket_key);
        self.in_bucket[i] = (new_bucket_key, new_index_key);
        self.counts[i] += 1;
    }

    pub fn dec(&mut self, i: usize) {
        let (old_bucket_key, old_index_key) = self.in_bucket[i];
        let new_bucket_key = self.get_dec_bucket(old_bucket_key);
        self.indices[old_bucket_key].remove(old_index_key);
        let new_index_key = self.indices[new_bucket_key].insert(i);
        self.maybe_collapse_bucket(old_bucket_key);
        self.in_bucket[i] = (new_bucket_key, new_index_key);
        self.counts[i] -= 1;
    }

    pub fn zero(&mut self, i: usize) {
        let (old_bucket_key, old_index_key) = self.in_bucket[i];
        self.indices[old_bucket_key].remove(old_index_key);
        let new_index_key = self.indices[self.zero_bucket].insert(i);
        self.maybe_collapse_bucket(old_bucket_key);
        self.in_bucket[i] = (self.zero_bucket, new_index_key);
        self.counts[i] = 0;
    }

    pub fn get(&self, i: usize) -> i64 {
        self.counts[i]
    }

    pub fn min_idx(&self) -> usize {
        let mut bucket = self.min_bucket;
        if self.indices[bucket].len() == 0 { // In case of empty zero bucket.
            bucket = self.buckets[bucket].next;
        }
        *self.indices[bucket].values().next().unwrap()
    }
    
    pub fn max_idx(&self) -> usize {
        let mut bucket = self.max_bucket;
        if self.indices[bucket].len() == 0 { // In case of empty zero bucket.
            bucket = self.buckets[bucket].prev;
        }
        *self.indices[bucket].values().next().unwrap()
    }
}


#[cfg(test)]
mod tests {
    #[test]
    fn test_count_tracker() {
        use super::CountTracker;
        use rand::Rng;

        for n in 1..10 {
            let mut c = CountTracker::new(n);
            let mut v = vec![0i64; n];
            
            let mut rng = rand::thread_rng();
            for _ in 0..10000 {
                let idx = rng.gen_range(0, n);
                if rng.gen() {
                    c.inc(idx);
                    v[idx] += 1;
                } else {
                    c.dec(idx);
                    v[idx] -= 1;
                }
                if rng.gen::<f32>() < 0.05 {
                    c.zero(idx);
                    v[idx] = 0;
                }
                assert!((0..n).all(|i| c.get(i) == v[i]));
                assert!(c.get(c.min_idx()) == *v.iter().min().unwrap());
                assert!(c.get(c.max_idx()) == *v.iter().max().unwrap());
            }
        }
    }
}