From b1ca4efe4227fbe86db17ee310802b7333e96cc3 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 13:23:11 +0200
Subject: [PATCH 01/13] feat(devshell): add `list-names`

---
 .utils/list-names.sh | 11 +++++++++++
 flake.nix            | 19 +++++++++++--------
 2 files changed, 22 insertions(+), 8 deletions(-)
 create mode 100755 .utils/list-names.sh

diff --git a/.utils/list-names.sh b/.utils/list-names.sh
new file mode 100755
index 000000000..e0c99c586
--- /dev/null
+++ b/.utils/list-names.sh
@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+function pager() {
+    if command -v bat &> /dev/null; then
+        bat -l ml
+    else
+        less
+    fi
+}
+
+hax-engine-names-extract | sed '/include .val/,$d' | pager
diff --git a/flake.nix b/flake.nix
index 9327e8582..19791054d 100644
--- a/flake.nix
+++ b/flake.nix
@@ -152,6 +152,15 @@
             })
           ];
         in let
+          utils = pkgs.stdenv.mkDerivation {
+            name = "hax-dev-scripts";
+            phases = ["installPhase"];
+            installPhase = ''
+                mkdir -p $out/bin
+                cp ${./.utils/rebuild.sh} $out/bin/rebuild
+                cp ${./.utils/list-names.sh} $out/bin/list-names
+              '';
+          };
           packages = [
             ocamlformat
             pkgs.ocamlPackages.ocaml-lsp
@@ -168,14 +177,7 @@
             rustfmt
             rustc
 
-            (pkgs.stdenv.mkDerivation {
-              name = "rebuild-script";
-              phases = ["installPhase"];
-              installPhase = ''
-                mkdir -p $out/bin
-                cp ${./.utils/rebuild.sh} $out/bin/rebuild
-              '';
-            })
+            utils
           ];
           LIBCLANG_PATH = "${pkgs.llvmPackages.libclang.lib}/lib";
         in {
@@ -191,6 +193,7 @@
           };
           default = pkgs.mkShell {
             inherit packages inputsFrom LIBCLANG_PATH;
+            shellHook = ''echo "Commands available: $(ls ${utils}/bin | tr '\n' ' ')"'';
           };
         };
       }

From fc65169c3000ac8fd52392d0604b6a93afa0f56c Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 13:23:37 +0200
Subject: [PATCH 02/13] feat(engine/ast_utils): add `Debug` helper module

---
 engine/lib/ast_utils.ml | 17 +++++++++++++++++
 engine/lib/utils.ml     | 14 ++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/engine/lib/ast_utils.ml b/engine/lib/ast_utils.ml
index 0675994b9..7632a7dfc 100644
--- a/engine/lib/ast_utils.ml
+++ b/engine/lib/ast_utils.ml
@@ -927,6 +927,23 @@ module Make (F : Features.T) = struct
     let item : AST.item -> Ast.Full.item = Stdlib.Obj.magic
   end
 
+  module Debug : sig
+    val expr : ?label:string -> AST.expr -> unit
+    (** Prints an expression pretty-printed as Rust, with its full
+        AST encoded as JSON, available as a file, so that one can
+        `jless` or `jq` into it. *)
+  end = struct
+    let expr ?(label = "") (e : AST.expr) : unit =
+      let path = tempfile_path ~suffix:".json" in
+      Core.Out_channel.write_all path
+        ~data:([%yojson_of: AST.expr] e |> Yojson.Safe.pretty_to_string);
+      let e = LiftToFullAst.expr e in
+      "```rust " ^ label ^ "\n" ^ Print_rust.pexpr_str e
+      ^ "\n```\x1b[34m JSON-encoded AST available at \x1b[1m" ^ path
+      ^ "\x1b[0m (hint: use `jless " ^ path ^ "`)"
+      |> Stdio.prerr_endline
+  end
+
   let unbox_expr' (next : expr -> expr) (e : expr) : expr =
     match e.e with
     | App { f = { e = GlobalVar f; _ }; args = [ e ]; _ }
diff --git a/engine/lib/utils.ml b/engine/lib/utils.ml
index 1243aa98a..aeef86cfd 100644
--- a/engine/lib/utils.ml
+++ b/engine/lib/utils.ml
@@ -93,3 +93,17 @@ module MyInt64 = struct
 
   let yojson_of_t (int64 : t) : Yojson.Safe.t = `Intlit (to_string int64)
 end
+
+include (
+  struct
+    let id = ref 0
+
+    let tempfile_path ~suffix =
+      id := !id + 1;
+      Core.Filename.(
+        concat temp_dir_name ("hax-debug-" ^ Int.to_string !id ^ suffix))
+  end :
+    sig
+      val tempfile_path : suffix:string -> string
+      (** Generates a temporary file path that ends with `suffix` *)
+    end)

From ccfe0d74f622dc19a3c5b166376ff00902e5d44d Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 13:26:35 +0200
Subject: [PATCH 03/13] feat(engine): recognize iterator combinators and loop
 invariants

---
 .../lib/phases/phase_functionalize_loops.ml   | 120 ++++++++++++++++--
 engine/names/src/lib.rs                       |  16 +++
 hax-lib/src/lib.rs                            |   3 +
 3 files changed, 128 insertions(+), 11 deletions(-)

diff --git a/engine/lib/phases/phase_functionalize_loops.ml b/engine/lib/phases/phase_functionalize_loops.ml
index 1e29f7959..fd6c67998 100644
--- a/engine/lib/phases/phase_functionalize_loops.ml
+++ b/engine/lib/phases/phase_functionalize_loops.ml
@@ -33,6 +33,95 @@ struct
       include Features.SUBTYPE.Id
     end
 
+    type body_and_invariant = { body : B.expr; invariant : B.expr option }
+
+    let extract_loop_invariant (body : B.expr) : body_and_invariant =
+      match body.e with
+      | Let
+          {
+            monadic = None;
+            lhs = { p = PWild; _ };
+            rhs =
+              {
+                e = App { f = { e = GlobalVar f; _ }; args = [ invariant ]; _ };
+                _;
+              };
+            body;
+          }
+        when Global_ident.eq_name Hax_lib__loop_invariant f ->
+          { body; invariant = Some invariant }
+      | _ -> { body; invariant = None }
+
+    type iterator =
+      | Range of { start : B.expr; end_ : B.expr }
+      | Slice of B.expr
+      | ChunksExact of { size : B.expr; slice : B.expr }
+      | Enumerate of iterator
+      | StepBy of { n : B.expr; it : iterator }
+    [@@deriving show]
+
+    let rec as_iterator (e : B.expr) : iterator option =
+      match e.e with
+      | Construct
+          {
+            constructor = `Concrete range_ctor;
+            is_record = true;
+            is_struct = true;
+            fields =
+              [ (`Concrete start_field, start); (`Concrete end_field, end_) ];
+            base = None;
+          }
+        when Concrete_ident.eq_name Core__ops__range__Range__start start_field
+             && Concrete_ident.eq_name Core__ops__range__Range range_ctor
+             && Concrete_ident.eq_name Core__ops__range__Range__end end_field ->
+          Some (Range { start; end_ })
+      | _ -> meth_as_iterator e
+
+    and meth_as_iterator (e : B.expr) : iterator option =
+      let* f, args =
+        match e.e with
+        | App { f = { e = GlobalVar f; _ }; args; _ } -> Some (f, args)
+        | _ -> None
+      in
+      let f_eq n = Global_ident.eq_name n f in
+      let one_arg () = match args with [ x ] -> Some x | _ -> None in
+      let two_args () = match args with [ x; y ] -> Some (x, y) | _ -> None in
+      if f_eq Core__iter__traits__iterator__Iterator__step_by then
+        let* it, n = two_args () in
+        let* it = as_iterator it in
+        Some (StepBy { n; it })
+      else if
+        f_eq Core__iter__traits__collect__IntoIterator__into_iter
+        || f_eq Core__slice__Impl__iter
+      then
+        let* iterable = one_arg () in
+        match iterable.typ with
+        | TSlice _ -> Some (Slice iterable)
+        | _ -> as_iterator iterable
+      else if f_eq Core__iter__traits__iterator__Iterator__enumerate then
+        let* iterable = one_arg () in
+        let* iterator = as_iterator iterable in
+        Some (Enumerate iterator)
+      else if f_eq Core__slice__Impl__chunks_exact then
+        let* slice, size = two_args () in
+        Some (ChunksExact { size; slice })
+      else None
+
+    let fn_args_of_iterator (it : iterator) :
+        (Concrete_ident.name * B.expr list) option =
+      let open Concrete_ident_generated in
+      match it with
+      | Enumerate (ChunksExact { size; slice }) ->
+          Some
+            ( Rust_primitives__hax__folds__fold_enumerated_chunked_slice,
+              [ size; slice ] )
+      | Enumerate (Slice slice) ->
+          Some (Rust_primitives__hax__folds__fold_enumerated_slice, [ slice ])
+      | StepBy { n; it = Range { start; end_ } } ->
+          Some
+            (Rust_primitives__hax__folds__fold_range_step_by, [ start; end_; n ])
+      | _ -> None
+
     [%%inline_defs dmutability]
 
     let rec dexpr_unwrapped (expr : A.expr) : B.expr =
