Preserve Lake declarations and traces when vendoring

anneal/rewrite-lake-vendor.py

@@ -13,6 +13,26 @@
 REMOTE_TOML_KEYS = {"git", "url", "rev", "scope", "subDir", "inputRev"}
 
 
+def validate_vendoring_roots(root: Path, packages_dir: Path) -> None:
+    """Reject links that could route a private rewrite into another installation."""
+    roots = [root.absolute(), packages_dir.absolute()]
+    for directory in roots:
+        if directory.is_symlink() or not directory.is_dir():
+            raise ValueError(f"Vendoring root must be a physical directory: {directory}")
+    allowed = [directory.resolve() for directory in roots]
+    def fail_walk(error: OSError) -> None:
+        raise error
+
+    for directory in roots:
+        for current, directories, files in os.walk(directory, followlinks=False, onerror=fail_walk):
+            for name in directories + files:
+                path = Path(current) / name
+                if path.is_symlink():
+                    target = path.resolve(strict=True)
+                    if not any(target.is_relative_to(parent) for parent in allowed):
+                        raise ValueError(f"Vendoring symlink leaves private roots: {path} -> {target}")
+
+
 def quote_name(name: str) -> str:
     if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_']*", name):
         return name
@@ -34,7 +54,9 @@ def package_dirs(packages_dir: Path) -> dict[str, Path]:
 
 def replace_require_toml_blocks(path: Path, packages: dict[str, Path]) -> bool:
     content = path.read_text()
-    block_re = re.compile(r"(?ms)^\s*\[\[require\]\]\s*\n(?:(?!^\s*\[\[).*\n?)*")
+    # Stop at every TOML table, including single-bracket option tables. A
+    # greedy dot with DOTALL would also consume later library/executable tables.
+    block_re = re.compile(r"(?ms)^[ \t]*\[\[require\]\][ \t]*\n.*?(?=^[ \t]*\[|\Z)")
 
     changed = False
 
@@ -151,7 +173,7 @@ def rewrite_trace_prefixes(
     prefixes = trace_prefixes(root, packages_dir, packages, extra_prefixes)
     package_root_patterns = [
         re.compile(
-            rf"(?<![A-Za-z0-9_.-])(?:/[^\s\"':]+)+/(?:packages|\.lake/packages)/{re.escape(name)}/"
+            rf"(?<![A-Za-z0-9_.-])(?:/[^/\s\"':]+)+/(?:packages|\.lake/packages)/{re.escape(name)}/"
         )
         for name in packages
     ]
@@ -215,6 +237,7 @@ def main() -> None:
     )
     args = parser.parse_args()
 
+    validate_vendoring_roots(args.root, args.packages_dir)
     root = args.root.resolve()
     packages_dir = args.packages_dir.resolve()
     packages = package_dirs(packages_dir)

anneal/tests/fixtures/lake-vendor/README.md

@@ -0,0 +1,10 @@
+These unmodified Lake configurations are regression inputs for vendoring, not
+dependency pins. They come from the dependency revisions in Anneal's pinned
+Aeneas archive:
+
+- `aesop.toml`: aesop `e3cb2f741431ce31bf73549fb52316a57368b06f`.
+- `importGraph.toml`: import-graph `5c7542ed018c78194f1e2b903eaf6a792b74c03d`.
+
+Both projects distribute these files under Apache-2.0; see the repository's
+`LICENSE-APACHE`. The regression checks that replacing a remote dependency does
+not remove the library and executable declarations that follow it.

anneal/tests/fixtures/lake-vendor/aesop.toml

@@ -0,0 +1,18 @@
+name = "aesop"
+defaultTargets = ["Aesop"]
+testDriver = "AesopTest"
+precompileModules = false # We would like to turn this on, but it breaks the Mathlib cache.
+platformIndependent = true
+
+[[require]]
+name = "batteries"
+scope = "leanprover-community"
+rev = "v4.31.0"
+
+[[lean_lib]]
+name = "Aesop"
+
+[[lean_lib]]
+name = "AesopTest"
+globs = ["AesopTest.+"]
+leanOptions = {linter.unusedVariables = false}

anneal/tests/fixtures/lake-vendor/importGraph.toml

@@ -0,0 +1,33 @@
+name = "importGraph"
+description = "Tools to analyse and visualise the import structure of Lean packages and their files."
+license = "Apache-2.0"
+licenseFiles = ["LICENSE"]
+
+defaultTargets = ["ImportGraph", "graph"]
+testRunner = "ImportGraphTest"
+
+[[require]]
+name = "Cli"
+scope = "leanprover"
+rev = "v4.31.0"
+
+[[lean_lib]]
+name = "ImportGraph"
+
+# `lake exe graph` constructs import graphs in `.dot` or graphical formats.
+[[lean_exe]]
+name = "graph"
+root = "MainGraph"
+# Enables the use of the Lean interpreter by the executable (e.g., `runFrontend`)
+# at the expense of increased binary size on Linux.
+# Remove this line if you do not need such functionality.
+supportInterpreter = true
+
+# `lake exe unused_transitive_imports` prints unused transitive imports from amongst a given list of modules.
+[[lean_exe]]
+name = "unused_transitive_imports"
+root = "MainUnusedTransitiveImports"
+supportInterpreter = true
+
+[[lean_lib]]
+name = "ImportGraphTest"

