elicitation 0.11.1

//! Kani harness generator for `elicitation generate kani`.
//!
//! Takes a `VsmDescriptor` (produced by the source scanner) and emits a Rust
//! source file containing Kani `proof_for_contract` harnesses that replicate
//! what `VerifiedStateMachine::vsm_kani_proof()` produces at build time.
//!
//! The generated file follows the two-shadow pattern:
//! 1. depth-2 state via `KaniCompose::kani_depth2()` + `kani::assume(inv_fn(&state))` + `mem::forget`
//! 2. depth-0 state via `KaniCompose::kani_depth0()` + proof credential + extra args + call + `mem::forget`

use std::collections::BTreeSet;
use std::path::Path;

use crate::cli::generate::{
    ImportStyle, TypeResolver, find_crate_name,
    scanner::{ArgKind, TransitionFn, VsmDescriptor},
};

// ─── Public API ───────────────────────────────────────────────────────────────

/// Generate the full Kani companion file content for `vsm`.
///
/// `crate_root` is used to discover the crate name via `Cargo.toml`.  The VSM
/// source file is parsed with [`TypeResolver`] to produce precise `use` imports
/// rather than glob wildcards.
///
/// Returns `Err` with an actionable message when a named invariant is present
/// but its `kani_invariant_fn` attribute is missing.
#[tracing::instrument(skip(vsm, crate_root), fields(machine = %vsm.machine))]
pub fn generate_kani_file(
    vsm: &VsmDescriptor,
    crate_root: impl AsRef<Path>,
) -> Result<String, String> {
    generate_kani_file_with_style(vsm, crate_root, ImportStyle::ExternalCrate)
}

