vir/src/ast_simplify.rs

Source
//! VIR-AST -> VIR-AST transformation to simplify away some complicated features

use crate::ast::CrateId;
use crate::ast::Quant;
use crate::ast::Typs;
use crate::ast::VarBinder;
use crate::ast::VarBinderX;
use crate::ast::VarBinders;
use crate::ast::VarIdent;
use crate::ast::{
    AssocTypeImpl, AutospecUsage, BinaryOp, Binder, BoundsCheck, BuiltinSpecFun, ByRef, CallTarget,
    ChainedOp, ClosureKind, Constant, CtorPrintStyle, CtorUpdateTail, Datatype,
    DatatypeTransparency, DatatypeX, Dt, Expr, ExprX, Exprs, Field, FieldOpr, Fun, Function,
    FunctionKind, Ident, IntRange, ItemKind, Krate, KrateX, LogicalOp, Mode, MultiOp, Path,
    Pattern, PatternBinding, PatternX, Place, PlaceX, SpannedTyped, Stmt, StmtX, TraitImpl, Typ,
    TypX, UnaryOp, UnaryOpr, Variant, VariantCheck, VirErr, Visibility,
};
use crate::ast_util::{
    conjoin, mk_eq, mk_implies, place_to_spec_expr, typ_args_for_datatype_typ, undecorate_typ,
    unit_typ, wrap_in_trigger,
};
use crate::ast_visitor::VisitorScopeMap;
use crate::context::GlobalCtx;
use crate::def::dummy_param_name;
use crate::def::is_dummy_param_name;
use crate::def::{
    Spanned, impl_fndef_path, positional_field_ident, prefix_tuple_param, prefix_tuple_variant,
    user_local_name,
};
use crate::messages::Span;
use crate::messages::{error, internal_error};
use crate::sst_util::subst_typ_for_datatype;
use crate::util::vec_map_result;
use air::ast_util::ident_binder;
use air::scope_map::ScopeMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

struct ClosureDatatype {
    enclosing_fun: Fun,
    args: Typs,
    output: Typ,
    kind: ClosureKind,
    path: Path,
}

struct State {
    // Counter to generate temporary variables
    next_var: u64,
    // Rename parameters to simplify their names
    rename_vars: HashMap<VarIdent, VarIdent>,
    // Rename parameters to simplify their names
    rename_vars_reverse: HashMap<VarIdent, VarIdent>,
    // Name of a datatype to represent each tuple arity
    tuple_typs: HashSet<usize>,
    // Name of a datatype to represent each closure
    closure_typs: HashMap<usize, ClosureDatatype>,
    // Functions for which the corresponding FnDef type is used
    fndef_typs: HashSet<Fun>,
}

impl State {
    fn new() -> Self {
        State {
            next_var: 0,
            rename_vars: HashMap::new(),
            rename_vars_reverse: HashMap::new(),
            tuple_typs: HashSet::new(),
            closure_typs: HashMap::new(),
            fndef_typs: HashSet::new(),
        }
    }

    fn reset_for_function(&mut self) {
        self.next_var = 0;
        self.rename_vars = HashMap::new();
    }

    fn next_temp(&mut self) -> VarIdent {
        self.next_var += 1;
        crate::def::simplify_temp_var(self.next_var)
    }

    fn tuple_type_name(&mut self, arity: usize) -> Dt {
        self.tuple_typs.insert(arity);
        Dt::Tuple(arity)
    }

    fn closure_type_name(
        &mut self,
        fun: Fun,
        typs: Typs,
        typ: Typ,
        kind: ClosureKind,
        id: usize,
    ) -> Path {
        let e = self.closure_typs.entry(id).or_insert(ClosureDatatype {
            enclosing_fun: fun.clone(),
            args: typs,
            output: typ,
            kind: kind,
            path: crate::def::prefix_closure_type(id),
        });
        assert!(&e.enclosing_fun == &fun);
        e.path.clone()
    }
}

struct LocalCtxt {
    span: Span,
    typ_params: Vec<Ident>,
    fun: Option<Fun>,
}

/// Should only return true if this expression is guaranteed constant
/// (i.e., does not depend on evaluation order, i.e., does not depend on any mutable variable)
fn is_small_expr(expr: &Expr) -> bool {
    match &expr.x {
        ExprX::Const(_) => true,
        ExprX::Unary(UnaryOp::Not | UnaryOp::Clip { .. }, e) => is_small_expr(e),
        ExprX::UnaryOpr(UnaryOpr::Box(_) | UnaryOpr::Unbox(_), _) => panic!("unexpected box"),
        _ => false,
    }
}

/// Create a temporary and return:
///  - A Stmt that assigns the given `expr` to the temporary
///  - The name of the temporary
fn temp_var(state: &mut State, expr: &Expr) -> (Stmt, VarIdent) {
    let temp = state.next_temp();
    let name = temp.clone();
    let pattern = PatternX::simple_var(name, &expr.span, &expr.typ);
    let decl = StmtX::Decl {
        pattern,
        mode: None,
        init: Some(PlaceX::spec_temporary(expr.clone())),
        els: None,
        assert_irrefutable: false,
    };
    let temp_decl = Spanned::new(expr.span.clone(), decl);
    (temp_decl, temp)
}

fn temp_expr(state: &mut State, expr: &Expr) -> (Stmt, Expr) {
    let (temp_decl, var_ident) = temp_var(state, expr);
    (temp_decl, SpannedTyped::new(&expr.span, &expr.typ, ExprX::Var(var_ident)))
}

fn small_or_temp(state: &mut State, expr: &Expr) -> (Vec<Stmt>, Expr) {
    if is_small_expr(&expr) {
        (vec![], expr.clone())
    } else {
        let (ts, te) = temp_expr(state, expr);
        (vec![ts], te)
    }
}

fn pattern_to_decls_with_no_initializer(pattern: &Pattern, stmts: &mut Vec<Stmt>) {
    match &pattern.x {
        PatternX::Wildcard(_) => {}
        PatternX::Var(binding) | PatternX::Binding { binding, sub_pat: _ } => {
            let v_patternx = PatternX::Var(PatternBinding {
                name: binding.name.clone(),
                user_mut: None,
                by_ref: ByRef::No,
                typ: binding.typ.clone(),
                copy: false,
            });
            let v_pattern = SpannedTyped::new(&pattern.span, &binding.typ, v_patternx);
            stmts.push(Spanned::new(
                pattern.span.clone(),
                StmtX::Decl {
                    pattern: v_pattern,
                    mode: None, // mode doesn't matter anymore
                    init: None,
                    els: None,
                    assert_irrefutable: false,
                },
            ));

            match &pattern.x {
                PatternX::Binding { sub_pat, .. } => {
                    pattern_to_decls_with_no_initializer(sub_pat, stmts);
                }
                _ => {}
            }
        }
        PatternX::Constructor(_path, _variant, patterns) => {
            for binder in patterns.iter() {
                pattern_to_decls_with_no_initializer(&binder.a, stmts);
            }
        }
        PatternX::Or(pat1, _pat2) => {
            pattern_to_decls_with_no_initializer(&pat1, stmts);
        }
        PatternX::Expr(_) => {}
        PatternX::Range(_, _) => {}
        PatternX::ImmutRef(p) | PatternX::MutRef(p) => {
            pattern_to_decls_with_no_initializer(p, stmts);
        }
    }
}

fn rename_var(state: &State, scope_map: &VisitorScopeMap, x: &VarIdent) -> VarIdent {
    if let Some(rename) = state.rename_vars.get(x) {
        if scope_map[x].is_outer_param_or_ret {
            return rename.clone();
        }
    }
    x.clone()
}

fn simplify_one_place(
    _ctx: &GlobalCtx,
    state: &mut State,
    scope_map: &VisitorScopeMap,
    place: &Place,
) -> Result<Place, VirErr> {
    match &place.x {
        PlaceX::Local(x) => Ok(place.new_x(PlaceX::Local(rename_var(state, scope_map, x)))),
        _ => Ok(place.clone()),
    }
}

/// Returns a "pure place", i.e., a Place with no-side effects, and which is rooted
/// at a Local (rather than a Temporary).
fn place_to_pure_place(state: &mut State, place: &Place) -> (Vec<Stmt>, Place) {
    let (mut stmts, place, wf) = place_to_pure_place_rec(state, place);
    stmts.extend(wf);
    (stmts, place)
}

