[WIP/experiment] formal proofs with aeneas/hax

.gitignore

@@ -9,3 +9,10 @@ executor/program_artifacts/
 # Shared cargo target directory for ELF builds
 executor/shared_target/
 
+
+# Lean/Lake build artifacts (formal proofs) at any depth
+**/.lake/
+
+# Python bytecode
+__pycache__/
+*.pyc

Makefile

@@ -1,7 +1,9 @@
 .PHONY: deps deps-linux deps-macos prepare-test-data compile-programs-asm compile-programs-rust compile-bench \
 compile-programs clean-asm clean-rust clean-bench clean-shared clean test test-asm test-no-compile \
 test-asm-no-compile test-rust test-rust-no-compile test-executor flamegraph-prover \
-test-fast test-prover test-prover-all test-disk-spill test-math-cuda test-cuda-integration bench-math-cuda bench-prover bench-prover-cuda build check clippy fmt lint
+test-fast test-prover test-prover-all test-disk-spill test-math-cuda test-cuda-integration bench-math-cuda bench-prover bench-prover-cuda build check clippy fmt lint \
+proofs proofs-charon proofs-aeneas proofs-hax proofs-check \
+proofs-lean-build proofs-lean-build-hax proofs-lean-build-aeneas
 
 UNAME := $(shell uname)
 
@@ -242,3 +244,212 @@ lint:
 
 flamegraph-prover:
 	cd crypto/stark && samply record cargo bench --bench profile_prover --features parallel
