use std::fs::File;
use std::io::BufWriter;
use std::sync::Mutex;
use std::time::Instant;
use clap::{Parser, ValueEnum};
use tilezz::dataset::{DEFAULT_TARGET_BLOCK_BYTES, RatDafsa};
use tilezz::enumerate::dfs::STREAM_RAT_LINES;
use tilezz::enumerate::output::{print_stats, rats_gif, run_rat_enum_polylines};
use tilezz::enumerate::prune::{
install_closure_table_prune, install_modular_prune, install_shadow_prune,
};
use tilezz::enumerate::run_rat_enum_seqs;
use tilezz::enumerate::seed::{dispatch_collect_seed_prefixes, dispatch_enumerate_from_seed};
use tilezz::vis::plotutils::P64;
static VERBOSE: Mutex<bool> = Mutex::new(false);
#[derive(Copy, Clone, Debug, ValueEnum)]
enum Mode {
Render,
Bench,
ListSeeds,
Dafsa,
DafsaBlocks,
Stream,
Merge,
Build,
Rehost,
}
#[derive(Parser, Debug)]
#[command(
version = tilezz::VERSION,
about = "Enumerate simple cyclotomic matchstick polygons (rats) up to a given perimeter",
long_about = "\
Enumerate every simple polygon (closed self-avoiding boundary) of\n\
unit-length edges on a cyclotomic ring `ZZn`, with perimeter up to `-n`.\n\
\n\
Output is selected via `--mode` (default: render):\n\
* render enumerate and render the polygons as a GIF (`-o file.gif`)\n\
* bench enumerate, time, report counts (no -o needed)\n\
* dafsa enumerate and write a gzipped DAFSA (`-o file.bin.gz`)\n\
* dafsa-blocks same, but as a directory of lazy-loadable blocks\n\
* stream stage 1 of the streaming pipeline: write per-thread runs\n\
to `<-o dir>/runs/run_tNN_rMM.bin` (bounded memory)\n\
* merge stage 2 of the streaming pipeline: k-way merge runs into\n\
`<-o dir>/unique.bin` + `certificate.json`\n\
* build stage 3 of the streaming pipeline: stream `unique.bin`\n\
through a DAFSA builder and write blocks to `<-o dir>/dafsa/`\n\
* list-seeds walk DFS to `--split-depth` and print SEED + RAT lines\n\
(for multi-host orchestration; see `--seed`)\n\
* rehost patch an existing dataset's RO-Crate host fields to a\n\
new `--base-url` / `--doi` (no recompute; see -o)\n\
\n\
Performance opts (combine for maximum speedup):\n\
--reachability-prune reachability prune, finite + archimedean places; up to 376x\n\
--closure-table-prune exact lattice closure-table prune; another 2-5x on top\n\
--threads N parallel DFS; near-linear in streaming modes,\n\
sub-linear in HashSet modes (set-merge overhead)\n\
--free free (= full dihedral reduction) output:\n\
one rep per chiral pair (lex-min over rotations and\n\
reflections); also accelerates DFS via\n\
complement-reflection prune\n\
\n\
For memory at large n: prefer the streaming pipeline (`--mode stream`\n\
then `--mode merge`) over `--mode bench --threads N`. See the file-level\n\
docstring in src/bin/rat_enum.rs for the full memory accounting.\n"
)]
struct Cli {
#[arg(short = 'r', long)]
ring: Option<u8>,
#[arg(short = 'n', long)]
max_steps: Option<usize>,
#[arg(long, value_enum, default_value_t = Mode::Render)]
mode: Mode,
#[arg(short = 'o', long)]
filename: Option<String>,
#[arg(short, long)]
verbose: bool,
#[arg(long)]
profile: Option<String>,
#[arg(long)]
domino: bool,
#[arg(long)]
stats: bool,
#[arg(long, default_value_t = 1)]
threads: usize,
#[arg(long, default_value_t = 1)]
step: i8,
#[arg(long)]
free: bool,
#[arg(long)]
paranoid: bool,
#[arg(long)]
reachability_prune: bool,
#[arg(long, value_delimiter = ',', num_args = 1..)]
reachability_moduli: Option<Vec<i64>>,
#[arg(long)]
closure_table_prune: bool,
#[arg(long, default_value_t = 4)]
closure_table_depth: usize,
#[arg(long, default_value_t = 0)]
heartbeat: u64,
#[arg(long, default_value_t = DEFAULT_TARGET_BLOCK_BYTES)]
target_block_bytes: u32,
#[arg(long)]
no_rocrate: bool,
#[arg(long)]
oeis_a_number: Option<String>,
#[arg(long)]
base_url: Option<String>,
#[arg(long)]
doi: Option<String>,
#[arg(long, value_delimiter = ',', allow_hyphen_values = true, num_args = 1)]
seed: Option<Vec<i8>>,
#[arg(long)]
split_depth: Option<usize>,
}
fn resolve_n_threads(requested: usize) -> usize {
let max = tilezz::util::available_workers();
if requested == 0 {
max
} else {
requested.min(max)
}
}
fn main() {
let cli = Cli::parse();
if cli.verbose {
let mut verbose = VERBOSE.lock().unwrap();
*verbose = true;
}
let n_threads = resolve_n_threads(cli.threads);
if matches!(cli.mode, Mode::Rehost) {
let Some(dir) = cli.filename.as_deref() else {
eprintln!("--mode rehost requires -o <existing dataset directory>");
std::process::exit(2);
};
let dir = std::path::Path::new(dir);
match tilezz::dataset::rehost_ro_crate(dir, cli.base_url.as_deref(), cli.doi.as_deref()) {
Ok(()) => {
let where_to = match (cli.base_url.as_deref(), cli.doi.as_deref()) {
(Some(b), Some(d)) => format!("base_url={b} doi={d}"),
(Some(b), None) => format!("base_url={b}"),
(None, Some(d)) => format!("doi={d}"),
(None, None) => "location-independent (relative) form".to_string(),
};
eprintln!(
"rehost: patched {}/ro-crate-metadata.json to {where_to}",
dir.display()
);
}
Err(e) => {
eprintln!("--mode rehost failed: {e}");
std::process::exit(2);
}
}
return;
}
let Some(ring) = cli.ring else {
eprintln!("--ring is required for every mode except --mode rehost");
std::process::exit(2);
};
let Some(max_steps) = cli.max_steps else {
eprintln!("-n / --max-steps is required for every mode except --mode rehost");
std::process::exit(2);
};
if cli.reachability_prune {
install_modular_prune(ring, max_steps, cli.reachability_moduli.as_deref());
install_shadow_prune(ring);
}
if cli.closure_table_prune {
install_closure_table_prune(ring, cli.closure_table_depth);
}
if matches!(cli.mode, Mode::ListSeeds) {
let split_depth = cli.split_depth.unwrap_or(3);
let (closed, seeds) =
dispatch_collect_seed_prefixes(ring, max_steps, cli.step, split_depth, cli.free);
for prefix in &seeds {
let s = prefix
.iter()
.map(|a| a.to_string())
.collect::<Vec<_>>()
.join(",");
println!("SEED {s}");
}
for seq in &closed {
println!("RAT {seq:?}");
}
eprintln!(
"list-seeds: {} alive prefixes at depth {}, {} polygons closed before split",
seeds.len(),
split_depth,
closed.len(),
);
return;
}
if let Some(seed) = cli.seed.as_deref() {
let t0 = Instant::now();
let rats = dispatch_enumerate_from_seed(
ring,
max_steps,
cli.step,
seed,
n_threads,
cli.free,
cli.paranoid,
);
let dt = t0.elapsed();
for seq in &rats {
println!("RAT {seq:?}");
}
eprintln!("seed {:?}: {} unique rats in {dt:?}", seed, rats.len());
return;
}
match cli.mode {
Mode::ListSeeds => unreachable!("handled above"),
Mode::Rehost => unreachable!("handled above"),
Mode::Build => {
let Some(out_dir) = cli.filename.as_deref() else {
eprintln!("--mode build requires -o <output directory>");
std::process::exit(2);
};
let out_dir = std::path::Path::new(out_dir);
let unique_path = out_dir.join(tilezz::enumerate::stream::UNIQUE_FILENAME);
if !unique_path.exists() {
eprintln!(
"--mode build: {} not found; run `--mode merge` first",
unique_path.display()
);
std::process::exit(2);
}
let t0 = Instant::now();
let records = tilezz::enumerate::stream::read_unique_records(&unique_path)
.expect("open unique.bin")
.map(|r| r.expect("read unique record"));
let dafsa = RatDafsa::from_sorted_unique_rats(records);
let t_build = t0.elapsed();
eprintln!(
"build: streamed {} rats into RatDafsa in {:?}",
dafsa.len(),
t_build
);
let blocks_dir = out_dir.join("dafsa");
std::fs::create_dir_all(&blocks_dir).expect("create dafsa/");
let t1 = Instant::now();
dafsa
.write_blocks(&blocks_dir, cli.target_block_bytes)
.expect("write_blocks");
if !cli.no_rocrate {
use tilezz::dataset::{
AssetParams, ProducedVia, SequenceCounts, write_archival_extras, write_ro_crate,
};
let params = AssetParams {
ring,
max_steps,
step: cli.step,
free: cli.free,
target_block_bytes: cli.target_block_bytes,
n_sequences: dafsa.len() as u64,
oeis_a_number: cli.oeis_a_number.as_deref(),
produced_via: ProducedVia::StreamingPipeline,
};
let counts = SequenceCounts::from_rats(dafsa.iter());
write_archival_extras(&blocks_dir, ¶ms).expect("write archival extras");
write_ro_crate(
&blocks_dir,
¶ms,
&counts,
cli.base_url.as_deref(),
cli.doi.as_deref(),
)
.expect("write ro-crate-metadata.json");
}
fn dir_size_recursive(p: &std::path::Path) -> u64 {
let mut total = 0u64;
if let Ok(rd) = std::fs::read_dir(p) {
for entry in rd.flatten() {
let path = entry.path();
if path.is_dir() {
total += dir_size_recursive(&path);
} else if let Ok(m) = entry.metadata() {
total += m.len();
}
}
}
total
}
let total_bytes = dir_size_recursive(&blocks_dir);
eprintln!(
"build: wrote {} ({} bytes) in {:?}",
blocks_dir.display(),
total_bytes,
t1.elapsed()
);
}
Mode::Merge => {
let Some(out_dir) = cli.filename.as_deref() else {
eprintln!("--mode merge requires -o <output directory>");
std::process::exit(2);
};
let t0 = Instant::now();
let cert = tilezz::enumerate::stream::merge_runs(
std::path::Path::new(out_dir),
ring,
max_steps,
cli.step,
cli.free,
)
.expect("merge_runs");
let dt = t0.elapsed();
println!(
"merge: ring={} step={} max_steps={} -> {} unique rats in {:?}",
ring, cli.step, max_steps, cert.unique_records, dt
);
}
Mode::Stream => {
let Some(out_dir) = cli.filename.as_deref() else {
eprintln!("--mode stream requires -o <output directory>");
std::process::exit(2);
};
STREAM_RAT_LINES.store(false, std::sync::atomic::Ordering::Relaxed);
let t0 = Instant::now();
let stats = tilezz::enumerate::stream::stream_enum_dispatch(
ring,
max_steps,
cli.step,
n_threads,
cli.free,
cli.paranoid,
cli.domino,
std::path::Path::new(out_dir),
cli.heartbeat,
)
.expect("stream_enum_dispatch");
let dt = t0.elapsed();
println!(
"stream: ring={} step={} max_steps={} -> wrote runs to {} in {:?}",
ring, cli.step, max_steps, out_dir, dt
);
println!("{stats}");
}
Mode::Dafsa => {
let Some(filename) = cli.filename.as_deref() else {
eprintln!("--mode dafsa requires -o <output path>");
std::process::exit(2);
};
STREAM_RAT_LINES.store(false, std::sync::atomic::Ordering::Relaxed);
let t0 = Instant::now();
let (rats, _stats) = run_rat_enum_seqs(
ring,
max_steps,
cli.step,
n_threads,
cli.free,
cli.paranoid,
cli.domino,
);
eprintln!("enumerated {} rats in {:?}", rats.len(), t0.elapsed());
let t1 = Instant::now();
let dafsa = RatDafsa::from_rats(rats.iter().map(|r| r.as_slice()));
eprintln!(
"built RatDafsa ({} entries) in {:?}",
dafsa.len(),
t1.elapsed()
);
let t2 = Instant::now();
let file = File::create(filename).expect("create output file");
dafsa
.write_json_gz(BufWriter::new(file))
.expect("write gzipped RatDafsa");
let bytes = std::fs::metadata(filename).map(|m| m.len()).unwrap_or(0);
eprintln!("wrote {filename} ({bytes} bytes) in {:?}", t2.elapsed());
if cli.stats {
print_stats(&rats);
}
}
Mode::DafsaBlocks => {
let Some(dir) = cli.filename.as_deref() else {
eprintln!("--mode dafsa-blocks requires -o <output directory>");
std::process::exit(2);
};
STREAM_RAT_LINES.store(false, std::sync::atomic::Ordering::Relaxed);
let t0 = Instant::now();
let (rats, _stats) = run_rat_enum_seqs(
ring,
max_steps,
cli.step,
n_threads,
cli.free,
cli.paranoid,
cli.domino,
);
eprintln!("enumerated {} rats in {:?}", rats.len(), t0.elapsed());
let t1 = Instant::now();
let dafsa = RatDafsa::from_rats(rats.iter().map(|r| r.as_slice()));
eprintln!(
"built RatDafsa ({} entries) in {:?}",
dafsa.len(),
t1.elapsed()
);
let t2 = Instant::now();
let path = std::path::Path::new(dir);
std::fs::create_dir_all(path).expect("create output dir");
dafsa
.write_blocks(path, cli.target_block_bytes)
.expect("write blocked RatDafsa");
if !cli.no_rocrate {
use tilezz::dataset::{
AssetParams, ProducedVia, SequenceCounts, write_archival_extras, write_ro_crate,
};
let params = AssetParams {
ring,
max_steps,
step: cli.step,
free: cli.free,
target_block_bytes: cli.target_block_bytes,
n_sequences: dafsa.len() as u64,
oeis_a_number: cli.oeis_a_number.as_deref(),
produced_via: ProducedVia::InMemory,
};
let counts = SequenceCounts::from_rats(dafsa.iter());
write_archival_extras(path, ¶ms).expect("write archival extras");
write_ro_crate(
path,
¶ms,
&counts,
cli.base_url.as_deref(),
cli.doi.as_deref(),
)
.expect("write ro-crate-metadata.json");
}
fn dir_size(p: &std::path::Path) -> u64 {
let mut total = 0u64;
if let Ok(rd) = std::fs::read_dir(p) {
for entry in rd.flatten() {
let path = entry.path();
if path.is_dir() {
total += dir_size(&path);
} else if let Ok(m) = entry.metadata() {
total += m.len();
}
}
}
total
}
let total_bytes = dir_size(path);
let extras = if cli.no_rocrate {
""
} else {
" + schemas + ro-crate-metadata"
};
eprintln!(
"wrote {dir}/ ({total_bytes} bytes across manifest + blocks{extras}) in {:?}",
t2.elapsed()
);
if cli.stats {
print_stats(&rats);
}
}
Mode::Bench => {
STREAM_RAT_LINES.store(false, std::sync::atomic::Ordering::Relaxed);
let profile = tilezz::util::profile::ProfileGuard::start(cli.profile.as_deref());
let t0 = Instant::now();
let (rats, stats) = run_rat_enum_seqs(
ring,
max_steps,
cli.step,
n_threads,
cli.free,
cli.paranoid,
cli.domino,
);
let dt = t0.elapsed();
let total_boundary_len: usize = rats.iter().map(|s| s.len()).sum();
println!(
"benchmark: backend={} ring={} step={} max_steps={} -> {} unique rats (total boundary len={}) in {:?}",
if cli.domino { "automaton" } else { "snake" },
ring,
cli.step,
max_steps,
rats.len(),
total_boundary_len,
dt
);
println!("{stats}");
profile.finish();
if cli.stats {
print_stats(&rats);
}
}
Mode::Render => {
let rats: Vec<Vec<P64>> = run_rat_enum_polylines(
ring,
max_steps,
cli.step,
n_threads,
cli.free,
cli.paranoid,
cli.domino,
);
let Some(filename) = cli.filename else {
return;
};
let gif_bytes = rats_gif(&rats, 500, 500);
std::fs::write(&filename, gif_bytes).expect("write GIF");
println!("wrote {filename}");
}
}
}