open-standards/proof/proof-rs/crates/proof-solvers/src/lib.rs

//! Solvers for the proof tier ladder (SPEC §4).
//!
//! The two bottom rungs are **real** and report energy with `ModelBased`
//! provenance (measured wall time × a power model):
//!
//! - [`Propositional`] (T0) → parses a boolean goal and decides validity by
//!   truth-table, in process.
//! - [`Z3Smt`] (T1) → shells out to `z3` on an SMT-LIB2 goal (assert the
//!   negation; `unsat` ⇒ valid). [`Z3Server`] keeps `z3` resident across goals.
//!
//! The model rungs are honest stand-ins that gate on `GoalKind` until the real
//! backends are wired behind the identical [`Solver`] trait:
//!
//! - [`SmallModel`] (T2) → a distilled in-domain prover.
//! - [`FrontierModel`] (T3) → a frontier reasoning model; the last resort.
//! - [`Automation`] / [`Smt`] → the original `GoalKind`-gated stand-ins, kept for
//!   [`reference_ladder`].
//!
//! A solver returns `Proved` only when a verifier would accept the artifact —
//! correctness is never traded for energy. The real rungs run the tool and
//! report what actually happened; the stand-ins model the rung's reach.

pub mod prop;

use std::io::{BufRead, BufReader, Write};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::Mutex;
use std::time::Instant;

use proof_schema::{Attempt, ContentId, Energy, GoalKind, ProofGoal, Solver, Tier, Verdict};

/// Convert elapsed nanoseconds to microjoules under a documented single-core
/// power model (~15 W active). The reading is `ModelBased`: a hardware interface
/// (RAPL/NVML) would make it `HwShunt`; from wall time alone it is honestly
/// modeled.
fn ns_to_uj(ns: u128) -> u64 {
    // 15 W = 15 µJ/µs = 0.015 µJ/ns  →  µJ = ns × 15 / 1000.
    (ns.saturating_mul(15) / 1000) as u64
}

fn proved(goal: &ProofGoal) -> Verdict {
    Verdict::Proved {
        artifact: ContentId::of_bytes(goal.statement.as_bytes()),
    }
}

/// T0 — proof-language automation stand-in. Discharges goals that close by
/// normalization or linear arithmetic.
pub struct Automation;

impl Solver for Automation {
    fn tier(&self) -> Tier {
        Tier::Automation
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        let verdict = match goal.kind {
            GoalKind::Trivial | GoalKind::Arithmetic => proved(goal),
            _ => Verdict::Unproved,
        };
        Attempt {
            tier: Tier::Automation,
            verdict,
            energy: Energy::estimate_tier(Tier::Automation),
        }
    }
}

/// T1 — SMT / decision-procedure stand-in. Discharges decidable / arithmetic goals.
pub struct Smt;

impl Solver for Smt {
    fn tier(&self) -> Tier {
        Tier::Smt
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        let verdict = match goal.kind {
            GoalKind::Decidable | GoalKind::Arithmetic => proved(goal),
            _ => Verdict::Unproved,
        };
        Attempt {
            tier: Tier::Smt,
            verdict,
            energy: Energy::estimate_tier(Tier::Smt),
        }
    }
}

/// T2 — a small distilled prover for goals that genuinely need learned search.
pub struct SmallModel;

impl Solver for SmallModel {
    fn tier(&self) -> Tier {
        Tier::SmallModel
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        let verdict = match goal.kind {
            GoalKind::General => proved(goal),
            _ => Verdict::Unproved,
        };
        Attempt {
            tier: Tier::SmallModel,
            verdict,
            energy: Energy::estimate_tier(Tier::SmallModel),
        }
    }
}

/// T3 — the frontier reasoning model, the last resort.
pub struct FrontierModel;

impl Solver for FrontierModel {
    fn tier(&self) -> Tier {
        Tier::BigModel
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        Attempt {
            tier: Tier::BigModel,
            verdict: proved(goal),
            energy: Energy::estimate_tier(Tier::BigModel),
        }
    }
}

/// T0 — a **real** propositional decider (no stand-in). Parses the goal and
/// decides validity by truth-table, with energy measured from execution time. A
/// goal beyond propositional logic is refused so the router escalates.
pub struct Propositional;

impl Solver for Propositional {
    fn tier(&self) -> Tier {
        Tier::Automation
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        let start = Instant::now();
        let verdict = match prop::parse(&goal.statement).and_then(|e| prop::is_valid(&e)) {
            Some(true) => Verdict::Proved {
                artifact: ContentId::of_bytes(goal.statement.as_bytes()),
            },
            _ => Verdict::Unproved,
        };
        let uj = ns_to_uj(start.elapsed().as_nanos()).max(1);
        Attempt {
            tier: Tier::Automation,
            verdict,
            energy: Energy::modeled(uj),
        }
    }
}

