"""Lean proof profiling via CLI trace output."""
import asyncio
import os
import re
import tempfile
from collections import defaultdict
from pathlib import Path
from lean_lsp_mcp.models import LineProfile, ProofProfileResult
_TRACE_RE = re.compile(r"^(\s*)\[([^\]]+)\]\s+\[([\d.]+)\]\s+(.+)$")
_CUMULATIVE_RE = re.compile(r"^\s+(\S+(?:\s+\S+)*)\s+([\d.]+)(ms|s)$")
_DECL_RE = re.compile(r"^\s*(?:private\s+)?(theorem|lemma|def)\s+(\S+)")
_HEADER_RE = re.compile(r"^(import|open|set_option|universe|variable)\s")
_SKIP_CATEGORIES = {"import", "initialization", "parsing", "interpretation", "linting"}
def _find_header_end(lines: list[str]) -> int:
"""Find where imports/header ends and declarations begin."""
header_end, in_block = 0, False
for i, line in enumerate(lines):
s = line.strip()
if "/-" in line:
in_block = True
if "-/" in line:
in_block = False
if in_block or not s or s.startswith("--") or _HEADER_RE.match(line):
header_end = i + 1
elif s.startswith(("namespace", "section")):
header_end = i + 1
elif _DECL_RE.match(line) or s.startswith(("@[", "private ", "protected ")):
break
else:
header_end = i + 1
return header_end
def _find_theorem_end(lines: list[str], start: int) -> int:
"""Find where theorem ends (next declaration or EOF)."""
for i in range(start + 1, len(lines)):
if _DECL_RE.match(lines[i]):
return i
return len(lines)
def _extract_theorem_source(lines: list[str], target_line: int) -> tuple[str, str, int]:
"""Extract imports/header + single theorem. Returns (source, name, theorem_start_in_source)."""
m = _DECL_RE.match(lines[target_line - 1])
if not m:
raise ValueError(f"No theorem/lemma/def at line {target_line}")
header_end = _find_header_end(lines)
theorem_end = _find_theorem_end(lines, target_line - 1)
header = "\n".join(lines[:header_end])
theorem = "\n".join(lines[target_line - 1 : theorem_end])
return f"{header}\n\n{theorem}\n", m.group(2), header_end + 2
def _parse_output(
output: str,
) -> tuple[list[tuple[int, str, float, str]], dict[str, float]]:
"""Parse trace output into (traces, cumulative). Traces are (depth, cls, ms, msg)."""
traces, cumulative, in_cumulative = [], {}, False
for line in output.splitlines():
if "cumulative profiling times:" in line:
in_cumulative = True
elif in_cumulative and (m := _CUMULATIVE_RE.match(line)):
cat, val, unit = m.groups()
cumulative[cat] = float(val) * (1000 if unit == "s" else 1)
elif not in_cumulative and (m := _TRACE_RE.match(line)):
indent, cls, time_s, msg = m.groups()
traces.append((len(indent) // 2, cls, float(time_s) * 1000, msg))
return traces, cumulative
def _build_proof_items(
source_lines: list[str], proof_start: int
) -> list[tuple[int, str, bool]]:
"""Build list of (line_no, content, is_bullet) for proof lines."""
items = []
for i in range(proof_start, len(source_lines)):
s = source_lines[i].strip()
if s and not s.startswith("--"):
is_bullet = s[0] in "·*-"
items.append((i + 1, s.lstrip("·*- \t"), is_bullet))
return items
def _match_line(
tactic: str, is_bullet: bool, items: list[tuple[int, str, bool]], used: set[int]
) -> int | None:
"""Find matching source line for a tactic trace. Returns line number or None."""
for ln, content, src_bullet in items:
if ln in used:
continue
if is_bullet and src_bullet:
return ln
if (
not is_bullet
and content
and (tactic.startswith(content[:25]) or content.startswith(tactic[:25]))
):
return ln
return None
def _extract_line_times(
traces: list[tuple[int, str, float, str]],
name: str,
proof_items: list[tuple[int, str, bool]],
) -> tuple[dict[int, float], float]:
"""Extract per-line timing from traces."""
line_times: dict[int, float] = defaultdict(float)
total, value_depth, in_value, tactic_depth = 0.0, 0, False, None
name_re = re.compile(rf"\b{re.escape(name)}\b")
used: set[int] = set()
for depth, cls, ms, msg in traces:
if cls == "Elab.definition.value" and name_re.search(msg):
in_value, value_depth, total = True, depth, ms
elif cls == "Elab.async" and f"proof of {name}" in msg:
total = max(total, ms)
elif in_value:
if depth <= value_depth:
break
if cls == "Elab.step" and not msg.startswith("expected type:"):
tactic_depth = tactic_depth or depth
if depth == tactic_depth:
tactic = msg.split("\n")[0].strip().lstrip("·*- \t")
if ln := _match_line(tactic, not tactic, proof_items, used):
line_times[ln] += ms
used.add(ln)
return dict(line_times), total
def _filter_categories(cumulative: dict[str, float]) -> dict[str, float]:
"""Filter to relevant categories >= 1ms."""
return {
k: round(v, 1)
for k, v in sorted(cumulative.items(), key=lambda x: -x[1])
if k not in _SKIP_CATEGORIES and v >= 1.0
}
async def _run_lean_profile(file_path: Path, project_path: Path, timeout: float) -> str:
"""Run lean --profile, return output."""
proc = await asyncio.create_subprocess_exec(
"lake",
"env",
"lean",
"--profile",
"-Dtrace.profiler=true",
"-Dtrace.profiler.threshold=0",
str(file_path.resolve()),
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.STDOUT,
cwd=project_path.resolve(),
env=os.environ.copy(),
)
try:
stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout)
return stdout.decode("utf-8", errors="replace")
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
raise TimeoutError(f"Profiling timed out after {timeout}s")
def _find_proof_start(source_lines: list[str]) -> int:
"""Find line after ':= by' in source."""
for i, line in enumerate(source_lines):
if ":= by" in line or line.rstrip().endswith(" by"):
return i + 1
raise ValueError("No 'by' proof found in theorem")
async def profile_theorem(
file_path: Path,
theorem_line: int,
project_path: Path,
timeout: float = 60.0,
top_n: int = 5,
) -> ProofProfileResult:
"""Profile a theorem via `lean --profile`. Returns per-line timing data."""
lines = file_path.read_text().splitlines()
if not (0 < theorem_line <= len(lines)):
raise ValueError(f"Line {theorem_line} out of range")
source, name, src_start = _extract_theorem_source(lines, theorem_line)
source_lines = source.splitlines()
line_offset = theorem_line - src_start
proof_start = _find_proof_start(source_lines)
proof_items = _build_proof_items(source_lines, proof_start)
with tempfile.NamedTemporaryFile(
mode="w", suffix=".lean", dir=project_path, delete=False
) as f:
f.write(source)
temp_path = Path(f.name)
try:
output = await _run_lean_profile(temp_path, project_path, timeout)
finally:
temp_path.unlink(missing_ok=True)
traces, cumulative = _parse_output(output)
line_times, total = _extract_line_times(traces, name, proof_items)
top_lines = sorted(
[(ln, ms) for ln, ms in line_times.items() if ms >= total * 0.01],
key=lambda x: -x[1],
)[:top_n]
return ProofProfileResult(
ms=round(total, 1),
lines=[
LineProfile(
line=ln + line_offset,
ms=round(ms, 1),
text=source_lines[ln - 1].strip()[:60]
if ln <= len(source_lines)
else "",
)
for ln, ms in top_lines
],
categories=_filter_categories(cumulative),
)