/// Returns (Stmts, place, wf)
/// wf are the remaining obligations to prove that accessing the place is safe
/// (i.e., those not already icnluded in Stmts)
fn place_to_pure_place_rec(state: &mut State, place: &Place) -> (Vec<Stmt>, Place, Vec<Stmt>) {
    match &place.x {
        PlaceX::Field(field_opr, p) => {
            let (stmts, p1, mut wf) = place_to_pure_place_rec(state, p);
            match field_opr.check {
                VariantCheck::None => {}
                VariantCheck::Union => {
                    let p1_expr = place_to_spec_expr(&p1);
                    let assert_stmt =
                        crate::place_preconditions::field_check(&place.span, &p1_expr, field_opr);
                    wf.push(assert_stmt);
                }
                // Handled later, at SST lowering - no AST-level encoding for this timing.
                VariantCheck::Recommends => {}
            }
            // Preserve Recommends for that later pass; Union is already discharged above.
            let check = if field_opr.check == VariantCheck::Recommends {
                VariantCheck::Recommends
            } else {
                VariantCheck::None
            };
            let field_opr = FieldOpr { check, ..field_opr.clone() };
            let p2 =
                SpannedTyped::new(&place.span, &place.typ, PlaceX::Field(field_opr.clone(), p1));
            (stmts, p2, wf)
        }
        PlaceX::DerefMut(p) => {
            let (stmts, p1, wf) = place_to_pure_place_rec(state, p);
            let p2 = SpannedTyped::new(&place.span, &place.typ, PlaceX::DerefMut(p1));
            (stmts, p2, wf)
        }
        PlaceX::ModeUnwrap(p, mwm) => {
            let (stmts, p1, wf) = place_to_pure_place_rec(state, p);
            let p2 = SpannedTyped::new(&place.span, &place.typ, PlaceX::ModeUnwrap(p1, *mwm));
            (stmts, p2, wf)
        }
        PlaceX::Local(_l) => (vec![], place.clone(), vec![]),
        PlaceX::Temporary(expr) => {
            let (ts, var_ident) = temp_var(state, expr);
            let p = SpannedTyped::new(&place.span, &place.typ, PlaceX::Local(var_ident));
            (vec![ts], p, vec![])
        }
        PlaceX::WithExpr(expr, p) => {
            let (mut stmts, p1, wf) = place_to_pure_place_rec(state, p);
            stmts.insert(0, Spanned::new(place.span.clone(), StmtX::Expr(expr.clone())));
            (stmts, p1, wf)
        }
        PlaceX::Index(p, idx, kind, bounds_check) => {
            let (mut stmts, p1, mut wf) = place_to_pure_place_rec(state, p);
            let (idx_decl, idx_expr) = temp_expr(state, idx);
            stmts.push(idx_decl);

            match bounds_check {
                BoundsCheck::Allow => {}
                BoundsCheck::Error => {
                    if kind.getting_len_requires_read() {
                        stmts.extend(wf);
                        wf = vec![];
                    }
                    let p1_expr = place_to_spec_expr(&p1);
                    let assert_stmt = crate::place_preconditions::index_bound(
                        &place.span,
                        &p1_expr,
                        &idx_expr,
                        *kind,
                    );
                    stmts.push(assert_stmt);
                }
            }

            let p = SpannedTyped::new(
                &place.span,
                &place.typ,
                PlaceX::Index(p1, idx_expr, *kind, BoundsCheck::Allow),
            );

            (stmts, p, wf)
        }
        PlaceX::UserDefinedTypInvariantObligation(..) => {
            panic!("Verus internal error: unexpected UserDefinedTypInvariantObligation");
        }
    }
}

// note that this gets called *bottom up*
// that is, if node A is the parent of children B and C,
// then simplify_one_expr is called first on B and C, and then on A

