diff --git a/Cargo.lock b/Cargo.lock index 6abda43..68c0e26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1339,6 +1339,7 @@ version = "0.1.0" dependencies = [ "anyhow", "openproof-lean", + "openproof-protocol", "reqwest", "serde", "serde_json", diff --git a/crates/openproof-cli/src/autonomous.rs b/crates/openproof-cli/src/autonomous.rs index 8bb2178..862d01c 100644 --- a/crates/openproof-cli/src/autonomous.rs +++ b/crates/openproof-cli/src/autonomous.rs @@ -18,7 +18,7 @@ use openproof_core::{AppEvent, AppState, AutonomousRunPatch}; use openproof_lean::lsp_mcp::LeanLspMcp; use openproof_protocol::{AgentRole, AgentStatus, BranchQueueState, SearchStrategy}; use openproof_search::config::TacticSearchConfig; -use openproof_search::search::best_first_search; +use openproof_search::lsp_search::best_first_search; use openproof_store::AppStore; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -1233,12 +1233,19 @@ fn spawn_tactic_search_for_sorrys( "[tactic-search] Pantograph search at line {line}: {}", &goal_type[..goal_type.len().min(60)] ); + let on_goal = { + let tx = tx.clone(); + move |goal: openproof_protocol::ProofGoal| { + let _ = tx.send(AppEvent::ProofGoalUpdated(goal)); + } + }; match openproof_search::search::pantograph_best_first_search( &pg, &propose_fn, &goal_type, "", &config, + Some(&on_goal), ) { Ok(result) => emit_search_result(&tx, &node_id, line, result), Err(e) => eprintln!("[tactic-search] Pantograph error at line {line}: {e}"), @@ -1254,7 +1261,21 @@ fn spawn_tactic_search_for_sorrys( let scratch = scratch_path.clone(); tokio::task::spawn_blocking(move || { eprintln!("[tactic-search] LSP search at line {line}"); - match best_first_search(&lsp, &propose_fn, &scratch, line, "", &config) { + let on_goal = { + let tx = tx.clone(); + move |goal: openproof_protocol::ProofGoal| { + let _ = tx.send(AppEvent::ProofGoalUpdated(goal)); + } + }; + match best_first_search( + &lsp, + &propose_fn, + &scratch, + line, + "", + &config, + Some(&on_goal), + ) { Ok(result) => emit_search_result(&tx, &node_id, line, result), Err(e) => eprintln!("[tactic-search] LSP error at line {line}: {e}"), } diff --git a/crates/openproof-core/src/apply.rs b/crates/openproof-core/src/apply.rs index 70febac..18ba70d 100644 --- a/crates/openproof-core/src/apply.rs +++ b/crates/openproof-core/src/apply.rs @@ -303,6 +303,15 @@ impl AppState { node.status = openproof_protocol::ProofNodeStatus::Proving; } } + // Mark matching proof goals as closed by BFS (solved_by = None) + for goal in &mut session.proof.proof_goals { + if goal.sorry_line == Some(sorry_line) + && goal.status != openproof_protocol::GoalStatus::Closed + { + goal.status = openproof_protocol::GoalStatus::Closed; + // solved_by stays None = BFS implicit + } + } session.proof.phase = "verifying".to_string(); } if let Some(session) = self.current_session().cloned() { diff --git a/crates/openproof-dashboard/static/graph.js b/crates/openproof-dashboard/static/graph.js index 0e6911f..de2e118 100644 --- a/crates/openproof-dashboard/static/graph.js +++ b/crates/openproof-dashboard/static/graph.js @@ -4,303 +4,283 @@ import { ReactFlow, Background, Controls, MiniMap, Handle, Position } from "http const h = htm.bind(React.createElement); -// ── Color helpers ────────────────────────────────────────────────────── - -const statusColor = (s) => { - const st = String(s || "").toLowerCase(); - if (st === "verified" || st === "done") return "#22c55e"; - if (st === "proving" || st === "running") return "#eab308"; - if (st === "failed" || st === "error" || st === "blocked") return "#ef4444"; - return "#525252"; +// ── Colors ───────────────────────────────────────────────────────────── + +const STATUS_COLORS = { + open: "#f59e0b", + in_progress: "#3b82f6", + closed: "#22c55e", + failed: "#ef4444", }; -const roleColor = (r) => { - const role = String(r || "").toLowerCase(); - if (role === "prover") return "#3b82f6"; - if (role === "repairer") return "#f59e0b"; - if (role === "planner") return "#8b5cf6"; - if (role === "retriever") return "#06b6d4"; - if (role === "critic") return "#ec4899"; - return "#6b7280"; +const ATTRIBUTION_COLORS = { + "agent:prover": "#3b82f6", + "agent:repairer": "#f59e0b", + "agent:planner": "#8b5cf6", + "agent:retriever": "#06b6d4", + "agent:critic": "#ec4899", + bfs: "#22c55e", }; -const kindIcon = (k) => { - const kind = String(k || "").toLowerCase(); - if (kind === "theorem") return "\u{1D4AF}"; - if (kind === "lemma") return "\u{2113}"; - if (kind === "def" || kind === "artifact") return "\u{1D49F}"; - if (kind === "axiom") return "\u{1D49C}"; - return "\u25CB"; +const STATUS_ICONS = { + open: "\u25CB", + in_progress: "\u25D4", + closed: "\u2713", + failed: "\u2717", }; -// ── Node components ──────────────────────────────────────────────────── +// ── Tree layout ──────────────────────────────────────────────────────── + +const H_GAP = 240; +const V_GAP = 120; + +function computeTreeLayout(goals) { + const childrenMap = new Map(); + const goalMap = new Map(); + const roots = []; + + for (const g of goals) { + goalMap.set(g.id, g); + const pid = g.parent_goal_id || g.parentGoalId; + if (pid) { + if (!childrenMap.has(pid)) childrenMap.set(pid, []); + childrenMap.get(pid).push(g.id); + } else { + roots.push(g.id); + } + } + + // Compute subtree widths (leaf = 1) + const widths = new Map(); + function subtreeWidth(id) { + if (widths.has(id)) return widths.get(id); + const ch = childrenMap.get(id) || []; + const w = ch.length === 0 ? 1 : ch.reduce((sum, c) => sum + subtreeWidth(c), 0); + widths.set(id, w); + return w; + } + for (const r of roots) subtreeWidth(r); + + // Assign positions + const positions = new Map(); + function layout(id, x, depth) { + positions.set(id, { x, y: depth * V_GAP }); + const ch = childrenMap.get(id) || []; + if (ch.length === 0) return; + const totalW = ch.reduce((s, c) => s + subtreeWidth(c), 0); + let curX = x - ((totalW - 1) * H_GAP) / 2; + for (const c of ch) { + const w = subtreeWidth(c); + layout(c, curX + ((w - 1) * H_GAP) / 2, depth + 1); + curX += w * H_GAP; + } + } + + // Layout each root with horizontal offset between them + let rootX = 0; + for (const r of roots) { + const w = subtreeWidth(r); + layout(r, rootX + ((w - 1) * H_GAP) / 2, 0); + rootX += w * H_GAP + H_GAP; + } + + return { positions, childrenMap }; +} + +// ── GoalNode component ───────────────────────────────────────────────── + +function GoalNodeComponent({ data }) { + const [expanded, setExpanded] = useState(false); + const status = String(data.status || "open").toLowerCase(); + const color = STATUS_COLORS[status] || "#737373"; + const isLeaf = !data._hasChildren; + const isFrontier = isLeaf && status === "open"; + const isActive = status === "in_progress"; + const attribution = data.solved_by || data.solvedBy || (status === "closed" ? "bfs" : null); + const attrColor = attribution ? (ATTRIBUTION_COLORS[attribution] || "#737373") : null; + + const goalText = data.goal_text || data.goalText || ""; + const failedTactics = data.failed_tactics || data.failedTactics || []; + const attempts = data.attempts || 0; -function ProofNodeComponent({ data }) { - const borderColor = statusColor(data.status); return h` -
setExpanded(!expanded)} style=${{ + background: "#0f0f0f", + border: "2px solid " + color, borderRadius: 8, padding: "8px 12px", - minWidth: 160, - maxWidth: 240, + minWidth: 180, + maxWidth: expanded ? 400 : 260, fontFamily: "system-ui, sans-serif", + cursor: "pointer", + transition: "all 0.2s", + animation: isFrontier ? "frontier-pulse 2s infinite" : isActive ? "active-pulse 1.5s infinite" : "none", }}> - <${Handle} type="target" position=${Position.Top} style=${{ background: "#555" }} /> + <${Handle} type="target" position=${Position.Top} style=${{ background: color }} /> +
- ${kindIcon(data.kind)} - ${data.label} + + ${STATUS_ICONS[status] || "\u25CB"} + + + ${status} + + ${attribution ? h` + ${attribution} + ` : null} + ${isFrontier ? h` + frontier + ` : null}
-
- ${data.kind || "node"} \u00b7 ${data.status} + +
+ ${goalText || "(empty goal)"}
- ${data.statement ? h` -
- ${data.statement} + + ${expanded && failedTactics.length > 0 ? h` +
+
Failed tactics (${failedTactics.length}):
+ ${failedTactics.slice(0, 5).map((t, i) => h` +
${t}
+ `)} + ${failedTactics.length > 5 ? h`
...and ${failedTactics.length - 5} more
` : null}
` : null} - <${Handle} type="source" position=${Position.Bottom} style=${{ background: "#555" }} /> -
- `; -} -function BranchNodeComponent({ data }) { - const color = roleColor(data.role); - const hasSnippet = !!(data.lean_snippet || data.leanSnippet || "").trim(); - return h` -
- <${Handle} type="target" position=${Position.Top} style=${{ background: color }} /> -
- ${data.role}${data.hidden ? " (hidden)" : ""} - ${hasSnippet ? h`\u25CF` : null} -
-
- ${String(data.status || "idle")} \u00b7 score ${(data.score || 0).toFixed(0)} \u00b7 ${data.attempt_count || data.attemptCount || 0} tries -
- ${data.summary ? h` -
- ${data.summary} + ${expanded && attempts > 0 ? h` +
+ ${attempts} tactic${attempts !== 1 ? "s" : ""} attempted
` : null} + + <${Handle} type="source" position=${Position.Bottom} style=${{ background: color }} />
`; } -function GoalNodeComponent({ data }) { - const colors = { open: "#f59e0b", in_progress: "#3b82f6", closed: "#22c55e", failed: "#ef4444" }; - const color = colors[data.status] || "#737373"; - const failCount = (data.failed_tactics || data.failedTactics || []).length; +const nodeTypes = { goalNode: GoalNodeComponent }; + +// ── Stats bar ────────────────────────────────────────────────────────── + +function StatsBar({ stats, proof }) { + if (!stats) return null; + const verification = proof?.last_verification; return h` -
- <${Handle} type="target" position=${Position.Top} style=${{ background: color }} /> -
- - ${data.status === "closed" ? "\u2713" : data.status === "failed" ? "\u2717" : "\u25CB"} goal - - ${failCount > 0 ? h` - - ${failCount} failed +
+ Phase: ${proof?.phase || "idle"} + \u00a0\u00b7\u00a0 Goals: ${stats.total} + \u00a0\u00b7\u00a0 Closed: ${stats.closed} + \u00a0\u00b7\u00a0 Open: ${stats.open} + ${stats.failed > 0 ? h`\u00a0\u00b7\u00a0 Failed: ${stats.failed}` : null} + ${stats.frontier > 0 ? h`\u00a0\u00b7\u00a0 BFS frontier: ${stats.frontier}` : null} + ${verification ? h` + \u00a0\u00b7\u00a0 + + ${verification.ok ? "Lean verified" : "Lean failed"} - ` : null} - ${data.attempts > 0 ? h` - ${data.attempts} tried - ` : null} -
-
- ${(data.goal_text || data.goalText || "").substring(0, 60)} -
- ${data.tactic_applied || data.tacticApplied ? h` -
- via: ${data.tactic_applied || data.tacticApplied} -
+ ` : null} - <${Handle} type="source" position=${Position.Bottom} style=${{ background: color }} />
`; } -const nodeTypes = { - proofNode: ProofNodeComponent, - branchNode: BranchNodeComponent, - goalNode: GoalNodeComponent, -}; - // ── GraphTab ─────────────────────────────────────────────────────────── export function GraphTab({ session }) { const proof = session?.proof; - const proofNodes = proof?.nodes || []; - const allBranches = proof?.branches || []; - const [showBranches, setShowBranches] = useState(false); - const branches = showBranches ? allBranches : []; + const goals = proof?.proof_goals || proof?.proofGoals || []; - if (proofNodes.length === 0 && allBranches.length === 0) { - return h`
No proof nodes to visualize
`; - } + const { flowNodes, flowEdges, stats } = useMemo(() => { + if (goals.length === 0) return { flowNodes: [], flowEdges: [], stats: null }; - const { flowNodes, flowEdges } = useMemo(() => { - const nodes = []; - const edges = []; - const hasTree = proofNodes.some(n => n.parent_id || n.parentId); - - if (hasTree) { - const byDepth = {}; - for (const n of proofNodes) { - const d = n.depth || 0; - if (!byDepth[d]) byDepth[d] = []; - byDepth[d].push(n); - } - for (const n of proofNodes) { - const d = n.depth || 0; - const siblings = byDepth[d] || []; - const idx = siblings.indexOf(n); - const totalWidth = siblings.length * 220; - const startX = -(totalWidth / 2) + 110; - nodes.push({ - id: n.id, type: "proofNode", - position: { x: startX + idx * 220, y: d * 120 }, - draggable: true, - data: { ...n, _nodeColor: statusColor(n.status) }, - }); - } - } else { - const cols = 2; - for (let i = 0; i < proofNodes.length; i++) { - const n = proofNodes[i]; - const col = i % cols; - const row = Math.floor(i / cols); - nodes.push({ - id: n.id, type: "proofNode", - position: { x: col * 280, y: row * 100 }, - draggable: true, - data: { ...n, _nodeColor: statusColor(n.status) }, - }); - } - } + const { positions, childrenMap } = computeTreeLayout(goals); - let prevId = null; - for (const n of proofNodes) { - const parentId = n.parent_id || n.parentId; - if (parentId) { - edges.push({ - id: "tree-" + n.id, source: parentId, target: n.id, - style: { stroke: "#3b82f6", strokeWidth: 2 }, - animated: n.status === "proving", - }); - } else if (prevId && !hasTree) { - edges.push({ - id: "flow-" + n.id, source: prevId, target: n.id, - style: { stroke: "#333", strokeWidth: 1 }, - type: "smoothstep", - }); - } - prevId = n.id; - - for (const depId of (n.depends_on || n.dependsOn || [])) { - edges.push({ - id: "dep-" + n.id + "-" + depId, source: depId, - target: n.id, - style: { stroke: "#6b7280", strokeWidth: 1, strokeDasharray: "4 3" }, - }); - } - } + const nodes = goals.map((g) => ({ + id: g.id, + type: "goalNode", + position: positions.get(g.id) || { x: 0, y: 0 }, + draggable: true, + data: { + ...g, + _hasChildren: childrenMap.has(g.id) && childrenMap.get(g.id).length > 0, + _nodeColor: STATUS_COLORS[String(g.status || "open").toLowerCase()] || "#737373", + }, + })); - const maxDepth = Math.max(0, ...proofNodes.map((n) => n.depth || 0)); - const branchY = (maxDepth + 1) * 100 + 40; - const branchCols = {}; - - for (const b of branches) { - const focusId = b.focus_node_id || b.focusNodeId || proofNodes[0]?.id; - if (!branchCols[focusId]) branchCols[focusId] = 0; - const col = branchCols[focusId]++; - const parentNode = nodes.find((n) => n.id === focusId); - const baseX = parentNode ? parentNode.position.x : col * 160; - const bId = "branch-" + b.id; - nodes.push({ - id: bId, type: "branchNode", - position: { x: baseX + col * 160, y: branchY }, - draggable: true, - data: { ...b, _nodeColor: roleColor(b.role) }, + const edges = goals + .filter((g) => g.parent_goal_id || g.parentGoalId) + .map((g) => { + const tactic = g.tactic_applied || g.tacticApplied || ""; + const st = String(g.status || "open").toLowerCase(); + return { + id: "tactic-" + g.id, + source: g.parent_goal_id || g.parentGoalId, + target: g.id, + label: tactic.length > 40 ? tactic.substring(0, 37) + "..." : tactic, + labelStyle: { fill: "#22c55e", fontSize: 9, fontFamily: "monospace" }, + style: { stroke: STATUS_COLORS[st] || "#525252", strokeWidth: 1.5 }, + type: "smoothstep", + animated: st === "in_progress", + }; }); - if (focusId) { - edges.push({ - id: "agent-" + b.id, source: focusId, target: bId, - style: { stroke: roleColor(b.role), strokeWidth: 1, strokeDasharray: "4 3", opacity: 0.5 }, - }); - } - } - const proofGoals = proof?.proof_goals || proof?.proofGoals || []; - const goalY = (maxDepth + 1) * 100 + (branches.length > 0 ? 180 : 40); - for (let i = 0; i < proofGoals.length; i++) { - const g = proofGoals[i]; - const gId = "goal-" + g.id; - nodes.push({ - id: gId, type: "goalNode", - position: { x: i * 200 - (proofGoals.length * 100), y: goalY + (g.parent_goal_id || g.parentGoalId ? 80 : 0) }, - draggable: true, data: g, - }); - const parentGoalId = g.parent_goal_id || g.parentGoalId; - if (parentGoalId) { - edges.push({ - id: "goaltree-" + g.id, source: "goal-" + parentGoalId, target: gId, - style: { stroke: "#3b82f6", strokeWidth: 1, strokeDasharray: "2 2" }, - label: g.tactic_applied || g.tacticApplied || "", - labelStyle: { fill: "#22c55e", fontSize: 8 }, - }); - } - if (!parentGoalId && proofNodes.length > 0) { - const activeNodeId = proof?.active_node_id || proof?.activeNodeId || proofNodes[0]?.id; - edges.push({ - id: "nodegoal-" + g.id, source: activeNodeId, target: gId, - style: { stroke: "#f59e0b", strokeWidth: 1, strokeDasharray: "3 3", opacity: 0.6 }, - }); - } + // Compute stats + let closed = 0, open = 0, failed = 0, inProgress = 0, frontier = 0; + for (const g of goals) { + const s = String(g.status || "open").toLowerCase(); + if (s === "closed") closed++; + else if (s === "failed") failed++; + else if (s === "in_progress") inProgress++; + else open++; + const isLeaf = !childrenMap.has(g.id) || childrenMap.get(g.id).length === 0; + if (s === "open" && isLeaf) frontier++; } - return { flowNodes: nodes, flowEdges: edges }; - }, [proofNodes, branches, proof]); + return { + flowNodes: nodes, + flowEdges: edges, + stats: { total: goals.length, closed, open, failed, inProgress, frontier }, + }; + }, [goals]); - const verification = proof?.last_verification; - const attemptNum = proof?.attempt_number || proof?.attemptNumber || 0; + if (goals.length === 0) { + return h`
+
+
No proof goals yet
+
Start a proof to see the tactic tree
+
+
`; + } return h`
-
- Phase: ${proof?.phase || "idle"} - \u00a0\u00b7\u00a0 Nodes: ${proofNodes.length} - \u00a0\u00b7\u00a0 Goals: ${(proof?.proof_goals || proof?.proofGoals || []).length} - \u00a0\u00b7\u00a0 Attempts: ${attemptNum} - - ${verification ? h` - \u00a0\u00b7\u00a0 - - ${verification.ok ? "Lean verified" : "Lean failed"} - - - ` : null} -
+ <${StatsBar} stats=${stats} proof=${proof} />
<${ReactFlow} nodes=${flowNodes} @@ -308,7 +288,7 @@ export function GraphTab({ session }) { nodeTypes=${nodeTypes} fitView fitViewOptions=${{ padding: 0.3 }} - minZoom=${0.2} + minZoom=${0.1} maxZoom=${2} defaultViewport=${{ x: 0, y: 0, zoom: 0.8 }} proOptions=${{ hideAttribution: true }} diff --git a/crates/openproof-dashboard/static/styles.css b/crates/openproof-dashboard/static/styles.css index 1dd1bba..4af52b1 100644 --- a/crates/openproof-dashboard/static/styles.css +++ b/crates/openproof-dashboard/static/styles.css @@ -378,3 +378,13 @@ body { /* Empty state */ .empty { display: flex; align-items: center; justify-content: center; height: 100%; color: var(--muted); } + +/* Proof tree graph animations */ +@keyframes frontier-pulse { + 0%, 100% { box-shadow: 0 0 0 0 rgba(251, 191, 36, 0.4); } + 50% { box-shadow: 0 0 0 6px rgba(251, 191, 36, 0); } +} +@keyframes active-pulse { + 0%, 100% { box-shadow: 0 0 0 0 rgba(59, 130, 246, 0.4); } + 50% { box-shadow: 0 0 0 6px rgba(59, 130, 246, 0); } +} diff --git a/crates/openproof-protocol/src/lib.rs b/crates/openproof-protocol/src/lib.rs index 5550383..a98f436 100644 --- a/crates/openproof-protocol/src/lib.rs +++ b/crates/openproof-protocol/src/lib.rs @@ -118,6 +118,8 @@ pub struct ProofGoal { pub sorry_line: Option, /// Pantograph state ID (for chaining further tactics). pub state_id: Option, + /// Who closed this goal: "agent:prover", "agent:repairer", etc. None = BFS. + pub solved_by: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] diff --git a/crates/openproof-search/Cargo.toml b/crates/openproof-search/Cargo.toml index 028eec9..154269f 100644 --- a/crates/openproof-search/Cargo.toml +++ b/crates/openproof-search/Cargo.toml @@ -10,3 +10,4 @@ reqwest = { workspace = true, features = ["blocking"] } serde.workspace = true serde_json.workspace = true openproof-lean = { path = "../openproof-lean" } +openproof-protocol = { path = "../openproof-protocol" } diff --git a/crates/openproof-search/src/lib.rs b/crates/openproof-search/src/lib.rs index 2c724a1..f682586 100644 --- a/crates/openproof-search/src/lib.rs +++ b/crates/openproof-search/src/lib.rs @@ -6,5 +6,6 @@ pub mod cache; pub mod config; +pub mod lsp_search; pub mod ollama; pub mod search; diff --git a/crates/openproof-search/src/lsp_search.rs b/crates/openproof-search/src/lsp_search.rs new file mode 100644 index 0000000..bfa181f --- /dev/null +++ b/crates/openproof-search/src/lsp_search.rs @@ -0,0 +1,259 @@ +//! LSP-based best-first tactic search. +//! +//! Operates by recompiling the file for each tactic test (~100ms per tactic). +//! Prefer `pantograph_best_first_search` when Pantograph is available (~3ms). + +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashSet}; +use std::path::Path; +use std::sync::Mutex; +use std::time::Instant; + +use anyhow::{bail, Result}; + +use openproof_lean::goal_state::AttemptResult; +use openproof_lean::lsp_mcp::LeanLspMcp; +use openproof_protocol::{GoalStatus, ProofGoal}; + +use crate::cache::{hash_goals, TacticCache}; +use crate::config::{SearchResult, TacticSearchConfig}; +use crate::search::{GoalUpdateFn, ProposeFn}; + +/// A node in the LSP search tree. +#[derive(Debug, Clone)] +struct SearchNode { + priority: u64, + score: usize, + tactics: Vec, + goals: Vec, + sorry_line: usize, +} + +impl SearchNode { + fn new( + goals: Vec, + tactics: Vec, + sorry_line: usize, + length_penalty: f64, + ) -> Self { + let score = goals.len(); + let priority = ((score as f64 + length_penalty * tactics.len() as f64) * 1000.0) as u64; + Self { + priority, + score, + tactics, + goals, + sorry_line, + } + } +} + +impl PartialEq for SearchNode { + fn eq(&self, other: &Self) -> bool { + self.priority == other.priority + } +} + +impl Eq for SearchNode {} + +impl PartialOrd for SearchNode { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for SearchNode { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + Reverse(self.priority) + .cmp(&Reverse(other.priority)) + .then_with(|| self.tactics.len().cmp(&other.tactics.len())) + } +} + +/// Run best-first search on a single sorry position using LSP. +pub fn best_first_search( + lsp: &Mutex, + propose_fn: &ProposeFn, + file_path: &Path, + sorry_line: usize, + retrieval_context: &str, + config: &TacticSearchConfig, + on_goal_update: Option<&GoalUpdateFn>, +) -> Result { + let start = Instant::now(); + let mut expansions: usize = 0; + let mut cache = TacticCache::new(); + let mut seen_states: HashSet = HashSet::new(); + let mut frontier: BinaryHeap = BinaryHeap::new(); + + // Get initial goal state + let initial_goals = { + let mut mcp = lsp.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?; + let goal_state = mcp.get_goals(file_path, sorry_line, None)?; + goal_state + .goals_before + .or(goal_state.goals) + .unwrap_or_default() + }; + + if initial_goals.is_empty() { + return Ok(SearchResult::Exhausted { expansions: 0 }); + } + + let initial_hash = hash_goals(&initial_goals); + seen_states.insert(initial_hash); + + let root_goal_id = format!("lsp-{sorry_line}-root"); + frontier.push(SearchNode::new( + initial_goals.clone(), + vec![], + sorry_line, + config.length_penalty, + )); + + if let Some(cb) = on_goal_update { + cb(ProofGoal { + id: root_goal_id.clone(), + goal_text: initial_goals.join("\n"), + status: GoalStatus::Open, + sorry_line: Some(sorry_line), + ..Default::default() + }); + } + + let mut goal_counter: usize = 0; + let mut best_partial = SearchNode { + priority: u64::MAX, + score: usize::MAX, + tactics: vec![], + goals: initial_goals, + sorry_line, + }; + + while let Some(node) = frontier.pop() { + if start.elapsed() > config.timeout { + return Ok(SearchResult::Timeout { + best_tactics: best_partial.tactics, + remaining_goals: best_partial.score, + }); + } + if expansions >= config.max_expansions { + return Ok(SearchResult::Exhausted { expansions }); + } + if node.score < best_partial.score { + best_partial = node.clone(); + } + + let goal_text = node.goals.first().map(|s| s.as_str()).unwrap_or(""); + if goal_text.is_empty() { + continue; + } + + let candidates = match propose_fn(goal_text, retrieval_context, config.beam_width) { + Ok(c) => c, + Err(_) => continue, + }; + if candidates.is_empty() { + continue; + } + + let mut to_screen: Vec = Vec::new(); + let mut cached_results: Vec<(String, AttemptResult)> = Vec::new(); + + for tactic in &candidates { + if let Some(cached) = cache.get(goal_text, tactic) { + cached_results.push((tactic.clone(), cached.clone())); + } else { + to_screen.push(tactic.clone()); + } + } + + let mut screen_results: Vec = Vec::new(); + if !to_screen.is_empty() { + let result = { + let mut mcp = lsp.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?; + if !mcp.is_alive() { + bail!("lean-lsp-mcp process died during search"); + } + mcp.screen_tactics(file_path, node.sorry_line, None, &to_screen)? + }; + for item in result.items { + cache.insert(goal_text, &item.snippet, item.clone()); + screen_results.push(item); + } + expansions += to_screen.len(); + } + + let all_results: Vec = cached_results + .into_iter() + .map(|(_, r)| r) + .chain(screen_results) + .collect(); + + for item in all_results { + if !item.succeeded() { + continue; + } + + if item.is_solved() { + let mut tactics = node.tactics.clone(); + tactics.push(item.snippet.clone()); + if let Some(cb) = on_goal_update { + goal_counter += 1; + cb(ProofGoal { + id: format!("lsp-{sorry_line}-{goal_counter}"), + goal_text: String::new(), + status: GoalStatus::Closed, + parent_goal_id: Some(root_goal_id.clone()), + tactic_applied: Some(item.snippet), + sorry_line: Some(sorry_line), + ..Default::default() + }); + } + return Ok(SearchResult::Solved { + tactics, + file_content: String::new(), + }); + } + + let new_goals = &item.goals; + let goals_hash = hash_goals(new_goals); + if config.dedup && !seen_states.insert(goals_hash) { + continue; + } + + let mut new_tactics = node.tactics.clone(); + new_tactics.push(item.snippet.clone()); + + if let Some(cb) = on_goal_update { + goal_counter += 1; + cb(ProofGoal { + id: format!("lsp-{sorry_line}-{goal_counter}"), + goal_text: new_goals.join("\n"), + status: GoalStatus::Open, + parent_goal_id: Some(root_goal_id.clone()), + tactic_applied: Some(item.snippet.clone()), + sorry_line: Some(sorry_line), + ..Default::default() + }); + } + + frontier.push(SearchNode::new( + new_goals.clone(), + new_tactics, + node.sorry_line, + config.length_penalty, + )); + } + } + + if best_partial.score < usize::MAX && !best_partial.tactics.is_empty() { + Ok(SearchResult::Partial { + tactics: best_partial.tactics, + remaining_goals: best_partial.score, + file_content: String::new(), + }) + } else { + Ok(SearchResult::Exhausted { expansions }) + } +} diff --git a/crates/openproof-search/src/search.rs b/crates/openproof-search/src/search.rs index 84b90b5..435a7fe 100644 --- a/crates/openproof-search/src/search.rs +++ b/crates/openproof-search/src/search.rs @@ -6,258 +6,24 @@ use std::cmp::Reverse; use std::collections::{BinaryHeap, HashSet}; -use std::path::Path; use std::sync::Mutex; use std::time::Instant; use anyhow::{bail, Result}; -use openproof_lean::goal_state::AttemptResult; -use openproof_lean::lsp_mcp::LeanLspMcp; use openproof_lean::pantograph::Pantograph; -use crate::cache::{hash_goals, TacticCache}; -use crate::config::{SearchResult, TacticSearchConfig}; - -/// A node in the search tree. -#[derive(Debug, Clone)] -struct SearchNode { - /// Length-normalized priority (milliunits): lower is better. - /// Computed as `(remaining_goals + length_penalty * depth) * 1000`. - priority: u64, - /// Raw remaining goal count (used for best-partial tracking). - score: usize, - /// Tactics applied so far from the initial state. - tactics: Vec, - /// Goal strings after applying these tactics. - goals: Vec, - /// The line in the file where the sorry is. - sorry_line: usize, -} - -impl SearchNode { - fn new( - goals: Vec, - tactics: Vec, - sorry_line: usize, - length_penalty: f64, - ) -> Self { - let score = goals.len(); - let priority = ((score as f64 + length_penalty * tactics.len() as f64) * 1000.0) as u64; - Self { - priority, - score, - tactics, - goals, - sorry_line, - } - } -} - -impl PartialEq for SearchNode { - fn eq(&self, other: &Self) -> bool { - self.priority == other.priority - } -} +use openproof_protocol::{GoalStatus, ProofGoal}; -impl Eq for SearchNode {} +use crate::cache::hash_goals; +use crate::config::{SearchResult, TacticSearchConfig}; -impl PartialOrd for SearchNode { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for SearchNode { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - // Min-heap by priority (lower = better), break ties by shorter proof - Reverse(self.priority) - .cmp(&Reverse(other.priority)) - .then_with(|| self.tactics.len().cmp(&other.tactics.len())) - } -} +/// Callback for emitting proof goal updates during search. +pub type GoalUpdateFn = dyn Fn(ProofGoal) + Send + Sync; /// Callback for LLM tactic proposal. Returns candidate tactics for a goal state. pub type ProposeFn = Box Result> + Send>; -/// Run best-first search on a single sorry position. -/// -/// Arguments: -/// - `lsp`: the lean-lsp-mcp client (mutex-protected for shared access) -/// - `propose_fn`: callback that asks the LLM to propose k tactics for a goal -/// - `file_path`: absolute path to the Lean file being worked on -/// - `sorry_line`: 1-indexed line number of the sorry to fill -/// - `retrieval_context`: corpus hits to include in the proposal prompt -/// - `config`: search parameters -pub fn best_first_search( - lsp: &Mutex, - propose_fn: &ProposeFn, - file_path: &Path, - sorry_line: usize, - retrieval_context: &str, - config: &TacticSearchConfig, -) -> Result { - let start = Instant::now(); - let mut expansions: usize = 0; - let mut cache = TacticCache::new(); - let mut seen_states: HashSet = HashSet::new(); - let mut frontier: BinaryHeap = BinaryHeap::new(); - - // Get initial goal state - let initial_goals = { - let mut mcp = lsp.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?; - let goal_state = mcp.get_goals(file_path, sorry_line, None)?; - goal_state - .goals_before - .or(goal_state.goals) - .unwrap_or_default() - }; - - if initial_goals.is_empty() { - // No goals at this position -- either already solved or not a sorry. - // Return Exhausted rather than Solved with empty tactics, since we - // didn't actually find the tactic that closes the goal. - return Ok(SearchResult::Exhausted { expansions: 0 }); - } - - let initial_hash = hash_goals(&initial_goals); - seen_states.insert(initial_hash); - - frontier.push(SearchNode::new( - initial_goals.clone(), - vec![], - sorry_line, - config.length_penalty, - )); - - let mut best_partial = SearchNode { - priority: u64::MAX, - score: usize::MAX, - tactics: vec![], - goals: initial_goals, - sorry_line, - }; - - while let Some(node) = frontier.pop() { - // Check timeout - if start.elapsed() > config.timeout { - return Ok(SearchResult::Timeout { - best_tactics: best_partial.tactics, - remaining_goals: best_partial.score, - }); - } - - // Check expansion budget - if expansions >= config.max_expansions { - return Ok(SearchResult::Exhausted { expansions }); - } - - // Track best partial result - if node.score < best_partial.score { - best_partial = node.clone(); - } - - // Use the first goal as the focus for tactic generation - let goal_text = node.goals.first().map(|s| s.as_str()).unwrap_or(""); - if goal_text.is_empty() { - continue; - } - - // Ask LLM for candidate tactics - let candidates = match propose_fn(goal_text, retrieval_context, config.beam_width) { - Ok(c) => c, - Err(_) => continue, - }; - - if candidates.is_empty() { - continue; - } - - // Filter out candidates we've already cached for this exact goal - let mut to_screen: Vec = Vec::new(); - let mut cached_results: Vec<(String, AttemptResult)> = Vec::new(); - - for tactic in &candidates { - if let Some(cached) = cache.get(goal_text, tactic) { - cached_results.push((tactic.clone(), cached.clone())); - } else { - to_screen.push(tactic.clone()); - } - } - - // Screen uncached tactics via LSP - let mut screen_results: Vec = Vec::new(); - if !to_screen.is_empty() { - let result = { - let mut mcp = lsp.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?; - if !mcp.is_alive() { - bail!("lean-lsp-mcp process died during search"); - } - mcp.screen_tactics(file_path, node.sorry_line, None, &to_screen)? - }; - - for item in result.items { - cache.insert(goal_text, &item.snippet, item.clone()); - screen_results.push(item); - } - expansions += to_screen.len(); - } - - // Process all results (cached + freshly screened) - let all_results: Vec = cached_results - .into_iter() - .map(|(_, r)| r) - .chain(screen_results) - .collect(); - - for item in all_results { - if !item.succeeded() { - continue; - } - - // Goal closed - if item.is_solved() { - let mut tactics = node.tactics.clone(); - tactics.push(item.snippet); - return Ok(SearchResult::Solved { - tactics, - file_content: String::new(), // caller fills this - }); - } - - // New sub-goals - let new_goals = &item.goals; - let goals_hash = hash_goals(new_goals); - - // Transposition table: skip if we've seen this state - if config.dedup && !seen_states.insert(goals_hash) { - continue; - } - - let mut new_tactics = node.tactics.clone(); - new_tactics.push(item.snippet.clone()); - - frontier.push(SearchNode::new( - new_goals.clone(), - new_tactics, - node.sorry_line, - config.length_penalty, - )); - } - } - - // Frontier exhausted - if best_partial.score < usize::MAX && !best_partial.tactics.is_empty() { - Ok(SearchResult::Partial { - tactics: best_partial.tactics, - remaining_goals: best_partial.score, - file_content: String::new(), - }) - } else { - Ok(SearchResult::Exhausted { expansions }) - } -} - // --------------------------------------------------------------------------- // Pantograph-native search (state-based, ~1000x faster) // --------------------------------------------------------------------------- @@ -326,6 +92,7 @@ pub fn pantograph_best_first_search( type_expr: &str, retrieval_context: &str, config: &TacticSearchConfig, + on_goal_update: Option<&GoalUpdateFn>, ) -> Result { let start = Instant::now(); let mut expansions: usize = 0; @@ -357,6 +124,7 @@ pub fn pantograph_best_first_search( }; seen_states.insert(initial_hash); + let root_goal_id = format!("pg-{initial_state_id}"); frontier.push(PantographNode::new( vec![type_expr.to_string()], vec![], @@ -364,6 +132,16 @@ pub fn pantograph_best_first_search( config.length_penalty, )); + if let Some(cb) = on_goal_update { + cb(ProofGoal { + id: root_goal_id.clone(), + goal_text: type_expr.to_string(), + status: GoalStatus::Open, + state_id: Some(initial_state_id), + ..Default::default() + }); + } + let mut best_partial = PantographNode { priority: u64::MAX, score: usize::MAX, @@ -428,10 +206,24 @@ pub fn pantograph_best_first_search( let new_state_id = result.new_state_id.unwrap(); allocated_states.push(new_state_id); + let child_goal_id = format!("pg-{new_state_id}"); + let parent_id = format!("pg-{}", node.state_id); + // Proof complete! if result.remaining_goals.is_empty() { let mut tactics = node.tactics.clone(); tactics.push(tactic.clone()); + if let Some(cb) = on_goal_update { + cb(ProofGoal { + id: child_goal_id, + goal_text: String::new(), + status: GoalStatus::Closed, + parent_goal_id: Some(parent_id), + tactic_applied: Some(tactic.clone()), + state_id: Some(new_state_id), + ..Default::default() + }); + } cleanup_states(pantograph, &allocated_states); return Ok(SearchResult::Solved { tactics, @@ -448,6 +240,18 @@ pub fn pantograph_best_first_search( let mut new_tactics = node.tactics.clone(); new_tactics.push(tactic.clone()); + if let Some(cb) = on_goal_update { + cb(ProofGoal { + id: child_goal_id, + goal_text: result.remaining_goals.join("\n"), + status: GoalStatus::Open, + parent_goal_id: Some(parent_id), + tactic_applied: Some(tactic.clone()), + state_id: Some(new_state_id), + ..Default::default() + }); + } + frontier.push(PantographNode::new( result.remaining_goals, new_tactics, diff --git a/crates/openproof-search/tests/minif2f_bench.rs b/crates/openproof-search/tests/minif2f_bench.rs index ba881dd..e21d031 100644 --- a/crates/openproof-search/tests/minif2f_bench.rs +++ b/crates/openproof-search/tests/minif2f_bench.rs @@ -265,7 +265,7 @@ fn minif2f_tactic_search_benchmark() { }; let start = Instant::now(); let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - pantograph_best_first_search(&pg, &propose_fn, type_expr, "", &config) + pantograph_best_first_search(&pg, &propose_fn, type_expr, "", &config, None) })); let elapsed = start.elapsed(); total_time += elapsed; diff --git a/crates/openproof-search/tests/search_integration.rs b/crates/openproof-search/tests/search_integration.rs index 6d253f6..aaf9f36 100644 --- a/crates/openproof-search/tests/search_integration.rs +++ b/crates/openproof-search/tests/search_integration.rs @@ -13,7 +13,8 @@ use openproof_lean::lsp_mcp::LeanLspMcp; use openproof_lean::pantograph::Pantograph; use openproof_lean::tools::{find_sorry_positions, run_lean_verify_raw}; use openproof_search::config::{SearchResult, TacticSearchConfig}; -use openproof_search::search::{best_first_search, pantograph_best_first_search, ProposeFn}; +use openproof_search::lsp_search::best_first_search; +use openproof_search::search::{pantograph_best_first_search, ProposeFn}; fn lean_project_dir() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")) @@ -175,7 +176,7 @@ theorem test_add_comm (a b : Nat) : a + b = b + a := by let (line, _) = sorrys[0]; println!("Searching sorry at line {} ...", line); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!("Completed in {:.2}s", start.elapsed().as_secs_f64()); print_result(&result); @@ -199,7 +200,7 @@ theorem test_ring (x : Int) : (x + 1) * (x + 1) = x * x + 2 * x + 1 := by let (line, _) = sorrys[0]; println!("Searching sorry at line {} ...", line); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!("Completed in {:.2}s", start.elapsed().as_secs_f64()); print_result(&result); @@ -223,7 +224,7 @@ theorem test_omega (n : Nat) : n < n + 1 := by let (line, _) = sorrys[0]; println!("Searching sorry at line {} ...", line); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!("Completed in {:.2}s", start.elapsed().as_secs_f64()); print_result(&result); @@ -255,7 +256,7 @@ theorem test_grind (a b c : Nat) (h1 : a = b) (h2 : b = c) : a = c := by let (line, _) = sorrys[0]; println!("Searching sorry at line {} with grind only ...", line); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!("Completed in {:.2}s", start.elapsed().as_secs_f64()); print_result(&result); @@ -285,7 +286,7 @@ theorem e2e_test (n : Nat) : n + 0 = n := by let (line, _) = sorrys[0]; println!("[e2e] Searching sorry at line {}...", line); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!( "[e2e] Search completed in {:.2}s", @@ -343,7 +344,7 @@ theorem sorry3 (n : Nat) : n * 1 = n := by for (i, &(line, _)) in sorrys.iter().enumerate() { println!("\n[multi] Sorry #{} at line {}...", i + 1, line); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!("[multi] Completed in {:.2}s", start.elapsed().as_secs_f64()); print_result(&result); @@ -382,7 +383,7 @@ theorem test_penalty (n : Nat) : 0 + n = n := by let (line, _) = sorrys[0]; println!("Searching with length_penalty=1.0 ..."); let start = Instant::now(); - let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config) + let result = best_first_search(&lsp, &propose_fn, &scratch_path, line, "", &config, None) .expect("search failed"); println!("Completed in {:.2}s", start.elapsed().as_secs_f64()); print_result(&result); @@ -423,8 +424,8 @@ fn pantograph_solves_simp_goal() { let goal = "forall (n : Nat), n + 0 = n"; println!("Goal: {goal}"); let start = Instant::now(); - let result = - pantograph_best_first_search(&pg, &propose_fn, goal, "", &config).expect("search failed"); + let result = pantograph_best_first_search(&pg, &propose_fn, goal, "", &config, None) + .expect("search failed"); println!("Completed in {:.3}s", start.elapsed().as_secs_f64()); print_result(&result); assert!(result.is_solved()); @@ -442,8 +443,8 @@ fn pantograph_solves_with_grind() { let goal = "forall (a b c : Nat), a = b -> b = c -> a = c"; println!("Goal (grind-only): {goal}"); let start = Instant::now(); - let result = - pantograph_best_first_search(&pg, &propose_fn, goal, "", &config).expect("search failed"); + let result = pantograph_best_first_search(&pg, &propose_fn, goal, "", &config, None) + .expect("search failed"); println!("Completed in {:.3}s", start.elapsed().as_secs_f64()); print_result(&result); assert!(result.is_solved(), "grind should solve transitivity"); @@ -471,7 +472,7 @@ fn pantograph_solves_multi_goals() { for (goal, label) in &goals { print!(" {label}: "); let start = Instant::now(); - let result = pantograph_best_first_search(&pg, &propose_fn, goal, "", &config) + let result = pantograph_best_first_search(&pg, &propose_fn, goal, "", &config, None) .expect("search failed"); let elapsed = start.elapsed(); match &result {