Tweak types of loop combinators for prover combinators

Split the variable (tuple) type into an input and output type.  They are
meant to be the same, but due to the way function types are
instantiated, unification can fail in the case of existential types, as
in the added test case (when trying to generate Lem definitions from
it).  The output of the loop will be checked against the expected type,
though, due to a type annotation outside the loop added by the rewrite
pass for variable updates.

2 files changed, 22 insertions, 5 deletions

diff --git a/src/rewrites.ml b/src/rewrites.ml
index ffa790f5..abb3e4ed 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -3785,11 +3785,11 @@ let rewrite_defs_remove_superfluous_returns env =
 
 let rewrite_defs_remove_e_assign env (Defs defs) =
   let (Defs loop_specs) = fst (Type_error.check initial_env (Defs (List.map gen_vs
-    [("foreach#", "forall ('vars : Type). (int, int, int, bool, 'vars, 'vars) -> 'vars");
-     ("while#", "forall ('vars : Type). (bool, 'vars, 'vars) -> 'vars");
-     ("until#", "forall ('vars : Type). (bool, 'vars, 'vars) -> 'vars");
-     ("while#t", "forall ('vars : Type). (bool, 'vars, 'vars, int) -> 'vars");
-     ("until#t", "forall ('vars : Type). (bool, 'vars, 'vars, int) -> 'vars")]))) in
+    [("foreach#", "forall ('vars_in 'vars_out : Type). (int, int, int, bool, 'vars_in, 'vars_out) -> 'vars_out");
+     ("while#", "forall ('vars_in 'vars_out : Type). (bool, 'vars_in, 'vars_out) -> 'vars_out");
+     ("until#", "forall ('vars_in 'vars_out : Type). (bool, 'vars_in, 'vars_out) -> 'vars_out");
+     ("while#t", "forall ('vars_in 'vars_out : Type). (bool, 'vars_in, 'vars_out, int) -> 'vars_out");
+     ("until#t", "forall ('vars_in 'vars_out : Type). (bool, 'vars_in, 'vars_out, int) -> 'vars_out")]))) in
  let rewrite_exp _ e =
     replace_memwrite_e_assign (remove_reference_types (rewrite_var_updates e)) in
   rewrite_defs_base
diff --git a/test/typecheck/pass/floor_pow2.sail b/test/typecheck/pass/floor_pow2.sail
new file mode 100644
index 00000000..fa8680cf
--- /dev/null
+++ b/test/typecheck/pass/floor_pow2.sail
@@ -0,0 +1,17 @@
+$include <arith.sail>
+
+val "pow2" : forall 'n, 'n >= 0. int('n) -> int(2 ^ 'n)
+
+val floor_pow2 : forall ('x : Int), 'x >= 0. int('x) -> int
+
+function floor_pow2 x = {
+    if x == 0 then {
+        return(0);
+    } else {
+        n : {'n, 'n >= 1. int('n)} = 1;
+        while x >= pow2(n) do {
+            n = n + 1
+        };
+        return(pow2(n - 1))
+    }
+}