/// Run `z3` on an SMT-LIB2 script (piped over stdin) and return the first
/// `check-sat` verdict. `None` if `z3` is absent or errored — the router reads
/// that as "escalate", never as proof.
fn run_z3(script: &str) -> Option<String> {
    let mut child = Command::new("z3")
        .args(["-in", "-smt2"])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
        .ok()?;
    let full = if script.contains("check-sat") {
        script.to_string()
    } else {
        format!("{script}\n(check-sat)\n")
    };
    {
        let mut stdin = child.stdin.take()?;
        stdin.write_all(full.as_bytes()).ok()?;
    }
    let out = child.wait_with_output().ok()?;
    String::from_utf8_lossy(&out.stdout)
        .lines()
        .map(str::trim)
        .find(|l| matches!(*l, "sat" | "unsat" | "unknown"))
        .map(str::to_string)
}

/// Fetch the counter-model for a refuted goal: re-run `z3` with `(get-model)`
/// appended so the witness assignment is captured for the `Refuted` verdict.
/// Best-effort — returns an empty string if `z3` declines a model (some goals
/// are `sat` without a printable model). Only called on the refutation path,
/// which is terminal, so the extra spawn costs nothing on the proved path.
fn z3_counter_model(script: &str) -> String {
    let full = if script.contains("check-sat") {
        format!("{script}\n(get-model)\n")
    } else {
        format!("{script}\n(check-sat)\n(get-model)\n")
    };
    let Ok(mut child) = Command::new("z3")
        .args(["-in", "-smt2"])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
    else {
        return String::new();
    };
    if let Some(mut stdin) = child.stdin.take() {
        let _ = stdin.write_all(full.as_bytes());
    }
    let Ok(out) = child.wait_with_output() else {
        return String::new();
    };
    let text = String::from_utf8_lossy(&out.stdout);
    // Keep the model block (everything from the first "sat" line onward),
    // trimmed; drop the leading bare "sat".
    let mut lines = text.lines().map(str::trim_end).skip_while(|l| l.trim() != "sat");
    lines.next(); // the "sat" line itself
    lines.collect::<Vec<_>>().join("\n").trim().to_string()
}

/// T1 — a **real** SMT backend that shells out to `z3`. The goal's statement is
/// an SMT-LIB2 script asserting the *negation* of the claim; `unsat` ⇒ the
/// claim is valid — a genuine decision by a trusted procedure. `sat` is a real
/// counter-model (refuse); `unknown`/missing `z3` escalates. Energy is measured
/// from the solver's wall time.
pub struct Z3Smt;

impl Solver for Z3Smt {
    fn tier(&self) -> Tier {
        Tier::Smt
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        let start = Instant::now();
        let verdict = match run_z3(&goal.statement).as_deref() {
            Some("unsat") => Verdict::Proved {
                artifact: ContentId::of_bytes(goal.statement.as_bytes()),
            },
            // `sat` = the negated claim is satisfiable = the claim is FALSE. A
            // decision, not a failure to decide — surface the witness and let
            // the router halt rather than escalate a disproof.
            Some("sat") => Verdict::Refuted {
                counter_model: z3_counter_model(&goal.statement),
            },
            _ => Verdict::Unproved,
        };
        let uj = ns_to_uj(start.elapsed().as_nanos()).max(1);
        Attempt {
            tier: Tier::Smt,
            verdict,
            energy: Energy::modeled(uj),
        }
    }
}

/// A long-lived `z3 -in` process that decides many goals across one startup —
/// the fix for the per-spawn startup tax. Each goal is isolated with `(reset)`
/// and delimited by an `(echo)` marker.
struct Z3Session {
    child: Child,
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
}

impl Z3Session {
    fn spawn() -> Option<Self> {
        let mut child = Command::new("z3")
            .args(["-in", "-smt2"])
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .ok()?;
        let stdin = child.stdin.take()?;
        let stdout = BufReader::new(child.stdout.take()?);
        Some(Self {
            child,
            stdin,
            stdout,
        })
    }

    fn decide(&mut self, script: &str) -> Option<String> {
        write!(
            self.stdin,
            "(reset)\n{script}\n(check-sat)\n(echo \"@@END@@\")\n"
        )
        .ok()?;
        self.stdin.flush().ok()?;
        let mut verdict = None;
        let mut line = String::new();
        loop {
            line.clear();
            if self.stdout.read_line(&mut line).ok()? == 0 {
                return None;
            }
            match line.trim() {
                "@@END@@" => break,
                v @ ("sat" | "unsat" | "unknown") => verdict = Some(v.to_string()),
                _ => {}
            }
        }
        verdict
    }
}

