diff --git a/src/lib.rs b/src/lib.rs index 227c4f1..ae109f5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![forbid(unsafe_code)] + use std::{ cmp::Reverse, f64::consts::{FRAC_1_SQRT_2, FRAC_2_SQRT_PI, SQRT_2}, diff --git a/src/time_slice.rs b/src/time_slice.rs index 0f001b1..b762a81 100644 --- a/src/time_slice.rs +++ b/src/time_slice.rs @@ -124,18 +124,19 @@ impl Event { .collect::>() } - /// Direct in-loop update: mutates self and `skills` inline with no - /// intermediate allocation. Used by both the sequential sweep path and, - /// via unsafe, by the parallel rayon path for events in the same color - /// group (which have disjoint agent sets — see `sweep_color_groups`). - fn iteration_direct>( - &mut self, - skills: &mut SkillStore, + /// Run inference for this event and return its per-item likelihoods. + /// + /// Reads `skills` immutably and does not touch `self`, so every event in + /// a color group can run concurrently without any aliasing question — + /// the mutation is deferred to `apply`. + fn compute>( + &self, + skills: &SkillStore, agents: &CompetitorStore, p_draw: f64, convergence: crate::ConvergenceOptions, arena: &mut ScratchArena, - ) { + ) -> EventUpdate { let teams = self.within_priors(false, false, skills, agents); let result = self.outputs(); let g = match self.kind { @@ -152,17 +153,47 @@ impl Event { ), }; + EventUpdate { + log_evidence: g.log_evidence, + likelihoods: g.likelihoods, + } + } + + /// Fold a computed update into the skill store and cache it on the items. + fn apply(&mut self, skills: &mut SkillStore, update: EventUpdate) { for (t, team) in self.teams.iter_mut().enumerate() { for (i, item) in team.items.iter_mut().enumerate() { + let fresh = update.likelihoods[t][i]; let old_likelihood = skills.get(item.agent).unwrap().likelihood; - let new_likelihood = (old_likelihood / item.likelihood) * g.likelihoods[t][i]; + let new_likelihood = (old_likelihood / item.likelihood) * fresh; skills.get_mut(item.agent).unwrap().likelihood = new_likelihood; - item.likelihood = g.likelihoods[t][i]; + item.likelihood = fresh; } } - self.log_evidence = g.log_evidence; + self.log_evidence = update.log_evidence; } + + /// Compute and apply in one step — the sequential sweep. + fn iteration_direct>( + &mut self, + skills: &mut SkillStore, + agents: &CompetitorStore, + p_draw: f64, + convergence: crate::ConvergenceOptions, + arena: &mut ScratchArena, + ) { + let update = self.compute(skills, agents, p_draw, convergence, arena); + self.apply(skills, update); + } +} + +/// The result of running inference for one event, before it is folded back +/// into the shared skill store. +#[derive(Debug)] +struct EventUpdate { + log_evidence: f64, + likelihoods: Vec>, } #[derive(Debug)] @@ -384,14 +415,13 @@ impl TimeSlice { /// Full event sweep using the color-group partition. Colors are processed /// sequentially; within each color the inner loop is parallel under rayon. /// - /// Events within each color group touch disjoint agent sets (guaranteed by - /// the greedy coloring). This lets each rayon thread write directly to its - /// events' skill likelihoods without a deferred-apply step, matching the - /// sequential path's allocation profile. The unsafe block is sound because: - /// 1. `self.events[range]` and `self.skills` are separate fields → disjoint. - /// 2. Events in the same color group access disjoint `Index` values in - /// `self.skills`, so concurrent writes land on different memory locations. - /// 3. Each event only writes to its own items' likelihoods (no sharing). + /// Events in one color group touch disjoint agent sets, so none of them + /// can observe another's writes. That makes the sweep separable: inference + /// runs concurrently over shared `&self.skills`, and the resulting updates + /// are folded in afterwards in index order. Splitting it this way needs no + /// `unsafe` and no aliasing argument, and it keeps results bit-identical + /// across thread counts because the apply order does not depend on which + /// worker finished first. #[cfg(feature = "rayon")] fn sweep_color_groups>(&mut self, agents: &CompetitorStore) { use rayon::prelude::*; @@ -411,29 +441,28 @@ impl TimeSlice { if group_len == 0 { continue; } + let range = self.color_groups.color_range(color_idx); let p_draw = self.p_draw; let convergence = self.convergence; if group_len >= RAYON_THRESHOLD { - // Obtain a raw pointer from the unique `&mut self.skills` reference. - // Casting back to `&mut` inside the closure is sound because: - // 1. The pointer originates from a `&mut` — no aliasing with shared refs. - // 2. Events in the same color group touch disjoint `Index` slots in the - // underlying Vec, so concurrent writes from different threads land on - // different memory locations — no data race. - // 3. `self.events[range]` and `self.skills` are separate struct fields, - // so the borrow splits cleanly. - let skills_addr: usize = (&mut self.skills as *mut SkillStore) as usize; - self.events[range].par_iter_mut().for_each(move |ev| { - // SAFETY: see above. - let skills: &mut SkillStore = unsafe { &mut *(skills_addr as *mut SkillStore) }; - ARENA.with(|cell| { - let mut arena = cell.borrow_mut(); - arena.reset(); - ev.iteration_direct(skills, agents, p_draw, convergence, &mut arena); - }); - }); + let skills = &self.skills; + let updates: Vec = self.events[range.clone()] + .par_iter() + .map(|ev| { + ARENA.with(|cell| { + let mut arena = cell.borrow_mut(); + arena.reset(); + + ev.compute(skills, agents, p_draw, convergence, &mut arena) + }) + }) + .collect(); + + for (ev, update) in self.events[range].iter_mut().zip(updates) { + ev.apply(&mut self.skills, update); + } } else { for ev in &mut self.events[range] { ev.iteration_direct(