aver-lang 0.26.0

use std::collections::{BTreeSet, HashMap, HashSet};

use crate::ast::{
    Expr, FnDef, MatchArm, Spanned, Stmt, StrPart, TailCallData, TopLevel, VerifyBlock, VerifyKind,
    VerifyLaw,
};
use crate::types::Type;

pub mod expand;

/// Legacy `name → (params, return_type, effects)` map carrying the fn /
/// ctor signatures the verify-law helpers query.
///
/// Kept as a type alias for callers (mostly tests) that still build the
/// map literally; production code reaches the verify-law surface
/// through [`FnSigOracle`] now, which `FnSigMap` also implements. Epic
/// #180 Phase 7 closed the side-channel storage path — see PRs #193 /
/// #194 / this PR's history.
pub type FnSigMap = HashMap<String, (Vec<Type>, Type, Vec<String>)>;

/// What every verify-law helper needs to know about a fn or ctor
/// reachable from a `verify ... law ...` block: its return type and
/// whether it has any declared effects (purity). Cheap to derive at
/// the call site — either from the legacy [`FnSigMap`] or directly
/// from a [`CodegenContext`]'s resolved view.
#[derive(Debug, Clone)]
pub struct FnSigInfo {
    pub return_type: Type,
    pub is_pure: bool,
}

/// Look up a fn / ctor signature by source-level name.
///
/// Epic #180 Phase 7 closing refactor: verify-law helpers used to take
/// `&FnSigMap` directly, which forced every caller to materialise the
/// map on demand. Switching to a trait lets a
/// [`crate::codegen::CodegenContext`] pass itself as the oracle — the
/// resolved-program view answers per-key queries without rebuilding
/// the map up front. `FnSigMap` keeps its `impl FnSigOracle` for
/// backward compat with tests that prefer to build a literal sig map.
pub trait FnSigOracle {
    fn fn_sig(&self, name: &str) -> Option<FnSigInfo>;
}

impl FnSigOracle for FnSigMap {
    fn fn_sig(&self, name: &str) -> Option<FnSigInfo> {
        let (_, return_type, effects) = self.get(name)?;
        Some(FnSigInfo {
            return_type: return_type.clone(),
            is_pure: effects.is_empty(),
        })
    }
}

/// Bare-name → `&FnDef` index for the **entry module only**.
///
/// # Why a typed wrapper instead of raw `HashMap<String, &FnDef>`?
///
/// **Entry-only verify exception** (epic #170 Phase 5 deferred follow-up):
/// every helper-law-hint and contextual-helper-hint walker in this
/// module operates exclusively on top-level entry fn defs, because
/// `verify <name>` parses `name` as a single source `Ident` (see
/// `src/parser/blocks.rs::parse_verify`). The parser does NOT accept
/// dotted verify targets, so dep-module fns are unreachable from the
/// verify-law surface today. Bare-name keying is therefore
/// identity-safe by the parser invariant.
///
/// The typed wrapper exists so any future change that tries to
/// thread dep-module fns through this surface trips the compiler: a
/// `HashMap<String, &FnDef>` from a dep module won't coerce into
/// `EntryFnIndex`, forcing the contributor to confront the keying
/// question explicitly. When module-scoped verify ships, the right
/// migration is to drop this wrapper in favour of an `FnId`-keyed
/// view sourced from `CodegenContext.resolved_program`.
///
/// **temporary-migration-bridge**: deferred per epic #170 audit (no
/// live trigger today; documented in
/// `project_phase_e_scope_b_deferred` memory).
pub struct EntryFnIndex<'a> {
    inner: HashMap<String, &'a FnDef>,
}

impl<'a> EntryFnIndex<'a> {
    /// Build the index by walking entry-scope `TopLevel` items.
    /// Dep modules are intentionally NOT considered — see type doc.
    pub fn from_entry_items(items: &'a [TopLevel]) -> Self {
        let inner = items
            .iter()
            .filter_map(|item| {
                if let TopLevel::FnDef(fd) = item {
                    Some((fd.name.clone(), fd))
                } else {
                    None
                }
            })
            .collect();
        Self { inner }
    }

