enforce matching Lean and LeanDojo versions

README.md

@@ -46,12 +46,15 @@ bash replace_files.sh
     - `ENCOUNTERED_THEOREMS_FILE`: Path for tracking encountered theorems
     - (Optional) `FISHER_DIR`: Directory for Fisher Information Matrices (FIMs)
 2. Adjust the options at the beginning of the `main()` function according to your requirements. The default options correspond to the LeanAgent configuration, and alternate options can be used to replicate ablation studies from the paper.
-3. Adjust the Lean toolchain paths in `generate_benchmark_lean4.py` if needed:
+3. Install the Lean toolchain required by each traced repository. LeanAgent reads
+   `ELAN_HOME` (or `~/.elan` by default), selects that exact toolchain, and
+   verifies it before tracing. For example:
 ```
-lean_dir2 = f"/.elan/toolchains/leanprover--lean4---{v}"
-lean_dir3 = f"~/.elan/toolchains/leanprover--lean4---{v}"
+elan toolchain install leanprover/lean4:v4.8.0
 ```
-These should match your system's Lean installation paths.
+   LeanAgent requires `lean_dojo==1.9.0`, which is pinned in
+   `requirements.txt`. Do not rely on the default `lean` on `PATH`: a different
+   Lean version makes the LeanDojo tracer fail to compile.
 
 ### Step 4: Install Models
 1. For ReProver's Tactic Generator:

generate_benchmark_lean4.py

@@ -1,4 +1,5 @@
 import json
+import os
 import shutil
 import random
 import networkx as nx
@@ -15,6 +16,7 @@
 import re
 import subprocess
 import sys
+from lean_toolchain import activate_lean_toolchain
 
 random.seed(3407)  # https://arxiv.org/abs/2109.08203
 
@@ -23,6 +25,7 @@
 SPLIT = Dict[SPLIT_NAME, List[TracedTheorem]]
 SPLIT_STRATEGY = str
 _LEAN4_VERSION_REGEX = re.compile(r"leanprover/lean4:(?P<version>.+?)")
+SUPPORTED_LEANDOJO_VERSION = "1.9.0"
 
 def get_lean4_version_from_config(toolchain: str) -> str:
     """Return the required Lean version given a ``lean-toolchain`` config."""
@@ -57,6 +60,20 @@ def is_supported_version(v) -> bool:
     else:
         return True
 