impl Drop for Z3Session {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

/// A balanced, paren-delimited SMT-LIB script is safe to feed the resident
/// session; anything else is refused without touching it, so a malformed goal
/// can never desync the long-lived process.
fn looks_like_smt(s: &str) -> bool {
    let t = s.trim();
    if !t.starts_with('(') {
        return false;
    }
    let mut depth: i32 = 0;
    for c in t.chars() {
        match c {
            '(' => depth += 1,
            ')' => {
                depth -= 1;
                if depth < 0 {
                    return false;
                }
            }
            _ => {}
        }
    }
    depth == 0
}

/// T1 — `z3` kept **resident** across goals. Same decision semantics as
/// [`Z3Smt`], but the per-goal energy excludes process startup.
pub struct Z3Server {
    session: Mutex<Option<Z3Session>>,
}

impl Z3Server {
    pub fn new() -> Self {
        Self {
            session: Mutex::new(Z3Session::spawn()),
        }
    }
}

impl Default for Z3Server {
    fn default() -> Self {
        Self::new()
    }
}

impl Solver for Z3Server {
    fn tier(&self) -> Tier {
        Tier::Smt
    }
    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        let start = Instant::now();
        let outcome = if looks_like_smt(&goal.statement) {
            self.session
                .lock()
                .ok()
                .and_then(|mut g| g.as_mut().and_then(|s| s.decide(&goal.statement)))
        } else {
            None
        };
        let verdict = match outcome.as_deref() {
            Some("unsat") => Verdict::Proved {
                artifact: ContentId::of_bytes(goal.statement.as_bytes()),
            },
            // `sat` ⇒ refuted. The resident session decided it; fetch the
            // witness with a one-shot `z3` (the refutation path is terminal, so
            // the extra spawn never touches the hot proved path).
            Some("sat") => Verdict::Refuted {
                counter_model: z3_counter_model(&goal.statement),
            },
            _ => Verdict::Unproved,
        };
        let uj = ns_to_uj(start.elapsed().as_nanos()).max(1);
        Attempt {
            tier: Tier::Smt,
            verdict,
            energy: Energy::modeled(uj),
        }
    }
}

/// **Lean 4 verifier backend** — the real proof-assistant rung, the SOTA
/// verifier substrate (Lean + Mathlib is what AlphaProof/Aristotle/Axiom build
/// on). It runs the `lean` typechecker on a goal whose `statement` *is* a
/// complete Lean theorem-with-proof and reports what Lean actually decided:
///
/// - typechecks (exit 0, no `error:`, no `sorry`) → `Proved` (Lean *is* the
///   verifier that makes `Proved` meaningful — the trust root, not a model).
/// - Lean's `decide` reports the proposition is false → `Refuted` (a decision,
///   carrying Lean's message as the counter-model witness — the cascade halts).
/// - any other failure, or `lean` not installed → `Unproved` (this rung did not
///   close it; the cascade escalates). Check [`LeanVerifier::available`] to skip
///   the rung when the toolchain is absent.
///
/// The goal→Lean encoding is the caller's (`context` carries the dialect, e.g.
/// `"lean4"`); the backend only attempts goals whose context names Lean. Energy
/// is `ModelBased` (wall time × the crate's power model), like the z3 rung.
pub struct LeanVerifier {
    tier: Tier,
    lean_bin: String,
}

static LEAN_TMP_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

impl Default for LeanVerifier {
    fn default() -> Self {
        Self::new()
    }
}

impl LeanVerifier {
    /// A Lean verifier invoking `lean`, declared at the automation tier (Lean's
    /// `decide`/`rfl`/`omega` are deterministic decision procedures; a
    /// Mathlib-heavy proof can be re-tiered with [`LeanVerifier::at_tier`]).
    pub fn new() -> Self {
        Self { tier: Tier::Automation, lean_bin: "lean".into() }
    }

    /// Place this verifier at a specific tier on the ladder.
    pub fn at_tier(mut self, tier: Tier) -> Self {
        self.tier = tier;
        self
    }

    /// Use a specific `lean` binary (e.g. a `lake env lean` wrapper for a
    /// Mathlib project).
    pub fn with_bin(mut self, bin: impl Into<String>) -> Self {
        self.lean_bin = bin.into();
        self
    }

