use crate::fmt::format_micros;
use std::fmt;
use std::fmt::Display;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Default)]
pub struct Average {
sum_and_count: AtomicU64,
count: AtomicU64,
}
impl Clone for Average {
fn clone(&self) -> Self {
Average {
sum_and_count: AtomicU64::new(self.sum_and_count.load(Ordering::Relaxed)),
count: AtomicU64::new(self.count.load(Ordering::Relaxed)),
}
}
}
impl Average {
pub fn new() -> Average {
Average {
sum_and_count: AtomicU64::new(0),
count: AtomicU64::new(0),
}
}
fn sum_and_count(&self) -> (i32, i32) {
let last_sum_and_count = self.sum_and_count.load(Ordering::Relaxed);
let count = (last_sum_and_count & 0xFFFFFFFF) as i32;
let sum = ((last_sum_and_count >> 32) & 0xFFFFFFFF) as i32;
(sum, count)
}
pub fn add(&self, value: i32) {
let _ = self.count.fetch_add(1, Ordering::Relaxed);
let (mut sum, mut count) = self.sum_and_count();
while count > 100 || sum as i64 + value as i64 > i32::MAX as i64 {
sum = count / 2 * sum / count;
count /= 2;
}
sum += value;
count += 1;
let next_sum_and_count = ((sum as u64 & 0xFFFFFFFF) << 32) | (count as u64 & 0xFFFFFFFF);
self.sum_and_count
.store(next_sum_and_count, Ordering::Relaxed);
}
pub fn count(&self) -> u64 {
self.count.load(Ordering::Relaxed)
}
pub fn avg(&self) -> i32 {
let (sum, count) = self.sum_and_count();
if sum == 0 {
0
} else {
sum / count
}
}
}
impl Display for Average {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
format_micros(self.avg(), f)?;
write!(f, " ({})", self.count())?;
Ok(())
}
}
#[cfg(test)]
mod test {
use crate::average::Average;
#[test]
fn empty_average_is_properly_initialized() {
let avg = Average::new();
assert_eq!(avg.avg(), 0);
assert_eq!(avg.count(), 0);
}
#[test]
fn average_with_some_values_works() {
let avg = Average::new();
for i in 1..=10 {
avg.add(i);
}
assert_eq!(avg.avg(), 5);
assert_eq!(avg.count(), 10);
}
#[test]
fn formatting_average_works() {
let avg = Average::new();
avg.add(10_123);
assert_eq!(format!("{}", avg), "10.1 ms (1)");
}
#[test]
fn average_with_many_values_keeps_count() {
let avg = Average::new();
for i in 1..=1000 {
avg.add(i);
}
assert_eq!(avg.avg(), 928);
assert_eq!(avg.count(), 1000);
}
#[test]
fn average_overflows_sanely() {
{
let avg = Average::new();
avg.add(i32::MAX);
assert_eq!(avg.avg(), i32::MAX);
avg.add(i32::MAX);
assert_eq!(avg.avg(), i32::MAX);
avg.add(i32::MAX / 2);
avg.add(i32::MAX / 2);
assert_eq!(avg.avg(), i32::MAX / 2);
}
{
let avg = Average::new();
avg.add(10);
avg.add(i32::MAX - 50);
avg.add(60);
let average_before_overflow = (i32::MAX - 50 + 10) / 2;
let expected_average = (average_before_overflow + 60) / 2;
assert_eq!(avg.avg(), expected_average);
}
}
}