From 0b8ce0840e10a819e86629251b57575899daf684 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Mon, 11 Dec 2023 20:34:43 -0800
Subject: [PATCH 01/19] In lifetime_generate, only select from
 all_impls(trait_id) on demand

---
 source/rust_verify/src/lifetime_generate.rs   | 203 ++++++++++++------
 .../tests/assoc_type_impls.rs                 |   7 +
 2 files changed, 146 insertions(+), 64 deletions(-)

diff --git a/source/rust_verify/src/lifetime_generate.rs b/source/rust_verify/src/lifetime_generate.rs
index 8e578edb36..8b7947b125 100644
--- a/source/rust_verify/src/lifetime_generate.rs
+++ b/source/rust_verify/src/lifetime_generate.rs
@@ -88,6 +88,13 @@ pub(crate) struct State {
     pub(crate) datatype_decls: Vec<DatatypeDecl>,
     pub(crate) assoc_type_impls: HashMap<AssocTypeImpl, Vec<AssocTypeImplType>>,
     pub(crate) fun_decls: Vec<FunDecl>,
+    // For an impl "bounds ==> trait T(...t...)", point some or all of t to impl:
+    // (We add each impl for T to this when we process T)
+    // (To avoid importing unnecessary impls, we delay processing impl until all t are used)
+    // t1 -> impl, ..., tn -> impl
+    typ_to_trait_impls: HashMap<DefId, Vec<DefId>>,
+    // impl -> (t1, ..., tn) and process impl when t1...tn is empty
+    impl_to_remaining_typs: HashMap<DefId, (Id, Vec<DefId>)>,
     enclosing_fun_id: Option<DefId>,
 }
 
@@ -110,6 +117,8 @@ impl State {
             datatype_decls: Vec::new(),
             assoc_type_impls: HashMap::new(),
             fun_decls: Vec::new(),
+            typ_to_trait_impls: HashMap::new(),
+            impl_to_remaining_typs: HashMap::new(),
             enclosing_fun_id: None,
         }
     }
@@ -205,6 +214,16 @@ impl State {
                 self.reached.insert((None, id));
                 self.datatype_worklist.push(id);
             }
+            if let Some(impl_ids) = self.typ_to_trait_impls.remove(&id) {
+                // Wake up any impls waiting for our type to be reached
+                for impl_id in impl_ids {
+                    if let Some((_, ref mut ts)) = self.impl_to_remaining_typs.get_mut(&impl_id) {
+                        // Remove ourself from what impl_id is waiting on
+                        ts.retain(|t| *t != id);
+                    }
+                    erase_impl_assocs(ctxt, self, impl_id);
+                }
+            }
         }
     }
 
@@ -316,6 +335,56 @@ fn erase_generic_const<'tcx>(ctxt: &Context<'tcx>, state: &mut State, cnst: &Con
     }
 }
 
+fn adt_args<'a, 'tcx>(
+    rust_item: Option<RustItem>,
+    args: &'a [rustc_middle::ty::GenericArg<'tcx>],
+) -> (bool, &'a [rustc_middle::ty::GenericArg<'tcx>]) {
+    if rust_item == Some(RustItem::Box)
+        || rust_item == Some(RustItem::Rc)
+        || rust_item == Some(RustItem::Arc)
+    {
+        // For Box, Rc, Arc, skip the second argument (the Allocator)
+        // which is currently restricted to always be `Global`.
+        assert!(args.len() == 2);
+        (false, &args[0..1])
+    } else {
+        (true, args)
+    }
+}
+
+// Collect some or all as-yet unreached types mentioned directly by ty
+// (It's ok to miss some, but the more we capture, the less extraneous code
+// we have to import.)
+fn collect_unreached_datatypes<'tcx>(
+    ctxt: &Context<'tcx>,
+    state: &State,
+    datatypes: &mut Vec<DefId>,
+    ty: &Ty<'tcx>,
+) {
+    match ty.kind() {
+        TyKind::Ref(_, t, _) | TyKind::Slice(t) | TyKind::Array(t, _) => {
+            collect_unreached_datatypes(ctxt, state, datatypes, t);
+        }
+        TyKind::Adt(AdtDef(adt_def_data), args) => {
+            let did = adt_def_data.did;
+            let rust_item = verus_items::get_rust_item(ctxt.tcx, did);
+            let (is_user_adt, args) = adt_args(rust_item, args);
+            if is_user_adt && !state.reached.contains(&(None, did)) {
+                datatypes.push(did);
+            }
+            for arg in args.iter() {
+                match arg.unpack() {
+                    rustc_middle::ty::GenericArgKind::Type(t) => {
+                        collect_unreached_datatypes(ctxt, state, datatypes, &t);
+                    }
+                    _ => {}
+                }
+            }
+        }
+        _ => {}
+    }
+}
+
 fn erase_ty<'tcx>(ctxt: &Context<'tcx>, state: &mut State, ty: &Ty<'tcx>) -> Typ {
     match ty.kind() {
         TyKind::Bool | TyKind::Uint(_) | TyKind::Int(_) | TyKind::Char | TyKind::Str => {
@@ -349,18 +418,7 @@ fn erase_ty<'tcx>(ctxt: &Context<'tcx>, state: &mut State, ty: &Ty<'tcx>) -> Typ
 
             let rust_item = verus_items::get_rust_item(ctxt.tcx, did);
             let mut typ_args: Vec<Typ> = Vec::new();
-
-            let args = if rust_item == Some(RustItem::Box)
-                || rust_item == Some(RustItem::Rc)
-                || rust_item == Some(RustItem::Arc)
-            {
-                // For Box, Rc, Arc, skip the second argument (the Allocator)
-                // which is currently restricted to always be `Global`.
-                assert!(args.len() == 2);
-                &args[0..1]
-            } else {
-                &args
-            };
+            let (_, args) = adt_args(rust_item, args);
 
             for arg in args.iter() {
                 match arg.unpack() {
@@ -1829,6 +1887,62 @@ fn erase_fn<'tcx>(
     );
 }
 
+fn erase_impl_assocs<'tcx>(ctxt: &Context<'tcx>, state: &mut State, impl_id: DefId) {
+    if !state.impl_to_remaining_typs.contains_key(&impl_id) {
+        // Already finished
+        return;
+    }
+    if state.impl_to_remaining_typs[&impl_id].1.len() > 0 {
+        // We haven't reached all the types we would need to justify this impl
+        return;
+    }
+    let (name, _) = state.impl_to_remaining_typs.remove(&impl_id).unwrap();
+    let trait_ref = ctxt.tcx.impl_trait_ref(impl_id).expect("impl_trait_ref");
+    for assoc_item in ctxt.tcx.associated_items(impl_id).in_definition_order() {
+        match assoc_item.kind {
+            rustc_middle::ty::AssocKind::Type => {
+                let impl_name = state.typ_param(&assoc_item.name.to_string(), None);
+
+                let mut trait_typ_args: Vec<Typ> = Vec::new();
+                for ty in trait_ref.skip_binder().args.types().skip(1) {
+                    trait_typ_args.push(erase_ty(ctxt, state, &ty));
+                }
+                let mut lifetimes: Vec<GenericParam> = Vec::new();
+                let mut typ_params: Vec<GenericParam> = Vec::new();
+                let mut generic_bounds: Vec<GenericBound> = Vec::new();
+                erase_mir_generics(
+                    ctxt,
+                    state,
+                    impl_id,
+                    false,
+                    &mut lifetimes,
+                    &mut typ_params,
+                    &mut generic_bounds,
+                );
+                let mut trait_lifetime_args: Vec<Id> = Vec::new();
+                for region in trait_ref.skip_binder().args.regions() {
+                    if let Some(id) = erase_hir_region(ctxt, state, &region.0) {
+                        trait_lifetime_args.push(id);
+                    }
+                }
+
+                let generic_params = lifetimes.into_iter().chain(typ_params.into_iter()).collect();
+                let self_ty = ctxt.tcx.type_of(impl_id).skip_binder();
+                let self_typ = erase_ty(ctxt, state, &self_ty);
+                let ty = ctxt.tcx.type_of(assoc_item.def_id).skip_binder();
+                let typ = erase_ty(ctxt, state, &ty);
+                let trait_as_datatype =
+                    Box::new(TypX::Datatype(name.clone(), trait_lifetime_args, trait_typ_args));
+                let assoc =
+                    AssocTypeImpl { generic_params, generic_bounds, self_typ, trait_as_datatype };
+                let assoc_type = AssocTypeImplType { name: impl_name, typ };
+                state.assoc_type_impls.entry(assoc).or_insert(Vec::new()).push(assoc_type);
+            }
+            _ => (),
+        }
+    }
+}
+
 fn erase_trait<'tcx>(ctxt: &Context<'tcx>, state: &mut State, trait_id: DefId) {
     if !state.reached.insert((None, trait_id)) {
         return;
@@ -1874,58 +1988,20 @@ fn erase_trait<'tcx>(ctxt: &Context<'tcx>, state: &mut State, trait_id: DefId) {
         state.trait_decls.push(decl);
 
         for impl_id in ctxt.tcx.all_impls(trait_id) {
+            let mut datatypes: Vec<DefId> = Vec::new();
             let trait_ref = ctxt.tcx.impl_trait_ref(impl_id).expect("impl_trait_ref");
-            for assoc_item in ctxt.tcx.associated_items(impl_id).in_definition_order() {
-                match assoc_item.kind {
-                    rustc_middle::ty::AssocKind::Type => {
-                        let impl_name = state.typ_param(&assoc_item.name.to_string(), None);
-
-                        let mut trait_typ_args: Vec<Typ> = Vec::new();
-                        for ty in trait_ref.skip_binder().args.types().skip(1) {
-                            trait_typ_args.push(erase_ty(ctxt, state, &ty));
-                        }
-                        let mut lifetimes: Vec<GenericParam> = Vec::new();
-                        let mut typ_params: Vec<GenericParam> = Vec::new();
-                        let mut generic_bounds: Vec<GenericBound> = Vec::new();
-                        erase_mir_generics(
-                            ctxt,
-                            state,
-                            impl_id,
-                            false,
-                            &mut lifetimes,
-                            &mut typ_params,
-                            &mut generic_bounds,
-                        );
-                        let mut trait_lifetime_args: Vec<Id> = Vec::new();
-                        for region in trait_ref.skip_binder().args.regions() {
-                            if let Some(id) = erase_hir_region(ctxt, state, &region.0) {
-                                trait_lifetime_args.push(id);
-                            }
-                        }
-
-                        let generic_params =
-                            lifetimes.into_iter().chain(typ_params.into_iter()).collect();
-                        let self_ty = ctxt.tcx.type_of(impl_id).skip_binder();
-                        let self_typ = erase_ty(ctxt, state, &self_ty);
-                        let ty = ctxt.tcx.type_of(assoc_item.def_id).skip_binder();
-                        let typ = erase_ty(ctxt, state, &ty);
-                        let trait_as_datatype = Box::new(TypX::Datatype(
-                            name.clone(),
-                            trait_lifetime_args,
-                            trait_typ_args,
-                        ));
-                        let assoc = AssocTypeImpl {
-                            generic_params,
-                            generic_bounds,
-                            self_typ,
-                            trait_as_datatype,
-                        };
-                        let assoc_type = AssocTypeImplType { name: impl_name, typ };
-                        state.assoc_type_impls.entry(assoc).or_insert(Vec::new()).push(assoc_type);
+            for ty in trait_ref.skip_binder().args.types() {
+                collect_unreached_datatypes(ctxt, state, &mut datatypes, &ty);
+                for t in &datatypes {
+                    if !state.typ_to_trait_impls.contains_key(t) {
+                        state.typ_to_trait_impls.insert(*t, Vec::new());
                     }
-                    _ => (),
+                    state.typ_to_trait_impls.get_mut(&t).unwrap().push(impl_id);
                 }
             }
+            assert!(!state.impl_to_remaining_typs.contains_key(&impl_id));
+            state.impl_to_remaining_typs.insert(impl_id, (name.clone(), datatypes));
+            erase_impl_assocs(ctxt, state, impl_id);
         }
     }
 }
@@ -2144,10 +2220,9 @@ fn erase_mir_datatype<'tcx>(ctxt: &Context<'tcx>, state: &mut State, id: DefId)
             DatatypeTransparency::Never => true,
             DatatypeTransparency::WhenVisible(_) => false,
         },
-        None => {
-            dbg!(id);
-            panic!("erase_mir_datatype: unrecognized datatype");
-        }
+        // We may see extra datatypes from imported libraries that we
+        // use for associated types:
+        None => true,
     };
     if is_external_body {
         erase_abstract_datatype(ctxt, state, span, id);
diff --git a/source/rust_verify_test/tests/assoc_type_impls.rs b/source/rust_verify_test/tests/assoc_type_impls.rs
index 37b0f89743..5ca319117e 100644
--- a/source/rust_verify_test/tests/assoc_type_impls.rs
+++ b/source/rust_verify_test/tests/assoc_type_impls.rs
@@ -376,3 +376,10 @@ test_verify_one_file! {
         }
     } => Err(err) => assert_vir_error_msg(err, "found a cyclic self-reference in a trait definition, which may result in nontermination")
 }
+
+test_verify_one_file! {
+    #[test] mention_external_trait_with_assoc_type verus_code! {
+        fn foo<A: IntoIterator>(a: &A) {
+        }
+    } => Ok(())
+}

From bda76f8a9b5a07d31fb2b37518fbf03db2ca1cf1 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Wed, 13 Dec 2023 13:40:11 -0800
Subject: [PATCH 02/19] Support into_iter for Range

---
 source/pervasive/std_specs/core.rs          |   9 +-
 source/pervasive/std_specs/mod.rs           |   1 +
 source/pervasive/std_specs/range.rs         |  90 +++++++++++++++++
 source/pervasive/vstd.rs                    |   1 +
 source/rust_verify/src/fn_call_to_vir.rs    |   6 --
 source/rust_verify_test/tests/integers.rs   | 104 ++++++++++++++++++++
 source/rust_verify_test/tests/loops.rs      |  51 ++++++++++
 source/rust_verify_test/tests/regression.rs |   2 +-
 8 files changed, 254 insertions(+), 10 deletions(-)
 create mode 100644 source/pervasive/std_specs/range.rs

diff --git a/source/pervasive/std_specs/core.rs b/source/pervasive/std_specs/core.rs
index f982f4d65d..029d51fadb 100644
--- a/source/pervasive/std_specs/core.rs
+++ b/source/pervasive/std_specs/core.rs
@@ -19,9 +19,12 @@ pub struct ExOption<V>(core::option::Option<V>);
 #[verifier::reject_recursive_types_in_ground_variants(E)]
 pub struct ExResult<T, E>(core::result::Result<T, E>);
 
-#[verifier(external_type_specification)]
-#[verifier::reject_recursive_types_in_ground_variants(Idx)]
-pub struct ExRange<Idx>(core::ops::Range<Idx>);
+#[verifier::external_fn_specification]
+pub fn ex_iter_into_iter<I: Iterator>(i: I) -> (r: I)
+    ensures r == i
+{
+    i.into_iter()
+}
 
 // I don't really expect this to be particularly useful;
 // this is mostly here because I wanted an easy way to test
diff --git a/source/pervasive/std_specs/mod.rs b/source/pervasive/std_specs/mod.rs
index 25afdc4ce5..7cde07157e 100644
--- a/source/pervasive/std_specs/mod.rs
+++ b/source/pervasive/std_specs/mod.rs
@@ -1,5 +1,6 @@
 pub mod core;
 pub mod num;
+pub mod range;
 pub mod result;
 pub mod option;
 pub mod atomic;
diff --git a/source/pervasive/std_specs/range.rs b/source/pervasive/std_specs/range.rs
new file mode 100644
index 0000000000..47a4647871
--- /dev/null
+++ b/source/pervasive/std_specs/range.rs
@@ -0,0 +1,90 @@
+use crate::prelude::*;
+use core::ops::Range;
+
+verus!{
+
+#[verifier(external_type_specification)]
+#[verifier::reject_recursive_types_in_ground_variants(Idx)]
+pub struct ExRange<Idx>(Range<Idx>);
+
+pub trait StepSpec where Self: Sized {
+    // REVIEW: it would be nice to be able to use SpecOrd::spec_lt (not yet supported)
+    spec fn spec_is_lt(self, other: Self) -> bool;
+    spec fn spec_steps_between(self, end: Self) -> Option<usize>;
+    spec fn spec_forward_checked(self, count: usize) -> Option<Self>;
+    spec fn spec_backward_checked(self, count: usize) -> Option<Self>;
+}
+
+pub spec fn spec_range_next<A>(a: Range<A>) -> (Range<A>, Option<A>);
+
+#[verifier::external_body]
+#[verifier::broadcast_forall]
+pub proof fn axiom_spec_range_next<A: StepSpec>(range: Range<A>)
+    ensures
+        range.start.spec_is_lt(range.end) ==>
+            // TODO (not important): use new "matches ==>" syntax here
+            (if let Some(n) = range.start.spec_forward_checked(1) {
+                spec_range_next(range) == (Range { start: n, ..range }, Some(range.start))
+            } else {
+                true
+            }),
+        !range.start.spec_is_lt(range.end) ==>
+            #[trigger] spec_range_next(range) == (range, None::<A>),
+{
+}
+
+#[verifier::external_fn_specification]
+pub fn ex_range_next<A: core::iter::Step>(range: &mut Range<A>) -> (r: Option<A>)
+    ensures (*range, r) == spec_range_next(*old(range))
+{
+    range.next()
+}
+
+} // verus!
+
+macro_rules! step_specs {
+    ($t: ty) => {
+        verus! {
+        impl StepSpec for $t {
+            open spec fn spec_is_lt(self, other: Self) -> bool {
+                self < other
+            }
+            open spec fn spec_steps_between(self, end: Self) -> Option<usize> {
+                let n = end - self;
+                if usize::MIN <= n <= usize::MAX {
+                    Some(n as usize)
+                } else {
+                    None
+                }
+            }
+            open spec fn spec_forward_checked(self, count: usize) -> Option<Self> {
+                if self + count <= $t::MAX {
+                    Some((self + count) as $t)
+                } else {
+                    None
+                }
+            }
+            open spec fn spec_backward_checked(self, count: usize) -> Option<Self> {
+                if self - count >= $t::MIN {
+                    Some((self - count) as $t)
+                } else {
+                    None
+                }
+            }
+        } // verus!
+        }
+    }
+}
+
+step_specs!(u8);
+step_specs!(u16);
+step_specs!(u32);
+step_specs!(u64);
+step_specs!(u128);
+step_specs!(usize);
+step_specs!(i8);
+step_specs!(i16);
+step_specs!(i32);
+step_specs!(i64);
+step_specs!(i128);
+step_specs!(isize);
diff --git a/source/pervasive/vstd.rs b/source/pervasive/vstd.rs
index 6d775067f9..27bc16441a 100644
--- a/source/pervasive/vstd.rs
+++ b/source/pervasive/vstd.rs
@@ -2,6 +2,7 @@
 //! Contains various utilities and datatypes for proofs,
 //! as well as runtime functionality with specifications.
 //! For an introduction to Verus, see [the tutorial](https://verus-lang.github.io/verus/guide/).
+#![feature(step_trait)]
 
 #![cfg_attr(not(feature = "std"), no_std)]
 #![allow(unused_parens)]
diff --git a/source/rust_verify/src/fn_call_to_vir.rs b/source/rust_verify/src/fn_call_to_vir.rs
index 01f1fb1b49..786c0cb514 100644
--- a/source/rust_verify/src/fn_call_to_vir.rs
+++ b/source/rust_verify/src/fn_call_to_vir.rs
@@ -99,12 +99,6 @@ pub(crate) fn fn_call_to_vir<'tcx>(
         Some(RustItem::TryTraitBranch) => {
             return err_span(expr.span, "Verus does not yet support the ? operator");
         }
-        Some(RustItem::IntoIterFn) => {
-            return err_span(
-                expr.span,
-                "Verus does not yet support IntoIterator::into_iter and for loops, use a while loop instead",
-            );
-        }
         Some(RustItem::Clone) => {
             // Special case `clone` for standard Rc and Arc types
             // (Could also handle it for other types where cloning is the identity
diff --git a/source/rust_verify_test/tests/integers.rs b/source/rust_verify_test/tests/integers.rs
index 79a6fa50e0..d79c72dc52 100644
--- a/source/rust_verify_test/tests/integers.rs
+++ b/source/rust_verify_test/tests/integers.rs
@@ -248,3 +248,107 @@ test_verify_one_file! {
         }
     } => Err(err) => assert_vir_error_msg(err, "integer literal out of range")
 }
+
+test_verify_one_file! {
+    #[test] test_step verus_code! {
+        use vstd::std_specs::range::*;
+        spec fn and_then<A, B>(o: Option<A>, f: FnSpec(A) -> Option<B>) -> Option<B> {
+            if let Some(a) = o {
+                f(a)
+            } else {
+                None
+            }
+        }
+        spec fn checked_add_usize(a: usize, b: usize) -> Option<usize> {
+            if 0 <= a + b <= usize::MAX { Some((a + b) as usize) } else { None }
+        }
+        proof fn step_properties_u8(a: u8, b: u8, n: usize, m: usize) {
+            // These are specified by `pub trait Step` in Rust's range.rs:
+            assert(a.spec_steps_between(b) == Some(n) <==> a.spec_forward_checked(n) == Some(b));
+            assert(a.spec_steps_between(b) == Some(n) <==> b.spec_backward_checked(n) == Some(a));
+            assert(a.spec_steps_between(b) == Some(n) ==> a <= b);
+            assert(a.spec_steps_between(b) == None::<usize> <==> a > b);
+            assert(and_then(a.spec_forward_checked(n), |x: u8| x.spec_forward_checked(m))
+                == and_then(a.spec_forward_checked(m), |x: u8| x.spec_forward_checked(n)));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_forward_checked(n), |x: u8| x.spec_forward_checked(m))
+                == a.spec_forward_checked((n + m) as usize));
+            assert(n + m > usize::MAX ==> and_then(a.spec_forward_checked(n), |x: u8| x.spec_forward_checked(m))
+                == None::<u8>);
+            assert(and_then(a.spec_backward_checked(n), |x: u8| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize| a.spec_backward_checked(x)));
+            assert(and_then(a.spec_backward_checked(n), |x: u8| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize|
+                    if let Some(z) = checked_add_usize(n, m) { a.spec_backward_checked(z) } else { None }));
+            // Additional checks:
+            assert(a < 255 ==> a.spec_forward_checked(1) == Some((a + 1) as u8));
+            assert(a >= 255 ==> a.spec_forward_checked(1) == None::<u8>);
+            assert(a >= 1 ==> a.spec_backward_checked(1) == Some((a - 1) as u8));
+            assert(a < 1 ==> a.spec_backward_checked(1) == None::<u8>);
+        }
+        proof fn step_properties_i8(a: i8, b: i8, n: usize, m: usize) {
+            // These are specified by `pub trait Step` in Rust's range.rs:
+            assert(a.spec_steps_between(b) == Some(n) <==> a.spec_forward_checked(n) == Some(b));
+            assert(a.spec_steps_between(b) == Some(n) <==> b.spec_backward_checked(n) == Some(a));
+            assert(a.spec_steps_between(b) == Some(n) ==> a <= b);
+            assert(a.spec_steps_between(b) == None::<usize> <==> a > b);
+            assert(and_then(a.spec_forward_checked(n), |x: i8| x.spec_forward_checked(m))
+                == and_then(a.spec_forward_checked(m), |x: i8| x.spec_forward_checked(n)));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_forward_checked(n), |x: i8| x.spec_forward_checked(m))
+                == a.spec_forward_checked((n + m) as usize));
+            assert(n + m > usize::MAX ==> and_then(a.spec_forward_checked(n), |x: i8| x.spec_forward_checked(m))
+                == None::<i8>);
+            assert(and_then(a.spec_backward_checked(n), |x: i8| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize| a.spec_backward_checked(x)));
+            assert(and_then(a.spec_backward_checked(n), |x: i8| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize|
+                    if let Some(z) = checked_add_usize(n, m) { a.spec_backward_checked(z) } else { None }));
+            // Additional checks:
+            assert(a < 127 ==> a.spec_forward_checked(1) == Some((a + 1) as i8));
+            assert(a >= 127 ==> a.spec_forward_checked(1) == None::<i8>);
+            assert(a >= -127 ==> a.spec_backward_checked(1) == Some((a - 1) as i8));
+            assert(a < -127 ==> a.spec_backward_checked(1) == None::<i8>);
+        }
+        proof fn step_properties_u128(a: u128, b: u128, n: usize, m: usize) {
+            // These are specified by `pub trait Step` in Rust's range.rs:
+            assert(a.spec_steps_between(b) == Some(n) <==> a.spec_forward_checked(n) == Some(b));
+            assert(a.spec_steps_between(b) == Some(n) <==> b.spec_backward_checked(n) == Some(a));
+            assert(a.spec_steps_between(b) == Some(n) ==> a <= b);
+            assert(a.spec_steps_between(b) == None::<usize> <== a > b);
+            assert(and_then(a.spec_forward_checked(n), |x: u128| x.spec_forward_checked(m))
+                == and_then(a.spec_forward_checked(m), |x: u128| x.spec_forward_checked(n)));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_forward_checked(n), |x: u128| x.spec_forward_checked(m))
+                == a.spec_forward_checked((n + m) as usize));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_backward_checked(n), |x: u128| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize| a.spec_backward_checked(x)));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_backward_checked(n), |x: u128| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize|
+                    if let Some(z) = checked_add_usize(n, m) { a.spec_backward_checked(z) } else { None }));
+            // Additional checks:
+            assert(a < 0xffff_ffff_ffff_ffff_ffff_ffff_ffff_ffff ==> a.spec_forward_checked(1) == Some((a + 1) as u128));
+            assert(a >= 0xffff_ffff_ffff_ffff_ffff_ffff_ffff_ffff ==> a.spec_forward_checked(1) == None::<u128>);
+            assert(a >= 1 ==> a.spec_backward_checked(1) == Some((a - 1) as u128));
+            assert(a < 1 ==> a.spec_backward_checked(1) == None::<u128>);
+        }
+        proof fn step_properties_i128(a: i128, b: i128, n: usize, m: usize) {
+            // These are specified by `pub trait Step` in Rust's range.rs:
+            assert(a.spec_steps_between(b) == Some(n) <==> a.spec_forward_checked(n) == Some(b));
+            assert(a.spec_steps_between(b) == Some(n) <==> b.spec_backward_checked(n) == Some(a));
+            assert(a.spec_steps_between(b) == Some(n) ==> a <= b);
+            assert(a.spec_steps_between(b) == None::<usize> <== a > b);
+            assert(and_then(a.spec_forward_checked(n), |x: i128| x.spec_forward_checked(m))
+                == and_then(a.spec_forward_checked(m), |x: i128| x.spec_forward_checked(n)));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_forward_checked(n), |x: i128| x.spec_forward_checked(m))
+                == a.spec_forward_checked((n + m) as usize));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_backward_checked(n), |x: i128| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize| a.spec_backward_checked(x)));
+            assert(n + m <= usize::MAX ==> and_then(a.spec_backward_checked(n), |x: i128| x.spec_backward_checked(m))
+                == and_then(checked_add_usize(n, m), |x: usize|
+                    if let Some(z) = checked_add_usize(n, m) { a.spec_backward_checked(z) } else { None }));
+            // Additional checks:
+            assert(a < 0x7fff_ffff_ffff_ffff_ffff_ffff_ffff_ffff ==> a.spec_forward_checked(1) == Some((a + 1) as i128));
+            assert(a >= 0x7fff_ffff_ffff_ffff_ffff_ffff_ffff_ffff ==> a.spec_forward_checked(1) == None::<i128>);
+            assert(a >= -0x7fff_ffff_ffff_ffff_ffff_ffff_ffff_ffff ==> a.spec_backward_checked(1) == Some((a - 1) as i128));
+            assert(a < -0x7fff_ffff_ffff_ffff_ffff_ffff_ffff_ffff ==> a.spec_backward_checked(1) == None::<i128>);
+        }
+    } => Ok(())
+}
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index ec0f66052c..fa3469e493 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -932,3 +932,54 @@ test_verify_one_file! {
         }
     } => Ok(())
 }
+
+test_verify_one_file! {
+    #[test] iter_loop verus_code! {
+        use vstd::prelude::*;
+        fn test_loop() {
+            let mut n: u64 = 0;
+            let mut iter = (0..10).into_iter();
+            loop
+                invariant_ensures
+                    iter.start <= 10,
+                    iter.end == 10,
+                    n == iter.start * 3,
+                ensures
+                    iter.start == 10,
+            {
+                if let Some(x) = iter.next() {
+                    assert(x < 10);
+                    assert(x == iter.start - 1);
+                    n += 3;
+                } else {
+                    break;
+                }
+            }
+            assert(iter.start == 10);
+            assert(n == 30);
+        }
+
+        fn test_loop_fail() {
+            let mut n: u64 = 0;
+            let mut iter = (0..10).into_iter();
+            loop
+                invariant_ensures
+                    iter.start <= 10,
+                    iter.end == 10,
+                    n == iter.start * 3,
+                ensures
+                    iter.start == 10,
+            {
+                if let Some(x) = iter.next() {
+                    assert(x < 9); // FAILS
+                    assert(x == iter.start - 1);
+                    n += 3;
+                } else {
+                    break;
+                }
+            }
+            assert(iter.start == 10);
+            assert(n == 30);
+        }
+    } => Err(e) => assert_one_fails(e)
+}
diff --git a/source/rust_verify_test/tests/regression.rs b/source/rust_verify_test/tests/regression.rs
index 7840dd39a0..8ed20d950e 100644
--- a/source/rust_verify_test/tests/regression.rs
+++ b/source/rust_verify_test/tests/regression.rs
@@ -543,7 +543,7 @@ test_verify_one_file! {
 }
 
 test_verify_one_file! {
-    #[test] test_for_loop_387_discussioncomment_5683342 verus_code! {
+    #[ignore] #[test] test_for_loop_387_discussioncomment_5683342 verus_code! {
         struct T{}
         fn f(v: Vec<T>) {
             for t in v {}

From d1085d2534fcdfe57590a026a1ed945fefd40600 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Wed, 13 Dec 2023 13:59:26 -0800
Subject: [PATCH 03/19] Mark feature(step_trait) as cfg_attr(verus_keep_ghost)

---
 source/pervasive/vstd.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/pervasive/vstd.rs b/source/pervasive/vstd.rs
index 27bc16441a..b9d4f028d9 100644
--- a/source/pervasive/vstd.rs
+++ b/source/pervasive/vstd.rs
@@ -2,7 +2,6 @@
 //! Contains various utilities and datatypes for proofs,
 //! as well as runtime functionality with specifications.
 //! For an introduction to Verus, see [the tutorial](https://verus-lang.github.io/verus/guide/).
-#![feature(step_trait)]
 
 #![cfg_attr(not(feature = "std"), no_std)]
 #![allow(unused_parens)]
@@ -13,6 +12,7 @@
 
 #![cfg_attr(verus_keep_ghost, feature(core_intrinsics))]
 #![cfg_attr(verus_keep_ghost, feature(allocator_api))]
+#![cfg_attr(verus_keep_ghost, feature(step_trait))]
 
 #[cfg(feature = "alloc")]
 extern crate alloc;

From 65847ae3d9781f5a45ff31ea4766fcb923a332fe Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Tue, 2 Jan 2024 23:05:34 -0800
Subject: [PATCH 04/19] Initial support for for-loops

---
 dependencies/syn/src/expr.rs                | 20 +++++
 dependencies/syn/src/gen/clone.rs           |  3 +
 dependencies/syn/src/gen/debug.rs           |  3 +
 dependencies/syn/src/gen/eq.rs              |  4 +-
 dependencies/syn/src/gen/fold.rs            |  7 ++
 dependencies/syn/src/gen/hash.rs            |  3 +
 dependencies/syn/src/gen/visit.rs           | 10 +++
 dependencies/syn/src/gen/visit_mut.rs       | 10 +++
 dependencies/syn/syn.json                   | 24 +++++
 dependencies/syn/tests/debug/gen.rs         | 96 ++++++++++++++++++++
 source/builtin/src/lib.rs                   |  7 ++
 source/builtin_macros/src/syntax.rs         | 99 +++++++++++++++++++++
 source/pervasive/pervasive.rs               | 14 +++
 source/pervasive/std_specs/core.rs          |  5 ++
 source/pervasive/std_specs/range.rs         | 24 +++++
 source/rust_verify/src/fn_call_to_vir.rs    |  6 ++
 source/rust_verify/src/lifetime_generate.rs | 59 ++----------
 source/rust_verify/src/rust_to_vir_base.rs  |  2 +
 source/rust_verify/src/verifier.rs          |  6 ++
 source/rust_verify/src/verus_items.rs       |  2 +
 source/rust_verify_test/tests/loops.rs      | 64 +++++++++++++
 source/vir/src/ast.rs                       |  9 ++
 source/vir/src/ast_simplify.rs              | 11 +++
 source/vir/src/ast_to_sst.rs                |  5 ++
 source/vir/src/context.rs                   | 13 ++-
 source/vir/src/def.rs                       |  7 ++
 source/vir/src/interpreter.rs               |  6 +-
 source/vir/src/lib.rs                       |  1 +
 source/vir/src/loop_inference.rs            | 65 ++++++++++++++
 source/vir/src/modes.rs                     | 19 +++-
 source/vir/src/poly.rs                      |  5 ++
 source/vir/src/prune.rs                     |  4 +
 source/vir/src/sst_to_air.rs                | 59 ++++++++----
 source/vir/src/sst_util.rs                  |  7 +-
 source/vir/src/triggers.rs                  |  6 +-
 source/vir/src/triggers_auto.rs             |  1 +
 36 files changed, 610 insertions(+), 76 deletions(-)
 create mode 100644 source/vir/src/loop_inference.rs

diff --git a/dependencies/syn/src/expr.rs b/dependencies/syn/src/expr.rs
index 235387743c..92b047e409 100644
--- a/dependencies/syn/src/expr.rs
+++ b/dependencies/syn/src/expr.rs
@@ -462,7 +462,10 @@ ast_struct! {
         pub for_token: Token![for],
         pub pat: Pat,
         pub in_token: Token![in],
+        pub expr_name: Option<Box<(Ident, Token![:])>>,
         pub expr: Box<Expr>,
+        pub invariant: Option<Invariant>,
+        pub decreases: Option<Decreases>,
         pub body: Block,
     }
 }
@@ -2435,7 +2438,14 @@ pub(crate) mod parsing {
             let pat = pat::parsing::multi_pat_with_leading_vert(input)?;
 
             let in_token: Token![in] = input.parse()?;
+            let expr_name = if input.peek2(Token![:]) && input.peek(Ident) {
+                Some(Box::new((input.parse()?, input.parse()?)))
+            } else {
+                None
+            };
             let expr: Expr = input.call(Expr::parse_without_eager_brace)?;
+            let invariant = input.parse()?;
+            let decreases = input.parse()?;
 
             let content;
             let brace_token = braced!(content in input);
@@ -2448,7 +2458,10 @@ pub(crate) mod parsing {
                 for_token,
                 pat,
                 in_token,
+                expr_name,
                 expr: Box::new(expr),
+                invariant,
+                decreases,
                 body: Block { brace_token, stmts },
             })
         }