    /// Whether `lean` is reachable on this host.
    pub fn available(&self) -> bool {
        Command::new(&self.lean_bin)
            .arg("--version")
            .output()
            .map(|o| o.status.success())
            .unwrap_or(false)
    }

    fn targets_lean(goal: &ProofGoal) -> bool {
        goal.context.to_ascii_lowercase().contains("lean")
    }

    /// Run `lean` on `source`; returns `(typechecks, combined_output)`.
    fn run(&self, source: &str) -> std::io::Result<(bool, String)> {
        let path = std::env::temp_dir().join(format!(
            "proof-lean-{}-{}.lean",
            std::process::id(),
            LEAN_TMP_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
        ));
        std::fs::write(&path, source)?;
        let result = Command::new(&self.lean_bin).arg(&path).output();
        let _ = std::fs::remove_file(&path);
        let out = result?;
        let mut combined = String::from_utf8_lossy(&out.stderr).into_owned();
        combined.push_str(&String::from_utf8_lossy(&out.stdout));
        // A real proof: clean exit, no error, and no `sorry` escape hatch.
        let ok = out.status.success()
            && !combined.contains("error:")
            && !combined.contains("sorry");
        Ok((ok, combined))
    }
}

impl Solver for LeanVerifier {
    fn tier(&self) -> Tier {
        self.tier
    }

    fn attempt(&self, goal: &ProofGoal) -> Attempt {
        // Not a Lean-dialect goal → this rung is inapplicable.
        if !Self::targets_lean(goal) {
            return Attempt {
                tier: self.tier,
                verdict: Verdict::Unproved,
                energy: Energy::estimate_tier(self.tier),
            };
        }
        let started = Instant::now();
        let verdict = match self.run(&goal.statement) {
            Ok((true, _)) => Verdict::Proved {
                artifact: ContentId::of_bytes(goal.statement.as_bytes()),
            },
            Ok((false, out)) => {
                // `decide` deciding the proposition false is a refutation, not a
                // mere failed proof — a decision the cascade must not escalate.
                if out.contains("proved that the proposition") && out.contains("is false") {
                    let witness = out.lines().find(|l| l.contains("error:")).unwrap_or("").trim();
                    Verdict::Refuted { counter_model: witness.to_string() }
                } else {
                    Verdict::Unproved
                }
            }
            // `lean` missing or a spawn failure: this rung simply didn't close it.
            Err(_) => Verdict::Unproved,
        };
        Attempt {
            tier: self.tier,
            verdict,
            energy: Energy::modeled(ns_to_uj(started.elapsed().as_nanos())),
        }
    }
}

/// The pure stand-in ladder, cheapest-first — every rung a `GoalKind`-gated
/// model, used by conformance to exercise routing deterministically.
pub fn reference_ladder() -> Vec<Box<dyn Solver>> {
    vec![
        Box::new(Automation),
        Box::new(Smt),
        Box::new(SmallModel),
        Box::new(FrontierModel),
    ]
}

/// A ladder with a real propositional decider at T0 over the stand-ins above.
pub fn real_t0_ladder() -> Vec<Box<dyn Solver>> {
    vec![
        Box::new(Propositional),
        Box::new(Smt),
        Box::new(SmallModel),
        Box::new(FrontierModel),
    ]
}

/// Two real rungs: the propositional decider (T0) and `z3` (T1), over the model
/// stand-ins.
pub fn real_ladder() -> Vec<Box<dyn Solver>> {
    vec![
        Box::new(Propositional),
        Box::new(Z3Smt),
        Box::new(SmallModel),
        Box::new(FrontierModel),
    ]
}

/// Like [`real_ladder`] but with `z3` kept resident at T1.
pub fn resident_z3_ladder() -> Vec<Box<dyn Solver>> {
    vec![
        Box::new(Propositional),
        Box::new(Z3Server::new()),
        Box::new(SmallModel),
        Box::new(FrontierModel),
    ]
}

