Loop Contracts Annotation for While-Loop

cprover_bindings/src/goto_program/stmt.rs

@@ -82,7 +82,12 @@ pub enum StmtBody {
         arguments: Vec<Expr>,
     },
     /// `goto dest;`
-    Goto(InternedString),
+    Goto {
+        dest: InternedString,
+        // The loop invariants annotated to the goto, which can be
+        // applied as loop contracts in CBMC if it is a backward goto.
+        loop_invariants: Option<Expr>,
+    },
     /// `if (i) { t } else { e }`
     Ifthenelse {
         i: Expr,
@@ -179,7 +184,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 +193,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 +288,7 @@ 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)
     }
 
     /// `if (i) { t } else { e }` or `if (i) { t }`
@@ -325,6 +330,16 @@ impl Stmt {
         assert!(!label.is_empty());
         stmt!(Label { label, body: self }, *self.location())
     }
+
+    /// `goto dest;` with loop invariant
+    pub fn with_loop_contracts(self, inv: Expr) -> Self {
+        if let Goto { dest, loop_invariants } = self.body() {
+            assert!(loop_invariants.is_none());
+            stmt!(Goto { dest: *dest, loop_invariants: Some(inv) }, *self.location())
+        } else {
+            unreachable!("Loop contracts should be annotated only to goto stmt")
+        }
+    }
 }
 
 /// Predicates

cprover_bindings/src/irep/to_irep.rs

@@ -520,8 +520,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 let Some(inv) = loop_invariants {
+                    stmt_goto.with_named_sub(
+                        IrepId::CSpecLoopInvariant,
+                        inv.clone().and(Expr::bool_true()).to_irep(mm),
+                    )
+                } else {
+                    stmt_goto
+                }
+            }
             StmtBody::Ifthenelse { i, t, e } => code_irep(IrepId::Ifthenelse, vec![
                 i.to_irep(mm),
                 t.to_irep(mm),

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

@@ -10,11 +10,14 @@
 
 use crate::codegen_cprover_gotoc::GotocCtx;
 use crate::codegen_cprover_gotoc::codegen::{PropertyClass, bb_label};
+use crate::kani_middle::attributes::KaniAttributes;
 use crate::kani_middle::attributes::matches_diagnostic as matches_function;
 use crate::unwrap_or_return_codegen_unimplemented_stmt;
+use cbmc::goto_program::CIntType;
 use cbmc::goto_program::{BuiltinFn, Expr, Stmt, Type};
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
+use rustc_span::Symbol;
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::{BasicBlockIdx, Place};
 use stable_mir::{CrateDef, ty::Span};
@@ -539,6 +542,43 @@ impl GotocHook for InitContracts {
     }
 }
 
+/// A loop contract register function call is assumed to be
+/// 1. of form `kani_register_loop_contract(inv)` where `inv`
+///    is the closure wrapping loop invariants
+/// 2. is the last statement in some loop, so that its `target`` is
+///    the head of the loop
+///
+/// Such call will be translate to
+/// ```c
+/// goto target
+/// ```
+/// with loop invariants (call to the register function) annotated as
+/// a named sub of the `goto`.
+pub struct LoopInvariantRegister;
+
+impl GotocHook for LoopInvariantRegister {
+    fn hook_applies(&self, tcx: TyCtxt, instance: Instance) -> bool {
+        KaniAttributes::for_instance(tcx, instance).fn_marker()
+            == Some(Symbol::intern("kani_register_loop_contract"))
+    }
+
+    fn handle(
+        &self,
+        gcx: &mut GotocCtx,
+        instance: Instance,
+        fargs: Vec<Expr>,
+        _assign_to: &Place,
+        target: Option<BasicBlockIdx>,
+        span: Span,
+    ) -> Stmt {
+        let loc = gcx.codegen_span_stable(span);
+        let func_exp = gcx.codegen_func_expr(instance, loc);
+
+        Stmt::goto(bb_label(target.unwrap()), loc)
+            .with_loop_contracts(func_exp.call(fargs).cast_to(Type::CInteger(CIntType::Bool)))
+    }
+}
+
 pub fn fn_hooks() -> GotocHooks {
     GotocHooks {
         hooks: vec![
@@ -555,6 +595,7 @@ pub fn fn_hooks() -> GotocHooks {
             Rc::new(MemCmp),
             Rc::new(UntrackedDeref),
             Rc::new(InitContracts),
+            Rc::new(LoopInvariantRegister),
         ],
     }
 }

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

@@ -59,6 +59,11 @@ impl MutableBody {
         self.arg_count
     }
 
+    #[allow(dead_code)]
+    pub fn var_debug_info(&self) -> &Vec<VarDebugInfo> {
+        &self.var_debug_info
+    }
+
     /// Create a mutable body from the original MIR body.
     pub fn from(body: Body) -> Self {
         MutableBody {
@@ -428,6 +433,15 @@ impl MutableBody {
         self.blocks.push(BasicBlock { statements: Vec::default(), terminator })
     }
 
+    /// Replace statements from the given basic block
+    pub fn replace_statements(
+        &mut self,
+        source_instruction: &SourceInstruction,
+        new_stmts: Vec<Statement>,
+    ) {
+        self.blocks.get_mut(source_instruction.bb()).unwrap().statements = new_stmts;
+    }
+
     /// Replace a terminator from the given basic block
     pub fn replace_terminator(
         &mut self,
@@ -559,7 +573,7 @@ pub trait MutMirVisitor {
             StatementKind::Assign(_, rvalue) => {
                 self.visit_rvalue(rvalue);
             }
-            StatementKind::Intrinsic(intrisic) => match intrisic {
+            StatementKind::Intrinsic(intrinsic) => match intrinsic {
                 NonDivergingIntrinsic::Assume(operand) => {
                     self.visit_operand(operand);
                 }