/// Generate the full Kani companion file content for `vsm` with a specific
/// import style.
#[tracing::instrument(skip(vsm, crate_root), fields(machine = %vsm.machine))]
pub fn generate_kani_file_with_style(
    vsm: &VsmDescriptor,
    crate_root: impl AsRef<Path>,
    import_style: ImportStyle,
) -> Result<String, String> {
    let machine = &vsm.machine;

    // Resolve invariant — error only when an invariant exists without kani_fn.
    let (inv_fn, consistent_ty) = match &vsm.invariant {
        Some(p) => {
            let fn_name = p.kani_fn.as_deref().ok_or_else(|| {
                format!(
                    "cannot generate Kani companion for {machine}: \
                     invariant prop `{name}` has no `kani_invariant_fn` attribute\n  \
                     hint: add `kani_invariant_fn = \"<fn_name>\"` to #[prop(...)] on `{name}`",
                    name = p.name,
                )
            })?;
            (fn_name, p.name.as_str())
        }
        None => ("", ""),
    };
    let has_invariant = !consistent_ty.is_empty();

    // Derive crate name from the VSM's own source file, not the (possibly workspace)
    // crate_root.  Walking up from the source file finds the crate's own Cargo.toml.
    let vsm_crate_root = vsm
        .source_file
        .parent()
        .unwrap_or_else(|| crate_root.as_ref());
    let crate_name = find_crate_name(vsm_crate_root);

    // ── Collect the bare type names this file will reference ─────────────────
    let mut needed: BTreeSet<String> = BTreeSet::new();
    if vsm.transition_fns.is_empty() {
        TypeResolver::collect_type(&derive_state_ty(machine), &mut needed);
        if has_invariant {
            TypeResolver::collect_type(consistent_ty, &mut needed);
        }
    } else {
        for tfn in &vsm.transition_fns {
            for arg in &tfn.args {
                match &arg.kind {
                    ArgKind::State => TypeResolver::collect_type(&arg.ty, &mut needed),
                    ArgKind::Proof { inner } => TypeResolver::collect_type(inner, &mut needed),
                    _ => TypeResolver::collect_type(&arg.ty, &mut needed),
                }
            }
        }
        // Import the transition functions themselves, not just their arg types.
        for name in &vsm.transitions {
            needed.insert(name.clone());
        }
        // inv_fn is a plain pub fn re-exported unconditionally at the crate root
        // via the generated KANI REEXPORTS section — use the flat crate-root path.
        // Do NOT add to `needed` here; TypeResolver would resolve it to its private
        // module path, which is unreachable from the proof crate.
    }

    let resolver = TypeResolver::build(&vsm.source_file, &crate_name, import_style);
    let resolved_imports = resolver.grouped_imports(&needed);

    // Build explicit imports for _kani_contracted wrappers and the invariant
    // predicate.  These are macro-generated under #[cfg(kani)] so the file
    // scanner cannot find them.  They are re-exported at the source-crate root
    // via the generated `kani_reexports` module — use the flat crate-root path.
    let contracted_imports: Vec<String> = if has_invariant {
        let mut imports: Vec<String> = vsm
            .transitions
            .iter()
            .map(|name| format!("{crate_name}::{name}_kani_contracted"))
            .collect();
        imports.push(format!("{crate_name}::{inv_fn}"));
        imports
    } else {
        vec![]
    };

    let mut out = String::new();

    // ── Header ──────────────────────────────────────────────────────────────
    out.push_str(&format!(
        "// AUTO-GENERATED by `elicitation generate kani` — DO NOT EDIT\n\
         //\n\
         // Kani proof_for_contract harnesses for {machine}.\n\
         // Regenerate: elicitation generate kani --crate-path <root>\n\
         // The `kani` cfg key is set by the Kani toolchain, not by cargo.\n\
         #![allow(unexpected_cfgs)]\n\
         \n"
    ));

    // ── Imports ─────────────────────────────────────────────────────────────
    out.push_str("#[cfg(kani)]\nuse elicitation::Established;\n");
    for import in &resolved_imports {
        out.push_str(&format!("#[cfg(kani)]\nuse {import};\n"));
    }
    for import in &contracted_imports {
        out.push_str(&format!("#[cfg(kani)]\nuse {import};\n"));
    }
    out.push('\n');

    // ── Marker proof ────────────────────────────────────────────────────────
    if has_invariant {
        out.push_str(&format!(
            "#[cfg(kani)]\n#[kani::proof]\nfn verify_{consistent_lower}_prop_marker() {{\n    let established: bool = true;\n    assert!(established);\n}}\n",
            consistent_lower = to_snake(consistent_ty)
        ));
    }

    // ── Per-transition harnesses ─────────────────────────────────────────────
    if vsm.transition_fns.is_empty() {
        let state_ty = derive_state_ty(machine);
        for name in &vsm.transitions {
            out.push('\n');
            out.push_str(&emit_minimal_harness(
                name,
                &state_ty,
                consistent_ty,
                inv_fn,
                has_invariant,
            ));
        }
    } else {
        for tfn in &vsm.transition_fns {
            out.push('\n');
            out.push_str(&emit_harness(tfn, inv_fn, has_invariant));
        }
    }

    Ok(out)
}

// ─── Harness emitters ─────────────────────────────────────────────────────────

