usenet-indexer/tools/edge-bench-rs/src/indexer.rs

//! lib/edge/scan.ts (SegmentScanner, scanOverviewText), lib/files/aggregator.ts
//! and lib/edge/range_indexer.ts (RangeIndexer, chunkOf, candidatesOf), in the
//! same order of operations, so the outputs can be compared one to one.

use crate::codec::{encode_chunk, ChunkFile};
use crate::js::{self, sha1_hex};
use crate::names::{self, is_random_file, JudgedFile, Release, Verdict};
use rustc_hash::{FxHashMap, FxHashSet};
use std::rc::Rc;

pub static mut PROF: [u64; 8] = [0; 8];

pub struct FileMeta {
    pub file_key: String,
    pub name: String,
    pub poster: String,
    pub subject: String,
    pub total_parts: u32,
    pub groups: Vec<String>,
    pub release: Release,
}

pub struct Seg {
    pub file: Rc<FileMeta>,
    pub number: u32,
    pub message_id: String,
    pub bytes: i64,
    pub article: i64,
    pub posted_at: Option<i64>,
}

pub struct ReleaseChunk {
    pub key: String,
    pub name: String,
    pub source: &'static str,
    pub file_total: Option<i64>,
    pub poster: String,
    pub groups: Vec<String>,
    pub posted_at: Option<i64>,
    pub files: usize,
    pub segments: usize,
    pub bytes: i64,
    pub first_article: i64,
    pub last_article: i64,
    pub readable: bool,
    pub has_par2: bool,
    pub payload: Vec<u8>,
}

pub struct Candidate {
    pub release_key: String,
    pub kind: &'static str,
    pub message_id: String,
    pub file: String,
}

pub struct Fragment {
    pub file_key: String,
    pub name: String,
    pub total_parts: u32,
    pub posted_at: Option<i64>,
    pub segments: Vec<(u32, String, i64, i64)>,
}

