Fix bug in rewriting of loops

Take into account existential types when determining bounds for the loop
variable

6 files changed, 31 insertions, 20 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index e003a1bc..2f2d27d8 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -818,14 +818,6 @@ let rec lexp_to_exp (LEXP_aux (lexp_aux, annot) as le) =
   | LEXP_memory (id, exps) -> rewrap (E_app (id, exps))
   | LEXP_deref exp -> rewrap (E_app (mk_id "reg_deref", [exp]))
 
-let destruct_range (Typ_aux (typ_aux, _)) =
-  match typ_aux with
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp n, _)])
-       when string_of_id f = "atom" -> Some (n, n)
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp n1, _); Typ_arg_aux (Typ_arg_nexp n2, _)])
-       when string_of_id f = "range" -> Some (n1, n2)
-  | _ -> None
-
 let is_unit_typ = function
   | Typ_aux (Typ_id u, _) -> string_of_id u = "unit"
   | _ -> false
diff --git a/src/ast_util.mli b/src/ast_util.mli
index 4ceeea44..c07732bf 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -287,7 +287,6 @@ val is_bitvector_typ : typ -> bool
 val typ_app_args_of : typ -> string * typ_arg_aux list * Ast.l
 val vector_typ_args_of : typ -> nexp * order * typ
 val vector_start_index : typ -> nexp
-val destruct_range : typ -> (nexp * nexp) option
 
 val is_order_inc : order -> bool
 
diff --git a/src/c_backend.ml b/src/c_backend.ml
index 6e1083df..5cf282f9 100644
--- a/src/c_backend.ml
+++ b/src/c_backend.ml
@@ -801,9 +801,9 @@ let rec ctyp_of_typ ctx typ =
   | Typ_id id when string_of_id id = "nat" -> CT_mpz
   | Typ_app (id, _) when string_of_id id = "range" || string_of_id id = "atom" ->
      begin
-       match destruct_range typ with
+       match destruct_range Env.empty typ with
        | None -> assert false (* Checked if range type in guard *)
-       | Some (n, m) ->
+       | Some (kids, constr, n, m) ->
           match nexp_simp n, nexp_simp m with
           | Nexp_aux (Nexp_constant n, _), Nexp_aux (Nexp_constant m, _)
                when Big_int.less_equal min_int64 n && Big_int.less_equal m max_int64 ->
diff --git a/src/rewrites.ml b/src/rewrites.ml
index 4a213970..cfeae1e3 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -2889,16 +2889,23 @@ let rec rewrite_var_updates ((E_aux (expaux,((l,_) as annot))) as exp) =
          |> mk_var_exps_pats pl env
        in
        let exp4 = rewrite_var_updates (add_vars overwrite exp4 vars) in
-       let ord_exp, lower, upper = match destruct_range (typ_of exp1), destruct_range (typ_of exp2) with
-         | None, _ | _, None ->
-            raise (Reporting_basic.err_unreachable el "Could not determine loop bounds")
-         | Some (l1, u1), Some (l2, u2) ->
-            if is_order_inc order
-            then (annot_exp (E_lit (mk_lit L_true)) el env bool_typ, l1, u2)
-            else (annot_exp (E_lit (mk_lit L_false)) el env bool_typ, l2, u1) in
+       let ord_exp, kids, constr, lower, upper =
+         match destruct_range env (typ_of exp1), destruct_range env (typ_of exp2) with
+           | None, _ | _, None ->
+              raise (Reporting_basic.err_unreachable el "Could not determine loop bounds")
+           | Some (kids1, constr1, l1, u1), Some (kids2, constr2, l2, u2) ->
+              let kids = kids1 @ kids2 in
+              let constr = nc_and constr1 constr2 in
+              let ord_exp, lower, upper =
+                if is_order_inc order
+                then (annot_exp (E_lit (mk_lit L_true)) el env bool_typ, l1, u2)
+                else (annot_exp (E_lit (mk_lit L_false)) el env bool_typ, l2, u1)
+              in
+              ord_exp, kids, constr, lower, upper
+       in
        let lvar_kid = mk_kid ("loop_" ^ string_of_id id) in
-       let lvar_nc = nc_and (nc_lteq lower (nvar lvar_kid)) (nc_lteq (nvar lvar_kid) upper) in
-       let lvar_typ = mk_typ (Typ_exist ([lvar_kid], lvar_nc, atom_typ (nvar lvar_kid))) in
+       let lvar_nc = nc_and constr (nc_and (nc_lteq lower (nvar lvar_kid)) (nc_lteq (nvar lvar_kid) upper)) in
+       let lvar_typ = mk_typ (Typ_exist (lvar_kid :: kids, lvar_nc, atom_typ (nvar lvar_kid))) in
        let lvar_pat = unaux_pat (add_p_typ lvar_typ (annot_pat (P_var (
          annot_pat (P_id id) el env (atom_typ (nvar lvar_kid)),
          TP_aux (TP_var lvar_kid, gen_loc el))) el env lvar_typ)) in
diff --git a/src/type_check.ml b/src/type_check.ml
index 52db43da..5c72983a 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -1097,6 +1097,17 @@ let bind_existential typ env =
             | Some (kids, nc, typ) -> typ, add_existential kids nc env
             | None -> typ, env
 
+let destruct_range env typ =
+  let kids, constr, (Typ_aux (typ_aux, _)) =
+    Util.option_default ([], nc_true, typ) (destruct_exist env typ)
+  in
+  match typ_aux with
+    | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp n, _)])
+         when string_of_id f = "atom" -> Some (kids, constr, n, n)
+    | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp n1, _); Typ_arg_aux (Typ_arg_nexp n2, _)])
+         when string_of_id f = "range" -> Some (kids, constr, n1, n2)
+    | _ -> None
+
 let destruct_vector env typ =
   let destruct_vector' = function
     | Typ_aux (Typ_app (id, [Typ_arg_aux (Typ_arg_nexp n1, _);
diff --git a/src/type_check.mli b/src/type_check.mli
index 82a30f4e..a047db9c 100644
--- a/src/type_check.mli
+++ b/src/type_check.mli
@@ -259,6 +259,8 @@ val destruct_atom_nexp : Env.t -> typ -> nexp option
    existential to ensure that no name-clashes occur. *)
 val destruct_exist : Env.t -> typ -> (kid list * n_constraint * typ) option
 
+val destruct_range : Env.t -> typ -> (kid list * n_constraint * nexp * nexp) option
+
 val destruct_vector : Env.t -> typ -> (nexp * order * typ) option
 
 type uvar =