Improve flow typing

Can now handle nexps such as (2**65 - 1). Uses big_ints for comparisons, and
keeps original nexps in the AST.

4 files changed, 88 insertions, 39 deletions

diff --git a/src/rewriter.ml b/src/rewriter.ml
index 79519af6..f906b597 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -1941,7 +1941,12 @@ let remove_bitvector_pat pat =
     (letexp, letbind) in
 
   let compose_guards guards =
-    List.fold_right (Util.option_binop bitwise_and_exp) guards None in
+    let conj g1 g2 = match g1, g2 with
+      | Some g1, Some g2 -> Some (bitwise_and_exp g1 g2)
+      | Some g1, None -> Some g1
+      | None, Some g2 -> Some g2
+      | None, None -> None in
+    List.fold_right conj guards None in
 
   let flatten_guards_decls gd =
     let (guards,decls,letbinds) = Util.split3 gd in
diff --git a/src/type_check.ml b/src/type_check.ml
index a3c4f767..14c91247 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -1686,13 +1686,31 @@ let pat_typ_of (P_aux (_, (l, tannot))) = typ_of_annot (l, tannot)
 
 (* Flow typing *)
 
+let rec big_int_of_nexp (Nexp_aux (nexp, _)) = match nexp with
+  | Nexp_constant c -> Some (big_int_of_int c)
+  | Nexp_times (n1, n2) ->
+     Util.option_binop add_big_int (big_int_of_nexp n1) (big_int_of_nexp n2)
+  | Nexp_sum (n1, n2) ->
+     Util.option_binop add_big_int (big_int_of_nexp n1) (big_int_of_nexp n2)
+  | Nexp_minus (n1, n2) ->
+     Util.option_binop add_big_int (big_int_of_nexp n1) (big_int_of_nexp n2)
+  | Nexp_exp n ->
+     Util.option_map (power_int_positive_big_int 2) (big_int_of_nexp n)
+  | _ -> None
+
 let destruct_atom (Typ_aux (typ_aux, _)) =
   match typ_aux with
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp (Nexp_aux (Nexp_constant c, _)), _)])
-       when string_of_id f = "atom" -> c
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp (Nexp_aux (Nexp_constant c1, _)), _); Typ_arg_aux (Typ_arg_nexp (Nexp_aux (Nexp_constant c2, _)), _)])
-       when string_of_id f = "range" && c1 = c2 -> c1
-  | _ -> assert false
+  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp, _)])
+       when string_of_id f = "atom" ->
+     Util.option_map (fun c -> (c, nexp)) (big_int_of_nexp nexp)
+  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp1, _); Typ_arg_aux (Typ_arg_nexp nexp2, _)])
+       when string_of_id f = "range" ->
+     begin
+       match big_int_of_nexp nexp1, big_int_of_nexp nexp2 with
+       | Some c1, Some c2 -> if eq_big_int c1 c2 then Some (c1, nexp1) else None
+       | _ -> None
+     end
+  | _ -> None
 
 let destruct_atom_nexp env typ =
   match Env.expand_synonyms env typ with
@@ -1702,20 +1720,6 @@ let destruct_atom_nexp env typ =
        when string_of_id f = "range" -> Some n
   | _ -> None
 
-let restrict_range_upper c1 (Typ_aux (typ_aux, l) as typ) =
-  match typ_aux with
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp, _); Typ_arg_aux (Typ_arg_nexp (Nexp_aux (Nexp_constant c2, _)), _)])
-     when string_of_id f = "range" ->
-     range_typ nexp (nconstant (min c1 c2))
-  | _ -> typ
-
-let restrict_range_lower c1 (Typ_aux (typ_aux, l) as typ) =
-  match typ_aux with
-  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp (Nexp_aux (Nexp_constant c2, _)), _); Typ_arg_aux (Typ_arg_nexp nexp, _)])
-     when string_of_id f = "range" ->
-     range_typ (nconstant (max c1 c2)) nexp
-  | _ -> typ
-
 exception Not_a_constraint;;
 
 let rec assert_nexp (E_aux (exp_aux, l)) =
@@ -1737,12 +1741,42 @@ let rec assert_constraint (E_aux (exp_aux, l)) =
   | _ -> nc_true
 
 type flow_constraint =
