fix(repsel): contain guarded for-of element facts

Author: Ralph Küpper

crates/perry-codegen/src/collectors/ptr_shape_elements.rs

@@ -41,9 +41,11 @@
 //!   `A` **dense** and monomorphic for its whole lifetime.
 //! * **E3 — array containment.** Every *other* use of `A` is an in-bounds
 //!   element read (E5), a `.length` read, or `return A`. The return exemption
-//!   is #7034 §4's, unchanged and for the same reason. Anything else — call
-//!   argument, closure capture, reassignment, `IndexSet`, an unrecognised
-//!   array method, being an element of another container — disqualifies `A`.
+//!   is #7034 §4's, unchanged and for the same reason. The one conditional
+//!   escape is the compiler-generated `ArrayIterationPatched` guard described
+//!   below. Anything else — call argument, closure capture, reassignment,
+//!   `IndexSet`, an unrecognised array method, being an element of another
+//!   container — disqualifies `A`.
 //! * **E4 — class admissibility.** `C` passes the same `chain_admissible`
 //!   gate rule 1 applies to a `new C(...)` local, and the module-wide rule-5
 //!   barrier scan is clear.
@@ -58,10 +60,17 @@
 //!   `A[i]` can be `undefined`, and a guard-free fixed-offset load masks a
 //!   NaN-boxed `undefined` into a wild pointer.
 //!
-//! `for (const r of A)` desugars to exactly the E5 shape
+//! `for (const r of A)` has an E5 index arm
 //! (`lower/stmt_loops.rs::lazy_or_index_elem` — a `__idx` local, `__idx <
-//! __arr.length`, `Let r = IndexGet(__arr, __idx)`), so the iterator form is
-//! covered by the indexed proof rather than by a second one.
+//! __arr.length`, `Let r = IndexGet(__arr, __idx)`) behind the
+//! `ArrayIterationPatched` runtime guard. Its lazy arm passes `A` to
+//! `GetIterator`, which is ordinarily an E3 escape: a custom iterator can
+//! reshape an element before returning. The index-arm facts remain sound only
+//! when that top-level guard is the last use of both the array and every
+//! element-group member. Then the mutating lazy arm and the proven index arm
+//! are mutually exclusive, and no fact crosses their join. A nested guard is
+//! refused because a loop backedge could bring the mutated array to an
+//! earlier proven access on the next iteration.
 //!
 //! ## What the facts are used for
 //!
