aprender-contracts 0.70.3

//! PVL-001 EV-7a (#4200): the STATEMENT of every contract-bound theorem, pinned apart from its proof.
//!
//! `pv challenge gen` writes `<lean>/Challenge/<contract-stem>.lean` for every contract with at least one bound
//! root ([`super::bind`]'s set): each root `X` is restated as `PvlChallenge.X` with its proof replaced by `sorry`.
//! `pv challenge check` renders the same set in memory and compares it with the directory; any extra, missing or
//! differing file is a failure. A comparator (EV-7b) then checks that `X` proves exactly `PvlChallenge.X`, so a
//! proof cannot quietly weaken the statement a contract names: that edit now shows as a diff to a committed file.
//!
//! The YAML carries no Lean types, so the statement is lifted from the solution: the text after the theorem's
//! name up to its first `:=` outside brackets, from blanked source (comments removed), whitespace collapsed. It
//! is elaborated in the scope it was written in. Each root gets its own `section … end` that reopens, verbatim,
//! every `namespace`/`section` still open at the declaration with the `open`/`variable`/`universe`/`set_option`
//! commands issued in it, and declares `theorem _root_.PvlChallenge.X`. So names resolve as in the source: the
//! current namespace keeps its priority (alibi's `dist` is `Alibi.dist`, not Mathlib's `Dist.dist`, which an
//! `open` of the namespace could not give), and a statement naming a sibling theorem `u` reads the solution's `u`,
//! not its challenge copy (both measured on lean v4.29.0-rc4, core only).
//!
//! Known limit: `local notation`, `local instance` and `attribute [local …]` are not copied; a statement that
//! needs one fails to elaborate, which is loud, never a silent pass.
//!
//! `Challenge/` sits outside the lakefile's `ProvableContracts` library and outside [`super::Tree::load`]'s scan,
//! so its `sorry`s are never part of the build or of the escape scan.

use std::collections::{BTreeMap, BTreeSet};
use std::path::Path;

use super::{bind, is_theorem, lex, Root, Tree};

pub const CHALLENGE_DIR: &str = "Challenge";
pub const CHALLENGE_NS: &str = "PvlChallenge";

/// The name the challenge copy of solution theorem `fqn` is declared under.
#[must_use]
pub fn challenge_decl(fqn: &str) -> String {
    format!("{CHALLENGE_NS}.{fqn}")
}

/// Everything `gen` writes and `check` compares.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Rendered {
    /// `Challenge/<stem>.lean` (relative to the lean dir) → its content.
    pub files: BTreeMap<String, String>,
    /// `(contract stem, fqn, why)`: a bound root whose statement could not be lifted. Each one is a failure: a
    /// root with no challenge is a statement nothing pins.
    pub unrestated: Vec<(String, String, String)>,
}

/// Render every challenge for the contracts under `contract_dir`. `Err` (a decline) when the tree cannot be loaded.
pub fn render(lean_dir: &Path, contract_dir: &Path) -> Result<Rendered, String> {
    let tree = Tree::load(lean_dir)?;
    let binding = bind(&tree, contract_dir);
    // Only roots in the root's import cone, as Axioms.lean: an ORPHANED-ROOT's module has no .olean, so its
    // challenge cannot elaborate and the comparator would withhold every row of that file (#4502).
    let cone = tree.cone();
    let mut out = Rendered::default();
    let mut sources: BTreeMap<&str, String> = BTreeMap::new();
    // Two obligations may bind one theorem; a second restatement is a duplicate declaration and the file fails to
    // elaborate (silu-kernel-v1: `sigmoid_lt_exp`, #4244). So may two CONTRACTS (attention-kernel-v1 and
    // attention-scaling-v1: `attention_weight_pos`, #4502): the root imports every Challenge file, so the second
    // file re-declares `PvlChallenge.X` -- 132 DUPLICATE rows. One challenge per root across the whole directory,
    // in the first contract (stem order) that binds it; the comparator matches by name, not by file.
    let mut seen = BTreeSet::new();
    for (stem, roots) in &binding.by_contract {
        let mut blocks = Vec::new();
        let mut imports = BTreeSet::new();
        for root in roots
            .iter()
            .filter(|r| cone.contains(&r.module))
            .filter(|r| seen.insert(r.fqn.as_str()))
        {
            match restate(&tree, root, &mut sources) {
                Ok(block) => {
                    imports.insert(root.module.clone());
                    blocks.push(block);
                }
                Err(why) => out.unrestated.push((stem.clone(), root.fqn.clone(), why)),
            }
        }
        if blocks.is_empty() {
            continue;
        }
        out.files.insert(
            format!("{CHALLENGE_DIR}/{stem}.lean"),
            file_text(stem, &imports, &blocks),
        );
    }
    Ok(out)
}

