Merge pull request #280 from FStarLang/protz_unroll

Enable macro-based loop unrolling

include/krml/internal/target.h

@@ -110,4 +110,188 @@ inline static int32_t krml_time() {
 #  define KRML_DEPRECATED(x) __declspec(deprecated(x))
 #endif
 
+/* Macros for prettier unrolling of loops */
+#define KRML_LOOP1(i, n, x) { \
+  x \
+  i += n; \
+}
+
+#define KRML_LOOP2(i, n, x) \
+  KRML_LOOP1(i, n, x) \
+  KRML_LOOP1(i, n, x)
+
+#define KRML_LOOP3(i, n, x) \
+  KRML_LOOP2(i, n, x) \
+  KRML_LOOP1(i, n, x)
+
+#define KRML_LOOP4(i, n, x) \
+  KRML_LOOP2(i, n, x) \
+  KRML_LOOP2(i, n, x)
+
+#define KRML_LOOP5(i, n, x) \
+  KRML_LOOP4(i, n, x) \
+  KRML_LOOP1(i, n, x)
+
+#define KRML_LOOP6(i, n, x) \
+  KRML_LOOP4(i, n, x) \
+  KRML_LOOP2(i, n, x)
+
+#define KRML_LOOP7(i, n, x) \
+  KRML_LOOP4(i, n, x) \
+  KRML_LOOP3(i, n, x)
+
+#define KRML_LOOP8(i, n, x) \
+  KRML_LOOP4(i, n, x) \
+  KRML_LOOP4(i, n, x)
+
+#define KRML_LOOP9(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP1(i, n, x)
+
+#define KRML_LOOP10(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP2(i, n, x)
+
+#define KRML_LOOP11(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP3(i, n, x)
+
+#define KRML_LOOP12(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP4(i, n, x)
+
+#define KRML_LOOP13(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP5(i, n, x)
+
+#define KRML_LOOP14(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP6(i, n, x)
+
+#define KRML_LOOP15(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP7(i, n, x)
+
+#define KRML_LOOP16(i, n, x) \
+  KRML_LOOP8(i, n, x) \
+  KRML_LOOP8(i, n, x)
+
+#define KRML_UNROLL_FOR(i, z, n, k, x) do { \
+  uint32_t i = z; \
+  KRML_LOOP##n(i, k, x) \
+} while (0)
+
+#define KRML_ACTUAL_FOR(i, z, n, k, x) \
+  do { \
+    for (uint32_t i = z; i < n; i += k) { \
+      x \
+    } \
+  } while (0)
+
+#ifndef KRML_UNROLL_MAX
+#define KRML_UNROLL_MAX 16
+#endif
+
+/* 1 is the number of loop iterations, i.e. (n - z)/k as evaluated by krml */
+#if 0 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR0(i, z, n, k, x)
+#else
+#define KRML_MAYBE_FOR0(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 1 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR1(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 1, k, x)
+#else
+#define KRML_MAYBE_FOR1(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 2 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR2(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 2, k, x)
+#else
+#define KRML_MAYBE_FOR2(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 3 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR3(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 3, k, x)
+#else
+#define KRML_MAYBE_FOR3(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 4 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR4(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 4, k, x)
+#else
+#define KRML_MAYBE_FOR4(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 5 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR5(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 5, k, x)
+#else
+#define KRML_MAYBE_FOR5(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 6 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR6(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 6, k, x)
+#else
+#define KRML_MAYBE_FOR6(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 7 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR7(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 7, k, x)
+#else
+#define KRML_MAYBE_FOR7(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 8 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR8(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 8, k, x)
+#else
+#define KRML_MAYBE_FOR8(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 9 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR9(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 9, k, x)
+#else
+#define KRML_MAYBE_FOR9(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 10 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR10(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 10, k, x)
+#else
+#define KRML_MAYBE_FOR10(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 11 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR11(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 11, k, x)
+#else
+#define KRML_MAYBE_FOR11(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 12 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR12(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 12, k, x)
+#else
+#define KRML_MAYBE_FOR12(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 13 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR13(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 13, k, x)
+#else
+#define KRML_MAYBE_FOR13(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 14 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR14(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 14, k, x)
+#else
+#define KRML_MAYBE_FOR14(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 15 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR15(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 15, k, x)
+#else
+#define KRML_MAYBE_FOR15(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
+
+#if 16 <= KRML_UNROLL_MAX
+#define KRML_MAYBE_FOR16(i, z, n, k, x) KRML_UNROLL_FOR(i, z, 16, k, x)
+#else
+#define KRML_MAYBE_FOR16(i, z, n, k, x) KRML_ACTUAL_FOR(i, z, n, k, x)
+#endif
 #endif