    /// Bare-name lookup. Safe by parser invariant — see type doc.
    pub fn get(&self, name: &str) -> Option<&&'a FnDef> {
        self.inner.get(name)
    }

    /// Existence check. Same identity contract as [`Self::get`].
    pub fn contains_key(&self, name: &str) -> bool {
        self.inner.contains_key(name)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NamedLawFunction {
    pub name: String,
    pub is_pure: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VerifyLawSpecRef {
    pub spec_fn_name: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MissingHelperLawHint {
    pub line: usize,
    pub fn_name: String,
    pub law_name: String,
    pub missing_helpers: Vec<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextualHelperLawHint {
    pub line: usize,
    pub fn_name: String,
    pub law_name: String,
    pub missing_helpers: Vec<String>,
}

pub fn named_law_function(law: &VerifyLaw, fn_sigs: &dyn FnSigOracle) -> Option<NamedLawFunction> {
    let info = fn_sigs.fn_sig(&law.name)?;
    Some(NamedLawFunction {
        name: law.name.clone(),
        is_pure: info.is_pure,
    })
}

pub fn declared_spec_ref(law: &VerifyLaw, fn_sigs: &dyn FnSigOracle) -> Option<VerifyLawSpecRef> {
    let named = named_law_function(law, fn_sigs)?;
    named.is_pure.then_some(VerifyLawSpecRef {
        spec_fn_name: named.name,
    })
}

pub fn law_spec_ref(law: &VerifyLaw, fn_sigs: &dyn FnSigOracle) -> Option<VerifyLawSpecRef> {
    let spec = declared_spec_ref(law, fn_sigs)?;
    law_calls_function(law, &spec.spec_fn_name).then_some(spec)
}

pub fn canonical_spec_ref(
    fn_name: &str,
    law: &VerifyLaw,
    fn_sigs: &dyn FnSigOracle,
) -> Option<VerifyLawSpecRef> {
    let spec = law_spec_ref(law, fn_sigs)?;
    canonical_spec_shape(fn_name, law, &spec.spec_fn_name).then_some(spec)
}

pub fn law_calls_function(law: &VerifyLaw, fn_name: &str) -> bool {
    expr_calls_function(&law.lhs, fn_name) || expr_calls_function(&law.rhs, fn_name)
}

pub fn canonical_spec_shape(fn_name: &str, law: &VerifyLaw, spec_fn_name: &str) -> bool {
    let try_side = |impl_side: &Spanned<Expr>, spec_side: &Spanned<Expr>| -> bool {
        let Some((impl_callee, impl_args)) = direct_call(impl_side) else {
            return false;
        };
        let Some((spec_callee, spec_args)) = direct_call(spec_side) else {
            return false;
        };
        impl_callee == fn_name && spec_callee == spec_fn_name && impl_args == spec_args
    };

    try_side(&law.lhs, &law.rhs) || try_side(&law.rhs, &law.lhs)
}

pub fn collect_missing_helper_law_hints(
    items: &[TopLevel],
    fn_sigs: &dyn FnSigOracle,
) -> Vec<MissingHelperLawHint> {
    let fn_defs = EntryFnIndex::from_entry_items(items);
    let verified_law_functions = items
        .iter()
        .filter_map(|item| {
            let TopLevel::Verify(vb) = item else {
                return None;
            };
            let VerifyKind::Law(law) = &vb.kind else {
                return None;
            };
            let mut covered = BTreeSet::new();
            covered.insert(vb.fn_name.clone());
            collect_direct_pure_user_calls(&law.lhs, &fn_defs, fn_sigs, &mut covered);
            collect_direct_pure_user_calls(&law.rhs, &fn_defs, fn_sigs, &mut covered);
            Some(covered)
        })
        .flatten()
        .collect::<HashSet<_>>();

    items
        .iter()
        .filter_map(|item| {
            let TopLevel::Verify(vb) = item else {
                return None;
            };
            let VerifyKind::Law(law) = &vb.kind else {
                return None;
            };
            missing_helper_law_hint_for_block(vb, law, &fn_defs, &verified_law_functions, fn_sigs)
        })
        .collect()
}

pub fn missing_helper_law_message(hint: &MissingHelperLawHint) -> String {
    format!(
        "verify law '{}.{}' uses helper functions without their own verify law: {}; add layered `verify ... law ...` blocks for those helpers before expecting a universal auto-proof",
        hint.fn_name,
        hint.law_name,
        hint.missing_helpers.join(", ")
    )
}

pub fn collect_contextual_helper_law_hints(
    items: &[TopLevel],
    fn_sigs: &dyn FnSigOracle,
) -> Vec<ContextualHelperLawHint> {
    let fn_defs = EntryFnIndex::from_entry_items(items);
    let contextual_law_targets = items
        .iter()
        .filter_map(|item| {
            let TopLevel::Verify(vb) = item else {
                return None;
            };
            let VerifyKind::Law(law) = &vb.kind else {
                return None;
            };
            top_level_direct_pure_call_in_law(law, &fn_defs, fn_sigs)
        })
        .collect::<HashSet<_>>();

    items
        .iter()
        .filter_map(|item| {
            let TopLevel::Verify(vb) = item else {
                return None;
            };
            let VerifyKind::Law(law) = &vb.kind else {
                return None;
            };
            // A law whose every given ranges over a closed finite
            // domain closes by exhaustive `cases` enumeration
            // (`ProofStrategy::FiniteDomainCases`) — the per-leaf
            // ground goals compute straight through helper fns, fuel
            // wrappers included, so the auto-proof does NOT stop at
            // helper boundaries and the ladder hint would be stale
            // advice.
            if law.when.is_none() && law_givens_are_closed_finite_domain(law, items) {
                return None;
            }
            contextual_helper_law_hint_for_block(
                vb,
                law,
                &fn_defs,
                &contextual_law_targets,
                fn_sigs,
            )
        })
        .collect()
}

/// True when `law` has at least one given and EVERY given ranges over
/// a closed finite domain — `Bool` (size 2) or a user-declared enum
/// whose constructors are ALL fieldless (size = constructor count) —
/// with the product of domain sizes ≤ 16. AST-level mirror of the
/// `ProofStrategy::FiniteDomainCases` detector in
/// `codegen::proof_lower` (kept in sync by the json.av
/// helper-law-ladder checker test); the caller must also mirror the
/// detector's `law.when.is_none()` gate — when-laws never get the
/// strategy, so their hints must not be suppressed. Such a law's
/// universal theorem is closed by exhaustive `cases` enumeration, not
/// by the helper-law ladder.
fn law_givens_are_closed_finite_domain(law: &VerifyLaw, items: &[TopLevel]) -> bool {
    if law.givens.is_empty() {
        return false;
    }
    let mut domain_product: usize = 1;
    for given in &law.givens {
        let size = if given.type_name == "Bool" {
            2
        } else {
            let fieldless_enum_size = items.iter().find_map(|item| match item {
                TopLevel::TypeDef(crate::ast::TypeDef::Sum { name, variants, .. })
                    if *name == given.type_name && variants.iter().all(|v| v.fields.is_empty()) =>
                {
                    Some(variants.len())
                }
                _ => None,
            });
            match fieldless_enum_size {
                Some(n) => n,
                None => return false,
            }
        };
        domain_product = match domain_product.checked_mul(size) {
            Some(p) if p <= 16 => p,
            _ => return false,
        };
    }
    true
}

pub fn contextual_helper_law_message(hint: &ContextualHelperLawHint) -> String {
    format!(
        "verify law '{}.{}' still lacks analogous `verify ... law ...` coverage for contextual helpers: {}; universal auto-proof will likely stop at those helper boundaries",
        hint.fn_name,
        hint.law_name,
        hint.missing_helpers.join(", ")
    )
}

fn missing_helper_law_hint_for_block(
    vb: &VerifyBlock,
    law: &VerifyLaw,
    fn_defs: &EntryFnIndex<'_>,
    verified_law_functions: &HashSet<String>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<MissingHelperLawHint> {
    if law.when.is_none() || law.givens.len() != 1 {
        return None;
    }

    let root_calls = direct_pure_user_calls_in_law(law, fn_defs, fn_sigs);
    if root_calls.is_empty() {
        return None;
    }

    let mut missing_helpers = BTreeSet::new();
    for root in root_calls {
        for helper in frontier_helper_calls(&root, fn_defs, fn_sigs) {
            if helper != vb.fn_name && !verified_law_functions.contains(&helper) {
                missing_helpers.insert(helper);
            }
        }
    }

    if missing_helpers.is_empty() {
        return None;
    }

    Some(MissingHelperLawHint {
        line: vb.line,
        fn_name: vb.fn_name.clone(),
        law_name: law.name.clone(),
        missing_helpers: missing_helpers.into_iter().collect(),
    })
}

fn contextual_helper_law_hint_for_block(
    vb: &VerifyBlock,
    law: &VerifyLaw,
    fn_defs: &EntryFnIndex<'_>,
    contextual_law_targets: &HashSet<String>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<ContextualHelperLawHint> {
    let parser_name = contextual_roundtrip_parser_name(law, fn_defs, fn_sigs)?;
    let root_parser_name = wrapper_dispatch_root(&parser_name, fn_defs, fn_sigs)
        .unwrap_or_else(|| parser_name.clone());
    if root_parser_name != parser_name {
        return None;
    }

    let missing_helpers = frontier_helper_calls(&root_parser_name, fn_defs, fn_sigs)
        .into_iter()
        .filter(|helper| helper != &vb.fn_name && !contextual_law_targets.contains(helper))
        .collect::<BTreeSet<_>>();

    if missing_helpers.is_empty() {
        return None;
    }

    Some(ContextualHelperLawHint {
        line: vb.line,
        fn_name: vb.fn_name.clone(),
        law_name: law.name.clone(),
        missing_helpers: missing_helpers.into_iter().collect(),
    })
}

fn direct_pure_user_calls_in_law(
    law: &VerifyLaw,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> BTreeSet<String> {
    let mut out = BTreeSet::new();
    collect_direct_pure_user_calls(&law.lhs, fn_defs, fn_sigs, &mut out);
    collect_direct_pure_user_calls(&law.rhs, fn_defs, fn_sigs, &mut out);
    out
}

fn top_level_direct_pure_call_in_law(
    law: &VerifyLaw,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<String> {
    direct_pure_user_call_name(&law.lhs, fn_defs, fn_sigs)
        .or_else(|| direct_pure_user_call_name(&law.rhs, fn_defs, fn_sigs))
}

fn contextual_roundtrip_parser_name(
    law: &VerifyLaw,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<String> {
    let given = law.givens.first()?;
    detect_roundtrip_layers(law, &given.name, fn_defs, fn_sigs).map(|(parser_name, _)| parser_name)
}

fn frontier_helper_calls(
    root_name: &str,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> BTreeSet<String> {
    let mut current =
        wrapper_dispatch_root(root_name, fn_defs, fn_sigs).unwrap_or_else(|| root_name.to_string());
    let mut visited = BTreeSet::new();

    for _ in 0..2 {
        if !visited.insert(current.clone()) {
            break;
        }
        let direct = direct_pure_fn_callees_matching_return(&current, fn_defs, fn_sigs);
        if direct.is_empty() {
            return BTreeSet::new();
        }
        if direct.len() == 1 {
            current = direct.iter().next().cloned().unwrap_or_default();
            continue;
        }
        return direct;
    }

    direct_pure_fn_callees_matching_return(&current, fn_defs, fn_sigs)
}

fn wrapper_dispatch_root(
    fn_name: &str,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<String> {
    let fd = fn_defs.get(fn_name)?;
    let tail = fd.body.tail_expr()?;
    match &tail.node {
        Expr::Match { subject, .. } => direct_pure_user_call_name(subject, fn_defs, fn_sigs),
        Expr::FnCall(_, _) => direct_pure_user_call_name(tail, fn_defs, fn_sigs),
        _ => None,
    }
}

fn direct_pure_fn_callees_matching_return(
    fn_name: &str,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> BTreeSet<String> {
    let Some(caller_info) = fn_sigs.fn_sig(fn_name) else {
        return BTreeSet::new();
    };
    let Some(fd) = fn_defs.get(fn_name) else {
        return BTreeSet::new();
    };

    let mut direct = BTreeSet::new();
    for stmt in fd.body.stmts() {
        match stmt {
            Stmt::Expr(expr) | Stmt::Binding(_, _, expr) => {
                collect_direct_pure_user_calls(expr, fn_defs, fn_sigs, &mut direct);
            }
        }
    }
    direct
        .into_iter()
        .filter(|callee| {
            callee != fn_name
                && fn_sigs.fn_sig(callee).is_some_and(|callee_info| {
                    callee_info.is_pure && callee_info.return_type == caller_info.return_type
                })
        })
        .collect()
}

fn collect_direct_pure_user_calls(
    expr: &Spanned<Expr>,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
    out: &mut BTreeSet<String>,
) {
    match &expr.node {
        Expr::FnCall(callee, args) => {
            if let Some(name) = direct_pure_user_call_name(expr, fn_defs, fn_sigs) {
                out.insert(name);
            }
            collect_direct_pure_user_calls(callee, fn_defs, fn_sigs, out);
            for arg in args {
                collect_direct_pure_user_calls(arg, fn_defs, fn_sigs, out);
            }
        }
        Expr::Attr(obj, _) => collect_direct_pure_user_calls(obj, fn_defs, fn_sigs, out),
        Expr::BinOp(_, left, right) => {
            collect_direct_pure_user_calls(left, fn_defs, fn_sigs, out);
            collect_direct_pure_user_calls(right, fn_defs, fn_sigs, out);
        }
        Expr::Neg(inner) => collect_direct_pure_user_calls(inner, fn_defs, fn_sigs, out),
        Expr::Match { subject, arms } => {
            collect_direct_pure_user_calls(subject, fn_defs, fn_sigs, out);
            for arm in arms {
                collect_direct_pure_user_calls(&arm.body, fn_defs, fn_sigs, out);
            }
        }
        Expr::Constructor(_, Some(inner)) | Expr::ErrorProp(inner) => {
            collect_direct_pure_user_calls(inner, fn_defs, fn_sigs, out);
        }
        Expr::InterpolatedStr(parts) => {
            for part in parts {
                if let StrPart::Parsed(inner) = part {
                    collect_direct_pure_user_calls(inner, fn_defs, fn_sigs, out);
                }
            }
        }
        Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
            for item in items {
                collect_direct_pure_user_calls(item, fn_defs, fn_sigs, out);
            }
        }
        Expr::MapLiteral(entries) => {
            for (key, value) in entries {
                collect_direct_pure_user_calls(key, fn_defs, fn_sigs, out);
                collect_direct_pure_user_calls(value, fn_defs, fn_sigs, out);
            }
        }
        Expr::RecordCreate { fields, .. } => {
            for (_, value) in fields {
                collect_direct_pure_user_calls(value, fn_defs, fn_sigs, out);
            }
        }
        Expr::RecordUpdate { base, updates, .. } => {
            collect_direct_pure_user_calls(base, fn_defs, fn_sigs, out);
            for (_, value) in updates {
                collect_direct_pure_user_calls(value, fn_defs, fn_sigs, out);
            }
        }
        Expr::TailCall(boxed) => {
            let TailCallData { target, args, .. } = boxed.as_ref();
            if fn_defs.contains_key(target)
                && fn_sigs.fn_sig(target).is_some_and(|info| info.is_pure)
            {
                out.insert(target.clone());
            }
            for arg in args {
                collect_direct_pure_user_calls(arg, fn_defs, fn_sigs, out);
            }
        }
        Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } | Expr::Constructor(_, None) => {}
    }
}

fn direct_pure_user_call_name(
    expr: &Spanned<Expr>,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<String> {
    let Expr::FnCall(callee, _) = &expr.node else {
        return None;
    };
    let name = dotted_name(callee)?;
    if !fn_defs.contains_key(&name) {
        return None;
    }
    fn_sigs
        .fn_sig(&name)
        .is_some_and(|info| info.is_pure)
        .then_some(name)
}

fn dotted_name(expr: &Spanned<Expr>) -> Option<String> {
    match &expr.node {
        Expr::Ident(name) => Some(name.clone()),
        Expr::Attr(base, field) => {
            let mut prefix = dotted_name(base)?;
            prefix.push('.');
            prefix.push_str(field);
            Some(prefix)
        }
        _ => None,
    }
}

fn detect_roundtrip_layers(
    law: &VerifyLaw,
    given_name: &str,
    fn_defs: &EntryFnIndex<'_>,
    fn_sigs: &dyn FnSigOracle,
) -> Option<(String, String)> {
    if law.givens.len() != 1 {
        return None;
    }

    fn detect_roundtrip_side(
        expr: &Spanned<Expr>,
        given_name: &str,
        fn_defs: &EntryFnIndex<'_>,
        fn_sigs: &dyn FnSigOracle,
    ) -> Option<(String, String)> {
        let Expr::FnCall(parser_callee, parser_args) = &expr.node else {
            return None;
        };
        if parser_args.is_empty() {
            return None;
        }
        let (serializer_callee, serializer_args) =
            extract_roundtrip_serializer_call(&parser_args[0], given_name)?;
        if !serializer_args
            .iter()
            .any(|arg| matches_ident(arg, given_name))
        {
            return None;
        }
        if serializer_args
            .iter()
            .filter(|arg| expr_mentions_ident(arg, given_name))
            .any(|arg| !matches_ident(arg, given_name))
        {
            return None;
        }
        if parser_args[1..]
            .iter()
            .any(|arg| expr_mentions_ident(arg, given_name))
        {
            return None;
        }

        let parser_name = dotted_name(parser_callee)?;
        let serializer_name = dotted_name(serializer_callee)?;
        if !fn_defs.contains_key(&parser_name) || !fn_defs.contains_key(&serializer_name) {
            return None;
        }
        if !fn_sigs
            .fn_sig(&parser_name)
            .is_some_and(|info| info.is_pure)
        {
            return None;
        }
        if !fn_sigs
            .fn_sig(&serializer_name)
            .is_some_and(|info| info.is_pure)
        {
            return None;
        }
        Some((parser_name, serializer_name))
    }

    detect_roundtrip_side(&law.lhs, given_name, fn_defs, fn_sigs)
        .or_else(|| detect_roundtrip_side(&law.rhs, given_name, fn_defs, fn_sigs))
}

fn extract_roundtrip_serializer_call<'a>(
    expr: &'a Spanned<Expr>,
    given_name: &str,
) -> Option<(&'a Spanned<Expr>, &'a [Spanned<Expr>])> {
    let mut candidates = Vec::new();
    collect_roundtrip_serializer_calls(expr, given_name, &mut candidates);
    if candidates.len() != 1 {
        return None;
    }
    let (callee, args) = candidates.pop()?;
    if expr_mentions_ident(expr, given_name)
        && args
            .iter()
            .filter(|arg| expr_mentions_ident(arg, given_name))
            .all(|arg| matches_ident(arg, given_name))
    {
        Some((callee, args))
    } else {
        None
    }
}

fn collect_roundtrip_serializer_calls<'a>(
    expr: &'a Spanned<Expr>,
    given_name: &str,
    out: &mut Vec<(&'a Spanned<Expr>, &'a [Spanned<Expr>])>,
) {
    match &expr.node {
        Expr::FnCall(callee, args) => {
            if args.iter().any(|arg| matches_ident(arg, given_name))
                && args
                    .iter()
                    .filter(|arg| expr_mentions_ident(arg, given_name))
                    .all(|arg| matches_ident(arg, given_name))
            {
                out.push((callee.as_ref(), args.as_slice()));
            }
            collect_roundtrip_serializer_calls(callee, given_name, out);
            for arg in args {
                collect_roundtrip_serializer_calls(arg, given_name, out);
            }
        }
        Expr::Attr(base, _) => collect_roundtrip_serializer_calls(base, given_name, out),
        Expr::BinOp(_, left, right) => {
            collect_roundtrip_serializer_calls(left, given_name, out);
            collect_roundtrip_serializer_calls(right, given_name, out);
        }
        Expr::Neg(inner) => collect_roundtrip_serializer_calls(inner, given_name, out),
        Expr::Match { subject, arms } => {
            collect_roundtrip_serializer_calls(subject, given_name, out);
            for arm in arms {
                collect_roundtrip_serializer_calls(&arm.body, given_name, out);
            }
        }
        Expr::Constructor(_, inner) => {
            if let Some(inner) = inner {
                collect_roundtrip_serializer_calls(inner, given_name, out);
            }
        }
        Expr::ErrorProp(inner) => collect_roundtrip_serializer_calls(inner, given_name, out),
        Expr::InterpolatedStr(parts) => {
            for part in parts {
                if let StrPart::Parsed(inner) = part {
                    collect_roundtrip_serializer_calls(inner, given_name, out);
                }
            }
        }
        Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
            for item in items {
                collect_roundtrip_serializer_calls(item, given_name, out);
            }
        }
        Expr::MapLiteral(entries) => {
            for (key, value) in entries {
                collect_roundtrip_serializer_calls(key, given_name, out);
                collect_roundtrip_serializer_calls(value, given_name, out);
            }
        }
        Expr::RecordCreate { fields, .. } => {
            for (_, value) in fields {
                collect_roundtrip_serializer_calls(value, given_name, out);
            }
        }
        Expr::RecordUpdate { base, updates, .. } => {
            collect_roundtrip_serializer_calls(base, given_name, out);
            for (_, value) in updates {
                collect_roundtrip_serializer_calls(value, given_name, out);
            }
        }
        Expr::TailCall(call) => {
            for arg in &call.args {
                collect_roundtrip_serializer_calls(arg, given_name, out);
            }
        }
        Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } => {}
    }
}

fn matches_ident(expr: &Spanned<Expr>, name: &str) -> bool {
    matches!(&expr.node, Expr::Ident(current) if current == name)
}

fn expr_mentions_ident(expr: &Spanned<Expr>, name: &str) -> bool {
    match &expr.node {
        Expr::Ident(current) => current == name,
        Expr::Attr(base, _) => expr_mentions_ident(base, name),
        Expr::FnCall(callee, args) => {
            expr_mentions_ident(callee, name)
                || args.iter().any(|arg| expr_mentions_ident(arg, name))
        }
        Expr::BinOp(_, left, right) => {
            expr_mentions_ident(left, name) || expr_mentions_ident(right, name)
        }
        Expr::Neg(inner) => expr_mentions_ident(inner, name),
        Expr::Match { subject, arms } => {
            expr_mentions_ident(subject, name)
                || arms.iter().any(|arm| expr_mentions_ident(&arm.body, name))
        }
        Expr::Constructor(_, inner) => inner
            .as_deref()
            .is_some_and(|inner| expr_mentions_ident(inner, name)),
        Expr::ErrorProp(inner) => expr_mentions_ident(inner, name),
        Expr::InterpolatedStr(parts) => parts.iter().any(|part| match part {
            StrPart::Literal(_) => false,
            StrPart::Parsed(inner) => expr_mentions_ident(inner, name),
        }),
        Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
            items.iter().any(|item| expr_mentions_ident(item, name))
        }
        Expr::MapLiteral(entries) => entries
            .iter()
            .any(|(key, value)| expr_mentions_ident(key, name) || expr_mentions_ident(value, name)),
        Expr::RecordCreate { fields, .. } => fields
            .iter()
            .any(|(_, value)| expr_mentions_ident(value, name)),
        Expr::RecordUpdate { base, updates, .. } => {
            expr_mentions_ident(base, name)
                || updates
                    .iter()
                    .any(|(_, value)| expr_mentions_ident(value, name))
        }
        Expr::TailCall(call) => call.args.iter().any(|arg| expr_mentions_ident(arg, name)),
        Expr::Literal(_) | Expr::Resolved { .. } => false,
    }
}

fn expr_calls_function(expr: &Spanned<Expr>, fn_name: &str) -> bool {
    match &expr.node {
        Expr::FnCall(callee, args) => {
            expr_is_function_name(callee, fn_name)
                || expr_calls_function(callee, fn_name)
                || args.iter().any(|arg| expr_calls_function(arg, fn_name))
        }
        Expr::Attr(obj, _) => expr_calls_function(obj, fn_name),
        Expr::BinOp(_, left, right) => {
            expr_calls_function(left, fn_name) || expr_calls_function(right, fn_name)
        }
        Expr::Neg(inner) => expr_calls_function(inner, fn_name),
        Expr::Match { subject, arms } => {
            expr_calls_function(subject, fn_name)
                || arms
                    .iter()
                    .any(|arm| match_arm_calls_function(arm, fn_name))
        }
        Expr::Constructor(_, Some(inner)) => expr_calls_function(inner, fn_name),
        Expr::ErrorProp(inner) => expr_calls_function(inner, fn_name),
        Expr::InterpolatedStr(parts) => parts.iter().any(|part| match part {
            StrPart::Literal(_) => false,
            StrPart::Parsed(expr) => expr_calls_function(expr, fn_name),
        }),
        Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
            items.iter().any(|item| expr_calls_function(item, fn_name))
        }
        Expr::MapLiteral(entries) => entries.iter().any(|(key, value)| {
            expr_calls_function(key, fn_name) || expr_calls_function(value, fn_name)
        }),
        Expr::RecordCreate { fields, .. } => fields
            .iter()
            .any(|(_, expr)| expr_calls_function(expr, fn_name)),
        Expr::RecordUpdate { base, updates, .. } => {
            expr_calls_function(base, fn_name)
                || updates
                    .iter()
                    .any(|(_, expr)| expr_calls_function(expr, fn_name))
        }
        Expr::TailCall(boxed) => {
            boxed.target == fn_name
                || boxed
                    .args
                    .iter()
                    .any(|arg| expr_calls_function(arg, fn_name))
        }
        Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } | Expr::Constructor(_, None) => {
            false
        }
    }
}

fn match_arm_calls_function(arm: &MatchArm, fn_name: &str) -> bool {
    expr_calls_function(&arm.body, fn_name)
}

fn expr_is_function_name(expr: &Spanned<Expr>, fn_name: &str) -> bool {
    matches!(&expr.node, Expr::Ident(name) if name == fn_name)
}

fn direct_call(expr: &Spanned<Expr>) -> Option<(&str, &[Spanned<Expr>])> {
    let Expr::FnCall(callee, args) = &expr.node else {
        return None;
    };
    let Expr::Ident(name) = &callee.node else {
        return None;
    };
    Some((name.as_str(), args.as_slice()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{Literal, Spanned, VerifyGiven, VerifyGivenDomain};

    fn int_sig() -> (Vec<Type>, Type, Vec<String>) {
        (vec![Type::Int], Type::Int, vec![])
    }

    fn law(lhs: Spanned<Expr>, rhs: Spanned<Expr>, name: &str) -> VerifyLaw {
        VerifyLaw {
            name: name.to_string(),
            givens: vec![VerifyGiven {
                name: "x".to_string(),
                type_name: "Int".to_string(),
                domain: VerifyGivenDomain::Explicit(vec![Spanned::bare(Expr::Literal(
                    Literal::Int(1),
                ))]),
            }],
            when: None,
            lhs,
            rhs,
            sample_guards: vec![],
        }
    }

    #[test]
    fn pure_named_law_function_becomes_declared_spec_ref() {
        let mut fn_sigs = FnSigMap::new();
        fn_sigs.insert("fooSpec".to_string(), int_sig());

        let verify_law = law(
            Spanned::bare(Expr::FnCall(
                Box::new(Spanned::bare(Expr::Ident("foo".to_string()))),
                vec![Spanned::bare(Expr::Ident("x".to_string()))],
            )),
            Spanned::bare(Expr::FnCall(
                Box::new(Spanned::bare(Expr::Ident("fooSpec".to_string()))),
                vec![Spanned::bare(Expr::Ident("x".to_string()))],
            )),
            "fooSpec",
        );

        assert_eq!(
            declared_spec_ref(&verify_law, &fn_sigs),
            Some(VerifyLawSpecRef {
                spec_fn_name: "fooSpec".to_string()
            })
        );
        assert_eq!(
            law_spec_ref(&verify_law, &fn_sigs),
            declared_spec_ref(&verify_law, &fn_sigs)
        );
        assert_eq!(
            canonical_spec_ref("foo", &verify_law, &fn_sigs),
            declared_spec_ref(&verify_law, &fn_sigs)
        );
    }

    #[test]
    fn effectful_named_law_function_is_not_a_spec_ref() {
        let mut fn_sigs = FnSigMap::new();
        fn_sigs.insert(
            "fooSpec".to_string(),
            (
                vec![Type::Int],
                Type::Int,
                vec!["Console.print".to_string()],
            ),
        );

        let verify_law = law(
            Spanned::bare(Expr::Ident("x".to_string())),
            Spanned::bare(Expr::Ident("x".to_string())),
            "fooSpec",
        );

        assert!(declared_spec_ref(&verify_law, &fn_sigs).is_none());
        assert_eq!(
            named_law_function(&verify_law, &fn_sigs),
            Some(NamedLawFunction {
                name: "fooSpec".to_string(),
                is_pure: false
            })
        );
    }

    #[test]
    fn canonical_spec_ref_requires_call_to_named_function() {
        let mut fn_sigs = FnSigMap::new();
        fn_sigs.insert("fooSpec".to_string(), int_sig());

        let verify_law = law(
            Spanned::bare(Expr::Ident("x".to_string())),
            Spanned::bare(Expr::Ident("x".to_string())),
            "fooSpec",
        );

        assert!(declared_spec_ref(&verify_law, &fn_sigs).is_some());
        assert!(law_spec_ref(&verify_law, &fn_sigs).is_none());
        assert!(!law_calls_function(&verify_law, "fooSpec"));
    }

    #[test]
    fn canonical_spec_ref_requires_same_arguments_on_both_sides() {
        let mut fn_sigs = FnSigMap::new();
        fn_sigs.insert("fooSpec".to_string(), int_sig());

        let verify_law = law(
            Spanned::bare(Expr::FnCall(
                Box::new(Spanned::bare(Expr::Ident("foo".to_string()))),
                vec![Spanned::bare(Expr::Ident("x".to_string()))],
            )),
            Spanned::bare(Expr::FnCall(
                Box::new(Spanned::bare(Expr::Ident("fooSpec".to_string()))),
                vec![
                    Spanned::bare(Expr::Literal(Literal::Int(5))),
                    Spanned::bare(Expr::Ident("x".to_string())),
                ],
            )),
            "fooSpec",
        );

        assert!(law_spec_ref(&verify_law, &fn_sigs).is_some());
        assert!(canonical_spec_ref("foo", &verify_law, &fn_sigs).is_none());
        assert!(!canonical_spec_shape("foo", &verify_law, "fooSpec"));
    }
}