@@ -3444,7 +3457,14 @@ pub(crate) mod printing {
             self.for_token.to_tokens(tokens);
             self.pat.to_tokens(tokens);
             self.in_token.to_tokens(tokens);
+            if let Some(expr_name) = &self.expr_name {
+                let (name, colon) = &**expr_name;
+                name.to_tokens(tokens);
+                colon.to_tokens(tokens);
+            }
             wrap_bare_struct(tokens, &self.expr);
+            self.invariant.to_tokens(tokens);
+            self.decreases.to_tokens(tokens);
             self.body.brace_token.surround(tokens, |tokens| {
                 inner_attrs_to_tokens(&self.attrs, tokens);
                 tokens.append_all(&self.body.stmts);
diff --git a/dependencies/syn/src/gen/clone.rs b/dependencies/syn/src/gen/clone.rs
index b0748b0812..b37f321aa2 100644
--- a/dependencies/syn/src/gen/clone.rs
+++ b/dependencies/syn/src/gen/clone.rs
@@ -574,7 +574,10 @@ impl Clone for ExprForLoop {
             for_token: self.for_token.clone(),
             pat: self.pat.clone(),
             in_token: self.in_token.clone(),
+            expr_name: self.expr_name.clone(),
             expr: self.expr.clone(),
+            invariant: self.invariant.clone(),
+            decreases: self.decreases.clone(),
             body: self.body.clone(),
         }
     }
diff --git a/dependencies/syn/src/gen/debug.rs b/dependencies/syn/src/gen/debug.rs
index ba00a49544..4868497659 100644
--- a/dependencies/syn/src/gen/debug.rs
+++ b/dependencies/syn/src/gen/debug.rs
@@ -983,7 +983,10 @@ impl Debug for ExprForLoop {
         formatter.field("for_token", &self.for_token);
         formatter.field("pat", &self.pat);
         formatter.field("in_token", &self.in_token);
+        formatter.field("expr_name", &self.expr_name);
         formatter.field("expr", &self.expr);
+        formatter.field("invariant", &self.invariant);
+        formatter.field("decreases", &self.decreases);
         formatter.field("body", &self.body);
         formatter.finish()
     }
diff --git a/dependencies/syn/src/gen/eq.rs b/dependencies/syn/src/gen/eq.rs
index da295c06b4..05db6aea91 100644
--- a/dependencies/syn/src/gen/eq.rs
+++ b/dependencies/syn/src/gen/eq.rs
@@ -568,7 +568,9 @@ impl Eq for ExprForLoop {}
 impl PartialEq for ExprForLoop {
     fn eq(&self, other: &Self) -> bool {
         self.attrs == other.attrs && self.label == other.label && self.pat == other.pat
-            && self.expr == other.expr && self.body == other.body
+            && self.expr_name == other.expr_name && self.expr == other.expr
+            && self.invariant == other.invariant && self.decreases == other.decreases
+            && self.body == other.body
     }
 }
 #[cfg(feature = "full")]
diff --git a/dependencies/syn/src/gen/fold.rs b/dependencies/syn/src/gen/fold.rs
index a4f4a6711f..37b320ff30 100644
--- a/dependencies/syn/src/gen/fold.rs
+++ b/dependencies/syn/src/gen/fold.rs
@@ -1603,7 +1603,14 @@ where
         for_token: Token![for](tokens_helper(f, &node.for_token.span)),
         pat: f.fold_pat(node.pat),
         in_token: Token![in](tokens_helper(f, &node.in_token.span)),
+        expr_name: (node.expr_name)
+            .map(|it| Box::new((
+                f.fold_ident((*it).0),
+                Token![:](tokens_helper(f, &(*it).1.spans)),
+            ))),
         expr: Box::new(f.fold_expr(*node.expr)),
+        invariant: (node.invariant).map(|it| f.fold_invariant(it)),
+        decreases: (node.decreases).map(|it| f.fold_decreases(it)),
         body: f.fold_block(node.body),
     }
 }
diff --git a/dependencies/syn/src/gen/hash.rs b/dependencies/syn/src/gen/hash.rs
index 79a4835c69..ae9a4b3a9d 100644
--- a/dependencies/syn/src/gen/hash.rs
+++ b/dependencies/syn/src/gen/hash.rs
@@ -865,7 +865,10 @@ impl Hash for ExprForLoop {
         self.attrs.hash(state);
         self.label.hash(state);
         self.pat.hash(state);
+        self.expr_name.hash(state);
         self.expr.hash(state);
+        self.invariant.hash(state);
+        self.decreases.hash(state);
         self.body.hash(state);
     }
 }
diff --git a/dependencies/syn/src/gen/visit.rs b/dependencies/syn/src/gen/visit.rs
index a2d5e51a65..75cfe46911 100644
--- a/dependencies/syn/src/gen/visit.rs
+++ b/dependencies/syn/src/gen/visit.rs
@@ -1738,7 +1738,17 @@ where
     tokens_helper(v, &node.for_token.span);
     v.visit_pat(&node.pat);
     tokens_helper(v, &node.in_token.span);
+    if let Some(it) = &node.expr_name {
+        v.visit_ident(&(**it).0);
+        tokens_helper(v, &(**it).1.spans);
+    }
     v.visit_expr(&*node.expr);
+    if let Some(it) = &node.invariant {
+        v.visit_invariant(it);
+    }
+    if let Some(it) = &node.decreases {
+        v.visit_decreases(it);
+    }
     v.visit_block(&node.body);
 }
 #[cfg(feature = "full")]
diff --git a/dependencies/syn/src/gen/visit_mut.rs b/dependencies/syn/src/gen/visit_mut.rs
index 849084e725..cc642a9e10 100644
--- a/dependencies/syn/src/gen/visit_mut.rs
+++ b/dependencies/syn/src/gen/visit_mut.rs
@@ -1739,7 +1739,17 @@ where
     tokens_helper(v, &mut node.for_token.span);
     v.visit_pat_mut(&mut node.pat);
     tokens_helper(v, &mut node.in_token.span);
+    if let Some(it) = &mut node.expr_name {
+        v.visit_ident_mut(&mut (**it).0);
+        tokens_helper(v, &mut (**it).1.spans);
+    }
     v.visit_expr_mut(&mut *node.expr);
+    if let Some(it) = &mut node.invariant {
+        v.visit_invariant_mut(it);
+    }
+    if let Some(it) = &mut node.decreases {
+        v.visit_decreases_mut(it);
+    }
     v.visit_block_mut(&mut node.body);
 }
 #[cfg(feature = "full")]
diff --git a/dependencies/syn/syn.json b/dependencies/syn/syn.json
index e17417bfd3..4db20fffdf 100644
--- a/dependencies/syn/syn.json
+++ b/dependencies/syn/syn.json
@@ -1560,11 +1560,35 @@
         "in_token": {
           "token": "In"
         },
+        "expr_name": {
+          "option": {
+            "box": {
+              "tuple": [
+                {
+                  "proc_macro2": "Ident"
+                },
+                {
+                  "token": "Colon"
+                }
+              ]
+            }
+          }
+        },
         "expr": {
           "box": {
             "syn": "Expr"
           }
         },
+        "invariant": {
+          "option": {
+            "syn": "Invariant"
+          }
+        },
+        "decreases": {
+          "option": {
+            "syn": "Decreases"
+          }
+        },
         "body": {
           "syn": "Block"
         }
diff --git a/dependencies/syn/tests/debug/gen.rs b/dependencies/syn/tests/debug/gen.rs
index fcfcdd36e5..e619e4cfd8 100644
--- a/dependencies/syn/tests/debug/gen.rs
+++ b/dependencies/syn/tests/debug/gen.rs
@@ -948,7 +948,55 @@ impl Debug for Lite<syn::Expr> {
                     formatter.field("label", Print::ref_cast(val));
                 }
                 formatter.field("pat", Lite(&_val.pat));
+                if let Some(val) = &_val.expr_name {
+                    #[derive(RefCast)]
+                    #[repr(transparent)]
+                    struct Print(Box<(proc_macro2::Ident, syn::token::Colon)>);
+                    impl Debug for Print {
+                        fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                            formatter.write_str("Some")?;
+                            let _val = &self.0;
+                            formatter.write_str("(")?;
+                            Debug::fmt(Lite(_val), formatter)?;
+                            formatter.write_str(")")?;
+                            Ok(())
+                        }
+                    }
+                    formatter.field("expr_name", Print::ref_cast(val));
+                }
                 formatter.field("expr", Lite(&_val.expr));
+                if let Some(val) = &_val.invariant {
+                    #[derive(RefCast)]
+                    #[repr(transparent)]
+                    struct Print(syn::Invariant);
+                    impl Debug for Print {
+                        fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                            formatter.write_str("Some")?;
+                            let _val = &self.0;
+                            formatter.write_str("(")?;
+                            Debug::fmt(Lite(_val), formatter)?;
+                            formatter.write_str(")")?;
+                            Ok(())
+                        }
+                    }
+                    formatter.field("invariant", Print::ref_cast(val));
+                }
+                if let Some(val) = &_val.decreases {
+                    #[derive(RefCast)]
+                    #[repr(transparent)]
+                    struct Print(syn::Decreases);
+                    impl Debug for Print {
+                        fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                            formatter.write_str("Some")?;
+                            let _val = &self.0;
+                            formatter.write_str("(")?;
+                            Debug::fmt(Lite(_val), formatter)?;
+                            formatter.write_str(")")?;
+                            Ok(())
+                        }
+                    }
+                    formatter.field("decreases", Print::ref_cast(val));
+                }
                 formatter.field("body", Lite(&_val.body));
                 formatter.finish()
             }
@@ -1890,7 +1938,55 @@ impl Debug for Lite<syn::ExprForLoop> {
             formatter.field("label", Print::ref_cast(val));
         }
         formatter.field("pat", Lite(&_val.pat));
+        if let Some(val) = &_val.expr_name {
+            #[derive(RefCast)]
+            #[repr(transparent)]
+            struct Print(Box<(proc_macro2::Ident, syn::token::Colon)>);
+            impl Debug for Print {
+                fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                    formatter.write_str("Some")?;
+                    let _val = &self.0;
+                    formatter.write_str("(")?;
+                    Debug::fmt(Lite(_val), formatter)?;
+                    formatter.write_str(")")?;
+                    Ok(())
+                }
+            }
+            formatter.field("expr_name", Print::ref_cast(val));
+        }
         formatter.field("expr", Lite(&_val.expr));
+        if let Some(val) = &_val.invariant {
+            #[derive(RefCast)]
+            #[repr(transparent)]
+            struct Print(syn::Invariant);
+            impl Debug for Print {
+                fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                    formatter.write_str("Some")?;
+                    let _val = &self.0;
+                    formatter.write_str("(")?;
+                    Debug::fmt(Lite(_val), formatter)?;
+                    formatter.write_str(")")?;
+                    Ok(())
+                }
+            }
+            formatter.field("invariant", Print::ref_cast(val));
+        }
+        if let Some(val) = &_val.decreases {
+            #[derive(RefCast)]
+            #[repr(transparent)]
+            struct Print(syn::Decreases);
+            impl Debug for Print {
+                fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                    formatter.write_str("Some")?;
+                    let _val = &self.0;
+                    formatter.write_str("(")?;
+                    Debug::fmt(Lite(_val), formatter)?;
+                    formatter.write_str(")")?;
+                    Ok(())
+                }
+            }
+            formatter.field("decreases", Print::ref_cast(val));
+        }
         formatter.field("body", Lite(&_val.body));
         formatter.finish()
     }
diff --git a/source/builtin/src/lib.rs b/source/builtin/src/lib.rs
index a02c8a02dd..aa6df05391 100644
--- a/source/builtin/src/lib.rs
+++ b/source/builtin/src/lib.rs
@@ -1313,6 +1313,13 @@ macro_rules! decreases_to {
     };
 }
 
+#[cfg(verus_keep_ghost)]
+#[rustc_diagnostic_item = "verus::builtin::infer_spec_for_loop_iter"]
+#[verifier::spec]
+pub fn infer_spec_for_loop_iter<A>(_: A) -> Option<A> {
+    unimplemented!()
+}
+
 #[cfg(verus_keep_ghost)]
 #[rustc_diagnostic_item = "verus::builtin::global_size_of"]
 #[verifier::spec]
diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index 8d918bb911..d6cd448f27 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -1305,6 +1305,101 @@ impl Visitor {
         self.inside_ghost -= 1;
     }
 
+    fn desugar_for_loop(&mut self, for_loop: syn_verus::ExprForLoop) -> Expr {
+        use syn_verus::parse_quote_spanned;
+        // The regular Rust for-loop doesn't give us direct access to the iterator,
+        // which we need for writing invariants.
+        // Therefore, rather than letting Rust desugar a for-loop into a loop with a break,
+        // we desugar the for-loop into a loop with a break here.
+        // Specifically, we desugar:
+        //  for x in y: e inv I { body }
+        // into:
+        //  {
+        //      let mut y = core::iter::IntoIterator::into_iter(e);
+        //      loop
+        //          invariant_ensures
+        //              vstd::pervasive::ForLoopSpec::invariant(&y,
+        //                  builtin::infer_spec_for_loop_iter(
+        //                      &core::iter::IntoIterator::into_iter(e))),
+        //              I,
+        //          ensures
+        //              !vstd::pervasive::ForLoopSpec::condition(&y),
+        //      {
+        //          if let Some(x) = y.next() { body } else { break }
+        //      }
+        //  }
+        let span = for_loop.span();
+
+        let syn_verus::ExprForLoop {
+            mut attrs,
+            label,
+            for_token: _,
+            pat,
+            in_token: _,
+            expr_name,
+            expr,
+            invariant,
+            decreases,
+            body,
+        } = for_loop;
+
+        // Note: in principle, the automatically generated loop invariant
+        // should always succeed.  In case something goes unexpectedly wrong, though,
+        // give people a reasonable way to disable it:
+        let no_auto_loop_invariant = attrs.iter().position(|attr| {
+            attr.path.segments.len() == 2
+                && attr.path.segments[0].ident.to_string() == "verifier"
+                && attr.path.segments[1].ident.to_string() == "no_auto_loop_invariant"
+        });
+        if let Some(i) = no_auto_loop_invariant {
+            attrs.remove(i);
+        }
+        let inv_msg = "Automatically generated loop invariant failed. \
+            You can disable the automatic generation by adding \
+            #[verifier::no_auto_loop_invariant] \
+            to the loop. \
+            You might also try storing the loop expression in a variable outside the loop \
+            (e.g. `let e = 0..10; for x in e { ... }`).";
+
+        let (x_iter, _) = *expr_name.expect("for loop expr_name");
+        let mut stmts: Vec<Stmt> = Vec::new();
+        let expr_inv = expr.clone();
+        let inv: Expr = parse_quote_spanned!(span =>
+            #[verifier::custom_err(#inv_msg)]
+            vstd::pervasive::ForLoopSpec::invariant(&#x_iter,
+                builtin::infer_spec_for_loop_iter(&core::iter::IntoIterator::into_iter(#expr_inv)))
+        );
+        let invariant_ensure = if let Some(mut invariant) = invariant {
+            if no_auto_loop_invariant.is_none() {
+                invariant.exprs.exprs.insert(0, inv);
+            }
+            InvariantEnsures { token: Token![invariant_ensures](span), exprs: invariant.exprs }
+        } else if no_auto_loop_invariant.is_none() {
+            parse_quote_spanned!(span => invariant_ensures #inv,)
+        } else {
+            parse_quote_spanned!(span => )
+        };
+        // REVIEW: we might also want no_auto_loop_invariant to suppress the ensures,
+        // but at the moment, user-supplied ensures aren't supported, so this would be hard to use.
+        let ensure = parse_quote_spanned!(span =>
+            ensures !vstd::pervasive::ForLoopSpec::condition(&#x_iter),
+        );
+        self.add_loop_specs(&mut stmts, None, Some(invariant_ensure), Some(ensure), decreases);
+        let mut body: Block = parse_quote_spanned!(span => {
+            if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_iter)
+            #body else { break }
+        });
+        body.stmts.splice(0..0, stmts);
+        let mut loop_expr: ExprLoop = parse_quote_spanned!(span => loop #body);
+        loop_expr.label = label;
+        loop_expr.attrs = attrs;
+        let full_loop: Expr = parse_quote_spanned!(span => {
+            let mut #x_iter = core::iter::IntoIterator::into_iter(#expr);
+            #loop_expr
+        });
+        full_loop
+    }
+
     fn extract_quant_triggers(
         &mut self,
         inner_attrs: Vec<Attribute>,
@@ -1698,6 +1793,7 @@ impl VisitMut for Visitor {
             Expr::Closure(..) => true,
             Expr::Is(..) => true,
             Expr::Has(..) => true,
+            Expr::ForLoop(syn_verus::ExprForLoop { expr_name: Some(_), .. }) => true,
             _ => false,
         };
         if do_replace && self.inside_type == 0 {
@@ -2026,6 +2122,9 @@ impl VisitMut for Visitor {
                         *expr = expr_replacement;
                     }
                 }
+                Expr::ForLoop(for_loop) => {
+                    *expr = self.desugar_for_loop(for_loop);
+                }
                 Expr::Closure(mut clos) => {
                     if is_inside_ghost {
                         let span = clos.span();
diff --git a/source/pervasive/pervasive.rs b/source/pervasive/pervasive.rs
index b592f23744..c17d2d9218 100644
--- a/source/pervasive/pervasive.rs
+++ b/source/pervasive/pervasive.rs
@@ -33,6 +33,20 @@ pub proof fn affirm(b: bool)
 {
 }
 
+pub trait ForLoopSpec {
+    type Decrease;
+    // Always true before and after each loop iteration
+    // (When the analysis can infer a spec initial value, it places it in init)
+    spec fn invariant(&self, init: Option<&Self>) -> bool;
+    // Is the loop condition satisfied?
+    spec fn condition(&self) -> bool;
+    // Value used by default for decreases clause when no explicit decreases clause is provided
+    // (the user can override this with an explicit decreases clause).
+    // (If there's no appropriate decrease, this can return an arbitrary value like 0,
+    // and the user will have to provide an explicit decreases clause.)
+    spec fn decrease(&self) -> Self::Decrease;
+}
+
 // Non-statically-determined function calls are translated *internally* (at the VIR level)
 // to this function call. This should not actually be called directly by the user.
 // That is, Verus treats `f(x, y)` as `exec_nonstatic_call(f, (x, y))`.
diff --git a/source/pervasive/std_specs/core.rs b/source/pervasive/std_specs/core.rs
index 029d51fadb..fd4c589e81 100644
--- a/source/pervasive/std_specs/core.rs
+++ b/source/pervasive/std_specs/core.rs
@@ -19,7 +19,12 @@ pub struct ExOption<V>(core::option::Option<V>);
 #[verifier::reject_recursive_types_in_ground_variants(E)]
 pub struct ExResult<T, E>(core::result::Result<T, E>);
 
+pub open spec fn iter_into_iter_spec<I: Iterator>(i: I) -> I {
+    i
+}
+
 #[verifier::external_fn_specification]
+#[verifier::when_used_as_spec(iter_into_iter_spec)]
 pub fn ex_iter_into_iter<I: Iterator>(i: I) -> (r: I)
     ensures r == i
 {
diff --git a/source/pervasive/std_specs/range.rs b/source/pervasive/std_specs/range.rs
index 47a4647871..3c43b3154c 100644
--- a/source/pervasive/std_specs/range.rs
+++ b/source/pervasive/std_specs/range.rs
@@ -11,6 +11,7 @@ pub trait StepSpec where Self: Sized {
     // REVIEW: it would be nice to be able to use SpecOrd::spec_lt (not yet supported)
     spec fn spec_is_lt(self, other: Self) -> bool;
     spec fn spec_steps_between(self, end: Self) -> Option<usize>;
+    spec fn spec_steps_between_int(self, end: Self) -> int;
     spec fn spec_forward_checked(self, count: usize) -> Option<Self>;
     spec fn spec_backward_checked(self, count: usize) -> Option<Self>;
 }
@@ -40,6 +41,26 @@ pub fn ex_range_next<A: core::iter::Step>(range: &mut Range<A>) -> (r: Option<A>
     range.next()
 }
 
+impl<A: StepSpec> crate::pervasive::ForLoopSpec for Range<A> {
+    type Decrease = int;
+    open spec fn invariant(&self, init: Option<&Self>) -> bool {
+        &&& self.condition() || self.start == self.end
+        // TODO (not important): use new "matches ==>" syntax here
+        &&& if let Some(init) = init {
+                &&& init.start.spec_is_lt(self.start) || init.start == self.start
+                &&& init.end == self.end
+            } else {
+                true
+            }
+    }
+    open spec fn condition(&self) -> bool {
+        self.start.spec_is_lt(self.end)
+    }
+    open spec fn decrease(&self) -> int {
+        self.start.spec_steps_between_int(self.end)
+    }
+}
+
 } // verus!
 
 macro_rules! step_specs {
@@ -57,6 +78,9 @@ macro_rules! step_specs {
                     None
                 }
             }
+            open spec fn spec_steps_between_int(self, end: Self) -> int {
+                end - self
+            }
             open spec fn spec_forward_checked(self, count: usize) -> Option<Self> {
                 if self + count <= $t::MAX {
                     Some((self + count) as $t)
diff --git a/source/rust_verify/src/fn_call_to_vir.rs b/source/rust_verify/src/fn_call_to_vir.rs
index 786c0cb514..e36ee38c9c 100644
--- a/source/rust_verify/src/fn_call_to_vir.rs
+++ b/source/rust_verify/src/fn_call_to_vir.rs
@@ -677,6 +677,12 @@ fn verus_item_to_vir<'tcx, 'a>(
                     expr_item == &ExprItem::IsSmallerThanRecursiveFunctionField,
                 )
             }
+            ExprItem::InferSpecForLoopIter => {
+                record_spec_fn_no_proof_args(bctx, expr);
+                assert!(args.len() == 1);
+                let arg = expr_to_vir(bctx, &args[0], ExprModifier::REGULAR)?;
+                mk_expr(ExprX::Unary(UnaryOp::InferSpecForLoopIter, arg))
+            }
             ExprItem::IsVariant => {
                 record_spec_fn_allow_proof_args(bctx, expr);
                 assert!(args.len() == 2);
diff --git a/source/rust_verify/src/lifetime_generate.rs b/source/rust_verify/src/lifetime_generate.rs
index 8b7947b125..f9a941ad78 100644
--- a/source/rust_verify/src/lifetime_generate.rs
+++ b/source/rust_verify/src/lifetime_generate.rs
@@ -1,8 +1,6 @@
 use crate::attributes::{get_ghost_block_opt, get_mode, get_verifier_attrs, GhostBlockAttr};
 use crate::erase::{ErasureHints, ResolvedCall};
-use crate::rust_to_vir_base::{
-    def_id_to_vir_path, local_to_var, mid_ty_const_to_vir, mid_ty_to_vir_datatype,
-};
+use crate::rust_to_vir_base::{def_id_to_vir_path, local_to_var, mid_ty_const_to_vir};
 use crate::rust_to_vir_expr::get_adt_res;
 use crate::verus_items::{PervasiveItem, RustItem, VerusItem, VerusItems};
 use crate::{lifetime_ast::*, verus_items};
@@ -227,9 +225,9 @@ impl State {
         }
     }
 
-    fn reach_fun(&mut self, self_path: Option<Path>, id: DefId) {
-        if id.as_local().is_none() && !self.reached.contains(&(self_path.clone(), id)) {
-            self.reached.insert((self_path.clone(), id));
+    fn reach_fun(&mut self, id: DefId) {
+        if id.as_local().is_none() && !self.reached.contains(&(None, id)) {
+            self.reached.insert((None, id));
             self.imported_fun_worklist.push(id);
         }
     }
@@ -669,7 +667,6 @@ fn erase_call<'tcx>(
     state: &mut State,
     expect_spec: bool,
     expr: &Expr<'tcx>,
-    self_path: Option<Path>,
     expr_fun: Option<&Expr<'tcx>>,
     fn_def_id: Option<DefId>,
     node_substs: &'tcx rustc_middle::ty::List<rustc_middle::ty::GenericArg<'tcx>>,
@@ -790,7 +787,7 @@ fn erase_call<'tcx>(
                 }
             }
 
-            state.reach_fun(self_path, fn_def_id);
+            state.reach_fun(fn_def_id);
 
             let typ_args = mk_typ_args(ctxt, state, node_substs);
             let mut exps: Vec<Exp> = Vec::new();
@@ -958,26 +955,6 @@ fn erase_inv_block<'tcx>(
     Box::new((span, ExpX::OpenInvariant(atomicity, inner_pat, arg, pat_typ, mid_body)))
 }
 
-fn call_self_path(ctxt: &Context, qpath: &rustc_hir::QPath) -> Option<vir::ast::Path> {
-    match qpath {
-        rustc_hir::QPath::Resolved(_, _) => None,
-        rustc_hir::QPath::LangItem(_, _, _) => None,
-        rustc_hir::QPath::TypeRelative(ty, _) => match &ty.kind {
-            rustc_hir::TyKind::Path(qpath) => match ctxt.types().qpath_res(&qpath, ty.hir_id) {
-                rustc_hir::def::Res::Def(_, def_id) => {
-                    crate::rust_to_vir_base::def_id_to_vir_path_option(
-                        ctxt.tcx,
-                        &ctxt.verus_items,
-                        def_id,
-                    )
-                }
-                _ => None,
-            },
-            _ => None,
-        },
-    }
-}
-
 fn erase_expr<'tcx>(
     ctxt: &Context<'tcx>,
     state: &mut State,
@@ -1096,23 +1073,17 @@ fn erase_expr<'tcx>(
                 }
                 _ => false,
             };
-            let (self_path, fn_def_id) = if let ExprKind::Path(qpath) = &e0.kind {
-                let self_path = call_self_path(ctxt, qpath);
+            let fn_def_id = if let ExprKind::Path(qpath) = &e0.kind {
                 let def = ctxt.types().qpath_res(&qpath, e0.hir_id);
-                if let Res::Def(_, fn_def_id) = def {
-                    (self_path, Some(fn_def_id))
-                } else {
-                    (self_path, None)
-                }
+                if let Res::Def(_, fn_def_id) = def { Some(fn_def_id) } else { None }
             } else {
-                (None, None)
+                None
             };
             erase_call(
                 ctxt,
                 state,
                 expect_spec,
                 expr,
-                self_path,
                 Some(e0),
                 fn_def_id,
                 ctxt.types().node_args(e0.hir_id),
@@ -1125,25 +1096,11 @@ fn erase_expr<'tcx>(
         }
         ExprKind::MethodCall(segment, receiver, args, _call_span) => {
             let fn_def_id = ctxt.types().type_dependent_def_id(expr.hir_id).expect("method id");
-            let rcvr_typ = mid_ty_to_vir_datatype(
-                ctxt.tcx,
-                &ctxt.verus_items,
-                state.enclosing_fun_id.expect("enclosing_fun_id"),
-                receiver.span,
-                ctxt.types().node_type(receiver.hir_id),
-                true,
-            )
-            .expect("type");
-            let self_path = match &*vir::ast_util::undecorate_typ(&rcvr_typ) {
-                vir::ast::TypX::Datatype(path, _, _) => Some(path.clone()),
-                _ => None,
-            };
             erase_call(
                 ctxt,
                 state,
                 expect_spec,
                 expr,
-                self_path,
                 None,
                 Some(fn_def_id),
                 ctxt.types().node_args(expr.hir_id),
diff --git a/source/rust_verify/src/rust_to_vir_base.rs b/source/rust_verify/src/rust_to_vir_base.rs
index f465270c6d..91adc1a163 100644
--- a/source/rust_verify/src/rust_to_vir_base.rs
+++ b/source/rust_verify/src/rust_to_vir_base.rs
@@ -620,6 +620,7 @@ pub(crate) fn mid_ty_to_vir_ghost<'tcx>(
     Ok(t)
 }
 
+/*
 pub(crate) fn mid_ty_to_vir_datatype<'tcx>(
     tcx: TyCtxt<'tcx>,
     verus_items: &crate::verus_items::VerusItems,
@@ -630,6 +631,7 @@ pub(crate) fn mid_ty_to_vir_datatype<'tcx>(
 ) -> Result<Typ, VirErr> {
     Ok(mid_ty_to_vir_ghost(tcx, verus_items, param_env_src, span, &ty, true, allow_mut_ref)?.0)
 }
+*/
 
 // TODO: rename this back to mid_ty_to_vir
 pub(crate) fn mid_ty_to_vir<'tcx>(
diff --git a/source/rust_verify/src/verifier.rs b/source/rust_verify/src/verifier.rs
index ad4bccb79b..a680e8c98e 100644
--- a/source/rust_verify/src/verifier.rs
+++ b/source/rust_verify/src/verifier.rs
@@ -1696,6 +1696,12 @@ impl Verifier {
             self.args.rlimit,
             interpreter_log_file,
             self.vstd_crate_name.clone(),
+            &self
+                .erasure_hints
+                .as_ref()
+                .expect("erasure_hints")
+                .erasure_modes
+                .infer_spec_for_loop_iter_modes,
         )?;
         vir::recursive_types::check_traits(&krate, &global_ctx)?;
         let krate = vir::ast_simplify::simplify_krate(&mut global_ctx, &krate)?;
diff --git a/source/rust_verify/src/verus_items.rs b/source/rust_verify/src/verus_items.rs
index 32332a4035..374f9c33e8 100644
--- a/source/rust_verify/src/verus_items.rs
+++ b/source/rust_verify/src/verus_items.rs
@@ -138,6 +138,7 @@ pub(crate) enum ExprItem {
     IsSmallerThan,
     IsSmallerThanLexicographic,
     IsSmallerThanRecursiveFunctionField,
+    InferSpecForLoopIter,
 }
 
 #[derive(PartialEq, Eq, Debug, Clone, Copy, Hash)]
@@ -371,6 +372,7 @@ fn verus_items_map() -> Vec<(&'static str, VerusItem)> {
         ("verus::builtin::is_smaller_than",         VerusItem::Expr(ExprItem::IsSmallerThan)),
         ("verus::builtin::is_smaller_than_lexicographic", VerusItem::Expr(ExprItem::IsSmallerThanLexicographic)),
         ("verus::builtin::is_smaller_than_recursive_function_field", VerusItem::Expr(ExprItem::IsSmallerThanRecursiveFunctionField)),
+        ("verus::builtin::infer_spec_for_loop_iter", VerusItem::Expr(ExprItem::InferSpecForLoopIter)),
 
         ("verus::builtin::imply",                   VerusItem::CompilableOpr(CompilableOprItem::Implies)),
         // TODO ("verus::builtin::smartptr_new",    VerusItem::CompilableOpr(CompilableOprItem::SmartPtrNew)),
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index fa3469e493..59a4f4f41c 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -983,3 +983,67 @@ test_verify_one_file! {
         }
     } => Err(e) => assert_one_fails(e)
 }
+
+test_verify_one_file! {
+    #[test] for_loop1 verus_code! {
+        use vstd::prelude::*;
+        fn test_loop() {
+            let mut n: u64 = 0;
+            for x in iter: 0..10
+                invariant n == iter.start * 3,
+            {
+                assert(x < 10);
+                assert(x == iter.start - 1);
+                n += 3;
+            }
+            assert(n == 30);
+        }
+
+        fn test_loop_fail() {
+            let mut n: u64 = 0;
+            for x in iter: 0..10
+                invariant n == iter.start * 3,
+            {
+                assert(x < 9); // FAILS
+                assert(x == iter.start - 1);
+                n += 3;
+            }
+            assert(n == 30);
+        }
+    } => Err(e) => assert_one_fails(e)
+}
+
+test_verify_one_file! {
+    #[test] for_loop2 verus_code! {
+        use vstd::prelude::*;
+        fn test_loop() {
+            let mut n: u64 = 0;
+            let mut end = 10;
+            for x in iter: 0..end
+                invariant
+                    n == iter.start * 3,
+                    end == 10,
+            {
+                assert(x < 10);
+                assert(x == iter.start - 1);
+                n += 3;
+            }
+            assert(n == 30);
+        }
+
+        fn test_loop_fail() {
+            let mut n: u64 = 0;
+            let mut end = 10;
+            for x in iter: 0..end
+                invariant
+                    n == iter.start * 3,
+                    end == 10,
+            {
+                assert(x < 10); // FAILS
+                assert(x == iter.start - 1);
+                n += 3;
+                end = end + 0; // causes end to be non-constant, so loop needs more invariants
+            }
+        }
+    } => Err(e) => assert_one_fails(e)
+}
diff --git a/source/vir/src/ast.rs b/source/vir/src/ast.rs
index c68fb6f455..3962de3b1b 100644
--- a/source/vir/src/ast.rs
+++ b/source/vir/src/ast.rs
@@ -254,6 +254,15 @@ pub enum UnaryOp {
     StrIsAscii,
     /// Used only for handling casts from chars to ints
     CharToInt,
+    /// Given an exec/proof expression used to construct a loop iterator,
+    /// try to infer a pure specification for the loop iterator.
+    /// Return Some(spec) if successful, None otherwise.
+    /// (Note: this is just used as a hint for loop invariants;
+    /// regardless of whether it is Some(spec) or None, it should not affect soundness.)
+    /// For an exec/proof expression e, the spec s should be chosen so that the value v
+    /// that e evaluates to is immutable and v == s, where v may contain local variables.
+    /// For example, if v == (n..m), then n and m must be immutable local variables.
+    InferSpecForLoopIter,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd, Ord, ToDebugSNode)]
diff --git a/source/vir/src/ast_simplify.rs b/source/vir/src/ast_simplify.rs
index 5db8c8bbd2..6ad120816d 100644
--- a/source/vir/src/ast_simplify.rs
+++ b/source/vir/src/ast_simplify.rs
@@ -358,6 +358,16 @@ fn simplify_one_expr(
             }
         }
         ExprX::Unary(UnaryOp::CoerceMode { .. }, expr0) => Ok(expr0.clone()),
+        ExprX::Unary(UnaryOp::InferSpecForLoopIter, _) => {
+            let mode_opt = ctx.infer_spec_for_loop_iter_modes.get(&expr.span.id);
+            if mode_opt == Some(&Mode::Spec) {
+                // InferSpecForLoopIter must be spec mode to be usable for invariant inference
+                Ok(expr.clone())
+            } else {
+                // Otherwise, abandon the expression and return None (no inference)
+                Ok(crate::loop_inference::make_none_expr(&expr.span, &expr.typ))
+            }
+        }
         ExprX::UnaryOpr(UnaryOpr::TupleField { tuple_arity, field }, expr0) => {
             Ok(tuple_get_field_expr(state, &expr.span, &expr.typ, expr0, *tuple_arity, *field))
         }
@@ -1014,6 +1024,7 @@ pub fn simplify_krate(ctx: &mut GlobalCtx, krate: &Krate) -> Result<Krate, VirEr
         ctx.rlimit,
         ctx.interpreter_log.clone(),
         ctx.vstd_crate_name.clone(),
+        &vec![],
     )?;
     Ok(krate)
 }
diff --git a/source/vir/src/ast_to_sst.rs b/source/vir/src/ast_to_sst.rs
index c4c40cf327..791fcdda52 100644
--- a/source/vir/src/ast_to_sst.rs
+++ b/source/vir/src/ast_to_sst.rs
@@ -1204,6 +1204,11 @@ pub(crate) fn expr_to_stm_opt(
         ExprX::NullaryOpr(op) => {
             Ok((vec![], ReturnValue::Some(mk_exp(ExpX::NullaryOpr(op.clone())))))
         }
+        ExprX::Unary(UnaryOp::InferSpecForLoopIter, spec_expr) => {
+            let spec_exp = expr_to_pure_exp_skip_checks(ctx, state, &spec_expr)?;
+            let infer_exp = mk_exp(ExpX::Unary(UnaryOp::InferSpecForLoopIter, spec_exp));
+            Ok((vec![], ReturnValue::Some(infer_exp)))
+        }
         ExprX::Unary(op, exprr) => {
             let (mut stms, exp) = expr_to_stm_opt(ctx, state, exprr)?;
             let exp = unwrap_or_return_never!(exp, stms);
diff --git a/source/vir/src/context.rs b/source/vir/src/context.rs
index 42ce24d293..eeecdeecc4 100644
--- a/source/vir/src/context.rs
+++ b/source/vir/src/context.rs
@@ -4,7 +4,7 @@ use crate::ast::{
 };
 use crate::datatype_to_air::is_datatype_transparent;
 use crate::def::FUEL_ID;
-use crate::messages::{error, Span};
+use crate::messages::{error, AstId, Span};
 use crate::poly::MonoTyp;
 use crate::recursion::Node;
 use crate::scc::Graph;
@@ -47,6 +47,7 @@ pub struct GlobalCtx {
     pub(crate) interpreter_log: Arc<std::sync::Mutex<Option<File>>>,
     pub(crate) vstd_crate_name: Option<Ident>, // already an arc
     pub arch: crate::ast::ArchWordBits,
+    pub(crate) infer_spec_for_loop_iter_modes: Arc<HashMap<AstId, Mode>>,
 }
 
 // Context for verifying one function
@@ -194,6 +195,7 @@ impl GlobalCtx {
         rlimit: f32,
         interpreter_log: Arc<std::sync::Mutex<Option<File>>>,
         vstd_crate_name: Option<Ident>,
+        infer_spec_for_loop_iter_modes_vec: &Vec<(Span, Mode)>,
     ) -> Result<Self, VirErr> {
         let chosen_triggers: std::cell::RefCell<Vec<ChosenTriggers>> =
             std::cell::RefCell::new(Vec::new());
@@ -291,6 +293,13 @@ impl GlobalCtx {
 
         let datatype_graph = crate::recursive_types::build_datatype_graph(krate);
 
+        let mut infer_spec_for_loop_iter_modes = HashMap::new();
+        for (span, mode) in infer_spec_for_loop_iter_modes_vec {
+            if *mode != Mode::Spec || !infer_spec_for_loop_iter_modes.contains_key(&span.id) {
+                infer_spec_for_loop_iter_modes.insert(span.id, *mode);
+            }
+        }
+
         Ok(GlobalCtx {
             chosen_triggers,
             datatypes: Arc::new(datatypes),
@@ -304,6 +313,7 @@ impl GlobalCtx {
             interpreter_log,
             vstd_crate_name,
             arch: krate.arch.word_bits,
+            infer_spec_for_loop_iter_modes: Arc::new(infer_spec_for_loop_iter_modes),
         })
     }
 
@@ -325,6 +335,7 @@ impl GlobalCtx {
             interpreter_log,
             vstd_crate_name: self.vstd_crate_name.clone(),
             arch: self.arch,
+            infer_spec_for_loop_iter_modes: self.infer_spec_for_loop_iter_modes.clone(),
         }
     }
 
diff --git a/source/vir/src/def.rs b/source/vir/src/def.rs
index ded5f80494..3bace970e2 100644
--- a/source/vir/src/def.rs
+++ b/source/vir/src/def.rs
@@ -741,3 +741,10 @@ pub fn array_index_path(vstd_crate_name: &Option<Ident>) -> Path {
         ]),
     })
 }
