Add invariants for while loops.

engine/lib/phases/phase_functionalize_loops.ml

@@ -74,6 +74,35 @@ struct
           { body; invariant = Some (pat, invariant) }
       | _ -> { body; invariant = None }
 
+    let extract_loop_variant (body : B.expr) : B.expr * B.expr =
+      match body.e with
+      | Let
+          {
+            monadic = None;
+            lhs = { p = PWild; _ };
+            rhs = { e = App { f = { e = GlobalVar f; _ }; args = [ e ]; _ }; _ };
+            body;
+          }
+        when Global_ident.eq_name Hax_lib___internal_loop_decreases f ->
+          (body, e)
+      | _ ->
+          let kind = { size = S32; signedness = Unsigned } in
+          let e =
+            UB.M.expr_Literal ~typ:(TInt kind) ~span:body.span
+              (Int { value = "0"; negative = false; kind })
+          in
+          let e =
+            UB.call Rust_primitives__hax__int__from_machine [ e ] e.span
+              (TApp
+                 {
+                   ident =
+                     `Concrete
+                       (Concrete_ident.of_name ~value:false Hax_lib__int__Int);
+                   args = [];
+                 })
+          in
+          (body, e)
+
     type iterator =
       | Range of { start : B.expr; end_ : B.expr }
       | Slice of B.expr
@@ -144,6 +173,16 @@ struct
             | None -> Rust_primitives__hax__folds__fold_enumerated_chunked_slice
           in
           Some (fold_op, [ size; slice ], usize)
+      | ChunksExact { size; slice } ->
+          let fold_op =
+            match cf with
+            | Some BreakOrReturn ->
+                Rust_primitives__hax__folds__fold_chunked_slice_return
+            | Some BreakOnly ->
+                Rust_primitives__hax__folds__fold_chunked_slice_cf
+            | None -> Rust_primitives__hax__folds__fold_chunked_slice
+          in
+          Some (fold_op, [ size; slice ], usize)
       | Enumerate (Slice slice) ->
           let fold_op =
             match cf with
@@ -259,6 +298,8 @@ struct
                 (M.pat_PWild ~span ~typ:unit.typ, unit)
           in
           let body = dexpr body in
+          let { body; invariant } = extract_loop_invariant body in
+          let body, variant = extract_loop_variant body in
           let condition = dexpr condition in
           let condition : B.expr =
             M.expr_Closure ~params:[ bpat ] ~body:condition ~captures:[]
@@ -276,8 +317,17 @@ struct
             | Some (BreakOnly, _) -> Rust_primitives__hax__while_loop_cf
             | None -> Rust_primitives__hax__while_loop
           in
-          UB.call fold_operator [ condition; init; body ] span
-            (dty span expr.typ)
+          let invariant : B.expr =
+            let default = MS.expr_Literal ~typ:TBool (Bool true) in
+            let invariant =
+              invariant |> Option.map ~f:snd |> Option.value ~default
+            in
+            UB.make_closure [ bpat ] invariant invariant.span
+          in
+          let variant = UB.make_closure [ bpat ] variant variant.span in
+          UB.call fold_operator
+            [ condition; invariant; variant; init; body ]
+            span (dty span expr.typ)
       | Loop { state = None; _ } ->
           Error.unimplemented ~issue_id:405 ~details:"Loop without mutation"
             span

engine/names/src/lib.rs

@@ -30,6 +30,7 @@ fn dummy_hax_concrete_ident_wrapper<I: core::iter::Iterator<Item = u8>>(x: I, mu
     assert_eq!(1, 1);
     hax_lib::assert!(true);
     hax_lib::_internal_loop_invariant(|_: usize| true);
+    hax_lib::_internal_loop_decreases(hax_lib::Int::_unsafe_from_str("0"));
 
     fn props() {
         use hax_lib::prop::*;
@@ -221,6 +222,9 @@ mod hax {
         fn fold_enumerated_chunked_slice() {}
         fn fold_enumerated_chunked_slice_cf() {}
         fn fold_enumerated_chunked_slice_return() {}
+        fn fold_chunked_slice() {}
+        fn fold_chunked_slice_cf() {}
+        fn fold_chunked_slice_return() {}
         fn fold_cf() {}
         fn fold_return() {}
     }

hax-lib/macros/src/dummy.rs

@@ -188,3 +188,8 @@ pub fn trait_fn_decoration(_attr: TokenStream, _item: TokenStream) -> TokenStrea
 pub fn loop_invariant(_predicate: TokenStream) -> TokenStream {
     quote! {}.into()
 }
+
+#[proc_macro]
+pub fn loop_decreases(_predicate: TokenStream) -> TokenStream {
+    quote! {}.into()
+}

hax-lib/macros/src/implementation.rs

@@ -77,6 +77,27 @@ pub fn loop_invariant(predicate: pm::TokenStream) -> pm::TokenStream {
     ts
 }
 
+/// Must be used to prove termination of while loops. This takes an
+/// expression that should be a usize that decreases at every iteration
+///
+/// This function must be called just after `loop_invariant`, or at the first
+/// line of the loop if there is no invariant.
+#[proc_macro]
+pub fn loop_decreases(predicate: pm::TokenStream) -> pm::TokenStream {
+    let predicate: TokenStream = predicate.into();
+    let ts: pm::TokenStream = quote! {
+        #[cfg(#HaxCfgOptionName)]
+        {
+            hax_lib::_internal_loop_decreases({
+                #HaxQuantifiers
+                #predicate
+            })
+        }
+    }
+    .into();
+    ts
+}
+
 /// When extracting to F*, inform about what is the current
 /// verification status for an item. It can either be `lax` or
 /// `panic_free`.

hax-lib/src/dummy.rs

@@ -44,7 +44,10 @@ pub fn inline_unsafe<T>(_: &str) -> T {
 }
 
 #[doc(hidden)]
-pub fn _internal_loop_invariant<T, R: Into<Prop>, P: FnOnce(T) -> R>(_: P) {}
+pub const fn _internal_loop_invariant<T, R: Into<Prop>, P: FnOnce(T) -> R>(_: &P) {}
+
+#[doc(hidden)]
+pub const fn _internal_loop_decreases(_: int::Int) {}
 
 pub trait Refinement {
     type InnerType;

hax-lib/src/implementation.rs

@@ -143,6 +143,10 @@ pub fn any_to_unit<T>(_: T) -> () {
 #[doc(hidden)]
 pub fn _internal_loop_invariant<T, R: Into<Prop>, P: FnOnce(T) -> R>(_: P) {}
 
+/// A dummy function that holds a loop variant.
+#[doc(hidden)]
+pub fn _internal_loop_decreases(_: Int) {}
+
 /// A type that implements `Refinement` should be a newtype for a
 /// type `T`. The field holding the value of type `T` should be
 /// private, and `Refinement` should be the only interface to the

hax-lib/src/proc_macros.rs

@@ -2,8 +2,9 @@
 //! proc-macro crate cannot export anything but procedural macros.
 
 pub use hax_lib_macros::{
-    attributes, decreases, ensures, exclude, impl_fn_decoration, include, lemma, loop_invariant,
-    opaque, opaque_type, refinement_type, requires, trait_fn_decoration, transparent,
+    attributes, decreases, ensures, exclude, impl_fn_decoration, include, lemma, loop_decreases,
+    loop_invariant, opaque, opaque_type, refinement_type, requires, trait_fn_decoration,
+    transparent,
 };
 
 pub use hax_lib_macros::{

proof-libs/fstar/core/Core.Ops.Arith.fsti

@@ -1,60 +1,60 @@
 module Core.Ops.Arith
 open Rust_primitives
-
+open Hax_lib.Prop
 
 class t_Add self rhs = {
    [@@@ Tactics.Typeclasses.no_method]
    f_Output: Type;
-   f_add_pre: self -> rhs -> bool;
-   f_add_post: self -> rhs -> f_Output -> bool;
+   f_add_pre: self -> rhs -> t_Prop;
+   f_add_post: self -> rhs -> f_Output -> t_Prop;
    f_add: x:self -> y:rhs -> Pure f_Output (f_add_pre x y) (fun r -> f_add_post x y r);
 }
 
 class t_Sub self rhs = {
    [@@@ Tactics.Typeclasses.no_method]
    f_Output: Type;
-   f_sub_pre: self -> rhs -> bool;
-   f_sub_post: self -> rhs -> f_Output -> bool;
+   f_sub_pre: self -> rhs -> t_Prop;
+   f_sub_post: self -> rhs -> f_Output -> t_Prop;
    f_sub: x:self -> y:rhs -> Pure f_Output (f_sub_pre x y) (fun r -> f_sub_post x y r);
 }
 
 class t_Mul self rhs = {
    [@@@ Tactics.Typeclasses.no_method]
    f_Output: Type;
-   f_mul_pre: self -> rhs -> bool;
-   f_mul_post: self -> rhs -> f_Output -> bool;
+   f_mul_pre: self -> rhs -> t_Prop;
+   f_mul_post: self -> rhs -> f_Output -> t_Prop;
    f_mul: x:self -> y:rhs -> Pure f_Output (f_mul_pre x y) (fun r -> f_mul_post x y r);
 }
 
 class t_Div self rhs = {
    [@@@ Tactics.Typeclasses.no_method]
    f_Output: Type;
-   f_div_pre: self -> rhs -> bool;
-   f_div_post: self -> rhs -> f_Output -> bool;
+   f_div_pre: self -> rhs -> t_Prop;
+   f_div_post: self -> rhs -> f_Output -> t_Prop;
    f_div: x:self -> y:rhs -> Pure f_Output (f_div_pre x y) (fun r -> f_div_post x y r);
 }
 
 class t_AddAssign self rhs = {
-  f_add_assign_pre: self -> rhs -> bool;
-  f_add_assign_post: self -> rhs -> self -> bool;
+  f_add_assign_pre: self -> rhs -> t_Prop;
+  f_add_assign_post: self -> rhs -> self -> t_Prop;
   f_add_assign: x:self -> y:rhs -> Pure self (f_add_assign_pre x y) (fun r -> f_add_assign_post x y r);
 }
 
 class t_SubAssign self rhs = {
-  f_sub_assign_pre: self -> rhs -> bool;
-  f_sub_assign_post: self -> rhs -> self -> bool;
+  f_sub_assign_pre: self -> rhs -> t_Prop;
+  f_sub_assign_post: self -> rhs -> self -> t_Prop;
   f_sub_assign: x:self -> y:rhs -> Pure self (f_sub_assign_pre x y) (fun r -> f_sub_assign_post x y r);
 }
 
 class t_MulAssign self rhs = {
-   f_mul_assign_pre: self -> rhs -> bool;
-   f_mul_assign_post: self -> rhs -> self -> bool;
+   f_mul_assign_pre: self -> rhs -> t_Prop;
+   f_mul_assign_post: self -> rhs -> self -> t_Prop;
    f_mul_assign: x:self -> y:rhs -> Pure self (f_mul_assign_pre x y) (fun r -> f_mul_assign_post x y r);
 }
 
 class t_DivAssign self rhs = {
-   f_div_assign_pre: self -> rhs -> bool;
-   f_div_assign_post: self -> rhs -> self -> bool;
+   f_div_assign_pre: self -> rhs -> t_Prop;
+   f_div_assign_post: self -> rhs -> self -> t_Prop;
    f_div_assign: x:self -> y:rhs -> Pure self (f_div_assign_pre x y) (fun r -> f_div_assign_post x y r);
 }
 

proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti

@@ -37,6 +37,26 @@ val fold_enumerated_chunked_slice
   )
   : result: acc_t {inv result (mk_int (Seq.length s / v chunk_size))}
 
+/// Fold function that is generated for `for` loops iterating on
+/// `s.chunks_exact(chunk_size)`-like iterators
+val fold_chunked_slice
+  (#t: Type0) (#acc_t: Type0)
+  (chunk_size: usize {v chunk_size > 0})
+  (s: t_Slice t)
+  (inv: acc_t -> (i:usize) -> Type0)
+  (init: acc_t {inv init (sz 0)})
+  (f: ( acc:acc_t
+      -> item:(t_Slice t) {
+        length item == chunk_size /\
+        inv acc (sz 0)
+      }
+      -> acc':acc_t {
+        inv acc' (sz 0)
+      }
+      )
+  )
+  : result: acc_t {inv result (mk_int 0)}
+
 (**** `s.enumerate()` *)
 /// Fold function that is generated for `for` loops iterating on
 /// `s.enumerate()`-like iterators

proof-libs/fstar/rust_primitives/Rust_primitives.Hax.fst

@@ -63,7 +63,22 @@ class iterator_return (self: Type u#0): Type u#1 = {
   parent_iterator: Core.Iter.Traits.Iterator.t_Iterator self;
   f_fold_return: #b:Type0 -> s:self -> b -> (b -> i:parent_iterator.f_Item{parent_iterator.f_contains s i} -> Core.Ops.Control_flow.t_ControlFlow b b) -> Core.Ops.Control_flow.t_ControlFlow b b;
 }
-let rec while_loop #s (condition: s -> bool) (init: s) (f: (i:s -> o:s{o << i})): s
-  = if condition init
-    then while_loop #s  condition (f init) f
-    else init
+let while_loop #acc_t 
+  (condition: acc_t -> bool) 
+  (inv: acc_t -> Type0)
+  (fuel: (a:acc_t{inv a} -> nat))
+  (init: acc_t {inv init}) 
+  (f: (i:acc_t{condition i /\ inv i} -> o:acc_t{inv o /\ fuel o < fuel i})): 
+  (res: acc_t {inv res /\ not (condition res)})
+  = 
+  let rec while_loop_internal
+  (current: acc_t {inv current}): 
+  Tot (res: acc_t {inv res /\ not (condition res)}) (decreases (fuel current))
+  = if condition current
+    then 
+      let next = f current in 
+      assert (fuel next < fuel current);
+      while_loop_internal next
+    else current in 
+  while_loop_internal init
+

test-harness/src/snapshots/toolchain__loops into-fstar.snap

@@ -208,6 +208,12 @@ let bigger_power_2_ (x: i32) : i32 =
   Rust_primitives.Hax.while_loop_cf (fun pow ->
         let pow:i32 = pow in
         pow <. mk_i32 1000000 <: bool)
+    (fun pow ->
+        let pow:i32 = pow in
+        true)
+    (fun pow ->
+        let pow:i32 = pow in
+        Rust_primitives.Hax.Int.from_machine (mk_u32 0) <: Hax_lib.Int.t_Int)
     pow
     (fun pow ->
         let pow:i32 = pow in
@@ -842,6 +848,32 @@ let f (_: Prims.unit) : u8 =
     Rust_primitives.Hax.while_loop (fun x ->
           let x:u8 = x in
           x <. mk_u8 10 <: bool)
+      (fun x ->
+          let x:u8 = x in
+          true)
+      (fun x ->
+          let x:u8 = x in
+          Rust_primitives.Hax.Int.from_machine (mk_u32 0) <: Hax_lib.Int.t_Int)
+      x
+      (fun x ->
+          let x:u8 = x in
+          let x:u8 = x +! mk_u8 3 in
+          x)
+  in
+  x +! mk_u8 12
+
+let while_invariant_decr (_: Prims.unit) : u8 =
+  let x:u8 = mk_u8 0 in
+  let x:u8 =
+    Rust_primitives.Hax.while_loop (fun x ->
+          let x:u8 = x in
+          x <. mk_u8 10 <: bool)
+      (fun x ->
+          let x:u8 = x in
+          x <=. mk_u8 10 <: bool)
+      (fun x ->
+          let x:u8 = x in
+          Rust_primitives.Hax.Int.from_machine (mk_u8 10 -! x <: u8) <: Hax_lib.Int.t_Int)
       x
       (fun x ->
           let x:u8 = x in

tests/loops/src/lib.rs

@@ -138,6 +138,16 @@ mod while_loops {
         }
         x + 12
     }
+    fn while_invariant_decr() -> u8 {
+        use hax_lib::ToInt;
+        let mut x = 0;
+        while x < 10 {
+            hax_lib::loop_invariant!(|_: usize| x <= 10);
+            hax_lib::loop_decreases!((10 - x).to_int());
+            x = x + 3;
+        }
+        x + 12
+    }
 }
 
 mod control_flow {