src/AstToCStar.ml

@@ -316,44 +316,56 @@ and mk_stmts env e ret_type =
         let e = CStar.While (mk_expr env false e1, mk_block env Not e2) in
         env, maybe_return (e :: acc)
 
-    | EFor (_,
-      { node = EConstant (K.UInt32, init); _ },
+    | EFor (binder,
+      ({ node = EConstant (K.UInt32, init as k_init); _ } as e_init),
       { node = EApp (
         { node = EOp (K.Lt, K.UInt32); _ },
         [{ node = EBound 0; _ };
-        { node = EConstant (K.UInt32, max); _ }]); _},
+        ({ node = EConstant (K.UInt32, max as k_max); _ } as e_max)]); _},
       { node = EAssign (
         { node = EBound 0; _ },
         { node = EApp (
           { node = EOp (K.Add, K.UInt32); _ },
           [{ node = EBound 0; _ };
-          { node = EConstant (K.UInt32, incr); _ }]); _}); _},
+          ({ node = EConstant (K.UInt32, incr as k_incr); _ } as e_incr)]); _}); _},
       body)
       when (
         let init = int_of_string init in
         let max = int_of_string max in
         let incr = int_of_string incr in
-        let len = (max - init) / incr in
-        len <= !Options.unroll_loops
+        let n_loops = (max - init + incr - 1) / incr in
+        n_loops <= !Options.unroll_loops
       )
       ->
         let init = int_of_string init in
         let max = int_of_string max in
         let incr = int_of_string incr in
-        let rec mk acc i =
-          if i < max then
-            let body = DeBruijn.subst (mk_uint32 i) 0 body in
-            mk (CStar.Block (mk_block env Not body) :: acc) (i + incr)
-          else
-            acc
-        in
-        env, maybe_return (mk [] init @ acc)
+        let n_loops = (max - init + incr - 1) / incr in
+
+        if n_loops = 0 then
+          env, maybe_return acc
+
+        else if n_loops = 1 then
+          let body = DeBruijn.subst e_init 0 body in
+          let body = mk_block env Not body in
+          env, (Block body) :: acc
+
+        else begin
+          assert (n_loops <= 16);
+          let env, { CStar.name; _ } = mk_and_push_binder env binder Local None [ e_init; e_max; e_incr; body ] in
+          let body = mk_block env Not body in
+
+          env, maybe_return (CStar.Ignore (Call (
+            Macro (["KRML_MAYBE"], "FOR" ^ string_of_int n_loops),
+            [ Var name; Constant k_init; Constant k_max; Constant k_incr; Stmt body ])) :: acc)
+        end
 
 
     | EFor (binder, e1, e2, e3, e4) ->
         (* Note: the arguments to mk_and_push_binder are solely for the purpose
          * of avoiding name collisions. *)
         let is_solo_assignment = binder.node.meta = Some MetaSequence in
+        (* TODO: shouldn't e1 be added here? *)
         let env', binder = mk_and_push_binder env binder Local (Some e1) [ e2; e3; e4 ] in
         let e2 = mk_expr env' false e2 in
         let e3 = KList.last (mk_block env' Not e3) in

src/C11.ml

@@ -62,7 +62,8 @@ and expr =
   | MemberAccessPointer of expr * ident
   | InlineComment of string * expr * string
   | Type of type_name
-    (** note: these two not in the C grammar *)
+  | Stmt of stmt list
+    (** note: last three in the C grammar *)
 
 (** this is a WILD approximation of the notion of "type name" in C _and_ a hack
  * because there's the invariant that the ident found at the bottom of the
@@ -98,7 +99,7 @@ and designator =
  * declarations can only be part of a compound statement... we do not enforce
  * this invariant via the type [stmt], but rather check it at runtime (see
  * [mk_compound_if]), at the risk of messing things up, naturally. *)
-type stmt =
+and stmt =
   | Compound of stmt list
   | Decl of declaration
   | Expr of expr

src/CStar.ml

@@ -78,6 +78,10 @@ and expr =
   | Any
   | AddrOf of expr
   | EAbort of typ * string
+  | Stmt of stmt list
+    (** Solely for the purposes of macro-calls, which are function-like and
+        therefore are more closely modeled as an expression, but truly may
+        contain anything as arguments, including statements. *)
   [@@deriving show]
 
 and block =