+
+pub(crate) fn option_type_path() -> Path {
+    Arc::new(PathX {
+        krate: Some(Arc::new("core".to_string())),
+        segments: Arc::new(vec![Arc::new("option".to_string()), Arc::new("Option".to_string())]),
+    })
+}
diff --git a/source/vir/src/interpreter.rs b/source/vir/src/interpreter.rs
index cbcad6b20f..aca5adcfe8 100644
--- a/source/vir/src/interpreter.rs
+++ b/source/vir/src/interpreter.rs
@@ -899,7 +899,8 @@ fn eval_expr_internal(ctx: &Ctx, state: &mut State, exp: &Exp) -> Result<Exp, Vi
                         | CoerceMode { .. }
                         | StrLen
                         | StrIsAscii
-                        | CharToInt => ok,
+                        | CharToInt
+                        | InferSpecForLoopIter => ok,
                         MustBeFinalized => {
                             panic!("Found MustBeFinalized op {:?} after calling finalize_exp", exp)
                         }
@@ -1006,7 +1007,8 @@ fn eval_expr_internal(ctx: &Ctx, state: &mut State, exp: &Exp) -> Result<Exp, Vi
                         | CoerceMode { .. }
                         | StrLen
                         | StrIsAscii
-                        | CharToInt => ok,
+                        | CharToInt
+                        | InferSpecForLoopIter => ok,
                     }
                 }
                 // !(!(e_inner)) == e_inner
diff --git a/source/vir/src/lib.rs b/source/vir/src/lib.rs
index 9029b10890..df3afbc510 100644
--- a/source/vir/src/lib.rs
+++ b/source/vir/src/lib.rs
@@ -45,6 +45,7 @@ pub mod headers;
 pub mod interpreter;
 mod inv_masks;
 pub mod layout;
+mod loop_inference;
 pub mod messages;
 pub mod modes;
 pub mod poly;