fn simplify_one_expr(
    ctx: &GlobalCtx,
    state: &mut State,
    scope_map: &VisitorScopeMap,
    expr: &Expr,
) -> Result<Expr, VirErr> {
    use crate::ast::CallTargetKind;
    match &expr.x {
        ExprX::Var(x) => Ok(expr.new_x(ExprX::Var(rename_var(state, scope_map, x)))),
        ExprX::VarAt(x, at) => Ok(expr.new_x(ExprX::VarAt(rename_var(state, scope_map, x), *at))),
        ExprX::Assign { place, .. }
        | ExprX::BorrowMut(place)
        | ExprX::TwoPhaseBorrowMut(place)
        | ExprX::BorrowMutTracked(place) => {
            // This check is no longer needed for soundness because ast_to_sst infers
            // mutability rather than relying on the user annotations.
            // Nonetheless, this check lets us pick up a few situations that wouldn't
            // get caught by borrowck (ghost variables).
            // However, much of the time it _does_ get caught by borrowck.
            // This check is in ast_simplify because it's after borrowck
            // (borrowck errors look nicer).
            if !crate::ast_util::place_has_deref_mut(place)
                && let Some(local) = crate::ast_util::place_get_local(place)
            {
                let PlaceX::Local(x) = &local.x else { unreachable!() };
                let x = match state.rename_vars_reverse.get(x) {
                    None => x,
                    Some(y) => y,
                };
                match scope_map.get(x) {
                    None => {
                        return Err(error(
                            &expr.span,
                            "Verus Internal Error: cannot find this variable",
                        ));
                    }
                    Some(entry) if entry.user_mut == Some(false) && entry.init => {
                        let name = user_local_name(x);
                        return Err(error(
                            &expr.span,
                            format!("variable `{name:}` is not marked mutable"),
                        ));
                    }
                    _ => {}
                }
            }
            Ok(expr.clone())
        }
        ExprX::ConstVar(x, autospec) => {
            let call_target_attrs = crate::ast::CallTargetAttrs {
                autospec: *autospec,
                const_var: true,
                assume_external_allowed: false,
            };
            let call = ExprX::Call {
                target: CallTarget::Fun(
                    CallTargetKind::Static,
                    x.clone(),
                    Arc::new(vec![]),
                    Arc::new(vec![]),
                    call_target_attrs,
                ),
                args: Arc::new(vec![]),
                post_args: None,
                body: None,
            };
            Ok(SpannedTyped::new(&expr.span, &expr.typ, call))
        }
        ExprX::Call {
            target: CallTarget::Fun(kind, tgt, typs, impl_paths, attrs),
            args,
            post_args,
            body,
        } => {
            assert!(attrs.autospec == AutospecUsage::Final);

            let is_trait_impl = match kind {
                CallTargetKind::Static => false,
                CallTargetKind::ProofFn(..) => false,
                CallTargetKind::Dynamic => true,
                CallTargetKind::DynamicResolved { .. } => true,
                CallTargetKind::ExternalTraitDefault => true,
            };
            let args = if typs.len() == 0 && args.len() == 0 && !is_trait_impl {
                // To simplify the AIR/SMT encoding, add a dummy argument to any function with 0 arguments
                let typ = Arc::new(TypX::Int(IntRange::Int));
                use num_traits::Zero;
                let argx = ExprX::Const(Constant::Int(num_bigint::BigInt::zero()));
                let arg = SpannedTyped::new(&expr.span, &typ, argx);
                Arc::new(vec![arg])
            } else {
                args.clone()
            };

            let call = ExprX::Call {
                target: CallTarget::Fun(
                    kind.clone(),
                    tgt.clone(),
                    typs.clone(),
                    impl_paths.clone(),
                    attrs.clone(),
                ),
                args,
                post_args: post_args.clone(),
                body: body.clone(),
            };
            Ok(SpannedTyped::new(&expr.span, &expr.typ, call))
        }
        ExprX::Ctor(name, variant, partial_binders, Some(update)) => {
            let CtorUpdateTail { place, taken_fields: _ } = update;
            let (stmts, update) = place_to_pure_place(state, place);
            // not really spec but that doesn't matter at this point
            let update = place_to_spec_expr(&update);
            let mut decls: Vec<Stmt> = Vec::new();
            let mut binders: Vec<Binder<Expr>> = Vec::new();
            if stmts.len() == 0 {
                for binder in partial_binders.iter() {
                    binders.push(binder.clone());
                }
            } else {
                // Because of Rust's order of evaluation here,
                // we have to put binders in temp vars, too.
                for binder in partial_binders.iter() {
                    let (temp_decl_inner, e) = small_or_temp(state, &binder.a);
                    decls.extend(temp_decl_inner.into_iter());
                    binders.push(binder.map_a(|_| e));
                }
                decls.extend(stmts.into_iter());
            }

            let path = match name {
                Dt::Path(p) => p,
                Dt::Tuple(_) => {
                    return Err(internal_error(
                        &expr.span,
                        "ExprX::Ctor with update and tuple type",
                    ));
                }
            };

            let (typ_positives, variants) = &ctx.datatypes[path];
            let fields = &crate::ast_util::get_variant(&variants, variant).fields;
            let typ_args = typ_args_for_datatype_typ(&expr.typ);
            // replace ..update
            // with f1: update.f1, f2: update.f2, ...
            for field in fields.iter() {
                if binders.iter().find(|b| b.name == field.name).is_none() {
                    let op = UnaryOpr::Field(FieldOpr {
                        datatype: name.clone(),
                        variant: variant.clone(),
                        field: field.name.clone(),
                        get_variant: false,
                        check: VariantCheck::None,
                    });
                    let exprx = ExprX::UnaryOpr(op, update.clone());
                    let ty = subst_typ_for_datatype(&typ_positives, typ_args, &field.a.0);
                    let field_exp = SpannedTyped::new(&expr.span, &ty, exprx);
                    binders.push(ident_binder(&field.name, &field_exp));
                }
            }
            let ctorx = ExprX::Ctor(name.clone(), variant.clone(), Arc::new(binders), None);
            let ctor = SpannedTyped::new(&expr.span, &expr.typ, ctorx);
            if decls.len() == 0 {
                Ok(ctor)
            } else {
                let block = ExprX::Block(Arc::new(decls), Some(ctor));
                Ok(SpannedTyped::new(&expr.span, &expr.typ, block))
            }
        }
        ExprX::ShrRefStructWrap(e1, e2, _t1, t2, variant, field) => {
            // Simplify as `Struct { field: e1, .. e2 }`
            let datatype_typ = undecorate_typ(&t2);
            let (dt, _typ_args) = match &*datatype_typ {
                TypX::Datatype(dt, typ_args, ..) => (dt, typ_args),
                _ => panic!("ShrRefStructWrap expects datatype"),
            };
            let Dt::Path(path) = dt else { panic!("ShrRefStructWrap expects Dt::Path") };
            let partial_binders =
                Arc::new(vec![Arc::new(air::ast::BinderX { name: field.clone(), a: e1.clone() })]);
            let place = SpannedTyped::new(&e2.span, &e2.typ, PlaceX::Temporary(e2.clone()));
            // taken_fields is ignored by this point
            let upd = CtorUpdateTail { place: place, taken_fields: Arc::new(vec![]) };
            let variant = if **variant == "" {
                let (_, variants) = &ctx.datatypes[path];
                assert!(variants.len() == 1);
                variants[0].name.clone()
            } else {
                variant.clone()
            };
            let ctor = ExprX::Ctor(dt.clone(), variant.clone(), partial_binders, Some(upd));
            let ctor = SpannedTyped::new(&expr.span, &expr.typ, ctor);
            simplify_one_expr(ctx, state, scope_map, &ctor)
        }
        ExprX::Unary(UnaryOp::CoerceMode { .. }, expr0) => Ok(expr0.clone()),
        ExprX::Multi(MultiOp::Chained(ops), args) => {
            use crate::ast::IeeeFloatBinaryOp;
            assert!(args.len() == ops.len() + 1);
            let mut stmts: Vec<Stmt> = Vec::new();
            let mut es: Vec<Expr> = Vec::new();
            let mut is_float = false;
            // Execute each argument in order; no short-circuiting
            for i in 0..args.len() {
                let t = crate::ast_util::undecorate_typ(&args[i].typ);
                if matches!(*t, TypX::Float(_)) {
                    is_float = true;
                }
                let (decl, e) = temp_expr(state, &args[i]);
                stmts.push(decl);
                es.push(e);
            }
            let mut conjunction: Expr = es[0].clone();
            for i in 0..ops.len() {
                let op = match (is_float, ops[i]) {
                    (false, ChainedOp::Inequality(a)) => BinaryOp::Inequality(a),
                    (true, ChainedOp::Inequality(a)) => {
                        BinaryOp::IeeeFloat(IeeeFloatBinaryOp::InEq(a))
                    }
                    (_, ChainedOp::MultiEq) => BinaryOp::Eq(Mode::Spec),
                };
                let left = es[i].clone();
                let right = es[i + 1].clone();
                let span = left.span.clone();
                let binary = SpannedTyped::new(&span, &expr.typ, ExprX::Binary(op, left, right));
                if i == 0 {
                    conjunction = binary;
                } else {
                    let exprx = ExprX::Logical(LogicalOp::And, conjunction, binary);
                    conjunction = SpannedTyped::new(&span, &expr.typ, exprx);
                }
            }
            if stmts.len() == 0 {
                Ok(conjunction)
            } else {
                let block = ExprX::Block(Arc::new(stmts), Some(conjunction));
                Ok(SpannedTyped::new(&expr.span, &expr.typ, block))
            }
        }
        ExprX::Match(place, arms1, assert_irrefutable) => {
            let (temp_decl, place) = place_to_pure_place(state, place);

            // Translate into If expression
            let t_bool = Arc::new(TypX::Bool);
            let mut if_expr: Option<Expr> = None;
            for arm in arms1.iter().rev() {
                let mut decls: Vec<Stmt> = Vec::new();
                let has_guard = arm.x.has_guard();

                let test_pattern = crate::patterns::pattern_to_exprs(
                    ctx,
                    &place,
                    &arm.x.pattern,
                    has_guard,
                    &mut decls,
                )?;

                let test = if !has_guard {
                    test_pattern
                } else {
                    assert!(!crate::patterns::pattern_has_or(&arm.x.pattern));

                    let mut guard = arm.x.guard.clone();
                    guard = SpannedTyped::new(
                        &guard.span,
                        &guard.typ,
                        ExprX::MatchGuardFreeze(place.clone(), guard.clone()),
                    );
                    let test_exp = ExprX::Logical(LogicalOp::And, test_pattern, guard);
                    let test = SpannedTyped::new(&arm.x.pattern.span, &t_bool, test_exp);
                    let block = ExprX::Block(Arc::new(decls.clone()), Some(test));
                    SpannedTyped::new(&arm.x.pattern.span, &t_bool, block)
                };

                let block = ExprX::Block(Arc::new(decls), Some(arm.x.body.clone()));
                let body = SpannedTyped::new(&arm.x.pattern.span, &expr.typ, block);
                if let Some(prev) = if_expr {
                    // if pattern && guard then body else prev
                    let ifx = ExprX::If(test.clone(), body, Some(prev));
                    if_expr = Some(SpannedTyped::new(&test.span, &expr.typ.clone(), ifx));
                } else if *assert_irrefutable {
                    if has_guard {
                        return Err(error(&arm.x.guard.span, "if-guard on final match arm"));
                    }
                    let assertion = SpannedTyped::new(
                        &arm.x.pattern.span,
                        &unit_typ(),
                        ExprX::AssertAssume {
                            is_assume: false,
                            expr: test,
                            msg: Some(irrefut_failure_msg(&arm.x.pattern.span)),
                        },
                    );
                    let block = SpannedTyped::new(
                        &body.span,
                        &body.typ,
                        ExprX::Block(
                            Arc::new(vec![Spanned::new(
                                assertion.span.clone(),
                                StmtX::Expr(assertion.clone()),
                            )]),
                            Some(body.clone()),
                        ),
                    );
                    if_expr = Some(block);
                } else {
                    // last arm is unconditional
                    if_expr = Some(body);
                }
            }
            if let Some(if_expr) = if_expr {
                let if_expr = if temp_decl.len() != 0 {
                    let block = ExprX::Block(Arc::new(temp_decl), Some(if_expr));
                    SpannedTyped::new(&expr.span, &expr.typ, block)
                } else {
                    if_expr
                };
                Ok(if_expr)
            } else {
                Err(error(&expr.span, "not yet implemented: zero-arm match expressions"))
            }
        }
        ExprX::Ghost { alloc_wrapper: _, tracked: _, expr: expr1 } => Ok(expr1.clone()),
        ExprX::NonSpecClosure {
            params,
            proof_fn_modes,
            body,
            requires,
            ensures,
            ret,
            external_spec,
        } => {
            assert!(external_spec.is_none());

            let closure_var_ident = state.next_temp();
            let closure_var = SpannedTyped::new(
                &expr.span,
                &expr.typ.clone(),
                ExprX::Var(closure_var_ident.clone()),
            );

            let external_spec_expr =
                exec_closure_spec(state, &expr.span, &closure_var, params, ret, requires, ensures)?;
            let external_spec = Some((closure_var_ident, external_spec_expr));

            Ok(SpannedTyped::new(
                &expr.span,
                &expr.typ,
                ExprX::NonSpecClosure {
                    params: params.clone(),
                    proof_fn_modes: proof_fn_modes.clone(),
                    body: body.clone(),
                    requires: requires.clone(),
                    ensures: ensures.clone(),
                    ret: ret.clone(),
                    external_spec,
                },
            ))
        }
        _ => Ok(expr.clone()),
    }
}