#[derive(Default)]
pub struct Batch {
    pub from: i64,
    pub next: i64,
    pub headers: u64,
    pub releases: Vec<ReleaseChunk>,
    pub samples: Vec<(ReleaseChunk, &'static str, String)>,
    pub candidates: Vec<Candidate>,
    /// (rule, key) → (articles, files)
    pub counts: Vec<(&'static str, String, u64, u64)>,
    pub fragments: Vec<Fragment>,
}

struct Counter {
    index: usize,
    length: usize,
    segment: i64,
    total: i64,
}

/// The last "(k/N)" after `after`, scanning back from the end (findCounter).
fn find_counter(subject: &[u8], after: usize) -> Option<Counter> {
    let mut open = subject.len();
    loop {
        // lastIndexOf('(', open - 1): at most open - 1.
        if open == 0 {
            return None;
        }
        open = subject[..open].iter().rposition(|&c| c == b'(')?;
        if open <= after {
            return None;
        }
        let mut i = open + 1;
        let mut segment: i64 = 0;
        let mut digits = 0;
        while let Some(c) = subject.get(i).filter(|c| c.is_ascii_digit()) {
            segment = segment * 10 + (c - b'0') as i64;
            digits += 1;
            i += 1;
        }
        if digits == 0 || subject.get(i) != Some(&b'/') {
            continue;
        }
        i += 1;
        let mut total: i64 = 0;
        digits = 0;
        while let Some(c) = subject.get(i).filter(|c| c.is_ascii_digit()) {
            total = total * 10 + (c - b'0') as i64;
            digits += 1;
            i += 1;
        }
        if digits == 0 || subject.get(i) != Some(&b')') {
            continue;
        }
        return Some(Counter { index: open, length: i + 1 - open, segment, total });
    }
}

fn xref_groups(xref: &str, group: &str) -> Vec<String> {
    let mut groups = Vec::new();
    for token in js::split_spaces(js::trim(xref)).into_iter().skip(1) {
        if let Some(colon) = token.rfind(':') {
            if colon > 0 {
                groups.push(token[..colon].to_string());
            }
        }
    }
    if groups.is_empty() {
        groups.push(group.to_string());
    }
    groups
}

#[derive(Default)]
struct Scanner {
    files: FxHashMap<String, Option<Rc<FileMeta>>>,
    key: String,
}

impl Scanner {
    fn owner<'a>(&self, xref: &'a str, group: &'a str, priorities: &FxHashMap<String, u32>) -> &'a str {
        let b = xref.as_bytes();
        let Some(mut at) = xref.find(' ') else { return group };
        let mut owner: Option<&str> = None;
        let mut best = u32::MAX;
        loop {
            while b.get(at) == Some(&b' ') {
                at += 1;
            }
            if at >= b.len() {
                break;
            }
            let end = xref[at..].find(' ').map(|i| i + at);
            // lastIndexOf(':', end ?? length): a colon at or before `end`.
            let limit = end.unwrap_or(b.len());
            let colon = xref[..limit.min(b.len())].rfind(':');
            if let Some(colon) = colon.filter(|&c| c > at) {
                let name = &xref[at..colon];
                if let Some(&p) = priorities.get(name) {
                    if p < best {
                        best = p;
                        owner = Some(name);
                    }
                }
            }
            match end {
                Some(e) => at = e,
                None => break,
            }
        }
        owner.unwrap_or(group)
    }

    #[allow(clippy::too_many_arguments)]
    fn segment(&mut self, article: i64, subject: &str, from: &str, date: &str, message_id: &str, bytes: &str, xref: &str, group: &str) -> Option<Seg> {
        if subject.is_empty() || message_id.is_empty() {
            return None;
        }
        let sb = subject.as_bytes();
        let name_end = subject.rfind('"')?;
        if name_end == 0 {
            return None;
        }
        let counter = find_counter(sb, name_end)?;
        self.key.clear();
        self.key.push_str(from);
        self.key.push('\u{1}');
        self.key.push_str(&subject[..counter.index]);
        self.key.push('\u{0}');
        self.key.push_str(&subject[counter.index + counter.length..]);
        let file = match self.files.get(self.key.as_str()) {
            Some(f) => f.clone(),
            None => {
                let f = describe(subject, name_end, &counter, from, xref, group).map(Rc::new);
                self.files.insert(self.key.clone(), f.clone());
                f
            }
        };
        let file = file?;
        if counter.segment < 1 || counter.segment > file.total_parts as i64 {
            return None;
        }
        let id = if message_id.starts_with('<') { message_id.get(1..message_id.len().saturating_sub(1)).unwrap_or("") } else { message_id };
        Some(Seg {
            number: counter.segment as u32,
            message_id: id.to_string(),
            bytes: js::parse_int(bytes).unwrap_or(0),
            article,
            posted_at: js::parse_date(date),
            file,
        })
    }
}

fn describe(subject: &str, name_end: usize, counter: &Counter, from: &str, xref: &str, group: &str) -> Option<FileMeta> {
    let name_start = subject[..name_end].rfind('"')?;
    let name = &subject[name_start + 1..name_end];
    if name.is_empty() || counter.total < 1 {
        return None;
    }
    let joined = format!("{}{}", &subject[..counter.index], &subject[counter.index + counter.length..]);
    let base = js::trim(&js::collapse_spaces(&joined)).to_string();
    let file_key = sha1_hex(&format!("{from}\n{base}\n{}", counter.total));
    let release = names::release_identity(from, &base, name);
    Some(FileMeta {
        file_key,
        name: name.to_string(),
        poster: from.to_string(),
        subject: base,
        total_parts: counter.total as u32,
        groups: xref_groups(xref, group),
        release,
    })
}

/// One line of an XOVER reply: number, subject, from, date, message-id, bytes, xref (scanOverviewText).
fn scan_overview_text(text: &str, mut line: impl FnMut(Option<i64>, &str, &str, &str, &str, &str, &str)) {
    let b = text.as_bytes();
    let mut start = 0;
    while start < b.len() {
        let end = memchr::memchr(b'\n', &b[start..]).map(|i| i + start).unwrap_or(b.len());
        let stop = if end > start && b[end - 1] == b'\r' { end - 1 } else { end };
        if stop > start {
            let l = &text[start..stop];
            let lb = l.as_bytes();
            let mut tabs = [usize::MAX; 9];
            let mut n = 0;
            for i in memchr::memchr_iter(b'\t', lb) {
                tabs[n] = i;
                n += 1;
                if n == 9 {
                    break;
                }
            }
            if n >= 5 {
                let field = |a: usize, b: usize| -> &str {
                    if a == usize::MAX {
                        ""
                    } else {
                        &l[a + 1..if b == usize::MAX { l.len() } else { b }]
                    }
                };
                let mut xref = if n >= 8 { field(tabs[7], tabs[8]) } else { "" };
                if let Some(rest) = xref.strip_prefix("Xref:") {
                    xref = rest.strip_prefix(' ').unwrap_or(rest);
                }
                line(js::parse_int(&l[..tabs[0]]), &l[tabs[0] + 1..tabs[1]], field(tabs[1], tabs[2]), field(tabs[2], tabs[3]), field(tabs[3], tabs[4]), field(tabs[5], tabs[6]), xref);
            }
        }
        start = end + 1;
    }
}

struct OpenFile {
    parts: Vec<Seg>,
    seen: FxHashSet<u32>,
    total: u32,
    touched: u64,
}

#[derive(Default)]
struct Aggregator {
    open: FxHashMap<String, OpenFile>,
    batch: u64,
}

impl Aggregator {
    fn add(&mut self, segments: Vec<Seg>) -> Vec<Vec<Seg>> {
        self.batch += 1;
        let mut touched: Vec<String> = Vec::new();
        let mut touched_set: FxHashSet<String> = FxHashSet::default();
        for s in segments {
            let key = &s.file.file_key;
            let file = match self.open.get_mut(key.as_str()) {
                Some(f) => f,
                None => self.open.entry(key.clone()).or_insert(OpenFile { parts: Vec::new(), seen: FxHashSet::default(), total: s.file.total_parts, touched: self.batch }),
            };
            if !touched_set.contains(key.as_str()) {
                touched_set.insert(key.clone());
                touched.push(key.clone());
            }
            file.touched = self.batch;
            if file.seen.insert(s.number) {
                file.parts.push(s);
            }
        }
        let mut complete = Vec::new();
        for key in touched {
            let f = self.open.get(&key).unwrap();
            if (f.parts.len() as u32) < f.total {
                continue;
            }
            let mut f = self.open.remove(&key).unwrap();
            f.parts.sort_by_key(|s| s.number);
            complete.push(f.parts);
        }
        complete
    }

    /// Files idle for `idle` batches (all when 0): their segments, file by file in arrival order.
    fn evict(&mut self, idle: u64) -> Vec<Vec<Seg>> {
        let batch = self.batch;
        let gone: Vec<String> = self.open.iter().filter(|(_, f)| !(idle > 0 && batch - f.touched < idle)).map(|(k, _)| k.clone()).collect();
        gone.into_iter().map(|k| self.open.remove(&k).unwrap().parts).collect()
    }
}

pub struct RangeIndexer {
    group: String,
    start: i64,
    end: i64,
    priorities: FxHashMap<String, u32>,
    idle: u64,
    chunk_articles: i64,
    scanner: Scanner,
    aggregator: Aggregator,
    batch: Batch,
    counts: FxHashMap<(&'static str, String), (u64, u64)>,
    count_order: Vec<(&'static str, String)>,
    completed: FxHashMap<String, Vec<Vec<Seg>>>,
    completed_order: Vec<String>,
}

impl RangeIndexer {
    pub fn new(group: &str, start: i64, end: i64, priorities: FxHashMap<String, u32>, idle: u64) -> Self {
        RangeIndexer {
            group: group.to_string(),
            start,
            end,
            priorities,
            idle,
            chunk_articles: 0,
            scanner: Scanner::default(),
            aggregator: Aggregator::default(),
            batch: Batch { from: start, next: start, ..Default::default() },
            counts: FxHashMap::default(),
            count_order: Vec::new(),
            completed: FxHashMap::default(),
            completed_order: Vec::new(),
        }
    }

    pub fn add_text(&mut self, text: &str, next: i64) {
        let mut segments: Vec<Seg> = Vec::new();
        let mut crossposts: Vec<String> = Vec::new();
        let mut obfuscated = 0u64;
        let mut headers = 0u64;
        {
            let group = self.group.clone();
            let scanner = &mut self.scanner;
            let priorities = &self.priorities;
            scan_overview_text(text, |n, subject, from, date, message_id, bytes, xref| {
                let Some(n) = n else { return };
                if message_id.is_empty() {
                    return;
                }
                headers += 1;
                let owner = if xref.is_empty() { group.as_str() } else { scanner.owner(xref, &group, priorities) };
                if owner != group {
                    crossposts.push(owner.to_string());
                    return;
                }
                match scanner.segment(n, subject, from, date, message_id, bytes, xref, &group) {
                    Some(s) => segments.push(s),
                    None => obfuscated += 1,
                }
            });
        }
        self.batch.headers += headers;
        for owner in crossposts {
            self.count("crosspost", owner, 1, 0);
        }
        if obfuscated > 0 {
            self.count("obfuscated", String::new(), obfuscated, 0);
        }
        self.end_chunk(segments, next);
    }

    fn end_chunk(&mut self, segments: Vec<Seg>, next: i64) {
        let t = std::time::Instant::now();
        let done = self.aggregator.add(segments);
        unsafe { PROF[3] += t.elapsed().as_nanos() as u64; }
        self.files(done);
        if self.chunk_articles == 0 {
            self.chunk_articles = next - self.batch.from;
        }
        let t = std::time::Instant::now();
        let gone = self.aggregator.evict(self.idle);
        self.cut_off(gone, false);
        unsafe { PROF[4] += t.elapsed().as_nanos() as u64; }
        self.batch.next = next;
    }

    pub fn finish(&mut self) {
        let gone = self.aggregator.evict(0);
        self.cut_off(gone, true);
        self.batch.next = self.end + 1;
    }

    pub fn take(&mut self) -> Batch {
        let t = std::time::Instant::now();
        self.judge_releases();
        unsafe { PROF[5] += t.elapsed().as_nanos() as u64; }
        self.scanner.files.clear();
        let next = self.batch.next;
        let mut out = std::mem::replace(&mut self.batch, Batch { from: next, next, ..Default::default() });
        for id in self.count_order.drain(..) {
            let (articles, files) = self.counts.remove(&id).unwrap();
            out.counts.push((id.0, id.1, articles, files));
        }
        out
    }

    fn files(&mut self, files: Vec<Vec<Seg>>) {
        for file in files {
            let key = release_key_of(&file[0]);
            match self.completed.get_mut(&key) {
                Some(list) => list.push(file),
                None => {
                    self.completed_order.push(key.clone());
                    self.completed.insert(key, vec![file]);
                }
            }
        }
    }

    fn judge_releases(&mut self) {
        let order = std::mem::take(&mut self.completed_order);
        let mut by_release = std::mem::take(&mut self.completed);
        for key in order {
            let list = by_release.remove(&key).unwrap();
            let judged: Vec<JudgedFile> = list
                .iter()
                .map(|f| JudgedFile { name: &f[0].file.name, total_parts: f[0].file.total_parts, groups: &f[0].file.groups, first_message_id: &f[0].message_id })
                .collect();
            match names::judge(&key, &judged) {
                Verdict::Keep => {
                    let candidates = candidates_of(&key, &list);
                    self.batch.releases.push(chunk_of(&key, &list));
                    self.batch.candidates.extend(candidates);
                }
                Verdict::Drop { rule, key: rule_key, sample } => {
                    let segments: usize = list.iter().map(|f| f.len()).sum();
                    self.count(rule, rule_key.clone(), segments as u64, list.len() as u64);
                    if sample {
                        self.batch.samples.push((chunk_of(&key, &list), rule, rule_key));
                    }
                }
            }
        }
    }

    fn cut_off(&mut self, files: Vec<Vec<Seg>>, at_end: bool) {
        let window = self.idle as i64 * self.chunk_articles.max(1);
        for list in files {
            if list.len() == 1 && is_random_file(&list[0].file.name) {
                let tool = names::upload_tool(&list[0].message_id);
                self.count("per-article", tool.to_string(), 1, 0);
                continue;
            }
            let min_part = list.iter().map(|s| s.number).min().unwrap();
            let max_part = list.iter().map(|s| s.number).max().unwrap();
            let first_article = list.iter().map(|s| s.article).min().unwrap();
            let cut_at_start = min_part > 1 && first_article < self.start + window;
            let cut_at_end = at_end && max_part < list[0].file.total_parts;
            if !cut_at_start && !cut_at_end {
                self.count("incomplete", String::new(), list.len() as u64, 1);
                continue;
            }
            let f = &list[0].file;
            let mut fragment = Fragment { file_key: f.file_key.clone(), name: f.name.clone(), total_parts: f.total_parts, posted_at: None, segments: Vec::with_capacity(list.len()) };
            for s in &list {
                fragment.segments.push((s.number, s.message_id.clone(), s.bytes, s.article));
                if let Some(t) = s.posted_at {
                    if fragment.posted_at.is_none_or(|p| t < p) {
                        fragment.posted_at = Some(t);
                    }
                }
            }
            self.batch.fragments.push(fragment);
        }
    }

    fn count(&mut self, rule: &'static str, key: String, articles: u64, files: u64) {
        match self.counts.get_mut(&(rule, key.clone())) {
            Some(c) => {
                c.0 += articles;
                c.1 += files;
            }
            None => {
                self.count_order.push((rule, key.clone()));
                self.counts.insert((rule, key), (articles, files));
            }
        }
    }
}

fn release_key_of(s: &Seg) -> String {
    // A named file always has a release (releaseIdentity isn't null unless weak).
    s.file.release.key.clone()
}

fn posted_of(segments: &[Seg]) -> Option<i64> {
    segments.iter().filter_map(|s| s.posted_at).min()
}

fn chunk_of(key: &str, files: &[Vec<Seg>]) -> ReleaseChunk {
    let first = &files[0][0].file;
    let chunk_files: Vec<ChunkFile> = files
        .iter()
        .map(|f| {
            let m = &f[0].file;
            ChunkFile { name: &m.name, poster: &m.poster, subject: &m.subject, posted_at: posted_of(f), groups: &m.groups, segments: f.iter().map(|s| (s.number, s.message_id.as_str(), s.bytes)).collect() }
        })
        .collect();
    let mut segments = 0;
    let mut bytes = 0;
    let mut first_article = i64::MAX;
    let mut last_article = 0;
    let mut posted_at: Option<i64> = None;
    for f in files {
        for s in f {
            segments += 1;
            bytes += s.bytes;
            first_article = first_article.min(s.article);
            last_article = last_article.max(s.article);
        }
        if let Some(t) = posted_of(f) {
            if posted_at.is_none_or(|p| t < p) {
                posted_at = Some(t);
            }
        }
    }
    ReleaseChunk {
        key: key.to_string(),
        name: first.release.name.clone(),
        source: first.release.source,
        file_total: files.iter().find_map(|f| f[0].file.release.file_total_truthy()),
        poster: first.poster.clone(),
        groups: first.groups.clone(),
        posted_at,
        files: files.len(),
        segments,
        bytes,
        first_article,
        last_article,
        readable: chunk_files.iter().any(|f| !is_random_file(f.name)),
        has_par2: chunk_files.iter().any(|f| names::is_par2(f.name)),
        payload: encode_chunk(&chunk_files),
    }
}

fn candidates_of(key: &str, files: &[Vec<Seg>]) -> Vec<Candidate> {
    let mut out = Vec::new();
    let mut tail: Option<(i64, &Vec<Seg>)> = None;
    let c = |kind, s: &Seg| Candidate { release_key: key.to_string(), kind, message_id: s.message_id.clone(), file: s.file.name.clone() };
    for f in files {
        let name = &f[0].file.name;
        if let Some(volume) = names::seven_zip_volume(name).filter(|&v| v > 1) {
            if tail.is_none_or(|(v, _)| volume > v) {
                tail = Some((volume, f));
            }
            continue;
        }
        if names::is_archive_start(name) {
            out.push(c("archive", &f[0]));
        } else if names::is_par2(name) && !names::is_later_volume(name) {
            out.push(c("par2", &f[0]));
        }
    }
    if let Some((_, f)) = tail {
        out.push(Candidate { release_key: key.to_string(), kind: "7z-tail", message_id: f.last().unwrap().message_id.clone(), file: f[0].file.name.clone() });
    }
    if out.is_empty() && files.len() == 1 && is_random_file(&files[0][0].file.name) {
        out.push(c("lone", &files[0][0]));
    }
    out.truncate(4);
    out
}