diff --git a/source/vir/src/loop_inference.rs b/source/vir/src/loop_inference.rs
new file mode 100644
index 0000000000..13ff035978
--- /dev/null
+++ b/source/vir/src/loop_inference.rs
@@ -0,0 +1,65 @@
+use crate::ast::{Expr, ExprX, SpannedTyped, Typ, TypX, UnaryOp};
+use crate::messages::Span;
+use crate::sst::{Exp, ExpX, UniqueIdent};
+use std::sync::Arc;
+
+pub(crate) fn make_none_expr(span: &Span, typ: &Typ) -> Expr {
+    let option_path = crate::def::option_type_path();
+    let option_typx =
+        TypX::Datatype(option_path.clone(), Arc::new(vec![typ.clone()]), Arc::new(vec![]));
+    let exprx =
+        ExprX::Ctor(option_path.clone(), Arc::new("None".to_string()), Arc::new(vec![]), None);
+    SpannedTyped::new(span, &Arc::new(option_typx), exprx)
+}
+
+pub(crate) fn make_option_exp(opt: Option<Exp>, span: &Span, typ: &Typ) -> Exp {
+    let option_path = crate::def::option_type_path();
+    let option_typx =
+        TypX::Datatype(option_path.clone(), Arc::new(vec![typ.clone()]), Arc::new(vec![]));
+    let expx = match opt {
+        None => ExpX::Ctor(option_path.clone(), Arc::new("None".to_string()), Arc::new(vec![])),
+        Some(exp) => {
+            let field_id = crate::def::positional_field_ident(0);
+            let field = air::ast_util::ident_binder(&field_id, &exp);
+            let fields = Arc::new(vec![field]);
+            ExpX::Ctor(option_path.clone(), Arc::new("Some".to_string()), fields)
+        }
+    };
+    SpannedTyped::new(span, &Arc::new(option_typx), expx)
+}
+
+// InferSpecForLoopIter produces None if any variables in the express are modified in the loop
+fn vars_unmodified(modified_vars: &Arc<Vec<UniqueIdent>>, exp: &Exp) -> bool {
+    let mut map = air::scope_map::ScopeMap::new();
+    let r = crate::sst_visitor::exp_visitor_check(exp, &mut map, &mut |e: &Exp, _| match &e.x {
+        ExpX::Var(x) => {
+            if modified_vars.contains(x) {
+                Err(())
+            } else {
+                Ok(())
+            }
+        }
+        _ => Ok(()),
+    });
+    match r {
+        Ok(()) => true,
+        Err(()) => false,
+    }
+}
+
+pub(crate) fn finalize_inv(modified_vars: &Arc<Vec<UniqueIdent>>, exp: &Exp) -> Exp {
+    crate::sst_visitor::map_exp_visitor(exp, &mut |e: &Exp| {
+        match &e.x {
+            ExpX::Unary(UnaryOp::InferSpecForLoopIter, e_infer) => {
+                if vars_unmodified(modified_vars, e_infer) {
+                    // promote to Some(e)
+                    make_option_exp(Some(e_infer.clone()), &e.span, &e.typ)
+                } else {
+                    // otherwise, leave as-is to be removed by sst_to_air
+                    e.clone()
+                }
+            }
+            _ => e.clone(),
+        }
+    })
+}
diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index db4d0dfe2f..c2be8bf7fa 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -49,6 +49,8 @@ pub struct ErasureModes {
     pub condition_modes: Vec<(Span, Mode)>,
     // Modes of variables in Var, Assign, Decl
     pub var_modes: Vec<(Span, Mode)>,
+    // Modes of InferSpecForLoopIter
+    pub infer_spec_for_loop_iter_modes: Vec<(Span, Mode)>,
 }
 
 impl Ghost {
@@ -623,6 +625,17 @@ fn check_expr_handle_mut_arg(
         ExprX::Unary(UnaryOp::HeightTrigger, _) => {
             panic!("direct access to 'height' is not allowed")
         }
+        ExprX::Unary(UnaryOp::InferSpecForLoopIter, e1) => {
+            // InferSpecForLoopIter is a loop-invariant hint that always has mode spec.
+            // If the expression already has mode spec (e.g. because the function calls
+            // are all autospec), then keep the expression.
+            // Otherwise, make a note that the expression had mode exec,
+            // so that ast_simplify can replace the expression with None.
+            let mode_opt = check_expr(typing, outer_mode, e1);
+            let mode = mode_opt.unwrap_or(Mode::Exec);
+            typing.erasure_modes.infer_spec_for_loop_iter_modes.push((expr.span.clone(), mode));
+            Ok(Mode::Spec)
+        }
         ExprX::Unary(_, e1) => check_expr(typing, outer_mode, e1),
         ExprX::UnaryOpr(UnaryOpr::Box(_), e1) => check_expr(typing, outer_mode, e1),
         ExprX::UnaryOpr(UnaryOpr::Unbox(_), e1) => check_expr(typing, outer_mode, e1),
@@ -1264,7 +1277,11 @@ pub fn check_crate(
     for datatype in krate.datatypes.iter() {
         datatypes.insert(datatype.x.path.clone(), datatype.clone());
     }
-    let erasure_modes = ErasureModes { condition_modes: vec![], var_modes: vec![] };
+    let erasure_modes = ErasureModes {
+        condition_modes: vec![],
+        var_modes: vec![],
+        infer_spec_for_loop_iter_modes: vec![],
+    };
     let mut typing = Typing {
         funs,
         datatypes,
diff --git a/source/vir/src/poly.rs b/source/vir/src/poly.rs
index af46d34baf..dfe518926d 100644
--- a/source/vir/src/poly.rs
+++ b/source/vir/src/poly.rs
@@ -410,6 +410,11 @@ fn poly_expr(ctx: &Ctx, state: &mut State, expr: &Expr) -> Expr {
                     let e1 = coerce_expr_to_native(ctx, &e1);
                     mk_expr(ExprX::Unary(*op, e1))
                 }
+                UnaryOp::InferSpecForLoopIter => {
+                    // e1 will be the argument to spec Option::Some(...)
+                    let e1 = coerce_expr_to_poly(ctx, &e1);
+                    mk_expr(ExprX::Unary(*op, e1))
+                }
                 UnaryOp::HeightTrigger => panic!("direct access to 'height' is not allowed"),
                 UnaryOp::Trigger(_) | UnaryOp::CoerceMode { .. } => {
                     mk_expr_typ(&e1.typ, ExprX::Unary(*op, e1.clone()))
diff --git a/source/vir/src/prune.rs b/source/vir/src/prune.rs
index 62b62534e1..ee282dde9e 100644
--- a/source/vir/src/prune.rs
+++ b/source/vir/src/prune.rs
@@ -284,6 +284,10 @@ fn traverse_reachable(ctxt: &Ctxt, state: &mut State) {
                             &fn_namespace_name(&ctxt.vstd_crate_name, *atomicity),
                         );
                     }
+                    ExprX::Unary(crate::ast::UnaryOp::InferSpecForLoopIter, _) => {
+                        let t = ReachedType::Datatype(crate::def::option_type_path());
+                        reach_type(ctxt, state, &t);
+                    }
                     _ => {}
                 }
                 Ok(e.clone())
diff --git a/source/vir/src/sst_to_air.rs b/source/vir/src/sst_to_air.rs
index bce64d6c4d..6d0767b040 100644
--- a/source/vir/src/sst_to_air.rs
+++ b/source/vir/src/sst_to_air.rs
@@ -574,6 +574,15 @@ pub(crate) fn constant_to_expr(ctx: &Ctx, constant: &crate::ast::Constant) -> Ex
     }
 }
 
+fn exp_get_custom_err(exp: &Exp) -> Option<Arc<String>> {
+    match &exp.x {
+        ExpX::UnaryOpr(UnaryOpr::Box(_), e) => exp_get_custom_err(e),
+        ExpX::UnaryOpr(UnaryOpr::Unbox(_), e) => exp_get_custom_err(e),
+        ExpX::UnaryOpr(UnaryOpr::CustomErr(s), _) => Some(s.clone()),
+        _ => None,
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub(crate) enum ExprMode {
     Spec,
@@ -915,6 +924,9 @@ pub(crate) fn exp_to_expr(ctx: &Ctx, exp: &Exp, expr_ctxt: &ExprCtxt) -> Result<
                     }
                 }
                 UnaryOp::HeightTrigger => panic!("internal error: unexpected HeightTrigger"),
+                UnaryOp::InferSpecForLoopIter => {
+                    panic!("internal error: unexpected InferSpecForLoopIter")
+                }
                 UnaryOp::Trigger(_) => exp_to_expr(ctx, arg, expr_ctxt)?,
                 UnaryOp::CoerceMode { .. } => {
                     panic!("internal error: TupleField should have been removed before here")
@@ -971,9 +983,13 @@ pub(crate) fn exp_to_expr(ctx: &Ctx, exp: &Exp, expr_ctxt: &ExprCtxt) -> Result<
             }
             UnaryOp::CharToInt => {
                 let expr = exp_to_expr(ctx, exp, expr_ctxt)?;
-
                 Arc::new(ExprX::Apply(str_ident(CHAR_TO_UNICODE), Arc::new(vec![expr])))
             }
+            UnaryOp::InferSpecForLoopIter => {
+                // loop_inference failed to promote to Some, so demote to None
+                let exp = crate::loop_inference::make_option_exp(None, &exp.span, &exp.typ);
+                exp_to_expr(ctx, &exp, expr_ctxt)?
+            }
         },
         (ExpX::UnaryOpr(UnaryOpr::Box(_) | UnaryOpr::Unbox(_), exp), true) => {
             exp_to_expr(ctx, exp, expr_ctxt)?
@@ -1424,8 +1440,8 @@ struct PostConditionInfo {
 struct LoopInfo {
     label: Option<String>,
     some_cond: bool,
-    invs_entry: Arc<Vec<(Span, Expr)>>,
-    invs_exit: Arc<Vec<(Span, Expr)>>,
+    invs_entry: Arc<Vec<(Span, Expr, Option<Arc<String>>)>>,
+    invs_exit: Arc<Vec<(Span, Expr, Option<Arc<String>>)>>,
 }
 
 struct State {
@@ -2039,8 +2055,11 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
                 } else {
                     error_with_label(&stm.span, "loop invariant not satisfied", "at this continue")
                 };
-                for (span, inv) in invs.iter() {
-                    let error = base_error.secondary_label(span, "failed this invariant");
+                for (span, inv, msg) in invs.iter() {
+                    let mut error = base_error.secondary_label(span, "failed this invariant");
+                    if let Some(msg) = msg {
+                        error = error.secondary_label(span, &**msg);
+                    }
                     stmts.push(Arc::new(StmtX::Assert(error, inv.clone())));
                 }
             }
@@ -2101,19 +2120,21 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
             } else {
                 (None, None, None)
             };
-            let mut invs_entry: Vec<(Span, Expr)> = Vec::new();
-            let mut invs_exit: Vec<(Span, Expr)> = Vec::new();
+            let mut invs_entry: Vec<(Span, Expr, Option<Arc<String>>)> = Vec::new();
+            let mut invs_exit: Vec<(Span, Expr, Option<Arc<String>>)> = Vec::new();
             for inv in invs.iter() {
-                let expr = exp_to_expr(ctx, &inv.inv, expr_ctxt)?;
+                let inv_exp = crate::loop_inference::finalize_inv(modified_vars, &inv.inv);
+                let msg_opt = exp_get_custom_err(&inv_exp);
+                let expr = exp_to_expr(ctx, &inv_exp, expr_ctxt)?;
                 if cond.is_some() {
                     assert!(inv.at_entry);
                     assert!(inv.at_exit);
                 }
                 if inv.at_entry {
-                    invs_entry.push((inv.inv.span.clone(), expr.clone()));
+                    invs_entry.push((inv.inv.span.clone(), expr.clone(), msg_opt.clone()));
                 }
                 if inv.at_exit {
-                    invs_exit.push((inv.inv.span.clone(), expr.clone()));
+                    invs_exit.push((inv.inv.span.clone(), expr.clone(), msg_opt.clone()));
                 }
             }
             let invs_entry = Arc::new(invs_entry);
@@ -2173,7 +2194,7 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
 
             // Assume invariants for the beginning of the loop body.
             // (These need to go after the above Havoc statements.)
-            for (_, inv) in invs_entry.iter() {
+            for (_, inv, _) in invs_entry.iter() {
                 air_body.push(Arc::new(StmtX::Assume(inv.clone())));
             }
 
@@ -2195,8 +2216,11 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
             state.loop_infos.pop();
 
             if !ctx.checking_spec_preconditions() {
-                for (span, inv) in invs_entry.iter() {
-                    let error = error(span, crate::def::INV_FAIL_LOOP_END);
+                for (span, inv, msg) in invs_entry.iter() {
+                    let mut error = error(span, crate::def::INV_FAIL_LOOP_END);
+                    if let Some(msg) = msg {
+                        error = error.secondary_label(span, &**msg);
+                    }
                     let inv_stmt = StmtX::Assert(error, inv.clone());
                     air_body.push(Arc::new(inv_stmt));
                 }
@@ -2231,8 +2255,11 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
             // At original site of while loop, assert invariant, havoc, assume invariant + neg_cond
             let mut stmts: Vec<Stmt> = Vec::new();
             if !ctx.checking_spec_preconditions() {
-                for (span, inv) in invs_entry.iter() {
-                    let error = error(span, crate::def::INV_FAIL_LOOP_FRONT);
+                for (span, inv, msg) in invs_entry.iter() {
+                    let mut error = error(span, crate::def::INV_FAIL_LOOP_FRONT);
+                    if let Some(msg) = msg {
+                        error = error.secondary_label(span, &**msg);
+                    }
                     let inv_stmt = StmtX::Assert(error, inv.clone());
                     stmts.push(Arc::new(inv_stmt));
                 }
@@ -2248,7 +2275,7 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
                     }
                 }
             }
-            for (_, inv) in invs_exit.iter() {
+            for (_, inv, _) in invs_exit.iter() {
                 let inv_stmt = StmtX::Assume(inv.clone());
                 stmts.push(Arc::new(inv_stmt));
             }
diff --git a/source/vir/src/sst_util.rs b/source/vir/src/sst_util.rs
index b606751f1d..19fdd837d0 100644
--- a/source/vir/src/sst_util.rs
+++ b/source/vir/src/sst_util.rs
@@ -299,9 +299,10 @@ impl ExpX {
                 UnaryOp::StrLen => (format!("{}.len()", exp.x.to_string_prec(99)), 90),
                 UnaryOp::StrIsAscii => (format!("{}.is_ascii()", exp.x.to_string_prec(99)), 90),
                 UnaryOp::CharToInt => (format!("{} as char", exp.x.to_string_prec(99)), 90),
-                UnaryOp::Trigger(..) | UnaryOp::CoerceMode { .. } | UnaryOp::MustBeFinalized => {
-                    ("".to_string(), 0)
-                }
+                UnaryOp::Trigger(..)
+                | UnaryOp::CoerceMode { .. }
+                | UnaryOp::MustBeFinalized
+                | UnaryOp::InferSpecForLoopIter => ("".to_string(), 0),
             },
             UnaryOpr(op, exp) => {
                 use crate::ast::UnaryOpr::*;
diff --git a/source/vir/src/triggers.rs b/source/vir/src/triggers.rs
index 6a01586abc..465da5320b 100644
--- a/source/vir/src/triggers.rs
+++ b/source/vir/src/triggers.rs
@@ -154,7 +154,8 @@ fn check_trigger_expr_arg(state: &State, expect_boxed: bool, arg: &Exp) -> Resul
             | UnaryOp::BitNot
             | UnaryOp::StrLen
             | UnaryOp::StrIsAscii
-            | UnaryOp::CharToInt => Ok(()),
+            | UnaryOp::CharToInt
+            | UnaryOp::InferSpecForLoopIter => Ok(()),
         },
         ExpX::UnaryOpr(op, arg) => match op {
             UnaryOpr::Box(_) | UnaryOpr::Unbox(_) | UnaryOpr::CustomErr(_) => {
@@ -257,6 +258,9 @@ fn check_trigger_expr(
                 | UnaryOp::HeightTrigger
                 | UnaryOp::CoerceMode { .. }
                 | UnaryOp::MustBeFinalized => Ok(()),
+                UnaryOp::InferSpecForLoopIter => {
+                    Err(error(&exp.span, "triggers cannot contain loop spec inference"))
+                }
                 UnaryOp::Not => Err(error(&exp.span, "triggers cannot contain boolean operators")),
             },
             ExpX::UnaryOpr(op, arg) => match op {
diff --git a/source/vir/src/triggers_auto.rs b/source/vir/src/triggers_auto.rs
index 2f18766d28..155302cb6a 100644
--- a/source/vir/src/triggers_auto.rs
+++ b/source/vir/src/triggers_auto.rs
@@ -331,6 +331,7 @@ fn gather_terms(ctxt: &mut Ctxt, ctx: &Ctx, exp: &Exp, depth: u64) -> (bool, Ter
                 | UnaryOp::CharToInt => 0,
                 UnaryOp::HeightTrigger => 1,
                 UnaryOp::Trigger(_) | UnaryOp::Clip { .. } | UnaryOp::BitNot => 1,
+                UnaryOp::InferSpecForLoopIter => 1,
                 UnaryOp::StrIsAscii | UnaryOp::StrLen => fail_on_strop(),
             };
             let (_, term1) = gather_terms(ctxt, ctx, e1, depth);

From 5b89140c7c18ea6df30822a5e4eb4f2893fedce5 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Thu, 4 Jan 2024 17:15:49 -0800
Subject: [PATCH 05/19] For for-loops, keep ghost iterator separate from exec
 iterator

---
 source/builtin_macros/src/syntax.rs        | 165 ++++++++++++++++-----
 source/pervasive/pervasive.rs              |  37 ++++-
 source/pervasive/std_specs/range.rs        |  73 ++++++++-
 source/pervasive/string.rs                 |   1 +
 source/rust_verify/example/recursion.rs    |   5 +-
 source/rust_verify/example/vectors.rs      |   5 +-
 source/rust_verify/src/attributes.rs       |   6 +
 source/rust_verify/src/rust_to_vir_expr.rs |  16 +-
 source/rust_verify_test/tests/loops.rs     |  62 +++++++-
 source/vir/src/ast.rs                      |   8 +-
 source/vir/src/ast_to_sst.rs               |   9 +-
 source/vir/src/ast_visitor.rs              |  12 +-
 source/vir/src/early_exit_cf.rs            |   2 +-
 source/vir/src/modes.rs                    |   2 +-
 source/vir/src/poly.rs                     |  10 +-
 source/vir/src/split_expression.rs         |   3 +-
 source/vir/src/sst.rs                      |   1 +
 source/vir/src/sst_to_air.rs               |   9 +-
 source/vir/src/sst_vars.rs                 |   3 +-
 source/vir/src/sst_visitor.rs              |  32 +++-
 20 files changed, 378 insertions(+), 83 deletions(-)

diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index d6cd448f27..752c225c3c 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -1312,30 +1312,38 @@ impl Visitor {
         // Therefore, rather than letting Rust desugar a for-loop into a loop with a break,
         // we desugar the for-loop into a loop with a break here.
         // Specifically, we desugar:
-        //  for x in y: e inv I { body }
+        //  for x in y: e invariant inv { body }
         // into:
         //  {
-        //      let mut y = core::iter::IntoIterator::into_iter(e);
+        //      #[allow(non_snake_case)]
+        //      let mut VERUS_exec_iter = core::iter::IntoIterator::into_iter(e);
+        //      #[allow(non_snake_case)]
+        //      let ghost mut y = vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(&VERUS_exec_iter);
         //      loop
         //          invariant_ensures
-        //              vstd::pervasive::ForLoopSpec::invariant(&y,
+        //              vstd::pervasive::ForLoopGhostIterator::exec_invariant(&y, &VERUS_exec_iter),
+        //              vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&y,
         //                  builtin::infer_spec_for_loop_iter(
-        //                      &core::iter::IntoIterator::into_iter(e))),
-        //              I,
+        //                      &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+        //                          &core::iter::IntoIterator::into_iter(e)))),
+        //              { let x = vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&y); inv },
         //          ensures
-        //              !vstd::pervasive::ForLoopSpec::condition(&y),
+        //              !vstd::pervasive::ForLoopGhostIterator::ghost_condition(&y),
         //      {
-        //          if let Some(x) = y.next() { body } else { break }
+        //          if let Some(x) = core::iter::Iterator::next(&mut y) { body } else { break }
+        //          proof { y = vstd::pervasive::ForLoopGhostIterator::ghost_advance(&y, &VERUS_exec_iter); }
         //      }
         //  }
+        // Note that "continue" and labels are not yet supported;
+        // continue would also need to call ghost_advance.
         let span = for_loop.span();
 
         let syn_verus::ExprForLoop {
             mut attrs,
             label,
-            for_token: _,
+            for_token,
             pat,
-            in_token: _,
+            in_token,
             expr_name,
             expr,
             invariant,
@@ -1343,6 +1351,14 @@ impl Visitor {
             body,
         } = for_loop;
 
+        let no_loop_invariant = attrs.iter().position(|attr| {
+            attr.path.segments.len() == 2
+                && attr.path.segments[0].ident.to_string() == "verifier"
+                && attr.path.segments[1].ident.to_string() == "no_loop_invariant"
+        });
+        if let Some(i) = no_loop_invariant {
+            attrs.remove(i);
+        }
         // Note: in principle, the automatically generated loop invariant
         // should always succeed.  In case something goes unexpectedly wrong, though,
         // give people a reasonable way to disable it:
@@ -1354,49 +1370,130 @@ impl Visitor {
         if let Some(i) = no_auto_loop_invariant {
             attrs.remove(i);
         }
-        let inv_msg = "Automatically generated loop invariant failed. \
+
+        if !self.erase_ghost.keep() {
+            return Expr::ForLoop(syn_verus::ExprForLoop {
+                attrs,
+                label,
+                for_token,
+                pat,
+                in_token,
+                expr_name: None,
+                expr,
+                invariant: None,
+                decreases: None,
+                body,
+            });
+        }
+
+        attrs.push(mk_verus_attr(span, quote! { for_loop }));
+        let exec_inv_msg = "For-loop iterator invariant failed. \
+            This may indicate a bug in the definition of the ForLoopGhostIterator. \
+            You might try using a `loop` instead of a `for`.";
+        let ghost_inv_msg = "Automatically generated loop invariant failed. \
             You can disable the automatic generation by adding \
             #[verifier::no_auto_loop_invariant] \
             to the loop. \
             You might also try storing the loop expression in a variable outside the loop \
             (e.g. `let e = 0..10; for x in e { ... }`).";
 
-        let (x_iter, _) = *expr_name.expect("for loop expr_name");
+        let x_exec_iter = Ident::new("VERUS_exec_iter", span);
+        let x_ghost_iter = if let Some(x_ghost_iter_box) = expr_name {
+            let (x_ghost_iter, _) = *x_ghost_iter_box;
+            x_ghost_iter
+        } else {
+            Ident::new("VERUS_ghost_iter", span)
+        };
         let mut stmts: Vec<Stmt> = Vec::new();
         let expr_inv = expr.clone();
-        let inv: Expr = parse_quote_spanned!(span =>
-            #[verifier::custom_err(#inv_msg)]
-            vstd::pervasive::ForLoopSpec::invariant(&#x_iter,
-                builtin::infer_spec_for_loop_iter(&core::iter::IntoIterator::into_iter(#expr_inv)))
+        //              vstd::pervasive::ForLoopGhostIterator::exec_invariant(&y, &VERUS_exec_iter),
+        //              vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&y,
+        //                  builtin::infer_spec_for_loop_iter(
+        //                      &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+        //                          &core::iter::IntoIterator::into_iter(e)))),
+        let exec_inv: Expr = parse_quote_spanned!(span =>
+            #[verifier::custom_err(#exec_inv_msg)]
+            vstd::pervasive::ForLoopGhostIterator::exec_invariant(&#x_ghost_iter, &#x_exec_iter)
+        );
+        let ghost_inv: Expr = parse_quote_spanned!(span =>
+            #[verifier::custom_err(#ghost_inv_msg)]
+            vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&#x_ghost_iter,
+                builtin::infer_spec_for_loop_iter(
+                    &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+                        &core::iter::IntoIterator::into_iter(#expr_inv))))
         );
         let invariant_ensure = if let Some(mut invariant) = invariant {
-            if no_auto_loop_invariant.is_none() {
-                invariant.exprs.exprs.insert(0, inv);
+            for inv in &mut invariant.exprs.exprs {
+                *inv = parse_quote_spanned!(span => {
+                    let #pat = vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&#x_ghost_iter);
+                    #inv
+                });
+            }
+            if no_loop_invariant.is_none() {
+                invariant.exprs.exprs.insert(0, exec_inv);
+                if no_auto_loop_invariant.is_none() {
+                    invariant.exprs.exprs.insert(1, ghost_inv);
+                }
             }
-            InvariantEnsures { token: Token![invariant_ensures](span), exprs: invariant.exprs }
-        } else if no_auto_loop_invariant.is_none() {
-            parse_quote_spanned!(span => invariant_ensures #inv,)
+            Some(InvariantEnsures {
+                token: Token![invariant_ensures](span),
+                exprs: invariant.exprs,
+            })
+        } else if no_loop_invariant.is_none() && no_auto_loop_invariant.is_none() {
+            Some(parse_quote_spanned!(span => invariant_ensures #exec_inv, #ghost_inv,))
+        } else if no_loop_invariant.is_none() && no_auto_loop_invariant.is_some() {
+            Some(parse_quote_spanned!(span => invariant_ensures #exec_inv,))
         } else {
-            parse_quote_spanned!(span => )
+            None
         };
         // REVIEW: we might also want no_auto_loop_invariant to suppress the ensures,
         // but at the moment, user-supplied ensures aren't supported, so this would be hard to use.
-        let ensure = parse_quote_spanned!(span =>
-            ensures !vstd::pervasive::ForLoopSpec::condition(&#x_iter),
-        );
-        self.add_loop_specs(&mut stmts, None, Some(invariant_ensure), Some(ensure), decreases);
-        let mut body: Block = parse_quote_spanned!(span => {
-            if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_iter)
-            #body else { break }
-        });
+        let ensure = if no_loop_invariant.is_none() {
+            Some(parse_quote_spanned!(span =>
+                ensures !vstd::pervasive::ForLoopGhostIterator::ghost_condition(&#x_ghost_iter),
+            ))
+        } else {
+            None
+        };
+        self.add_loop_specs(&mut stmts, None, invariant_ensure, ensure, decreases);
+        let mut body: Block = if no_loop_invariant.is_none() {
+            parse_quote_spanned!(span => {
+                if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_exec_iter)
+                #body else { break }
+                #[verifier::proof_block] {
+                    #x_ghost_iter = vstd::pervasive::ForLoopGhostIterator::ghost_advance(
+                        &#x_ghost_iter, &#x_exec_iter);
+                }
+            })
+        } else {
+            parse_quote_spanned!(span => {
+                if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_exec_iter)
+                #body else { break }
+            })
+        };
         body.stmts.splice(0..0, stmts);
         let mut loop_expr: ExprLoop = parse_quote_spanned!(span => loop #body);
         loop_expr.label = label;
         loop_expr.attrs = attrs;
-        let full_loop: Expr = parse_quote_spanned!(span => {
-            let mut #x_iter = core::iter::IntoIterator::into_iter(#expr);
-            #loop_expr
-        });
+        let full_loop: Expr = if no_loop_invariant.is_none() {
+            parse_quote_spanned!(span => {
+                #[allow(non_snake_case)]
+                let mut #x_exec_iter = core::iter::IntoIterator::into_iter(#expr);
+                #[allow(non_snake_case)]
+                #[verus::internal(spec)]
+                let mut #x_ghost_iter;
+                #[verifier::proof_block] {
+                    #x_ghost_iter = vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(&#x_exec_iter);
+                }
+                #loop_expr
+            })
+        } else {
+            parse_quote_spanned!(span => {
+                #[allow(non_snake_case)]
+                let mut #x_exec_iter = core::iter::IntoIterator::into_iter(#expr);
+                #loop_expr
+            })
+        };
         full_loop
     }
 
@@ -1793,7 +1890,7 @@ impl VisitMut for Visitor {
             Expr::Closure(..) => true,
             Expr::Is(..) => true,
             Expr::Has(..) => true,
-            Expr::ForLoop(syn_verus::ExprForLoop { expr_name: Some(_), .. }) => true,
+            Expr::ForLoop(..) => true,
             _ => false,
         };
         if do_replace && self.inside_type == 0 {
diff --git a/source/pervasive/pervasive.rs b/source/pervasive/pervasive.rs
index c17d2d9218..6a9846029d 100644
--- a/source/pervasive/pervasive.rs
+++ b/source/pervasive/pervasive.rs
@@ -33,18 +33,45 @@ pub proof fn affirm(b: bool)
 {
 }
 
-pub trait ForLoopSpec {
+pub trait ForLoopGhostIterator {
+    type ExecIter;
+    type Item;
     type Decrease;
+
+    // Connect the ExecIter to the GhostIter
+    // Always enabled
     // Always true before and after each loop iteration
-    // (When the analysis can infer a spec initial value, it places it in init)
-    spec fn invariant(&self, init: Option<&Self>) -> bool;
+    spec fn exec_invariant(&self, exec_iter: &Self::ExecIter) -> bool;
+
+    // Additional optional invariants about the GhostIter
+    // May be disabled with #[verifier::no_auto_loop_invariant]
+    // If enabled, always true before and after each loop iteration
+    // (When the analysis can infer a spec initial value, the analysis places the value in init)
+    spec fn ghost_invariant(&self, init: Option<&Self>) -> bool;
+
     // Is the loop condition satisfied?
-    spec fn condition(&self) -> bool;
+    spec fn ghost_condition(&self) -> bool;
+
     // Value used by default for decreases clause when no explicit decreases clause is provided
     // (the user can override this with an explicit decreases clause).
     // (If there's no appropriate decrease, this can return an arbitrary value like 0,
     // and the user will have to provide an explicit decreases clause.)
-    spec fn decrease(&self) -> Self::Decrease;
+    spec fn ghost_decrease(&self) -> Self::Decrease;
+
+    // If there will be Some next value, and we can make a useful guess as to what the next value
+    // will be, return it.
+    // Otherwise, return an arbitrary value.
+    spec fn ghost_peek_next(&self) -> Self::Item;
+
+    // At the end of the for loop, advance to the next position.
+    spec fn ghost_advance(&self, exec_iter: &Self::ExecIter) -> Self where Self: Sized;
+}
+
+pub trait ForLoopGhostIteratorNew {
+    type GhostIter;
+
+    // Create a new ghost iterator from an exec iterator
+    spec fn ghost_iter(&self) -> Self::GhostIter;
 }
 
 // Non-statically-determined function calls are translated *internally* (at the VIR level)
diff --git a/source/pervasive/std_specs/range.rs b/source/pervasive/std_specs/range.rs
index 3c43b3154c..6548e7e68a 100644
--- a/source/pervasive/std_specs/range.rs
+++ b/source/pervasive/std_specs/range.rs
@@ -13,7 +13,9 @@ pub trait StepSpec where Self: Sized {
     spec fn spec_steps_between(self, end: Self) -> Option<usize>;
     spec fn spec_steps_between_int(self, end: Self) -> int;
     spec fn spec_forward_checked(self, count: usize) -> Option<Self>;
+    spec fn spec_forward_checked_int(self, count: int) -> Option<Self>;
     spec fn spec_backward_checked(self, count: usize) -> Option<Self>;
+    spec fn spec_backward_checked_int(self, count: int) -> Option<Self>;
 }
 
 pub spec fn spec_range_next<A>(a: Range<A>) -> (Range<A>, Option<A>);
@@ -41,23 +43,72 @@ pub fn ex_range_next<A: core::iter::Step>(range: &mut Range<A>) -> (r: Option<A>
     range.next()
 }
 
-impl<A: StepSpec> crate::pervasive::ForLoopSpec for Range<A> {
+pub struct RangeGhostIterator<A> {
+    pub start: A,
+    pub cur: A,
+    pub end: A,
+}
+
+impl<A: StepSpec> crate::pervasive::ForLoopGhostIteratorNew for Range<A> {
+    type GhostIter = RangeGhostIterator<A>;
+
+    open spec fn ghost_iter(&self) -> RangeGhostIterator<A> {
+        RangeGhostIterator {
+            start: self.start,
+            cur: self.start,
+            end: self.end,
+        }
+    }
+}
+
+impl<A: StepSpec + core::iter::Step> crate::pervasive::ForLoopGhostIterator for RangeGhostIterator<A> {
+    type ExecIter = Range<A>;
+    type Item = A;
     type Decrease = int;
-    open spec fn invariant(&self, init: Option<&Self>) -> bool {
-        &&& self.condition() || self.start == self.end
+
+    open spec fn exec_invariant(&self, exec_iter: &Range<A>) -> bool {
+        &&& self.cur == exec_iter.start
+        &&& self.end == exec_iter.end
+    }
+
+    open spec fn ghost_invariant(&self, init: Option<&Self>) -> bool {
+        &&& self.start.spec_is_lt(self.cur) || self.start == self.cur
+        &&& self.cur.spec_is_lt(self.end) || self.cur == self.end
         // TODO (not important): use new "matches ==>" syntax here
         &&& if let Some(init) = init {
-                &&& init.start.spec_is_lt(self.start) || init.start == self.start
+                &&& init.start == init.cur
+                &&& init.start == self.start
                 &&& init.end == self.end
             } else {
                 true
             }
     }
-    open spec fn condition(&self) -> bool {
-        self.start.spec_is_lt(self.end)
+
+    open spec fn ghost_condition(&self) -> bool {
+        self.cur.spec_is_lt(self.end)
     }
-    open spec fn decrease(&self) -> int {
-        self.start.spec_steps_between_int(self.end)
+
+    open spec fn ghost_decrease(&self) -> int {
+        self.cur.spec_steps_between_int(self.end)
+    }
+
+    open spec fn ghost_peek_next(&self) -> A {
+        self.cur
+    }
+
+    open spec fn ghost_advance(&self, _exec_iter: &Range<A>) -> RangeGhostIterator<A> {
+        RangeGhostIterator { cur: self.cur.spec_forward_checked(1).unwrap(), ..*self }
+    }
+}
+
+impl<A: StepSpec + core::iter::Step> crate::view::View for RangeGhostIterator<A> {
+    type V = Seq<A>;
+    // generate seq![start, start + 1, start + 2, ..., cur - 1]
+    open spec fn view(&self) -> Seq<A> {
+        Seq::new(
+            self.start.spec_steps_between_int(self.cur) as nat,
+            |i: int| self.start.spec_forward_checked_int(i).unwrap(),
+        )
     }
 }
 
@@ -82,6 +133,9 @@ macro_rules! step_specs {
                 end - self
             }
             open spec fn spec_forward_checked(self, count: usize) -> Option<Self> {
+                self.spec_forward_checked_int(count as int)
+            }
+            open spec fn spec_forward_checked_int(self, count: int) -> Option<Self> {
                 if self + count <= $t::MAX {
                     Some((self + count) as $t)
                 } else {
@@ -89,6 +143,9 @@ macro_rules! step_specs {
                 }
             }
             open spec fn spec_backward_checked(self, count: usize) -> Option<Self> {
+                self.spec_backward_checked_int(count as int)
+            }
+            open spec fn spec_backward_checked_int(self, count: int) -> Option<Self> {
                 if self - count >= $t::MIN {
                     Some((self - count) as $t)
                 } else {
diff --git a/source/pervasive/string.rs b/source/pervasive/string.rs
index b60df941b1..b9f357fba3 100644
--- a/source/pervasive/string.rs
+++ b/source/pervasive/string.rs
@@ -141,6 +141,7 @@ impl<'a> StrSlice<'a> {
             ret.view() == Seq::new(self.view().len(), |i| self.view().index(i) as u8)
     {
         let mut v = alloc::vec::Vec::new();
+        #[verifier::no_loop_invariant]
         for c in self.inner.as_bytes().iter() {
             v.push(*c);
         }
diff --git a/source/rust_verify/example/recursion.rs b/source/rust_verify/example/recursion.rs
index be805d19bc..8637e26b7d 100644
--- a/source/rust_verify/example/recursion.rs
+++ b/source/rust_verify/example/recursion.rs
@@ -101,7 +101,8 @@ fn exec_with_decreases(n: u64) -> u64
     }
 }
 
-#[verifier(external)]
+} // verus!
+
 fn main() {
     let args = std::env::args();
     for arg in args {
@@ -110,5 +111,3 @@ fn main() {
         }
     }
 }
-
-} // verus!
diff --git a/source/rust_verify/example/vectors.rs b/source/rust_verify/example/vectors.rs
index 9c01fe6e24..6219d41c3f 100644
--- a/source/rust_verify/example/vectors.rs
+++ b/source/rust_verify/example/vectors.rs
@@ -101,7 +101,8 @@ fn pust_test(t: Vec<u64>, y: u64)
     );
 }
 
-#[verifier(external)]
+} // verus!
+
 fn main() {
     let mut v = vec![0, 10, 20, 30, 40, 50, 60, 70, 80, 90];
     println!("{}", binary_search(&v, 70));
@@ -116,5 +117,3 @@ fn main() {
         println!("{}", x);
     }
 }
-
-} // verus!
diff --git a/source/rust_verify/src/attributes.rs b/source/rust_verify/src/attributes.rs
index 328f889beb..1a1e670601 100644
--- a/source/rust_verify/src/attributes.rs
+++ b/source/rust_verify/src/attributes.rs
@@ -244,6 +244,8 @@ pub(crate) enum Attr {
     Atomic,
     // specifies an invariant block
     InvariantBlock,
+    // mark that a loop was desugared from a for-loop in the syntax macro
+    ForLoop,
     // this proof function is a termination proof
     DecreasesBy,
     // in a spec function, check the body for violations of recommends
@@ -525,6 +527,7 @@ pub(crate) fn parse_attrs(
                     AttrTree::Fun(_, arg, None) if arg == "reveal_fn" => {
                         v.push(Attr::InternalRevealFn)
                     }
+                    AttrTree::Fun(_, arg, None) if arg == "for_loop" => v.push(Attr::ForLoop),
                     AttrTree::Fun(_, arg, Some(box [AttrTree::Fun(_, ident, None)]))
                         if arg == "prover" =>
                     {
@@ -669,6 +672,7 @@ pub(crate) struct VerifierAttrs {
     pub(crate) autospec: Option<String>,
     pub(crate) custom_req_err: Option<String>,
     pub(crate) bit_vector: bool,
+    pub(crate) for_loop: bool,
     pub(crate) atomic: bool,
     pub(crate) integer_ring: bool,
     pub(crate) decreases_by: bool,
@@ -723,6 +727,7 @@ pub(crate) fn get_verifier_attrs(
         autospec: None,
         custom_req_err: None,
         bit_vector: false,
+        for_loop: false,
         atomic: false,
         integer_ring: false,
         decreases_by: false,
@@ -772,6 +777,7 @@ pub(crate) fn get_verifier_attrs(
             Attr::Autospec(method_ident) => vs.autospec = Some(method_ident),
             Attr::CustomReqErr(s) => vs.custom_req_err = Some(s.clone()),
             Attr::BitVector => vs.bit_vector = true,
+            Attr::ForLoop => vs.for_loop = true,
             Attr::Atomic => vs.atomic = true,
             Attr::IntegerRing => vs.integer_ring = true,
             Attr::DecreasesBy => vs.decreases_by = true,
diff --git a/source/rust_verify/src/rust_to_vir_expr.rs b/source/rust_verify/src/rust_to_vir_expr.rs
index 283f60c73b..1102295d2d 100644
--- a/source/rust_verify/src/rust_to_vir_expr.rs
+++ b/source/rust_verify/src/rust_to_vir_expr.rs
@@ -1564,7 +1564,13 @@ pub(crate) fn expr_to_vir_innermost<'tcx>(
             let mut body = block_to_vir(bctx, block, &expr.span, &typ, ExprModifier::REGULAR)?;
             let header = vir::headers::read_header(&mut body)?;
             let label = label.map(|l| l.ident.to_string());
-            mk_expr(ExprX::Loop { label, cond: None, body, invs: header.loop_invariants() })
+            mk_expr(ExprX::Loop {
+                is_for_loop: expr_vattrs.for_loop,
+                label,
+                cond: None,
+                body,
+                invs: header.loop_invariants(),
+            })
         }
         ExprKind::Loop(
             Block {
@@ -1617,7 +1623,13 @@ pub(crate) fn expr_to_vir_innermost<'tcx>(
                 return err_span(expr.span, "termination checking of loops is not supported");
             }
             let label = label.map(|l| l.ident.to_string());
-            mk_expr(ExprX::Loop { label, cond, body, invs: header.loop_invariants() })
+            mk_expr(ExprX::Loop {
+                is_for_loop: false,
+                label,
+                cond,
+                body,
+                invs: header.loop_invariants(),
+            })
         }
         ExprKind::Ret(expr) => {
             let expr = match expr {
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index 59a4f4f41c..20bd1c7fc5 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -990,10 +990,21 @@ test_verify_one_file! {
         fn test_loop() {
             let mut n: u64 = 0;
             for x in iter: 0..10
-                invariant n == iter.start * 3,
+                invariant n == iter.cur * 3,
+            {
+                assert(x < 10);
+                assert(x == iter.cur);
+                n += 3;
+            }
+            assert(n == 30);
+        }
+
+        fn test_loop_peek() {
+            let mut n: u64 = 0;
+            for x in 0..10
+                invariant n == x * 3,
             {
                 assert(x < 10);
-                assert(x == iter.start - 1);
                 n += 3;
             }
             assert(n == 30);
@@ -1002,10 +1013,10 @@ test_verify_one_file! {
         fn test_loop_fail() {
             let mut n: u64 = 0;
             for x in iter: 0..10
-                invariant n == iter.start * 3,
+                invariant n == iter.cur * 3,
             {
                 assert(x < 9); // FAILS
-                assert(x == iter.start - 1);
+                assert(x == iter.cur);
                 n += 3;
             }
             assert(n == 30);
@@ -1021,11 +1032,11 @@ test_verify_one_file! {
             let mut end = 10;
             for x in iter: 0..end
                 invariant
-                    n == iter.start * 3,
+                    n == iter.cur * 3,
                     end == 10,
             {
                 assert(x < 10);
-                assert(x == iter.start - 1);
+                assert(x == iter.cur);
                 n += 3;
             }
             assert(n == 30);
@@ -1036,14 +1047,49 @@ test_verify_one_file! {
             let mut end = 10;
             for x in iter: 0..end
                 invariant
-                    n == iter.start * 3,
+                    n == iter.cur * 3,
                     end == 10,
             {
                 assert(x < 10); // FAILS
-                assert(x == iter.start - 1);
+                assert(x == iter.cur);
                 n += 3;
                 end = end + 0; // causes end to be non-constant, so loop needs more invariants
             }
         }
     } => Err(e) => assert_one_fails(e)
 }
+
+test_verify_one_file! {
+    #[test] for_loop3 verus_code! {
+        use vstd::prelude::*;
+        fn test_loop(n: u32) -> (v: Vec<u32>)
+            ensures
+                v.len() == n,
+                forall|i: int| 0 <= i < n ==> v[i] == i,
+        {
+            let mut v: Vec<u32> = Vec::new();
+            for i in iter: 0..n
+                invariant
+                    v@ =~= iter@,
+            {
+                v.push(i);
+            }
+            v
+        }
+
+        fn test_loop_fail(n: u32) -> (v: Vec<u32>)
+            ensures
+                v.len() == n,
+                forall|i: int| 0 <= i < n ==> v[i] == i,
+        {
+            let mut v: Vec<u32> = Vec::new();
+            for i in iter: 0..n
+                invariant
+                    v@ =~= iter@, // FAILS
+            {
+                v.push(0);
+            }
+            v
+        }
+    } => Err(e) => assert_one_fails(e)
+}
diff --git a/source/vir/src/ast.rs b/source/vir/src/ast.rs
index 3962de3b1b..070b3f5e87 100644
--- a/source/vir/src/ast.rs
+++ b/source/vir/src/ast.rs
@@ -681,7 +681,13 @@ pub enum ExprX {
     /// Match (Note: ast_simplify replaces Match with other expressions)
     Match(Expr, Arms),
     /// Loop (either "while", cond = Some(...), or "loop", cond = None), with invariants
-    Loop { label: Option<String>, cond: Option<Expr>, body: Expr, invs: LoopInvariants },
+    Loop {
+        is_for_loop: bool,
+        label: Option<String>,
+        cond: Option<Expr>,
+        body: Expr,
+        invs: LoopInvariants,
+    },
     /// Open invariant
     OpenInvariant(Expr, Binder<Typ>, Expr, InvAtomicity),
     /// Return from function
diff --git a/source/vir/src/ast_to_sst.rs b/source/vir/src/ast_to_sst.rs
index 791fcdda52..0ad72d17a0 100644
--- a/source/vir/src/ast_to_sst.rs
+++ b/source/vir/src/ast_to_sst.rs
@@ -492,8 +492,10 @@ fn loop_body_has_break(loop_label: &Option<String>, body: &Expr) -> bool {
 
 pub fn can_control_flow_reach_after_loop(expr: &Expr) -> bool {
     match &expr.x {
-        ExprX::Loop { label, cond: None, body, invs: _ } => loop_body_has_break(label, body),
-        ExprX::Loop { label: _, cond: Some(_), body: _, invs: _ } => true,
+        ExprX::Loop { is_for_loop: _, label, cond: None, body, invs: _ } => {
+            loop_body_has_break(label, body)
+        }
+        ExprX::Loop { is_for_loop: _, label: _, cond: Some(_), body: _, invs: _ } => true,
         _ => {
             panic!("expected while loop");
         }
@@ -1873,7 +1875,7 @@ pub(crate) fn expr_to_stm_opt(
         ExprX::Match(..) => {
             panic!("internal error: Match should have been simplified by ast_simplify")
         }
-        ExprX::Loop { label, cond, body, invs } => {
+        ExprX::Loop { is_for_loop, label, cond, body, invs } => {
             let has_break = loop_body_has_break(label, body);
             let simple_invs = invs.iter().all(|inv| inv.kind == LoopInvariantKind::Invariant);
             let simple_while = !has_break && simple_invs && cond.is_some();
@@ -1930,6 +1932,7 @@ pub(crate) fn expr_to_stm_opt(
             let while_stm = Spanned::new(
                 expr.span.clone(),
                 StmX::Loop {
+                    is_for_loop: *is_for_loop,
                     label: label.clone(),
                     cond: cnd,
                     body: stms_to_one_stm(&body.span, stms1),
diff --git a/source/vir/src/ast_visitor.rs b/source/vir/src/ast_visitor.rs
index 7427582e94..e3c316bbfa 100644
--- a/source/vir/src/ast_visitor.rs
+++ b/source/vir/src/ast_visitor.rs
@@ -433,7 +433,7 @@ where
                         map.pop_scope();
                     }
                 }
-                ExprX::Loop { label: _, cond, body, invs } => {
+                ExprX::Loop { is_for_loop: _, label: _, cond, body, invs } => {
                     if let Some(cond) = cond {
                         expr_visitor_control_flow!(expr_visitor_dfs(cond, map, mf));
                     }
@@ -867,7 +867,7 @@ where
             });
             ExprX::Match(expr1, Arc::new(arms?))
         }
-        ExprX::Loop { label, cond, body, invs } => {
+        ExprX::Loop { is_for_loop, label, cond, body, invs } => {
             let cond =
                 cond.as_ref().map(|e| map_expr_visitor_env(e, map, env, fe, fs, ft)).transpose()?;
             let body = map_expr_visitor_env(body, map, env, fe, fs, ft)?;
@@ -876,7 +876,13 @@ where
                 let e1 = map_expr_visitor_env(&inv.inv, map, env, fe, fs, ft)?;
                 invs1.push(crate::ast::LoopInvariant { inv: e1, ..inv.clone() });
             }
-            ExprX::Loop { label: label.clone(), cond, body, invs: Arc::new(invs1) }
+            ExprX::Loop {
+                is_for_loop: *is_for_loop,
+                label: label.clone(),
+                cond,
+                body,
+                invs: Arc::new(invs1),
+            }
         }
         ExprX::Return(e1) => {
             let e1 = match e1 {
diff --git a/source/vir/src/early_exit_cf.rs b/source/vir/src/early_exit_cf.rs
index bab6851dbc..ee477acecd 100644
--- a/source/vir/src/early_exit_cf.rs
+++ b/source/vir/src/early_exit_cf.rs
@@ -85,7 +85,7 @@ fn expr_get_early_exits_rec(
             | ExprX::AssertBy { .. }
             | ExprX::RevealString(_) => VisitorControlFlow::Return,
             ExprX::AssertQuery { .. } => VisitorControlFlow::Return,
-            ExprX::Loop { label: _, cond, body, invs: _ } => {
+            ExprX::Loop { is_for_loop: _, label: _, cond, body, invs: _ } => {
                 if let Some(cond) = cond {
                     expr_get_early_exits_rec(cond, in_loop, scope_map, results);
                 }
diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index c2be8bf7fa..65f19ea2a3 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -941,7 +941,7 @@ fn check_expr_handle_mut_arg(
             }
             Ok(final_mode)
         }
-        ExprX::Loop { label: _, cond, body, invs } => {
+        ExprX::Loop { is_for_loop: _, label: _, cond, body, invs } => {
             // We could also allow this for proof, if we check it for termination
             if typing.check_ghost_blocks && typing.block_ghostness != Ghost::Exec {
                 return Err(error(&expr.span, "cannot use while in proof or spec mode"));
diff --git a/source/vir/src/poly.rs b/source/vir/src/poly.rs
index dfe518926d..41e0e04816 100644
--- a/source/vir/src/poly.rs
+++ b/source/vir/src/poly.rs
@@ -672,7 +672,7 @@ fn poly_expr(ctx: &Ctx, state: &mut State, expr: &Expr) -> Expr {
             mk_expr_typ(&t, ExprX::If(e0, e1.clone(), Some(e2)))
         }
         ExprX::Match(..) => panic!("Match should already be removed"),
-        ExprX::Loop { label, cond, body, invs } => {
+        ExprX::Loop { is_for_loop, label, cond, body, invs } => {
             let cond = cond.as_ref().map(|e| coerce_expr_to_native(ctx, &poly_expr(ctx, state, e)));
             let body = poly_expr(ctx, state, body);
             let invs = invs.iter().map(|inv| crate::ast::LoopInvariant {
@@ -680,7 +680,13 @@ fn poly_expr(ctx: &Ctx, state: &mut State, expr: &Expr) -> Expr {
                 ..inv.clone()
             });
             let invs = Arc::new(invs.collect());
-            mk_expr(ExprX::Loop { label: label.clone(), cond, body, invs })
+            mk_expr(ExprX::Loop {
+                is_for_loop: *is_for_loop,
+                label: label.clone(),
+                cond,
+                body,
+                invs,
+            })
         }
         ExprX::OpenInvariant(inv, binder, body, atomicity) => {
             let inv = coerce_expr_to_poly(ctx, &poly_expr(ctx, state, inv));
diff --git a/source/vir/src/split_expression.rs b/source/vir/src/split_expression.rs
index 9a665e92e3..68aace94cc 100644
--- a/source/vir/src/split_expression.rs
+++ b/source/vir/src/split_expression.rs
@@ -672,7 +672,7 @@ fn visit_split_stm(
             state.pop_fuel_scope();
             Ok(Spanned::new(stm.span.clone(), StmX::If(cond.clone(), lhs, rhs)))
         }
-        StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+        StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
             let cond = if let Some((cond_stm, cond_exp)) = cond {
                 let cond_stm = visit_split_stm(ctx, state, diagnostics, cond_stm)?;
                 Some((cond_stm, cond_exp.clone()))
@@ -685,6 +685,7 @@ fn visit_split_stm(
             Ok(Spanned::new(
                 stm.span.clone(),
                 StmX::Loop {
+                    is_for_loop: *is_for_loop,
                     label: label.clone(),
                     cond,
                     body,
diff --git a/source/vir/src/sst.rs b/source/vir/src/sst.rs
index 864c66ab95..9e8fe0dbea 100644
--- a/source/vir/src/sst.rs
+++ b/source/vir/src/sst.rs
@@ -171,6 +171,7 @@ pub enum StmX {
         // 1. cond = Some(...), all invs are true on entry and exit, no break statements
         // 2. cond = None, invs may have false at_entry/at_exit, may have break statements
         // Any while loop not satisfying (1) is converted to (2).
+        is_for_loop: bool,
         label: Option<String>,
         cond: Option<(Stm, Exp)>,
         body: Stm,
diff --git a/source/vir/src/sst_to_air.rs b/source/vir/src/sst_to_air.rs
index 6d0767b040..4c1df9df8c 100644
--- a/source/vir/src/sst_to_air.rs
+++ b/source/vir/src/sst_to_air.rs
@@ -1438,6 +1438,7 @@ struct PostConditionInfo {
 
 #[derive(Debug)]
 struct LoopInfo {
+    is_for_loop: bool,
     label: Option<String>,
     some_cond: bool,
     invs_entry: Arc<Vec<(Span, Expr, Option<Arc<String>>)>>,
@@ -2045,6 +2046,11 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
                     state.loop_infos.last().expect("inside loop")
                 };
                 assert!(!loop_info.some_cond); // AST-to-SST conversion must eliminate the cond
+                if loop_info.is_for_loop && !*is_break {
+                    // At the very least, the syntax macro will need to advance the ghost iterator
+                    // at each continue.
+                    return Err(error(&stm.span, "for-loops do not yet support continue"));
+                }
                 let invs = if *is_break {
                     loop_info.invs_exit.clone()
                 } else {
@@ -2110,7 +2116,7 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
             }
             stmts
         }
-        StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+        StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
             let (cond_stm, pos_assume, neg_assume) = if let Some((cond_stm, cond_exp)) = cond {
                 let pos_cond = exp_to_expr(ctx, &cond_exp, expr_ctxt)?;
                 let neg_cond = Arc::new(ExprX::Unary(air::ast::UnaryOp::Not, pos_cond.clone()));
@@ -2206,6 +2212,7 @@ fn stm_to_stmts(ctx: &Ctx, state: &mut State, stm: &Stm) -> Result<Vec<Stmt>, Vi
                 air_body.push(pos_assume);
             }
             let loop_info = LoopInfo {
+                is_for_loop: *is_for_loop,
                 label: label.clone(),
                 some_cond: cond.is_some(),
                 invs_entry: invs_entry.clone(),
diff --git a/source/vir/src/sst_vars.rs b/source/vir/src/sst_vars.rs
index c2fa9633c4..3c6a3bef6c 100644
--- a/source/vir/src/sst_vars.rs
+++ b/source/vir/src/sst_vars.rs
@@ -116,7 +116,7 @@ pub(crate) fn stm_assign(
             *assigned = pre_assigned;
             Spanned::new(stm.span.clone(), StmX::If(cond.clone(), lhs, rhs))
         }
-        StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+        StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
             let mut pre_modified = modified.clone();
             *modified = IndexSet::new();
             let cond = if let Some((cond_stm, cond_exp)) = cond {
@@ -146,6 +146,7 @@ pub(crate) fn stm_assign(
                 typ_inv_vars.push((x.clone(), declared[x].clone()));
             }
             let loop_x = StmX::Loop {
+                is_for_loop: *is_for_loop,
                 label: label.clone(),
                 cond,
                 body,
diff --git a/source/vir/src/sst_visitor.rs b/source/vir/src/sst_visitor.rs
index 14198a5d18..10eb7db027 100644
--- a/source/vir/src/sst_visitor.rs
+++ b/source/vir/src/sst_visitor.rs
@@ -175,7 +175,15 @@ where
                 StmX::AssertQuery { body, mode: _, typ_inv_vars: _ } => {
                     expr_visitor_control_flow!(stm_visitor_dfs(body, f));
                 }
-                StmX::Loop { label: _, cond, body, invs: _, typ_inv_vars: _, modified_vars: _ } => {
+                StmX::Loop {
+                    is_for_loop: _,
+                    label: _,
+                    cond,
+                    body,
+                    invs: _,
+                    typ_inv_vars: _,
+                    modified_vars: _,
+                } => {
                     if let Some((cond_stm, _cond_exp)) = cond {
                         expr_visitor_control_flow!(stm_visitor_dfs(cond_stm, f));
                     }
@@ -244,7 +252,15 @@ where
             StmX::If(exp, _s1, _s2) => {
                 expr_visitor_control_flow!(exp_visitor_dfs(exp, &mut ScopeMap::new(), f))
             }
-            StmX::Loop { label: _, cond, body: _, invs, typ_inv_vars: _, modified_vars: _ } => {
+            StmX::Loop {
+                is_for_loop: _,
+                label: _,
+                cond,
+                body: _,
+                invs,
+                typ_inv_vars: _,
+                modified_vars: _,
+            } => {
                 if let Some((_cond_stm, cond_exp)) = cond {
                     expr_visitor_control_flow!(exp_visitor_dfs(cond_exp, &mut ScopeMap::new(), f));
                 }
@@ -601,7 +617,7 @@ where
             let stm = Spanned::new(stm.span.clone(), StmX::If(cond.clone(), lhs, rhs));
             fs(&stm)
         }
-        StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+        StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
             let cond = if let Some((cond_stm, cond_exp)) = cond {
                 let cond_stm = map_stm_visitor(cond_stm, fs)?;
                 Some((cond_stm, cond_exp.clone()))
@@ -612,6 +628,7 @@ where
             let stm = Spanned::new(
                 stm.span.clone(),
                 StmX::Loop {
+                    is_for_loop: *is_for_loop,
                     label: label.clone(),
                     cond,
                     body,
@@ -680,7 +697,7 @@ where
             let rhs = rhs.as_ref().map(|rhs| fs(rhs)).transpose()?;
             Ok(Spanned::new(stm.span.clone(), StmX::If(cond.clone(), lhs, rhs)))
         }
-        StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+        StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
             let cond = if let Some((cond_stm, cond_exp)) = cond {
                 let cond_stm = fs(cond_stm)?;
                 Some((cond_stm, cond_exp.clone()))
@@ -691,6 +708,7 @@ where
             Ok(Spanned::new(
                 stm.span.clone(),
                 StmX::Loop {
+                    is_for_loop: *is_for_loop,
                     label: label.clone(),
                     cond,
                     body,
@@ -761,7 +779,7 @@ where
                 },
             ))
         }
-        StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+        StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
             let mut typ_inv_vars2 = vec![];
             for (uid, typ) in typ_inv_vars.iter() {
                 typ_inv_vars2.push((uid.clone(), ft(typ)?));
@@ -769,6 +787,7 @@ where
             Ok(Spanned::new(
                 stm.span.clone(),
                 StmX::Loop {
+                    is_for_loop: *is_for_loop,
                     label: label.clone(),
                     cond: cond.clone(),
                     body: body.clone(),
@@ -867,7 +886,7 @@ where
                 let exp = fe(exp)?;
                 Spanned::new(span, StmX::If(exp, s1.clone(), s2.clone()))
             }
-            StmX::Loop { label, cond, body, invs, typ_inv_vars, modified_vars } => {
+            StmX::Loop { is_for_loop, label, cond, body, invs, typ_inv_vars, modified_vars } => {
                 let cond = if let Some((cond_stm, cond_exp)) = cond {
                     let cond_exp = fe(cond_exp)?;
                     Some((cond_stm.clone(), cond_exp))
@@ -881,6 +900,7 @@ where
                 Spanned::new(
                     span,
                     StmX::Loop {
+                        is_for_loop: *is_for_loop,
                         label: label.clone(),
                         cond,
                         body: body.clone(),

From 31cd01f1e77bb6dacd3dda03ddf4502bccd0d173 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Thu, 4 Jan 2024 21:16:46 -0800
Subject: [PATCH 06/19] Refactor syntax macro visit order on loop invariants

---
 source/builtin_macros/src/syntax.rs     | 34 +++++++++++--------------
 source/rust_verify/example/recursion.rs |  7 ++---
 source/rust_verify/example/vectors.rs   |  5 +---
 3 files changed, 18 insertions(+), 28 deletions(-)

diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index 752c225c3c..5ad5e4dfe1 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -14,8 +14,8 @@ use syn_verus::token::{Brace, Bracket, Paren, Semi};
 use syn_verus::visit_mut::{
     visit_block_mut, visit_expr_loop_mut, visit_expr_mut, visit_expr_while_mut, visit_field_mut,
     visit_impl_item_method_mut, visit_item_const_mut, visit_item_enum_mut, visit_item_fn_mut,
-    visit_item_static_mut, visit_item_struct_mut, visit_local_mut, visit_trait_item_method_mut,
-    VisitMut,
+    visit_item_static_mut, visit_item_struct_mut, visit_local_mut, visit_specification_mut,
+    visit_trait_item_method_mut, VisitMut,
 };
 use syn_verus::{
     braced, bracketed, parenthesized, parse_macro_input, AttrStyle, Attribute, BareFnArg, BinOp,
@@ -1263,37 +1263,27 @@ impl Visitor {
     ) {
         // TODO: wrap specs inside ghost blocks
         self.inside_ghost += 1;
-        if let Some(Invariant { token, mut exprs }) = invariants {
+        if let Some(Invariant { token, exprs }) = invariants {
             if exprs.exprs.len() > 0 {
-                for expr in exprs.exprs.iter_mut() {
-                    self.visit_expr_mut(expr);
-                }
                 stmts.push(stmt_with_semi!(token.span => ::builtin::invariant([#exprs])));
             }
         }
-        if let Some(InvariantEnsures { token, mut exprs }) = invariant_ensures {
+        if let Some(InvariantEnsures { token, exprs }) = invariant_ensures {
             if exprs.exprs.len() > 0 {
-                for expr in exprs.exprs.iter_mut() {
-                    self.visit_expr_mut(expr);
-                }
                 stmts.push(stmt_with_semi!(token.span => ::builtin::invariant_ensures([#exprs])));
             }
         }
-        if let Some(Ensures { token, mut exprs, attrs }) = ensures {
+        if let Some(Ensures { token, exprs, attrs }) = ensures {
             if attrs.len() > 0 {
                 let err = "outer attributes only allowed on function's ensures";
                 let expr = Expr::Verbatim(quote_spanned!(token.span => compile_error!(#err)));
                 stmts.push(Stmt::Semi(expr, Semi { spans: [token.span] }));
             } else if exprs.exprs.len() > 0 {
-                for expr in exprs.exprs.iter_mut() {
-                    self.visit_expr_mut(expr);
-                }
                 stmts.push(stmt_with_semi!(token.span => ::builtin::ensures([#exprs])));
             }
         }
-        if let Some(Decreases { token, mut exprs }) = decreases {
-            for expr in exprs.exprs.iter_mut() {
-                self.visit_expr_mut(expr);
+        if let Some(Decreases { token, exprs }) = decreases {
+            for expr in exprs.exprs.iter() {
                 if matches!(expr, Expr::Tuple(..)) {
                     let err = "decreases cannot be a tuple; use `decreases x, y` rather than `decreases (x, y)`";
                     let expr = Expr::Verbatim(quote_spanned!(token.span => compile_error!(#err)));
@@ -2433,6 +2423,7 @@ impl VisitMut for Visitor {
     }
 
     fn visit_expr_while_mut(&mut self, expr_while: &mut ExprWhile) {
+        visit_expr_while_mut(self, expr_while);
         let invariants = self.take_ghost(&mut expr_while.invariant);
         let invariant_ensures = self.take_ghost(&mut expr_while.invariant_ensures);
         let ensures = self.take_ghost(&mut expr_while.ensures);
@@ -2440,10 +2431,10 @@ impl VisitMut for Visitor {
         let mut stmts: Vec<Stmt> = Vec::new();
         self.add_loop_specs(&mut stmts, invariants, invariant_ensures, ensures, decreases);
         expr_while.body.stmts.splice(0..0, stmts);
-        visit_expr_while_mut(self, expr_while);
     }
 
     fn visit_expr_loop_mut(&mut self, expr_loop: &mut ExprLoop) {
+        visit_expr_loop_mut(self, expr_loop);
         let invariants = self.take_ghost(&mut expr_loop.invariant);
         let invariant_ensures = self.take_ghost(&mut expr_loop.invariant_ensures);
         let ensures = self.take_ghost(&mut expr_loop.ensures);
@@ -2451,7 +2442,12 @@ impl VisitMut for Visitor {
         let mut stmts: Vec<Stmt> = Vec::new();
         self.add_loop_specs(&mut stmts, invariants, invariant_ensures, ensures, decreases);
         expr_loop.body.stmts.splice(0..0, stmts);
-        visit_expr_loop_mut(self, expr_loop);
+    }
+
+    fn visit_specification_mut(&mut self, spec: &mut syn_verus::Specification) {
+        self.inside_ghost += 1;
+        visit_specification_mut(self, spec);
+        self.inside_ghost -= 1;
     }
 
     fn visit_local_mut(&mut self, local: &mut Local) {
diff --git a/source/rust_verify/example/recursion.rs b/source/rust_verify/example/recursion.rs
index 8637e26b7d..7b82043178 100644
--- a/source/rust_verify/example/recursion.rs
+++ b/source/rust_verify/example/recursion.rs
@@ -68,17 +68,14 @@ fn compute_arith_sum(n: u64) -> (sum: u64)
     ensures
         arith_sum_int(n as int) == sum,
 {
-    let mut i: u64 = 0;
     let mut sum: u64 = 0;
-    while i < n
+    for i in 0..n
         invariant
             n < 100,
-            i <= n,
             arith_sum_int(i as int) == sum,
             sum <= 100 * i,
     {
-        i = i + 1;
-        sum = sum + i;
+        sum = sum + (i + 1);
     }
     sum
 }
diff --git a/source/rust_verify/example/vectors.rs b/source/rust_verify/example/vectors.rs
index 6219d41c3f..fb69d74669 100644
--- a/source/rust_verify/example/vectors.rs
+++ b/source/rust_verify/example/vectors.rs
@@ -39,8 +39,7 @@ fn reverse(v: &mut Vec<u64>)
 {
     let length = v.len();
     let ghost v1 = v@;
-    let mut n: usize = 0;
-    while n < length / 2
+    for n in 0..(length / 2)
         invariant
             length == v.len(),
             forall|i: int| 0 <= i < n ==> v[i] == v1[length - i - 1],
@@ -51,8 +50,6 @@ fn reverse(v: &mut Vec<u64>)
         let y = v[length - 1 - n];
         v.set(n, y);
         v.set(length - 1 - n, x);
-
-        n = n + 1;
     }
 }
 

From 01bac5f1f79a6c9da456ca97f30080c32d623c0e Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Fri, 5 Jan 2024 10:48:11 -0800
Subject: [PATCH 07/19] Switch some parse_quote_spanned to Verbatim
 quote_spanned

---
 source/builtin_macros/src/syntax.rs | 35 ++++++++++++++++-------------
 source/pervasive/pervasive.rs       |  2 ++
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index 5ad5e4dfe1..669199b88f 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -1401,23 +1401,23 @@ impl Visitor {
         //                  builtin::infer_spec_for_loop_iter(
         //                      &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
         //                          &core::iter::IntoIterator::into_iter(e)))),
-        let exec_inv: Expr = parse_quote_spanned!(span =>
+        let exec_inv: Expr = Expr::Verbatim(quote_spanned!(span =>
             #[verifier::custom_err(#exec_inv_msg)]
             vstd::pervasive::ForLoopGhostIterator::exec_invariant(&#x_ghost_iter, &#x_exec_iter)
-        );
-        let ghost_inv: Expr = parse_quote_spanned!(span =>
+        ));
+        let ghost_inv: Expr = Expr::Verbatim(quote_spanned!(span =>
             #[verifier::custom_err(#ghost_inv_msg)]
             vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&#x_ghost_iter,
                 builtin::infer_spec_for_loop_iter(
                     &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
                         &core::iter::IntoIterator::into_iter(#expr_inv))))
-        );
+        ));
         let invariant_ensure = if let Some(mut invariant) = invariant {
             for inv in &mut invariant.exprs.exprs {
-                *inv = parse_quote_spanned!(span => {
+                *inv = Expr::Verbatim(quote_spanned!(span => {
                     let #pat = vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&#x_ghost_iter);
                     #inv
-                });
+                }));
             }
             if no_loop_invariant.is_none() {
                 invariant.exprs.exprs.insert(0, exec_inv);
@@ -1446,19 +1446,24 @@ impl Visitor {
             None
         };
         self.add_loop_specs(&mut stmts, None, invariant_ensure, ensure, decreases);
+        let body_exec = Expr::Verbatim(quote_spanned!(span =>
+            if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_exec_iter)
+            #body else { break }
+        ));
         let mut body: Block = if no_loop_invariant.is_none() {
-            parse_quote_spanned!(span => {
-                if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_exec_iter)
-                #body else { break }
+            let body_ghost = Expr::Verbatim(quote_spanned!(span =>
                 #[verifier::proof_block] {
                     #x_ghost_iter = vstd::pervasive::ForLoopGhostIterator::ghost_advance(
                         &#x_ghost_iter, &#x_exec_iter);
                 }
+            ));
+            parse_quote_spanned!(span => {
+                #body_exec
+                #body_ghost
             })
         } else {
             parse_quote_spanned!(span => {
-                if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_exec_iter)
-                #body else { break }
+                #body_exec
             })
         };
         body.stmts.splice(0..0, stmts);
@@ -1466,7 +1471,7 @@ impl Visitor {
         loop_expr.label = label;
         loop_expr.attrs = attrs;
         let full_loop: Expr = if no_loop_invariant.is_none() {
-            parse_quote_spanned!(span => {
+            Expr::Verbatim(quote_spanned!(span => {
                 #[allow(non_snake_case)]
                 let mut #x_exec_iter = core::iter::IntoIterator::into_iter(#expr);
                 #[allow(non_snake_case)]
@@ -1476,13 +1481,13 @@ impl Visitor {
                     #x_ghost_iter = vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(&#x_exec_iter);
                 }
                 #loop_expr
-            })
+            }))
         } else {
-            parse_quote_spanned!(span => {
+            Expr::Verbatim(quote_spanned!(span => {
                 #[allow(non_snake_case)]
                 let mut #x_exec_iter = core::iter::IntoIterator::into_iter(#expr);
                 #loop_expr
-            })
+            }))
         };
         full_loop
     }
diff --git a/source/pervasive/pervasive.rs b/source/pervasive/pervasive.rs
index 6a9846029d..aa54647055 100644
--- a/source/pervasive/pervasive.rs
+++ b/source/pervasive/pervasive.rs
@@ -64,6 +64,7 @@ pub trait ForLoopGhostIterator {
     spec fn ghost_peek_next(&self) -> Self::Item;
 
     // At the end of the for loop, advance to the next position.
+    // Future TODO: this may be better as a proof function
     spec fn ghost_advance(&self, exec_iter: &Self::ExecIter) -> Self where Self: Sized;
 }
 
@@ -71,6 +72,7 @@ pub trait ForLoopGhostIteratorNew {
     type GhostIter;
 
     // Create a new ghost iterator from an exec iterator
+    // Future TODO: this may be better as a proof function
     spec fn ghost_iter(&self) -> Self::GhostIter;
 }
 

From 7187d6b0400580423f026a4698923af1779573f8 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Fri, 5 Jan 2024 14:41:42 -0800
Subject: [PATCH 08/19] Add Vec custom iterator example

---
 source/rust_verify_test/tests/loops.rs | 118 +++++++++++++++++++++++++
 1 file changed, 118 insertions(+)

diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index 20bd1c7fc5..c0fe841f88 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -1093,3 +1093,121 @@ test_verify_one_file! {
         }
     } => Err(e) => assert_one_fails(e)
 }
+
+test_verify_one_file! {
+    #[test] for_loop_vec_custom_iterator verus_code! {
+        use vstd::prelude::*;
+        pub struct VecIterCopy<'a, T: 'a> {
+            pub vec: &'a Vec<T>,
+            pub cur: usize,
+        }
+
+        impl<'a, T: Copy> Iterator for VecIterCopy<'a, T> {
+            type Item = T;
+            fn next(&mut self) -> (item: Option<T>)
+                ensures
+                    self.vec == old(self).vec,
+                    old(self).cur < self.vec.len() ==> self.cur == old(self).cur + 1,
+                    old(self).cur < self.vec.len() ==> item == Some(self.vec[old(self).cur as int]),
+                    old(self).cur >= self.vec.len() ==> item.is_none() && self.cur == old(self).cur,
+            {
+                if self.cur < self.vec.len() {
+                    let item = self.vec[self.cur];
+                    self.cur = self.cur + 1;
+                    Some(item)
+                } else {
+                    None
+                }
+            }
+        }
+
+        pub struct VecGhostIterCopy<'a, T> {
+            pub seq: Seq<T>,
+            pub cur: int,
+            pub unused_phantom_dummy: &'a Vec<T>, // use the 'a parameter, which I don't even want; ugh
+        }
+
+        impl<'a, T: 'a> vstd::pervasive::ForLoopGhostIteratorNew for VecIterCopy<'a, T> {
+            type GhostIter = VecGhostIterCopy<'a, T>;
+
+            open spec fn ghost_iter(&self) -> VecGhostIterCopy<'a, T> {
+                VecGhostIterCopy {
+                    seq: self.vec@,
+                    cur: 0,
+                    unused_phantom_dummy: &self.vec,
+                }
+            }
+        }
+
+        impl<'a, T: 'a> vstd::pervasive::ForLoopGhostIterator for VecGhostIterCopy<'a, T> {
+            type ExecIter = VecIterCopy<'a, T>;
+            type Item = T;
+            type Decrease = int;
+
+            open spec fn exec_invariant(&self, exec_iter: &VecIterCopy<'a, T>) -> bool {
+                &&& self.seq == exec_iter.vec@
+                &&& self.cur == exec_iter.cur
+            }
+
+            open spec fn ghost_invariant(&self, init: Option<&Self>) -> bool {
+                &&& 0 <= self.cur <= self.seq.len()
+                &&& if let Some(init) = init {
+                        init.seq == self.seq
+                    } else {
+                        true
+                    }
+            }
+
+            open spec fn ghost_condition(&self) -> bool {
+                self.cur < self.seq.len()
+            }
+
+            open spec fn ghost_decrease(&self) -> int {
+                self.seq.len() - self.cur
+            }
+
+            open spec fn ghost_peek_next(&self) -> T {
+                self.seq[self.cur]
+            }
+
+            open spec fn ghost_advance(&self, _exec_iter: &VecIterCopy<T>) -> VecGhostIterCopy<'a, T> {
+                VecGhostIterCopy { cur: self.cur + 1, ..*self }
+            }
+        }
+
+        impl<'a, T: 'a> vstd::view::View for VecGhostIterCopy<'a, T> {
+            type V = Seq<T>;
+
+            open spec fn view(&self) -> Seq<T> {
+                self.seq.subrange(0, self.cur)
+            }
+        }
+
+        spec fn vec_iter_copy_spec<'a, T: 'a>(vec: &'a Vec<T>) -> VecIterCopy<'a, T> {
+            VecIterCopy { vec, cur: 0 }
+        }
+
+        #[verifier::when_used_as_spec(vec_iter_copy_spec)]
+        fn vec_iter_copy<'a, T: 'a>(vec: &'a Vec<T>) -> (iter: VecIterCopy<'a, T>)
+            ensures
+                iter == (VecIterCopy { vec, cur: 0 }),
+        {
+            VecIterCopy { vec, cur: 0 }
+        }
+
+        fn all_positive(v: &Vec<u8>) -> (b: bool)
+            ensures
+                b <==> (forall|i: int| 0 <= i < v.len() ==> v[i] > 0),
+        {
+            let mut b: bool = true;
+
+            for x in iter: vec_iter_copy(v)
+                invariant
+                    b <==> (forall|i: int| 0 <= i < iter.cur ==> v[i] > 0),
+            {
+                b = b && x > 0;
+            }
+            b
+        }
+    } => Ok(())
+}

From 6d66cc1d15187636aa014e7b4e3f4122418652ad Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Tue, 9 Jan 2024 10:47:35 -0800
Subject: [PATCH 09/19] Remove old ignored test

---
 source/rust_verify_test/tests/regression.rs | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/source/rust_verify_test/tests/regression.rs b/source/rust_verify_test/tests/regression.rs
index 8ed20d950e..319f06c7eb 100644
--- a/source/rust_verify_test/tests/regression.rs
+++ b/source/rust_verify_test/tests/regression.rs
@@ -542,15 +542,6 @@ test_verify_one_file! {
     } => Err(err) => assert_vir_error_msg(err, "unrecognized verifier attribute")
 }
 
-test_verify_one_file! {
-    #[ignore] #[test] test_for_loop_387_discussioncomment_5683342 verus_code! {
-        struct T{}
-        fn f(v: Vec<T>) {
-            for t in v {}
-        }
-    } => Err(err) => assert_vir_error_msg(err, "Verus does not yet support IntoIterator::into_iter")
-}
-
 test_verify_one_file! {
     #[test] test_empty_recommends_387_discussioncomment_5670055 verus_code! {
         pub open spec fn foo() -> bool

From 7a360307988117f6b886fde4dd7124d23365e64d Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Tue, 9 Jan 2024 14:03:00 -0800
Subject: [PATCH 10/19] Rewrite InferSpecForLoopIter in modes not in
 ast_simplify

---
 source/rust_verify/src/verifier.rs     |  8 +---
 source/rust_verify_test/tests/loops.rs | 10 ++++-
 source/vir/src/ast.rs                  |  2 +-
 source/vir/src/ast_simplify.rs         | 11 -----
 source/vir/src/context.rs              | 13 +-----
 source/vir/src/def.rs                  |  4 ++
 source/vir/src/modes.rs                | 56 ++++++++++++++++++++------
 7 files changed, 59 insertions(+), 45 deletions(-)

diff --git a/source/rust_verify/src/verifier.rs b/source/rust_verify/src/verifier.rs
index a680e8c98e..1a75feaaf2 100644
--- a/source/rust_verify/src/verifier.rs
+++ b/source/rust_verify/src/verifier.rs
@@ -1696,12 +1696,6 @@ impl Verifier {
             self.args.rlimit,
             interpreter_log_file,
             self.vstd_crate_name.clone(),
-            &self
-                .erasure_hints
-                .as_ref()
-                .expect("erasure_hints")
-                .erasure_modes
-                .infer_spec_for_loop_iter_modes,
         )?;
         vir::recursive_types::check_traits(&krate, &global_ctx)?;
         let krate = vir::ast_simplify::simplify_krate(&mut global_ctx, &krate)?;
@@ -2375,7 +2369,7 @@ impl Verifier {
         }
         check_crate_result.map_err(|e| (e, Vec::new()))?;
         let vir_crate = vir::autospec::resolve_autospec(&vir_crate).map_err(|e| (e, Vec::new()))?;
-        let erasure_modes =
+        let (vir_crate, erasure_modes) =
             vir::modes::check_crate(&vir_crate, true).map_err(|e| (e, Vec::new()))?;
 
         self.vir_crate = Some(vir_crate.clone());
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index c0fe841f88..8806b93775 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -1097,6 +1097,12 @@ test_verify_one_file! {
 test_verify_one_file! {
     #[test] for_loop_vec_custom_iterator verus_code! {
         use vstd::prelude::*;
+
+        #[verifier::external_body]
+        pub closed spec fn spec_phantom_data<V: ?Sized>() -> core::marker::PhantomData<V> {
+            core::marker::PhantomData::default()
+        }
+
         pub struct VecIterCopy<'a, T: 'a> {
             pub vec: &'a Vec<T>,
             pub cur: usize,
@@ -1124,7 +1130,7 @@ test_verify_one_file! {
         pub struct VecGhostIterCopy<'a, T> {
             pub seq: Seq<T>,
             pub cur: int,
-            pub unused_phantom_dummy: &'a Vec<T>, // use the 'a parameter, which I don't even want; ugh
+            pub phantom: core::marker::PhantomData<&'a T>,
         }
 
         impl<'a, T: 'a> vstd::pervasive::ForLoopGhostIteratorNew for VecIterCopy<'a, T> {
@@ -1134,7 +1140,7 @@ test_verify_one_file! {
                 VecGhostIterCopy {
                     seq: self.vec@,
                     cur: 0,
-                    unused_phantom_dummy: &self.vec,
+                    phantom: spec_phantom_data(),
                 }
             }
         }
diff --git a/source/vir/src/ast.rs b/source/vir/src/ast.rs
index 070b3f5e87..98cbf10b06 100644
--- a/source/vir/src/ast.rs
+++ b/source/vir/src/ast.rs
@@ -256,7 +256,7 @@ pub enum UnaryOp {
     CharToInt,
     /// Given an exec/proof expression used to construct a loop iterator,
     /// try to infer a pure specification for the loop iterator.
-    /// Return Some(spec) if successful, None otherwise.
+    /// Evaluate to Some(spec) if successful, None otherwise.
     /// (Note: this is just used as a hint for loop invariants;
     /// regardless of whether it is Some(spec) or None, it should not affect soundness.)
     /// For an exec/proof expression e, the spec s should be chosen so that the value v
diff --git a/source/vir/src/ast_simplify.rs b/source/vir/src/ast_simplify.rs
index 6ad120816d..5db8c8bbd2 100644
--- a/source/vir/src/ast_simplify.rs
+++ b/source/vir/src/ast_simplify.rs
@@ -358,16 +358,6 @@ fn simplify_one_expr(
             }
         }
         ExprX::Unary(UnaryOp::CoerceMode { .. }, expr0) => Ok(expr0.clone()),
-        ExprX::Unary(UnaryOp::InferSpecForLoopIter, _) => {
-            let mode_opt = ctx.infer_spec_for_loop_iter_modes.get(&expr.span.id);
-            if mode_opt == Some(&Mode::Spec) {
-                // InferSpecForLoopIter must be spec mode to be usable for invariant inference
-                Ok(expr.clone())
-            } else {
-                // Otherwise, abandon the expression and return None (no inference)
-                Ok(crate::loop_inference::make_none_expr(&expr.span, &expr.typ))
-            }
-        }
         ExprX::UnaryOpr(UnaryOpr::TupleField { tuple_arity, field }, expr0) => {
             Ok(tuple_get_field_expr(state, &expr.span, &expr.typ, expr0, *tuple_arity, *field))
         }
@@ -1024,7 +1014,6 @@ pub fn simplify_krate(ctx: &mut GlobalCtx, krate: &Krate) -> Result<Krate, VirEr
         ctx.rlimit,
         ctx.interpreter_log.clone(),
         ctx.vstd_crate_name.clone(),
-        &vec![],
     )?;
     Ok(krate)
 }
diff --git a/source/vir/src/context.rs b/source/vir/src/context.rs
index eeecdeecc4..42ce24d293 100644
--- a/source/vir/src/context.rs
+++ b/source/vir/src/context.rs
@@ -4,7 +4,7 @@ use crate::ast::{
 };
 use crate::datatype_to_air::is_datatype_transparent;
 use crate::def::FUEL_ID;
-use crate::messages::{error, AstId, Span};
+use crate::messages::{error, Span};
 use crate::poly::MonoTyp;
 use crate::recursion::Node;
 use crate::scc::Graph;
@@ -47,7 +47,6 @@ pub struct GlobalCtx {
     pub(crate) interpreter_log: Arc<std::sync::Mutex<Option<File>>>,
     pub(crate) vstd_crate_name: Option<Ident>, // already an arc
     pub arch: crate::ast::ArchWordBits,
-    pub(crate) infer_spec_for_loop_iter_modes: Arc<HashMap<AstId, Mode>>,
 }
 
 // Context for verifying one function
@@ -195,7 +194,6 @@ impl GlobalCtx {
         rlimit: f32,
         interpreter_log: Arc<std::sync::Mutex<Option<File>>>,
         vstd_crate_name: Option<Ident>,
-        infer_spec_for_loop_iter_modes_vec: &Vec<(Span, Mode)>,
     ) -> Result<Self, VirErr> {
         let chosen_triggers: std::cell::RefCell<Vec<ChosenTriggers>> =
             std::cell::RefCell::new(Vec::new());
@@ -293,13 +291,6 @@ impl GlobalCtx {
 
         let datatype_graph = crate::recursive_types::build_datatype_graph(krate);
 
-        let mut infer_spec_for_loop_iter_modes = HashMap::new();
-        for (span, mode) in infer_spec_for_loop_iter_modes_vec {
-            if *mode != Mode::Spec || !infer_spec_for_loop_iter_modes.contains_key(&span.id) {
-                infer_spec_for_loop_iter_modes.insert(span.id, *mode);
-            }
-        }
-
         Ok(GlobalCtx {
             chosen_triggers,
             datatypes: Arc::new(datatypes),
@@ -313,7 +304,6 @@ impl GlobalCtx {
             interpreter_log,
             vstd_crate_name,
             arch: krate.arch.word_bits,
-            infer_spec_for_loop_iter_modes: Arc::new(infer_spec_for_loop_iter_modes),
         })
     }
 
@@ -335,7 +325,6 @@ impl GlobalCtx {
             interpreter_log,
             vstd_crate_name: self.vstd_crate_name.clone(),
             arch: self.arch,
-            infer_spec_for_loop_iter_modes: self.infer_spec_for_loop_iter_modes.clone(),
         }
     }
 
diff --git a/source/vir/src/def.rs b/source/vir/src/def.rs
index 3bace970e2..1d24bb2acb 100644
--- a/source/vir/src/def.rs
+++ b/source/vir/src/def.rs
@@ -560,6 +560,10 @@ impl<X> Spanned<X> {
         Arc::new(Spanned { span: span, x: x })
     }
 
+    pub fn new_x(&self, x: X) -> Arc<Spanned<X>> {
+        Arc::new(Spanned { span: self.span.clone(), x })
+    }
+
     pub fn map_x<Y>(&self, f: impl FnOnce(&X) -> Y) -> Arc<Spanned<Y>> {
         Arc::new(Spanned { span: self.span.clone(), x: f(&self.x) })
     }
diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index 65f19ea2a3..2664bd06e9 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -12,6 +12,7 @@ use air::scope_map::ScopeMap;
 use std::cmp::min;
 use std::collections::{HashMap, HashSet};
 use std::mem::swap;
+use std::sync::Arc;
 
 // Exec <= Proof <= Spec
 pub fn mode_le(m1: Mode, m2: Mode) -> bool {
@@ -49,8 +50,6 @@ pub struct ErasureModes {
     pub condition_modes: Vec<(Span, Mode)>,
     // Modes of variables in Var, Assign, Decl
     pub var_modes: Vec<(Span, Mode)>,
-    // Modes of InferSpecForLoopIter
-    pub infer_spec_for_loop_iter_modes: Vec<(Span, Mode)>,
 }
 
 impl Ghost {
@@ -88,6 +87,8 @@ struct Typing {
     pub(crate) fun_mode: Mode,
     pub(crate) ret_mode: Option<Mode>,
     pub(crate) atomic_insts: Option<AtomicInstCollector>,
+    // Modes of InferSpecForLoopIter
+    infer_spec_for_loop_iter_modes: Option<Vec<(Span, Mode)>>,
 }
 
 impl Typing {
@@ -633,7 +634,14 @@ fn check_expr_handle_mut_arg(
             // so that ast_simplify can replace the expression with None.
             let mode_opt = check_expr(typing, outer_mode, e1);
             let mode = mode_opt.unwrap_or(Mode::Exec);
-            typing.erasure_modes.infer_spec_for_loop_iter_modes.push((expr.span.clone(), mode));
+            if let Some(infer_spec) = typing.infer_spec_for_loop_iter_modes.as_mut() {
+                infer_spec.push((expr.span.clone(), mode));
+            } else {
+                return Err(error(
+                    &expr.span,
+                    "infer_spec_for_loop_iter is only allowed in function body",
+                ));
+            }
             Ok(Mode::Spec)
         }
         ExprX::Unary(_, e1) => check_expr(typing, outer_mode, e1),
@@ -1160,7 +1168,7 @@ fn check_stmt(typing: &mut Typing, outer_mode: Mode, stmt: &Stmt) -> Result<(),
     }
 }
 
-fn check_function(typing: &mut Typing, function: &Function) -> Result<(), VirErr> {
+fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), VirErr> {
     typing.vars.push_scope(true);
 
     if let FunctionKind::TraitMethodImpl { method, trait_path, .. } = &function.x.kind {
@@ -1257,7 +1265,33 @@ fn check_function(typing: &mut Typing, function: &Function) -> Result<(), VirErr
     }
     if let Some(body) = &function.x.body {
         typing.block_ghostness = Ghost::of_mode(function.x.mode);
+        assert!(typing.infer_spec_for_loop_iter_modes.is_none());
+        typing.infer_spec_for_loop_iter_modes = Some(Vec::new());
         check_expr_has_mode(typing, function.x.mode, body, function.x.ret.x.mode)?;
+
+        // Replace any InferSpecForLoopIter with Some(...expr...) or None
+        let infer_spec = typing.infer_spec_for_loop_iter_modes.as_ref().expect("infer_spec");
+        if infer_spec.len() > 0 {
+            let mut functionx = function.x.clone();
+            functionx.body = Some(crate::ast_visitor::map_expr_visitor(body, &|expr: &Expr| {
+                match &expr.x {
+                    ExprX::Unary(UnaryOp::InferSpecForLoopIter, _) => {
+                        let mode_opt = infer_spec.iter().find(|(span, _)| span.id == expr.span.id);
+                        if let Some((_, Mode::Spec)) = mode_opt {
+                            // InferSpecForLoopIter must be spec mode
+                            // to be usable for invariant inference
+                            Ok(expr.clone())
+                        } else {
+                            // Otherwise, abandon the expression and return None (no inference)
+                            Ok(crate::loop_inference::make_none_expr(&expr.span, &expr.typ))
+                        }
+                    }
+                    _ => Ok(expr.clone()),
+                }
+            })?);
+            *function = function.new_x(functionx);
+        }
+        typing.infer_spec_for_loop_iter_modes = None;
     }
     typing.ret_mode = None;
     typing.vars.pop_scope();
@@ -1268,7 +1302,7 @@ fn check_function(typing: &mut Typing, function: &Function) -> Result<(), VirErr
 pub fn check_crate(
     krate: &Krate,
     macro_erasure: bool, // TODO(main_new) remove, always true
-) -> Result<ErasureModes, VirErr> {
+) -> Result<(Krate, ErasureModes), VirErr> {
     let mut funs: HashMap<Fun, Function> = HashMap::new();
     let mut datatypes: HashMap<Path, Datatype> = HashMap::new();
     for function in krate.functions.iter() {
@@ -1277,11 +1311,7 @@ pub fn check_crate(
     for datatype in krate.datatypes.iter() {
         datatypes.insert(datatype.x.path.clone(), datatype.clone());
     }
-    let erasure_modes = ErasureModes {
-        condition_modes: vec![],
-        var_modes: vec![],
-        infer_spec_for_loop_iter_modes: vec![],
-    };
+    let erasure_modes = ErasureModes { condition_modes: vec![], var_modes: vec![] };
     let mut typing = Typing {
         funs,
         datatypes,
@@ -1295,8 +1325,10 @@ pub fn check_crate(
         fun_mode: Mode::Exec,
         ret_mode: None,
         atomic_insts: None,
+        infer_spec_for_loop_iter_modes: None,
     };
-    for function in krate.functions.iter() {
+    let mut kratex = (**krate).clone();
+    for function in kratex.functions.iter_mut() {
         typing.check_ghost_blocks = function.x.attrs.uses_ghost_blocks;
         typing.fun_mode = function.x.mode;
         if function.x.attrs.atomic {
@@ -1311,5 +1343,5 @@ pub fn check_crate(
             check_function(&mut typing, function)?;
         }
     }
-    Ok(typing.erasure_modes)
+    Ok((Arc::new(kratex), typing.erasure_modes))
 }

From 8e0b755454aa9ba0769b5dace2f8f3d6c7e0b0ba Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Tue, 9 Jan 2024 15:21:21 -0800
Subject: [PATCH 11/19] Push/pop modes.rs state around InferSpecForLoopIter

---
 source/rust_verify/src/verifier.rs     |  2 +-
 source/rust_verify_test/tests/loops.rs | 16 +++++++
 source/vir/src/modes.rs                | 66 ++++++++++++++++++++++----
 3 files changed, 74 insertions(+), 10 deletions(-)

diff --git a/source/rust_verify/src/verifier.rs b/source/rust_verify/src/verifier.rs
index 1a75feaaf2..feabfd8af7 100644
--- a/source/rust_verify/src/verifier.rs
+++ b/source/rust_verify/src/verifier.rs
@@ -2370,7 +2370,7 @@ impl Verifier {
         check_crate_result.map_err(|e| (e, Vec::new()))?;
         let vir_crate = vir::autospec::resolve_autospec(&vir_crate).map_err(|e| (e, Vec::new()))?;
         let (vir_crate, erasure_modes) =
-            vir::modes::check_crate(&vir_crate, true).map_err(|e| (e, Vec::new()))?;
+            vir::modes::check_crate(&vir_crate).map_err(|e| (e, Vec::new()))?;
 
         self.vir_crate = Some(vir_crate.clone());
         self.crate_names = Some(crate_names);
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index 8806b93775..b002a0c18e 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -1056,6 +1056,22 @@ test_verify_one_file! {
                 end = end + 0; // causes end to be non-constant, so loop needs more invariants
             }
         }
+
+        fn non_spec() {}
+
+        fn test_loop_modes_transient_state() {
+            let mut n: u64 = 0;
+            let mut end = 10;
+            // test Typing::snapshot_transient_state
+            for x in iter: 0..({let z = end; non_spec(); z})
+                invariant
+                    n == iter.cur * 3,
+                    end == 10,
+            {
+                n += 3;
+                end = end + 0; // causes end to be non-constant
+            }
+        }
     } => Err(e) => assert_one_fails(e)
 }
 
diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index 2664bd06e9..ca1b2bade9 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -77,7 +77,6 @@ struct Typing {
     pub(crate) erasure_modes: ErasureModes,
     pub(crate) in_forall_stmt: bool,
     pub(crate) check_ghost_blocks: bool,
-    pub(crate) macro_erasure: bool,
     // Are we in a syntactic ghost block?
     // If not, Ghost::Exec (corresponds to exec mode).
     // If yes (corresponding to proof/spec mode), say whether variables are tracked or not.
@@ -91,6 +90,16 @@ struct Typing {
     infer_spec_for_loop_iter_modes: Option<Vec<(Span, Mode)>>,
 }
 
+// If we need to catch an error and continue,
+// keep a snapshot of the old transient Typing state fields so we can restore the fields.
+// (used by InferSpecForLoopIter)
+struct TypingSnapshot {
+    vars_scope_count: usize,
+    in_forall_stmt: bool,
+    block_ghostness: Ghost,
+    atomic_insts: Option<AtomicInstCollector>,
+}
+
 impl Typing {
     fn get(&self, x: &Ident) -> (bool, Mode) {
         *self.vars.get(x).expect("internal error: missing mode")
@@ -99,6 +108,43 @@ impl Typing {
     fn insert(&mut self, _span: &Span, x: &Ident, mutable: bool, mode: Mode) {
         self.vars.insert(x.clone(), (mutable, mode)).expect("internal error: Typing insert");
     }
+
+    fn snapshot_transient_state(&mut self) -> TypingSnapshot {
+        // mention every field of Typing to make sure TypingSnapshot stays up to date with Typing:
+        let Typing {
+            funs: _,
+            datatypes: _,
+            traits: _,
+            vars,
+            erasure_modes: _,
+            in_forall_stmt,
+            check_ghost_blocks: _,
+            block_ghostness,
+            fun_mode: _,
+            ret_mode: _,
+            atomic_insts,
+            infer_spec_for_loop_iter_modes: _,
+        } = &self;
+        let snapshot = TypingSnapshot {
+            vars_scope_count: vars.num_scopes(),
+            in_forall_stmt: *in_forall_stmt,
+            block_ghostness: *block_ghostness,
+            atomic_insts: atomic_insts.clone(),
+        };
+        self.vars.push_scope(true);
+        snapshot
+    }
+
+    fn pop_transient_state(&mut self, snapshot: TypingSnapshot) {
+        let TypingSnapshot { vars_scope_count, in_forall_stmt, block_ghostness, atomic_insts } =
+            snapshot;
+        while self.vars.num_scopes() != vars_scope_count {
+            self.vars.pop_scope();
+        }
+        self.in_forall_stmt = in_forall_stmt;
+        self.block_ghostness = block_ghostness;
+        self.atomic_insts = atomic_insts;
+    }
 }
 
 // One tricky thing in mode-checking is that an invariant block needs to have
@@ -125,7 +171,7 @@ impl Typing {
 // In principle, we could do something fancy and allow 1 atomic instruction in each branch,
 // but for now we just error if there is more than 1 atomic call in the AST.
 
-#[derive(Default)]
+#[derive(Default, Clone)]
 struct AtomicInstCollector {
     atomics: Vec<Span>,
     non_atomics: Vec<Span>,
@@ -632,7 +678,13 @@ fn check_expr_handle_mut_arg(
             // are all autospec), then keep the expression.
             // Otherwise, make a note that the expression had mode exec,
             // so that ast_simplify can replace the expression with None.
+
+            // Since we may catch a Result::Err,
+            // explicitly push and pop transient state around check_expr
+            let snapshot = typing.snapshot_transient_state();
             let mode_opt = check_expr(typing, outer_mode, e1);
+            typing.pop_transient_state(snapshot);
+
             let mode = mode_opt.unwrap_or(Mode::Exec);
             if let Some(infer_spec) = typing.infer_spec_for_loop_iter_modes.as_mut() {
                 infer_spec.push((expr.span.clone(), mode));
@@ -828,14 +880,14 @@ fn check_expr_handle_mut_arg(
             Ok(x_mode)
         }
         ExprX::Fuel(_, _) => {
-            if typing.macro_erasure && typing.block_ghostness == Ghost::Exec {
+            if typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use reveal/hide in exec mode")
                     .help("wrap the reveal call in a `proof` block"));
             }
             Ok(outer_mode)
         }
         ExprX::RevealString(_) => {
-            if typing.macro_erasure && typing.block_ghostness == Ghost::Exec {
+            if typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use reveal_strlit in exec mode")
                     .help("wrap the reveal_strlit call in a `proof` block"));
             }
@@ -1299,10 +1351,7 @@ fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), Vi
     Ok(())
 }
 
-pub fn check_crate(
-    krate: &Krate,
-    macro_erasure: bool, // TODO(main_new) remove, always true
-) -> Result<(Krate, ErasureModes), VirErr> {
+pub fn check_crate(krate: &Krate) -> Result<(Krate, ErasureModes), VirErr> {
     let mut funs: HashMap<Fun, Function> = HashMap::new();
     let mut datatypes: HashMap<Path, Datatype> = HashMap::new();
     for function in krate.functions.iter() {
@@ -1320,7 +1369,6 @@ pub fn check_crate(
         erasure_modes,
         in_forall_stmt: false,
         check_ghost_blocks: false,
-        macro_erasure,
         block_ghostness: Ghost::Exec,
         fun_mode: Mode::Exec,
         ret_mode: None,

From 635e34ce0d52536c70d1d7e7dd0cee9ab4548e65 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Tue, 9 Jan 2024 16:57:44 -0800
Subject: [PATCH 12/19] Limit axiom_spec_range_next to specific types

---
 source/pervasive/std_specs/range.rs | 61 +++++++++++++++--------------
 1 file changed, 31 insertions(+), 30 deletions(-)

diff --git a/source/pervasive/std_specs/range.rs b/source/pervasive/std_specs/range.rs
index 6548e7e68a..ac3e2d5e95 100644
--- a/source/pervasive/std_specs/range.rs
+++ b/source/pervasive/std_specs/range.rs
@@ -20,22 +20,6 @@ pub trait StepSpec where Self: Sized {
 
 pub spec fn spec_range_next<A>(a: Range<A>) -> (Range<A>, Option<A>);
 
-#[verifier::external_body]
-#[verifier::broadcast_forall]
-pub proof fn axiom_spec_range_next<A: StepSpec>(range: Range<A>)
-    ensures
-        range.start.spec_is_lt(range.end) ==>
-            // TODO (not important): use new "matches ==>" syntax here
-            (if let Some(n) = range.start.spec_forward_checked(1) {
-                spec_range_next(range) == (Range { start: n, ..range }, Some(range.start))
-            } else {
-                true
-            }),
-        !range.start.spec_is_lt(range.end) ==>
-            #[trigger] spec_range_next(range) == (range, None::<A>),
-{
-}
-
 #[verifier::external_fn_specification]
 pub fn ex_range_next<A: core::iter::Step>(range: &mut Range<A>) -> (r: Option<A>)
     ensures (*range, r) == spec_range_next(*old(range))
@@ -115,7 +99,7 @@ impl<A: StepSpec + core::iter::Step> crate::view::View for RangeGhostIterator<A>
 } // verus!
 
 macro_rules! step_specs {
-    ($t: ty) => {
+    ($t: ty, $axiom: ident) => {
         verus! {
         impl StepSpec for $t {
             open spec fn spec_is_lt(self, other: Self) -> bool {
@@ -152,20 +136,37 @@ macro_rules! step_specs {
                     None
                 }
             }
-        } // verus!
         }
+        // TODO: we might be able to make this generic over A: StepSpec
+        // once we settle on a way to connect std traits like Step with spec traits like StepSpec.
+        #[verifier::external_body]
+        #[verifier::broadcast_forall]
+        pub proof fn $axiom(range: Range<$t>)
+            ensures
+                range.start.spec_is_lt(range.end) ==>
+                    // TODO (not important): use new "matches ==>" syntax here
+                    (if let Some(n) = range.start.spec_forward_checked(1) {
+                        spec_range_next(range) == (Range { start: n, ..range }, Some(range.start))
+                    } else {
+                        true
+                    }),
+                !range.start.spec_is_lt(range.end) ==>
+                    #[trigger] spec_range_next(range) == (range, None::<$t>),
+        {
+        }
+        } // verus!
     }
 }
 
-step_specs!(u8);
-step_specs!(u16);
-step_specs!(u32);
-step_specs!(u64);
-step_specs!(u128);
-step_specs!(usize);
-step_specs!(i8);
-step_specs!(i16);
-step_specs!(i32);
-step_specs!(i64);
-step_specs!(i128);
-step_specs!(isize);
+step_specs!(u8, axiom_spec_range_next_u8);
+step_specs!(u16, axiom_spec_range_next_u16);
+step_specs!(u32, axiom_spec_range_next_u32);
+step_specs!(u64, axiom_spec_range_next_u64);
+step_specs!(u128, axiom_spec_range_next_u128);
+step_specs!(usize, axiom_spec_range_next_usize);
+step_specs!(i8, axiom_spec_range_next_i8);
+step_specs!(i16, axiom_spec_range_next_i16);
+step_specs!(i32, axiom_spec_range_next_i32);
+step_specs!(i64, axiom_spec_range_next_i64);
+step_specs!(i128, axiom_spec_range_next_i128);
+step_specs!(isize, axiom_spec_range_next_isize);

From f6ab69e6ad3ff836a7c6c4bbf999c82cb34d8274 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Tue, 9 Jan 2024 17:32:04 -0800
Subject: [PATCH 13/19] Replace ghost_condition with ghost_ensures

---
 source/builtin_macros/src/syntax.rs    | 4 ++--
 source/pervasive/pervasive.rs          | 4 ++--
 source/pervasive/std_specs/range.rs    | 4 ++--
 source/rust_verify_test/tests/loops.rs | 4 ++--
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index 669199b88f..403ac149fa 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -1318,7 +1318,7 @@ impl Visitor {
         //                          &core::iter::IntoIterator::into_iter(e)))),
         //              { let x = vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&y); inv },
         //          ensures
-        //              !vstd::pervasive::ForLoopGhostIterator::ghost_condition(&y),
+        //              vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&y),
         //      {
         //          if let Some(x) = core::iter::Iterator::next(&mut y) { body } else { break }
         //          proof { y = vstd::pervasive::ForLoopGhostIterator::ghost_advance(&y, &VERUS_exec_iter); }
@@ -1440,7 +1440,7 @@ impl Visitor {
         // but at the moment, user-supplied ensures aren't supported, so this would be hard to use.
         let ensure = if no_loop_invariant.is_none() {
             Some(parse_quote_spanned!(span =>
-                ensures !vstd::pervasive::ForLoopGhostIterator::ghost_condition(&#x_ghost_iter),
+                ensures vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&#x_ghost_iter),
             ))
         } else {
             None
diff --git a/source/pervasive/pervasive.rs b/source/pervasive/pervasive.rs
index aa54647055..11f0e776ff 100644
--- a/source/pervasive/pervasive.rs
+++ b/source/pervasive/pervasive.rs
@@ -49,8 +49,8 @@ pub trait ForLoopGhostIterator {
     // (When the analysis can infer a spec initial value, the analysis places the value in init)
     spec fn ghost_invariant(&self, init: Option<&Self>) -> bool;
 
-    // Is the loop condition satisfied?
-    spec fn ghost_condition(&self) -> bool;
+    // True upon loop exit
+    spec fn ghost_ensures(&self) -> bool;
 
     // Value used by default for decreases clause when no explicit decreases clause is provided
     // (the user can override this with an explicit decreases clause).
diff --git a/source/pervasive/std_specs/range.rs b/source/pervasive/std_specs/range.rs
index ac3e2d5e95..f46314499d 100644
--- a/source/pervasive/std_specs/range.rs
+++ b/source/pervasive/std_specs/range.rs
@@ -68,8 +68,8 @@ impl<A: StepSpec + core::iter::Step> crate::pervasive::ForLoopGhostIterator for
             }
     }
 
-    open spec fn ghost_condition(&self) -> bool {
-        self.cur.spec_is_lt(self.end)
+    open spec fn ghost_ensures(&self) -> bool {
+        !self.cur.spec_is_lt(self.end)
     }
 
     open spec fn ghost_decrease(&self) -> int {
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index b002a0c18e..c44a8029c5 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -1180,8 +1180,8 @@ test_verify_one_file! {
                     }
             }
 
-            open spec fn ghost_condition(&self) -> bool {
-                self.cur < self.seq.len()
+            open spec fn ghost_ensures(&self) -> bool {
+                self.cur >= self.seq.len()
             }
 
             open spec fn ghost_decrease(&self) -> int {

From 569b87058270d1ac4267135b028b4205c4f5c765 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Thu, 11 Jan 2024 08:16:45 -0800
Subject: [PATCH 14/19] Add ret_mode to TypingSnapshot

---
 source/vir/src/modes.rs | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index ca1b2bade9..1d8c3eee5f 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -97,6 +97,7 @@ struct TypingSnapshot {
     vars_scope_count: usize,
     in_forall_stmt: bool,
     block_ghostness: Ghost,
+    ret_mode: Option<Mode>,
     atomic_insts: Option<AtomicInstCollector>,
 }
 
@@ -121,7 +122,7 @@ impl Typing {
             check_ghost_blocks: _,
             block_ghostness,
             fun_mode: _,
-            ret_mode: _,
+            ret_mode,
             atomic_insts,
             infer_spec_for_loop_iter_modes: _,
         } = &self;
@@ -129,6 +130,7 @@ impl Typing {
             vars_scope_count: vars.num_scopes(),
             in_forall_stmt: *in_forall_stmt,
             block_ghostness: *block_ghostness,
+            ret_mode: *ret_mode,
             atomic_insts: atomic_insts.clone(),
         };
         self.vars.push_scope(true);
@@ -136,13 +138,19 @@ impl Typing {
     }
 
     fn pop_transient_state(&mut self, snapshot: TypingSnapshot) {
-        let TypingSnapshot { vars_scope_count, in_forall_stmt, block_ghostness, atomic_insts } =
-            snapshot;
+        let TypingSnapshot {
+            vars_scope_count,
+            in_forall_stmt,
+            block_ghostness,
+            ret_mode,
+            atomic_insts,
+        } = snapshot;
         while self.vars.num_scopes() != vars_scope_count {
             self.vars.pop_scope();
         }
         self.in_forall_stmt = in_forall_stmt;
         self.block_ghostness = block_ghostness;
+        self.ret_mode = ret_mode;
         self.atomic_insts = atomic_insts;
     }
 }

From 5b2a059562240686d8273bb8e21bd935aaff5d47 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Thu, 11 Jan 2024 08:28:29 -0800
Subject: [PATCH 15/19] Fix comment and panic message

---
 source/vir/src/bitvector_to_air.rs | 4 +---
 source/vir/src/modes.rs            | 3 ++-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/source/vir/src/bitvector_to_air.rs b/source/vir/src/bitvector_to_air.rs
index e6fb474bbd..49ae040984 100644
--- a/source/vir/src/bitvector_to_air.rs
+++ b/source/vir/src/bitvector_to_air.rs
@@ -131,9 +131,7 @@ pub(crate) fn bv_exp_to_expr(ctx: &Ctx, exp: &Exp, expr_ctxt: &BvExprCtxt) -> Re
                     "internal error: matching for bit vector ops on this match should be impossible"
                 ),
                 UnaryOp::InferSpecForLoopIter { .. } => {
-                    panic!(
-                        "internal error: InferSpecForLoopIter should have been removed before here"
-                    )
+                    panic!("internal error: unexpected Option type (from InferSpecForLoopIter)")
                 }
             }
         }
diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index 1d8c3eee5f..682a31075e 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -1329,7 +1329,8 @@ fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), Vi
         typing.infer_spec_for_loop_iter_modes = Some(Vec::new());
         check_expr_has_mode(typing, function.x.mode, body, function.x.ret.x.mode)?;
 
-        // Replace any InferSpecForLoopIter with Some(...expr...) or None
+        // Replace InferSpecForLoopIter None if it fails to have mode spec
+        // (if it's mode spec, leave as is to be processed by sst_to_air and loop_inference)
         let infer_spec = typing.infer_spec_for_loop_iter_modes.as_ref().expect("infer_spec");
         if infer_spec.len() > 0 {
             let mut functionx = function.x.clone();

From b602bfe805069da47ebcee9417afd90181ae554d Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Wed, 17 Jan 2024 11:42:32 -0800
Subject: [PATCH 16/19] Return Option from ghost_peek_next and ghost_decrease

---
 source/builtin_macros/src/syntax.rs    |  9 +++++++--
 source/pervasive/pervasive.rs          | 16 +++++++++++-----
 source/pervasive/std_specs/range.rs    |  8 ++++----
 source/rust_verify_test/tests/loops.rs |  8 ++++----
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index e280a97927..34634d1427 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -1327,7 +1327,10 @@ impl Visitor {
         //                  builtin::infer_spec_for_loop_iter(
         //                      &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
         //                          &core::iter::IntoIterator::into_iter(e)))),
-        //              { let x = vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&y); inv },
+        //              { let x =
+        //                  vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&y)
+        //                  .unwrap_or(vstd::pervasive::arbitrary());
+        //                inv },
         //          ensures
         //              vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&y),
         //      {
@@ -1426,7 +1429,9 @@ impl Visitor {
         let invariant_ensure = if let Some(mut invariant) = invariant {
             for inv in &mut invariant.exprs.exprs {
                 *inv = Expr::Verbatim(quote_spanned!(span => {
-                    let #pat = vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&#x_ghost_iter);
+                    let #pat =
+                        vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&#x_ghost_iter)
+                        .unwrap_or(vstd::pervasive::arbitrary());
                     #inv
                 }));
             }
diff --git a/source/pervasive/pervasive.rs b/source/pervasive/pervasive.rs
index 85a5750c3b..81c95815b2 100644
--- a/source/pervasive/pervasive.rs
+++ b/source/pervasive/pervasive.rs
@@ -32,6 +32,7 @@ pub proof fn affirm(b: bool)
 {
 }
 
+// TODO: when default trait methods are supported, most of these should be given defaults
 pub trait ForLoopGhostIterator {
     type ExecIter;
     type Item;
@@ -53,14 +54,19 @@ pub trait ForLoopGhostIterator {
 
     // Value used by default for decreases clause when no explicit decreases clause is provided
     // (the user can override this with an explicit decreases clause).
-    // (If there's no appropriate decrease, this can return an arbitrary value like 0,
+    // (If there's no appropriate decrease, this can return None,
     // and the user will have to provide an explicit decreases clause.)
-    spec fn ghost_decrease(&self) -> Self::Decrease;
+    spec fn ghost_decrease(&self) -> Option<Self::Decrease>;
 
     // If there will be Some next value, and we can make a useful guess as to what the next value
-    // will be, return it.
-    // Otherwise, return an arbitrary value.
-    spec fn ghost_peek_next(&self) -> Self::Item;
+    // will be, return Some of it.
+    // Otherwise, return None.
+    // TODO: in the long term, we could have VIR insert an assertion (or warning)
+    // that ghost_peek_next returns non-null if it is used in the invariants.
+    // (this will take a little bit of engineering since the syntax macro blindly inserts
+    // let bindings using ghost_peek_next, even if they aren't needed, and we only learn
+    // what is actually needed later in VIR.)
+    spec fn ghost_peek_next(&self) -> Option<Self::Item>;
 
     // At the end of the for loop, advance to the next position.
     // Future TODO: this may be better as a proof function
diff --git a/source/pervasive/std_specs/range.rs b/source/pervasive/std_specs/range.rs
index f46314499d..3911e86489 100644
--- a/source/pervasive/std_specs/range.rs
+++ b/source/pervasive/std_specs/range.rs
@@ -72,12 +72,12 @@ impl<A: StepSpec + core::iter::Step> crate::pervasive::ForLoopGhostIterator for
         !self.cur.spec_is_lt(self.end)
     }
 
-    open spec fn ghost_decrease(&self) -> int {
-        self.cur.spec_steps_between_int(self.end)
+    open spec fn ghost_decrease(&self) -> Option<int> {
+        Some(self.cur.spec_steps_between_int(self.end))
     }
 
-    open spec fn ghost_peek_next(&self) -> A {
-        self.cur
+    open spec fn ghost_peek_next(&self) -> Option<A> {
+        Some(self.cur)
     }
 
     open spec fn ghost_advance(&self, _exec_iter: &Range<A>) -> RangeGhostIterator<A> {
diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index c44a8029c5..63fa3f0d19 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -1184,12 +1184,12 @@ test_verify_one_file! {
                 self.cur >= self.seq.len()
             }
 
-            open spec fn ghost_decrease(&self) -> int {
-                self.seq.len() - self.cur
+            open spec fn ghost_decrease(&self) -> Option<int> {
+                Some(self.seq.len() - self.cur)
             }
 
-            open spec fn ghost_peek_next(&self) -> T {
-                self.seq[self.cur]
+            open spec fn ghost_peek_next(&self) -> Option<T> {
+                Some(self.seq[self.cur])
             }
 
             open spec fn ghost_advance(&self, _exec_iter: &VecIterCopy<T>) -> VecGhostIterCopy<'a, T> {

From 7393ae1d70de0d5458753f67224c03825528a0b2 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Wed, 17 Jan 2024 11:52:01 -0800
Subject: [PATCH 17/19] Update loops.rs test with more precise ghost_peek_next

---
 source/rust_verify_test/tests/loops.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/source/rust_verify_test/tests/loops.rs b/source/rust_verify_test/tests/loops.rs
index 63fa3f0d19..9f0d832bb1 100644
--- a/source/rust_verify_test/tests/loops.rs
+++ b/source/rust_verify_test/tests/loops.rs
@@ -1189,7 +1189,11 @@ test_verify_one_file! {
             }
 
             open spec fn ghost_peek_next(&self) -> Option<T> {
-                Some(self.seq[self.cur])
+                if 0 <= self.cur < self.seq.len() {
+                    Some(self.seq[self.cur])
+                } else {
+                    None
+                }
             }
 
             open spec fn ghost_advance(&self, _exec_iter: &VecIterCopy<T>) -> VecGhostIterCopy<'a, T> {

From 1fe732a94e208af09b92103f4f8417cd35ab8d68 Mon Sep 17 00:00:00 2001
From: Chris Hawblitzel <Chris.Hawblitzel@microsoft.com>
Date: Fri, 19 Jan 2024 17:16:56 -0800
Subject: [PATCH 18/19] Match for loop translation from
 https://doc.rust-lang.org/reference/expressions/loop-expr.html

---
 source/builtin_macros/src/syntax.rs | 121 +++++++++++++++++-----------
 1 file changed, 75 insertions(+), 46 deletions(-)

diff --git a/source/builtin_macros/src/syntax.rs b/source/builtin_macros/src/syntax.rs
index 34634d1427..17b5b59fc6 100644
--- a/source/builtin_macros/src/syntax.rs
+++ b/source/builtin_macros/src/syntax.rs
@@ -1312,31 +1312,50 @@ impl Visitor {
         // which we need for writing invariants.
         // Therefore, rather than letting Rust desugar a for-loop into a loop with a break,
         // we desugar the for-loop into a loop with a break here.
+        // (See https://doc.rust-lang.org/reference/expressions/loop-expr.html for the
+        // official definition of the desugaring that we follow.)
         // Specifically, we desugar:
-        //  for x in y: e invariant inv { body }
+        //  'label: for x in y: e invariant inv { body }
         // into:
         //  {
         //      #[allow(non_snake_case)]
-        //      let mut VERUS_exec_iter = core::iter::IntoIterator::into_iter(e);
-        //      #[allow(non_snake_case)]
-        //      let ghost mut y = vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(&VERUS_exec_iter);
-        //      loop
-        //          invariant_ensures
-        //              vstd::pervasive::ForLoopGhostIterator::exec_invariant(&y, &VERUS_exec_iter),
-        //              vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&y,
-        //                  builtin::infer_spec_for_loop_iter(
-        //                      &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
-        //                          &core::iter::IntoIterator::into_iter(e)))),
-        //              { let x =
-        //                  vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&y)
-        //                  .unwrap_or(vstd::pervasive::arbitrary());
-        //                inv },
-        //          ensures
-        //              vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&y),
-        //      {
-        //          if let Some(x) = core::iter::Iterator::next(&mut y) { body } else { break }
-        //          proof { y = vstd::pervasive::ForLoopGhostIterator::ghost_advance(&y, &VERUS_exec_iter); }
-        //      }
+        //      let VERUS_loop_result = match ::core::iter::IntoIterator::into_iter(e) {
+        //          #[allow(non_snake_case)]
+        //          mut VERUS_exec_iter => {
+        //              #[allow(non_snake_case)]
+        //              let ghost mut y = ::vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+        //                  &VERUS_exec_iter);
+        //              'label: loop
+        //                  invariant_ensures
+        //                      ::vstd::pervasive::ForLoopGhostIterator::exec_invariant(&y,
+        //                          &VERUS_exec_iter),
+        //                      ::vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&y,
+        //                          builtin::infer_spec_for_loop_iter(
+        //                              &::vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+        //                                  &::core::iter::IntoIterator::into_iter(e)))),
+        //                      { let x =
+        //                          ::vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&y)
+        //                          .unwrap_or(vstd::pervasive::arbitrary());
+        //                        inv },
+        //                  ensures
+        //                      ::vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&y),
+        //              {
+        //                  #[allow(non_snake_case)]
+        //                  let mut VERUS_loop_next;
+        //                  match ::core::iter::Iterator::next(&mut VERUS_exec_iter) {
+        //                      ::core::option::Option::Some(VERUS_loop_val) => {
+        //                          VERUS_loop_next = VERUS_loop_val;
+        //                      }
+        //                      ::core::option::Option::None => break,
+        //                  };
+        //                  let x = VERUS_loop_next;
+        //                  let () = { body };
+        //                  proof { y = ::vstd::pervasive::ForLoopGhostIterator::ghost_advance(
+        //                      &y, &VERUS_exec_iter); }
+        //              }
+        //          }
+        //      };
+        //      VERUS_loop_result
         //  }
         // Note that "continue" and labels are not yet supported;
         // continue would also need to call ghost_advance.
@@ -1410,28 +1429,28 @@ impl Visitor {
         };
         let mut stmts: Vec<Stmt> = Vec::new();
         let expr_inv = expr.clone();
-        //              vstd::pervasive::ForLoopGhostIterator::exec_invariant(&y, &VERUS_exec_iter),
-        //              vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&y,
+        //              ::vstd::pervasive::ForLoopGhostIterator::exec_invariant(&y, &VERUS_exec_iter),
+        //              ::vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&y,
         //                  builtin::infer_spec_for_loop_iter(
-        //                      &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
-        //                          &core::iter::IntoIterator::into_iter(e)))),
+        //                      &::vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+        //                          &::core::iter::IntoIterator::into_iter(e)))),
         let exec_inv: Expr = Expr::Verbatim(quote_spanned!(span =>
             #[verifier::custom_err(#exec_inv_msg)]
-            vstd::pervasive::ForLoopGhostIterator::exec_invariant(&#x_ghost_iter, &#x_exec_iter)
+            ::vstd::pervasive::ForLoopGhostIterator::exec_invariant(&#x_ghost_iter, &#x_exec_iter)
         ));
         let ghost_inv: Expr = Expr::Verbatim(quote_spanned!(span =>
             #[verifier::custom_err(#ghost_inv_msg)]
-            vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&#x_ghost_iter,
+            ::vstd::pervasive::ForLoopGhostIterator::ghost_invariant(&#x_ghost_iter,
                 builtin::infer_spec_for_loop_iter(
-                    &vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
-                        &core::iter::IntoIterator::into_iter(#expr_inv))))
+                    &::vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(
+                        &::core::iter::IntoIterator::into_iter(#expr_inv))))
         ));
         let invariant_ensure = if let Some(mut invariant) = invariant {
             for inv in &mut invariant.exprs.exprs {
                 *inv = Expr::Verbatim(quote_spanned!(span => {
                     let #pat =
-                        vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&#x_ghost_iter)
-                        .unwrap_or(vstd::pervasive::arbitrary());
+                        ::vstd::pervasive::ForLoopGhostIterator::ghost_peek_next(&#x_ghost_iter)
+                        .unwrap_or(::vstd::pervasive::arbitrary());
                     #inv
                 }));
             }
@@ -1456,20 +1475,28 @@ impl Visitor {
         // but at the moment, user-supplied ensures aren't supported, so this would be hard to use.
         let ensure = if no_loop_invariant.is_none() {
             Some(parse_quote_spanned!(span =>
-                ensures vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&#x_ghost_iter),
+                ensures ::vstd::pervasive::ForLoopGhostIterator::ghost_ensures(&#x_ghost_iter),
             ))
         } else {
             None
         };
         self.add_loop_specs(&mut stmts, None, invariant_ensure, ensure, decreases);
-        let body_exec = Expr::Verbatim(quote_spanned!(span =>
-            if let core::option::Option::Some(#pat) = core::iter::Iterator::next(&mut #x_exec_iter)
-            #body else { break }
-        ));
+        let body_exec = Expr::Verbatim(quote_spanned!(span => {
+            #[allow(non_snake_case)]
+            let mut VERUS_loop_next;
+            match ::core::iter::Iterator::next(&mut #x_exec_iter) {
+                ::core::option::Option::Some(VERUS_loop_val) => {
+                    VERUS_loop_next = VERUS_loop_val;
+                }
+                ::core::option::Option::None => break,
+            };
+            let #pat = VERUS_loop_next;
+            let () = #body;
+        }));
         let mut body: Block = if no_loop_invariant.is_none() {
             let body_ghost = Expr::Verbatim(quote_spanned!(span =>
                 #[verifier::proof_block] {
-                    #x_ghost_iter = vstd::pervasive::ForLoopGhostIterator::ghost_advance(
+                    #x_ghost_iter = ::vstd::pervasive::ForLoopGhostIterator::ghost_advance(
                         &#x_ghost_iter, &#x_exec_iter);
                 }
             ));
@@ -1488,24 +1515,26 @@ impl Visitor {
         loop_expr.attrs = attrs;
         let full_loop: Expr = if no_loop_invariant.is_none() {
             Expr::Verbatim(quote_spanned!(span => {
-                #[allow(non_snake_case)]
-                let mut #x_exec_iter = core::iter::IntoIterator::into_iter(#expr);
                 #[allow(non_snake_case)]
                 #[verus::internal(spec)]
                 let mut #x_ghost_iter;
                 #[verifier::proof_block] {
-                    #x_ghost_iter = vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(&#x_exec_iter);
+                    #x_ghost_iter =
+                        ::vstd::pervasive::ForLoopGhostIteratorNew::ghost_iter(&#x_exec_iter);
                 }
                 #loop_expr
             }))
         } else {
-            Expr::Verbatim(quote_spanned!(span => {
-                #[allow(non_snake_case)]
-                let mut #x_exec_iter = core::iter::IntoIterator::into_iter(#expr);
-                #loop_expr
-            }))
+            Expr::Verbatim(quote_spanned!(span => { #loop_expr }))
         };
-        full_loop
+        Expr::Verbatim(quote_spanned!(span => {
+            #[allow(non_snake_case)]
+            let VERUS_loop_result = match ::core::iter::IntoIterator::into_iter(#expr) {
+                #[allow(non_snake_case)]
+                mut #x_exec_iter => #full_loop
+            };
+            VERUS_loop_result
+        }))
     }
 
     fn extract_quant_triggers(

From e062d0ac50d6b282852d23c6c3fb62946e6d19ec Mon Sep 17 00:00:00 2001
From: "Chris Hawblitzel (Microsoft)" <Chris.Hawblitzel@microsoft.com>
Date: Wed, 24 Jan 2024 13:06:51 -0800
Subject: [PATCH 19/19] Refactor state in modes.rs (#961)

* Refactor state in modes.rs

* Move Typing into submodule
---
 source/vir/src/modes.rs | 697 +++++++++++++++++++++++-----------------
 1 file changed, 400 insertions(+), 297 deletions(-)

diff --git a/source/vir/src/modes.rs b/source/vir/src/modes.rs
index 682a31075e..48052d5161 100644
--- a/source/vir/src/modes.rs
+++ b/source/vir/src/modes.rs
@@ -68,90 +68,168 @@ impl Ghost {
     }
 }
 
-struct Typing {
+struct Ctxt {
     pub(crate) funs: HashMap<Fun, Function>,
     pub(crate) datatypes: HashMap<Path, Datatype>,
     pub(crate) traits: HashSet<Path>,
+    pub(crate) check_ghost_blocks: bool,
+    pub(crate) fun_mode: Mode,
+}
+
+// Accumulated data recorded during mode checking
+struct Record {
+    pub(crate) erasure_modes: ErasureModes,
+    // Modes of InferSpecForLoopIter
+    infer_spec_for_loop_iter_modes: Option<Vec<(Span, Mode)>>,
+}
+
+// Temporary state pushed and popped during mode checking
+struct State {
     // for each variable: (is_mutable, mode)
     pub(crate) vars: ScopeMap<Ident, (bool, Mode)>,
-    pub(crate) erasure_modes: ErasureModes,
     pub(crate) in_forall_stmt: bool,
-    pub(crate) check_ghost_blocks: bool,
     // Are we in a syntactic ghost block?
     // If not, Ghost::Exec (corresponds to exec mode).
     // If yes (corresponding to proof/spec mode), say whether variables are tracked or not.
     // (note that tracked may be false even in proof mode,
     // and tracked is allowed in spec mode, although that would be needlessly constraining)
     pub(crate) block_ghostness: Ghost,
-    pub(crate) fun_mode: Mode,
     pub(crate) ret_mode: Option<Mode>,
     pub(crate) atomic_insts: Option<AtomicInstCollector>,
-    // Modes of InferSpecForLoopIter
-    infer_spec_for_loop_iter_modes: Option<Vec<(Span, Mode)>>,
 }
 
-// If we need to catch an error and continue,
-// keep a snapshot of the old transient Typing state fields so we can restore the fields.
-// (used by InferSpecForLoopIter)
-struct TypingSnapshot {
-    vars_scope_count: usize,
-    in_forall_stmt: bool,
-    block_ghostness: Ghost,
-    ret_mode: Option<Mode>,
-    atomic_insts: Option<AtomicInstCollector>,
-}
+mod typing {
+    use super::*;
 
-impl Typing {
-    fn get(&self, x: &Ident) -> (bool, Mode) {
-        *self.vars.get(x).expect("internal error: missing mode")
+    pub(super) struct Typing<'a> {
+        // don't use these fields directly; use * and push_*
+        internal_state: &'a mut State,
+        internal_undo: Option<Box<dyn for<'b> FnOnce(&'b mut State)>>,
     }
 
-    fn insert(&mut self, _span: &Span, x: &Ident, mutable: bool, mode: Mode) {
-        self.vars.insert(x.clone(), (mutable, mode)).expect("internal error: Typing insert");
+    impl Drop for Typing<'_> {
+        fn drop(&mut self) {
+            let f: Box<dyn for<'b> FnOnce(&'b mut State)> =
+                self.internal_undo.take().expect("drop-undo");
+            f(&mut self.internal_state);
+        }
     }
 
-    fn snapshot_transient_state(&mut self) -> TypingSnapshot {
-        // mention every field of Typing to make sure TypingSnapshot stays up to date with Typing:
-        let Typing {
-            funs: _,
-            datatypes: _,
-            traits: _,
-            vars,
-            erasure_modes: _,
-            in_forall_stmt,
-            check_ghost_blocks: _,
-            block_ghostness,
-            fun_mode: _,
-            ret_mode,
-            atomic_insts,
-            infer_spec_for_loop_iter_modes: _,
-        } = &self;
-        let snapshot = TypingSnapshot {
-            vars_scope_count: vars.num_scopes(),
-            in_forall_stmt: *in_forall_stmt,
-            block_ghostness: *block_ghostness,
-            ret_mode: *ret_mode,
-            atomic_insts: atomic_insts.clone(),
-        };
-        self.vars.push_scope(true);
-        snapshot
+    impl Typing<'_> {
+        pub(super) fn new<'a>(state: &'a mut State) -> Typing<'a> {
+            Typing { internal_state: state, internal_undo: Some(Box::new(|_| {})) }
+        }
+
+        pub(super) fn push_var_scope<'a>(&'a mut self) -> Typing<'a> {
+            self.internal_state.vars.push_scope(true);
+            Typing {
+                internal_state: self.internal_state,
+                internal_undo: Some(Box::new(|state| {
+                    state.vars.pop_scope();
+                })),
+            }
+        }
+
+        pub(super) fn push_var_multi_scope<'a>(&'a mut self) -> Typing<'a> {
+            let vars_scope_count = self.internal_state.vars.num_scopes();
+            Typing {
+                internal_state: self.internal_state,
+                internal_undo: Some(Box::new(move |state: &mut State| {
+                    while state.vars.num_scopes() != vars_scope_count {
+                        state.vars.pop_scope();
+                    }
+                })),
+            }
+        }
+
+        // For use after push_var_multi_scope (otherwise, use push_var_scope)
+        pub(super) fn add_var_multi_scope<'a>(&mut self) {
+            self.internal_state.vars.push_scope(true);
+        }
+
+        pub(super) fn push_in_forall_stmt<'a>(
+            &'a mut self,
+            mut in_forall_stmt: bool,
+        ) -> Typing<'a> {
+            swap(&mut in_forall_stmt, &mut self.internal_state.in_forall_stmt);
+            Typing {
+                internal_state: self.internal_state,
+                internal_undo: Some(Box::new(move |state| {
+                    state.in_forall_stmt = in_forall_stmt;
+                })),
+            }
+        }
+
+        pub(super) fn push_block_ghostness<'a>(
+            &'a mut self,
+            mut block_ghostness: Ghost,
+        ) -> Typing<'a> {
+            swap(&mut block_ghostness, &mut self.internal_state.block_ghostness);
+            Typing {
+                internal_state: self.internal_state,
+                internal_undo: Some(Box::new(move |state| {
+                    state.block_ghostness = block_ghostness;
+                })),
+            }
+        }
+
+        pub(super) fn push_ret_mode<'a>(&'a mut self, mut ret_mode: Option<Mode>) -> Typing<'a> {
+            swap(&mut ret_mode, &mut self.internal_state.ret_mode);
+            Typing {
+                internal_state: self.internal_state,
+                internal_undo: Some(Box::new(move |state| {
+                    state.ret_mode = ret_mode;
+                })),
+            }
+        }
+
+        pub(super) fn push_atomic_insts<'a>(
+            &'a mut self,
+            mut atomic_insts: Option<AtomicInstCollector>,
+        ) -> Typing<'a> {
+            swap(&mut atomic_insts, &mut self.internal_state.atomic_insts);
+            Typing {
+                internal_state: self.internal_state,
+                internal_undo: Some(Box::new(move |state| {
+                    state.atomic_insts = atomic_insts;
+                })),
+            }
+        }
+
+        // If we want to catch a VirErr, use this to make sure state is restored upon catching the error
+        pub(super) fn push_restore_on_error<'a>(&'a mut self) -> Typing<'a> {
+            self.push_var_scope()
+        }
+
+        pub(super) fn assert_zero_scopes(&self) {
+            assert_eq!(self.internal_state.vars.num_scopes(), 0);
+        }
+
+        pub(super) fn insert(&mut self, _span: &Span, x: &Ident, mutable: bool, mode: Mode) {
+            self.internal_state
+                .vars
+                .insert(x.clone(), (mutable, mode))
+                .expect("internal error: Typing insert");
+        }
+
+        pub(super) fn update_atomic_insts<'a>(&'a mut self) -> &'a mut Option<AtomicInstCollector> {
+            &mut self.internal_state.atomic_insts
+        }
+    }
+
+    impl std::ops::Deref for Typing<'_> {
+        type Target = State;
+
+        fn deref(&self) -> &State {
+            &self.internal_state
+        }
     }
+}
+use typing::Typing;
 
-    fn pop_transient_state(&mut self, snapshot: TypingSnapshot) {
-        let TypingSnapshot {
-            vars_scope_count,
-            in_forall_stmt,
-            block_ghostness,
-            ret_mode,
-            atomic_insts,
-        } = snapshot;
-        while self.vars.num_scopes() != vars_scope_count {
-            self.vars.pop_scope();
-        }
-        self.in_forall_stmt = in_forall_stmt;
-        self.block_ghostness = block_ghostness;
-        self.ret_mode = ret_mode;
-        self.atomic_insts = atomic_insts;
+impl State {
+    fn get(&self, x: &Ident) -> (bool, Mode) {
+        *self.vars.get(x).expect("internal error: missing mode")
     }
 }
 
@@ -237,9 +315,15 @@ impl AtomicInstCollector {
     }
 }
 
-fn add_pattern(typing: &mut Typing, mode: Mode, pattern: &Pattern) -> Result<(), VirErr> {
+fn add_pattern(
+    ctxt: &Ctxt,
+    record: &mut Record,
+    typing: &mut Typing,
+    mode: Mode,
+    pattern: &Pattern,
+) -> Result<(), VirErr> {
     let mut decls = vec![];
-    add_pattern_rec(typing, &mut decls, mode, pattern, false)?;
+    add_pattern_rec(ctxt, record, typing, &mut decls, mode, pattern, false)?;
     for decl in decls {
         let PatternBoundDecl { span, name, mutable, mode } = decl;
         typing.insert(&span, &name, mutable, mode);
@@ -255,6 +339,8 @@ struct PatternBoundDecl {
 }
 
 fn add_pattern_rec(
+    ctxt: &Ctxt,
+    record: &mut Record,
     typing: &mut Typing,
     decls: &mut Vec<PatternBoundDecl>,
     mode: Mode,
@@ -266,7 +352,7 @@ fn add_pattern_rec(
     if !(in_or && matches!(&pattern.x, PatternX::Or(..)))
         && !matches!(&pattern.x, PatternX::Wildcard(true))
     {
-        typing.erasure_modes.var_modes.push((pattern.span.clone(), mode));
+        record.erasure_modes.var_modes.push((pattern.span.clone(), mode));
     }
 
     match &pattern.x {
@@ -282,26 +368,34 @@ fn add_pattern_rec(
         }
         PatternX::Tuple(patterns) => {
             for p in patterns.iter() {
-                add_pattern_rec(typing, decls, mode, p, false)?;
+                add_pattern_rec(ctxt, record, typing, decls, mode, p, false)?;
             }
             Ok(())
         }
         PatternX::Constructor(datatype, variant, patterns) => {
-            let datatype = typing.datatypes[datatype].clone();
+            let datatype = ctxt.datatypes[datatype].clone();
             let variant =
                 datatype.x.variants.iter().find(|v| v.name == *variant).expect("missing variant");
             for binder in patterns.iter() {
                 let field = get_field(&variant.a, &binder.name);
                 let (_, field_mode, _) = field.a;
-                add_pattern_rec(typing, decls, mode_join(field_mode, mode), &binder.a, false)?;
+                add_pattern_rec(
+                    ctxt,
+                    record,
+                    typing,
+                    decls,
+                    mode_join(field_mode, mode),
+                    &binder.a,
+                    false,
+                )?;
             }
             Ok(())
         }
         PatternX::Or(pat1, pat2) => {
             let mut decls1 = vec![];
             let mut decls2 = vec![];
-            add_pattern_rec(typing, &mut decls1, mode, pat1, true)?;
-            add_pattern_rec(typing, &mut decls2, mode, pat2, true)?;
+            add_pattern_rec(ctxt, record, typing, &mut decls1, mode, pat1, true)?;
+            add_pattern_rec(ctxt, record, typing, &mut decls2, mode, pat2, true)?;
 
             // Rust type-checking should have made sure that both sides
             // of the pattern bound the same variables with the same types.
@@ -334,6 +428,8 @@ fn add_pattern_rec(
 }
 
 fn get_var_loc_mode(
+    ctxt: &Ctxt,
+    record: &mut Record,
     typing: &mut Typing,
     outer_mode: Mode,
     expr_inner_mode: Option<Mode>,
@@ -344,7 +440,7 @@ fn get_var_loc_mode(
         ExprX::VarLoc(x) => {
             let (_, x_mode) = typing.get(x);
 
-            if typing.check_ghost_blocks
+            if ctxt.check_ghost_blocks
                 && typing.block_ghostness == Ghost::Exec
                 && x_mode != Mode::Exec
             {
@@ -358,7 +454,7 @@ fn get_var_loc_mode(
             e1,
         ) => {
             assert!(!init_not_mut);
-            if typing.check_ghost_blocks {
+            if ctxt.check_ghost_blocks {
                 if (*op_mode == Mode::Exec) != (typing.block_ghostness == Ghost::Exec) {
                     return Err(error(
                         &expr.span,
@@ -372,7 +468,15 @@ fn get_var_loc_mode(
                     format!("cannot perform operation with mode {}", op_mode),
                 ));
             }
-            let mode1 = get_var_loc_mode(typing, outer_mode, Some(*to_mode), e1, init_not_mut)?;
+            let mode1 = get_var_loc_mode(
+                ctxt,
+                record,
+                typing,
+                outer_mode,
+                Some(*to_mode),
+                e1,
+                init_not_mut,
+            )?;
             if !mode_le(mode1, *from_mode) {
                 return Err(error(
                     &expr.span,
@@ -385,9 +489,16 @@ fn get_var_loc_mode(
             UnaryOpr::Field(FieldOpr { datatype, variant: _, field, get_variant }),
             rcvr,
         ) => {
-            let rcvr_mode =
-                get_var_loc_mode(typing, outer_mode, expr_inner_mode, rcvr, init_not_mut)?;
-            let datatype = &typing.datatypes[datatype].x;
+            let rcvr_mode = get_var_loc_mode(
+                ctxt,
+                record,
+                typing,
+                outer_mode,
+                expr_inner_mode,
+                rcvr,
+                init_not_mut,
+            )?;
+            let datatype = &ctxt.datatypes[datatype].x;
             assert!(datatype.variants.len() == 1);
             let (_, field_mode, _) = &datatype.variants[0]
                 .a
@@ -400,14 +511,13 @@ fn get_var_loc_mode(
         }
         ExprX::Block(stmts, Some(e1)) if stmts.len() == 0 => {
             // For now, only support the special case for Tracked::borrow_mut.
-            get_var_loc_mode(typing, outer_mode, None, e1, init_not_mut)?
+            get_var_loc_mode(ctxt, record, typing, outer_mode, None, e1, init_not_mut)?
         }
         ExprX::Ghost { alloc_wrapper: false, tracked: true, expr: e1 } => {
             // For now, only support the special case for Tracked::borrow_mut.
-            let prev = typing.block_ghostness;
-            typing.block_ghostness = Ghost::Ghost;
-            let mode = get_var_loc_mode(typing, outer_mode, None, e1, init_not_mut)?;
-            typing.block_ghostness = prev;
+            let mut typing = typing.push_block_ghostness(Ghost::Ghost);
+            let mode =
+                get_var_loc_mode(ctxt, record, &mut typing, outer_mode, None, e1, init_not_mut)?;
             mode
         }
         _ => {
@@ -424,19 +534,21 @@ fn get_var_loc_mode(
         ExprX::Ghost { .. } => {}
         _ => {
             let push_mode = expr_inner_mode.unwrap_or(x_mode);
-            typing.erasure_modes.var_modes.push((expr.span.clone(), push_mode));
+            record.erasure_modes.var_modes.push((expr.span.clone(), push_mode));
         }
     }
     Ok(x_mode)
 }
 
 fn check_expr_has_mode(
+    ctxt: &Ctxt,
+    record: &mut Record,
     typing: &mut Typing,
     outer_mode: Mode,
     expr: &Expr,
     expected: Mode,
 ) -> Result<(), VirErr> {
-    let mode = check_expr(typing, outer_mode, expr)?;
+    let mode = check_expr(ctxt, record, typing, outer_mode, expr)?;
     match &*expr.typ {
         crate::ast::TypX::Tuple(ts) if ts.len() == 0 => return Ok(()),
         _ => {}
@@ -448,11 +560,19 @@ fn check_expr_has_mode(
     }
 }
 
-fn check_expr(typing: &mut Typing, outer_mode: Mode, expr: &Expr) -> Result<Mode, VirErr> {
-    Ok(check_expr_handle_mut_arg(typing, outer_mode, expr)?.0)
+fn check_expr(
+    ctxt: &Ctxt,
+    record: &mut Record,
+    typing: &mut Typing,
+    outer_mode: Mode,
+    expr: &Expr,
+) -> Result<Mode, VirErr> {
+    Ok(check_expr_handle_mut_arg(ctxt, record, typing, outer_mode, expr)?.0)
 }
 
 fn check_expr_handle_mut_arg(
+    ctxt: &Ctxt,
+    record: &mut Record,
     typing: &mut Typing,
     outer_mode: Mode,
     expr: &Expr,
@@ -469,7 +589,7 @@ fn check_expr_handle_mut_arg(
 
             let x_mode = typing.get(x).1;
 
-            if typing.check_ghost_blocks
+            if ctxt.check_ghost_blocks
                 && typing.block_ghostness == Ghost::Exec
                 && x_mode != Mode::Exec
             {
@@ -481,16 +601,13 @@ fn check_expr_handle_mut_arg(
 
             let mode = mode_join(outer_mode, x_mode);
 
-            let mode = if typing.check_ghost_blocks {
-                typing.block_ghostness.join_mode(mode)
-            } else {
-                mode
-            };
-            typing.erasure_modes.var_modes.push((expr.span.clone(), mode));
+            let mode =
+                if ctxt.check_ghost_blocks { typing.block_ghostness.join_mode(mode) } else { mode };
+            record.erasure_modes.var_modes.push((expr.span.clone(), mode));
             return Ok((mode, Some(x_mode)));
         }
         ExprX::ConstVar(x, _) | ExprX::StaticVar(x) => {
-            let function = match typing.funs.get(x) {
+            let function = match ctxt.funs.get(x) {
                 None => {
                     let name = crate::ast_util::path_as_friendly_rust_name(&x.path);
                     return Err(error(&expr.span, format!("cannot find constant {}", name)));
@@ -498,7 +615,7 @@ fn check_expr_handle_mut_arg(
                 Some(f) => f.clone(),
             };
             let kind = if matches!(&expr.x, ExprX::ConstVar(..)) { "const" } else { "static" };
-            if typing.check_ghost_blocks {
+            if ctxt.check_ghost_blocks {
                 if function.x.mode == Mode::Exec && typing.block_ghostness != Ghost::Exec {
                     return Err(error(
                         &expr.span,
@@ -519,18 +636,15 @@ fn check_expr_handle_mut_arg(
                 ));
             }
             let mode = function.x.ret.x.mode;
-            let mode = if typing.check_ghost_blocks {
-                typing.block_ghostness.join_mode(mode)
-            } else {
-                mode
-            };
-            typing.erasure_modes.var_modes.push((expr.span.clone(), mode));
+            let mode =
+                if ctxt.check_ghost_blocks { typing.block_ghostness.join_mode(mode) } else { mode };
+            record.erasure_modes.var_modes.push((expr.span.clone(), mode));
             Ok(mode)
         }
         ExprX::Call(CallTarget::Fun(_, x, _, _, autospec_usage), es) => {
             assert!(*autospec_usage == AutospecUsage::Final);
 
-            let function = match typing.funs.get(x) {
+            let function = match ctxt.funs.get(x) {
                 None => {
                     let name = crate::ast_util::path_as_friendly_rust_name(&x.path);
                     return Err(error(&expr.span, format!("cannot find function {}", name)));
@@ -539,7 +653,7 @@ fn check_expr_handle_mut_arg(
             };
 
             if function.x.mode == Mode::Exec {
-                match &mut typing.atomic_insts {
+                match typing.update_atomic_insts() {
                     None => {}
                     Some(ai) => {
                         if function.x.attrs.atomic {
@@ -559,7 +673,7 @@ fn check_expr_handle_mut_arg(
                     format!("cannot call function with mode {}", function.x.mode)
                 }
             };
-            if typing.check_ghost_blocks {
+            if ctxt.check_ghost_blocks {
                 if (function.x.mode == Mode::Exec) != (typing.block_ghostness == Ghost::Exec) {
                     return Err(error(&expr.span, mode_error_msg()));
                 }
@@ -577,7 +691,7 @@ fn check_expr_handle_mut_arg(
                         ));
                     }
                     let (arg_mode_read, arg_mode_write) =
-                        check_expr_handle_mut_arg(typing, outer_mode, arg)?;
+                        check_expr_handle_mut_arg(ctxt, record, typing, outer_mode, arg)?;
                     let arg_mode_write = if let Some(arg_mode_write) = arg_mode_write {
                         arg_mode_write
                     } else {
@@ -605,44 +719,45 @@ fn check_expr_handle_mut_arg(
                         ));
                     }
                 } else {
-                    check_expr_has_mode(typing, param_mode, arg, param.x.mode)?;
+                    check_expr_has_mode(ctxt, record, typing, param_mode, arg, param.x.mode)?;
                 }
             }
             Ok(function.x.ret.x.mode)
         }
         ExprX::Call(CallTarget::FnSpec(e0), es) => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot call spec function from exec mode"));
             }
-            check_expr_has_mode(typing, Mode::Spec, e0, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, e0, Mode::Spec)?;
             for arg in es.iter() {
-                check_expr_has_mode(typing, Mode::Spec, arg, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, typing, Mode::Spec, arg, Mode::Spec)?;
             }
             Ok(Mode::Spec)
         }
         ExprX::Call(CallTarget::BuiltinSpecFun(_f, _typs), es) => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot call spec function from exec mode"));
             }
             for arg in es.iter() {
-                check_expr_has_mode(typing, Mode::Spec, arg, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, typing, Mode::Spec, arg, Mode::Spec)?;
             }
             Ok(Mode::Spec)
         }
         ExprX::Tuple(es) | ExprX::ArrayLiteral(es) => {
-            let modes = vec_map_result(es, |e| check_expr(typing, outer_mode, e))?;
+            let modes = vec_map_result(es, |e| check_expr(ctxt, record, typing, outer_mode, e))?;
             Ok(modes.into_iter().fold(outer_mode, mode_join))
         }
         ExprX::Ctor(path, variant, binders, update) => {
-            let datatype = &typing.datatypes[path].clone();
+            let datatype = &ctxt.datatypes[path].clone();
             let variant = datatype.x.get_variant(variant);
             let mut mode = mode_join(outer_mode, datatype.x.mode);
             if let Some(update) = update {
-                mode = mode_join(mode, check_expr(typing, outer_mode, update)?);
+                mode = mode_join(mode, check_expr(ctxt, record, typing, outer_mode, update)?);
             }
             for arg in binders.iter() {
                 let (_, field_mode, _) = get_field(&variant.a, &arg.name).a;
-                let mode_arg = check_expr(typing, mode_join(outer_mode, field_mode), &arg.a)?;
+                let mode_arg =
+                    check_expr(ctxt, record, typing, mode_join(outer_mode, field_mode), &arg.a)?;
                 if !mode_le(mode_arg, field_mode) {
                     // allow this arg by weakening whole struct's mode
                     mode = mode_join(mode, mode_arg);
@@ -655,7 +770,7 @@ fn check_expr_handle_mut_arg(
         ExprX::NullaryOpr(crate::ast::NullaryOpr::TraitBound(..)) => Ok(Mode::Spec),
         ExprX::Unary(UnaryOp::CoerceMode { op_mode, from_mode, to_mode, kind }, e1) => {
             // same as a call to an op_mode function with parameter from_mode and return to_mode
-            if typing.check_ghost_blocks {
+            if ctxt.check_ghost_blocks {
                 if (*op_mode == Mode::Exec) != (typing.block_ghostness == Ghost::Exec) {
                     return Err(error(
                         &expr.span,
@@ -670,7 +785,7 @@ fn check_expr_handle_mut_arg(
                 ));
             }
             let param_mode = mode_join(outer_mode, *from_mode);
-            check_expr_has_mode(typing, param_mode, e1, *from_mode)?;
+            check_expr_has_mode(ctxt, record, typing, param_mode, e1, *from_mode)?;
             if *kind == ModeCoercion::BorrowMut {
                 return Ok((*to_mode, Some(*to_mode)));
             } else {
@@ -686,15 +801,10 @@ fn check_expr_handle_mut_arg(
             // are all autospec), then keep the expression.
             // Otherwise, make a note that the expression had mode exec,
             // so that ast_simplify can replace the expression with None.
-
-            // Since we may catch a Result::Err,
-            // explicitly push and pop transient state around check_expr
-            let snapshot = typing.snapshot_transient_state();
-            let mode_opt = check_expr(typing, outer_mode, e1);
-            typing.pop_transient_state(snapshot);
-
+            let mut typing = typing.push_restore_on_error();
+            let mode_opt = check_expr(ctxt, record, &mut typing, outer_mode, e1);
             let mode = mode_opt.unwrap_or(Mode::Exec);
-            if let Some(infer_spec) = typing.infer_spec_for_loop_iter_modes.as_mut() {
+            if let Some(infer_spec) = record.infer_spec_for_loop_iter_modes.as_mut() {
                 infer_spec.push((expr.span.clone(), mode));
             } else {
                 return Err(error(
@@ -704,28 +814,29 @@ fn check_expr_handle_mut_arg(
             }
             Ok(Mode::Spec)
         }
-        ExprX::Unary(_, e1) => check_expr(typing, outer_mode, e1),
-        ExprX::UnaryOpr(UnaryOpr::Box(_), e1) => check_expr(typing, outer_mode, e1),
-        ExprX::UnaryOpr(UnaryOpr::Unbox(_), e1) => check_expr(typing, outer_mode, e1),
+        ExprX::Unary(_, e1) => check_expr(ctxt, record, typing, outer_mode, e1),
+        ExprX::UnaryOpr(UnaryOpr::Box(_), e1) => check_expr(ctxt, record, typing, outer_mode, e1),
+        ExprX::UnaryOpr(UnaryOpr::Unbox(_), e1) => check_expr(ctxt, record, typing, outer_mode, e1),
         ExprX::UnaryOpr(UnaryOpr::HasType(_), _) => panic!("internal error: HasType in modes.rs"),
         ExprX::UnaryOpr(UnaryOpr::IsVariant { .. }, e1) => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot test variant in exec mode"));
             }
-            check_expr(typing, outer_mode, e1)
+            check_expr(ctxt, record, typing, outer_mode, e1)
         }
         ExprX::UnaryOpr(UnaryOpr::TupleField { .. }, e1) => {
-            return check_expr_handle_mut_arg(typing, outer_mode, e1);
+            return check_expr_handle_mut_arg(ctxt, record, typing, outer_mode, e1);
         }
         ExprX::UnaryOpr(
             UnaryOpr::Field(FieldOpr { datatype, variant, field, get_variant }),
             e1,
         ) => {
-            if *get_variant && typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if *get_variant && ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot get variant in exec mode"));
             }
-            let (e1_mode_read, e1_mode_write) = check_expr_handle_mut_arg(typing, outer_mode, e1)?;
-            let datatype = &typing.datatypes[datatype];
+            let (e1_mode_read, e1_mode_write) =
+                check_expr_handle_mut_arg(ctxt, record, typing, outer_mode, e1)?;
+            let datatype = &ctxt.datatypes[datatype];
             let field = get_field(&datatype.x.get_variant(variant).a, field);
             let field_mode = field.a.1;
             let mode_read = mode_join(e1_mode_read, field_mode);
@@ -737,15 +848,15 @@ fn check_expr_handle_mut_arg(
         }
         ExprX::UnaryOpr(UnaryOpr::IntegerTypeBound(_kind, min_mode), e1) => {
             let joined_mode = mode_join(outer_mode, *min_mode);
-            let mode = check_expr(typing, joined_mode, e1)?;
+            let mode = check_expr(ctxt, record, typing, joined_mode, e1)?;
             Ok(mode_join(*min_mode, mode))
         }
         ExprX::UnaryOpr(UnaryOpr::CustomErr(_), e1) => {
-            check_expr_has_mode(typing, Mode::Spec, e1, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, e1, Mode::Spec)?;
             Ok(Mode::Spec)
         }
         ExprX::Loc(e) => {
-            return check_expr_handle_mut_arg(typing, outer_mode, e);
+            return check_expr_handle_mut_arg(ctxt, record, typing, outer_mode, e);
         }
         ExprX::Binary(op, e1, e2) => {
             let op_mode = match op {
@@ -759,50 +870,42 @@ fn check_expr_handle_mut_arg(
                 BinaryOp::HeightCompare { .. } => Mode::Spec,
                 _ => outer_mode,
             };
-            let mode1 = check_expr(typing, outer_mode, e1)?;
-            let mode2 = check_expr(typing, outer_mode, e2)?;
+            let mode1 = check_expr(ctxt, record, typing, outer_mode, e1)?;
+            let mode2 = check_expr(ctxt, record, typing, outer_mode, e2)?;
             Ok(mode_join(op_mode, mode_join(mode1, mode2)))
         }
         ExprX::BinaryOpr(crate::ast::BinaryOpr::ExtEq(..), e1, e2) => {
-            check_expr_has_mode(typing, Mode::Spec, e1, Mode::Spec)?;
-            check_expr_has_mode(typing, Mode::Spec, e2, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, e1, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, e2, Mode::Spec)?;
             Ok(Mode::Spec)
         }
         ExprX::Multi(MultiOp::Chained(_), es) => {
             for e in es.iter() {
-                check_expr_has_mode(typing, Mode::Spec, e, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, typing, Mode::Spec, e, Mode::Spec)?;
             }
             Ok(Mode::Spec)
         }
         ExprX::Quant(_, binders, e1) => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use forall/exists in exec mode"));
             }
-            typing.vars.push_scope(true);
+            let mut typing = typing.push_var_scope();
             for binder in binders.iter() {
                 typing.insert(&expr.span, &binder.name, false, Mode::Spec);
             }
-            check_expr_has_mode(typing, Mode::Spec, e1, Mode::Spec)?;
-            typing.vars.pop_scope();
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, e1, Mode::Spec)?;
             Ok(Mode::Spec)
         }
         ExprX::Closure(params, body) => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use FnSpec closure in 'exec' mode"));
             }
-            typing.vars.push_scope(true);
+            let mut typing = typing.push_var_scope();
             for binder in params.iter() {
                 typing.insert(&expr.span, &binder.name, false, Mode::Spec);
             }
-
-            let mut inner_atomic_insts = None;
-            swap(&mut inner_atomic_insts, &mut typing.atomic_insts);
-
-            check_expr_has_mode(typing, Mode::Spec, body, Mode::Spec)?;
-
-            swap(&mut inner_atomic_insts, &mut typing.atomic_insts);
-
-            typing.vars.pop_scope();
+            let mut typing = typing.push_atomic_insts(None);
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, body, Mode::Spec)?;
             Ok(Mode::Spec)
         }
         ExprX::ExecClosure { params, ret, requires, ensures, body, external_spec } => {
@@ -815,59 +918,49 @@ fn check_expr_handle_mut_arg(
                     "closure in ghost code must be marked as a FnSpec by wrapping it in `closure_to_fn_spec` (this should happen automatically in the Verus syntax macro)",
                 ));
             }
-            typing.vars.push_scope(true);
+            let mut typing = typing.push_var_scope();
             for binder in params.iter() {
                 typing.insert(&expr.span, &binder.name, false, Mode::Exec);
             }
+            let mut typing = typing.push_atomic_insts(None);
+            let mut typing = typing.push_ret_mode(Some(Mode::Exec));
 
-            let mut inner_atomic_insts = None;
-            let mut ret_mode = Some(Mode::Exec);
-            let mut block_ghostness = Ghost::Ghost;
-            swap(&mut inner_atomic_insts, &mut typing.atomic_insts);
-            swap(&mut ret_mode, &mut typing.ret_mode);
-            swap(&mut block_ghostness, &mut typing.block_ghostness);
-
+            let mut ghost_typing = typing.push_block_ghostness(Ghost::Ghost);
             for req in requires.iter() {
-                check_expr_has_mode(typing, Mode::Spec, req, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, &mut ghost_typing, Mode::Spec, req, Mode::Spec)?;
             }
 
-            typing.vars.push_scope(true);
-            typing.insert(&expr.span, &ret.name, false, Mode::Exec);
+            let mut ens_typing = ghost_typing.push_var_scope();
+            ens_typing.insert(&expr.span, &ret.name, false, Mode::Exec);
             for ens in ensures.iter() {
-                check_expr_has_mode(typing, Mode::Spec, ens, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, &mut ens_typing, Mode::Spec, ens, Mode::Spec)?;
             }
-            typing.vars.pop_scope();
-
-            swap(&mut block_ghostness, &mut typing.block_ghostness);
-
-            check_expr_has_mode(typing, Mode::Exec, body, Mode::Exec)?;
+            drop(ens_typing);
+            drop(ghost_typing);
 
-            swap(&mut inner_atomic_insts, &mut typing.atomic_insts);
-            swap(&mut ret_mode, &mut typing.ret_mode);
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Exec, body, Mode::Exec)?;
 
-            typing.vars.pop_scope();
             Ok(Mode::Exec)
         }
         ExprX::Choose { params, cond, body } => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use choose in exec mode"));
             }
-            typing.vars.push_scope(true);
+            let mut typing = typing.push_var_scope();
             for binder in params.iter() {
                 typing.insert(&expr.span, &binder.name, false, Mode::Spec);
             }
-            check_expr_has_mode(typing, Mode::Spec, cond, Mode::Spec)?;
-            check_expr_has_mode(typing, Mode::Spec, body, Mode::Spec)?;
-            typing.vars.pop_scope();
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, cond, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, body, Mode::Spec)?;
             Ok(Mode::Spec)
         }
         ExprX::WithTriggers { triggers, body } => {
             for trigger in triggers.iter() {
                 for term in trigger.iter() {
-                    check_expr_has_mode(typing, Mode::Spec, term, Mode::Spec)?;
+                    check_expr_has_mode(ctxt, record, typing, Mode::Spec, term, Mode::Spec)?;
                 }
             }
-            check_expr_has_mode(typing, Mode::Spec, body, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, body, Mode::Spec)?;
             Ok(Mode::Spec)
         }
         ExprX::Assign { init_not_mut, lhs, rhs, op: _ } => {
@@ -877,14 +970,15 @@ fn check_expr_handle_mut_arg(
                     "assignment is not allowed in 'assert ... by' statement",
                 ));
             }
-            let x_mode = get_var_loc_mode(typing, outer_mode, None, lhs, *init_not_mut)?;
+            let x_mode =
+                get_var_loc_mode(ctxt, record, typing, outer_mode, None, lhs, *init_not_mut)?;
             if !mode_le(outer_mode, x_mode) {
                 return Err(error(
                     &expr.span,
                     format!("cannot assign to {x_mode} variable from {outer_mode} mode"),
                 ));
             }
-            check_expr_has_mode(typing, outer_mode, rhs, x_mode)?;
+            check_expr_has_mode(ctxt, record, typing, outer_mode, rhs, x_mode)?;
             Ok(x_mode)
         }
         ExprX::Fuel(_, _) => {
@@ -903,84 +997,81 @@ fn check_expr_handle_mut_arg(
         }
         ExprX::Header(_) => panic!("internal error: Header shouldn't exist here"),
         ExprX::AssertAssume { is_assume: _, expr: e } => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use assert or assume in exec mode"));
             }
-            check_expr_has_mode(typing, Mode::Spec, e, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, e, Mode::Spec)?;
             Ok(outer_mode)
         }
         ExprX::AssertBy { vars, require, ensure, proof } => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use 'assert ... by' in exec mode")
                     .help("use a `proof` block"));
             }
-            let in_forall_stmt = typing.in_forall_stmt;
             // REVIEW: we could allow proof vars when vars.len() == 0,
             // but we'd have to implement the proper lifetime checking in erase.rs
-            typing.in_forall_stmt = true;
-            typing.vars.push_scope(true);
+            let mut typing = typing.push_in_forall_stmt(true);
+            let mut typing = typing.push_var_scope();
             for var in vars.iter() {
                 typing.insert(&expr.span, &var.name, false, Mode::Spec);
             }
-            check_expr_has_mode(typing, Mode::Spec, require, Mode::Spec)?;
-            check_expr_has_mode(typing, Mode::Spec, ensure, Mode::Spec)?;
-            check_expr_has_mode(typing, Mode::Proof, proof, Mode::Proof)?;
-            typing.vars.pop_scope();
-            typing.in_forall_stmt = in_forall_stmt;
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, require, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, ensure, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, &mut typing, Mode::Proof, proof, Mode::Proof)?;
             Ok(Mode::Proof)
         }
         ExprX::AssertQuery { requires, ensures, proof, mode: _ } => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use assert in exec mode"));
             }
             for req in requires.iter() {
-                check_expr_has_mode(typing, Mode::Spec, req, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, typing, Mode::Spec, req, Mode::Spec)?;
             }
             for ens in ensures.iter() {
-                check_expr_has_mode(typing, Mode::Spec, ens, Mode::Spec)?;
+                check_expr_has_mode(ctxt, record, typing, Mode::Spec, ens, Mode::Spec)?;
             }
-            check_expr_has_mode(typing, Mode::Proof, proof, Mode::Proof)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Proof, proof, Mode::Proof)?;
             Ok(Mode::Proof)
         }
         ExprX::AssertCompute(e, _) => {
-            if typing.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness == Ghost::Exec {
                 return Err(error(&expr.span, "cannot use assert in exec mode"));
             }
-            check_expr_has_mode(typing, Mode::Spec, e, Mode::Spec)?;
+            check_expr_has_mode(ctxt, record, typing, Mode::Spec, e, Mode::Spec)?;
             Ok(Mode::Proof)
         }
         ExprX::If(e1, e2, e3) => {
-            let mode1 = check_expr(typing, outer_mode, e1)?;
-            if typing.check_ghost_blocks
+            let mode1 = check_expr(ctxt, record, typing, outer_mode, e1)?;
+            if ctxt.check_ghost_blocks
                 && typing.block_ghostness == Ghost::Exec
                 && mode1 != Mode::Exec
             {
                 return Err(error(&expr.span, "condition must have mode exec"));
             }
-            typing.erasure_modes.condition_modes.push((expr.span.clone(), mode1));
+            record.erasure_modes.condition_modes.push((expr.span.clone(), mode1));
 
             let mode_branch = match (outer_mode, mode1) {
                 (Mode::Exec, Mode::Spec) => Mode::Proof,
                 _ => outer_mode,
             };
-            let mode2 = check_expr(typing, mode_branch, e2)?;
+            let mode2 = check_expr(ctxt, record, typing, mode_branch, e2)?;
             match e3 {
                 None => Ok(mode2),
                 Some(e3) => {
-                    let mode3 = check_expr(typing, mode_branch, e3)?;
+                    let mode3 = check_expr(ctxt, record, typing, mode_branch, e3)?;
                     Ok(mode_join(mode2, mode3))
                 }
             }
         }
         ExprX::Match(e1, arms) => {
-            let mode1 = check_expr(typing, outer_mode, e1)?;
-            if typing.check_ghost_blocks
+            let mode1 = check_expr(ctxt, record, typing, outer_mode, e1)?;
+            if ctxt.check_ghost_blocks
                 && typing.block_ghostness == Ghost::Exec
                 && mode1 != Mode::Exec
             {
                 return Err(error(&expr.span, "exec code cannot match on non-exec value"));
             }
-            typing.erasure_modes.condition_modes.push((expr.span.clone(), mode1));
+            record.erasure_modes.condition_modes.push((expr.span.clone(), mode1));
 
             match (mode1, arms.len()) {
                 (Mode::Spec, 0) => {
@@ -992,47 +1083,45 @@ fn check_expr_handle_mut_arg(
             }
             let mut final_mode = outer_mode;
             for arm in arms.iter() {
-                typing.vars.push_scope(true);
-                add_pattern(typing, mode1, &arm.x.pattern)?;
+                let mut typing = typing.push_var_scope();
+                add_pattern(ctxt, record, &mut typing, mode1, &arm.x.pattern)?;
                 let arm_outer_mode = match (outer_mode, mode1) {
                     (Mode::Exec, Mode::Spec | Mode::Proof) => Mode::Proof,
                     (m, _) => m,
                 };
-                let guard_mode = check_expr(typing, arm_outer_mode, &arm.x.guard)?;
+                let guard_mode =
+                    check_expr(ctxt, record, &mut typing, arm_outer_mode, &arm.x.guard)?;
                 let arm_outer_mode = match (arm_outer_mode, guard_mode) {
                     (Mode::Exec, Mode::Spec | Mode::Proof) => Mode::Proof,
                     (m, _) => m,
                 };
-                let arm_mode = check_expr(typing, arm_outer_mode, &arm.x.body)?;
+                let arm_mode = check_expr(ctxt, record, &mut typing, arm_outer_mode, &arm.x.body)?;
                 final_mode = mode_join(final_mode, arm_mode);
-                typing.vars.pop_scope();
             }
             Ok(final_mode)
         }
         ExprX::Loop { is_for_loop: _, label: _, cond, body, invs } => {
             // We could also allow this for proof, if we check it for termination
-            if typing.check_ghost_blocks && typing.block_ghostness != Ghost::Exec {
+            if ctxt.check_ghost_blocks && typing.block_ghostness != Ghost::Exec {
                 return Err(error(&expr.span, "cannot use while in proof or spec mode"));
             }
-            match &mut typing.atomic_insts {
+            match typing.update_atomic_insts() {
                 None => {}
                 Some(ai) => ai.add_loop(&expr.span),
             }
             if let Some(cond) = cond {
-                check_expr_has_mode(typing, outer_mode, cond, Mode::Exec)?;
+                check_expr_has_mode(ctxt, record, typing, outer_mode, cond, Mode::Exec)?;
             }
-            check_expr_has_mode(typing, outer_mode, body, Mode::Exec)?;
+            check_expr_has_mode(ctxt, record, typing, outer_mode, body, Mode::Exec)?;
             for inv in invs.iter() {
-                let prev = typing.block_ghostness;
-                typing.block_ghostness = Ghost::Ghost;
-                check_expr_has_mode(typing, Mode::Spec, &inv.inv, Mode::Spec)?;
-                typing.block_ghostness = prev;
+                let mut typing = typing.push_block_ghostness(Ghost::Ghost);
+                check_expr_has_mode(ctxt, record, &mut typing, Mode::Spec, &inv.inv, Mode::Spec)?;
             }
             Ok(Mode::Exec)
         }
         ExprX::Return(e1) => {
-            if typing.check_ghost_blocks {
-                match (typing.fun_mode, typing.block_ghostness) {
+            if ctxt.check_ghost_blocks {
+                match (ctxt.fun_mode, typing.block_ghostness) {
                     (Mode::Exec, Ghost::Exec) => {}
                     (Mode::Proof, _) => {}
                     (Mode::Spec, _) => {}
@@ -1044,7 +1133,7 @@ fn check_expr_handle_mut_arg(
                     }
                 }
             } else {
-                match (typing.fun_mode, outer_mode) {
+                match (ctxt.fun_mode, outer_mode) {
                     (Mode::Exec, Mode::Exec) => {}
                     (Mode::Proof, _) => {}
                     (Mode::Spec, _) => {}
@@ -1072,15 +1161,14 @@ fn check_expr_handle_mut_arg(
                     } => {}
                 (_, None) => panic!("internal error: missing return type"),
                 (Some(e1), Some(ret_mode)) => {
-                    check_expr_has_mode(typing, outer_mode, e1, ret_mode)?;
+                    check_expr_has_mode(ctxt, record, typing, outer_mode, e1, ret_mode)?;
                 }
             }
             Ok(Mode::Exec)
         }
         ExprX::BreakOrContinue { label: _, is_break: _ } => Ok(Mode::Exec),
         ExprX::Ghost { alloc_wrapper, tracked, expr: e1 } => {
-            let prev = typing.block_ghostness;
-            let block_ghostness = match (prev, alloc_wrapper, tracked) {
+            let block_ghostness = match (typing.block_ghostness, alloc_wrapper, tracked) {
                 (Ghost::Exec, false, false) => match &*e1.typ {
                     crate::ast::TypX::Tuple(ts) if ts.len() == 0 => Ghost::Ghost,
                     _ => {
@@ -1105,12 +1193,12 @@ fn check_expr_handle_mut_arg(
                     ));
                 }
             };
-            typing.block_ghostness = block_ghostness;
+            let mut typing = typing.push_block_ghostness(block_ghostness);
             let outer_mode = match (outer_mode, block_ghostness) {
                 (Mode::Exec, Ghost::Ghost) => Mode::Proof,
                 _ => outer_mode,
             };
-            let inner_mode = check_expr_handle_mut_arg(typing, outer_mode, e1)?;
+            let inner_mode = check_expr_handle_mut_arg(ctxt, record, &mut typing, outer_mode, e1)?;
             let mode = if *alloc_wrapper {
                 let (inner_read, inner_write) = inner_mode;
                 let target_mode = if *tracked { Mode::Proof } else { Mode::Spec };
@@ -1132,24 +1220,21 @@ fn check_expr_handle_mut_arg(
             } else {
                 inner_mode
             };
-            typing.block_ghostness = prev;
             return Ok(mode);
         }
         ExprX::Block(ss, Some(e1)) if ss.len() == 0 => {
-            return check_expr_handle_mut_arg(typing, outer_mode, e1);
+            return check_expr_handle_mut_arg(ctxt, record, typing, outer_mode, e1);
         }
         ExprX::Block(ss, e1) => {
+            let mut typing = typing.push_var_multi_scope();
             for stmt in ss.iter() {
-                typing.vars.push_scope(true);
-                check_stmt(typing, outer_mode, stmt)?;
+                typing.add_var_multi_scope();
+                check_stmt(ctxt, record, &mut typing, outer_mode, stmt)?;
             }
             let mode = match e1 {
                 None => outer_mode,
-                Some(expr) => check_expr(typing, outer_mode, expr)?,
+                Some(expr) => check_expr(ctxt, record, &mut typing, outer_mode, expr)?,
             };
-            for _ in ss.iter() {
-                typing.vars.pop_scope();
-            }
             Ok(mode)
         }
         ExprX::OpenInvariant(inv, binder, body, atomicity) => {
@@ -1157,15 +1242,14 @@ fn check_expr_handle_mut_arg(
                 return Err(error(&expr.span, "Cannot open invariant in Spec mode."));
             }
 
-            let prev = typing.block_ghostness;
-            typing.block_ghostness = Ghost::Ghost;
-            let mode1 = check_expr(typing, outer_mode, inv)?;
-            typing.block_ghostness = prev;
+            let mut ghost_typing = typing.push_block_ghostness(Ghost::Ghost);
+            let mode1 = check_expr(ctxt, record, &mut ghost_typing, outer_mode, inv)?;
+            drop(ghost_typing);
 
             if mode1 != Mode::Proof {
                 return Err(error(&inv.span, "Invariant must be Proof mode."));
             }
-            typing.vars.push_scope(true);
+            let mut typing = typing.push_var_scope();
             typing.insert(&expr.span, &binder.name, /* mutable */ true, Mode::Proof);
 
             if *atomicity == InvAtomicity::NonAtomic
@@ -1178,26 +1262,33 @@ fn check_expr_handle_mut_arg(
                 // mode, and we don't need to do the atomicity check at all.
                 // And of course, we don't do atomicity checks for the 'NonAtomic'
                 // invariant type.
-                let _ = check_expr(typing, outer_mode, body)?;
+                let _ = check_expr(ctxt, record, &mut typing, outer_mode, body)?;
             } else {
-                let mut my_atomic_insts = Some(AtomicInstCollector::new());
-                swap(&mut my_atomic_insts, &mut typing.atomic_insts);
-                let _ = check_expr(typing, outer_mode, body)?;
-                swap(&mut my_atomic_insts, &mut typing.atomic_insts);
-                my_atomic_insts.expect("my_atomic_insts").validate(&body.span, false)?;
+                let mut typing = typing.push_atomic_insts(Some(AtomicInstCollector::new()));
+                let _ = check_expr(ctxt, record, &mut typing, outer_mode, body)?;
+                typing
+                    .atomic_insts
+                    .as_ref()
+                    .expect("my_atomic_insts")
+                    .validate(&body.span, false)?;
             }
 
-            typing.vars.pop_scope();
             Ok(Mode::Exec)
         }
     };
     Ok((mode?, None))
 }
 
-fn check_stmt(typing: &mut Typing, outer_mode: Mode, stmt: &Stmt) -> Result<(), VirErr> {
+fn check_stmt(
+    ctxt: &Ctxt,
+    record: &mut Record,
+    typing: &mut Typing,
+    outer_mode: Mode,
+    stmt: &Stmt,
+) -> Result<(), VirErr> {
     match &stmt.x {
         StmtX::Expr(e) => {
-            let _ = check_expr(typing, outer_mode, e)?;
+            let _ = check_expr(ctxt, record, typing, outer_mode, e)?;
             Ok(())
         }
         StmtX::Decl { pattern, mode, init } => {
@@ -1206,7 +1297,7 @@ fn check_stmt(typing: &mut Typing, outer_mode: Mode, stmt: &Stmt) -> Result<(),
             } else {
                 *mode
             };
-            if typing.check_ghost_blocks
+            if ctxt.check_ghost_blocks
                 && typing.block_ghostness == Ghost::Exec
                 && mode != Mode::Exec
                 && init.is_some()
@@ -1216,11 +1307,11 @@ fn check_stmt(typing: &mut Typing, outer_mode: Mode, stmt: &Stmt) -> Result<(),
             if !mode_le(outer_mode, mode) {
                 return Err(error(&stmt.span, format!("pattern cannot have mode {}", mode)));
             }
-            add_pattern(typing, mode, pattern)?;
+            add_pattern(ctxt, record, typing, mode, pattern)?;
             match init.as_ref() {
                 None => {}
                 Some(expr) => {
-                    check_expr_has_mode(typing, outer_mode, expr, mode)?;
+                    check_expr_has_mode(ctxt, record, typing, outer_mode, expr, mode)?;
                 }
             }
             Ok(())
@@ -1228,13 +1319,18 @@ fn check_stmt(typing: &mut Typing, outer_mode: Mode, stmt: &Stmt) -> Result<(),
     }
 }
 
-fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), VirErr> {
-    typing.vars.push_scope(true);
+fn check_function(
+    ctxt: &Ctxt,
+    record: &mut Record,
+    typing: &mut Typing,
+    function: &mut Function,
+) -> Result<(), VirErr> {
+    let mut fun_typing = typing.push_var_scope();
 
     if let FunctionKind::TraitMethodImpl { method, trait_path, .. } = &function.x.kind {
-        let our_trait = typing.traits.contains(trait_path);
+        let our_trait = ctxt.traits.contains(trait_path);
         let (expected_params, expected_ret_mode): (Vec<Mode>, Mode) = if our_trait {
-            let trait_method = &typing.funs[method];
+            let trait_method = &ctxt.funs[method];
             let expect_mode = trait_method.x.mode;
             if function.x.mode != expect_mode {
                 return Err(error(
@@ -1273,30 +1369,30 @@ fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), Vi
         }
         let inner_param_mode =
             if let Some((mode, _)) = param.x.unwrapped_info { mode } else { param.x.mode };
-        typing.insert(&function.span, &param.x.name, param.x.is_mut, inner_param_mode);
+        fun_typing.insert(&function.span, &param.x.name, param.x.is_mut, inner_param_mode);
     }
 
     for expr in function.x.require.iter() {
-        typing.block_ghostness = Ghost::Ghost;
-        check_expr_has_mode(typing, Mode::Spec, expr, Mode::Spec)?;
+        let mut req_typing = fun_typing.push_block_ghostness(Ghost::Ghost);
+        check_expr_has_mode(ctxt, record, &mut req_typing, Mode::Spec, expr, Mode::Spec)?;
     }
 
-    typing.vars.push_scope(true);
+    let mut ens_typing = fun_typing.push_var_scope();
     if function.x.has_return() {
-        typing.insert(&function.span, &function.x.ret.x.name, false, function.x.ret.x.mode);
+        ens_typing.insert(&function.span, &function.x.ret.x.name, false, function.x.ret.x.mode);
     }
     for expr in function.x.ensure.iter() {
-        typing.block_ghostness = Ghost::Ghost;
-        check_expr_has_mode(typing, Mode::Spec, expr, Mode::Spec)?;
+        let mut ens_typing = ens_typing.push_block_ghostness(Ghost::Ghost);
+        check_expr_has_mode(ctxt, record, &mut ens_typing, Mode::Spec, expr, Mode::Spec)?;
     }
-    typing.vars.pop_scope();
+    drop(ens_typing);
 
     for expr in function.x.decrease.iter() {
-        typing.block_ghostness = Ghost::Ghost;
-        check_expr_has_mode(typing, Mode::Spec, expr, Mode::Spec)?;
+        let mut dec_typing = fun_typing.push_block_ghostness(Ghost::Ghost);
+        check_expr_has_mode(ctxt, record, &mut dec_typing, Mode::Spec, expr, Mode::Spec)?;
     }
 
-    if function.x.has_return() {
+    let ret_mode = if function.x.has_return() {
         let ret_mode = function.x.ret.x.mode;
         if !matches!(function.x.item_kind, ItemKind::Const) && !mode_le(function.x.mode, ret_mode) {
             return Err(error(
@@ -1321,17 +1417,27 @@ fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), Vi
                 ));
             }
         }
-        typing.ret_mode = Some(ret_mode);
-    }
+        Some(ret_mode)
+    } else {
+        None
+    };
     if let Some(body) = &function.x.body {
-        typing.block_ghostness = Ghost::of_mode(function.x.mode);
-        assert!(typing.infer_spec_for_loop_iter_modes.is_none());
-        typing.infer_spec_for_loop_iter_modes = Some(Vec::new());
-        check_expr_has_mode(typing, function.x.mode, body, function.x.ret.x.mode)?;
+        let mut body_typing = fun_typing.push_ret_mode(ret_mode);
+        let mut body_typing = body_typing.push_block_ghostness(Ghost::of_mode(function.x.mode));
+        assert!(record.infer_spec_for_loop_iter_modes.is_none());
+        record.infer_spec_for_loop_iter_modes = Some(Vec::new());
+        check_expr_has_mode(
+            ctxt,
+            record,
+            &mut body_typing,
+            function.x.mode,
+            body,
+            function.x.ret.x.mode,
+        )?;
 
         // Replace InferSpecForLoopIter None if it fails to have mode spec
         // (if it's mode spec, leave as is to be processed by sst_to_air and loop_inference)
-        let infer_spec = typing.infer_spec_for_loop_iter_modes.as_ref().expect("infer_spec");
+        let infer_spec = record.infer_spec_for_loop_iter_modes.as_ref().expect("infer_spec");
         if infer_spec.len() > 0 {
             let mut functionx = function.x.clone();
             functionx.body = Some(crate::ast_visitor::map_expr_visitor(body, &|expr: &Expr| {
@@ -1352,11 +1458,10 @@ fn check_function(typing: &mut Typing, function: &mut Function) -> Result<(), Vi
             })?);
             *function = function.new_x(functionx);
         }
-        typing.infer_spec_for_loop_iter_modes = None;
+        record.infer_spec_for_loop_iter_modes = None;
     }
-    typing.ret_mode = None;
-    typing.vars.pop_scope();
-    assert_eq!(typing.vars.num_scopes(), 0);
+    drop(fun_typing);
+    typing.assert_zero_scopes();
     Ok(())
 }
 
@@ -1370,35 +1475,33 @@ pub fn check_crate(krate: &Krate) -> Result<(Krate, ErasureModes), VirErr> {
         datatypes.insert(datatype.x.path.clone(), datatype.clone());
     }
     let erasure_modes = ErasureModes { condition_modes: vec![], var_modes: vec![] };
-    let mut typing = Typing {
+    let mut ctxt = Ctxt {
         funs,
         datatypes,
         traits: krate.traits.iter().map(|t| t.x.name.clone()).collect(),
+        check_ghost_blocks: false,
+        fun_mode: Mode::Exec,
+    };
+    let mut record = Record { erasure_modes, infer_spec_for_loop_iter_modes: None };
+    let mut state = State {
         vars: ScopeMap::new(),
-        erasure_modes,
         in_forall_stmt: false,
-        check_ghost_blocks: false,
         block_ghostness: Ghost::Exec,
-        fun_mode: Mode::Exec,
         ret_mode: None,
         atomic_insts: None,
-        infer_spec_for_loop_iter_modes: None,
     };
+    let mut typing = Typing::new(&mut state);
     let mut kratex = (**krate).clone();
     for function in kratex.functions.iter_mut() {
-        typing.check_ghost_blocks = function.x.attrs.uses_ghost_blocks;
-        typing.fun_mode = function.x.mode;
+        ctxt.check_ghost_blocks = function.x.attrs.uses_ghost_blocks;
+        ctxt.fun_mode = function.x.mode;
         if function.x.attrs.atomic {
-            let mut my_atomic_insts = Some(AtomicInstCollector::new());
-            swap(&mut my_atomic_insts, &mut typing.atomic_insts);
-
-            check_function(&mut typing, function)?;
-
-            swap(&mut my_atomic_insts, &mut typing.atomic_insts);
-            my_atomic_insts.expect("my_atomic_insts").validate(&function.span, true)?;
+            let mut typing = typing.push_atomic_insts(Some(AtomicInstCollector::new()));
+            check_function(&ctxt, &mut record, &mut typing, function)?;
+            typing.atomic_insts.as_ref().expect("atomic_insts").validate(&function.span, true)?;
         } else {
-            check_function(&mut typing, function)?;
+            check_function(&ctxt, &mut record, &mut typing, function)?;
         }
     }
-    Ok((Arc::new(kratex), typing.erasure_modes))
+    Ok((Arc::new(kratex), record.erasure_modes))
 }