From 8dc17959f25325028c8340c5c0d82a25f32232b7 Mon Sep 17 00:00:00 2001
From: Li-yao Xia <lysxia@gmail.com>
Date: Wed, 20 Nov 2024 16:41:49 +0100
Subject: [PATCH] Move loop invariants

---
 creusot-contracts-proc/src/invariant.rs | 36 +++++++++++++++----------
 creusot/src/gather_spec_closures.rs     | 26 +-----------------
 2 files changed, 23 insertions(+), 39 deletions(-)

diff --git a/creusot-contracts-proc/src/invariant.rs b/creusot-contracts-proc/src/invariant.rs
index 909b7b2504..7aec4ed786 100644
--- a/creusot-contracts-proc/src/invariant.rs
+++ b/creusot-contracts-proc/src/invariant.rs
@@ -3,7 +3,7 @@ use proc_macro2::{Span, TokenStream};
 use quote::{quote_spanned, ToTokens};
 use syn::{
     parse_quote_spanned, spanned::Spanned, AttrStyle, Attribute, Error, Expr, ExprForLoop,
-    ExprLoop, ExprWhile, Ident, Meta, Result,
+    ExprLoop, ExprWhile, Ident, Meta, Result, Stmt,
 };
 
 #[derive(Debug)]
@@ -99,23 +99,31 @@ pub fn lower(loop_: Loop) -> TokenStream {
     let invariants = loop_.invariants;
     match loop_.kind {
         LoopKind::For(floop) => desugar_for(invariants, floop),
-        LoopKind::While(l) => {
-            let mut tokens = TokenStream::new();
-            for i in invariants {
-                i.to_tokens(&mut tokens);
-            }
+        LoopKind::While(mut l) => {
             let sp = loop_.span;
-            quote_spanned! {sp=>{
-                #tokens
+            let cond = l.cond;
+            l.cond = Box::new(Expr::Verbatim(quote_spanned! {sp=> {
+                #(#invariants)*
+                #cond
+            }}));
+            quote_spanned! {sp=>
                 #l
-            }}
+            }
         }
-        LoopKind::Loop(l) => {
+        LoopKind::Loop(mut l) => {
             let sp = loop_.span;
-            quote_spanned! {sp=> {
-                #(#invariants)*
+            l.body.stmts.insert(
+                0,
+                Stmt::Expr(
+                    Expr::Verbatim(parse_quote_spanned! {sp=> {
+                        #(#invariants)*
+                    }}),
+                    None,
+                ),
+            );
+            quote_spanned! {sp=>
                 #l
-            }}
+            }
         }
     }
 }
@@ -172,10 +180,10 @@ fn desugar_for(mut invariants: Vec<Invariant>, f: ExprForLoop) -> TokenStream {
         let mut #it = ::std::iter::IntoIterator::into_iter(#iter);
         let #iter_old = snapshot! { #it };
         let mut #produced = snapshot! { ::creusot_contracts::logic::Seq::EMPTY };
-        #(#invariants)*
         #(#outer)*
         #lbl
         loop {
+            #(#invariants)*
             #(#inner)*
             match ::std::iter::Iterator::next(&mut #it) {
                 Some(#elem) => {
diff --git a/creusot/src/gather_spec_closures.rs b/creusot/src/gather_spec_closures.rs
index a946ed15fa..2ee14a7761 100644
--- a/creusot/src/gather_spec_closures.rs
+++ b/creusot/src/gather_spec_closures.rs
@@ -115,31 +115,7 @@ pub(crate) fn corrected_invariant_names_and_locations<'tcx>(
     invs_gather.visit_body(body);
 
     for (loc, (clos, kind)) in invs_gather.invariants.into_iter() {
-        let mut target: BasicBlock = loc.block;
-
-        loop {
-            let mut succs = body.basic_blocks.successors(target);
-
-            target = succs.next().unwrap();
-
-            // Check if `taget_block` is a loop header by testing if it dominates
-            // one of its predecessors.
-            if let Some(preds) = body.basic_blocks.predecessors().get(target) {
-                let is_loop_header = preds
-                    .iter()
-                    .any(|pred| body.basic_blocks.dominators().dominates(target, *pred));
-
-                if is_loop_header {
-                    break;
-                }
-            };
-
-            // If we've hit a switch then stop trying to push the invariants down.
-            if body[target].terminator().kind.as_switch().is_some() {
-                panic!("Could not find loop header")
-            }
-        }
-
+        let target: BasicBlock = loc.block;
         let term = ctx.term(clos).unwrap().clone();
         results.entry(target).or_insert_with(Vec::new).push((kind, term));
     }