@@ -46,23 +135,32 @@ struct
             _;
           } ->
           let body = dexpr body in
+          let { body; invariant } = extract_loop_invariant body in
           let it = dexpr it in
           let pat = dpat pat in
           let bpat = dpat bpat in
-          let fn : B.expr' =
-            Closure { params = [ bpat; pat ]; body; captures = [] }
-          in
-          let fn : B.expr =
+          let as_lhs_closure e : B.expr =
             {
-              e = fn;
-              typ = TArrow ([ bpat.typ; pat.typ ], body.typ);
-              span = body.span;
+              e = Closure { params = [ bpat; pat ]; body = e; captures = [] };
+              typ = TArrow ([ bpat.typ; pat.typ ], e.typ);
+              span = e.span;
             }
           in
-          UB.call ~kind:(AssociatedItem Value)
-            Core__iter__traits__iterator__Iterator__fold
-            [ it; dexpr init; fn ]
-            span (dty span expr.typ)
+          let fn : B.expr = as_lhs_closure body in
+          let invariant : B.expr =
+            let default : B.expr =
+              { e = Literal (Bool true); typ = TBool; span = expr.span }
+            in
+            Option.value ~default invariant |> as_lhs_closure
+          in
+          let init = dexpr init in
+          let f, args =
+            match as_iterator it |> Option.bind ~f:fn_args_of_iterator with
+            | Some (f, args) -> (f, args @ [ init; invariant; fn ])
+            | None ->
+                (Core__iter__traits__iterator__Iterator__fold, [ it; init; fn ])
+          in
+          UB.call ~kind:(AssociatedItem Value) f args span (dty span expr.typ)
       | Loop
           {
             body;
diff --git a/engine/names/src/lib.rs b/engine/names/src/lib.rs
index 6e53dd5f0..27c923921 100644
--- a/engine/names/src/lib.rs
+++ b/engine/names/src/lib.rs
@@ -29,6 +29,7 @@ fn dummy_hax_concrete_ident_wrapper<I: core::iter::Iterator<Item = u8>>(x: I, mu
     assert!(true);
     assert_eq!(1, 1);
     hax_lib::assert!(true);
+    hax_lib::loop_invariant(true);
 
     let _ = [()].into_iter();
     let _: u16 = 6u8.into();
@@ -37,6 +38,14 @@ fn dummy_hax_concrete_ident_wrapper<I: core::iter::Iterator<Item = u8>>(x: I, mu
     let _ = ..;
     let _ = ..1;
 
+    fn iterator_functions<It: Iterator + Clone>(it: It) {
+        let _ = it.clone().step_by(2);
+        let _ = it.clone().enumerate();
+        let _ = [()].chunks_exact(2);
+        let _ = [()].iter();
+        let _ = (&[()] as &[()]).iter();
+    }
+
     {
         use hax_lib::int::*;
         let a: Int = 3u8.lift();
@@ -163,6 +172,13 @@ mod hax {
     fn array_of_list() {}
     fn never_to_any() {}
 
+    mod folds {
+        fn fold_range() {}
+        fn fold_range_step_by() {}
+        fn fold_enumerated_slice() {}
+        fn fold_enumerated_chunked_slice() {}
+    }
+
     /// The engine uses this `dropped_body` symbol as a marker value
     /// to signal that a item was extracted without body.
     fn dropped_body() {}
diff --git a/hax-lib/src/lib.rs b/hax-lib/src/lib.rs
index 44eba43ab..0b8341e8b 100644
--- a/hax-lib/src/lib.rs
+++ b/hax-lib/src/lib.rs
@@ -147,6 +147,9 @@ pub fn inline_unsafe<T>(_: &str) -> T {
     unreachable!()
 }
 
+#[doc(hidden)]
+pub fn loop_invariant(_: bool) {}
+
 /// 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

From e22677aa134a71b600fef62e5355f614c4db8156 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 13:28:24 +0200
Subject: [PATCH 04/13] test(engine): recognize iterator combinators and loop
 invariants

---
 .../toolchain__loops into-fstar.snap          | 65 ++++++++++++++++---
 tests/Cargo.lock                              |  3 +
 tests/loops/Cargo.toml                        |  1 +
 tests/loops/src/lib.rs                        | 23 +++++++
 4 files changed, 82 insertions(+), 10 deletions(-)

diff --git a/test-harness/src/snapshots/toolchain__loops into-fstar.snap b/test-harness/src/snapshots/toolchain__loops into-fstar.snap
index ca9d2fa70..24a50c837 100644
--- a/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
+++ b/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
@@ -148,17 +148,12 @@ let enumerate_chunks (arr: Alloc.Vec.t_Vec usize Alloc.Alloc.t_Global) : usize =
       (fun acc temp_1_ ->
           let acc:usize = acc in
           let i, chunk:(usize & t_Slice usize) = temp_1_ in
-          Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Iter.Adapters.Enumerate.t_Enumerate
-                  (Core.Slice.Iter.t_Iter usize))
-                #FStar.Tactics.Typeclasses.solve
-                (Core.Iter.Traits.Iterator.f_enumerate #(Core.Slice.Iter.t_Iter usize)
-                    #FStar.Tactics.Typeclasses.solve
-                    (Core.Slice.impl__iter #usize chunk <: Core.Slice.Iter.t_Iter usize)
-                  <:
-                  Core.Iter.Adapters.Enumerate.t_Enumerate (Core.Slice.Iter.t_Iter usize))
-              <:
-              Core.Iter.Adapters.Enumerate.t_Enumerate (Core.Slice.Iter.t_Iter usize))
+          Rust_primitives.Hax.f_fold_enumerated_slice chunk
             acc
+            (fun acc temp_1_ ->
+                let acc:usize = acc in
+                let j, x:(usize & usize) = temp_1_ in
+                true)
             (fun acc temp_1_ ->
                 let acc:usize = acc in
                 let j, x:(usize & usize) = temp_1_ in
@@ -356,6 +351,56 @@ let rev_range (n: usize) : usize =
   in
   acc
 '''
+"Loops.Recognized_loops.fst" = '''
+module Loops.Recognized_loops
+#set-options "--fuel 0 --ifuel 1 --z3rlimit 15"
+open Core
+open FStar.Mul
+
+let enumerated_chunked_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
+  let count:u64 = 0uL in
+  Rust_primitives.Hax.f_fold_enumerated_chunked_slice (sz 3)
+    slice
+    count
+    (fun count i ->
+        let count:u64 = count in
+        let i:(usize & t_Slice v_T) = i in
+        true)
+    (fun count i ->
+        let count:u64 = count in
+        let i:(usize & t_Slice v_T) = i in
+        count +! 3uL <: u64)
+
+let enumerated_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
+  let count:u64 = 0uL in
+  Rust_primitives.Hax.f_fold_enumerated_slice slice
+    count
+    (fun count i ->
+        let count:u64 = count in
+        let i:(usize & v_T) = i in
+        false)
+    (fun count i ->
+        let count:u64 = count in
+        let i:(usize & v_T) = i in
+        let count:u64 = count +! 2uL in
+        count)
+
+let range_step_by (_: Prims.unit) : u64 =
+  let count:u64 = 0uL in
+  Rust_primitives.Hax.f_fold_range_step_by 0l
+    10l
+    (sz 2)
+    count
+    (fun count i ->
+        let count:u64 = count in
+        let i:i32 = i in
+        i <. 20l <: bool)
+    (fun count i ->
+        let count:u64 = count in
+        let i:i32 = i in
+        let count:u64 = count +! 1uL in
+        count)
+'''
 "Loops.While_loops.fst" = '''
 module Loops.While_loops
 #set-options "--fuel 0 --ifuel 1 --z3rlimit 15"
diff --git a/tests/Cargo.lock b/tests/Cargo.lock
index f48220811..d3d491045 100644
--- a/tests/Cargo.lock
+++ b/tests/Cargo.lock
@@ -499,6 +499,9 @@ checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f"
 [[package]]
 name = "loops"
 version = "0.1.0"
+dependencies = [
+ "hax-lib",
+]
 
 [[package]]
 name = "memchr"
diff --git a/tests/loops/Cargo.toml b/tests/loops/Cargo.toml
index a6139c47a..96dfe7dc3 100644
--- a/tests/loops/Cargo.toml
+++ b/tests/loops/Cargo.toml
@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+hax-lib = { path = "../../hax-lib" }
 
 [package.metadata.hax-tests]
 into."fstar" = { }
diff --git a/tests/loops/src/lib.rs b/tests/loops/src/lib.rs
index 22ee1cab6..7eb0eaf50 100644
--- a/tests/loops/src/lib.rs
+++ b/tests/loops/src/lib.rs
@@ -1,3 +1,26 @@
+mod recognized_loops {
+    fn range_step_by() {
+        let mut count = 0u64;
+        for i in (0..10).step_by(2) {
+            hax_lib::loop_invariant(i < 20);
+            count += 1;
+        }
+    }
+    fn enumerated_slice<T>(slice: &[T]) {
+        let mut count = 0u64;
+        for i in slice.into_iter().enumerate() {
+            hax_lib::loop_invariant(false);
+            count += 2;
+        }
+    }
+    fn enumerated_chunked_slice<T>(slice: &[T]) {
+        let mut count = 0u64;
+        for i in slice.chunks_exact(3).enumerate() {
+            count += 3;
+        }
+    }
+}
+
 mod for_loops {
     fn range1() -> usize {
         let mut acc = 0;

From 214365a79ebc554b553c323bb4eecaa677b53e91 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 14:23:55 +0200
Subject: [PATCH 05/13] refactor(engine): `loop_invariant!` instead of
 `loop_invariant`

---
 .../lib/phases/phase_functionalize_loops.ml   |  2 +-
 engine/names/src/lib.rs                       |  2 +-
 hax-lib-macros/src/lib.rs                     | 19 +++++++++++++++++++
 hax-lib/src/lib.rs                            |  3 ++-
 hax-lib/src/proc_macros.rs                    |  4 ++--
 tests/loops/src/lib.rs                        |  5 +++--
 6 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/engine/lib/phases/phase_functionalize_loops.ml b/engine/lib/phases/phase_functionalize_loops.ml
index fd6c67998..3c9959bf4 100644
--- a/engine/lib/phases/phase_functionalize_loops.ml
+++ b/engine/lib/phases/phase_functionalize_loops.ml
@@ -48,7 +48,7 @@ struct
               };
             body;
           }
-        when Global_ident.eq_name Hax_lib__loop_invariant f ->
+        when Global_ident.eq_name Hax_lib___internal_loop_invariant f ->
           { body; invariant = Some invariant }
       | _ -> { body; invariant = None }
 
diff --git a/engine/names/src/lib.rs b/engine/names/src/lib.rs
index 27c923921..db3fd44e0 100644
--- a/engine/names/src/lib.rs
+++ b/engine/names/src/lib.rs
@@ -29,7 +29,7 @@ fn dummy_hax_concrete_ident_wrapper<I: core::iter::Iterator<Item = u8>>(x: I, mu
     assert!(true);
     assert_eq!(1, 1);
     hax_lib::assert!(true);
-    hax_lib::loop_invariant(true);
+    hax_lib::_internal_loop_invariant(true);
 
     let _ = [()].into_iter();
     let _: u16 = 6u8.into();
diff --git a/hax-lib-macros/src/lib.rs b/hax-lib-macros/src/lib.rs
index bcf624326..f2cfac8e9 100644
--- a/hax-lib-macros/src/lib.rs
+++ b/hax-lib-macros/src/lib.rs
@@ -40,6 +40,25 @@ pub fn fstar_options(attr: pm::TokenStream, item: pm::TokenStream) -> pm::TokenS
     .into()
 }
 
+/// Add an invariant to a loop. This function must be called on the
+/// first line of a loop body to be effective. Note that in the
+/// invariant expression, `forall`, `exists`, and `BACKEND!`
+/// (`BACKEND` can be `fstar`, `proverif`, `coq`...) are in scope.
+#[proc_macro]
+pub fn loop_invariant(predicate: pm::TokenStream) -> pm::TokenStream {
+    let predicate: TokenStream = predicate.into();
+    quote! {
+        #[cfg(#HaxCfgOptionName)]
+        {
+            hax_lib::_internal_loop_invariant({
+                #HaxQuantifiers
+                #predicate
+            })
+        }
+    }
+    .into()
+}
+
 /// When extracting to F*, inform about what is the current
 /// verification status for an item. It can either be `lax` or
 /// `panic_free`.
diff --git a/hax-lib/src/lib.rs b/hax-lib/src/lib.rs
index 0b8341e8b..a6dbd9ea7 100644
--- a/hax-lib/src/lib.rs
+++ b/hax-lib/src/lib.rs
@@ -147,8 +147,9 @@ pub fn inline_unsafe<T>(_: &str) -> T {
     unreachable!()
 }
 
+/// A dummy function that holds a loop invariant.
 #[doc(hidden)]
-pub fn loop_invariant(_: bool) {}
+pub fn _internal_loop_invariant(_: bool) {}
 
 /// A type that implements `Refinement` should be a newtype for a
 /// type `T`. The field holding the value of type `T` should be
diff --git a/hax-lib/src/proc_macros.rs b/hax-lib/src/proc_macros.rs
index f18274df2..866c7cfda 100644
--- a/hax-lib/src/proc_macros.rs
+++ b/hax-lib/src/proc_macros.rs
@@ -2,8 +2,8 @@
 //! proc-macro crate cannot export anything but procedural macros.
 
 pub use hax_lib_macros::{
-    attributes, ensures, exclude, impl_fn_decoration, include, lemma, opaque_type, refinement_type,
-    requires, trait_fn_decoration,
+    attributes, ensures, exclude, impl_fn_decoration, include, lemma, loop_invariant, opaque_type,
+    refinement_type, requires, trait_fn_decoration,
 };
 
 pub use hax_lib_macros::{
diff --git a/tests/loops/src/lib.rs b/tests/loops/src/lib.rs
index 7eb0eaf50..24a737899 100644
--- a/tests/loops/src/lib.rs
+++ b/tests/loops/src/lib.rs
@@ -2,20 +2,21 @@ mod recognized_loops {
     fn range_step_by() {
         let mut count = 0u64;
         for i in (0..10).step_by(2) {
-            hax_lib::loop_invariant(i < 20);
+            hax_lib::loop_invariant!(i < 20);
             count += 1;
         }
     }
     fn enumerated_slice<T>(slice: &[T]) {
         let mut count = 0u64;
         for i in slice.into_iter().enumerate() {
-            hax_lib::loop_invariant(false);
+            hax_lib::loop_invariant!(false);
             count += 2;
         }
     }
     fn enumerated_chunked_slice<T>(slice: &[T]) {
         let mut count = 0u64;
         for i in slice.chunks_exact(3).enumerate() {
+            hax_lib::loop_invariant!(fstar!("forall (x: nat). x >= 0"));
             count += 3;
         }
     }

From 9a82ffcf2f173943a94e76f548ab7b5ce6b205d7 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 14:31:24 +0200
Subject: [PATCH 06/13] refactor(engine): fix `rust_primitives::hax::folds`
 name

---
 engine/lib/phases/phase_functionalize_loops.ml      |  8 ++++----
 .../src/snapshots/toolchain__loops into-fstar.snap  | 13 +++++++------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/engine/lib/phases/phase_functionalize_loops.ml b/engine/lib/phases/phase_functionalize_loops.ml
index 3c9959bf4..6bd8d8737 100644
--- a/engine/lib/phases/phase_functionalize_loops.ml
+++ b/engine/lib/phases/phase_functionalize_loops.ml
@@ -154,13 +154,13 @@ struct
             Option.value ~default invariant |> as_lhs_closure
           in
           let init = dexpr init in
-          let f, args =
+          let f, kind, args =
             match as_iterator it |> Option.bind ~f:fn_args_of_iterator with
-            | Some (f, args) -> (f, args @ [ init; invariant; fn ])
+            | Some (f, args) -> (f, Concrete_ident.Kind.Value, args @ [ init; invariant; fn ])
             | None ->
-                (Core__iter__traits__iterator__Iterator__fold, [ it; init; fn ])
+                (Core__iter__traits__iterator__Iterator__fold, AssociatedItem Value, [ it; init; fn ])
           in
-          UB.call ~kind:(AssociatedItem Value) f args span (dty span expr.typ)
+          UB.call ~kind f args span (dty span expr.typ)
       | Loop
           {
             body;
diff --git a/test-harness/src/snapshots/toolchain__loops into-fstar.snap b/test-harness/src/snapshots/toolchain__loops into-fstar.snap
index 24a50c837..89101b5c7 100644
--- a/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
+++ b/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
@@ -148,7 +148,7 @@ let enumerate_chunks (arr: Alloc.Vec.t_Vec usize Alloc.Alloc.t_Global) : usize =
       (fun acc temp_1_ ->
           let acc:usize = acc in
           let i, chunk:(usize & t_Slice usize) = temp_1_ in
-          Rust_primitives.Hax.f_fold_enumerated_slice chunk
+          Rust_primitives.Hax.Folds.fold_enumerated_slice chunk
             acc
             (fun acc temp_1_ ->
                 let acc:usize = acc in
@@ -359,21 +359,22 @@ open FStar.Mul
 
 let enumerated_chunked_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
   let count:u64 = 0uL in
-  Rust_primitives.Hax.f_fold_enumerated_chunked_slice (sz 3)
+  Rust_primitives.Hax.Folds.fold_enumerated_chunked_slice (sz 3)
     slice
     count
     (fun count i ->
         let count:u64 = count in
         let i:(usize & t_Slice v_T) = i in
-        true)
+        forall (x: nat). x >= 0)
     (fun count i ->
         let count:u64 = count in
         let i:(usize & t_Slice v_T) = i in
-        count +! 3uL <: u64)
+        let count:u64 = count +! 3uL in
+        count)
 
 let enumerated_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
   let count:u64 = 0uL in
-  Rust_primitives.Hax.f_fold_enumerated_slice slice
+  Rust_primitives.Hax.Folds.fold_enumerated_slice slice
     count
     (fun count i ->
         let count:u64 = count in
@@ -387,7 +388,7 @@ let enumerated_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
 
 let range_step_by (_: Prims.unit) : u64 =
   let count:u64 = 0uL in
-  Rust_primitives.Hax.f_fold_range_step_by 0l
+  Rust_primitives.Hax.Folds.fold_range_step_by 0l
     10l
     (sz 2)
     count

From 14cd97f12f7ad7c40477c785f2832bb33447aa64 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 17:06:26 +0200
Subject: [PATCH 07/13] misc(engine): more on loops

---
 engine/lib/ast_utils.ml                       |   7 ++
 .../lib/phases/phase_functionalize_loops.ml   |  71 +++++++++----
 engine/names/src/lib.rs                       |   2 +-
 hax-lib-macros/src/lib.rs                     |  28 +++--
 hax-lib/src/lib.rs                            |   2 +-
 .../toolchain__generics into-fstar.snap       |  14 ++-
 .../toolchain__loops into-fstar.snap          | 100 +++++++++++-------
 ..._mut-ref-functionalization into-fstar.snap |  28 +++--
 tests/loops/src/lib.rs                        |  15 ++-
 9 files changed, 171 insertions(+), 96 deletions(-)

diff --git a/engine/lib/ast_utils.ml b/engine/lib/ast_utils.ml
index 7632a7dfc..88610da31 100644
--- a/engine/lib/ast_utils.ml
+++ b/engine/lib/ast_utils.ml
@@ -905,6 +905,13 @@ module Make (F : Features.T) = struct
       (`Concrete (Concrete_ident.of_name kind f_name))
       args span ret_typ
 
+  let make_closure (params : pat list) (body : expr) (span : span) : expr =
+    let params =
+      match params with [] -> [ make_wild_pat unit_typ span ] | _ -> params
+    in
+    let e = Closure { params; body; captures = [] } in
+    { e; typ = TArrow (List.map ~f:(fun p -> p.typ) params, body.typ); span }
+
   let string_lit span (s : string) : expr =
     { span; typ = TStr; e = Literal (String s) }
 
diff --git a/engine/lib/phases/phase_functionalize_loops.ml b/engine/lib/phases/phase_functionalize_loops.ml
index 6bd8d8737..35591e89f 100644
--- a/engine/lib/phases/phase_functionalize_loops.ml
+++ b/engine/lib/phases/phase_functionalize_loops.ml
@@ -33,7 +33,10 @@ struct
       include Features.SUBTYPE.Id
     end
 
-    type body_and_invariant = { body : B.expr; invariant : B.expr option }
+    type body_and_invariant = {
+      body : B.expr;
+      invariant : (B.pat * B.expr) option;
+    }
 
     let extract_loop_invariant (body : B.expr) : body_and_invariant =
       match body.e with
@@ -43,13 +46,26 @@ struct
             lhs = { p = PWild; _ };
             rhs =
               {
-                e = App { f = { e = GlobalVar f; _ }; args = [ invariant ]; _ };
+                e =
+                  App
+                    {
+                      f = { e = GlobalVar f; _ };
+                      args =
+                        [
+                          {
+                            e =
+                              Closure { params = [ pat ]; body = invariant; _ };
+                            _;
+                          };
+                        ];
+                      _;
+                    };
                 _;
               };
             body;
           }
         when Global_ident.eq_name Hax_lib___internal_loop_invariant f ->
-          { body; invariant = Some invariant }
+          { body; invariant = Some (pat, invariant) }
       | _ -> { body; invariant = None }
 
     type iterator =
@@ -108,18 +124,28 @@ struct
       else None
 
     let fn_args_of_iterator (it : iterator) :
-        (Concrete_ident.name * B.expr list) option =
+        (Concrete_ident.name * B.expr list * B.ty) option =
       let open Concrete_ident_generated in
+      let usize = B.TInt { size = SSize; signedness = Unsigned } in
       match it with
       | Enumerate (ChunksExact { size; slice }) ->
           Some
             ( Rust_primitives__hax__folds__fold_enumerated_chunked_slice,
-              [ size; slice ] )
+              [ size; slice ],
+              usize )
       | Enumerate (Slice slice) ->
-          Some (Rust_primitives__hax__folds__fold_enumerated_slice, [ slice ])
+          Some
+            ( Rust_primitives__hax__folds__fold_enumerated_slice,
+              [ slice ],
+              usize )
       | StepBy { n; it = Range { start; end_ } } ->
           Some
-            (Rust_primitives__hax__folds__fold_range_step_by, [ start; end_; n ])
+            ( Rust_primitives__hax__folds__fold_range_step_by,
+              [ start; end_; n ],
+              start.typ )
+      | Range { start; end_ } ->
+          Some
+            (Rust_primitives__hax__folds__fold_range, [ start; end_ ], start.typ)
       | _ -> None
 
     [%%inline_defs dmutability]
@@ -139,26 +165,25 @@ struct
           let it = dexpr it in
           let pat = dpat pat in
           let bpat = dpat bpat in
-          let as_lhs_closure e : B.expr =
-            {
-              e = Closure { params = [ bpat; pat ]; body = e; captures = [] };
-              typ = TArrow ([ bpat.typ; pat.typ ], e.typ);
-              span = e.span;
-            }
-          in
-          let fn : B.expr = as_lhs_closure body in
-          let invariant : B.expr =
-            let default : B.expr =
-              { e = Literal (Bool true); typ = TBool; span = expr.span }
-            in
-            Option.value ~default invariant |> as_lhs_closure
-          in
+          let fn : B.expr = UB.make_closure [ bpat; pat ] body body.span in
           let init = dexpr init in
           let f, kind, args =
             match as_iterator it |> Option.bind ~f:fn_args_of_iterator with
-            | Some (f, args) -> (f, Concrete_ident.Kind.Value, args @ [ init; invariant; fn ])
+            | Some (f, args, typ) ->
+                let invariant : B.expr =
+                  let default =
+                    let span = expr.span in
+                    let pat = UB.make_wild_pat typ span in
+                    (pat, B.{ e = Literal (Bool true); typ = TBool; span })
+                  in
+                  let pat, invariant = Option.value ~default invariant in
+                  UB.make_closure [ bpat; pat ] invariant invariant.span
+                in
+                (f, Concrete_ident.Kind.Value, args @ [ invariant; init; fn ])
             | None ->
-                (Core__iter__traits__iterator__Iterator__fold, AssociatedItem Value, [ it; init; fn ])
+                ( Core__iter__traits__iterator__Iterator__fold,
+                  AssociatedItem Value,
+                  [ it; init; fn ] )
           in
           UB.call ~kind f args span (dty span expr.typ)
       | Loop
diff --git a/engine/names/src/lib.rs b/engine/names/src/lib.rs
index db3fd44e0..438dd41f0 100644
--- a/engine/names/src/lib.rs
+++ b/engine/names/src/lib.rs
@@ -29,7 +29,7 @@ fn dummy_hax_concrete_ident_wrapper<I: core::iter::Iterator<Item = u8>>(x: I, mu
     assert!(true);
     assert_eq!(1, 1);
     hax_lib::assert!(true);
-    hax_lib::_internal_loop_invariant(true);
+    hax_lib::_internal_loop_invariant(|_: usize| true);
 
     let _ = [()].into_iter();
     let _: u16 = 6u8.into();
diff --git a/hax-lib-macros/src/lib.rs b/hax-lib-macros/src/lib.rs
index f2cfac8e9..22261bf84 100644
--- a/hax-lib-macros/src/lib.rs
+++ b/hax-lib-macros/src/lib.rs
@@ -40,14 +40,29 @@ pub fn fstar_options(attr: pm::TokenStream, item: pm::TokenStream) -> pm::TokenS
     .into()
 }
 
-/// Add an invariant to a loop. This function must be called on the
-/// first line of a loop body to be effective. Note that in the
-/// invariant expression, `forall`, `exists`, and `BACKEND!`
-/// (`BACKEND` can be `fstar`, `proverif`, `coq`...) are in scope.
+/// Add an invariant to a loop which deals with an index. The
+/// invariant cannot refer to any variable introduced within the
+/// loop. An invariant is a closure that takes one argument, the
+/// index, and returns a boolean.
+///
+/// Note that loop invariants are unstable (this will be handled in a
+/// better way in the future, see
+/// https://github.com/hacspec/hax/issues/858) and only supported on
+/// specific `for` loops with specific iterators:
+///
+///  - `for i in start..end {...}`
+///  - `for i in (start..end).step_by(n) {...}`
+///  - `for i in slice.enumerate() {...}`
+///  - `for i in slice.chunks_exact(n).enumerate() {...}`
+///
+/// This function must be called on the first line of a loop body to
+/// be effective. Note that in the invariant expression, `forall`,
+/// `exists`, and `BACKEND!` (`BACKEND` can be `fstar`, `proverif`,
+/// `coq`...) are in scope.
 #[proc_macro]
 pub fn loop_invariant(predicate: pm::TokenStream) -> pm::TokenStream {
     let predicate: TokenStream = predicate.into();
-    quote! {
+    let ts: pm::TokenStream = quote! {
         #[cfg(#HaxCfgOptionName)]
         {
             hax_lib::_internal_loop_invariant({
@@ -56,7 +71,8 @@ pub fn loop_invariant(predicate: pm::TokenStream) -> pm::TokenStream {
             })
         }
     }
-    .into()
+    .into();
+    ts
 }
 
 /// When extracting to F*, inform about what is the current
diff --git a/hax-lib/src/lib.rs b/hax-lib/src/lib.rs
index a6dbd9ea7..1b2b1b4d1 100644
--- a/hax-lib/src/lib.rs
+++ b/hax-lib/src/lib.rs
@@ -149,7 +149,7 @@ pub fn inline_unsafe<T>(_: &str) -> T {
 
 /// A dummy function that holds a loop invariant.
 #[doc(hidden)]
-pub fn _internal_loop_invariant(_: bool) {}
+pub fn _internal_loop_invariant<T, P: FnOnce(T) -> bool>(_: P) {}
 
 /// A type that implements `Refinement` should be a newtype for a
 /// type `T`. The field holding the value of type `T` should be
diff --git a/test-harness/src/snapshots/toolchain__generics into-fstar.snap b/test-harness/src/snapshots/toolchain__generics into-fstar.snap
index caa4a97f7..c875fcfa5 100644
--- a/test-harness/src/snapshots/toolchain__generics into-fstar.snap	
+++ b/test-harness/src/snapshots/toolchain__generics into-fstar.snap	
@@ -111,14 +111,12 @@ let call_g (_: Prims.unit) : usize =
 let foo (v_LEN: usize) (arr: t_Array usize v_LEN) : usize =
   let acc:usize = v_LEN +! sz 9 in
   let acc:usize =
-    Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Ops.Range.t_Range
-            usize)
-          #FStar.Tactics.Typeclasses.solve
-          ({ Core.Ops.Range.f_start = sz 0; Core.Ops.Range.f_end = v_LEN }
-            <:
-            Core.Ops.Range.t_Range usize)
-        <:
-        Core.Ops.Range.t_Range usize)
+    Rust_primitives.Hax.Folds.fold_range (sz 0)
+      v_LEN
+      (fun acc temp_1_ ->
+          let acc:usize = acc in
+          let _:usize = temp_1_ in
+          true)
       acc
       (fun acc i ->
           let acc:usize = acc in
diff --git a/test-harness/src/snapshots/toolchain__loops into-fstar.snap b/test-harness/src/snapshots/toolchain__loops into-fstar.snap
index 89101b5c7..62b4fad08 100644
--- a/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
+++ b/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
@@ -21,7 +21,22 @@ info:
     backend_options: ~
 ---
 exit = 0
-stderr = 'Finished `dev` profile [unoptimized + debuginfo] target(s) in XXs'
+stderr = '''
+warning: unnecessary parentheses around `for` iterator expression
+ --> loops/src/lib.rs:4:18
+  |
+4 |         for i in (0u8..10u8) {
+  |                  ^         ^
+  |
+  = note: `#[warn(unused_parens)]` on by default
+help: remove these parentheses
+  |
+4 -         for i in (0u8..10u8) {
+4 +         for i in 0u8..10u8 {
+  |
+
+warning: `loops` (lib) generated 1 warning (run `cargo fix --lib -p loops` to apply 1 suggestion)
+    Finished `dev` profile [unoptimized + debuginfo] target(s) in XXs'''
 
 [stdout]
 diagnostics = []
@@ -149,11 +164,11 @@ let enumerate_chunks (arr: Alloc.Vec.t_Vec usize Alloc.Alloc.t_Global) : usize =
           let acc:usize = acc in
           let i, chunk:(usize & t_Slice usize) = temp_1_ in
           Rust_primitives.Hax.Folds.fold_enumerated_slice chunk
-            acc
             (fun acc temp_1_ ->
                 let acc:usize = acc in
-                let j, x:(usize & usize) = temp_1_ in
+                let _:usize = temp_1_ in
                 true)
+            acc
             (fun acc temp_1_ ->
                 let acc:usize = acc in
                 let j, x:(usize & usize) = temp_1_ in
@@ -165,12 +180,12 @@ let enumerate_chunks (arr: Alloc.Vec.t_Vec usize Alloc.Alloc.t_Global) : usize =
 
 let f (_: Prims.unit) : u8 =
   let acc:u8 = 0uy in
-  Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Ops.Range.t_Range u8
-        )
-        #FStar.Tactics.Typeclasses.solve
-        ({ Core.Ops.Range.f_start = 1uy; Core.Ops.Range.f_end = 10uy } <: Core.Ops.Range.t_Range u8)
-      <:
-      Core.Ops.Range.t_Range u8)
+  Rust_primitives.Hax.Folds.fold_range 1uy
+    10uy
+    (fun acc temp_1_ ->
+        let acc:u8 = acc in
+        let _:u8 = temp_1_ in
+        true)
     acc
     (fun acc i ->
         let acc:u8 = acc in
@@ -293,14 +308,12 @@ let pattern (arr: Alloc.Vec.t_Vec (usize & usize) Alloc.Alloc.t_Global) : usize
 let range1 (_: Prims.unit) : usize =
   let acc:usize = sz 0 in
   let acc:usize =
-    Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Ops.Range.t_Range
-            usize)
-          #FStar.Tactics.Typeclasses.solve
-          ({ Core.Ops.Range.f_start = sz 0; Core.Ops.Range.f_end = sz 15 }
-            <:
-            Core.Ops.Range.t_Range usize)
-        <:
-        Core.Ops.Range.t_Range usize)
+    Rust_primitives.Hax.Folds.fold_range (sz 0)
+      (sz 15)
+      (fun acc temp_1_ ->
+          let acc:usize = acc in
+          let _:usize = temp_1_ in
+          true)
       acc
       (fun acc i ->
           let acc:usize = acc in
@@ -312,14 +325,12 @@ let range1 (_: Prims.unit) : usize =
 let range2 (n: usize) : usize =
   let acc:usize = sz 0 in
   let acc:usize =
-    Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Ops.Range.t_Range
-            usize)
-          #FStar.Tactics.Typeclasses.solve
-          ({ Core.Ops.Range.f_start = sz 0; Core.Ops.Range.f_end = n +! sz 10 <: usize }
-            <:
-            Core.Ops.Range.t_Range usize)
-        <:
-        Core.Ops.Range.t_Range usize)
+    Rust_primitives.Hax.Folds.fold_range (sz 0)
+      (n +! sz 10 <: usize)
+      (fun acc temp_1_ ->
+          let acc:usize = acc in
+          let _:usize = temp_1_ in
+          true)
       acc
       (fun acc i ->
           let acc:usize = acc in
@@ -361,11 +372,11 @@ let enumerated_chunked_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
   let count:u64 = 0uL in
   Rust_primitives.Hax.Folds.fold_enumerated_chunked_slice (sz 3)
     slice
-    count
     (fun count i ->
         let count:u64 = count in
-        let i:(usize & t_Slice v_T) = i in
-        forall (x: nat). x >= 0)
+        let i:usize = i in
+        i <= Core.Slice.impl__len #v_T slice)
+    count
     (fun count i ->
         let count:u64 = count in
         let i:(usize & t_Slice v_T) = i in
@@ -375,30 +386,45 @@ let enumerated_chunked_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
 let enumerated_slice (#v_T: Type0) (slice: t_Slice v_T) : u64 =
   let count:u64 = 0uL in
   Rust_primitives.Hax.Folds.fold_enumerated_slice slice
-    count
     (fun count i ->
         let count:u64 = count in
-        let i:(usize & v_T) = i in
-        false)
+        let i:usize = i in
+        i <=. sz 10 <: bool)
+    count
     (fun count i ->
         let count:u64 = count in
         let i:(usize & v_T) = i in
         let count:u64 = count +! 2uL in
         count)
 
+let range (_: Prims.unit) : u64 =
+  let count:u64 = 0uL in
+  Rust_primitives.Hax.Folds.fold_range 0uy
+    10uy
+    (fun count i ->
+        let count:u64 = count in
+        let i:u8 = i in
+        i <=. 10uy <: bool)
+    count
+    (fun count i ->
+        let count:u64 = count in
+        let i:u8 = i in
+        let count:u64 = count +! 1uL in
+        count)
+
 let range_step_by (_: Prims.unit) : u64 =
   let count:u64 = 0uL in
-  Rust_primitives.Hax.Folds.fold_range_step_by 0l
-    10l
+  Rust_primitives.Hax.Folds.fold_range_step_by 0uy
+    10uy
     (sz 2)
-    count
     (fun count i ->
         let count:u64 = count in
-        let i:i32 = i in
-        i <. 20l <: bool)
+        let i:u8 = i in
+        i <=. 10uy <: bool)
+    count
     (fun count i ->
         let count:u64 = count in
-        let i:i32 = i in
+        let i:u8 = i in
         let count:u64 = count +! 1uL in
         count)
 '''
diff --git a/test-harness/src/snapshots/toolchain__mut-ref-functionalization into-fstar.snap b/test-harness/src/snapshots/toolchain__mut-ref-functionalization into-fstar.snap
index 5782cadea..955a69474 100644
--- a/test-harness/src/snapshots/toolchain__mut-ref-functionalization into-fstar.snap	
+++ b/test-harness/src/snapshots/toolchain__mut-ref-functionalization into-fstar.snap	
@@ -195,14 +195,12 @@ let impl__S__update (self: t_S) (x: u8) : t_S =
 
 let foo (lhs rhs: t_S) : t_S =
   let lhs:t_S =
-    Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Ops.Range.t_Range
-            usize)
-          #FStar.Tactics.Typeclasses.solve
-          ({ Core.Ops.Range.f_start = sz 0; Core.Ops.Range.f_end = sz 1 }
-            <:
-            Core.Ops.Range.t_Range usize)
-        <:
-        Core.Ops.Range.t_Range usize)
+    Rust_primitives.Hax.Folds.fold_range (sz 0)
+      (sz 1)
+      (fun lhs temp_1_ ->
+          let lhs:t_S = lhs in
+          let _:usize = temp_1_ in
+          true)
       lhs
       (fun lhs i ->
           let lhs:t_S = lhs in
@@ -244,14 +242,12 @@ let g (x: t_Pair (Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global))
     : Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global =
   let x:t_Pair (Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global) = x in
   let x:t_Pair (Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global) =
-    Core.Iter.Traits.Iterator.f_fold (Core.Iter.Traits.Collect.f_into_iter #(Core.Ops.Range.t_Range
-            u8)
-          #FStar.Tactics.Typeclasses.solve
-          ({ Core.Ops.Range.f_start = 1uy; Core.Ops.Range.f_end = 10uy }
-            <:
-            Core.Ops.Range.t_Range u8)
-        <:
-        Core.Ops.Range.t_Range u8)
+    Rust_primitives.Hax.Folds.fold_range 1uy
+      10uy
+      (fun x temp_1_ ->
+          let x:t_Pair (Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global) = x in
+          let _:u8 = temp_1_ in
+          true)
       x
       (fun x i ->
           let x:t_Pair (Alloc.Vec.t_Vec u8 Alloc.Alloc.t_Global) = x in
diff --git a/tests/loops/src/lib.rs b/tests/loops/src/lib.rs
index 24a737899..7c209eee1 100644
--- a/tests/loops/src/lib.rs
+++ b/tests/loops/src/lib.rs
@@ -1,22 +1,29 @@
 mod recognized_loops {
+    fn range() {
+        let mut count = 0u64;
+        for i in 0u8..10u8 {
+            hax_lib::loop_invariant!(|i: u8| i <= 10);
+            count += 1;
+        }
+    }
     fn range_step_by() {
         let mut count = 0u64;
-        for i in (0..10).step_by(2) {
-            hax_lib::loop_invariant!(i < 20);
+        for i in (0u8..10u8).step_by(2) {
+            hax_lib::loop_invariant!(|i: u8| i <= 10);
             count += 1;
         }
     }
     fn enumerated_slice<T>(slice: &[T]) {
         let mut count = 0u64;
         for i in slice.into_iter().enumerate() {
-            hax_lib::loop_invariant!(false);
+            hax_lib::loop_invariant!(|i: usize| i <= 10);
             count += 2;
         }
     }
     fn enumerated_chunked_slice<T>(slice: &[T]) {
         let mut count = 0u64;
         for i in slice.chunks_exact(3).enumerate() {
-            hax_lib::loop_invariant!(fstar!("forall (x: nat). x >= 0"));
+            hax_lib::loop_invariant!(|i: usize| { fstar!("$i <= ${slice.len()}") });
             count += 3;
         }
     }

From c27e57e4539684a23c2243664a8dd1976d545269 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Mon, 19 Aug 2024 17:06:52 +0200
Subject: [PATCH 08/13] misc(proof-libs/f*): intro. `Rust_primitives.Hax.Folds`

---
 .../Rust_primitives.Hax.Folds.fsti            | 103 ++++++++++++++++++
 1 file changed, 103 insertions(+)
 create mode 100644 proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti

diff --git a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
new file mode 100644
index 000000000..f2b7d363e
--- /dev/null
+++ b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
@@ -0,0 +1,103 @@
+module Rust_primitives.Hax.Folds
+
+open Rust_primitives
+open Core.Ops.Range
+open FStar.Mul
+
+(**** `s.chunks_exact(chunk_size).enumerate()` *)
+/// Predicate that asserts a slice `s_chunk` is exactly the nth chunk
+/// of the sequence `s`
+let nth_chunk_of #t
+  (s: Seq.seq t)
+  (s_chunk: Seq.seq t {Seq.length s_chunk > 0})
+  (chunk_nth: nat {chunk_nth < Seq.length s / Seq.length s_chunk})
+  =  Seq.slice s (Seq.length s_chunk * chunk_nth) (Seq.length s_chunk * (chunk_nth + 1))
+  == s_chunk
+
+/// Fold function that is generated for `for` loops iterating on
+/// `s.chunks_exact(chunk_size).enumerate()`-like iterators
+val fold_enumerated_chunked_slice
+  (#t: eqtype) (#acc_t: eqtype)
+  (chunk_size: usize {v chunk_size > 0})
+  (s: t_Slice t)
+  (inv: acc_t -> (i:usize{v i <= v (length s)}) -> Type0)
+  (init: acc_t {inv init (sz 0)})
+  (f: ( acc:acc_t
+      -> item:(usize & t_Array t chunk_size) {
+        let (i, s_chunk) = item in
+          v i < Seq.length s / v chunk_size
+        /\ nth_chunk_of s s_chunk (v i)
+        /\ v i < v (length s) 
+        /\ inv acc i
+      }
+      -> acc':acc_t {
+        v (fst item) < v (length s) /\ inv acc' (fst item)
+      }
+      )
+  )
+  : result: acc_t {inv result (length s)}
+
+(**** `s.enumerate()` *)
+/// Fold function that is generated for `for` loops iterating on
+/// `s.enumerate()`-like iterators
+val fold_enumerated_slice
+  (#t: eqtype) (#acc_t: eqtype)
+  (s: t_Slice t)
+  (inv: acc_t -> (i:usize{v i <= v (length s)}) -> Type0)
+  (init: acc_t {inv init (sz 0)})
+  (f: (acc:acc_t -> i:(usize & t) {v (fst i) < v (length s) /\ inv acc  (fst i)}
+                 -> acc':acc_t    {v (fst i) < v (length s) /\ inv acc' (fst i)}))
+  : result: acc_t {inv result (length s)}
+
+(**** `(start..end_).step_by(step)` *)
+unfold let fold_range_step_by_wf_index (#u: Lib.IntTypes.inttype)
+  (start: int_t u) (end_: int_t u {v start <= v end_})
+  (step: usize {v step > 0}) (strict: bool) (i: int)
+  = i >= v start 
+  /\ (if strict then i < v end_ else i <= v end_ + v step)
+  // /\ i % v step == v start % v step
+
+#push-options "--z3rlimit 80"
+unfold let fold_range_step_by_upper_bound (#u: Lib.IntTypes.inttype)
+  (start: int_t u) (end_: int_t u {v start <= v end_})
+  (step: usize {v step > 0})
+  : end':int {fold_range_step_by_wf_index start end_ step false end'}
+  = if end_ = start 
+    then v end_
+    else
+      let range: nat = v end_ - v start in
+      let k: nat = range / v step in
+      let end' = v start + k * v step in
+      FStar.Math.Lemmas.division_propriety range (v step);
+      end'
+#pop-options
+
+/// Fold function that is generated for `for` loops iterating on
+/// `s.enumerate()`-like iterators
+val fold_range_step_by
+  (#acc_t: eqtype) (#u: Lib.IntTypes.inttype)
+  (start: int_t u)
+  (end_: int_t u {v start <= v end_})
+  (step: usize {v step > 0 /\ range (v end_ + v step) u})
+  (inv: acc_t -> (i:int_t u{fold_range_step_by_wf_index start end_ step false (v i)}) -> Type0)
+  (init: acc_t {inv init start})
+  (f: (acc:acc_t -> i:int_t u  {fold_range_step_by_wf_index start end_ step true (v i) /\ inv acc i}
+                 -> acc':acc_t {(inv acc (mk_int (v i + v step)))}))
+  : result: acc_t {inv result (mk_int (fold_range_step_by_upper_bound start end_ step))}
+
+(**** `start..end_` *)
+unfold let fold_range_wf_index (#u: Lib.IntTypes.inttype)
+  (start: int_t u) (end_: int_t u {v start <= v end_})
+  (strict: bool) (i: int)
+  = i >= v start 
+  /\ (if strict then i < v end_ else i <= v end_)
+
+val fold_range
+  (#acc_t: eqtype) (#u: Lib.IntTypes.inttype)
+  (start: int_t u)
+  (end_: int_t u {v start <= v end_})
+  (inv: acc_t -> (i:int_t u{fold_range_wf_index start end_ false (v i)}) -> Type0)
+  (init: acc_t {inv init start})
+  (f: (acc:acc_t -> i:int_t u  {fold_range_wf_index start end_ true (v i) /\ inv acc i}
+                 -> acc':acc_t {(inv acc (mk_int (v i + 1)))}))
+  : result: acc_t {inv result end_}

From 719886a08099600db7d6fc7dab431289f92a6dfe Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Tue, 20 Aug 2024 08:55:07 +0200
Subject: [PATCH 09/13] fix test

---
 .../snapshots/toolchain__loops into-fstar.snap  | 17 +----------------
 1 file changed, 1 insertion(+), 16 deletions(-)

diff --git a/test-harness/src/snapshots/toolchain__loops into-fstar.snap b/test-harness/src/snapshots/toolchain__loops into-fstar.snap
index 62b4fad08..695f8bf5b 100644
--- a/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
+++ b/test-harness/src/snapshots/toolchain__loops into-fstar.snap	
@@ -21,22 +21,7 @@ info:
     backend_options: ~
 ---
 exit = 0
-stderr = '''
-warning: unnecessary parentheses around `for` iterator expression
- --> loops/src/lib.rs:4:18
-  |
-4 |         for i in (0u8..10u8) {
-  |                  ^         ^
-  |
-  = note: `#[warn(unused_parens)]` on by default
-help: remove these parentheses
-  |
-4 -         for i in (0u8..10u8) {
-4 +         for i in 0u8..10u8 {
-  |
-
-warning: `loops` (lib) generated 1 warning (run `cargo fix --lib -p loops` to apply 1 suggestion)
-    Finished `dev` profile [unoptimized + debuginfo] target(s) in XXs'''
+stderr = 'Finished `dev` profile [unoptimized + debuginfo] target(s) in XXs'
 
 [stdout]
 diagnostics = []

From c240db045f02c54add4e3b7cff17ee0e6936d0f6 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Tue, 20 Aug 2024 09:30:22 +0200
Subject: [PATCH 10/13] fix(proof-libs): empty range are alloed

---
 .../Rust_primitives.Hax.Folds.fsti            | 25 ++++++++++---------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
index f2b7d363e..3ece3bf2d 100644
--- a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
+++ b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
@@ -51,18 +51,18 @@ val fold_enumerated_slice
 
 (**** `(start..end_).step_by(step)` *)
 unfold let fold_range_step_by_wf_index (#u: Lib.IntTypes.inttype)
-  (start: int_t u) (end_: int_t u {v start <= v end_})
+  (start: int_t u) (end_: int_t u)
   (step: usize {v step > 0}) (strict: bool) (i: int)
-  = i >= v start 
-  /\ (if strict then i < v end_ else i <= v end_ + v step)
+  = v start <= v end_ ==> ( i >= v start 
+                        /\ (if strict then i < v end_ else i <= v end_ + v step))
   // /\ i % v step == v start % v step
 
 #push-options "--z3rlimit 80"
 unfold let fold_range_step_by_upper_bound (#u: Lib.IntTypes.inttype)
-  (start: int_t u) (end_: int_t u {v start <= v end_})
+  (start: int_t u) (end_: int_t u)
   (step: usize {v step > 0})
   : end':int {fold_range_step_by_wf_index start end_ step false end'}
-  = if end_ = start 
+  = if v end_ <= v start 
     then v end_
     else
       let range: nat = v end_ - v start in
@@ -77,27 +77,28 @@ unfold let fold_range_step_by_upper_bound (#u: Lib.IntTypes.inttype)
 val fold_range_step_by
   (#acc_t: eqtype) (#u: Lib.IntTypes.inttype)
   (start: int_t u)
-  (end_: int_t u {v start <= v end_})
+  (end_: int_t u)
   (step: usize {v step > 0 /\ range (v end_ + v step) u})
   (inv: acc_t -> (i:int_t u{fold_range_step_by_wf_index start end_ step false (v i)}) -> Type0)
   (init: acc_t {inv init start})
-  (f: (acc:acc_t -> i:int_t u  {fold_range_step_by_wf_index start end_ step true (v i) /\ inv acc i}
+  (f: (acc:acc_t -> i:int_t u  {v i < v end_ /\ fold_range_step_by_wf_index start end_ step true (v i) /\ inv acc i}
                  -> acc':acc_t {(inv acc (mk_int (v i + v step)))}))
   : result: acc_t {inv result (mk_int (fold_range_step_by_upper_bound start end_ step))}
 
 (**** `start..end_` *)
 unfold let fold_range_wf_index (#u: Lib.IntTypes.inttype)
-  (start: int_t u) (end_: int_t u {v start <= v end_})
+  (start: int_t u) (end_: int_t u)
   (strict: bool) (i: int)
-  = i >= v start 
-  /\ (if strict then i < v end_ else i <= v end_)
+  = v start <= v end_
+  ==> ( i >= v start 
+     /\ (if strict then i < v end_ else i <= v end_))
 
 val fold_range
   (#acc_t: eqtype) (#u: Lib.IntTypes.inttype)
   (start: int_t u)
-  (end_: int_t u {v start <= v end_})
+  (end_: int_t u)
   (inv: acc_t -> (i:int_t u{fold_range_wf_index start end_ false (v i)}) -> Type0)
   (init: acc_t {inv init start})
-  (f: (acc:acc_t -> i:int_t u  {fold_range_wf_index start end_ true (v i) /\ inv acc i}
+  (f: (acc:acc_t -> i:int_t u  {v i <= v end_ /\ fold_range_wf_index start end_ true (v i) /\ inv acc i}
                  -> acc':acc_t {(inv acc (mk_int (v i + 1)))}))
   : result: acc_t {inv result end_}

From 6f7236574d498039c9fc80c49b273d9aef2f3990 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Tue, 20 Aug 2024 10:39:37 +0200
Subject: [PATCH 11/13] fix(proof-libs): fix bad invariant we discovered with
 @mamonet

---
 .../rust_primitives/Rust_primitives.Hax.Folds.fsti   | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
index 3ece3bf2d..d68ff936d 100644
--- a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
+++ b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
@@ -17,7 +17,7 @@ let nth_chunk_of #t
 /// Fold function that is generated for `for` loops iterating on
 /// `s.chunks_exact(chunk_size).enumerate()`-like iterators
 val fold_enumerated_chunked_slice
-  (#t: eqtype) (#acc_t: eqtype)
+  (#t: Type0) (#acc_t: Type0)
   (chunk_size: usize {v chunk_size > 0})
   (s: t_Slice t)
   (inv: acc_t -> (i:usize{v i <= v (length s)}) -> Type0)
@@ -41,7 +41,7 @@ val fold_enumerated_chunked_slice
 /// Fold function that is generated for `for` loops iterating on
 /// `s.enumerate()`-like iterators
 val fold_enumerated_slice
-  (#t: eqtype) (#acc_t: eqtype)
+  (#t: Type0) (#acc_t: Type0)
   (s: t_Slice t)
   (inv: acc_t -> (i:usize{v i <= v (length s)}) -> Type0)
   (init: acc_t {inv init (sz 0)})
@@ -75,14 +75,14 @@ unfold let fold_range_step_by_upper_bound (#u: Lib.IntTypes.inttype)
 /// Fold function that is generated for `for` loops iterating on
 /// `s.enumerate()`-like iterators
 val fold_range_step_by
-  (#acc_t: eqtype) (#u: Lib.IntTypes.inttype)
+  (#acc_t: Type0) (#u: Lib.IntTypes.inttype)
   (start: int_t u)
   (end_: int_t u)
   (step: usize {v step > 0 /\ range (v end_ + v step) u})
   (inv: acc_t -> (i:int_t u{fold_range_step_by_wf_index start end_ step false (v i)}) -> Type0)
   (init: acc_t {inv init start})
   (f: (acc:acc_t -> i:int_t u  {v i < v end_ /\ fold_range_step_by_wf_index start end_ step true (v i) /\ inv acc i}
-                 -> acc':acc_t {(inv acc (mk_int (v i + v step)))}))
+                 -> acc':acc_t {(inv acc' (mk_int (v i + v step)))}))
   : result: acc_t {inv result (mk_int (fold_range_step_by_upper_bound start end_ step))}
 
 (**** `start..end_` *)
@@ -94,11 +94,11 @@ unfold let fold_range_wf_index (#u: Lib.IntTypes.inttype)
      /\ (if strict then i < v end_ else i <= v end_))
 
 val fold_range
-  (#acc_t: eqtype) (#u: Lib.IntTypes.inttype)
+  (#acc_t: Type0) (#u: Lib.IntTypes.inttype)
   (start: int_t u)
   (end_: int_t u)
   (inv: acc_t -> (i:int_t u{fold_range_wf_index start end_ false (v i)}) -> Type0)
   (init: acc_t {inv init start})
   (f: (acc:acc_t -> i:int_t u  {v i <= v end_ /\ fold_range_wf_index start end_ true (v i) /\ inv acc i}
-                 -> acc':acc_t {(inv acc (mk_int (v i + 1)))}))
+                 -> acc':acc_t {(inv acc' (mk_int (v i + 1)))}))
   : result: acc_t {inv result end_}

From 07841180bc8883689d8c46adaed6919eb9cc43b0 Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas@cryspen.com>
Date: Wed, 21 Aug 2024 11:17:02 +0200
Subject: [PATCH 12/13] fix(engine/iterators): support `into_iter` on arrays

---
 engine/lib/phases/phase_functionalize_loops.ml | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/engine/lib/phases/phase_functionalize_loops.ml b/engine/lib/phases/phase_functionalize_loops.ml
index 35591e89f..07b010714 100644
--- a/engine/lib/phases/phase_functionalize_loops.ml
+++ b/engine/lib/phases/phase_functionalize_loops.ml
@@ -76,7 +76,7 @@ struct
       | StepBy of { n : B.expr; it : iterator }
     [@@deriving show]
 
-    let rec as_iterator (e : B.expr) : iterator option =
+    let rec as_iterator' (e : B.expr) : iterator option =
       match e.e with
       | Construct
           {
@@ -93,6 +93,12 @@ struct
           Some (Range { start; end_ })
       | _ -> meth_as_iterator e
 
+    and as_iterator (e : B.expr) : iterator option =
+      let result = as_iterator' e in
+      (* UB.Debug.expr ~label:"as_iterator" e; *)
+      (* " = " ^ [%show: iterator option] result |> Stdio.prerr_endline; *)
+      result
+
     and meth_as_iterator (e : B.expr) : iterator option =
       let* f, args =
         match e.e with
@@ -112,7 +118,7 @@ struct
       then
         let* iterable = one_arg () in
         match iterable.typ with
-        | TSlice _ -> Some (Slice iterable)
+        | TSlice _ | TArray _ -> Some (Slice iterable)
         | _ -> as_iterator iterable
       else if f_eq Core__iter__traits__iterator__Iterator__enumerate then
         let* iterable = one_arg () in

From 1140b8bccbb6132e7109fc53497424f04787c102 Mon Sep 17 00:00:00 2001
From: mamonet <maamoun.tk@gmail.com>
Date: Thu, 22 Aug 2024 16:21:30 +0000
Subject: [PATCH 13/13] Fix fold_enumerated_chunked_slice signature

---
 .../rust_primitives/Rust_primitives.Hax.Folds.fsti     | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
index d68ff936d..484cf9186 100644
--- a/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
+++ b/proof-libs/fstar/rust_primitives/Rust_primitives.Hax.Folds.fsti
@@ -20,22 +20,22 @@ val fold_enumerated_chunked_slice
   (#t: Type0) (#acc_t: Type0)
   (chunk_size: usize {v chunk_size > 0})
   (s: t_Slice t)
-  (inv: acc_t -> (i:usize{v i <= v (length s)}) -> Type0)
+  (inv: acc_t -> (i:usize{v i <= Seq.length s / v chunk_size}) -> Type0)
   (init: acc_t {inv init (sz 0)})
   (f: ( acc:acc_t
-      -> item:(usize & t_Array t chunk_size) {
+      -> item:(usize & t_Slice t) {
         let (i, s_chunk) = item in
           v i < Seq.length s / v chunk_size
+        /\ length s_chunk == chunk_size
         /\ nth_chunk_of s s_chunk (v i)
-        /\ v i < v (length s) 
         /\ inv acc i
       }
       -> acc':acc_t {
-        v (fst item) < v (length s) /\ inv acc' (fst item)
+        inv acc' (fst item +! sz 1)
       }
       )
   )
-  : result: acc_t {inv result (length s)}
+  : result: acc_t {inv result (mk_int (Seq.length s / v chunk_size))}
 
 (**** `s.enumerate()` *)
 /// Fold function that is generated for `for` loops iterating on