fn file_text(stem: &str, imports: &BTreeSet<String>, blocks: &[String]) -> String {
    let mut s = String::new();
    for m in imports {
        s.push_str(&format!("import {m}\n"));
    }
    s.push_str(&format!(
        "\n-- GENERATED by `pv challenge gen` (PVL-001 EV-7a) for contract `{stem}`. Do not edit: \
         `pv challenge check` regenerates it and fails on any difference.\n\
         -- Each `{CHALLENGE_NS}.X` restates the contract-bound theorem `X` with its proof replaced by `sorry`.\n"
    ));
    for b in blocks {
        s.push('\n');
        s.push_str(b);
    }
    s
}

/// One `section … end` block restating `root` inside the scope frames still open at its declaration.
fn restate<'t>(
    tree: &'t Tree,
    root: &Root,
    sources: &mut BTreeMap<&'t str, String>,
) -> Result<String, String> {
    let file = tree
        .files
        .iter()
        .find(|f| f.module == root.module)
        .ok_or_else(|| format!("module {} is not in the tree", root.module))?;
    let decl = file
        .decls
        .iter()
        .find(|d| is_theorem(d) && d.fqn == root.fqn)
        .ok_or_else(|| format!("no theorem {} in {}", root.fqn, file.rel))?;
    if !sources.contains_key(file.rel.as_str()) {
        let src = std::fs::read_to_string(tree.dir.join(&file.rel))
            .map_err(|e| format!("{}: {e}", file.rel))?;
        sources.insert(file.rel.as_str(), lex::blank(&src));
    }
    let blanked = &sources[file.rel.as_str()];
    let (universes, stmt) = statement(blanked, decl.line, &decl.keyword).ok_or_else(|| {
        format!(
            "{}:{}: no `:=` closes the statement of {}",
            file.rel, decl.line, root.fqn
        )
    })?;
    let scope = preamble(blanked, decl.line);
    let mut b = String::from("section\n");
    for cmd in &scope.top {
        b.push_str(cmd);
        b.push('\n');
    }
    for f in &scope.frames {
        b.push_str(&f.opener);
        b.push('\n');
        for cmd in &f.cmds {
            b.push_str(cmd);
            b.push('\n');
        }
    }
    b.push_str(&format!(
        "theorem _root_.{}{universes} {stmt} := sorry\n",
        challenge_decl(&root.fqn)
    ));
    for f in scope.frames.iter().rev() {
        b.push_str(&f.closer);
        b.push('\n');
    }
    b.push_str("end\n");
    Ok(b)
}

const OPENERS: &[char] = &['(', '[', '{', '⟨', '⦃'];
const CLOSERS: &[char] = &[')', ']', '}', '⟩', '⦄'];

/// The statement of the theorem declared at 1-based `line`, from blanked source: `(universe params, text)`, the
/// text running from after the name to the first `:=` outside brackets, whitespace collapsed. `None` when the
/// declaration is not found, or a new command starts before any `:=`. An equation-style theorem's statement ends
/// at its first `| … =>` arm.
fn statement(blanked: &str, line: usize, keyword: &str) -> Option<(String, String)> {
    let start_of_line: usize = blanked
        .split_inclusive('\n')
        .take(line.checked_sub(1)?)
        .map(str::len)
        .sum();
    let rest = &blanked[start_of_line..];
    let first_line = rest.split('\n').next()?;
    let kw = find_word(first_line, keyword)?;
    let after_kw = &rest[kw + keyword.len()..];
    let name_at = after_kw.find(|c: char| !c.is_whitespace())?;
    let named = &after_kw[name_at..];
    let end = name_len(named);
    if end == 0 {
        return None;
    }
    let after_name = &named[end..];
    let (universes, body) = match after_name.strip_prefix(".{") {
        Some(u) => {
            let close = u.find('}')?;
            (format!(".{{{}}}", &u[..close]), &u[close + 1..])
        }
        None => (String::new(), after_name),
    };
    statement_text(body).map(|text| (universes, text))
}