-  | Flow_lteq of int
-  | Flow_gteq of int
+  | Flow_lteq of big_int * nexp
+  | Flow_gteq of big_int * nexp
+
+let restrict_range_upper c1 nexp1 (Typ_aux (typ_aux, l) as typ) =
+  match typ_aux with
+  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp, _); Typ_arg_aux (Typ_arg_nexp nexp2, _)])
+     when string_of_id f = "range" ->
+     begin
+       match big_int_of_nexp nexp2 with
+       | Some c2 ->
+          let upper = if (lt_big_int c1 c2) then nexp1 else nexp2 in
+          range_typ nexp upper
+       | _ -> typ
+     end
+  | _ -> typ
+
+let restrict_range_lower c1 nexp1 (Typ_aux (typ_aux, l) as typ) =
+  match typ_aux with
+  | Typ_app (f, [Typ_arg_aux (Typ_arg_nexp nexp2, _); Typ_arg_aux (Typ_arg_nexp nexp, _)])
+     when string_of_id f = "range" ->
+     begin
+       match big_int_of_nexp nexp2 with
+       | Some c2 ->
+          let lower = if (gt_big_int c1 c2) then nexp1 else nexp2 in
+          range_typ lower nexp
+       | _ -> typ
+     end
+  | _ -> typ
 
 let apply_flow_constraint = function
-  | Flow_lteq c -> (restrict_range_upper c, restrict_range_lower (c + 1))
-  | Flow_gteq c -> (restrict_range_lower c, restrict_range_upper (c - 1))
+  | Flow_lteq (c, nexp) ->
+     (restrict_range_upper c nexp,
+      restrict_range_lower (succ_big_int c) (nexp_simp (nsum nexp (nconstant 1))))
+  | Flow_gteq (c, nexp) ->
+     (restrict_range_lower c nexp,
+      restrict_range_upper (pred_big_int c) (nexp_simp (nminus nexp (nconstant 1))))
 
 let rec infer_flow env (E_aux (exp_aux, (l, _))) =
   match exp_aux with
@@ -1764,20 +1798,34 @@ let rec infer_flow env (E_aux (exp_aux, (l, _))) =
      [], [nc_gt n1 n2]
   | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "lt_range_atom" ->
      let kid = Env.fresh_kid env in
-     let c = destruct_atom (typ_of y) in
-     [(v, Flow_lteq (c - 1))], []
+     begin
+       match destruct_atom (typ_of y) with
+       | Some (c, nexp) ->
+          [(v, Flow_lteq (pred_big_int c, nexp_simp (nminus nexp (nconstant 1))))], []
+       | _ -> [], []
+     end
   | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "lteq_range_atom" ->
      let kid = Env.fresh_kid env in
-     let c = destruct_atom (typ_of y) in
-     [(v, Flow_lteq c)], []
+     begin
+       match destruct_atom (typ_of y) with
+       | Some (c, nexp) -> [(v, Flow_lteq (c, nexp))], []
+       | _ -> [], []
+     end
   | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "gt_range_atom" ->
      let kid = Env.fresh_kid env in
-     let c = destruct_atom (typ_of y) in
-     [(v, Flow_gteq (c + 1))], []
+     begin
+       match destruct_atom (typ_of y) with
+       | Some (c, nexp) ->
+          [(v, Flow_gteq (succ_big_int c, nexp_simp (nsum nexp (nconstant 1))))], []
+       | _ -> [], []
+     end
   | E_app (f, [E_aux (E_id v, _); y]) when string_of_id f = "gteq_range_atom" ->
      let kid = Env.fresh_kid env in
-     let c = destruct_atom (typ_of y) in
-     [(v, Flow_gteq c)], []
+     begin
+       match destruct_atom (typ_of y) with
+       | Some (c, nexp) -> [(v, Flow_gteq (c, nexp))], []
+       | _ -> [], []
+     end
   | _ -> [], []
 
 let rec add_flows b flows env =
diff --git a/src/util.ml b/src/util.ml
index 75732376..c89cc1ef 100644
--- a/src/util.ml
+++ b/src/util.ml
@@ -204,10 +204,8 @@ let option_bind f = function
   | Some(o) -> f o
 
 let rec option_binop f x y = match x, y with
-  | None, None -> None
-  | Some x, None -> Some x
-  | None, Some y -> Some y
   | Some x, Some y -> Some (f x y)
+  | _ -> None
 
 let changed2 f g x h y =
   match (g x, h y) with
diff --git a/src/util.mli b/src/util.mli
index aa442ada..f1182c61 100644
--- a/src/util.mli
+++ b/src/util.mli
@@ -77,10 +77,8 @@ val option_bind : ('a -> 'b option) -> 'a option -> 'b option
     whereas [option_default d (Some x)] returns [x]. *)
 val option_default : 'a -> 'a option -> 'a
 
-(** [option_binop f None None] returns [None], while
-    [option_binop f (Some x) None] and [option_binop f None (Some x)]
-    return [Some x], and [option_binop f (Some x) (Some y)] returns
-    [Some (f x y)] *)
+(** [option_binop f (Some x) (Some y)] returns [Some (f x y)],
+    and in all other cases, [option_binop] returns [None]. *)
 val option_binop : ('a -> 'a -> 'a) -> 'a option -> 'a option -> 'a option
 
 (** [option_get_exn exn None] throws the exception [exn],