/// Whether a working `z3` is on PATH — lets callers/tests adapt honestly.
pub fn z3_available() -> bool {
    Command::new("z3")
        .arg("--version")
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn each_rung_reaches_its_kind() {
        let triv = ProofGoal::new("x = x", "lean4", GoalKind::Trivial);
        assert!(Automation.attempt(&triv).verdict.is_proved());
        let dec = ProofGoal::new("decidable thing", "lean4", GoalKind::Decidable);
        assert!(!Automation.attempt(&dec).verdict.is_proved());
        assert!(Smt.attempt(&dec).verdict.is_proved());
        let gen = ProofGoal::new("needs search", "lean4", GoalKind::General);
        assert!(SmallModel.attempt(&gen).verdict.is_proved());
        let hard = ProofGoal::new("research conjecture", "lean4", GoalKind::Hard);
        assert!(FrontierModel.attempt(&hard).verdict.is_proved());
        assert_eq!(reference_ladder().len(), 4);
    }

    #[test]
    fn real_t0_proves_tautologies_and_refuses_the_rest() {
        let taut = ProofGoal::new("a | ~a", "prop", GoalKind::Trivial);
        let a = Propositional.attempt(&taut);
        assert!(a.verdict.is_proved());
        assert!(a.energy.measured_uj.is_none()); // ModelBased: no hardware shunt
        assert!(a.energy.uj() >= 1);
        assert!(!Propositional
            .attempt(&ProofGoal::new("a -> b", "prop", GoalKind::Trivial))
            .verdict
            .is_proved());
        assert!(!Propositional
            .attempt(&ProofGoal::new("∀ n, n+0 = n", "lean4", GoalKind::General))
            .verdict
            .is_proved());
        assert_eq!(real_t0_ladder().len(), 4);
    }

    #[test]
    fn z3_decides_real_smt_goals_when_available() {
        let valid = ProofGoal::new(
            "(declare-const x Int)\n(assert (not (=> (> x 0) (>= x 0))))",
            "smt2",
            GoalKind::Decidable,
        );
        let a = Z3Smt.attempt(&valid);
        assert_eq!(real_ladder().len(), 4);
        if z3_available() {
            assert!(a.verdict.is_proved());
            assert!(a.energy.uj() >= 1);
            let not_thm = ProofGoal::new(
                "(declare-const x Int)\n(assert (not (> x 0)))",
                "smt2",
                GoalKind::Decidable,
            );
            assert!(!Z3Smt.attempt(&not_thm).verdict.is_proved());
        } else {
            assert!(!a.verdict.is_proved());
        }
    }

    #[test]
    fn lean_verifies_proves_refutes_and_skips_non_lean() {
        let lean = LeanVerifier::new();
        // A non-Lean goal is inapplicable regardless of toolchain.
        let smt = ProofGoal::new("(assert true)", "smt2", GoalKind::Decidable);
        assert!(!lean.attempt(&smt).verdict.is_proved());

        if !lean.available() {
            eprintln!("lean not installed — skipping the live verification smoke");
            return;
        }
        // A real, checkable Lean proof → Proved (Lean is the verifier).
        let ok = ProofGoal::new("example : 1 + 1 = 2 := by decide", "lean4", GoalKind::Arithmetic);
        let a = lean.attempt(&ok);
        assert!(a.verdict.is_proved(), "valid lean proof must verify");
        assert!(a.energy.uj() >= 1);

        // `decide` deciding the proposition false → Refuted (a decision).
        let f = ProofGoal::new("example : 1 + 1 = 3 := by decide", "lean4", GoalKind::Arithmetic);
        assert!(lean.attempt(&f).verdict.is_refuted(), "decide-false must refute");

        // A failing-but-not-decided-false proof → Unproved.
        let u = ProofGoal::new("example : 2 + 2 = 5 := rfl", "lean4", GoalKind::Arithmetic);
        assert_eq!(lean.attempt(&u).verdict, Verdict::Unproved);

        // The `sorry` escape hatch must NOT count as proved.
        let s = ProofGoal::new("example (p : Prop) : p := by sorry", "lean4", GoalKind::General);
        assert!(!lean.attempt(&s).verdict.is_proved(), "sorry is not a proof");
    }

    #[test]
    fn resident_z3_decides_many_goals_on_one_session() {
        let server = Z3Server::new();
        let thm = ProofGoal::new(
            "(declare-const x Int)\n(assert (not (=> (> x 0) (>= x 0))))",
            "smt2",
            GoalKind::Decidable,
        );
        let comm = ProofGoal::new(
            "(declare-const a Int)(declare-const b Int)\n(assert (not (= (+ a b) (+ b a))))",
            "smt2",
            GoalKind::Decidable,
        );
        let prose = ProofGoal::new("∃ continuous f with property P", "lean4", GoalKind::General);
        assert!(!server.attempt(&prose).verdict.is_proved());
        if z3_available() {
            assert!(server.attempt(&thm).verdict.is_proved());
            assert!(server.attempt(&comm).verdict.is_proved());
            assert!(server.attempt(&thm).verdict.is_proved()); // still healthy after prose
        }
        assert_eq!(resident_z3_ladder().len(), 4);
        assert!(looks_like_smt("(assert (> x 0))"));
        assert!(!looks_like_smt("∃ continuous f"));
        assert!(!looks_like_smt("(unbalanced"));
    }
}