Loop Contracts Annotation for While-Loop

cprover_bindings/src/goto_program/stmt.rs

@@ -82,7 +82,10 @@ pub enum StmtBody {
         arguments: Vec<Expr>,
     },
     /// `goto dest;`
-    Goto(InternedString),
+    Goto {
+        dest: InternedString,
+        loop_invariants: Option<Expr>,
+    },
     /// `if (i) { t } else { e }`
     Ifthenelse {
         i: Expr,
@@ -179,7 +182,7 @@ impl Stmt {
         stmt!(Assign { lhs, rhs }, loc)
     }
 
-    /// `assert(cond, property_class, commment);`
+    /// `assert(cond, property_class, comment);`
     pub fn assert(cond: Expr, property_name: &str, message: &str, loc: Location) -> Self {
         assert!(cond.typ().is_bool());
         assert!(!property_name.is_empty() && !message.is_empty());
@@ -188,7 +191,7 @@ impl Stmt {
         let loc_with_property =
             Location::create_location_with_property(message, property_name, loc);
 
-        // Chose InternedString to seperate out codegen from the cprover_bindings logic
+        // Chose InternedString to separate out codegen from the cprover_bindings logic
         let property_class = property_name.intern();
         let msg = message.into();
 
@@ -283,7 +286,18 @@ impl Stmt {
     pub fn goto<T: Into<InternedString>>(dest: T, loc: Location) -> Self {
         let dest = dest.into();
         assert!(!dest.is_empty());
-        stmt!(Goto(dest), loc)
+        stmt!(Goto { dest, loop_invariants: None }, loc)
+    }
+
+    /// `goto dest;` with loop invariant
+    pub fn goto_with_loop_inv<T: Into<InternedString>>(
+        dest: T,
+        loop_invariants: Expr,
+        loc: Location,
+    ) -> Self {
+        let dest = dest.into();
+        assert!(!dest.is_empty());
+        stmt!(Goto { dest, loop_invariants: Some(loop_invariants) }, loc)
     }
 
     /// `if (i) { t } else { e }` or `if (i) { t }`

cprover_bindings/src/irep/to_irep.rs

@@ -516,8 +516,18 @@ impl ToIrep for StmtBody {
                     arguments_irep(arguments.iter(), mm),
                 ],
             ),
-            StmtBody::Goto(dest) => code_irep(IrepId::Goto, vec![])
-                .with_named_sub(IrepId::Destination, Irep::just_string_id(dest.to_string())),
+            StmtBody::Goto { dest, loop_invariants } => {
+                let stmt_goto = code_irep(IrepId::Goto, vec![])
+                    .with_named_sub(IrepId::Destination, Irep::just_string_id(dest.to_string()));
+                if loop_invariants.is_some() {
+                    stmt_goto.with_named_sub(
+                        IrepId::CSpecLoopInvariant,
+                        loop_invariants.clone().unwrap().to_irep(mm),
+                    )
+                } else {
+                    stmt_goto
+                }
+            }
             StmtBody::Ifthenelse { i, t, e } => code_irep(
                 IrepId::Ifthenelse,
                 vec![

kani-compiler/src/codegen_cprover_gotoc/codegen/block.rs

@@ -21,12 +21,23 @@ impl<'tcx> GotocCtx<'tcx> {
         debug!(?bb, "codegen_block");
         let label = bb_label(bb);
         let check_coverage = self.queries.args().check_coverage;
+
+        // record the seen bbidx if loop contracts enabled
+        if self.loop_contracts_ctx.loop_contracts_enabled() {
+            self.loop_contracts_ctx.add_new_seen_bbidx(bb);
+        }
+
         // the first statement should be labelled. if there is no statements, then the
         // terminator should be labelled.
         match bbd.statements.len() {
             0 => {
                 let term = &bbd.terminator;
-                let tcode = self.codegen_terminator(term);
+                let tcode = if self.loop_contracts_ctx.loop_contracts_enabled() {
+                    let codegen_result = self.codegen_terminator(term);
+                    self.loop_contracts_ctx.push_onto_block(codegen_result)
+                } else {
+                    self.codegen_terminator(term)
+                };
                 // When checking coverage, the `coverage` check should be
                 // labelled instead.
                 if check_coverage {
@@ -40,7 +51,12 @@ impl<'tcx> GotocCtx<'tcx> {
             }
             _ => {
                 let stmt = &bbd.statements[0];
-                let scode = self.codegen_statement(stmt);
+                let scode = if self.loop_contracts_ctx.loop_contracts_enabled() {
+                    let codegen_result = self.codegen_statement(stmt);
+                    self.loop_contracts_ctx.push_onto_block(codegen_result)
+                } else {
+                    self.codegen_statement(stmt)
+                };
                 // When checking coverage, the `coverage` check should be
                 // labelled instead.
                 if check_coverage {
@@ -58,7 +74,12 @@ impl<'tcx> GotocCtx<'tcx> {
                         let cover = self.codegen_coverage(span);
                         self.current_fn_mut().push_onto_block(cover);
                     }
-                    let stmt = self.codegen_statement(s);
+                    let stmt = if self.loop_contracts_ctx.loop_contracts_enabled() {
+                        let codegen_result = self.codegen_statement(s);
+                        self.loop_contracts_ctx.push_onto_block(codegen_result)
+                    } else {
+                        self.codegen_statement(s)
+                    };
                     self.current_fn_mut().push_onto_block(stmt);
                 }
                 let term = &bbd.terminator;
@@ -67,7 +88,12 @@ impl<'tcx> GotocCtx<'tcx> {
                     let cover = self.codegen_coverage(span);
                     self.current_fn_mut().push_onto_block(cover);
                 }
-                let tcode = self.codegen_terminator(term);
+                let tcode = if self.loop_contracts_ctx.loop_contracts_enabled() {
+                    let codegen_result = self.codegen_terminator(term);
+                    self.loop_contracts_ctx.push_onto_block(codegen_result)
+                } else {
+                    self.codegen_terminator(term)
+                };
                 self.current_fn_mut().push_onto_block(tcode);
             }
         }

kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs

@@ -56,6 +56,9 @@ impl<'tcx> GotocCtx<'tcx> {
         let old_sym = self.symbol_table.lookup(&name).unwrap();
 
         let _trace_span = debug_span!("CodegenFunction", name = instance.name()).entered();
+        if self.loop_contracts_ctx.loop_contracts_enabled() {
+            self.loop_contracts_ctx.enter_new_function();
+        }
         if old_sym.is_function_definition() {
             debug!("Double codegen of {:?}", old_sym);
         } else {

kani-compiler/src/codegen_cprover_gotoc/codegen/statement.rs

@@ -175,12 +175,24 @@ impl<'tcx> GotocCtx<'tcx> {
     ///
     /// See also [`GotocCtx::codegen_statement`] for ordinary [Statement]s.
     pub fn codegen_terminator(&mut self, term: &Terminator) -> Stmt {
-        let loc = self.codegen_span_stable(term.span);
+        let loc: Location = self.codegen_span_stable(term.span);
         let _trace_span = debug_span!("CodegenTerminator", statement = ?term.kind).entered();
         debug!("handling terminator {:?}", term);
         //TODO: Instead of doing location::none(), and updating, just putit in when we make the stmt.
         match &term.kind {
-            TerminatorKind::Goto { target } => Stmt::goto(bb_label(*target), loc),
+            TerminatorKind::Goto { target } => {
+                if self.loop_contracts_ctx.loop_contracts_enabled()
+                    && self.loop_contracts_ctx.is_loop_latch(target)
+                {
+                    Stmt::goto_with_loop_inv(
+                        bb_label(*target),
+                        self.loop_contracts_ctx.extract_block(loc),
+                        loc,
+                    )
+                } else {
+                    Stmt::goto(bb_label(*target), loc)
+                }
+            }
             TerminatorKind::SwitchInt { discr, targets } => {
                 self.codegen_switch_int(discr, targets, loc)
             }

kani-compiler/src/codegen_cprover_gotoc/context/goto_ctx.rs

@@ -15,6 +15,7 @@
 //! Any MIR specific functionality (e.g. codegen etc) should live in specialized files that use
 //! this structure as input.
 use super::current_fn::CurrentFnCtx;
+use super::loop_contracts_ctx::LoopContractsCtx;
 use super::vtable_ctx::VtableCtx;
 use crate::codegen_cprover_gotoc::overrides::{fn_hooks, GotocHooks};
 use crate::codegen_cprover_gotoc::utils::full_crate_name;
@@ -74,6 +75,8 @@ pub struct GotocCtx<'tcx> {
     pub concurrent_constructs: UnsupportedConstructs,
     /// The body transformation agent.
     pub transformer: BodyTransformation,
+    /// The context for loop contracts code generation.
+    pub loop_contracts_ctx: LoopContractsCtx,
 }
 
 /// Constructor
@@ -87,6 +90,8 @@ impl<'tcx> GotocCtx<'tcx> {
         let fhks = fn_hooks();
         let symbol_table = SymbolTable::new(machine_model.clone());
         let emit_vtable_restrictions = queries.args().emit_vtable_restrictions;
+        let loop_contracts_enabled =
+            queries.args().unstable_features.contains(&"loop-contracts".to_string());
         GotocCtx {
             tcx,
             queries,
@@ -103,6 +108,7 @@ impl<'tcx> GotocCtx<'tcx> {
             unsupported_constructs: FxHashMap::default(),
             concurrent_constructs: FxHashMap::default(),
             transformer,
+            loop_contracts_ctx: LoopContractsCtx::new(loop_contracts_enabled),
         }
     }
 }

kani-compiler/src/codegen_cprover_gotoc/context/loop_contracts_ctx.rs

@@ -0,0 +1,140 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+use crate::codegen_cprover_gotoc::codegen::bb_label;
+use cbmc::goto_program::{CIntType, Expr, Location, Stmt, StmtBody, Type};
+use stable_mir::mir::BasicBlockIdx;
+use std::collections::HashSet;
+
+pub struct LoopContractsCtx {
+    /// the GOTO block compiled from the corresponding loop invariants
+    invariants_block: Vec<Stmt>,
+    /// Which codegen state
+    stage: LoopContractsStage,
+    /// If enable loop contracts
+    loop_contracts_enabled: bool,
+    /// Seen basic block indexes. Used to decide if a jump is backward
+    seen_bbidx: HashSet<BasicBlockIdx>,
+    /// Current unused bbidx label
+    current_bbidx_label: Option<String>,
+    /// The lhs of evaluation of the loop invariant
+    loop_invariant_lhs: Option<Stmt>,
+}
+
+/// We define two states:
+/// 1. loop invariants block
+///     In this state, we push all codegen stmts into the invariant block.
+///     We enter this state when codegen for `KaniLoopInvariantBegin`.
+///     We exit this state when codegen for `KaniLoopInvariantEnd`.
+/// 2. loop latch block
+///     In this state, we codegen a statement expression from the
+///     invariant_block annotate the statement expression to the named sub
+///     of the next backward jumping we codegen.
+///     We enter this state when codegen for `KaniLoopInvariantEnd`.
+///     We exit this state when codegen for the first backward jumping.
+#[allow(dead_code)]
+#[derive(Debug, PartialEq)]
+enum LoopContractsStage {
+    /// Codegen for user code as usual
+    UserCode,
+    /// Codegen for loop invariants
+    InvariantBlock,
+    /// Codegen for loop latch node
+    FindingLatchNode,
+}
+
+/// Constructor
+impl LoopContractsCtx {
+    pub fn new(loop_contracts_enabled: bool) -> Self {
+        Self {
+            invariants_block: Vec::new(),
+            stage: LoopContractsStage::UserCode,
+            loop_contracts_enabled,
+            seen_bbidx: HashSet::new(),
+            current_bbidx_label: None,
+            loop_invariant_lhs: None,
+        }
+    }
+}
+
+/// Getters
+impl LoopContractsCtx {
+    pub fn loop_contracts_enabled(&self) -> bool {
+        self.loop_contracts_enabled
+    }
+
+    /// decide if a GOTO with `target` is backward jump
+    pub fn is_loop_latch(&self, target: &BasicBlockIdx) -> bool {
+        self.stage == LoopContractsStage::FindingLatchNode && self.seen_bbidx.contains(target)
+    }
+}
+
+/// Setters
+impl LoopContractsCtx {
+    /// Returns the current block as a statement expression.
+    /// Exit loop latch block.
+    pub fn extract_block(&mut self, loc: Location) -> Expr {
+        assert!(self.loop_invariant_lhs.is_some());
+        self.stage = LoopContractsStage::UserCode;
+        self.invariants_block.push(self.loop_invariant_lhs.as_ref().unwrap().clone());
+
+        // The first statement is the GOTO in the rhs of __kani_loop_invariant_begin()
+        // Ignore it
+        self.invariants_block.remove(0);
+
+        Expr::statement_expression(
+            std::mem::take(&mut self.invariants_block),
+            Type::CInteger(CIntType::Bool),
+            loc,
+        )
+        .cast_to(Type::bool())
+        .and(Expr::bool_true())
+    }
+
+    /// Push the `s` onto the block if it is in the loop invariant block
+    /// and return `skip`. Otherwise, do nothing and return `s`.
+    pub fn push_onto_block(&mut self, s: Stmt) -> Stmt {
+        if self.stage == LoopContractsStage::InvariantBlock {
+            // Attach the lable to the first Stmt in that block and reset it.
+            let to_push = if self.current_bbidx_label.is_none() {
+                s.clone()
+            } else {
+                s.clone().with_label(self.current_bbidx_label.clone().unwrap())
+            };
+            self.current_bbidx_label = None;
+
+            match s.body() {
+                StmtBody::Assign { lhs, rhs: _ } => {
+                    let lhs_stmt = lhs.clone().as_stmt(*s.location());
+                    self.loop_invariant_lhs = Some(lhs_stmt.clone());
+                    self.invariants_block.push(to_push);
+                }
+                _ => {
+                    self.invariants_block.push(to_push);
+                }
+            };
+            Stmt::skip(*s.location())
+        } else {
+            s
+        }
+    }
+
+    pub fn enter_loop_invariant_block(&mut self) {
+        assert!(self.invariants_block.is_empty());
+        self.stage = LoopContractsStage::InvariantBlock;
+    }
+
+    pub fn exit_loop_invariant_block(&mut self) {
+        self.stage = LoopContractsStage::FindingLatchNode;
+    }
+
+    /// Enter a new function, reset the seen_bbidx set
+    pub fn enter_new_function(&mut self) {
+        self.seen_bbidx = HashSet::new()
+    }
+
+    pub fn add_new_seen_bbidx(&mut self, bbidx: BasicBlockIdx) {
+        self.seen_bbidx.insert(bbidx);
+        self.current_bbidx_label = Some(bb_label(bbidx));
+    }
+}

kani-compiler/src/codegen_cprover_gotoc/context/mod.rs

@@ -8,6 +8,7 @@
 
 mod current_fn;
 mod goto_ctx;
+mod loop_contracts_ctx;
 mod vtable_ctx;
 
 pub use goto_ctx::GotocCtx;

kani-compiler/src/codegen_cprover_gotoc/overrides/hooks.rs

@@ -8,6 +8,7 @@
 //! It would be too nasty if we spread around these sort of undocumented hooks in place, so
 //! this module addresses this issue.
 
+use super::loop_contracts_hooks::{LoopInvariantBegin, LoopInvariantEnd};
 use crate::codegen_cprover_gotoc::codegen::{bb_label, PropertyClass};
 use crate::codegen_cprover_gotoc::GotocCtx;
 use crate::kani_middle::attributes::matches_diagnostic as matches_function;
@@ -555,6 +556,8 @@ pub fn fn_hooks() -> GotocHooks {
             Rc::new(MemCmp),
             Rc::new(UntrackedDeref),
             Rc::new(InitContracts),
+            Rc::new(LoopInvariantBegin),
+            Rc::new(LoopInvariantEnd),
         ],
     }
 }