/// Byte length of the name as written (`foo`, `Foo.bar`, `«x»`-free); a `.` continues it only when an identifier
/// char follows, so `foo.{u}` stops before `.{`.
fn name_len(named: &str) -> usize {
    let mut end = 0;
    let mut it = named.char_indices().peekable();
    while let Some((i, c)) = it.next() {
        let next = it.peek().map(|&(_, n)| n);
        if c == '.' && !next.is_some_and(|n| lex::is_ident_char(n) && n != '.') {
            break;
        }
        if !lex::is_ident_char(c) {
            break;
        }
        end = i + c.len_utf8();
    }
    end
}

/// The statement text of `body` (everything after the name and universes): up to the first `:=` outside brackets,
/// or the first equation arm, whitespace collapsed. `None` when empty or a new command starts first.
fn statement_text(body: &str) -> Option<String> {
    let mut depth = 0usize;
    let mut prev = ' ';
    let chars: Vec<(usize, char)> = body.char_indices().collect();
    for (k, &(_, c)) in chars.iter().enumerate() {
        if OPENERS.contains(&c) {
            depth += 1;
        } else if CLOSERS.contains(&c) {
            depth = depth.saturating_sub(1);
        } else if depth == 0 {
            if let Some(end) = statement_end(&chars, k, body, prev) {
                let text = collapse(&body[..end?]);
                return (!text.is_empty()).then_some(text);
            }
        }
        prev = c;
    }
    None
}

/// At bracket depth 0, char `k` of `chars` ends the statement: `Some(Some(byte end))` at `:=` or an equation arm,
/// `Some(None)` at a new command in column 0 before any `:=`, `None` to keep scanning.
fn statement_end(
    chars: &[(usize, char)],
    k: usize,
    body: &str,
    prev: char,
) -> Option<Option<usize>> {
    let (i, c) = chars[k];
    if c == '=' && prev == ':' {
        return Some(Some(i - 1));
    }
    if prev == '\n' && !c.is_whitespace() {
        return Some(None); // a new command at column 0 before any `:=`
    }
    // Equation style: the type before the first `| pat => …` arm is the whole statement.
    if c == '|' && line_is_blank_before(chars, k) && is_arm(&body[i..]) {
        return Some(Some(i));
    }
    None
}

/// A line opening with `|` is an equation arm when it carries a `=>` outside brackets; otherwise the `|` opens an
/// absolute value (`|f x - x| ≤ e`).
fn is_arm(line_from_bar: &str) -> bool {
    let line = line_from_bar.split('\n').next().unwrap_or("");
    let mut depth = 0usize;
    let mut prev = ' ';
    for c in line.chars() {
        if OPENERS.contains(&c) {
            depth += 1;
        } else if CLOSERS.contains(&c) {
            depth = depth.saturating_sub(1);
        } else if depth == 0 && c == '>' && prev == '=' {
            return true;
        }
        prev = c;
    }
    false
}

fn line_is_blank_before(chars: &[(usize, char)], k: usize) -> bool {
    chars[..k]
        .iter()
        .rev()
        .take_while(|(_, c)| *c != '\n')
        .all(|(_, c)| c.is_whitespace())
}

/// The byte offset of `word` in `line` as a whole identifier.
fn find_word(line: &str, word: &str) -> Option<usize> {
    line.match_indices(word).map(|(at, _)| at).find(|&at| {
        let before = line[..at].chars().next_back();
        let after = line[at + word.len()..].chars().next();
        !before.is_some_and(lex::is_ident_char) && !after.is_some_and(lex::is_ident_char)
    })
}

fn collapse(s: &str) -> String {
    s.split_whitespace().collect::<Vec<_>>().join(" ")
}

/// One `namespace`/`section` still open at the declaration: the line that opened it, the `end` that closes it, and
/// the scoped commands issued inside it.
#[derive(Debug, Clone, PartialEq, Eq)]
struct Frame {
    opener: String,
    closer: String,
    cmds: Vec<String>,
}

/// The scope at a declaration: file-level commands, then every frame still open, outermost first.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct Scope {
    top: Vec<String>,
    frames: Vec<Frame>,
}

