tensors-rs 0.1.2

Compact NumPy-like dense tensor primitives for safe numerical Rust.
Documentation
//! Reusable scratch allocation for numerical algorithms.

use crate::array2::Array2;
use crate::numeric::Float;

/// Reusable one-dimensional scratch storage.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ScratchBuffer<T> {
    data: Vec<T>,
}

impl<T> ScratchBuffer<T> {
    /// Create an empty scratch buffer.
    pub fn new() -> Self {
        Self { data: Vec::new() }
    }

    /// Number of currently initialized elements.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Whether the buffer is empty.
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Current allocation capacity.
    pub fn capacity(&self) -> usize {
        self.data.capacity()
    }

    /// Borrow the initialized buffer contents.
    pub fn as_slice(&self) -> &[T] {
        &self.data
    }

    /// Borrow the initialized buffer contents mutably.
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        &mut self.data
    }

    /// Drop initialized elements but keep allocated capacity.
    pub fn clear(&mut self) {
        self.data.clear();
    }
}

impl<T: Clone> ScratchBuffer<T> {
    /// Resize to `len` and fill new elements with `value`.
    pub fn resize(&mut self, len: usize, value: T) -> &mut [T] {
        self.data.resize(len, value);
        &mut self.data
    }
}

impl<T: Float> ScratchBuffer<T> {
    /// Resize to `len`, filling new elements with zero.
    pub fn zeros(&mut self, len: usize) -> &mut [T] {
        self.resize(len, T::zero())
    }
}

/// Workspace for repeated dense matrix algorithms.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Workspace<T> {
    buffers: Vec<ScratchBuffer<T>>,
}

impl<T> Workspace<T> {
    /// Create an empty workspace.
    pub fn new() -> Self {
        Self {
            buffers: Vec::new(),
        }
    }

    /// Number of scratch buffers currently owned by the workspace.
    pub fn buffers(&self) -> usize {
        self.buffers.len()
    }

    /// Get a named scratch buffer, creating it if needed.
    pub fn buffer(&mut self, index: usize) -> &mut ScratchBuffer<T> {
        while self.buffers.len() <= index {
            self.buffers.push(ScratchBuffer::new());
        }
        &mut self.buffers[index]
    }

    /// Get two distinct named scratch buffers.
    pub fn two_buffers_mut(
        &mut self,
        left: usize,
        right: usize,
    ) -> (&mut ScratchBuffer<T>, &mut ScratchBuffer<T>) {
        assert_ne!(left, right, "scratch buffer indices must be distinct");
        let max = left.max(right);
        while self.buffers.len() <= max {
            self.buffers.push(ScratchBuffer::new());
        }
        if left < right {
            let (head, tail) = self.buffers.split_at_mut(right);
            (&mut head[left], &mut tail[0])
        } else {
            let (head, tail) = self.buffers.split_at_mut(left);
            (&mut tail[0], &mut head[right])
        }
    }

    /// Clear initialized contents from every buffer while retaining capacity.
    pub fn clear(&mut self) {
        for buffer in &mut self.buffers {
            buffer.clear();
        }
    }
}

impl<T: Float> Workspace<T> {
    /// Allocate a zero-filled row-major matrix.
    pub fn matrix(&mut self, shape: [usize; 2]) -> Array2<T> {
        Array2::zeros(shape)
    }
}