fn tuple_get_field_expr(
    state: &mut State,
    span: &Span,
    typ: &Typ,
    tuple_expr: &Expr,
    tuple_arity: usize,
    field: usize,
) -> Expr {
    let datatype = state.tuple_type_name(tuple_arity);

    let variant = prefix_tuple_variant(tuple_arity);
    let field = positional_field_ident(field);
    let op = UnaryOpr::Field(FieldOpr {
        datatype,
        variant,
        field,
        get_variant: false,
        check: VariantCheck::None,
    });
    let field_expr = SpannedTyped::new(span, typ, ExprX::UnaryOpr(op, tuple_expr.clone()));
    field_expr
}

fn irrefut_failure_msg(pattern_span: &Span) -> crate::messages::Message {
    error(pattern_span, "unable to prove this pattern will successfully match")
}

fn simplify_one_stmt(ctx: &GlobalCtx, state: &mut State, stmt: &Stmt) -> Result<Vec<Stmt>, VirErr> {
    match &stmt.x {
        StmtX::Decl { pattern, mode: _, init: None, els: None, assert_irrefutable } => {
            assert!(!assert_irrefutable);
            match &pattern.x {
                PatternX::Var(PatternBinding {
                    by_ref: ByRef::No,
                    name: _,
                    user_mut: _,
                    typ: _,
                    copy: _,
                }) => Ok(vec![stmt.clone()]),
                _ => {
                    let mut stmts: Vec<Stmt> = Vec::new();
                    pattern_to_decls_with_no_initializer(pattern, &mut stmts);
                    Ok(stmts)
                }
            }
        }
        StmtX::Decl { pattern, mode: _, init: None, els: Some(_), .. } => Err(error(
            &pattern.span,
            "Verus Internal Error: Decl with else-block but no initializer",
        )),
        StmtX::Decl { pattern, mode: _, init: Some(_init), els: None, assert_irrefutable: _ }
            if matches!(
                pattern.x,
                PatternX::Var(PatternBinding {
                    by_ref: ByRef::No,
                    name: _,
                    user_mut: _,
                    typ: _,
                    copy: _
                })
            ) =>
        {
            Ok(vec![stmt.clone()])
        }
        StmtX::Decl { pattern, mode: _, init: Some(init), els, assert_irrefutable } => {
            let (mut stmts, place) = place_to_pure_place(state, init);
            let mut stmts2: Vec<Stmt> = vec![];
            let pattern_check =
                crate::patterns::pattern_to_exprs(ctx, &place, pattern, false, &mut stmts2)?;
            if let Some(els) = &els {
                assert!(!assert_irrefutable);
                let checkx = ExprX::Unary(UnaryOp::Not, pattern_check.clone());
                let check = SpannedTyped::new(&pattern_check.span, &pattern_check.typ, checkx);
                let neverx = ExprX::NeverToAny(els.clone());
                let never = SpannedTyped::new(&els.span, &unit_typ(), neverx);
                let ifx = ExprX::If(check.clone(), never, None);
                let ife = SpannedTyped::new(&stmt.span, &unit_typ(), ifx);
                let ifstmtx = StmtX::Expr(ife);
                let ifstmt = Spanned::new(stmt.span.clone(), ifstmtx);
                stmts.push(ifstmt);
            } else if *assert_irrefutable {
                stmts.push(Spanned::new(
                    stmt.span.clone(),
                    StmtX::Expr(SpannedTyped::new(
                        &stmt.span,
                        &unit_typ(),
                        ExprX::AssertAssume {
                            is_assume: false,
                            expr: pattern_check,
                            msg: Some(irrefut_failure_msg(&pattern.span)),
                        },
                    )),
                ));
            }
            stmts.extend(stmts2);
            Ok(stmts)
        }
        StmtX::Expr(_) => Ok(vec![stmt.clone()]),
    }
}

fn simplify_one_typ(local: &LocalCtxt, state: &mut State, typ: &Typ) -> Result<Typ, VirErr> {
    match &**typ {
        TypX::Datatype(Dt::Tuple(i), ..) => {
            state.tuple_type_name(*i);
            Ok(typ.clone())
        }
        TypX::AnonymousClosure(typs, typ, kind, id) => {
            let Some(fun) = local.fun.clone() else {
                return Err(error(
                    &local.span,
                    format!("Verus Internal Error: found AnonymousClosure type outside function"),
                ));
            };
            let path =
                Dt::Path(state.closure_type_name(fun, typs.clone(), typ.clone(), *kind, *id));
            let typ_args: Vec<Typ> = local
                .typ_params
                .iter()
                .map(|name| Arc::new(TypX::TypParam(name.clone())))
                .collect();
            Ok(Arc::new(TypX::Datatype(path, Arc::new(typ_args), Arc::new(vec![]))))
        }
        TypX::FnDef(fun, _typs, resolved) => {
            state.fndef_typs.insert(fun.clone());
            if let Some(resolved_fun) = resolved {
                state.fndef_typs.insert(resolved_fun.clone());
            }
            Ok(typ.clone())
        }
        TypX::TypParam(x) => {
            if !local.typ_params.contains(&x) {
                return Err(error(
                    &local.span,
                    format!("type parameter {} used before being declared", x),
                ));
            }
            Ok(typ.clone())
        }
        _ => Ok(typ.clone()),
    }
}

// TODO: a lot of this closure stuff could get its own file
// rename to apply to all fn types, not just closure types

fn closure_trait_call_typ_args(state: &mut State, fn_val: &Expr, params: &VarBinders<Typ>) -> Typs {
    let path = state.tuple_type_name(params.len());

    let param_typs: Vec<Typ> = params.iter().map(|p| p.a.clone()).collect();
    let tup_typ = Arc::new(TypX::Datatype(path, Arc::new(param_typs), Arc::new(vec![])));

    Arc::new(vec![fn_val.typ.clone(), tup_typ])
}

fn mk_closure_req_call(
    state: &mut State,
    span: &Span,
    params: &VarBinders<Typ>,
    fn_val: &Expr,
    arg_tuple: &Expr,
) -> Expr {
    let bool_typ = Arc::new(TypX::Bool);
    SpannedTyped::new(
        span,
        &bool_typ,
        ExprX::Call {
            target: CallTarget::BuiltinSpecFun(
                BuiltinSpecFun::ClosureReq,
                closure_trait_call_typ_args(state, fn_val, params),
                Arc::new(vec![]),
            ),
            args: Arc::new(vec![fn_val.clone(), arg_tuple.clone()]),
            post_args: None,
            body: None,
        },
    )
}

fn mk_closure_ens_call(
    state: &mut State,
    span: &Span,
    params: &VarBinders<Typ>,
    fn_val: &Expr,
    arg_tuple: &Expr,
    ret_arg: &Expr,
    builtin_spec_fun: BuiltinSpecFun,
) -> Expr {
    let bool_typ = Arc::new(TypX::Bool);
    SpannedTyped::new(
        span,
        &bool_typ,
        ExprX::Call {
            target: CallTarget::BuiltinSpecFun(
                builtin_spec_fun,
                closure_trait_call_typ_args(state, fn_val, params),
                Arc::new(vec![]),
            ),
            args: Arc::new(vec![fn_val.clone(), arg_tuple.clone(), ret_arg.clone()]),
            post_args: None,
            body: None,
        },
    )
}