/// The scope still active at 1-based `line`: `open`/`variable`/`universe`/`set_option` commands (continuation lines,
/// indented, joined; `… in` forms skipped) in the frame they were issued in. `namespace`/`section`/
/// `noncomputable section` open a frame; `end` closes one.
fn preamble(blanked: &str, line: usize) -> Scope {
    let mut scope = Scope::default();
    // An iterator, not an index: no mutated cursor can stop the scan advancing.
    let mut rest = blanked.lines().take(line.saturating_sub(1));
    while let Some(first) = rest.next() {
        let t = first.trim();
        let words: Vec<&str> = t.split_whitespace().collect();
        let head = words.first().copied().unwrap_or("");
        // Indented, non-blank lines continue the command; they are consumed only when the command is kept below.
        let cont: Vec<&str> = rest
            .clone()
            .take_while(|l| l.starts_with(char::is_whitespace) && !l.trim().is_empty())
            .map(str::trim)
            .collect();
        let mut cmd = t.to_string();
        for c in &cont {
            cmd.push(' ');
            cmd.push_str(c);
        }
        let opened = match words.as_slice() {
            ["namespace", name, ..] => Some(format!("end {name}")),
            ["section", name, ..] | ["noncomputable", "section", name, ..] => {
                Some(format!("end {name}"))
            }
            ["section"] | ["noncomputable", "section"] => Some("end".to_string()),
            _ => None,
        };
        if let Some(closer) = opened {
            scope.frames.push(Frame {
                opener: collapse(t),
                closer,
                cmds: Vec::new(),
            });
        } else if head == "end" {
            scope.frames.pop();
        } else if matches!(head, "open" | "variable" | "universe" | "set_option") {
            if !cmd.ends_with(" in") {
                let cmds = match scope.frames.last_mut() {
                    Some(f) => &mut f.cmds,
                    None => &mut scope.top,
                };
                cmds.push(collapse(&cmd));
            }
            rest.by_ref().take(cont.len()).for_each(drop);
        }
    }
    scope
}

/// Write the rendered set under `lean_dir`, removing any other file in `Challenge/`. Returns the paths written.
pub fn write(lean_dir: &Path, r: &Rendered) -> std::io::Result<Vec<String>> {
    let dir = lean_dir.join(CHALLENGE_DIR);
    std::fs::create_dir_all(&dir)?;
    for rel in on_disk(lean_dir) {
        if !r.files.contains_key(&rel) {
            std::fs::remove_file(lean_dir.join(&rel))?;
        }
    }
    let mut written = Vec::new();
    for (rel, text) in &r.files {
        let p = lean_dir.join(rel);
        if std::fs::read_to_string(&p).ok().as_deref() != Some(text.as_str()) {
            std::fs::write(&p, text)?;
            written.push(rel.clone());
        }
    }
    Ok(written)
}

/// Every regular file under `<lean_dir>/Challenge`, relative to `lean_dir`, sorted.
fn on_disk(lean_dir: &Path) -> Vec<String> {
    fn walk(dir: &Path, base: &Path, out: &mut Vec<String>) {
        let Ok(rd) = std::fs::read_dir(dir) else {
            return;
        };
        for e in rd.flatten() {
            let p = e.path();
            if p.is_dir() {
                walk(&p, base, out);
            } else if let Ok(rel) = p.strip_prefix(base) {
                out.push(rel.to_string_lossy().replace('\\', "/"));
            }
        }
    }
    let mut out = Vec::new();
    walk(&lean_dir.join(CHALLENGE_DIR), lean_dir, &mut out);
    out.sort();
    out
}

/// The differences between the rendered set and `<lean_dir>/Challenge`, one line each (empty = fresh).
#[must_use]
pub fn diff(lean_dir: &Path, r: &Rendered) -> Vec<String> {
    let disk = on_disk(lean_dir);
    let mut out = Vec::new();
    for rel in &disk {
        if !r.files.contains_key(rel) {
            out.push(format!("extra    {rel}: not generated from any contract"));
        }
    }
    for (rel, want) in &r.files {
        if !disk.contains(rel) {
            out.push(format!("missing  {rel}"));
            continue;
        }
        let got = std::fs::read_to_string(lean_dir.join(rel)).unwrap_or_default();
        if &got != want {
            let at = want
                .lines()
                .zip(got.lines())
                .position(|(a, b)| a != b)
                .unwrap_or_else(|| want.lines().count().min(got.lines().count()));
            out.push(format!(
                "differs  {rel}:{}: want {:?}, got {:?}",
                at + 1,
                want.lines().nth(at).unwrap_or("<eof>"),
                got.lines().nth(at).unwrap_or("<eof>")
            ));
        }
    }
    out
}

#[cfg(test)]
#[path = "challenge_tests.rs"]
mod tests;