(ml5717) Refactored PrimeableRng

juntyr · juntyr · commit 1f535b1a3ef9 · 2020-11-10T10:36:41.000Z
diff --git a/necsim-core/src/cogs/rng.rs b/necsim-core/src/cogs/rng.rs
@@ -1,4 +1,6 @@
+use crate::cogs::HabitatToU64Injection;
 use crate::intrinsics::{floor, ln};
+use crate::landscape::IndexedLocation;
 
 use core::convert::AsMut;
 use core::default::Default;
@@ -97,8 +99,6 @@ pub trait RngSampler: RngCore {
 impl<R: RngCore> RngSampler for R {}
 
 #[allow(clippy::module_name_repetitions)]
-pub trait PrimeableRng: RngCore {
-    type Prime: AsMut<[u8]> + Sized;
-
-    fn prime_with(&mut self, prime: Self::Prime);
+pub trait PrimeableRng<H: HabitatToU64Injection>: RngCore {
+    fn prime_with(&mut self, habitat: &H, indexed_location: &IndexedLocation, time_index: u64);
 }
diff --git a/necsim-cuda/kernel/src/lib.rs b/necsim-cuda/kernel/src/lib.rs
@@ -40,8 +40,8 @@ impl core::fmt::Debug for F64 {
 }
 
 use necsim_core::cogs::{
-    ActiveLineageSampler, CoalescenceSampler, DispersalSampler, EventSampler, Habitat,
-    IncoherentLineageStore, LineageReference, PrimeableRng,
+    ActiveLineageSampler, CoalescenceSampler, DispersalSampler, EventSampler,
+    HabitatToU64Injection, IncoherentLineageStore, LineageReference, PrimeableRng,
 };
 use necsim_core::reporter::NullReporter;
 use necsim_core::simulation::Simulation;
@@ -84,8 +84,8 @@ pub unsafe extern "ptx-kernel" fn simulate(
 }
 
 unsafe fn simulate_generic<
-    H: Habitat + RustToCuda,
-    G: PrimeableRng<Prime = [u8; 16]> + RustToCuda,
+    H: HabitatToU64Injection + RustToCuda,
+    G: PrimeableRng<H> + RustToCuda,
     D: DispersalSampler<H, G> + RustToCuda,
     R: LineageReference<H> + DeviceCopy,
     S: IncoherentLineageStore<H, R> + RustToCuda,
diff --git a/necsim-cuda/src/lib.rs b/necsim-cuda/src/lib.rs
@@ -149,7 +149,7 @@ impl CudaSimulation {
         sample_percentage <= 1.0_f64,
         "0.0 <= sample_percentage <= 1.0"
     )]
-    pub fn simulate<G: PrimeableRng<Prime = [u8; 16]>>(
+    pub fn simulate<G: PrimeableRng<InMemoryHabitat>>(
         habitat: &Array2D<u32>,
         dispersal: &Array2D<f64>,
         speciation_probability_per_generation: f64,
diff --git a/necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/mod.rs b/necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/mod.rs
@@ -16,7 +16,7 @@ mod sampler;
 #[derive(Debug)]
 pub struct IndependentActiveLineageSampler<
     H: HabitatToU64Injection,
-    G: PrimeableRng<Prime = [u8; 16]>,
+    G: PrimeableRng<H>,
     D: DispersalSampler<H, G>,
     R: LineageReference<H>,
     S: IncoherentLineageStore<H, R>,
@@ -28,7 +28,7 @@ pub struct IndependentActiveLineageSampler<
 
 impl<
         H: HabitatToU64Injection,
-        G: PrimeableRng<Prime = [u8; 16]>,
+        G: PrimeableRng<H>,
         D: DispersalSampler<H, G>,
         R: LineageReference<H>,
         S: IncoherentLineageStore<H, R>,
@@ -47,7 +47,7 @@ impl<
 
 impl<
         H: HabitatToU64Injection,
-        G: PrimeableRng<Prime = [u8; 16]>,
+        G: PrimeableRng<H>,
         D: DispersalSampler<H, G>,
         R: LineageReference<H>,
         S: IncoherentLineageStore<H, R>,
diff --git a/necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/sampler.rs b/necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/sampler.rs
@@ -16,7 +16,7 @@ use super::IndependentActiveLineageSampler;
 #[contract_trait]
 impl<
         H: HabitatToU64Injection,
-        G: PrimeableRng<Prime = [u8; 16]>,
+        G: PrimeableRng<H>,
         D: DispersalSampler<H, G>,
         R: LineageReference<H>,
         S: IncoherentLineageStore<H, R>,
@@ -83,30 +83,7 @@ impl<
         let mut time_step = floor(time / delta_t) as u64 + 1;
 
         loop {
-            let location_bytes = simulation
-                .habitat
-                .map_indexed_location_to_u64_injective(&lineage_indexed_location)
-                .to_le_bytes();
-            let time_step_bytes = time_step.to_le_bytes();
-
-            rng.prime_with([
-                location_bytes[0],
-                location_bytes[1],
-                location_bytes[2],
-                location_bytes[3],
-                location_bytes[4],
-                location_bytes[5],
-                location_bytes[6],
-                location_bytes[7],
-                time_step_bytes[0],
-                time_step_bytes[1],
-                time_step_bytes[2],
-                time_step_bytes[3],
-                time_step_bytes[4],
-                time_step_bytes[5],
-                time_step_bytes[6],
-                time_step_bytes[7],
-            ]);
+            rng.prime_with(simulation.habitat, &lineage_indexed_location, time_step);
 
             if rng.sample_event(p) {
                 break;
diff --git a/necsim-impls-no-std/src/cogs/rng/aes.rs b/necsim-impls-no-std/src/cogs/rng/aes.rs
@@ -1,3 +1,6 @@
+use necsim_core::cogs::HabitatToU64Injection;
+use necsim_core::landscape::IndexedLocation;
+
 use aes_soft::cipher::generic_array::GenericArray;
 use aes_soft::cipher::{BlockCipher, NewBlockCipher};
 use aes_soft::Aes128;
@@ -63,11 +66,32 @@ impl necsim_core::cogs::RngCore for AesRng {
     }
 }
 
-impl necsim_core::cogs::PrimeableRng for AesRng {
-    type Prime = [u8; 16];
+impl<H: HabitatToU64Injection> necsim_core::cogs::PrimeableRng<H> for AesRng {
+    fn prime_with(&mut self, habitat: &H, indexed_location: &IndexedLocation, time_index: u64) {
+        let location_bytes = habitat
+            .map_indexed_location_to_u64_injective(indexed_location)
+            .to_le_bytes();
+
+        self.state[0] = location_bytes[0];
+        self.state[1] = location_bytes[1];
+        self.state[2] = location_bytes[2];
+        self.state[3] = location_bytes[3];
+        self.state[4] = location_bytes[4];
+        self.state[5] = location_bytes[5];
+        self.state[6] = location_bytes[6];
+        self.state[7] = location_bytes[7];
+
+        let time_index_bytes = time_index.to_le_bytes();
+
+        self.state[8] = time_index_bytes[0];
+        self.state[9] = time_index_bytes[1];
+        self.state[10] = time_index_bytes[2];
+        self.state[11] = time_index_bytes[3];
+        self.state[12] = time_index_bytes[4];
+        self.state[13] = time_index_bytes[5];
+        self.state[14] = time_index_bytes[6];
+        self.state[15] = time_index_bytes[7];
 
-    fn prime_with(&mut self, prime: Self::Prime) {
-        self.state = prime;
         self.cached = false;
     }
 }
diff --git a/necsim-impls-no-std/src/cogs/rng/cuda.rs b/necsim-impls-no-std/src/cogs/rng/cuda.rs
@@ -1,4 +1,5 @@
-use necsim_core::cogs::{PrimeableRng, RngCore};
+use necsim_core::cogs::{HabitatToU64Injection, PrimeableRng, RngCore};
+use necsim_core::landscape::IndexedLocation;
 
 #[allow(clippy::module_name_repetitions)]
 #[derive(Clone, Debug, RustToCuda, LendToCuda)]
@@ -28,10 +29,8 @@ impl<R: RngCore> RngCore for CudaRng<R> {
     }
 }
 
-impl<R: PrimeableRng> PrimeableRng for CudaRng<R> {
-    type Prime = <R as PrimeableRng>::Prime;
-
-    fn prime_with(&mut self, prime: Self::Prime) {
-        self.0.prime_with(prime)
+impl<H: HabitatToU64Injection, R: PrimeableRng<H>> PrimeableRng<H> for CudaRng<R> {
+    fn prime_with(&mut self, habitat: &H, indexed_location: &IndexedLocation, time_index: u64) {
+        self.0.prime_with(habitat, indexed_location, time_index)
     }
 }
diff --git a/rustcoalescence/src/simulation.rs b/rustcoalescence/src/simulation.rs
@@ -12,7 +12,7 @@ use necsim_skipping_gillespie::SkippingGillespieSimulation;
 
 use super::args::{Algorithm, CommandLineArguments};
 
-pub fn simulate<G: PrimeableRng<Prime = [u8; 16]>>(
+pub fn simulate<G: PrimeableRng<InMemoryHabitat>>(
     args: &CommandLineArguments,
     habitat: &Array2D<u32>,
     dispersal: &Array2D<f64>,

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`		`-use necsim_core::cogs::{PrimeableRng, RngCore};`
	`1`	`+use necsim_core::cogs::{HabitatToU64Injection, PrimeableRng, RngCore};`
	`2`	`+use necsim_core::landscape::IndexedLocation;`
`2`	`3`
`3`	`4`	`#[allow(clippy::module_name_repetitions)]`
`4`	`5`	`#[derive(Clone, Debug, RustToCuda, LendToCuda)]`
`@@ -28,10 +29,8 @@ impl<R: RngCore> RngCore for CudaRng<R> {`
`28`	`29`	`}`
`29`	`30`	`}`
`30`	`31`
`31`		`-impl<R: PrimeableRng> PrimeableRng for CudaRng<R> {`
`32`		`- type Prime = <R as PrimeableRng>::Prime;`
`33`		`-`
`34`		`- fn prime_with(&mut self, prime: Self::Prime) {`
`35`		`- self.0.prime_with(prime)`
	`32`	`+impl<H: HabitatToU64Injection, R: PrimeableRng<H>> PrimeableRng<H> for CudaRng<R> {`
	`33`	`+ fn prime_with(&mut self, habitat: &H, indexed_location: &IndexedLocation, time_index: u64) {`
	`34`	`+ self.0.prime_with(habitat, indexed_location, time_index)`
`36`	`35`	`}`
`37`	`36`	`}`