fn exec_closure_spec_param(
    state: &mut State,
    span: &Span,
    params: &VarBinders<Typ>,
) -> (VarIdent, Expr) {
    let param_typs: Vec<Typ> = params.iter().map(|p| p.a.clone()).collect();
    let tuple_path = state.tuple_type_name(params.len());
    let tuple_typ = Arc::new(TypX::Datatype(tuple_path, Arc::new(param_typs), Arc::new(vec![])));
    let tuple_ident = crate::def::closure_param_var();
    let tuple_var = SpannedTyped::new(span, &tuple_typ, ExprX::Var(tuple_ident.clone()));
    (tuple_ident, tuple_var)
}

fn exec_closure_spec_requires(
    state: &mut State,
    span: &Span,
    closure_var: &Expr,
    params: &VarBinders<Typ>,
    requires: &Exprs,
) -> Result<Expr, VirErr> {
    // For requires:

    // If the closure has `|a0, a1, a2| requires f(a0, a1, a2)`
    // then we emit a spec of the form
    //
    //      forall x :: f(x.0, x.1, x.2) ==> closure.requires(x)
    //
    // with `closure.requires(x)` as the trigger

    // (Since the user doesn't have the option to specify a trigger here,
    // we need to use the most general one, and that means we need to
    // quantify over a tuple.)

    let (tuple_ident, tuple_var) = exec_closure_spec_param(state, span, params);
    let tuple_typ = tuple_var.typ.clone();

    let reqs = conjoin(span, requires);

    // Supply 'let' statements of the form 'let a0 = x.0; let a1 = x.1; ...' etc.

    let mut decls: Vec<Stmt> = Vec::new();
    for (i, p) in params.iter().enumerate() {
        let typ = &p.a;
        let pattern = PatternX::simple_var(p.name.clone(), span, typ);
        let tuple_field = tuple_get_field_expr(state, span, typ, &tuple_var, params.len(), i);
        let decl = StmtX::Decl {
            pattern,
            mode: None,
            init: Some(PlaceX::spec_temporary(tuple_field)),
            els: None,
            assert_irrefutable: false,
        };
        decls.push(Spanned::new(span.clone(), decl));
    }

    let reqs_body =
        SpannedTyped::new(&reqs.span, &reqs.typ, ExprX::Block(Arc::new(decls), Some(reqs.clone())));

    let closure_req_call =
        wrap_in_trigger(&mk_closure_req_call(state, span, params, closure_var, &tuple_var));

    let bool_typ = Arc::new(TypX::Bool);
    let req_quant_body = mk_implies(span, &reqs_body, &closure_req_call);

    let forall = Quant { quant: air::ast::Quant::Forall };
    let binders = Arc::new(vec![Arc::new(VarBinderX { name: tuple_ident, a: tuple_typ })]);
    let req_forall =
        SpannedTyped::new(span, &bool_typ, ExprX::Quant(forall, binders, req_quant_body));

    Ok(req_forall)
}

fn exec_closure_spec_ensures(
    state: &mut State,
    span: &Span,
    closure_var: &Expr,
    params: &VarBinders<Typ>,
    ret: &VarBinder<Typ>,
    ensures: &Vec<Expr>,
    default_ens: bool,
) -> Result<Expr, VirErr> {
    // For ensures:

    // If the closure has `|a0, a1, a2| ensures |b| f(a0, a1, a2, b)`
    // then we emit a spec of the form
    //
    //      forall x, y :: closure.ensures(x, y) ==> f(x.0, x.1, x.2, y)
    //
    // with `closure.ensures(x)` as the trigger

    let (tuple_ident, tuple_var) = exec_closure_spec_param(state, span, params);
    let tuple_typ = tuple_var.typ.clone();

    let ret_ident = ret.clone();
    let ret_var = SpannedTyped::new(span, &ret.a, ExprX::Var(ret_ident.name.clone()));

    let enss = conjoin(span, ensures);

    // Supply 'let' statements of the form 'let a0 = x.0; let a1 = x.1; ...' etc.

    let mut decls: Vec<Stmt> = Vec::new();
    for (i, p) in params.iter().enumerate() {
        let typ = &p.a;
        let pattern = PatternX::simple_var(p.name.clone(), span, typ);
        let tuple_field = tuple_get_field_expr(state, span, typ, &tuple_var, params.len(), i);
        let decl = StmtX::Decl {
            pattern,
            mode: None,
            init: Some(PlaceX::spec_temporary(tuple_field)),
            els: None,
            assert_irrefutable: false,
        };
        decls.push(Spanned::new(span.clone(), decl));
    }

    let enss_body =
        SpannedTyped::new(&enss.span, &enss.typ, ExprX::Block(Arc::new(decls), Some(enss.clone())));

    let closure_ens_call = wrap_in_trigger(&mk_closure_ens_call(
        state,
        span,
        params,
        closure_var,
        &tuple_var,
        &ret_var,
        if default_ens { BuiltinSpecFun::DefaultEns } else { BuiltinSpecFun::ClosureEns },
    ));

    let bool_typ = Arc::new(TypX::Bool);
    let ens_quant_body = mk_implies(span, &closure_ens_call, &enss_body);

    let forall = Quant { quant: air::ast::Quant::Forall };
    let binders =
        Arc::new(vec![Arc::new(VarBinderX { name: tuple_ident, a: tuple_typ }), ret.clone()]);
    let ens_forall =
        SpannedTyped::new(span, &bool_typ, ExprX::Quant(forall, binders, ens_quant_body));

    Ok(ens_forall)
}

fn exec_closure_spec(
    state: &mut State,
    span: &Span,
    closure_var: &Expr,
    params: &VarBinders<Typ>,
    ret: &VarBinder<Typ>,
    requires: &Exprs,
    ensures: &Exprs,
) -> Result<Expr, VirErr> {
    let req_forall = exec_closure_spec_requires(state, span, closure_var, params, requires)?;

    if ensures.len() > 0 {
        let ens_forall =
            exec_closure_spec_ensures(state, span, closure_var, params, ret, ensures, false)?;
        Ok(conjoin(span, &vec![req_forall, ens_forall]))
    } else {
        Ok(req_forall)
    }
}

pub(crate) fn need_fndef_axiom(fndef_typs: &HashSet<Fun>, f: &Function) -> bool {
    if fndef_typs.contains(&f.x.name) {
        return true;
    }
    match &f.x.kind {
        FunctionKind::TraitMethodImpl { method, .. } => fndef_typs.contains(method),
        _ => false,
    }
}

