use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::composite::CompositeScore;
use crate::run_identity::RunKey;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CensusError {
EmptyAxis { axis: &'static str },
DuplicateDeclaration { axis: &'static str, value: String },
}
impl fmt::Display for CensusError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::EmptyAxis { axis } => write!(
f,
"trial census: the declared {axis} axis is empty, so the roster expects no trials"
),
Self::DuplicateDeclaration { axis, value } => write!(
f,
"trial census: {axis} `{value}` is declared more than once"
),
}
}
}
impl std::error::Error for CensusError {}
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct TrialRoster {
agents: Vec<String>,
windows: Vec<String>,
seeds: Vec<u64>,
}
#[derive(Deserialize)]
#[serde(rename = "TrialRoster")]
struct RosterWire {
agents: Vec<String>,
windows: Vec<String>,
seeds: Vec<u64>,
}
impl<'de> Deserialize<'de> for TrialRoster {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let wire = RosterWire::deserialize(deserializer)?;
Self::declare(&wire.agents, &wire.windows, &wire.seeds).map_err(serde::de::Error::custom)
}
}
impl TrialRoster {
pub fn declare(
agents: &[String],
windows: &[String],
seeds: &[u64],
) -> Result<Self, CensusError> {
check_unique("entrant", agents.iter().cloned())?;
check_unique("window", windows.iter().cloned())?;
check_unique("seed", seeds.iter().map(|s| s.to_string()))?;
if agents.is_empty() {
return Err(CensusError::EmptyAxis { axis: "entrant" });
}
if windows.is_empty() {
return Err(CensusError::EmptyAxis { axis: "window" });
}
if seeds.is_empty() {
return Err(CensusError::EmptyAxis { axis: "seed" });
}
let mut windows = windows.to_vec();
windows.sort();
let mut seeds = seeds.to_vec();
seeds.sort_unstable();
Ok(Self {
agents: agents.to_vec(),
windows,
seeds,
})
}
pub fn agents(&self) -> &[String] {
&self.agents
}
pub fn cells(&self) -> Vec<RunKey> {
self.windows
.iter()
.flat_map(|window| {
self.seeds.iter().map(move |seed| RunKey {
window: window.clone(),
seed: *seed,
})
})
.collect()
}
pub fn expected_per_agent(&self) -> usize {
self.windows.len() * self.seeds.len()
}
pub fn expected(&self) -> usize {
self.agents.len() * self.expected_per_agent()
}
pub fn cohort(&self) -> CohortIdentity {
let mut agents = self.agents.clone();
agents.sort();
CohortIdentity {
agents,
windows: self.windows.clone(),
seeds: self.seeds.clone(),
}
}
}
fn check_unique<I: Iterator<Item = String>>(
axis: &'static str,
values: I,
) -> Result<(), CensusError> {
let mut seen: BTreeSet<String> = BTreeSet::new();
for value in values {
if !seen.insert(value.clone()) {
return Err(CensusError::DuplicateDeclaration { axis, value });
}
}
Ok(())
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CohortIdentity {
pub agents: Vec<String>,
pub windows: Vec<String>,
pub seeds: Vec<u64>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum TrialOutcome {
Completed,
Failed { reason: String },
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TrialReport {
pub agent_id: String,
#[serde(flatten)]
pub key: RunKey,
#[serde(flatten)]
pub outcome: TrialOutcome,
}
impl TrialReport {
pub fn completed(agent_id: &str, window: &str, seed: u64) -> Self {
Self {
agent_id: agent_id.to_string(),
key: RunKey {
window: window.to_string(),
seed,
},
outcome: TrialOutcome::Completed,
}
}
pub fn failed(agent_id: &str, window: &str, seed: u64, reason: &str) -> Self {
Self {
agent_id: agent_id.to_string(),
key: RunKey {
window: window.to_string(),
seed,
},
outcome: TrialOutcome::Failed {
reason: reason.to_string(),
},
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "gate", rename_all = "snake_case")]
pub enum CensusGate {
TrialFailed {
window: String,
seed: u64,
reason: String,
},
TrialUnreported { window: String, seed: u64 },
TrialReportedTwice { window: String, seed: u64 },
OffRosterCell { window: String, seed: u64 },
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentCensus {
pub agent_id: String,
pub expected: usize,
pub completed: usize,
pub failed: usize,
pub unreported: usize,
pub gates: Vec<CensusGate>,
}
impl AgentCensus {
pub fn complete(&self) -> bool {
self.completed == self.expected
}
pub fn ungated(&self) -> bool {
self.gates.is_empty()
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TrialCensus {
pub cohort: CohortIdentity,
pub expected: usize,
pub completed: usize,
pub failed: usize,
pub unreported: usize,
pub per_agent: Vec<AgentCensus>,
pub off_roster_agents: Vec<String>,
}
impl TrialCensus {
pub fn agent(&self, agent_id: &str) -> Option<&AgentCensus> {
self.per_agent.iter().find(|a| a.agent_id == agent_id)
}
pub fn suite_complete(&self) -> bool {
self.completed == self.expected
&& self.off_roster_agents.is_empty()
&& self.per_agent.iter().all(AgentCensus::ungated)
}
}
pub fn census(roster: &TrialRoster, reports: &[TrialReport]) -> TrialCensus {
let declared_cells: BTreeSet<RunKey> = roster.cells().into_iter().collect();
let declared_agents: BTreeSet<&str> = roster.agents().iter().map(String::as_str).collect();
let mut by_agent: BTreeMap<&str, BTreeMap<&RunKey, Vec<&TrialOutcome>>> = BTreeMap::new();
let mut off_roster_agents: BTreeSet<String> = BTreeSet::new();
for report in reports {
if !declared_agents.contains(report.agent_id.as_str()) {
off_roster_agents.insert(report.agent_id.clone());
continue;
}
by_agent
.entry(report.agent_id.as_str())
.or_default()
.entry(&report.key)
.or_default()
.push(&report.outcome);
}
let expected_per_agent = roster.expected_per_agent();
let mut per_agent: Vec<AgentCensus> = Vec::with_capacity(roster.agents().len());
for agent_id in roster.agents() {
let reported = by_agent.remove(agent_id.as_str()).unwrap_or_default();
let mut completed = 0usize;
let mut failed = 0usize;
let mut unreported = 0usize;
let mut gates: Vec<CensusGate> = Vec::new();
for key in roster.cells() {
match reported.get(&key).map(Vec::as_slice) {
None => {
unreported += 1;
gates.push(CensusGate::TrialUnreported {
window: key.window.clone(),
seed: key.seed,
});
}
Some([TrialOutcome::Completed]) => completed += 1,
Some([TrialOutcome::Failed { reason }]) => {
failed += 1;
gates.push(CensusGate::TrialFailed {
window: key.window.clone(),
seed: key.seed,
reason: reason.to_string(),
});
}
Some(_) => gates.push(CensusGate::TrialReportedTwice {
window: key.window.clone(),
seed: key.seed,
}),
}
}
for key in reported.keys() {
if !declared_cells.contains(*key) {
gates.push(CensusGate::OffRosterCell {
window: key.window.clone(),
seed: key.seed,
});
}
}
per_agent.push(AgentCensus {
agent_id: agent_id.clone(),
expected: expected_per_agent,
completed,
failed,
unreported,
gates,
});
}
TrialCensus {
cohort: roster.cohort(),
expected: roster.expected(),
completed: per_agent.iter().map(|a| a.completed).sum(),
failed: per_agent.iter().map(|a| a.failed).sum(),
unreported: per_agent.iter().map(|a| a.unreported).sum(),
per_agent,
off_roster_agents: off_roster_agents.into_iter().collect(),
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CensusedScore {
pub agent_id: String,
pub score: Option<CompositeScore>,
pub trials: AgentCensus,
pub reported_eligible: bool,
}
pub fn attach_census(board: &[CompositeScore], census: &TrialCensus) -> Vec<CensusedScore> {
census
.per_agent
.iter()
.map(|trials| {
let score = board
.iter()
.find(|s| s.agent_id == trials.agent_id)
.cloned();
let reported_eligible = score.as_ref().is_some_and(|s| s.rank_eligible)
&& trials.complete()
&& trials.ungated();
CensusedScore {
agent_id: trials.agent_id.clone(),
score,
trials: trials.clone(),
reported_eligible,
}
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::composite::{rank, AgentSubmission, Run, ScoreConfig};
fn ids(names: &[&str]) -> Vec<String> {
names.iter().map(|n| (*n).to_string()).collect()
}
fn roster() -> TrialRoster {
TrialRoster::declare(&ids(&["alpha", "beta"]), &ids(&["w1", "w2"]), &[7, 9, 11])
.expect("a well formed roster")
}
fn all_completed() -> Vec<TrialReport> {
let mut out = Vec::new();
for agent in ["alpha", "beta"] {
for window in ["w1", "w2"] {
for seed in [7, 9, 11] {
out.push(TrialReport::completed(agent, window, seed));
}
}
}
out
}
#[test]
fn expected_comes_from_the_roster_not_from_the_rows() {
let reports: Vec<TrialReport> = all_completed().into_iter().take(5).collect();
let c = census(&roster(), &reports);
assert_eq!(c.expected, 12, "{c:?}");
assert_eq!(c.completed, 5, "{c:?}");
assert_eq!(c.unreported, 7, "{c:?}");
}
#[test]
fn a_failed_trial_is_not_a_completion() {
let mut reports = all_completed();
reports[0] = TrialReport::failed("alpha", "w1", 7, "gateway refused the journal");
let c = census(&roster(), &reports);
assert_eq!(c.completed, 11, "{c:?}");
assert_eq!(c.failed, 1, "{c:?}");
let alpha = c.agent("alpha").expect("declared entrant");
assert_eq!(
alpha.gates,
vec![CensusGate::TrialFailed {
window: "w1".to_string(),
seed: 7,
reason: "gateway refused the journal".to_string(),
}],
"the reason is carried, not just the count"
);
}
#[test]
fn removing_a_failed_trial_cannot_improve_eligibility_or_the_cohort() {
let mut with_failure = all_completed();
with_failure[0] = TrialReport::failed("alpha", "w1", 7, "sim crashed");
let dropped: Vec<TrialReport> = with_failure.iter().skip(1).cloned().collect();
let before = census(&roster(), &with_failure);
let after = census(&roster(), &dropped);
assert_eq!(
before.expected, after.expected,
"the denominator is declared"
);
assert_eq!(before.completed, after.completed, "no completion appeared");
assert_eq!(before.cohort, after.cohort, "the declared cohort is fixed");
assert_eq!((before.failed, before.unreported), (1, 0));
assert_eq!((after.failed, after.unreported), (0, 1));
let before_alpha = before.agent("alpha").expect("declared entrant");
let after_alpha = after.agent("alpha").expect("declared entrant");
assert!(!before_alpha.complete());
assert!(
!after_alpha.complete(),
"deleting the failed report must not complete the cohort: {after_alpha:?}"
);
}
#[test]
fn an_undeclared_cell_is_not_a_completion() {
let mut reports = all_completed();
reports[0] = TrialReport::failed("alpha", "w1", 7, "sim crashed");
reports.push(TrialReport::completed("alpha", "w-easy", 7));
let c = census(&roster(), &reports);
assert_eq!(c.expected, 12);
assert_eq!(c.completed, 11, "{c:?}");
let alpha = c.agent("alpha").expect("declared entrant");
assert!(!alpha.ungated(), "{alpha:?}");
assert!(
alpha.gates.contains(&CensusGate::OffRosterCell {
window: "w-easy".to_string(),
seed: 7,
}),
"{alpha:?}"
);
}
#[test]
fn an_off_roster_entrant_is_counted_nowhere() {
let mut reports = all_completed();
reports.push(TrialReport::completed("gatecrasher", "w1", 7));
let c = census(&roster(), &reports);
assert_eq!(c.off_roster_agents, vec!["gatecrasher".to_string()]);
assert_eq!(
(c.expected, c.completed),
(12, 12),
"an entrant the roster never declared moves no count: {c:?}"
);
assert!(c.agent("gatecrasher").is_none());
assert!(
!c.suite_complete(),
"every declared trial completed, but an undeclared entrant reported into the suite, so the suite is not clean: {c:?}"
);
}
#[test]
fn a_cell_reported_twice_is_not_established() {
let mut reports = all_completed();
reports.push(TrialReport::completed("alpha", "w1", 7));
let c = census(&roster(), &reports);
assert_eq!(
c.completed, 11,
"the ambiguous cell counts as no completion"
);
let alpha = c.agent("alpha").expect("declared entrant");
assert!(alpha.gates.contains(&CensusGate::TrialReportedTwice {
window: "w1".to_string(),
seed: 7,
}));
assert!(!alpha.complete());
}
#[test]
fn the_declared_cohort_ignores_declaration_order() {
let a = TrialRoster::declare(&ids(&["alpha", "beta"]), &ids(&["w1", "w2"]), &[7, 9])
.expect("roster");
let b = TrialRoster::declare(&ids(&["beta", "alpha"]), &ids(&["w2", "w1"]), &[9, 7])
.expect("roster");
assert_eq!(a.cohort(), b.cohort());
assert_eq!(census(&a, &[]).per_agent[0].agent_id, "alpha");
assert_eq!(census(&b, &[]).per_agent[0].agent_id, "beta");
}
#[test]
fn a_roster_refuses_an_empty_axis_and_a_repeat() {
assert_eq!(
TrialRoster::declare(&ids(&["alpha"]), &[], &[7]),
Err(CensusError::EmptyAxis { axis: "window" })
);
assert_eq!(
TrialRoster::declare(&ids(&["alpha", "alpha"]), &ids(&["w1"]), &[7]),
Err(CensusError::DuplicateDeclaration {
axis: "entrant",
value: "alpha".to_string(),
})
);
}
fn submission(id: &str, mean: f64) -> AgentSubmission {
AgentSubmission {
agent_id: id.to_string(),
runs: (0..4)
.map(|_| Run {
returns: (0..60)
.map(|i| mean + 0.0005 * (i as f64 * 0.7).sin())
.collect(),
..Run::default()
})
.collect(),
..AgentSubmission::default()
}
}
#[test]
fn counts_travel_beside_the_score_and_a_shortfall_withholds_eligibility() {
let board = rank(
&[submission("alpha", 0.002), submission("beta", 0.0019)],
&ScoreConfig::default(),
);
assert!(
board
.iter()
.find(|s| s.agent_id == "alpha")
.expect("the board carried alpha")
.rank_eligible,
"the fixture must make alpha host-eligible, or withholding proves nothing: {board:?}"
);
let whole = attach_census(&board, &census(&roster(), &all_completed()));
let alpha_whole = &whole[0];
assert_eq!(alpha_whole.agent_id, "alpha");
assert_eq!(alpha_whole.trials.expected, 6);
assert_eq!(alpha_whole.trials.completed, 6);
assert!(
alpha_whole.reported_eligible,
"with a whole declared cohort the host verdict stands unchanged: {alpha_whole:?}"
);
let short: Vec<TrialReport> = all_completed()
.into_iter()
.filter(|r| !(r.agent_id == "alpha" && r.key.window == "w2" && r.key.seed == 9))
.collect();
let gated = attach_census(&board, &census(&roster(), &short));
let alpha_gated = &gated[0];
assert_eq!(
alpha_gated.score, alpha_whole.score,
"the score does not move"
);
assert!(
!alpha_gated.reported_eligible,
"an incomplete declared cohort cannot report as eligible: {alpha_gated:?}"
);
}
#[test]
fn a_whole_cohort_cannot_confer_the_eligibility_the_host_withheld() {
let board = rank(&[submission("loser", -0.003)], &ScoreConfig::default());
let host = &board[0];
assert!(
!host.rank_eligible,
"the fixture must be host-ineligible, or the join proves nothing: {host:?}"
);
let one = TrialRoster::declare(&ids(&["loser"]), &ids(&["w1", "w2"]), &[7, 9, 11])
.expect("roster");
let reports: Vec<TrialReport> = one
.cells()
.into_iter()
.map(|k| TrialReport::completed("loser", &k.window, k.seed))
.collect();
let rows = attach_census(&board, &census(&one, &reports));
assert!(rows[0].trials.complete() && rows[0].trials.ungated());
assert!(
!rows[0].reported_eligible,
"a complete cohort cannot promote a host-ineligible row: {:?}",
rows[0]
);
}
#[test]
fn a_declared_entrant_with_no_row_still_appears_with_its_counts() {
let board = rank(&[submission("alpha", 0.002)], &ScoreConfig::default());
let reports: Vec<TrialReport> = all_completed()
.into_iter()
.filter(|r| r.agent_id == "alpha")
.collect();
let rows = attach_census(&board, &census(&roster(), &reports));
let beta = rows.iter().find(|r| r.agent_id == "beta").expect(
"a declared entrant absent from the board must still be published with its counts",
);
assert!(beta.score.is_none());
assert_eq!(beta.trials.expected, 6);
assert_eq!(beta.trials.completed, 0);
assert_eq!(beta.trials.unreported, 6);
assert!(!beta.reported_eligible);
}
#[test]
fn a_whole_suite_reports_complete() {
let c = census(&roster(), &all_completed());
assert!(c.suite_complete(), "{c:?}");
assert_eq!(
(c.expected, c.completed, c.failed, c.unreported),
(12, 12, 0, 0)
);
}
#[test]
fn a_roster_deserialized_with_an_empty_axis_is_refused() {
let err =
serde_json::from_str::<TrialRoster>(r#"{"agents":["a"],"windows":[],"seeds":[1]}"#)
.expect_err("an empty declared axis expects no trials and must be refused");
assert!(
err.to_string().contains("window axis is empty"),
"the refusal must name the axis it refused on: {err}"
);
}
#[test]
fn a_declared_roster_survives_the_json_round_trip() {
let declared = roster();
let json = serde_json::to_string(&declared).expect("serializes");
let parsed: TrialRoster = serde_json::from_str(&json).expect("a declared roster parses");
assert_eq!(parsed, declared);
assert_eq!(parsed.expected(), 12);
let c = census(&parsed, &all_completed());
assert!(c.suite_complete(), "{c:?}");
}
}