[anneal] Fix type ambiguity in spec generation

When a Rust function uses a type via use import, Anneal's spec generator only sees the short name from the syn AST. For names like Ordering, that can conflict with Lean's native Ordering brought in by open Aeneas.Std and cause type mismatches in generated Pre/Post structs.

After Aeneas generates Funs.lean, parse its def signatures and use the fully-qualified Aeneas parameter and return types as overrides in spec generation. Keep escaped Lean parameter names, such as show1, in the same namespace as the generated specs. Fall back to map_type when no Aeneas signature is available.

Author: nnunley

anneal/src/aeneas.rs

@@ -346,8 +346,10 @@ pub fn generate_lean_workspace(roots: &LockedRoots, artifacts: &[AnnealArtifact]
         let slug = artifact.artifact_slug();
         let output_dir = lean_generated_root.join(&slug);
 
+        let funs_types = parse_funs_types_in_dir(&output_dir)?;
+
         // Generate Anneal specs
-        let generated = generate::generate_artifact(artifact);
+        let generated = generate::generate_artifact_with_funs_types(artifact, &funs_types);
         let specs_path = output_dir.join(artifact.lean_spec_file_name());
         let map_path = output_dir.join(format!("{}.lean.map", artifact.artifact_slug()));
 
@@ -591,6 +593,17 @@ pub fn generate_lean_workspace(roots: &LockedRoots, artifacts: &[AnnealArtifact]
     Ok(())
 }
 
+pub(crate) fn parse_funs_types_in_dir(output_dir: &Path) -> Result<crate::funs_types::FunsTypeMap> {
+    let funs_path = output_dir.join("Funs.lean");
+    if !funs_path.exists() {
+        return Ok(crate::funs_types::FunsTypeMap::new());
+    }
+
+    let content = std::fs::read_to_string(&funs_path)
+        .with_context(|| format!("Failed to read {}", funs_path.display()))?;
+    Ok(crate::funs_types::parse_funs_types(&content))
+}
+
 /// Completes Lean verification by generating Anneal `Specs.lean`, writing `Generated.lean`,
 /// and running `lake build` + diagnostics.
 pub fn verify_lean_workspace(roots: &LockedRoots, artifacts: &[AnnealArtifact]) -> Result<()> {

anneal/src/funs_types.rs

@@ -0,0 +1,211 @@
+//! Parses Aeneas-generated `Funs.lean` signatures.
+//!
+//! Anneal's spec generator re-derives types from the Rust AST, losing
+//! qualification for `use`-imported names (e.g., `Ordering` instead of
+//! `core.sync.atomic.Ordering`). This module provides a lookup table from the
+//! Aeneas output as a corrective.
+
+use std::collections::HashMap;
+
+/// A function signature parsed from Aeneas-generated Lean.
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct FunsSignature {
+    pub params: Vec<(String, String)>,
+    pub ret: Option<String>,
+}
+
+/// Function name → Aeneas-generated function signature.
+pub type FunsTypeMap = HashMap<String, FunsSignature>;
+
+/// Parses `def` signatures from `Funs.lean`.
+///
+/// This extracts explicit parameter types and return types while skipping
+/// implicit `{}` parameters. Malformed signatures are ignored.
+pub fn parse_funs_types(content: &str) -> FunsTypeMap {
+    let mut map = FunsTypeMap::new();
+    let lines: Vec<&str> = content.lines().collect();
+
+    for i in 0..lines.len() {
+        let Some(rest) = lines[i].trim().strip_prefix("def ") else {
+            continue;
+        };
+        let name = rest.split([' ', '(', '{', ':']).next().unwrap_or("");
+        if name.is_empty() {
+            continue;
+        }
+
+        // Collect signature lines until `:=`.
+        let mut sig = String::from(rest);
+        for j in (i + 1)..lines.len() {
+            if sig.contains(":=") {
+                break;
+            }
+            sig.push(' ');
+            sig.push_str(lines[j].trim());
+        }
+
+        let params = extract_params(&sig);
+        let ret = extract_return_type(&sig);
+        if !params.is_empty() || ret.is_some() {
+            map.insert(name.to_string(), FunsSignature { params, ret });
+        }
+    }
+    map
+}
+
+/// Extracts `(name : type)` bindings, skipping `{implicit}` params.
+fn extract_params(sig: &str) -> Vec<(String, String)> {
+    let mut params = Vec::new();
+    let mut chars = sig.chars().peekable();
+
+    while let Some(&c) = chars.peek() {
+        match c {
+            '(' => {
+                chars.next();
+                if let Some(pair) = parse_binding(&collect_delimited(&mut chars, ')')) {
+                    params.push(pair);
+                }
+            }
+            '{' => {
+                chars.next();
+                collect_delimited(&mut chars, '}');
+            }
+            ':' => break,
+            _ => {
+                chars.next();
+            }
+        }
+    }
+    params
+}
+
+/// Extracts the function return type.
+fn extract_return_type(sig: &str) -> Option<String> {
+    let return_start = top_level_return_colon(sig)? + 1;
+    let return_end = sig[return_start..].find(":=").map(|i| return_start + i).unwrap_or(sig.len());
+    let ret = sig[return_start..return_end].trim();
+    if ret.is_empty() {
+        return None;
+    }
+    Some(ret.strip_prefix("Result ").unwrap_or(ret).trim().to_string())
+}
+
+/// Finds the colon that separates the parameter list from the return type.
+fn top_level_return_colon(sig: &str) -> Option<usize> {
+    let mut paren_depth = 0u32;
+    let mut brace_depth = 0u32;
+    for (i, c) in sig.char_indices() {
+        match c {
+            '(' => paren_depth += 1,
+            ')' => paren_depth = paren_depth.saturating_sub(1),
+            '{' => brace_depth += 1,
+            '}' => brace_depth = brace_depth.saturating_sub(1),
+            ':' if paren_depth == 0 && brace_depth == 0 => return Some(i),
+            _ => {}
+        }
+    }
+    None
+}
+
+/// Reads chars until the matching `close` delimiter, handling nesting.
+fn collect_delimited(chars: &mut std::iter::Peekable<std::str::Chars<'_>>, close: char) -> String {
+    let open = if close == ')' { '(' } else { '{' };
+    let mut depth = 1u32;
+    let mut buf = String::new();
+    for c in chars.by_ref() {
+        if c == open {
+            depth += 1;
+        } else if c == close {
+            depth -= 1;
+            if depth == 0 {
+                return buf;
+            }
+        }
+        buf.push(c);
+    }
+    buf
+}
+
+/// Splits `"name : type"` on the first ` : `.
+fn parse_binding(s: &str) -> Option<(String, String)> {
+    let s = s.trim();
+    let i = s.find(" : ")?;
+    let (name, ty) = (s[..i].trim(), s[i + 3..].trim());
+    (!name.is_empty() && !ty.is_empty()).then_some((name.to_string(), ty.to_string()))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn simple() {
+        let m = parse_funs_types("def hash_key (k : Std.Usize) : Result Std.Usize := do\n");
+        assert_eq!(m["hash_key"].params, [("k".into(), "Std.Usize".into())]);
+        assert_eq!(m["hash_key"].ret.as_deref(), Some("Std.Usize"));
+    }
+
+    #[test]
+    fn multiline() {
+        let m = parse_funs_types(
+            "def frame.AtomicFrameState.load\n\
+             \x20 (self : frame.AtomicFrameState) (order : core.sync.atomic.Ordering) :\n\
+             \x20 Result frame.FrameState\n\
+             \x20 := do\n",
+        );
+        let p = &m["frame.AtomicFrameState.load"].params;
+        assert_eq!(p[0], ("self".into(), "frame.AtomicFrameState".into()));
+        assert_eq!(p[1], ("order".into(), "core.sync.atomic.Ordering".into()));
+        assert_eq!(m["frame.AtomicFrameState.load"].ret.as_deref(), Some("frame.FrameState"));
+    }
+
+    #[test]
+    fn skips_implicits() {
+        let m = parse_funs_types(
+            "def HashMap.alloc {T : Type} (slots : alloc.vec.Vec T) (n : Std.Usize) :\n\
+             \x20 Result Unit := do\n",
+        );
+        let p = &m["HashMap.alloc"].params;
+        assert_eq!(p.len(), 2);
+        assert_eq!(p[0].0, "slots");
+        assert_eq!(p[1], ("n".into(), "Std.Usize".into()));
+    }
+
+    #[test]
+    fn multiple_ordering_params() {
+        let m = parse_funs_types(
+            "def f.compare_exchange\n\
+             \x20 (self : f.T) (expected : f.S)\n\
+             \x20 (success : core.sync.atomic.Ordering)\n\
+             \x20 (failure : core.sync.atomic.Ordering) :\n\
+             \x20 Result Unit := do\n",
+        );
+        let p = &m["f.compare_exchange"].params;
+        assert_eq!(p.len(), 4);
+        assert_eq!(p[2].1, "core.sync.atomic.Ordering");
+        assert_eq!(p[3].1, "core.sync.atomic.Ordering");
+    }
+
+    #[test]
+    fn preserves_escaped_keyword_params() {
+        let m = parse_funs_types("def f (show1 : core.sync.atomic.Ordering) : Result Unit := do\n");
+        assert_eq!(m["f"].params, [("show1".into(), "core.sync.atomic.Ordering".into())]);
+    }
+
+    #[test]
+    fn parses_return_type() {
+        let m = parse_funs_types(
+            "def f (order : core.sync.atomic.Ordering) :\n\
+             \x20 Result core.sync.atomic.Ordering\n\
+             \x20 := do\n",
+        );
+        assert_eq!(m["f"].ret.as_deref(), Some("core.sync.atomic.Ordering"));
+    }
+
+    #[test]
+    fn trait_instance_no_params() {
+        let m = parse_funs_types("def Foo.Clone : core.clone.Clone Foo := {\n  clone := x\n}\n");
+        assert!(m["Foo.Clone"].params.is_empty());
+        assert_eq!(m["Foo.Clone"].ret.as_deref(), Some("core.clone.Clone Foo"));
+    }
+}