+
+def validate_tracing_environment(version: str) -> Path:
+    """Validate the Lean/LeanDojo pair before compiling LeanDojo's tracer."""
+    installed_leandojo_version = getattr(lean_dojo, "__version__", "unknown")
+    if installed_leandojo_version != SUPPORTED_LEANDOJO_VERSION:
+        raise RuntimeError(
+            "LeanAgent requires lean_dojo=="
+            f"{SUPPORTED_LEANDOJO_VERSION} for its supported Lean toolchains; "
+            f"found {installed_leandojo_version}. Reinstall it with: "
+            f"python -m pip install --upgrade --force-reinstall "
+            f"lean_dojo=={SUPPORTED_LEANDOJO_VERSION}"
+        )
+    return activate_lean_toolchain(version)
+
 def _split_sequentially(
     traced_theorems: List[TracedTheorem],
     num_val: int,
@@ -510,24 +527,18 @@ def main(url, commit, dst_dir):
     logger.info(f"lean toolchain version: {config}")
     v = get_lean4_version_from_config(config["content"])
     logger.info(f"lean version v: {v}")
-    logger.info(f"is supported: {is_supported_version(v)}")
-    if not is_supported_version(v):  # Won't get here since we checked for a compatible commit, but sanity check in case
-        logger.info("Unsupported version")
-    v = v[1:] # ignore "v" at beginning
-
-    lean_dir2 = f"/.elan/toolchains/leanprover--lean4---{v}"
-    lean_dir3 = f"~/.elan/toolchains/leanprover--lean4---{v}"
-    logger.info(f"lean path2 {lean_dir2}")
-    logger.info(f"lean path3 {lean_dir3}")
-    if not os.path.exists(lean_dir2):
-        logger.info(f"Lean toolchain path 2 does not exist: {lean_dir2}")
-    if not os.path.exists(lean_dir3):
-        logger.info(f"Lean toolchain path 3 does not exist: {lean_dir3}")
-    os.environ['LEAN4_PATH'] = lean_dir2
-    os.environ['PATH'] = f"{lean_dir2}/bin:{os.environ.get('PATH', '')}"
-    logger.info(f"Switched to Lean toolchain at: {lean_dir2}")
-
-    logger.info(f"lean --version: {subprocess.run(['lean', '--version'], capture_output=True).stdout.decode('utf-8')}")
+    if not is_supported_version(v):
+        logger.error(f"Unsupported Lean version: {v}")
+        return None, 0, 0, 10
+
+    try:
+        lean_binary = validate_tracing_environment(v)
+    except RuntimeError as exc:
+        logger.error(f"Cannot trace {repo}: {exc}")
+        return None, 0, 0, 10
+
+    logger.info(f"Switched to Lean toolchain at: {lean_binary.parent.parent}")
+    logger.info(f"lean --version: {subprocess.run([str(lean_binary), '--version'], capture_output=True, text=True).stdout.strip()}")
     logger.info(f"repo: {repo}")
 
     logger.info("Configuring LeanDojo again...")

lean_toolchain.py

@@ -0,0 +1,59 @@
+"""Resolve and activate the exact Lean toolchain required for tracing."""
+
+import os
+from pathlib import Path
+import re
+import subprocess
+
+
+_LEAN_VERSION_OUTPUT_REGEX = re.compile(r"Lean \(version (?P<version>[^,\s)]+)")
+
+
+def get_lean_toolchain_dir(version: str) -> Path:
+    """Return elan's installation directory for a Lean toolchain version."""
+    if not version.startswith("v"):
+        raise ValueError(f"Lean version must start with 'v', got {version!r}")
+
+    elan_home = Path(os.environ.get("ELAN_HOME", "~/.elan")).expanduser()
+    return elan_home / "toolchains" / f"leanprover--lean4---{version[1:]}"
+
+
+def activate_lean_toolchain(version: str) -> Path:
+    """Put the requested installed Lean toolchain first on ``PATH``.
+
+    LeanDojo compiles its tracer with the Lean executable available on ``PATH``.
+    Failing fast here prevents an installed, newer default Lean from being used when
+    the repository requires an older toolchain.
+    """
+    toolchain_dir = get_lean_toolchain_dir(version)
+    lean_binary = toolchain_dir / "bin" / "lean"
+    if not lean_binary.is_file():
+        raise RuntimeError(
+            f"Lean toolchain {version} is not installed at {toolchain_dir}. "
+            f"Install it with: elan toolchain install leanprover/lean4:{version}"
+        )
+
+    result = subprocess.run(
+        [str(lean_binary), "--version"],
+        capture_output=True,
+        check=False,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"Could not run Lean toolchain {version} at {lean_binary}: "
+            f"{result.stderr.strip()}"
+        )
+
+    match = _LEAN_VERSION_OUTPUT_REGEX.search(result.stdout)
+    actual_version = match["version"] if match else "unknown"
+    expected_version = version[1:]
+    if actual_version != expected_version:
+        raise RuntimeError(
+            f"Lean toolchain at {toolchain_dir} reports {actual_version}, "
+            f"but the repository requires {expected_version}."
+        )
+
+    os.environ["LEAN4_PATH"] = str(toolchain_dir)
+    os.environ["PATH"] = f"{lean_binary.parent}{os.pathsep}{os.environ.get('PATH', '')}"
+    return lean_binary

tests/test_lean_toolchain.py

@@ -0,0 +1,62 @@
+import os
+
+import pytest
+
+import lean_toolchain
+from lean_toolchain import activate_lean_toolchain, get_lean_toolchain_dir
+
+
+def test_get_lean_toolchain_dir_uses_elan_home(monkeypatch, tmp_path):
+    monkeypatch.setenv("ELAN_HOME", str(tmp_path / "elan"))
+
+    toolchain_dir = get_lean_toolchain_dir("v4.8.0")
+
+    assert toolchain_dir == (
+        tmp_path
+        / "elan"
+        / "toolchains"
+        / "leanprover--lean4---4.8.0"
+    )
+
+
+def test_get_lean_toolchain_dir_rejects_unprefixed_version():
+    with pytest.raises(ValueError, match="must start with 'v'"):
+        get_lean_toolchain_dir("4.8.0")
+
+
+def test_activate_lean_toolchain_uses_the_requested_binary(monkeypatch, tmp_path):
+    toolchain_dir = tmp_path / "elan" / "toolchains" / "leanprover--lean4---4.8.0"
+    lean_binary = toolchain_dir / "bin" / "lean"
+    lean_binary.parent.mkdir(parents=True)
+    lean_binary.touch()
+    monkeypatch.setenv("ELAN_HOME", str(tmp_path / "elan"))
+    monkeypatch.setenv("PATH", "existing-path")
+
+    class Result:
+        returncode = 0
+        stdout = "Lean (version 4.8.0, test)"
+        stderr = ""
+
+    monkeypatch.setattr(lean_toolchain.subprocess, "run", lambda *args, **kwargs: Result())
+
+    assert activate_lean_toolchain("v4.8.0") == lean_binary
+    assert os.environ["LEAN4_PATH"] == str(toolchain_dir)
+    assert os.environ["PATH"].split(os.pathsep) == [str(lean_binary.parent), "existing-path"]
+
+
+def test_activate_lean_toolchain_rejects_a_wrong_lean_version(monkeypatch, tmp_path):
+    toolchain_dir = tmp_path / "elan" / "toolchains" / "leanprover--lean4---4.8.0"
+    lean_binary = toolchain_dir / "bin" / "lean"
+    lean_binary.parent.mkdir(parents=True)
+    lean_binary.touch()
+    monkeypatch.setenv("ELAN_HOME", str(tmp_path / "elan"))
+
+    class Result:
+        returncode = 0
+        stdout = "Lean (version 4.9.0, test)"
+        stderr = ""
+
+    monkeypatch.setattr(lean_toolchain.subprocess, "run", lambda *args, **kwargs: Result())
+
+    with pytest.raises(RuntimeError, match="reports 4.9.0"):
+        activate_lean_toolchain("v4.8.0")