yum/vfft

git clone https://git.yummers.dev/yum/vfft

yumRemove hot-path allocation from fft drainc1d59be2026-10-07

use num::complex::Complex;
use num::traits::{Float, FloatConst};

// Evaluates the base-`k` logarithm of `n`.
// Precondition: `n` is a power of `k`.
pub const fn log_k_of<const K: usize>(mut n: usize) -> usize {
    let mut res = 0;
    while n > 1 {
        n /= K;
        res += 1;
    }
    res
}

pub fn is_power_of_k(n: usize, k: usize) -> bool {
    match n {
        0 => false,
        1 => true,
        _ => n % k == 0 && is_power_of_k(n / k, k),
    }
}

pub fn usize_to_float<T: Float>(value: usize) -> T {
    num::cast(value).unwrap()
}

// Calculates the "twiddle factors" for an n-element FFT, aka all of the nth roots of unity.
pub fn precompute_twiddles<T: Float + FloatConst>(n: usize) -> Vec<Complex<T>> {
    let mut result = vec![Complex::<T>::new(T::zero(), T::zero()); n];

    let n_f64 = usize_to_float::<f64>(n);
    for i in 0..n {
        let tw_f64 = Complex::<f64>::cis(-f64::TAU() * usize_to_float::<f64>(i) / (n_f64));
        result[i] = Complex::new(T::from(tw_f64.re).unwrap(), T::from(tw_f64.im).unwrap());
    }

    result
}

// A FIFO data structure that stores a fixed number of elements. When it fills
// up, pushing overwrites the oldest element.
pub struct RingBuffer<T: Clone + Default> {
    data: Vec<T>,
    // Start points to the oldest element. End points to the next available
    // slot.
    start: usize,
    end: usize,
    // The number of elements currently in the buffer.
    n_elems: usize,
}

impl<T: Clone + Default> RingBuffer<T> {
    pub fn new(capacity: usize) -> Self {
        assert!(capacity > 0);
        Self {
            data: vec![T::default(); capacity],
            start: 0,
            end: 0,
            n_elems: 0,
        }
    }

    fn advance(&self, index: usize) -> usize {
        (index + 1) % self.data.len()
    }

    // Pop one element. Returns an error if the buffer is empty.
    pub fn pop(&mut self) -> anyhow::Result<&T> {
        if self.n_elems == 0 {
            std::hint::cold_path();
            return Err(anyhow::anyhow!("RingBuffer is empty"));
        }
        let result = &self.data[self.start];
        self.start = self.advance(self.start);
        self.n_elems -= 1;
        Ok(result)
    }

    // Pop elements until `cb` return false. If `cb` returns an error, we
    // propagate it immediately and do not remove the element.
    pub fn pop_each(
        &mut self,
        mut cb: impl FnMut(&T) -> anyhow::Result<bool>,
    ) -> anyhow::Result<()> {
        while self.n_elems > 0 {
            let keep_going = cb(&self.data[self.start])?;
            self.pop()?;
            if !keep_going {
                break;
            }
        }
        Ok(())
    }

    // Push one element.
    pub fn push(&mut self, value: T) {
        self.data[self.end] = value;
        self.end = self.advance(self.end);
        if self.n_elems == self.data.len() {
            self.start = self.advance(self.start);
        } else {
            self.n_elems += 1;
        }
    }

    // Push a slice of elements.
    pub fn push_slice(&mut self, values: &[T]) {
        let capacity = self.data.len();
        // Ensure that `values` is not longer than our ringbuf capacity.
        let values = &values[values.len().saturating_sub(capacity)..];

        // Calculate number of elements to copy over after the end pointer.
        let n_after_end = values.len().min(capacity - self.end);
        let values_after_end = &values[..n_after_end];
        self.data[self.end..self.end + n_after_end].clone_from_slice(values_after_end);

        let values_before_start = &values[n_after_end..];
        self.data[..values_before_start.len()].clone_from_slice(values_before_start);

        let n_overwritten = values.len().saturating_sub(capacity - self.n_elems);
        self.start = (self.start + n_overwritten) % capacity;
        self.end = (self.end + values.len()) % capacity;
        self.n_elems += values.len() - n_overwritten;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ringbuf_tests() {
        // Sanity.
        {
            let mut rb = RingBuffer::<i32>::new(3);
            assert!(rb.pop().is_err());
            rb.push(1);
            assert_eq!(*rb.pop().unwrap(), 1);
            rb.push(2);
            rb.push(3);
            rb.push(4);
            assert_eq!(*rb.pop().unwrap(), 2);
            assert_eq!(*rb.pop().unwrap(), 3);
            assert_eq!(*rb.pop().unwrap(), 4);
            assert!(rb.pop().is_err());
            rb.push_slice(&[5, 6, 7, 8]);
            assert_eq!(*rb.pop().unwrap(), 6);
            assert_eq!(*rb.pop().unwrap(), 7);
            assert_eq!(*rb.pop().unwrap(), 8);
            assert!(rb.pop().is_err());
        }
        // Callback access.
        {
            let mut rb = RingBuffer::<i32>::new(3);
            rb.push(1);
            rb.push(2);
            rb.push(3);
            rb.push(4);
            let mut iter = 2;
            assert!(
                rb.pop_each(|x| {
                    assert_eq!(*x, iter);
                    iter += 1;
                    Ok(true)
                })
                .is_ok()
            );
        }
        // Size-1 edge cases.
        {
            let mut rb = RingBuffer::<i32>::new(1);
            rb.push(1);
            assert_eq!(*rb.pop().unwrap(), 1);
            assert!(rb.pop().is_err());
        }
    }
}