Copy a loop-body local to the loop head only when a nested loop needs it and it can be moved

kani-compiler/src/kani_middle/transform/loop_contracts.rs

@@ -978,13 +978,16 @@ impl LoopContractPass {
         let mut found_local_list = found_local_list;
         let localvars = self.get_storage_moving_variables(body);
         let forloopvars = self.get_kaniiter_variables(body);
+        let firstpat_vars: Vec<usize> =
+            self.get_first_pats_and_nth_pats(body).iter().map(|(first, ..)| *first).collect();
         let mut current_user_local = 0;
         // The loop head of the block that declares `current_user_local`.
         let mut current_user_local_loop_head = 0;
         let mut current_local_decl_blocks: Vec<BasicBlock> = Vec::new();
         let mut move_call_list: Vec<(usize, Vec<BasicBlock>)> = Vec::new();
         let mut moved_locals: Vec<(usize, usize, Vec<BasicBlock>)> = Vec::new();
         let mut kaniiter_blocks: Vec<usize> = Vec::new();
+        let mut moved_kaniiter: HashSet<usize> = HashSet::new();
         for (block_idx, block) in body.blocks().iter().enumerate() {
             let mut decl_current_user_local = false;
             let mut storage_live_block_stmt: Vec<Statement> = Vec::new();
@@ -1023,23 +1026,62 @@ impl LoopContractPass {
                 && dest.local == current_user_local
                 && current_user_local != 0
             {
-                // Only a local that is declared in the body of this loop only exists for the
-                // loop (a local declared in an outer loop, e.g. `let x;`, can be named in the
-                // clause of this loop).
-                if current_user_local_loop_head == closest_loop_head {
-                    moved_locals.push((
-                        closest_loop_head,
-                        current_user_local,
-                        current_local_decl_blocks.clone(),
-                    ));
+                // The splice helpers re-target each moved block's terminator via
+                // `get_mut_target_ref`, which only supports terminators with a single
+                // `target` (`Call` with a return target, `Goto`, `Assert`, `Drop`).
+                // A group spanning any other terminator (e.g. the `SwitchInt` of a
+                // branching initializer) cannot be moved; before this guard,
+                // attempting to move it panicked there.
+                let retargetable = current_local_decl_blocks.iter().all(|b| {
+                    matches!(
+                        b.terminator.kind,
+                        TerminatorKind::Call { target: Some(_), .. }
+                            | TerminatorKind::Goto { .. }
+                            | TerminatorKind::Assert { .. }
+                            | TerminatorKind::Drop { .. }
+                    )
+                });
+                // The for-loop machinery (kaniiter, firstpat tuples) keeps its
+                // move unconditionally. A nested `for` whose range reads
+                // outer-loop state computes its iterator once from the outer
+                // variable's value at entry.
+                let machinery = forloopvars.contains(&current_user_local)
+                    || firstpat_vars.contains(&current_user_local);
+                // The copy exists only so CBMC can infer an inner loop's assigns
+                // clause, so copy iff an inner loop writes the local.
+                let written_in_inner_loop = body.blocks().iter().enumerate().any(|(bb, blk)| {
+                    loop_head_map.get(&bb).is_some_and(|head| *head != closest_loop_head)
+                        && (blk.statements.iter().any(|s| {
+                            matches!(&s.kind, StatementKind::Assign(p, _) if p.local == current_user_local)
+                        }) || matches!(&blk.terminator.kind,
+                            TerminatorKind::Call { destination, .. } if destination.local == current_user_local))
+                });
+                if retargetable && (machinery || written_in_inner_loop) {
+                    // Only a local that is declared in the body of this loop only exists for the
+                    // loop (a local declared in an outer loop, e.g. `let x;`, can be named in the
+                    // clause of this loop). Recorded only for groups actually moved, since an
+                    // unmoved group is never assigned at the loop head.
+                    if current_user_local_loop_head == closest_loop_head {
+                        moved_locals.push((
+                            closest_loop_head,
+                            current_user_local,
+                            current_local_decl_blocks.clone(),
+                        ));
+                    }
+                    move_call_list.push((closest_loop_head, current_local_decl_blocks.clone()));
+                    // Record only the kaniiter groups actually moved, so the goto below
+                    // leaves an unmoved `kani_into_iter` in place.
+                    if machinery {
+                        moved_kaniiter.insert(current_user_local);
+                    }
                 }
-                move_call_list.push((closest_loop_head, current_local_decl_blocks.clone()));
                 current_local_decl_blocks = Vec::new();
                 current_user_local = 0;
             }
 
             if let TerminatorKind::Call { destination: dest, .. } = terminatorkind
                 && forloopvars.contains(&dest.local)
+                && moved_kaniiter.contains(&dest.local)
             {
                 kaniiter_blocks.push(block_idx);
             }

tests/expected/loop-contract/loop_head_call_computed_index.expected

@@ -0,0 +1,10 @@
+	 - Status: SUCCESS\
+	 - Description: "Check invariant before entry for loop search::<{closure@
+
+	 - Status: SUCCESS\
+	 - Description: "Check invariant after step for loop search::<{closure@
+
+<usize as std::slice::SliceIndex<[u8]>>::get_unchecked.assume.1\
+	 - Status: SUCCESS
+
+Complete - 3 successfully verified harnesses, 0 failures, 3 total.

tests/expected/loop-contract/loop_head_call_computed_index.rs

@@ -0,0 +1,91 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! A call in the loop body whose index is computed in the body (a binary-search-shaped loop
+//! passing `s.get_unchecked(mid)` into a closure) must not have its initialization copied to
+//! the loop head: at the head the computed index does not exist yet and `get_unchecked`'s
+//! precondition check fails spuriously. The two sibling harnesses pin the shapes that already
+//! worked: an element discarded with no binding, and the value copied out before the call.
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+use std::cmp::Ordering;
+
+// The #4911 shape: computed index, element ref flows into the closure.
+fn search<F: FnMut(&u8) -> Ordering>(s: &[u8], mut f: F) {
+    let len = s.len();
+    if len == 0 {
+        return;
+    }
+    let mut size = len;
+    let mut base = 0usize;
+    #[kani::loop_invariant(1 <= size && size <= len && base <= len - size)]
+    while size > 1 {
+        let half = size / 2;
+        let mid = base + half;
+        let cmp = f(unsafe { s.get_unchecked(mid) });
+        base = if cmp == Ordering::Greater { base } else { mid };
+        size -= half;
+    }
+}
+
+// Sibling pin 1: element discarded, no user binding starts a group.
+fn search_no_closure(s: &[u8]) {
+    let len = s.len();
+    if len == 0 {
+        return;
+    }
+    let mut size = len;
+    let mut base = 0usize;
+    #[kani::loop_invariant(1 <= size && size <= len && base <= len - size)]
+    while size > 1 {
+        let half = size / 2;
+        let mid = base + half;
+        let _ = unsafe { s.get_unchecked(mid) };
+        base = if kani::any() { base } else { mid };
+        size -= half;
+    }
+}
+
+// Sibling pin 2: value copied out before the closure call.
+fn search_copyval<F: FnMut(&u8) -> Ordering>(s: &[u8], mut f: F) {
+    let len = s.len();
+    if len == 0 {
+        return;
+    }
+    let mut size = len;
+    let mut base = 0usize;
+    #[kani::loop_invariant(1 <= size && size <= len && base <= len - size)]
+    while size > 1 {
+        let half = size / 2;
+        let mid = base + half;
+        let v = unsafe { *s.get_unchecked(mid) };
+        let cmp = f(&v);
+        base = if cmp == Ordering::Greater { base } else { mid };
+        size -= half;
+    }
+}
+
+#[kani::proof]
+fn check_with_closure() {
+    let a: [u8; 8] = kani::any();
+    let s: &[u8] = kani::slice::any_slice_of_array(&a);
+    search(s, |_| if kani::any() { Ordering::Less } else { Ordering::Greater });
+}
+
+#[kani::proof]
+fn check_no_closure() {
+    let a: [u8; 8] = kani::any();
+    let s: &[u8] = kani::slice::any_slice_of_array(&a);
+    search_no_closure(s);
+}
+
+#[kani::proof]
+fn check_copyval() {
+    let a: [u8; 8] = kani::any();
+    let s: &[u8] = kani::slice::any_slice_of_array(&a);
+    search_copyval(s, |_| if kani::any() { Ordering::Less } else { Ordering::Greater });
+}

tests/expected/loop-contract/loop_head_call_computed_index_fail.expected

@@ -0,0 +1,3 @@
+Rust intrinsic assumption failed
+
+VERIFICATION:- FAILED

tests/expected/loop-contract/loop_head_call_computed_index_fail.rs

@@ -0,0 +1,38 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Negative twin of loop_head_call_computed_index.rs: with the spurious loop-head copy gone, a
+//! genuinely out-of-bounds `get_unchecked` in the loop body must still fail its precondition
+//! check at the real call site.
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+use std::cmp::Ordering;
+
+fn search_oob<F: FnMut(&u8) -> Ordering>(s: &[u8], mut f: F) {
+    let len = s.len();
+    if len == 0 {
+        return;
+    }
+    let mut size = len;
+    let mut base = 0usize;
+    #[kani::loop_invariant(1 <= size && size <= len && base <= len - size)]
+    while size > 1 {
+        let half = size / 2;
+        let mid = base + half;
+        // mid + 1 can equal len: a real precondition violation at the body call site.
+        let cmp = f(unsafe { s.get_unchecked(mid + 1) });
+        base = if cmp == Ordering::Greater { base } else { mid };
+        size -= half;
+    }
+}
+
+#[kani::proof]
+fn check_oob_still_fails() {
+    let a: [u8; 8] = kani::any();
+    let s: &[u8] = kani::slice::any_slice_of_array(&a);
+    search_oob(s, |_| if kani::any() { Ordering::Less } else { Ordering::Greater });
+}

tests/expected/loop-contract/loop_local_var_branching_init.expected

@@ -0,0 +1 @@
+VERIFICATION:- SUCCESSFUL

tests/expected/loop-contract/loop_local_var_branching_init.rs

@@ -0,0 +1,30 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Regression: moving a loop-local's declaration blocks to the loop head panicked with
+//! "Kani can only insert instructions after terminators that have a `target` field"
+//! when the local is initialized through branching control flow (the collected init
+//! group spans a `SwitchInt`). Such groups are now left in place.
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+fn pick(x: u32) -> u32 {
+    x + 1
+}
+
+#[kani::proof]
+fn loop_local_var_branching_init() {
+    let n: u32 = kani::any_where(|x| *x <= 8);
+    let mut i: u32 = 0;
+    #[kani::loop_invariant(i <= n)]
+    while i < n {
+        // Both arms yield 1 so the loop stays provable; the `if` is only here to
+        // force a branching (SwitchInt) initializer for `step`.
+        let step = if i % 2 == 0 { pick(0) } else { pick(0) };
+        i += step;
+    }
+    assert!(i == n);
+}

tests/expected/loop-contract/loop_local_var_branching_init_external_cond.expected

@@ -0,0 +1 @@
+VERIFICATION:- SUCCESSFUL

tests/expected/loop-contract/loop_local_var_branching_init_external_cond.rs

@@ -0,0 +1,31 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Same panic as loop_local_var_branching_init, with the branch condition reading
+//! a variable declared before the loop. The collected init group mentions no loop
+//! local, so the reads-loop-locals skip does not apply; only the terminator-shape
+//! check keeps the group from being moved and panicking.
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+fn pick(x: u32) -> u32 {
+    x + 1
+}
+
+#[kani::proof]
+fn loop_local_var_branching_init_external_cond() {
+    let gate: bool = kani::any();
+    let n: u32 = kani::any_where(|x| *x <= 8);
+    let mut i: u32 = 0;
+    #[kani::loop_invariant(i <= n)]
+    while i < n {
+        // Both arms yield 1 so the loop stays provable; the `if` is only here to
+        // force a branching (SwitchInt) initializer for `step`.
+        let step = if gate { pick(0) } else { pick(0) };
+        i += step;
+    }
+    assert!(i == n);
+}

tests/expected/loop-contract/nested_loop_local_var_func_call.expected

@@ -19,3 +19,5 @@ main.assertion.6\
 	 - Description: "assertion failed: s == 100"
 
 VERIFICATION:- SUCCESSFUL
+
+Complete - 6 successfully verified harnesses, 0 failures, 6 total.