use std::collections::BTreeMap;
use std::path::PathBuf;
use crate::domain::config::Config;
use crate::domain::template::RuleSpec;
#[derive(Debug, Clone)]
pub struct ResolvedRule {
pub name: String,
pub description: String,
pub judges: u32,
pub agent: String,
pub files: Vec<PathBuf>,
pub rationale: bool,
pub relevance: Option<String>,
pub require_line_attribution: bool,
}
#[derive(Debug, Clone)]
pub struct JudgeRun {
pub agent: String,
pub harness: Option<String>,
pub model: Option<String>,
pub schema_max_retries: Option<u32>,
pub judge_index: u32,
pub template: String,
pub files: Vec<PathBuf>,
pub rules: Vec<RuleSpec>,
}
#[derive(Debug, Default)]
pub struct Plan {
pub runs: Vec<JudgeRun>,
pub skipped: Vec<String>,
}
pub fn build(
config: &Config,
master_template: &str,
default_batch_size: usize,
resolved: Vec<ResolvedRule>,
) -> Plan {
let mut plan = Plan::default();
let mut by_agent: BTreeMap<String, Vec<ResolvedRule>> = BTreeMap::new();
for r in resolved {
by_agent.entry(r.agent.clone()).or_default().push(r);
}
for (agent_name, rules) in by_agent {
let agent = config.agent_or_default(&agent_name);
let harness = agent.harness.clone();
let batch_size = agent.batch_size.unwrap_or(default_batch_size).max(1);
let template = match &agent.prompt_template {
Some(extra) => format!("{master_template}\n\n{extra}"),
None => master_template.to_string(),
};
let (eligible, empty): (Vec<ResolvedRule>, Vec<ResolvedRule>) =
rules.into_iter().partition(|r| !r.files.is_empty());
plan.skipped.extend(empty.into_iter().map(|r| r.name));
if eligible.is_empty() {
continue;
}
let max_judges = eligible.iter().map(|r| r.judges).max().unwrap_or(1);
for j in 1..=max_judges {
let subset: Vec<&ResolvedRule> = eligible.iter().filter(|r| r.judges >= j).collect();
for chunk in balanced_chunks(&subset, batch_size) {
let mut files: Vec<PathBuf> =
chunk.iter().flat_map(|r| r.files.iter().cloned()).collect();
files.sort();
files.dedup();
plan.runs.push(JudgeRun {
agent: agent_name.clone(),
harness: harness.clone(),
model: agent.model.clone(),
schema_max_retries: config.oneharness.schema_max_retries,
judge_index: j,
template: template.clone(),
files,
rules: chunk
.iter()
.map(|r| RuleSpec {
name: r.name.clone(),
description: r.description.clone(),
rationale: r.rationale,
relevance: r.relevance.clone(),
require_line_attribution: r.require_line_attribution,
files: r.files.iter().map(|p| crate::domain::to_slash(p)).collect(),
})
.collect(),
});
}
}
}
plan
}
fn balanced_chunks<T>(items: &[T], batch_size: usize) -> Vec<&[T]> {
if items.is_empty() {
return Vec::new();
}
let batch_size = batch_size.max(1);
let num_batches = items.len().div_ceil(batch_size);
let base = items.len() / num_batches;
let remainder = items.len() % num_batches;
let mut chunks = Vec::with_capacity(num_batches);
let mut start = 0;
for i in 0..num_batches {
let size = base + usize::from(i < remainder);
chunks.push(&items[start..start + size]);
start += size;
}
chunks
}
#[cfg(test)]
mod tests {
use super::*;
use crate::domain::config::Agent;
fn rr(name: &str, judges: u32, agent: &str, files: &[&str]) -> ResolvedRule {
ResolvedRule {
name: name.into(),
description: format!("desc {name}"),
judges,
agent: agent.into(),
files: files.iter().map(PathBuf::from).collect(),
rationale: true,
relevance: None,
require_line_attribution: false,
}
}
#[test]
fn single_judge_rules_run_once() {
let cfg = Config::default();
let plan = build(
&cfg,
"T",
20,
vec![
rr("a", 1, "default", &["f.rs"]),
rr("b", 1, "default", &["f.rs"]),
],
);
assert_eq!(plan.runs.len(), 1);
assert_eq!(plan.runs[0].rules.len(), 2);
assert_eq!(plan.runs[0].harness, None);
assert!(plan.skipped.is_empty());
}
#[test]
fn multi_judge_expands_into_one_run_per_judge_index() {
let cfg = Config::default();
let plan = build(
&cfg,
"T",
20,
vec![
rr("a", 3, "default", &["f.rs"]),
rr("b", 1, "default", &["f.rs"]),
],
);
assert_eq!(plan.runs.len(), 3);
let j1 = plan.runs.iter().find(|r| r.judge_index == 1).unwrap();
assert_eq!(j1.rules.len(), 2);
assert_eq!(plan.runs.iter().filter(|r| r.judge_index == 2).count(), 1);
assert_eq!(plan.runs.iter().filter(|r| r.judge_index == 3).count(), 1);
}
#[test]
fn batches_respect_agent_batch_size() {
let mut cfg = Config::default();
cfg.agents.insert(
"small".into(),
Agent {
batch_size: Some(2),
prompt_template: Some("be terse".into()),
..Default::default()
},
);
let rules = vec![
rr("a", 1, "small", &["f.rs"]),
rr("b", 1, "small", &["f.rs"]),
rr("c", 1, "small", &["f.rs"]),
];
let plan = build(&cfg, "MASTER", 20, rules);
assert_eq!(plan.runs.len(), 2); assert!(plan.runs[0].template.contains("MASTER"));
assert!(plan.runs[0].template.contains("be terse"));
}
#[test]
fn batches_are_balanced_not_packed() {
let mut cfg = Config::default();
cfg.agents.insert(
"big".into(),
Agent {
batch_size: Some(20),
..Default::default()
},
);
let rules: Vec<ResolvedRule> = (0..21)
.map(|i| rr(&format!("r{i}"), 1, "big", &["f.rs"]))
.collect();
let plan = build(&cfg, "T", 20, rules);
let mut sizes: Vec<usize> = plan.runs.iter().map(|r| r.rules.len()).collect();
sizes.sort_unstable();
assert_eq!(sizes, vec![10, 11]);
let total: usize = plan.runs.iter().map(|r| r.rules.len()).sum();
assert_eq!(total, 21);
}
#[test]
fn balanced_chunks_covers_items_without_overlap() {
let items: Vec<usize> = (0..25).collect();
let chunks = balanced_chunks(&items, 10);
assert_eq!(chunks.len(), 3);
let mut sizes: Vec<usize> = chunks.iter().map(|c| c.len()).collect();
sizes.sort_unstable();
assert_eq!(sizes, vec![8, 8, 9]);
let flat: Vec<usize> = chunks.iter().flat_map(|c| c.iter().copied()).collect();
assert_eq!(flat, items);
}
#[test]
fn balanced_chunks_single_full_batch() {
let items: Vec<usize> = (0..20).collect();
let chunks = balanced_chunks(&items, 20);
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].len(), 20);
}
#[test]
fn balanced_chunks_empty_is_no_batches() {
let items: Vec<usize> = Vec::new();
assert!(balanced_chunks(&items, 5).is_empty());
}
#[test]
fn per_rule_rationale_flows_into_the_rule_spec() {
let cfg = Config::default();
let mut on = rr("on", 1, "default", &["f.rs"]);
let mut off = rr("off", 1, "default", &["f.rs"]);
on.rationale = true;
off.rationale = false;
let plan = build(&cfg, "T", 20, vec![on, off]);
let specs = &plan.runs[0].rules;
let find = |n: &str| specs.iter().find(|r| r.name == n).unwrap().rationale;
assert!(find("on"));
assert!(!find("off"));
}
#[test]
fn relevance_condition_flows_into_the_rule_spec() {
let cfg = Config::default();
let mut conditional = rr("conditional", 1, "default", &["f.rs"]);
let always = rr("always", 1, "default", &["f.rs"]);
conditional.relevance = Some("the change touches SQL".into());
let plan = build(&cfg, "T", 20, vec![conditional, always]);
let specs = &plan.runs[0].rules;
let find = |n: &str| {
specs
.iter()
.find(|r| r.name == n)
.unwrap()
.relevance
.clone()
};
assert_eq!(
find("conditional").as_deref(),
Some("the change touches SQL")
);
assert_eq!(find("always"), None);
}
#[test]
fn require_line_attribution_flows_into_the_rule_spec() {
let cfg = Config::default();
let mut strict = rr("strict", 1, "default", &["f.rs"]);
let lax = rr("lax", 1, "default", &["f.rs"]);
strict.require_line_attribution = true;
let plan = build(&cfg, "T", 20, vec![strict, lax]);
let specs = &plan.runs[0].rules;
let find = |n: &str| {
specs
.iter()
.find(|r| r.name == n)
.unwrap()
.require_line_attribution
};
assert!(find("strict"));
assert!(!find("lax"));
}
#[test]
fn distinct_file_sets_merge_into_one_call_over_the_union() {
let cfg = Config::default();
let plan = build(
&cfg,
"T",
20,
vec![
rr("a", 1, "default", &["x.rs"]),
rr("b", 1, "default", &["y.rs"]),
],
);
assert_eq!(plan.runs.len(), 1);
let run = &plan.runs[0];
assert_eq!(
run.files,
vec![PathBuf::from("x.rs"), PathBuf::from("y.rs")]
);
let a = run.rules.iter().find(|r| r.name == "a").unwrap();
let b = run.rules.iter().find(|r| r.name == "b").unwrap();
assert_eq!(a.files, vec!["x.rs"]);
assert_eq!(b.files, vec!["y.rs"]);
}
#[test]
fn batch_size_splits_the_union_merge_too() {
let mut cfg = Config::default();
cfg.agents.insert(
"small".into(),
Agent {
batch_size: Some(1),
..Default::default()
},
);
let plan = build(
&cfg,
"T",
20,
vec![
rr("a", 1, "small", &["x.rs"]),
rr("b", 1, "small", &["y.rs"]),
],
);
assert_eq!(plan.runs.len(), 2, "batch_size 1 -> one rule per call");
}
#[test]
fn empty_file_set_is_skipped_not_run() {
let cfg = Config::default();
let plan = build(&cfg, "T", 20, vec![rr("a", 1, "default", &[])]);
assert!(plan.runs.is_empty());
assert_eq!(plan.skipped, vec!["a".to_string()]);
}
#[test]
fn agent_harness_override_is_used() {
let mut cfg = Config::default();
cfg.agents.insert(
"arch".into(),
Agent {
harness: Some("codex".into()),
..Default::default()
},
);
let plan = build(&cfg, "T", 20, vec![rr("a", 1, "arch", &["f.rs"])]);
assert_eq!(plan.runs[0].harness.as_deref(), Some("codex"));
}
#[test]
fn no_agent_harness_leaves_it_unset() {
let mut cfg = Config::default();
cfg.agents.insert(
"arch".into(),
Agent {
model: Some("gpt-5".into()),
..Default::default()
},
);
let plan = build(&cfg, "T", 20, vec![rr("a", 1, "arch", &["f.rs"])]);
assert_eq!(plan.runs[0].harness, None);
assert_eq!(plan.runs[0].model.as_deref(), Some("gpt-5"));
}
}