Change sample to accept all distributions

Author: junniest

phylo/src/random/fake_random.rs

@@ -1,7 +1,7 @@
 use std::any::{Any, TypeId};
 use std::sync::Mutex;
 
-use rand::distributions::{Standard, WeightedIndex};
+use rand::distributions::Standard;
 use rand::prelude::Distribution;
 
 use crate::random::RandomSource;
@@ -203,14 +203,22 @@ impl RandomSource for FakeGenerator {
         *self.seed.lock().unwrap() = seed;
     }
 
-    fn sample(&self, _dist: &WeightedIndex<f64>) -> usize {
-        self.next_u64() as usize
+    fn sample<D, T>(&self, _dist: &D) -> T
+    where
+        T: 'static,
+        D: Distribution<T>,
+        Standard: Distribution<T>,
+    {
+        // Will return indicies provided as input, cannot check if the index is within the range.
+        self.gen::<T>()
     }
 }
 
 #[cfg(test)]
 #[cfg_attr(coverage, coverage(off))]
 mod tests {
+    use rand::distributions::WeightedIndex;
+
     use crate::random::RandomSource;
 
     use super::*;

phylo/src/random/mod.rs

@@ -2,7 +2,7 @@ use std::sync::Mutex;
 
 use ntimestamp::Timestamp;
 use rand::distributions::uniform::{SampleRange, SampleUniform};
-use rand::distributions::{Distribution, Standard, WeightedIndex};
+use rand::distributions::{Distribution, Standard};
 use rand::prelude::SliceRandom;
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
@@ -34,7 +34,11 @@ pub trait RandomSource {
     fn reseed(&self, seed: u64);
 
     /// Sample from a weighted distribution.
-    fn sample(&self, dist: &WeightedIndex<f64>) -> usize;
+    fn sample<D, T>(&self, dist: &D) -> T
+    where
+        T: 'static,
+        D: Distribution<T>,
+        Standard: Distribution<T>;
 }
 
 pub struct SeededRng<R>
@@ -145,9 +149,14 @@ where
     }
 
     /// Sample from a weighted distribution.
-    fn sample(&self, dist: &WeightedIndex<f64>) -> usize {
+    fn sample<D, T>(&self, dist: &D) -> T
+    where
+        T: 'static,
+        D: Distribution<T>,
+        Standard: Distribution<T>,
+    {
         let mut r = self.r.lock().unwrap();
-        dist.sample(&mut r.rng)
+        r.rng.sample(dist)
     }
 }
 
@@ -170,6 +179,7 @@ where
 #[cfg(test)]
 mod tests {
     use itertools::repeat_n;
+    use rand::distributions::WeightedIndex;
 
     use super::*;