use crate::decode::{Kind, Probe};
use image::ImageFormat;
use indicatif::{ProgressBar, ProgressFinish, ProgressState, ProgressStyle};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
const RULE: &str = " \u{2502} ";
const LEN: u64 = 1_000_000;
const FRAME: Duration = Duration::from_millis(50);
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Stage {
Identical,
CacheRead,
Headers,
Describe,
CacheWrite,
Vocabulary,
Quantise,
InvertedFile,
Candidates,
Verify,
Variants,
Propagate,
}
const STAGES: usize = 12;
impl Stage {
fn label(self) -> &'static str {
match self {
Stage::Identical => "finding identical files",
Stage::CacheRead => "reading the cache",
Stage::Headers => "reading headers",
Stage::Describe => "describing",
Stage::CacheWrite => "writing the cache",
Stage::Vocabulary => "building the vocabulary",
Stage::Quantise => "quantising",
Stage::InvertedFile => "building the inverted file",
Stage::Candidates => "finding candidates",
Stage::Verify => "verifying",
Stage::Variants => "mirrored and inverted",
Stage::Propagate => "propagating",
}
}
}
#[derive(Clone, Debug, Default)]
pub struct Forecast {
pub files: usize,
pub cache_bytes: u64,
pub to_describe: Option<usize>,
pub describe: Option<u64>,
pub cache_write: bool,
pub images: Option<usize>,
pub descriptors: Option<usize>,
pub vocab_sample: Option<usize>,
pub live_words: Option<usize>,
pub candidates_per_query: usize,
pub candidate_pairs: Option<usize>,
pub variants: Option<u64>,
pub edges: Option<usize>,
}
impl Forecast {
const HASH_PER_FILE: f64 = 20_000.0;
const CACHE_READ_PER_BYTE: f64 = 7.4;
const HEADER_PER_FILE: f64 = 60_000.0;
const DESCRIBE_PER_FILE: f64 = 20e6;
const KEYPOINTS_PER_IMAGE: f64 = 500.0;
const CACHE_WRITE_PER_KEYPOINT: f64 = 300.0;
const VOCAB_PER_SAMPLE: f64 = 26_000.0;
const VOCAB_SAMPLE: f64 = 160_000.0;
const QUANTISE_PER_DESC: f64 = 2_700.0;
const INVERT_PER_DESC: f64 = 150.0;
const QUERY_PER_POSTING: f64 = 55.0;
const VERIFY_PER_PAIR: f64 = 10_400.0;
const PAIRS_PER_QUERY: f64 = 0.75;
const LONELY_GUESS: f64 = 0.8;
const DESCRIBE_SCALE: f64 = 0.8;
const VARIANT_OVERHEAD: f64 = 1.3;
const EDGES_PER_IMAGE_GUESS: f64 = 2.0;
const PROPAGATE_PER_EDGE: f64 = 20_000.0;
fn images(&self) -> f64 {
self.images.unwrap_or(self.files) as f64
}
fn descriptors(&self) -> f64 {
self.descriptors.map(|d| d as f64).unwrap_or(self.images() * Self::KEYPOINTS_PER_IMAGE)
}
fn postings_per_word(&self) -> f64 {
let words = self.live_words.map(|w| w as f64).unwrap_or(self.descriptors().min(Self::VOCAB_SAMPLE));
(self.descriptors() / words.max(1.0)).max(1.0)
}
pub fn quantise_cost(&self, keypoints: usize) -> u64 {
(keypoints as f64 * Self::QUANTISE_PER_DESC) as u64
}
pub fn query_cost(&self, keypoints: usize) -> u64 {
(keypoints as f64 * self.postings_per_word() * Self::QUERY_PER_POSTING) as u64
}
pub fn variant_cost(&self, keypoints: usize) -> u64 {
let verify = (self.candidates_per_query as f64).min(self.images()) * Self::VERIFY_PER_PAIR;
let cost = 3 * (self.quantise_cost(keypoints) + self.query_cost(keypoints)) + verify as u64;
(cost as f64 * Self::VARIANT_OVERHEAD) as u64
}
fn estimates(&self) -> [f64; STAGES] {
let files = self.files as f64;
let images = self.images();
let desc = self.descriptors();
let keypoints = (desc / images.max(1.0)) as usize;
let pairs = self
.candidate_pairs
.map(|p| p as f64)
.unwrap_or(images * (self.candidates_per_query as f64).min(images - 1.0).max(0.0) * Self::PAIRS_PER_QUERY);
let variants = self
.variants
.map(|v| v as f64)
.unwrap_or(images * Self::LONELY_GUESS * self.variant_cost(keypoints) as f64);
let edges = self.edges.map(|e| e as f64).unwrap_or(images * Self::EDGES_PER_IMAGE_GUESS);
let mut e = [0.0; STAGES];
e[Stage::Identical as usize] = files * Self::HASH_PER_FILE;
e[Stage::CacheRead as usize] = self.cache_bytes as f64 * Self::CACHE_READ_PER_BYTE;
e[Stage::Headers as usize] = files * Self::HEADER_PER_FILE;
e[Stage::Describe as usize] = match (self.describe, self.to_describe) {
(Some(d), _) => d as f64 * Self::DESCRIBE_SCALE,
(None, Some(n)) => n as f64 * Self::DESCRIBE_PER_FILE * Self::DESCRIBE_SCALE,
(None, None) => {
(files - (self.cache_bytes as f64 / 50_000.0).min(files)) * Self::DESCRIBE_PER_FILE * Self::DESCRIBE_SCALE
}
};
e[Stage::CacheWrite as usize] =
if self.cache_write { desc * Self::CACHE_WRITE_PER_KEYPOINT } else { 0.0 };
e[Stage::Vocabulary as usize] =
self.vocab_sample.map(|s| s as f64).unwrap_or(desc.min(Self::VOCAB_SAMPLE)) * Self::VOCAB_PER_SAMPLE;
e[Stage::Quantise as usize] = desc * Self::QUANTISE_PER_DESC;
e[Stage::InvertedFile as usize] = desc * Self::INVERT_PER_DESC;
e[Stage::Candidates as usize] = desc * self.postings_per_word() * Self::QUERY_PER_POSTING;
e[Stage::Verify as usize] = pairs * Self::VERIFY_PER_PAIR;
e[Stage::Variants as usize] = variants;
e[Stage::Propagate as usize] = edges * Self::PROPAGATE_PER_EDGE;
e
}
}
#[derive(Default)]
pub struct Counter {
weight: AtomicU64,
items: AtomicU64,
}
impl Counter {
#[inline]
pub fn add(&self, weight: u64) {
self.weight.fetch_add(weight, Ordering::Relaxed);
self.items.fetch_add(1, Ordering::Relaxed);
}
#[inline]
pub fn tick(&self) {
self.add(1);
}
}
struct Counted {
counter: Arc<Counter>,
weight: u64,
items: u64,
unit: &'static str,
}
struct Current {
stage: Stage,
from: f64,
span: f64,
started: Instant,
counted: Option<Counted>,
expected: f64,
}
struct State {
estimates: [f64; STAGES],
at: f64,
current: Option<Current>,
measured_cost: f64,
measured_secs: f64,
}
impl State {
fn rate(&self) -> f64 {
if self.measured_secs >= 0.5 && self.measured_cost > 0.0 {
self.measured_cost / self.measured_secs
} else {
3e8 * rayon::current_num_threads() as f64
}
}
fn close(&mut self) {
if let Some(c) = self.current.take() {
self.at = self.at.max(c.from + c.span);
if c.counted.is_some() {
self.measured_cost += self.estimates[c.stage as usize];
self.measured_secs += c.started.elapsed().as_secs_f64();
}
}
}
fn frame(&mut self) -> Option<(f64, String)> {
let c = self.current.as_ref()?;
let secs = c.started.elapsed().as_secs_f64();
let (p, msg) = match &c.counted {
Some(k) => {
let w = k.counter.weight.load(Ordering::Relaxed);
let n = k.counter.items.load(Ordering::Relaxed).min(k.items);
let p = if k.weight == 0 { 1.0 } else { (w as f64 / k.weight as f64).min(1.0) };
let mut msg = format!("{} {}/{} {}", c.stage.label(), n, k.items, k.unit);
if c.stage == Stage::Describe {
msg.push_str(RULE);
msg.push_str(&image_rate(n as f64 / secs));
}
(p, msg)
}
None => (creep(secs / c.expected.max(1e-3)), c.stage.label().to_string()),
};
let at = (c.from + c.span * p).max(self.at);
self.at = at;
Some((at, msg))
}
}
fn creep(x: f64) -> f64 {
const KNEE: f64 = 0.9;
if x <= KNEE {
x
} else {
KNEE + (1.0 - KNEE) * (1.0 - (-(x - KNEE) / (1.0 - KNEE)).exp())
}
}
pub struct Progress {
bar: ProgressBar,
state: Arc<Mutex<State>>,
forecast: Mutex<Forecast>,
stop: Arc<AtomicBool>,
drawer: Mutex<Option<JoinHandle<()>>>,
}
impl Progress {
pub fn new() -> Progress {
let bar = ProgressBar::new(LEN).with_finish(ProgressFinish::AndClear);
let _ = BAR.set(bar.clone());
bar.set_style(
ProgressStyle::with_template(&format!(
"{{elapsed_precise}}{RULE}[{{bar:28.cyan/blue}}]{RULE}{{pct}}{RULE}{{msg}}"
))
.unwrap()
.with_key("pct", |s: &ProgressState, w: &mut dyn std::fmt::Write| {
let _ = write!(w, "{:5.1}%", s.fraction() * 100.0);
})
.progress_chars("=>-"),
);
bar.set_message("scanning");
bar.enable_steady_tick(Duration::from_millis(200));
let state = Arc::new(Mutex::new(State {
estimates: [0.0; STAGES],
at: 0.0,
current: None,
measured_cost: 0.0,
measured_secs: 0.0,
}));
let stop = Arc::new(AtomicBool::new(false));
let drawer = {
let (bar, state, stop) = (bar.clone(), state.clone(), stop.clone());
std::thread::spawn(move || {
while !stop.load(Ordering::Relaxed) {
std::thread::sleep(FRAME);
let frame = state.lock().unwrap().frame();
if let Some((at, msg)) = frame {
bar.set_position((at * LEN as f64) as u64);
bar.set_message(msg);
}
}
})
};
Progress { bar, state, forecast: Mutex::new(Forecast::default()), stop, drawer: Mutex::new(Some(drawer)) }
}
pub fn println(&self, line: &str) {
self.bar.suspend(|| eprintln!("{line}"));
}
pub fn forecast(&self, f: impl FnOnce(&mut Forecast)) {
let mut fc = self.forecast.lock().unwrap();
f(&mut fc);
self.state.lock().unwrap().estimates = fc.estimates();
}
pub fn forecast_now(&self) -> Forecast {
self.forecast.lock().unwrap().clone()
}
fn start(&self, stage: Stage, counted: Option<Counted>) {
let mut s = self.state.lock().unwrap();
s.close();
let remaining: f64 = s.estimates[stage as usize..].iter().sum();
let own = s.estimates[stage as usize];
let span = if remaining > 0.0 { (1.0 - s.at) * own / remaining } else { 0.0 };
let expected = own / s.rate();
let (from, label) = (s.at, stage.label());
s.current = Some(Current { stage, from, span, started: Instant::now(), counted, expected });
drop(s);
self.bar.set_message(label);
}
pub fn begin(&self, stage: Stage) {
self.start(stage, None);
}
pub fn begin_counted(&self, stage: Stage, weight: u64, items: usize, unit: &'static str) -> Arc<Counter> {
let counter = Arc::new(Counter::default());
self.start(stage, Some(Counted { counter: counter.clone(), weight, items: items as u64, unit }));
counter
}
pub fn finish(&self) {
self.state.lock().unwrap().close();
self.stop.store(true, Ordering::Relaxed);
if let Some(h) = self.drawer.lock().unwrap().take() {
let _ = h.join();
}
self.bar.finish_and_clear();
}
}
static BAR: std::sync::OnceLock<ProgressBar> = std::sync::OnceLock::new();
pub fn clear_for_exit() {
if let Some(bar) = BAR.get() {
bar.finish_and_clear();
}
}
impl Drop for Progress {
fn drop(&mut self) {
self.finish();
}
}
pub fn image_rate(per_sec: f64) -> String {
if !per_sec.is_finite() || per_sec <= 0.0 {
return "-".to_string();
}
let (scaled, unit) = if per_sec >= 1e4 { (per_sec / 1e3, "k img/s") } else { (per_sec, " img/s") };
let number = if scaled >= 100.0 {
format!("{scaled:.0}")
} else if scaled >= 10.0 {
format!("{scaled:.1}")
} else {
format!("{scaled:.2}")
};
if number.len() > 4 {
return "-".to_string();
}
format!("{number}{unit}")
}
pub fn analysis_cost(probe: Option<&Probe>, bytes: u64, work: usize, upsample_below: usize) -> u64 {
let (kind, w, h) = match probe {
Some(p) => (p.kind, p.w as f64, p.h as f64),
None => {
let side = (bytes as f64 / 0.48 / 0.75).sqrt();
(Kind::Image(ImageFormat::Jpeg), side, side * 0.75)
}
};
let (per_px, per_byte) = match kind {
Kind::Image(ImageFormat::Jpeg) => (4.4, 22.0),
Kind::Image(ImageFormat::Png) => (4.5, 8.5),
Kind::Image(ImageFormat::WebP) => (35.0, 0.0),
Kind::Image(ImageFormat::Tiff) => (37.0, 0.0),
Kind::Heif => (40.0, 290.0),
Kind::Jxl => (123.0, 0.0),
_ => (10.0, 0.0),
};
let decode = per_px * w * h + per_byte * bytes as f64;
let long = w.max(h).max(1.0);
let s = if work > 0 && long > work as f64 { work as f64 / long } else { 1.0 };
let (bw, bh) = ((w * s).round().max(1.0), (h * s).round().max(1.0));
let mut factor = 1.0;
while bw.max(bh) * factor * 2.0 <= upsample_below.max(2) as f64 {
factor *= 2.0;
}
let extract = 118.0 * bw * bh * factor * factor;
const PER_FILE: f64 = 200_000.0;
(decode + extract + PER_FILE) as u64
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rate_holds_three_figures_and_climbs_the_ladder() {
assert_eq!(image_rate(412.3), "412 img/s");
assert_eq!(image_rate(56.24), "56.2 img/s");
assert_eq!(image_rate(2.5), "2.50 img/s");
assert_eq!(image_rate(12_345.0), "12.3k img/s");
assert_eq!(image_rate(0.0), "-");
assert_eq!(image_rate(f64::NAN), "-");
}
#[test]
fn cost_follows_the_format_and_the_size() {
let jpeg = |w, h| Probe { kind: Kind::Image(ImageFormat::Jpeg), w, h };
let jxl = Probe { kind: Kind::Jxl, w: 4000, h: 3000 };
let big = analysis_cost(Some(&jpeg(4000, 3000)), 3_000_000, 384, 512);
let small = analysis_cost(Some(&jpeg(224, 224)), 20_000, 384, 512);
let jx = analysis_cost(Some(&jxl), 1_000_000, 384, 512);
assert!(big > 5 * small, "{big} {small}");
assert!(jx > 2 * big, "{jx} {big}");
assert!(small > 118 * 448 * 448, "{small}");
}
}