anneal/tests/test_rewrite_lake_vendor.py

@@ -14,6 +14,8 @@
 
 import importlib.util
 import tempfile
+import subprocess
+import sys
 import unittest
 from pathlib import Path
 
@@ -31,6 +33,93 @@ def write(path: Path, contents: str = "") -> None:
 
 
 class RewriteLakeVendorTests(unittest.TestCase):
+    def test_toml_require_rewrite_preserves_following_tables(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            lakefile = root / "lakefile.toml"
+            following = ('[[lean_lib]]\nname = "Library"\n\n'
+                         '[[lean_exe]]\nname = "driver"\nroot = "Main"\n'
+                         '[leanOptions]\nautoImplicit = false\n')
+            write(lakefile, 'name = "consumer"\n\n[[require]]\nname = "dep"\n'
+                  'scope = "owner"\nrev = "main"\n\n' + following)
+            self.assertTrue(rewrite_lake_vendor.replace_require_toml_blocks(
+                lakefile, {"dep": root / "packages/dep"}))
+            rewritten = lakefile.read_text()
+            self.assertTrue(rewritten.endswith(following))
+            self.assertIn('path = "packages/dep"', rewritten)
+            self.assertNotIn('rev =', rewritten)
+            self.assertFalse(rewrite_lake_vendor.replace_require_toml_blocks(
+                lakefile, {"dep": root / "packages/dep"}))
+
+    def test_toml_require_stops_before_single_bracket_table(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / 'lakefile.toml'
+            following = '[leanOptions]\nautoImplicit = false\n'
+            write(path, '[[require]]\nname = "dep"\nrev = "main"\n' + following)
+            rewrite_lake_vendor.replace_require_toml_blocks(path, {"dep": Path(tmp) / 'dep'})
+            self.assertTrue(path.read_text().endswith(following))
+
+    def test_pinned_aesop_and_importgraph_keep_all_declared_targets(self) -> None:
+        fixtures = SCRIPT.parent / "tests/fixtures/lake-vendor"
+        for name, dependency in [("aesop", "batteries"), ("importGraph", "Cli")]:
+            with self.subTest(package=name), tempfile.TemporaryDirectory() as tmp:
+                path = Path(tmp) / "lakefile.toml"
+                original = (fixtures / (name + ".toml")).read_text()
+                write(path, original)
+                following = original[original.index('[[lean_lib]]'):]
+                self.assertTrue(rewrite_lake_vendor.replace_require_toml_blocks(
+                    path, {dependency: Path(tmp) / dependency}))
+                self.assertTrue(path.read_text().endswith(following))
+                self.assertIn('path = "' + dependency + '"', path.read_text())
+
+    def test_external_symlink_rejected_before_any_rewrite(self) -> None:
+        for entry in [".lake", "lakefile.lean"]:
+            with self.subTest(entry=entry), tempfile.TemporaryDirectory() as tmp:
+                base = Path(tmp)
+                root, packages, shared = base / "private", base / "packages", base / "shared"
+                root.mkdir()
+                packages.mkdir()
+                shared.mkdir()
+                original = 'require dep from git "https://example.invalid/dep" @ "rev"\n'
+                write(root / "lakefile.lean", original)
+                write(shared / "lakefile.lean", original)
+                if entry == ".lake":
+                    (root / entry).symlink_to(shared, target_is_directory=True)
+                else:
+                    (root / entry).unlink()
+                    (root / entry).symlink_to(shared / entry)
+                (packages / "dep").mkdir()
+                result = subprocess.run([sys.executable, str(SCRIPT), "--root", str(root),
+                                         "--packages-dir", str(packages)],
+                                        capture_output=True, text=True)
+                self.assertNotEqual(result.returncode, 0)
+                self.assertIn("symlink leaves private roots", result.stderr)
+                self.assertEqual((shared / "lakefile.lean").read_text(), original)
+                self.assertEqual((root / "lakefile.lean").read_text(), original)
+
+    def test_internal_payload_links_are_allowed(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root, packages = Path(tmp) / "private", Path(tmp) / "packages"
+            root.mkdir()
+            packages.mkdir()
+            write(root / "real.lean", "def value := 1\n")
+            (root / "linked.lean").symlink_to("real.lean")
+            rewrite_lake_vendor.validate_vendoring_roots(root, packages)
+
+    def test_long_unmatched_trace_path_does_not_backtrack_exponentially(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root, packages = Path(tmp) / 'root', Path(tmp) / 'packages'
+            root.mkdir()
+            (packages / 'mathlib').mkdir(parents=True)
+            unrelated = '/' + '/'.join(['longsegment'] * 80) + '/Unrelated.lean'
+            trace = root / 'Module.trace'
+            write(trace, unrelated + '\n/upstream/packages/mathlib/Mathlib.lean\n')
+            result = subprocess.run([sys.executable, str(SCRIPT), '--root', str(root),
+                                     '--packages-dir', str(packages), '--rewrite-traces'],
+                                    capture_output=True, text=True, timeout=5)
+            self.assertEqual(result.returncode, 0, result.stderr)
+            self.assertEqual(trace.read_text(), unrelated + '\nMathlib.lean\n')
+
     def test_rewrites_upstream_aeneas_backend_trace_prefix(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp) / "aeneas" / "backends" / "lean"