Add trait ConstMultiDistribution

Author: dhardy

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Additions
 - `MultiDistribution` trait to sample more efficiently from multi-dimensional distributions (#18)
+- `ConstMultiDistribution` trait as support for fixed-dimension distributions (#29)
 
 ### Changes
 - Moved `Dirichlet` into the new `multi` module and implement `MultiDistribution` for it (#18)

src/multi/mod.rs

@@ -8,19 +8,40 @@
 
 //! Contains Multi-dimensional distributions.
 //!
-//! We provide a trait `MultiDistribution` which allows to sample from a multi-dimensional distribution without extra allocations.
-//! All multi-dimensional distributions implement `MultiDistribution` instead of the `Distribution` trait.
+//! The trait [`MultiDistribution`] supports multi-dimensional sampling without
+//! allocating a [`Vec`](std::vec::Vec) for each sample.
+//! [`ConstMultiDistribution`] is an extension for distributions with constant
+//! dimension.
+//!
+//! Multi-dimensional distributions implement `MultiDistribution<T>` and (where
+//! the dimension is fixed) `ConstMultiDistribution<T>` for some scalar type
+//! `T`. They may also implement `Distribution<Vec<T>>` and (where the
+//! dimension, `N`, is fixed) `Distribution<[T; N]>`.
 
 use rand::Rng;
 
 /// A standard abstraction for distributions with multi-dimensional results
+///
+/// Implementations may also implement `Distribution<Vec<T>>`.
 pub trait MultiDistribution<T> {
-    /// returns the length of one sample (dimension of the distribution)
+    /// The length of a sample (dimension of the distribution)
     fn sample_len(&self) -> usize;
-    /// samples from the distribution and writes the result to `output`
+
+    /// Sample a multi-dimensional result from the distribution
+    ///
+    /// The result is written to `output`. Implementations should assert that
+    /// `output.len()` equals the result of [`Self::sample_len`].
     fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [T]);
 }
 
+/// An extension of [`MultiDistribution`] for multi-dimensional distributions of fixed dimension
+///
+/// Implementations may also implement `Distribution<[T; SAMPLE_LEN]>`.
+pub trait ConstMultiDistribution<T>: MultiDistribution<T> {
+    /// Constant sample length (dimension of the distribution)
+    const SAMPLE_LEN: usize;
+}
+
 macro_rules! distribution_impl {
     ($scalar:ident) => {
         fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<$scalar> {
@@ -32,6 +53,22 @@ macro_rules! distribution_impl {
     };
 }
 
+#[allow(unused)]
+macro_rules! const_distribution_impl {
+    ($scalar:ident) => {
+        fn sample<R: Rng + ?Sized>(
+            &self,
+            rng: &mut R,
+        ) -> [$scalar; <Self as crate::multi::MultiDistribution>::SAMPLE_LEN] {
+            use crate::multi::MultiDistribution;
+            let mut buf =
+                [Default::default(); <Self as crate::multi::MultiDistribution>::SAMPLE_LEN];
+            self.sample_to_slice(rng, &mut buf);
+            buf
+        }
+    };
+}
+
 pub use dirichlet::Dirichlet;
 
 mod dirichlet;