+
+# =============================================================================
+# Formal proofs — generated Lean code from Rust sources
+#
+# Layout:
+#   proofs/charon/   — LLBC files (charon output, one per crate)
+#   proofs/aeneas/   — Lean extraction (aeneas output, split-files per crate)
+#   proofs/hax/      — Lean extraction (hax output, per crate)
+#
+# Manually maintained (never overwritten by generation):
+#   proofs/aeneas/lakefile.toml, proofs/aeneas/lean-toolchain
+#   proofs/aeneas/{Math,Crypto,Executor}.lean   (static entry-point, 1 line each)
+#   proofs/hax/lakefile.toml, proofs/hax/lean-toolchain  (one lean_lib per crate)
+#
+# Targets: proofs (all), proofs-charon, proofs-aeneas, proofs-hax, proofs-check (CI).
+# =============================================================================
+
+PROOFS_DIR    := proofs
+CHARON_DIR    := $(PROOFS_DIR)/charon
+AENEAS_DIR    := $(PROOFS_DIR)/aeneas
+HAX_DIR       := $(PROOFS_DIR)/hax
+
+CHARON_MATH     := $(CHARON_DIR)/math.llbc
+CHARON_CRYPTO   := $(CHARON_DIR)/crypto.llbc
+CHARON_EXECUTOR := $(CHARON_DIR)/executor.llbc
+
+# Dedicated cargo target dir for charon. Charon only emits an .llbc when the crate
+# actually (re)compiles; sharing ./target with `cargo build` means a cache hit skips
+# the charon-driver and produces no output. An isolated target dir guarantees a clean
+# compile, hence deterministic emission (mirrors how hax uses ./target/hax).
+CHARON_TARGET_DIR := $(CURDIR)/target/charon
+
+proofs: proofs-charon proofs-aeneas proofs-hax
+
+# --- charon: Rust → LLBC ---
+
+$(CHARON_DIR):
+	mkdir -p $(CHARON_DIR)
+
+# `math`'s default features include `parallel` (rayon), whose `par_chunks_mut`
+# paths are not extractable (hax #420). Extract with `--no-default-features
+# --features alloc` to drop only the parallelism — the crate ships a real
+# sequential path for each, so this matches the hax invocation below and keeps
+# all three backends extracting the same code. `crypto`'s default (`asm`,`std`,
+# and `std` now pulls `alloc`) does NOT enable rayon (`crypto/parallel` is
+# opt-in), so it extracts under defaults with no flag change; `executor` has no
+# features to gate.
+$(CHARON_MATH): $(CHARON_DIR)
+	CARGO_TARGET_DIR=$(CHARON_TARGET_DIR) \
+		charon cargo --preset=aeneas --dest-file=$(CURDIR)/$(CHARON_MATH) -- -p math --lib --no-default-features --features alloc
+
+$(CHARON_CRYPTO): $(CHARON_DIR)
+	CARGO_TARGET_DIR=$(CHARON_TARGET_DIR) \
+		charon cargo --preset=aeneas --dest-file=$(CURDIR)/$(CHARON_CRYPTO) -- -p crypto --lib
+
+$(CHARON_EXECUTOR): $(CHARON_DIR)
+	CARGO_TARGET_DIR=$(CHARON_TARGET_DIR) \
+		charon cargo --preset=aeneas --dest-file=$(CURDIR)/$(CHARON_EXECUTOR) -- -p executor --lib
+
+proofs-charon: $(CHARON_MATH) $(CHARON_CRYPTO) $(CHARON_EXECUTOR)
+
+# --- aeneas: LLBC → Lean (split-files, one subdir per crate) ---
+# Entry-point files (Math.lean, Crypto.lean, Executor.lean) are static and not regenerated.
+# The _Template files are generated fresh each time; rename to FunsExternal.lean /
+# TypesExternal.lean once and fill them in — they will not be overwritten.
+
+$(AENEAS_DIR)/Math/Funs.lean: $(CHARON_MATH)
+	aeneas -backend lean $(CHARON_MATH) \
+		-dest $(AENEAS_DIR) -subdir Math -split-files; \
+	test -f $@
+	# Patch aeneas codegen defects in the generated Math/Types.lean (duplicate
+	# struct fields; mutually-recursive rand Rng/Fill -> opaque axioms) so the
+	# Lean compiles. Idempotent; fails loudly if aeneas output changes shape.
+	python3 $(CURDIR)/proofs/scripts/patch_aeneas_math_types.py $(AENEAS_DIR)/Math/Types.lean
+
+$(AENEAS_DIR)/Crypto/Funs.lean: $(CHARON_CRYPTO)
+	aeneas -backend lean $(CHARON_CRYPTO) \
+		-dest $(AENEAS_DIR) -subdir Crypto -split-files; \
+	test -f $@
+	# Patch the duplicate-field codegen defects in Crypto/{Types,Funs}.lean
+	# (same class as Math; see the script), then carve the self-contained,
+	# zero-maintenance single-leaf Merkle `Proof::verify` subset
+	# (Crypto/MerkleVerify.lean) by dependency closure — the full Crypto/Funs.lean
+	# does not compile (out-of-scope upstream-blocked code), but the carved subset
+	# does.
+	python3 $(CURDIR)/proofs/scripts/patch_aeneas_crypto_types.py $(AENEAS_DIR)/Crypto
+	python3 $(CURDIR)/proofs/scripts/carve_merkle_verify.py $(AENEAS_DIR)/Crypto
+
+$(AENEAS_DIR)/Executor/Funs.lean: $(CHARON_EXECUTOR)
+	aeneas -backend lean $(CHARON_EXECUTOR) \
+		-dest $(AENEAS_DIR) -subdir Executor -split-files; \
+	test -f $@
+
+proofs-aeneas: $(AENEAS_DIR)/Math/Funs.lean $(AENEAS_DIR)/Crypto/Funs.lean $(AENEAS_DIR)/Executor/Funs.lean
+
+# UPSTREAM-BLOCKED (#4): associated-type EQUALITY constraints.
+# `IsUnsignedInteger: Shr<usize, Output = Self> + BitAnd<Output = Self> + ...`
+# (crypto/math/src/unsigned_integer/traits.rs:6) puts `Output = Self` equalities
+# in supertrait bounds; `IsField`/`IsPrimeField` depend on it transitively
+# (field/traits.rs:232,254,269). hax (#1921) cannot emit these to Lean (the root
+# of the `IsField.AssociatedTypes`/`IsUnsignedInteger.AssociatedTypes` "unknown
+# identifier" failures in math.lean) and aeneas warns "Found an associated type
+# in a trait declaration ... Aeneas cannot handle such types today". Aeneas has
+# an internal `parameterize_trait_types` flag built for exactly this, but it is
+# NOT exposed as a CLI option, so it can't be enabled without rebuilding aeneas
+# from OCaml source. These traits are kept opaque/external; no production change.
+
+# --- hax: Rust → Lean (annotated items only) ---
+# hax writes one <crate>.lean per crate directly into HAX_DIR, all covered by the
+# single hand-maintained $(HAX_DIR)/lakefile.toml (one lean_lib per crate).
+# Note: only the math crate is currently targeted. math.lean emits but does NOT
+# yet fully compile: remaining blockers are #4 assoc-type equality (HAX0001, see
+# above), the sequential-FFT `&mut input[range]` slice borrows in bowers_fft.rs
+# (HAX0003/HAX0010, hax #420), and missing `core_models.*` stdlib models in the
+# hax Lean proof-lib. crypto/executor are not yet extractable.
+#
+# We extract with `--no-default-features --features alloc` (i.e. WITHOUT `parallel`).
+# The rayon `par_chunks_mut` paths (FFT, batch-inverse) are gated behind the
+# `parallel` feature and are NOT extractable (hax issue #420). Disabling the
+# feature removes them from extraction; the crate ships a real, separately
+# compiled sequential path for each (the same code used by no-parallel/wasm
+# builds), so this excludes only parallelism, not the verified computation —
+# e.g. `inplace_batch_inverse` extracts via its sequential form, which is the
+# exact code the verifier runs at its (sub-threshold) slice sizes.
+
+proofs-hax:
+	mkdir -p $(HAX_DIR)
+	cargo hax -C '-p' 'math' '--no-default-features' '--features' 'alloc' ';' into \
+		--output-dir $(CURDIR)/$(HAX_DIR) lean; \
+	test -f $(HAX_DIR)/math.lean
+	# Inject `import CoreModelsSupplement` so math.lean sees our opaque stubs for
+	# core_models.* stdlib models missing from the pinned Hax proof-lib.
+	python3 $(CURDIR)/proofs/scripts/patch_hax_math.py $(HAX_DIR)/math.lean
+	# crypto: extract ONLY the in-scope Merkle items via hax's native item
+	# selection `-i` (glob + `+` transitive-dependency closure), rather than
+	# gating the source with cfg(hax). `+...::verify` pulls in Proof::verify and
+	# its dependency closure (the Proof struct + IsMerkleTreeBackend trait);
+	# everything else (poseidon, transcript, batch, concrete backends) is left
+	# out. See proofs/hax/HAX_INCLUDE below.
+	cargo hax -C '-p' 'crypto' ';' into \
+		-i '-** +crypto::merkle_tree::proof::**::verify' \
+		--output-dir $(CURDIR)/$(HAX_DIR) lean; \
+	test -f $(HAX_DIR)/crypto.lean
+	python3 $(CURDIR)/proofs/scripts/patch_hax_math.py $(HAX_DIR)/crypto.lean
+
+# --- CI check: regenerate into a temp dir and diff against committed output ---
+
+# Regenerates everything into a throwaway dir and diffs against the committed Lean.
+# _Template files are excluded: they are advisory starting points for the manually
+# maintained *External.lean files, and aeneas may skip them on partial extraction.
+# MerkleVerifyProofs.lean is excluded too: it is the hand-written proof file (not
+# produced by extraction), so it must not count as "stale generated output".
+proofs-check:
+	$(eval TMPDIR := $(shell mktemp -d))
+	@trap 'rm -rf $(TMPDIR)' EXIT; \
+	mkdir -p $(TMPDIR)/charon $(TMPDIR)/aeneas $(TMPDIR)/hax; \
+	CARGO_TARGET_DIR=$(TMPDIR)/target charon cargo --preset=aeneas --dest-file=$(TMPDIR)/charon/math.llbc -- -p math --lib --no-default-features --features alloc; \
+	CARGO_TARGET_DIR=$(TMPDIR)/target charon cargo --preset=aeneas --dest-file=$(TMPDIR)/charon/crypto.llbc -- -p crypto --lib; \
+	CARGO_TARGET_DIR=$(TMPDIR)/target charon cargo --preset=aeneas --dest-file=$(TMPDIR)/charon/executor.llbc -- -p executor --lib; \
+	aeneas -backend lean $(TMPDIR)/charon/math.llbc -dest $(TMPDIR)/aeneas -subdir Math -split-files; true; \
+	python3 $(CURDIR)/proofs/scripts/patch_aeneas_math_types.py $(TMPDIR)/aeneas/Math/Types.lean; \
+	aeneas -backend lean $(TMPDIR)/charon/crypto.llbc -dest $(TMPDIR)/aeneas -subdir Crypto -split-files; true; \
+	python3 $(CURDIR)/proofs/scripts/patch_aeneas_crypto_types.py $(TMPDIR)/aeneas/Crypto; \
+	python3 $(CURDIR)/proofs/scripts/carve_merkle_verify.py $(TMPDIR)/aeneas/Crypto; \
+	aeneas -backend lean $(TMPDIR)/charon/executor.llbc -dest $(TMPDIR)/aeneas -subdir Executor -split-files; true; \
+	cargo hax -C '-p' 'math' '--no-default-features' '--features' 'alloc' ';' into --output-dir $(TMPDIR)/hax lean; \
+	python3 $(CURDIR)/proofs/scripts/patch_hax_math.py $(TMPDIR)/hax/math.lean; \
+	diff -r --brief --exclude="*_Template.lean" --exclude="*External.lean" \
+		$(TMPDIR)/aeneas/Math     $(AENEAS_DIR)/Math     || { echo "FAIL: Math aeneas output is stale"; exit 1; }; \
+	diff -r --brief --exclude="*_Template.lean" --exclude="*External.lean" --exclude="MerkleVerifyProofs.lean" \
+		$(TMPDIR)/aeneas/Crypto   $(AENEAS_DIR)/Crypto   || { echo "FAIL: Crypto aeneas output is stale"; exit 1; }; \
+	diff -r --brief --exclude="*_Template.lean" \
+		$(TMPDIR)/aeneas/Executor $(AENEAS_DIR)/Executor || { echo "FAIL: Executor aeneas output is stale"; exit 1; }; \
+	diff --brief \
+		$(TMPDIR)/hax/math.lean   $(HAX_DIR)/math.lean   || { echo "FAIL: hax math output is stale"; exit 1; }; \
+	echo "proofs-check: all generated files are up to date"
+	# Compile the buildable extracted Lean (aeneas Math). Catches a regression
+	# where the committed Lean is up-to-date but no longer well-formed. hax is
+	# excluded here — it is #4-blocked (run `make proofs-lean-build-hax`).
+	$(MAKE) proofs-lean-build
+
+# --- lake build: compile the extracted Lean ---
+# Extraction emitting a file is NOT proof the Lean is well-formed; a partial or
+# garbage extraction can still write output. `lake build` is the real check that
+# the generated definitions and the hand-maintained *External.lean stubs compile
+# and that the opaque (#4 trait-bound / #6 aeneas) externals line up.
+# Each subdir is an independent lake project (its own lakefile.toml +
+# lean-toolchain + manifest, pulling Hax / aeneas backends from git).
+#
+# `proofs-lean-build` aggregates ONLY the targets that currently compile, so it
+# is a green gate suitable for `proofs-check`/CI. The hax build is BLOCKED on the
+# upstream assoc-type-equality limit (#4; see the comment above proofs-aeneas)
+# and is kept as a separately-runnable target, NOT in the aggregate, so a known
+# upstream gap doesn't turn the whole proofs flow red.
+proofs-lean-build: proofs-lean-build-aeneas
+
+# Runnable but expected to FAIL on #4 (IsField/IsUnsignedInteger assoc-type
+# equality). All non-#4 errors are resolved (asm, core_models, iterators).
+proofs-lean-build-hax:
+	cd $(HAX_DIR) && lake build
+
+proofs-lean-build-aeneas:
+	# `Math`: full crate. `MerkleVerify`: the carved single-leaf Merkle
+	# `Proof::verify` subset (the rest of `Crypto`, and `Executor`, are not
+	# buildable — out-of-scope upstream-blocked code: assoc-type equality #4,
+	# poseidon/transcript/batch codegen limits). `MerkleVerifyProofs`: the
+	# hand-written panic-freedom, index-algebra and completeness proofs about the
+	# carved defs — built here so a proof regression turns the green gate red.
+	cd $(AENEAS_DIR) && lake build Math MerkleVerify MerkleVerifyProofs