fn add_fndef_axioms_to_function(
    _ctx: &GlobalCtx,
    state: &mut State,
    function: &Function,
    fn_once_trait_in_scope: bool,
) -> Result<(Function, Vec<TraitImpl>, Option<AssocTypeImpl>), VirErr> {
    state.reset_for_function();

    let params: Vec<_> = function
        .x
        .params
        .iter()
        .filter(|p| !is_dummy_param_name(&p.x.name))
        .map(|p| Arc::new(VarBinderX { name: p.x.name.clone(), a: p.x.typ.clone() }))
        .collect();
    let params = Arc::new(params);

    let (fun, typ_args, is_trait_method_impl, inherit) = match &function.x.kind {
        FunctionKind::TraitMethodImpl { method, trait_typ_args, inherit_body_from, .. } => {
            (method, trait_typ_args.clone(), true, inherit_body_from.is_some())
        }
        _ => {
            let typ_args: Vec<_> = function
                .x
                .typ_params
                .iter()
                .map(|tp| Arc::new(TypX::TypParam(tp.clone())))
                .collect();
            let typ_args = Arc::new(typ_args);
            (&function.x.name, typ_args, false, false)
        }
    };

    let fndef_singleton = SpannedTyped::new(
        &function.span,
        &Arc::new(TypX::FnDef(fun.clone(), typ_args.clone(), None)),
        ExprX::ExecFnByName(fun.clone()),
    );

    // Emit `FnDef : {Fn, FnMut, FnOnce}<Args>` and `<FnDef as FnOnce<Args>>::Output = Ret`.
    //
    // We emit a TraitImpl for each of the three Fn-family traits (not just Fn), because
    // code that mentions only one through an associated-type projection (e.g. `Map::Item = F::Output`)
    // never creates a Fn term for the Fn-related axioms to trigger on.
    let (trait_impls_out, assoc_type_impl) = if fn_once_trait_in_scope {
        let self_typ = Arc::new(TypX::FnDef(fun.clone(), typ_args.clone(), None));
        let arg_typs: Vec<Typ> = params.iter().map(|p| p.a.clone()).collect();
        let tuple_dt = state.tuple_type_name(arg_typs.len());
        let args_tuple_typ =
            Arc::new(TypX::Datatype(tuple_dt, Arc::new(arg_typs), Arc::new(vec![])));
        let trait_typ_args = Arc::new(vec![self_typ, args_tuple_typ]);

        let mut trait_impls_out: Vec<TraitImpl> = Vec::new();
        for kind in [ClosureKind::Fn, ClosureKind::FnMut, ClosureKind::FnOnce] {
            let trait_implx = crate::ast::TraitImplX {
                impl_path: impl_fndef_path(&function.x.name, kind),
                typ_params: function.x.typ_params.clone(),
                typ_bounds: function.x.typ_bounds.clone(),
                trait_path: kind.trait_path(),
                trait_typ_args: trait_typ_args.clone(),
                trait_typ_arg_impls: Spanned::new(function.span.clone(), Arc::new(vec![])),
                owning_module: None,
                auto_imported: true,
                external_trait_blanket: false,
            };
            trait_impls_out.push(Spanned::new(function.span.clone(), trait_implx));
        }

        let assoc_typ_implx = crate::ast::AssocTypeImplX {
            name: Arc::new("Output".to_string()),
            impl_path: impl_fndef_path(&function.x.name, ClosureKind::FnOnce),
            typ_params: function.x.typ_params.clone(),
            typ_bounds: function.x.typ_bounds.clone(),
            trait_path: ClosureKind::FnOnce.trait_path(),
            trait_typ_args,
            typ: function.x.ret.x.typ.clone(),
            impl_paths: Arc::new(vec![]),
        };

        (trait_impls_out, Some(Spanned::new(function.span.clone(), assoc_typ_implx)))
    } else {
        (Vec::new(), None)
    };

    let mut fndef_axioms = vec![];

    // Don't need to repeat the 'requires' for a trait impl fn because requires can't change
    if !is_trait_method_impl {
        let req_forall = exec_closure_spec_requires(
            state,
            &function.span,
            &fndef_singleton,
            &params,
            &function.x.require,
        )?;
        fndef_axioms.push(req_forall);
    }

    let ret = Arc::new(VarBinderX {
        name: function.x.ret.x.name.clone(),
        a: function.x.ret.x.typ.clone(),
    });
    let (mut closure_enss, default_enss) = function.x.ensure.clone();
    if inherit {
        assert!(closure_enss.len() + default_enss.len() == 0);
        let (_, tuple_var) = exec_closure_spec_param(state, &function.span, &params);
        let ret_var = SpannedTyped::new(&function.span, &ret.a, ExprX::Var(ret.name.clone()));
        let default_expr = mk_closure_ens_call(
            state,
            &function.span,
            &params,
            &fndef_singleton,
            &tuple_var,
            &ret_var,
            BuiltinSpecFun::DefaultEns,
        );
        closure_enss = Arc::new(vec![default_expr]);
    }
    for (default_ens, enss) in [(false, closure_enss), (true, default_enss)] {
        if enss.len() > 0 {
            let ens_forall = exec_closure_spec_ensures(
                state,
                &function.span,
                &fndef_singleton,
                &params,
                &ret,
                &*enss,
                default_ens,
            )?;
            fndef_axioms.push(ens_forall);
        }
    }

    let mut functionx = function.x.clone();
    assert!(functionx.fndef_axioms.is_none());
    functionx.fndef_axioms = Some(Arc::new(fndef_axioms));
    Ok((Spanned::new(function.span.clone(), functionx), trait_impls_out, assoc_type_impl))
}

fn simplify_function(
    ctx: &GlobalCtx,
    state: &mut State,
    function: &Function,
) -> Result<Function, VirErr> {
    state.reset_for_function();
    let mut functionx = function.x.clone();

    if let Some(r) = functionx.returns.clone() {
        functionx.returns = None;

        if functionx.ens_has_return {
            let var = SpannedTyped::new(
                &r.span,
                &functionx.ret.x.typ,
                ExprX::Var(functionx.ret.x.name.clone()),
            );
            let eq = mk_eq(&r.span, &var, &r);
            Arc::make_mut(&mut functionx.ensure.0).push(eq);
        } else {
            // For a unit return type, any returns clause is tautological so we
            // can just skip appending to the postconditions.
        }
    }

    let local = LocalCtxt {
        span: function.span.clone(),
        typ_params: (*functionx.typ_params).clone(),
        fun: Some(functionx.name.clone()),
    };

    let is_trait_impl = matches!(functionx.kind, FunctionKind::TraitMethodImpl { .. });

    // If possible, rename parameters to drop the rustc id
    let mut param_ids: HashSet<Ident> = HashSet::new();
    let mut rename_ok = true;
    for p in functionx.params.iter() {
        let x = &p.x.name.0;
        if param_ids.contains(x) {
            rename_ok = false;
        }
        param_ids.insert(x.clone());
    }
    let mut param_names: Vec<VarIdent> = Vec::new();
    for param in functionx.params.iter() {
        let prev = param.x.name.clone();
        let name = if rename_ok {
            let name = VarIdent(prev.0.clone(), crate::ast::VarIdentDisambiguate::VirParam);
            state.rename_vars.insert(prev, name.clone()).map(|_| panic!("rename params"));
            name
        } else {
            prev
        };
        param_names.push(name);
    }
    let ret_name = if rename_ok && !param_ids.contains(&functionx.ret.x.name.0) {
        let prev = functionx.ret.x.name.clone();
        let name = VarIdent(prev.0.clone(), crate::ast::VarIdentDisambiguate::VirParam);
        state.rename_vars.insert(prev, name.clone()).map(|_| panic!("rename ret"));
        name
    } else {
        functionx.ret.x.name.clone()
    };

    for (a, b) in state.rename_vars.iter() {
        state.rename_vars_reverse.insert(b.clone(), a.clone());
    }

    // To simplify the AIR/SMT encoding, add a dummy argument to any function with 0 arguments
    if functionx.typ_params.len() == 0
        && functionx.params.len() == 0
        && !matches!(functionx.item_kind, ItemKind::Const)
        && !matches!(functionx.item_kind, ItemKind::Static)
        && !functionx.attrs.broadcast_forall
        && !is_trait_impl
    {
        let paramx = crate::ast::ParamX {
            name: dummy_param_name(),
            typ: Arc::new(TypX::Int(IntRange::Int)),
            mode: Mode::Spec,
            user_mut: false,
            unwrapped_info: None,
        };
        param_names.push(paramx.name.clone());
        let param = Spanned::new(function.span.clone(), paramx);
        functionx.params = Arc::new(vec![param]);
    }

    let function = Spanned::new(function.span.clone(), functionx);
    let mut map: VisitorScopeMap = ScopeMap::new();
    let function = crate::ast_visitor::map_function_visitor_env(
        &function,
        &mut map,
        state,
        &|state, map, expr| simplify_one_expr(ctx, state, map, expr),
        &|state, _, stmt| simplify_one_stmt(ctx, state, stmt),
        &|state, typ| simplify_one_typ(&local, state, typ),
        &|state, map, place| simplify_one_place(ctx, state, map, place),
    )?;
    let mut functionx = function.x.clone();
    assert!(functionx.params.len() == param_names.len());
    functionx.params = Arc::new(
        functionx
            .params
            .iter()
            .zip(param_names.iter())
            .map(|(p, x)| p.new_x(crate::ast::ParamX { name: x.clone(), ..p.x.clone() }))
            .collect(),
    );
    functionx.ret =
        functionx.ret.new_x(crate::ast::ParamX { name: ret_name, ..functionx.ret.x.clone() });

    Ok(Spanned::new(function.span.clone(), functionx))
}

fn simplify_datatype(state: &mut State, datatype: &Datatype) -> Result<Datatype, VirErr> {
    let mut local = LocalCtxt { span: datatype.span.clone(), typ_params: Vec::new(), fun: None };
    for (x, _strict_pos) in datatype.x.typ_params.iter() {
        local.typ_params.push(x.clone());
    }
    crate::ast_visitor::map_datatype_visitor_env(datatype, state, &|state, typ| {
        simplify_one_typ(&local, state, typ)
    })
}

