use std::cell::RefCell;
use std::collections::HashMap;
use std::fs;
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicI32, Ordering};
#[derive(Debug, Hash, Eq, PartialEq, Clone, Copy)]
pub enum DiagCategory {
Kkt,
Iterate,
Step,
Mu,
Ls,
Resto,
Convergence,
Timing,
}
impl DiagCategory {
pub fn as_str(self) -> &'static str {
match self {
DiagCategory::Kkt => "kkt",
DiagCategory::Iterate => "iterate",
DiagCategory::Step => "step",
DiagCategory::Mu => "mu",
DiagCategory::Ls => "ls",
DiagCategory::Resto => "resto",
DiagCategory::Convergence => "convergence",
DiagCategory::Timing => "timing",
}
}
pub fn parse(s: &str) -> Result<Self, String> {
match s {
"kkt" => Ok(DiagCategory::Kkt),
"iterate" | "iterates" => Ok(DiagCategory::Iterate),
"step" => Ok(DiagCategory::Step),
"mu" => Ok(DiagCategory::Mu),
"ls" => Ok(DiagCategory::Ls),
"resto" => Ok(DiagCategory::Resto),
"convergence" => Ok(DiagCategory::Convergence),
"timing" => Ok(DiagCategory::Timing),
other => Err(format!(
"unknown dump category '{other}' (expected one of: kkt, iterate, step, mu, ls, resto, convergence, timing)"
)),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IterSpec {
All,
Single(i32),
Range(Option<i32>, Option<i32>),
}
impl IterSpec {
pub fn includes(&self, iter: i32) -> bool {
match self {
IterSpec::All => true,
IterSpec::Single(n) => iter == *n,
IterSpec::Range(lo, hi) => lo.is_none_or(|l| iter >= l) && hi.is_none_or(|h| iter <= h),
}
}
pub fn parse(s: &str) -> Result<Self, String> {
let s = s.trim();
if s.is_empty() || s == "all" {
return Ok(IterSpec::All);
}
if let Some(rest) = s.strip_prefix('-') {
let hi: i32 = rest.parse().map_err(|_| {
format!("invalid iter-spec '{s}': expected '-M' with non-negative integer M")
})?;
if hi < 0 {
return Err(format!(
"invalid iter-spec '{s}': iter must be non-negative"
));
}
return Ok(IterSpec::Range(None, Some(hi)));
}
if let Some((a, b)) = s.split_once('-') {
let lo: i32 = a
.parse()
.map_err(|_| format!("invalid iter-spec '{s}': '{a}' is not an integer"))?;
if lo < 0 {
return Err(format!(
"invalid iter-spec '{s}': iter must be non-negative"
));
}
if b.is_empty() {
return Ok(IterSpec::Range(Some(lo), None));
}
let hi: i32 = b
.parse()
.map_err(|_| format!("invalid iter-spec '{s}': '{b}' is not an integer"))?;
if hi < 0 {
return Err(format!(
"invalid iter-spec '{s}': iter must be non-negative"
));
}
if hi < lo {
return Err(format!(
"invalid iter-spec '{s}': end ({hi}) is below start ({lo})"
));
}
return Ok(IterSpec::Range(Some(lo), Some(hi)));
}
let n: i32 = s.parse().map_err(|_| {
format!("invalid iter-spec '{s}': expected 'all', 'N', 'N-M', 'N-', or '-M'")
})?;
if n < 0 {
return Err(format!(
"invalid iter-spec '{s}': iter must be non-negative"
));
}
Ok(IterSpec::Single(n))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DumpFormat {
Jsonl,
}
impl DumpFormat {
pub fn parse(s: &str) -> Result<Self, String> {
match s {
"jsonl" => Ok(DumpFormat::Jsonl),
other => Err(format!("unknown dump format '{other}' (expected: jsonl)")),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum IterateVariant {
#[default]
Summary,
Full,
}
impl IterateVariant {
pub fn as_str(self) -> &'static str {
match self {
IterateVariant::Summary => "summary",
IterateVariant::Full => "full",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum KktVariant {
#[default]
KOnly,
WithLPattern,
WithLValues,
}
impl KktVariant {
pub fn as_str(self) -> &'static str {
match self {
KktVariant::KOnly => "k-only",
KktVariant::WithLPattern => "with-l-pattern",
KktVariant::WithLValues => "with-l-values",
}
}
pub fn wants_l_pattern(self) -> bool {
matches!(self, KktVariant::WithLPattern | KktVariant::WithLValues)
}
pub fn wants_l_values(self) -> bool {
matches!(self, KktVariant::WithLValues)
}
}
pub fn parse_kkt_spec(s: &str) -> Result<(IterSpec, KktVariant), String> {
let s = s.trim();
let (rest, has_lvals) = match s.strip_suffix("+Lvals") {
Some(r) => (r, true),
None => (s, false),
};
let (filter_str, has_l) = match rest.strip_suffix("+L") {
Some(r) => (r, true),
None => (rest, false),
};
if has_lvals && !has_l {
return Err(format!(
"invalid kkt-spec '{s}': '+Lvals' requires '+L' (use '+L+Lvals' for L pattern with values)"
));
}
let variant = if has_lvals {
KktVariant::WithLValues
} else if has_l {
KktVariant::WithLPattern
} else {
KktVariant::KOnly
};
let filter_str = if filter_str.is_empty() {
"all"
} else {
filter_str
};
let filter = IterSpec::parse(filter_str)?;
Ok((filter, variant))
}
pub fn parse_iterate_spec(s: &str) -> Result<(IterSpec, IterateVariant), String> {
let s = s.trim();
if s == "summary" {
return Ok((IterSpec::All, IterateVariant::Summary));
}
if s == "full" {
return Ok((IterSpec::All, IterateVariant::Full));
}
let (filter_str, variant) = if let Some(rest) = s.strip_suffix(":summary") {
(rest, IterateVariant::Summary)
} else if let Some(rest) = s.strip_suffix(":full") {
(rest, IterateVariant::Full)
} else {
(s, IterateVariant::Summary)
};
let filter_str = if filter_str.is_empty() {
"all"
} else {
filter_str
};
let filter = IterSpec::parse(filter_str)?;
Ok((filter, variant))
}
#[derive(Debug, Clone)]
pub struct DiagnosticsConfig {
pub dump_dir: PathBuf,
pub format: DumpFormat,
pub categories: HashMap<DiagCategory, IterSpec>,
pub iterate_variant: IterateVariant,
pub kkt_variant: KktVariant,
}
impl DiagnosticsConfig {
pub fn new(dump_dir: PathBuf) -> Self {
Self {
dump_dir,
format: DumpFormat::Jsonl,
categories: HashMap::new(),
iterate_variant: IterateVariant::Summary,
kkt_variant: KktVariant::KOnly,
}
}
pub fn with_category(mut self, cat: DiagCategory, spec: IterSpec) -> Self {
self.categories.insert(cat, spec);
self
}
pub fn with_iterate_variant(mut self, v: IterateVariant) -> Self {
self.iterate_variant = v;
self
}
pub fn with_kkt_variant(mut self, v: KktVariant) -> Self {
self.kkt_variant = v;
self
}
pub fn is_empty(&self) -> bool {
self.categories.is_empty()
}
}
pub struct DiagnosticsState {
pub config: DiagnosticsConfig,
current_iter: AtomicI32,
solves_this_iter: AtomicI32,
in_restoration: AtomicBool,
resto_parent_iter: AtomicI32,
resto_inner_iter: AtomicI32,
resto_solves_this_iter: AtomicI32,
iterates_writer: RefCell<Option<BufWriter<fs::File>>>,
}
impl DiagnosticsState {
pub fn new(config: DiagnosticsConfig) -> std::io::Result<Self> {
fs::create_dir_all(&config.dump_dir)?;
Ok(Self {
config,
current_iter: AtomicI32::new(-1),
solves_this_iter: AtomicI32::new(0),
in_restoration: AtomicBool::new(false),
resto_parent_iter: AtomicI32::new(-1),
resto_inner_iter: AtomicI32::new(-1),
resto_solves_this_iter: AtomicI32::new(0),
iterates_writer: RefCell::new(None),
})
}
pub fn want(&self, cat: DiagCategory) -> bool {
let iter = self.effective_iter();
if iter < 0 {
return false;
}
self.config
.categories
.get(&cat)
.map(|spec| spec.includes(iter))
.unwrap_or(false)
}
pub fn bump_iter(&self) {
if self.in_restoration.load(Ordering::SeqCst) {
self.resto_inner_iter.fetch_add(1, Ordering::SeqCst);
self.resto_solves_this_iter.store(0, Ordering::SeqCst);
} else {
self.current_iter.fetch_add(1, Ordering::SeqCst);
self.solves_this_iter.store(0, Ordering::SeqCst);
}
}
pub fn next_solve_index(&self) -> i32 {
let counter = if self.in_restoration.load(Ordering::SeqCst) {
&self.resto_solves_this_iter
} else {
&self.solves_this_iter
};
counter.fetch_add(1, Ordering::SeqCst) + 1
}
pub fn enter_restoration(&self) {
let parent = self.current_iter.load(Ordering::SeqCst);
self.resto_parent_iter.store(parent, Ordering::SeqCst);
self.resto_inner_iter.store(-1, Ordering::SeqCst);
self.resto_solves_this_iter.store(0, Ordering::SeqCst);
self.in_restoration.store(true, Ordering::SeqCst);
}
pub fn exit_restoration(&self) {
self.in_restoration.store(false, Ordering::SeqCst);
}
pub fn current_iter(&self) -> i32 {
self.effective_iter()
}
pub fn in_restoration(&self) -> bool {
self.in_restoration.load(Ordering::SeqCst)
}
fn effective_iter(&self) -> i32 {
if self.in_restoration.load(Ordering::SeqCst) {
self.resto_inner_iter.load(Ordering::SeqCst)
} else {
self.current_iter.load(Ordering::SeqCst)
}
}
pub fn iter_dir(&self) -> Option<PathBuf> {
let dir = if self.in_restoration.load(Ordering::SeqCst) {
let parent = self.resto_parent_iter.load(Ordering::SeqCst);
let inner = self.resto_inner_iter.load(Ordering::SeqCst).max(0);
self.config
.dump_dir
.join(format!("resto/parent_iter_{parent:03}/iter_{inner:03}"))
} else {
let iter = self.current_iter.load(Ordering::SeqCst).max(0);
self.config.dump_dir.join(format!("iter_{iter:03}"))
};
fs::create_dir_all(&dir).ok()?;
Some(dir)
}
pub fn open_writer(&self, filename: &str) -> Option<BufWriter<fs::File>> {
let dir = self.iter_dir()?;
let path = dir.join(filename);
fs::File::create(path).ok().map(BufWriter::new)
}
pub fn write_top_level(&self, filename: &str, contents: &str) -> std::io::Result<()> {
let path = self.config.dump_dir.join(filename);
let mut f = fs::File::create(path)?;
f.write_all(contents.as_bytes())?;
f.flush()
}
pub fn append_iterate_line(&self, json: &str) -> std::io::Result<()> {
let mut slot = self.iterates_writer.borrow_mut();
if slot.is_none() {
let path = self.config.dump_dir.join("iterates.jsonl");
let f = fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(path)?;
*slot = Some(BufWriter::new(f));
}
let w = slot.as_mut().expect("just initialized");
w.write_all(json.as_bytes())?;
w.write_all(b"\n")?;
w.flush()
}
pub fn dump_dir(&self) -> &Path {
&self.config.dump_dir
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn iter_spec_parses_all_grammar_forms() {
assert_eq!(IterSpec::parse("").unwrap(), IterSpec::All);
assert_eq!(IterSpec::parse("all").unwrap(), IterSpec::All);
assert_eq!(IterSpec::parse("5").unwrap(), IterSpec::Single(5));
assert_eq!(
IterSpec::parse("5-10").unwrap(),
IterSpec::Range(Some(5), Some(10))
);
assert_eq!(
IterSpec::parse("5-").unwrap(),
IterSpec::Range(Some(5), None)
);
assert_eq!(
IterSpec::parse("-10").unwrap(),
IterSpec::Range(None, Some(10))
);
}
#[test]
fn iter_spec_rejects_malformed_input() {
assert!(IterSpec::parse("abc").is_err());
assert!(IterSpec::parse("5-3").is_err()); assert!(IterSpec::parse("-x").is_err());
assert!(IterSpec::parse("5--10").is_err()); }
#[test]
fn iter_spec_includes_matches_grammar() {
assert!(IterSpec::All.includes(0));
assert!(IterSpec::All.includes(1000));
assert!(IterSpec::Single(5).includes(5));
assert!(!IterSpec::Single(5).includes(4));
let r = IterSpec::Range(Some(5), Some(10));
assert!(!r.includes(4));
assert!(r.includes(5));
assert!(r.includes(7));
assert!(r.includes(10));
assert!(!r.includes(11));
assert!(IterSpec::Range(Some(5), None).includes(1_000_000));
assert!(IterSpec::Range(None, Some(5)).includes(0));
}
#[test]
fn category_parses_known_names() {
assert_eq!(DiagCategory::parse("kkt").unwrap(), DiagCategory::Kkt);
assert_eq!(
DiagCategory::parse("iterate").unwrap(),
DiagCategory::Iterate
);
assert!(DiagCategory::parse("bogus").is_err());
}
#[test]
fn iterate_spec_parses_all_combinations() {
assert_eq!(
parse_iterate_spec("summary").unwrap(),
(IterSpec::All, IterateVariant::Summary)
);
assert_eq!(
parse_iterate_spec("full").unwrap(),
(IterSpec::All, IterateVariant::Full)
);
assert_eq!(
parse_iterate_spec("all").unwrap(),
(IterSpec::All, IterateVariant::Summary)
);
assert_eq!(
parse_iterate_spec("5").unwrap(),
(IterSpec::Single(5), IterateVariant::Summary)
);
assert_eq!(
parse_iterate_spec("5-10").unwrap(),
(IterSpec::Range(Some(5), Some(10)), IterateVariant::Summary)
);
assert_eq!(
parse_iterate_spec("all:summary").unwrap(),
(IterSpec::All, IterateVariant::Summary)
);
assert_eq!(
parse_iterate_spec("all:full").unwrap(),
(IterSpec::All, IterateVariant::Full)
);
assert_eq!(
parse_iterate_spec("5-:full").unwrap(),
(IterSpec::Range(Some(5), None), IterateVariant::Full)
);
assert_eq!(
parse_iterate_spec("10-20:full").unwrap(),
(IterSpec::Range(Some(10), Some(20)), IterateVariant::Full)
);
}
#[test]
fn append_iterate_line_streams_rows_to_top_level() {
let tmp = tempdir();
let cfg =
DiagnosticsConfig::new(tmp.clone()).with_category(DiagCategory::Iterate, IterSpec::All);
let state = DiagnosticsState::new(cfg).unwrap();
state.append_iterate_line("{\"iter\":0}").unwrap();
state.append_iterate_line("{\"iter\":1}").unwrap();
state.enter_restoration();
state
.append_iterate_line("{\"iter\":0,\"restoration\":true}")
.unwrap();
state.exit_restoration();
state.append_iterate_line("{\"iter\":2}").unwrap();
let path = tmp.join("iterates.jsonl");
let contents = fs::read_to_string(&path).unwrap();
let lines: Vec<&str> = contents.lines().collect();
assert_eq!(lines.len(), 4);
assert_eq!(lines[0], "{\"iter\":0}");
assert_eq!(lines[2], "{\"iter\":0,\"restoration\":true}");
fs::remove_dir_all(tmp).ok();
}
#[test]
fn kkt_spec_parses_all_combinations() {
assert_eq!(
parse_kkt_spec("").unwrap(),
(IterSpec::All, KktVariant::KOnly)
);
assert_eq!(
parse_kkt_spec("all").unwrap(),
(IterSpec::All, KktVariant::KOnly)
);
assert_eq!(
parse_kkt_spec("5-10").unwrap(),
(IterSpec::Range(Some(5), Some(10)), KktVariant::KOnly)
);
assert_eq!(
parse_kkt_spec("+L").unwrap(),
(IterSpec::All, KktVariant::WithLPattern)
);
assert_eq!(
parse_kkt_spec("5-10+L").unwrap(),
(IterSpec::Range(Some(5), Some(10)), KktVariant::WithLPattern)
);
assert_eq!(
parse_kkt_spec("3+L").unwrap(),
(IterSpec::Single(3), KktVariant::WithLPattern)
);
assert_eq!(
parse_kkt_spec("+L+Lvals").unwrap(),
(IterSpec::All, KktVariant::WithLValues)
);
assert_eq!(
parse_kkt_spec("5-10+L+Lvals").unwrap(),
(IterSpec::Range(Some(5), Some(10)), KktVariant::WithLValues)
);
}
#[test]
fn kkt_spec_rejects_lvals_without_l() {
assert!(parse_kkt_spec("+Lvals").is_err());
assert!(parse_kkt_spec("5-10+Lvals").is_err());
}
#[test]
fn iterate_spec_rejects_garbage_and_unknown_variants() {
assert!(parse_iterate_spec("5-:bogus").is_err());
assert!(parse_iterate_spec("abc").is_err());
}
#[test]
fn state_gates_on_iter_spec() {
let tmp = tempdir();
let cfg = DiagnosticsConfig::new(tmp.clone())
.with_category(DiagCategory::Kkt, IterSpec::Range(Some(2), Some(4)));
let state = DiagnosticsState::new(cfg).unwrap();
assert!(!state.want(DiagCategory::Kkt));
state.bump_iter(); assert!(!state.want(DiagCategory::Kkt));
state.bump_iter(); assert!(!state.want(DiagCategory::Kkt));
state.bump_iter(); assert!(state.want(DiagCategory::Kkt));
state.bump_iter(); assert!(state.want(DiagCategory::Kkt));
state.bump_iter(); assert!(state.want(DiagCategory::Kkt));
state.bump_iter(); assert!(!state.want(DiagCategory::Kkt));
assert!(!state.want(DiagCategory::Iterate));
fs::remove_dir_all(tmp).ok();
}
#[test]
fn state_emits_solve_indices_and_iter_dirs() {
let tmp = tempdir();
let cfg =
DiagnosticsConfig::new(tmp.clone()).with_category(DiagCategory::Kkt, IterSpec::All);
let state = DiagnosticsState::new(cfg).unwrap();
state.bump_iter(); assert_eq!(state.next_solve_index(), 1);
assert_eq!(state.next_solve_index(), 2);
state.bump_iter(); assert_eq!(state.next_solve_index(), 1);
let dir = state.iter_dir().unwrap();
assert!(dir.ends_with("iter_001"));
fs::remove_dir_all(tmp).ok();
}
#[test]
fn restoration_dumps_live_under_resto_subtree() {
let tmp = tempdir();
let cfg =
DiagnosticsConfig::new(tmp.clone()).with_category(DiagCategory::Kkt, IterSpec::All);
let state = DiagnosticsState::new(cfg).unwrap();
state.bump_iter(); state.bump_iter(); state.enter_restoration();
state.bump_iter(); let dir = state.iter_dir().unwrap();
assert!(
dir.ends_with("resto/parent_iter_001/iter_000"),
"got {dir:?}"
);
assert_eq!(state.next_solve_index(), 1);
state.exit_restoration();
let dir = state.iter_dir().unwrap();
assert!(dir.ends_with("iter_001"), "got {dir:?}");
fs::remove_dir_all(tmp).ok();
}
fn tempdir() -> PathBuf {
let p = std::env::temp_dir().join(format!(
"pounce-diag-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
fs::create_dir_all(&p).unwrap();
p
}
}