/// Emit a full `proof_for_contract` harness for a `TransitionFn` with known args.
fn emit_harness(tfn: &TransitionFn, inv_fn: &str, has_invariant: bool) -> String {
    let fn_name = &tfn.name;
    let closure_name = format!("{fn_name}_kani_closure");

    let state_arg = tfn.state_arg();
    let state_ty = state_arg
        .map(|a| a.ty.as_str())
        .unwrap_or("/* TODO: StateTy */");
    let state_pat = state_arg.map(|a| a.name.as_str()).unwrap_or("_state");

    let call_args: Vec<&str> = tfn.args.iter().map(|a| a.name.as_str()).collect();

    let mut lines: Vec<String> = Vec::new();

    // Shadow 1: depth-2 non-det state + assume invariant + forget.
    lines.push(format!(
        "    let {state_pat}: {state_ty} = <{state_ty} as ::elicitation::KaniCompose>::kani_depth2();"
    ));
    if has_invariant {
        lines.push(format!("    ::kani::assume({inv_fn}(&{state_pat}));"));
    }
    lines.push(format!("    ::std::mem::forget({state_pat});"));

    // Shadow 2: depth-0 state for the actual call.
    lines.push(format!(
        "    let {state_pat}: {state_ty} = <{state_ty} as ::elicitation::KaniCompose>::kani_depth0();"
    ));

    // Proof credential and extra args.
    for arg in &tfn.args {
        match &arg.kind {
            ArgKind::State => {}
            ArgKind::Proof { inner } => {
                let name = &arg.name;
                lines.push(format!("    let {name}: Established<{inner}> = {{"));
                lines.push(format!(
                    "        let _cred = {inner}::kani_proof_credential();"
                ));
                lines.push("        ::elicitation::Established::prove(&_cred)".to_string());
                lines.push("    };".to_string());
            }
            ArgKind::StringArg => {
                lines.push(format!(
                    "    let {name}: {ty} = <::std::string::String as ::elicitation::KaniCompose>::kani_depth1();",
                    name = arg.name,
                    ty = arg.ty,
                ));
            }
            ArgKind::VecArg => {
                lines.push(format!(
                    "    let {name}: {ty} = ::std::vec::Vec::new();",
                    name = arg.name,
                    ty = arg.ty,
                ));
            }
            ArgKind::OptionArg => {
                lines.push(format!(
                    "    let {name}: {ty} = ::core::option::Option::None;",
                    name = arg.name,
                    ty = arg.ty,
                ));
            }
            ArgKind::Other => {
                lines.push(format!(
                    "    let {name}: {ty} = <{ty} as ::elicitation::KaniCompose>::kani_depth0();",
                    name = arg.name,
                    ty = arg.ty,
                ));
            }
        }
    }

    let call_arg_str = call_args.join(", ");
    let contracted = format!("{fn_name}_kani_contracted");
    lines.push(format!("    let _result = {contracted}({call_arg_str});"));
    lines.push("    ::std::mem::forget(_result);".to_string());

    let body = lines.join("\n");
    format!(
        "#[cfg(kani)]\n#[::kani::proof_for_contract({contracted})]\nfn {closure_name}() {{\n{body}\n}}\n"
    )
}

/// Emit a minimal placeholder harness when function signatures are not available.
fn emit_minimal_harness(
    fn_name: &str,
    state_ty: &str,
    consistent_ty: &str,
    inv_fn: &str,
    has_invariant: bool,
) -> String {
    let closure_name = format!("{fn_name}_kani_closure");
    let assume_line = if has_invariant {
        format!("    ::kani::assume({inv_fn}(&_state));")
    } else {
        String::new()
    };
    let (proof_lines, call_arg) = if has_invariant {
        (
            vec![
                format!("    let proof: Established<{consistent_ty}> = {{"),
                format!("        let _cred = {consistent_ty}::kani_proof_credential();"),
                "        ::elicitation::Established::prove(&_cred)".to_string(),
                "    };".to_string(),
            ],
            ", proof".to_string(),
        )
    } else {
        (vec![], String::new())
    };
    let contracted = format!("{fn_name}_kani_contracted");
    let mut lines = vec![
        format!("// TODO: verify argument types for {fn_name} and regenerate"),
        "#[cfg(kani)]".to_string(),
        format!("#[::kani::proof_for_contract({contracted})]"),
        format!("fn {closure_name}() {{"),
        format!(
            "    let _state: {state_ty} = <{state_ty} as ::elicitation::KaniCompose>::kani_depth2();"
        ),
        assume_line,
        "    ::std::mem::forget(_state);".to_string(),
        format!(
            "    let _state: {state_ty} = <{state_ty} as ::elicitation::KaniCompose>::kani_depth0();"
        ),
    ];
    lines.extend(proof_lines);
    lines.push(format!("    let _result = {contracted}(_state{call_arg});"));
    lines.push("    ::std::mem::forget(_result);".to_string());
    lines.push("}".to_string());
    lines.join("\n") + "\n"
}

// ─── Helpers ──────────────────────────────────────────────────────────────────

/// Infer the state enum type name from the machine struct name.
///
/// `ArchiveConnectionMachine` → `ArchiveConnectionState`
fn derive_state_ty(machine: &str) -> String {
    machine.replace("Machine", "State")
}

/// Convert a PascalCase or camelCase identifier to snake_case for use in fn names.
fn to_snake(s: &str) -> String {
    let mut out = String::new();
    for (i, ch) in s.chars().enumerate() {
        if ch.is_uppercase() && i > 0 {
            out.push('_');
        }
        out.push(ch.to_ascii_lowercase());
    }
    out
}