fn simplify_trait_impl(state: &mut State, imp: &TraitImpl) -> Result<TraitImpl, VirErr> {
    let mut local = LocalCtxt { span: imp.span.clone(), typ_params: Vec::new(), fun: None };
    for x in imp.x.typ_params.iter() {
        local.typ_params.push(x.clone());
    }
    crate::ast_visitor::map_trait_impl_visitor_env(imp, state, &|state, typ| {
        simplify_one_typ(&local, state, typ)
    })
}

fn simplify_assoc_type_impl(
    state: &mut State,
    assoc: &AssocTypeImpl,
) -> Result<AssocTypeImpl, VirErr> {
    let mut local = LocalCtxt { span: assoc.span.clone(), typ_params: Vec::new(), fun: None };
    for x in assoc.x.typ_params.iter() {
        local.typ_params.push(x.clone());
    }
    crate::ast_visitor::map_assoc_type_impl_visitor_env(assoc, state, &|state, typ| {
        simplify_one_typ(&local, state, typ)
    })
}

/*
fn mk_fun_decl(
    span: &Span,
    path: &Path,
    typ_params: &Idents,
    params: &Params,
    ret: &Param,
) -> Function {
    let mut attrs: crate::ast::FunctionAttrsX = Default::default();
    attrs.no_auto_trigger = true;
    Spanned::new(
        span.clone(),
        FunctionX {
            name: Arc::new(FunX { path: path.clone() }),
            visibility: Visibility { owning_module: None, restricted_to: None },
            mode: Mode::Spec,
            fuel: 0,
            typ_params: typ_params.clone(),
            typ_bounds: Arc::new(vec![]),
            params: params.clone(),
            ret: ret.clone(),
            require: Arc::new(vec![]),
            ensure: Arc::new(vec![]),
            decrease: None,
            is_const: false,
            is_abstract: false,
            attrs: Arc::new(attrs),
            body: None,
        },
    )
}
*/

fn add_tuple_auto_impl(
    ctx: &mut GlobalCtx,
    traits: &Vec<crate::ast::Trait>,
    trait_impls: &mut Vec<TraitImpl>,
    arity: usize,
    needs_bounds: bool,
    trait_path: Path,
) {
    use crate::ast::{GenericBound, GenericBoundX};
    use crate::def::impl_tuple;

    if !traits.iter().any(|t| &t.x.name == &trait_path) {
        return;
    }
    // Rust doesn't seem to give us the TraitImpl for tuples for the traits
    // Tuple, Clone, Copy, Send, Sync, Unpin, so we make it ourselves:
    let typ_params: Vec<Ident> = (0..arity).map(|i| prefix_tuple_param(i)).collect();
    let typ_args: Vec<Typ> =
        typ_params.iter().map(|x| Arc::new(TypX::TypParam(x.clone()))).collect();
    let mut bounds: Vec<GenericBound> = Vec::new();
    if needs_bounds {
        for i in 0..arity {
            let id = crate::ast::TraitId::Path(trait_path.clone());
            let typs = Arc::new(vec![typ_args[i].clone()]);
            bounds.push(Arc::new(GenericBoundX::Trait(id, typs)));
        }
    }
    let self_ty = Arc::new(TypX::Datatype(Dt::Tuple(arity), Arc::new(typ_args), Arc::new(vec![])));
    let impl_path = Arc::new(crate::ast::PathX {
        krate: CrateId::Internal,
        segments: Arc::new(vec![impl_tuple(trait_path.segments.last().unwrap(), arity)]),
    });
    let trait_implx = crate::ast::TraitImplX {
        impl_path,
        typ_params: Arc::new(typ_params),
        typ_bounds: Arc::new(bounds),
        trait_path,
        trait_typ_args: Arc::new(vec![self_ty]),
        trait_typ_arg_impls: Spanned::new(ctx.no_span.clone(), Arc::new(vec![])),
        owning_module: None,
        auto_imported: true,
        external_trait_blanket: false,
    };
    trait_impls.push(Spanned::new(ctx.no_span.clone(), trait_implx));
}