crypto/crypto/Cargo.toml

@@ -33,8 +33,18 @@ bincode = "1"
 [features]
 default = ["asm", "std"]
 asm = ["sha3/asm"]
-std = ["math/std", "sha3/std", "serde?/std"]
+std = ["alloc", "math/std", "sha3/std", "serde?/std"]
 serde = ["dep:serde"]
 parallel = ["dep:rayon"]
 disk-spill = ["std", "dep:memmap2", "dep:tempfile", "dep:libc"]
-alloc = []
+alloc = []
+
+[dependencies.hax-lib]
+git = "https://github.com/hacspec/hax"
+version = "0.3.7"
+# hax-lib's requires/ensures/attributes proc-macros are no-ops in native builds;
+# kept as a normal (not cfg(hax)-gated) dep so contract annotations compile both
+# natively and under hax extraction (cf. hax's own examples).
+
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(hax)'] }

crypto/crypto/src/merkle_tree/merkle.rs

@@ -172,6 +172,8 @@ where
     ///
     /// This skips the `hash_leaves` step, useful when leaves have already been
     /// hashed externally (e.g., to avoid materializing large intermediate data).
+    #[cfg_attr(hax, hax_lib::requires(hashed_leaves.len().is_power_of_two()))]
+    #[cfg_attr(hax, hax_lib::ensures(|res| true))]
     pub fn build_from_hashed_leaves(hashed_leaves: Vec<B::Node>) -> Option<Self> {
         if hashed_leaves.is_empty() {
             return None;

crypto/crypto/src/merkle_tree/proof.rs

@@ -1,6 +1,4 @@
 use alloc::{collections::BTreeMap, vec::Vec};
-#[cfg(feature = "alloc")]
-use math::traits::Serializable;
 use math::{errors::DeserializationError, traits::Deserializable};
 
 use super::{
@@ -19,15 +17,30 @@ pub struct Proof<T: PartialEq + Eq> {
     pub merkle_path: Vec<T>,
 }
 
+#[hax_lib::attributes]
 impl<T: PartialEq + Eq> Proof<T> {
     /// Verifies a Merkle inclusion proof for the value contained at leaf index.
+    // Panic-freedom contract: with no precondition, `verify` always returns a
+    // value (no panic, no divergence). The loop is a bounded fold over the
+    // finite `merkle_path`; the only fallible ops (`index % 2`, `index >> 1`)
+    // are total for the constant operands. Proved in Lean via the generated
+    // `Impl.verify` spec (see proofs/hax). hax-lib macros are no-ops in native
+    // builds, so these annotations are unconditional (cf. hax examples).
+    #[hax_lib::requires(true)]
+    #[hax_lib::ensures(|_res| true)]
+    #[hax_lib::lean::proof_method::grind]
     pub fn verify<B>(&self, root_hash: &B::Node, mut index: usize, value: &B::Data) -> bool
     where
         B: IsMerkleTreeBackend<Node = T>,
     {
         let mut hashed_value = B::hash_data(value);
 
-        for sibling_node in self.merkle_path.iter() {
+        // Range loop (not `.iter()`) so the Rust->Lean extractors model it as a
+        // bounded `fold_range` rather than an `Iterator::fold` over a slice
+        // iterator, which the hax Lean proof-lib does not support. Equivalent:
+        // visits each `merkle_path` node once, in order.
+        for i in 0..self.merkle_path.len() {
+            let sibling_node = &self.merkle_path[i];
             if index.is_multiple_of(2) {
                 hashed_value = B::hash_new_parent(&hashed_value, sibling_node);
             } else {
@@ -41,19 +54,6 @@ impl<T: PartialEq + Eq> Proof<T> {
     }
 }
 
-#[cfg(feature = "alloc")]
-impl<T> Serializable for Proof<T>
-where
-    T: Serializable + PartialEq + Eq,
-{
-    fn serialize(&self) -> Vec<u8> {
-        self.merkle_path
-            .iter()
-            .flat_map(|node| node.serialize())
-            .collect()
-    }
-}
-
 impl<T> Deserializable for Proof<T>
 where
     T: Deserializable + PartialEq + Eq,