@@ -314,14 +323,15 @@ pub(crate) fn collect_element_shape_facts(
     // E3/E5: the array use walk.
     let mut walk = ArrayWalk {
         roots: &array_roots,
+        alias_edges: &alias_edges,
         disqualified: HashSet::new(),
         pushes: HashMap::new(),
         reads: Vec::new(),
         idx_writes: HashMap::new(),
         bounded: Vec::new(),
         in_closure: false,
     };
-    walk.walk_stmts(stmts);
+    walk.walk_region_stmts(stmts);
     let ArrayWalk {
         mut disqualified,
         pushes,
@@ -487,6 +497,7 @@ struct ReadSite {
 
 struct ArrayWalk<'a> {
     roots: &'a HashMap<u32, u32>,
+    alias_edges: &'a [(u32, u32)],
     disqualified: HashSet<u32>,
     pushes: HashMap<u32, Vec<PushValue>>,
     reads: Vec<ReadSite>,
@@ -520,6 +531,110 @@ impl<'a> ArrayWalk<'a> {
         }
     }
 
+    /// Walk the region's outer statement list, where a one-shot temporal
+    /// boundary can be proved. Nested statement lists deliberately use
+    /// `walk_stmts`: admitting a guarded escape inside a loop would let its
+    /// lazy arm reshape the array before a backedge reaches an earlier fact.
+    fn walk_region_stmts(&mut self, stmts: &[Stmt]) {
+        for (index, stmt) in stmts.iter().enumerate() {
+            if self.walk_terminal_array_iteration_guard(stmt, &stmts[index + 1..]) {
+                continue;
+            }
+            self.walk_stmt(stmt);
+        }
+    }
+
+    /// Admit the compiler-generated guarded `for…of` shape without treating
+    /// its one `GetIterator(A)` as an unconditional E3 escape.
+    ///
+    /// A patched iterator is arbitrary code and may transition any element's
+    /// shape. Consequently the exception is temporal, not semantic: the lazy
+    /// and index arms must be the final use of the array and of every producer
+    /// or licensed reader in its element group. Otherwise the whole root is
+    /// disqualified exactly as a normal bare escape would be.
+    fn walk_terminal_array_iteration_guard(&mut self, stmt: &Stmt, following: &[Stmt]) -> bool {
+        let Stmt::If {
+            condition: Expr::ArrayIterationPatched,
+            then_branch,
+            else_branch: Some(index_branch),
+        } = stmt
+        else {
+            return false;
+        };
+        let Some(Stmt::Let {
+            id: iterator_id,
+            init: Some(Expr::GetIterator(source)),
+            ..
+        }) = then_branch.first()
+        else {
+            return false;
+        };
+        let Expr::LocalGet(source_id) = source.as_ref() else {
+            return false;
+        };
+        let Some(root) = self.root_of(*source_id) else {
+            return false;
+        };
+
+        // Preserve the Let write, but exempt exactly its GetIterator source.
+        // Every later lazy-arm statement is walked normally, so a second use
+        // of the array still disqualifies it through the ordinary E3 rules.
+        self.note_write(*iterator_id);
+        self.walk_stmts(&then_branch[1..]);
+        self.walk_stmts(index_branch);
+
+        let array_aliases: HashSet<u32> = self
+            .roots
+            .iter()
+            .filter_map(|(id, candidate_root)| (*candidate_root == root).then_some(*id))
+            .collect();
+        let mut group_members: HashSet<u32> = self
+            .pushes
+            .get(&root)
+            .into_iter()
+            .flatten()
+            .filter_map(|push| match push {
+                PushValue::Local(id) => Some(*id),
+                PushValue::Fresh(_) | PushValue::Other => None,
+            })
+            .chain(
+                self.reads
+                    .iter()
+                    .filter_map(|read| (read.root == root).then_some(read.local)),
+            )
+            .collect();
+        // `ptr_shape` promotes immutable aliases with their root. A use of an
+        // alias after the iterator escape is therefore a use of the same
+        // potentially-reshaped object and must participate in this boundary.
+        loop {
+            let mut changed = false;
+            for (alias, source) in self.alias_edges {
+                if group_members.contains(source) {
+                    changed |= group_members.insert(*alias);
+                }
+            }
+            if !changed {
+                break;
+            }
+        }
+
+        let lazy_refs = local_refs(then_branch);
+        let following_refs = local_refs(following);
+        let lazy_array_uses = lazy_refs
+            .iter()
+            .filter(|id| array_aliases.contains(id))
+            .count();
+        let unsafe_after_escape = lazy_array_uses != 1
+            || lazy_refs.iter().any(|id| group_members.contains(id))
+            || following_refs
+                .iter()
+                .any(|id| array_aliases.contains(id) || group_members.contains(id));
+        if unsafe_after_escape {
+            self.disqualified.insert(root);
+        }
+        true
+    }
+
     fn walk_stmt(&mut self, s: &Stmt) {
         match s {
             Stmt::Let { id, init, .. } => {
@@ -864,6 +979,18 @@ impl<'a> ArrayWalk<'a> {
     }
 }
 
+/// Every local referenced from `stmts`, including id-keyed array operations
+/// and nested closure bodies. Reuse HIR's exhaustive local-id walker rather
+/// than maintaining another list of expression variants in this proof pass.
+fn local_refs(stmts: &[Stmt]) -> Vec<u32> {
+    let mut refs = Vec::new();
+    let mut visited_closures = HashSet::new();
+    for stmt in stmts {
+        perry_hir::collect_local_refs_stmt(stmt, &mut refs, &mut visited_closures);
+    }
+    refs
+}
+
 /// Statement walker over one region's statement tree.
 ///
 /// It does NOT descend into closure bodies — those live inside `Expr`s, not

crates/perry-codegen/src/collectors/ptr_shape_elements_tests.rs

@@ -165,6 +165,41 @@ fn bounded_loop_cond(idx: u32, condition: Expr, body: Vec<Stmt>) -> Stmt {
     }
 }
 
+/// The #7761 guarded lowering of `for (const r of rows)`: the lazy arm starts
+/// with `GetIterator(rows)`, while the byte-identical fast arm aliases `rows`
+/// and performs the ordinary E5 index loop.
+fn guarded_for_of(
+    root: u32,
+    iterator: u32,
+    fast_alias: u32,
+    fast_index: u32,
+    lazy_tail: Vec<Stmt>,
+    fast_body: Vec<Stmt>,
+) -> Stmt {
+    let mut lazy = vec![Stmt::Let {
+        id: iterator,
+        name: format!("__iterator_{iterator}"),
+        ty: Type::Any,
+        mutable: false,
+        init: Some(Expr::GetIterator(Box::new(Expr::LocalGet(root)))),
+    }];
+    lazy.extend(lazy_tail);
+    Stmt::If {
+        condition: Expr::ArrayIterationPatched,
+        then_branch: lazy,
+        else_branch: Some(vec![
+            Stmt::Let {
+                id: fast_alias,
+                name: format!("__arr_{fast_alias}"),
+                ty: Type::Array(Box::new(Type::Named("C".to_string()))),
+                mutable: false,
+                init: Some(Expr::LocalGet(root)),
+            },
+            bounded_loop(fast_index, fast_alias, fast_body),
+        ]),
+    }
+}
+
 /// `const <name> = <arr>[<idx>];`
 fn let_elem(id: u32, name: &str, arr: u32, idx: u32) -> Stmt {
     let_elem_ty(id, name, arr, idx, Type::Named("C".to_string()))
@@ -850,6 +885,137 @@ fn reads_through_an_array_alias_are_licensed() {
     assert!(promoted.contains_key(&2), "and the producer with it");
 }
 
+/// #7777: #7761 wrapped a proven-array `for…of` in a runtime iterator-patch
+/// branch. The lazy `GetIterator(rows)` is an escape, but it is mutually
+/// exclusive with the E5 index arm, and here neither the array nor any group
+/// member is used after the join. The producer, explicit indexed reader, and
+/// guarded fast-arm reader therefore remain the census fixture's three facts.
+///
+/// Sabotage: route the outer statement list through ordinary `walk_stmts`, or
+/// delete `walk_terminal_array_iteration_guard`, and the GetIterator source
+/// voids all three facts exactly as main did in #7777.
+#[test]
+fn terminal_guarded_for_of_retains_the_index_arm_facts() {
+    let c = class_c();
+    let cs = [c];
+    let classes = classes_of(&cs);
+    let stmts = vec![
+        let_arr(1, "rows"),
+        let_c(2, "row"),
+        store_x(2),
+        push(1, Expr::LocalGet(2)),
+        bounded_loop(4, 1, vec![let_elem(5, "indexed", 1, 4), read_x(5)]),
+        guarded_for_of(
+            1,
+            20,
+            30,
+            31,
+            Vec::new(),
+            vec![let_elem(32, "iterated", 30, 31), read_x(32)],
+        ),
+        Stmt::Return(Some(Expr::Number(0.0))),
+    ];
+    let promoted = promote(&stmts, &classes);
+    assert!(promoted.contains_key(&2), "the pushed producer");
+    assert!(promoted.contains_key(&5), "the explicit indexed reader");
+    assert!(
+        promoted.contains_key(&32),
+        "the mutually-exclusive guarded index reader"
+    );
+}
+
+/// A custom iterator receives the actual array and may transition an element
+/// before the branch rejoins. A later indexed reader would then consume a
+/// stale shape proof, even if the iterator patch were restored before that
+/// later loop.
+///
+/// Sabotage: remove the `following_refs` half of the temporal-boundary check
+/// and both guarded and post-join readers promote.
+#[test]
+fn guarded_for_of_does_not_export_facts_across_its_join() {
+    let c = class_c();
+    let cs = [c];
+    let classes = classes_of(&cs);
+    let stmts = vec![
+        let_arr(1, "rows"),
+        let_c(2, "row"),
+        push(1, Expr::LocalGet(2)),
+        guarded_for_of(
+            1,
+            20,
+            30,
+            31,
+            Vec::new(),
+            vec![let_elem(32, "guarded", 30, 31), read_x(32)],
+        ),
+        bounded_loop(40, 1, vec![let_elem(41, "after", 1, 40), read_x(41)]),
+    ];
+    assert!(
+        elements(&stmts, &classes).is_empty(),
+        "GetIterator can reshape an element before the post-join indexed read"
+    );
+    assert!(promote(&stmts, &classes).is_empty());
+}
+
+/// The array can reach the producer object before `GetIterator` returns. A
+/// direct producer/alias access in the lazy arm is therefore just as unsafe as
+/// an array access after the join.
+///
+/// Sabotage: remove the lazy-arm `group_members` intersection and the producer
+/// retains a fact across the arbitrary iterator call.
+#[test]
+fn guarded_for_of_lazy_arm_cannot_reuse_a_group_member() {
+    let c = class_c();
+    let cs = [c];
+    let classes = classes_of(&cs);
+    let stmts = vec![
+        let_arr(1, "rows"),
+        let_c(2, "row"),
+        push(1, Expr::LocalGet(2)),
+        guarded_for_of(
+            1,
+            20,
+            30,
+            31,
+            vec![read_x(2)],
+            vec![let_elem(32, "guarded", 30, 31), read_x(32)],
+        ),
+    ];
+    assert!(elements(&stmts, &classes).is_empty());
+    assert!(promote(&stmts, &classes).is_empty());
+}
+
+/// A guarded escape inside a loop is not terminal: after the lazy arm mutates
+/// an element, the backedge can reach facts from the next iteration. Only the
+/// outer region statement list is eligible for the temporal exception.
+///
+/// Sabotage: invoke the special guard walker from recursive `walk_stmts` and
+/// this nested form starts issuing facts.
+#[test]
+fn a_nested_guarded_for_of_remains_an_escape() {
+    let c = class_c();
+    let cs = [c];
+    let classes = classes_of(&cs);
+    let stmts = vec![
+        let_arr(1, "rows"),
+        let_c(2, "row"),
+        push(1, Expr::LocalGet(2)),
+        Stmt::While {
+            condition: Expr::Bool(true),
+            body: vec![guarded_for_of(
+                1,
+                20,
+                30,
+                31,
+                Vec::new(),
+                vec![let_elem(32, "guarded", 30, 31), read_x(32)],
+            )],
+        },
+    ];
+    assert!(elements(&stmts, &classes).is_empty());
+    assert!(promote(&stmts, &classes).is_empty());
+}
+
 // ── CodeRabbit review reproducers (PR #7149) ───────────────────────────────
 
 /// **CodeRabbit 🔴 #1** (`ptr_shape_elements.rs:710`, review of `816a5a3`):