pub fn simplify_krate(ctx: &mut GlobalCtx, krate: &Krate) -> Result<Krate, VirErr> {
    let KrateX {
        functions,
        reveal_groups,
        datatypes,
        opaque_types,
        traits,
        trait_impls,
        assoc_type_impls,
        modules: module_ids,
        external_fns,
        external_types,
        path_as_rust_names,
        arch,
    } = &**krate;
    let mut state = State::new();

    // Always add this because unit values might be added later, after ast_simplify.
    state.tuple_type_name(0);

    let mut datatypes = vec_map_result(&datatypes, |d| simplify_datatype(&mut state, d))?;
    ctx.datatypes = Arc::new(
        datatypes
            .iter()
            .map(|d| (d.x.name.expect_path(), (d.x.typ_params.clone(), d.x.variants.clone())))
            .collect(),
    );
    let functions = vec_map_result(functions, |f| simplify_function(ctx, &mut state, f))?;
    let mut trait_impls = vec_map_result(&trait_impls, |t| simplify_trait_impl(&mut state, t))?;
    let mut assoc_type_impls =
        vec_map_result(&assoc_type_impls, |a| simplify_assoc_type_impl(&mut state, a))?;

    let fn_once_trait_path = ClosureKind::FnOnce.trait_path();
    let fn_once_trait_in_scope = traits.iter().any(|t| t.x.name == fn_once_trait_path);

    let mut new_functions: Vec<Function> = Vec::with_capacity(functions.len());
    for f in functions.iter() {
        if need_fndef_axiom(&state.fndef_typs, f) {
            let (f2, tis, ai) =
                add_fndef_axioms_to_function(ctx, &mut state, f, fn_once_trait_in_scope)?;
            trait_impls.extend(tis);
            if let Some(ai) = ai {
                assoc_type_impls.push(ai);
            }
            new_functions.push(f2);
        } else {
            new_functions.push(f.clone());
        }
    }
    let functions = new_functions;

    // Add a generic datatype to represent each tuple arity
    // Iterate in sorted order to get consistent output
    let mut tuples: Vec<usize> = state.tuple_typs.into_iter().collect();
    tuples.sort();
    for arity in tuples {
        let visibility = Visibility { restricted_to: None };
        let transparency = DatatypeTransparency::WhenVisible(visibility.clone());
        let acc = crate::ast::AcceptRecursiveType::RejectInGround;
        let typ_params = Arc::new((0..arity).map(|i| (prefix_tuple_param(i), acc)).collect());
        let mut fields: Vec<Field> = Vec::new();
        for i in 0..arity {
            let typ = Arc::new(TypX::TypParam(prefix_tuple_param(i)));
            let vis = Visibility { restricted_to: None };
            // Note: the mode is irrelevant at this stage, so we arbitrarily use Mode::Exec
            fields.push(ident_binder(&positional_field_ident(i), &(typ, Mode::Exec, vis)));
        }
        let variant = Variant {
            name: prefix_tuple_variant(arity),
            fields: Arc::new(fields),
            ctor_style: CtorPrintStyle::Tuple,
        };
        let variants = Arc::new(vec![variant]);
        let datatypex = DatatypeX {
            name: Dt::Tuple(arity),
            proxy: None,
            visibility,
            owning_module: None,
            transparency,
            typ_params,
            typ_bounds: Arc::new(vec![]),
            variants,
            mode: Mode::Exec,
            ext_equal: arity > 0,
            user_defined_invariant_fn: None,
            sized_constraint: if arity == 0 {
                None
            } else {
                Some(Arc::new(TypX::TypParam(prefix_tuple_param(arity - 1))))
            },
            destructor: false,
        };
        datatypes.push(Spanned::new(ctx.no_span.clone(), datatypex));

        let tuple_path = crate::path![CrateId::Core => "marker", "Tuple"];
        let clone_path = crate::path![CrateId::Core => "clone", "Clone"];
        let copy_path = crate::path![CrateId::Core => "marker", "Copy"];
        //let send_path = crate::path![CrateId::Core => "marker", "Send"];
        //let sync_path = crate::path![CrateId::Core => "marker", "Sync"];
        //let unpin_path = crate::path![CrateId::Core => "marker", "Unpin"];
        add_tuple_auto_impl(ctx, &traits, &mut trait_impls, arity, false, tuple_path);
        add_tuple_auto_impl(ctx, &traits, &mut trait_impls, arity, true, clone_path);
        add_tuple_auto_impl(ctx, &traits, &mut trait_impls, arity, true, copy_path);
        // TODO when when we have full support for auto traits:
        //add_tuple_auto_impl(ctx, &traits, &mut trait_impls, arity, true, send_path);
        //add_tuple_auto_impl(ctx, &traits, &mut trait_impls, arity, true, sync_path);
        //add_tuple_auto_impl(ctx, &traits, &mut trait_impls, arity, true, unpin_path);
    }

    let mut closures: Vec<_> = state.closure_typs.into_iter().collect();
    closures.sort_by_key(|kv| kv.0);
    for (id, closure) in closures {
        // Right now, we translate the closure type into an a global datatype.
        //
        // However, I'm pretty sure an anonymous closure can't actually be referenced
        // from outside the item that defines it (Rust seems to represent it as an
        // "opaque type" if it escapes through an existential type, which Verus currently
        // doesn't support anyway.)
        // So in principle, we could make the type private to the item and not emit any
        // global declarations for it.

        // Also, note that Rust already prohibits a closure type from depending on itself
        // (not even via reference types, which would be allowed for other types).
        // As such, we don't have to worry about any kind of recursion-checking:
        // a closure type cannot possibly be involved in any type cycle.
        // (In principle, the closure should depend negatively on its param and return types,
        // since they are arguments to the 'requires' and 'ensures' predicates, but thanks
        // to Rust's restrictions, we don't have to do any additional checks.)

        let visibility = Visibility { restricted_to: None };
        let transparency = DatatypeTransparency::Never;

        let variants = Arc::new(vec![]);

        let function = functions.iter().find(|f| f.x.name == closure.enclosing_fun).unwrap();

        let typ_params: crate::ast::TypPositives = Arc::new(
            function
                .x
                .typ_params
                .iter()
                .map(|tb| (tb.clone(), crate::ast::AcceptRecursiveType::Accept))
                .collect(),
        );
        let datatypex = DatatypeX {
            name: Dt::Path(closure.path.clone()),
            proxy: None,
            visibility,
            owning_module: None,
            transparency,
            typ_params: typ_params.clone(),
            typ_bounds: function.x.typ_bounds.clone(),
            variants,
            mode: Mode::Exec,
            ext_equal: false,
            user_defined_invariant_fn: None,
            sized_constraint: None,
            destructor: false,
        };
        datatypes.push(Spanned::new(ctx.no_span.clone(), datatypex));

        // Add a trait bound, `ClosureType: {Fn, FnMut, FnOnce}`, and the corresponding
        // `<ClosureType as FnOnce<Args>>::Output = ReturnType` associated-type impl.

        let typ_args: Typs = Arc::new(
            function.x.typ_params.iter().map(|tb| Arc::new(TypX::TypParam(tb.clone()))).collect(),
        );
        let self_typ =
            Arc::new(TypX::Datatype(Dt::Path(closure.path.clone()), typ_args, Arc::new(vec![])));
        let args_tuple_typ = Arc::new(TypX::Datatype(
            Dt::Tuple(closure.args.len()),
            closure.args.clone(),
            Arc::new(vec![]),
        ));
        let impl_path = Arc::new(crate::ast::PathX {
            krate: CrateId::Internal,
            segments: Arc::new(vec![crate::def::impl_closure(closure.kind, id)]),
        });
        let trait_typ_args = Arc::new(vec![self_typ.clone(), args_tuple_typ.clone()]);
        let trait_implx = crate::ast::TraitImplX {
            impl_path: impl_path.clone(),
            typ_params: function.x.typ_params.clone(),
            typ_bounds: function.x.typ_bounds.clone(),
            trait_path: closure.kind.trait_path(),
            trait_typ_args: trait_typ_args.clone(),
            trait_typ_arg_impls: Spanned::new(ctx.no_span.clone(), Arc::new(vec![])),
            owning_module: None,
            auto_imported: true,
            external_trait_blanket: false,
        };
        trait_impls.push(Spanned::new(ctx.no_span.clone(), trait_implx));

        // The `Output` associated type is defined on the `FnOnce` trait (and
        // inherited by `Fn` and `FnMut`), so we always use the `FnOnce` trait
        // path for the associated-type impl, regardless of the closure's kind.
        let fn_once_trait_path = crate::ast::ClosureKind::FnOnce.trait_path();
        if traits.iter().any(|t| &t.x.name == &fn_once_trait_path) {
            let assoc_typ_implx = crate::ast::AssocTypeImplX {
                name: Arc::new("Output".to_string()),
                impl_path,
                typ_params: function.x.typ_params.clone(),
                typ_bounds: function.x.typ_bounds.clone(),
                trait_path: fn_once_trait_path,
                trait_typ_args,
                typ: closure.output.clone(),
                impl_paths: Arc::new(vec![]),
            };
            assoc_type_impls.push(Spanned::new(ctx.no_span.clone(), assoc_typ_implx));
        }
    }

    let traits = traits.clone();
    let module_ids = module_ids.clone();
    let external_fns = external_fns.clone();
    let external_types = external_types.clone();
    let krate = Arc::new(KrateX {
        functions,
        reveal_groups: reveal_groups.clone(),
        datatypes,
        opaque_types: opaque_types.clone(),
        traits,
        trait_impls,
        assoc_type_impls,
        modules: module_ids,
        external_fns,
        external_types,
        path_as_rust_names: path_as_rust_names.clone(),
        arch: arch.clone(),
    });
    *ctx = crate::context::GlobalCtx::new(
        &krate,
        ctx.crate_name.clone(),
        ctx.no_span.clone(),
        ctx.rlimit,
        ctx.interpreter_log.clone(),
        ctx.func_call_graph_log.clone(),
        ctx.warning_ctx.clone(),
        ctx.solver.clone(),
        true,
        ctx.check_api_safety,
        ctx.axiom_usage_info,
        ctx.no_bv_simplify,
        ctx.report_long_running,
    )?;
    Ok(krate)
}

pub fn merge_krates(krates: Vec<Krate>) -> Result<Krate, VirErr> {
    let mut krates = krates.into_iter();
    let mut kratex: KrateX = (*krates.next().expect("at least one crate")).clone();
    for k in krates {
        let KrateX {
            functions,
            reveal_groups,
            datatypes,
            opaque_types,
            traits,
            trait_impls,
            assoc_type_impls,
            modules,
            external_fns,
            external_types,
            path_as_rust_names,
            arch,
        } = &*k;
        kratex.functions.extend(functions.clone());
        kratex.reveal_groups.extend(reveal_groups.clone());
        kratex.datatypes.extend(datatypes.clone());
        kratex.opaque_types.extend(opaque_types.clone());
        kratex.traits.extend(traits.clone());
        kratex.trait_impls.extend(trait_impls.clone());
        kratex.assoc_type_impls.extend(assoc_type_impls.clone());
        kratex.modules.extend(modules.clone());
        kratex.external_fns.extend(external_fns.clone());
        kratex.external_types.extend(external_types.clone());
        kratex.path_as_rust_names.extend(path_as_rust_names.clone());
        kratex.arch.word_bits = {
            let word_bits = match (arch.word_bits, kratex.arch.word_bits) {
                (crate::ast::ArchWordBits::Exactly(l), crate::ast::ArchWordBits::Exactly(r)) => {
                    if l != r {
                        return Err(crate::messages::error_bare(
                            "all crates must have compatible arch_word_bits (set via `global size_of usize`",
                        ));
                    } else {
                        crate::ast::ArchWordBits::Exactly(l)
                    }
                }
                (crate::ast::ArchWordBits::Either32Or64, other)
                | (other, crate::ast::ArchWordBits::Either32Or64) => other,
            };
            if let crate::ast::ArchWordBits::Exactly(e) = &word_bits {
                assert!(*e == 32 || *e == 64);
            }
            word_bits
        };
    }
    Ok(Arc::new(kratex))
}

rust_verify_test coverage (fbbbbcf)

Coverage Report

Created: 2026-08-23 08:37