Improve pretty-printing of register declaration and assignment

7 files changed, 58 insertions, 58 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index ade13fe3..df8f5af3 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -397,6 +397,13 @@ let rec nexp_identical (Nexp_aux (nexp1, _)) (Nexp_aux (nexp2, _)) =
   | Nexp_neg n1, Nexp_neg n2 -> nexp_identical n1 n2
   | _, _ -> false
 
+let rec is_nexp_constant (Nexp_aux (nexp, _)) = match nexp with
+| Nexp_id _ | Nexp_var _ -> false
+| Nexp_constant _ -> true
+| Nexp_times (n1, n2) | Nexp_sum (n1, n2) | Nexp_minus (n1, n2) ->
+  is_nexp_constant n1 && is_nexp_constant n2
+| Nexp_exp n | Nexp_neg n -> is_nexp_constant n
+
 let rec is_number (Typ_aux (t,_)) =
   match t with
   | Typ_app (Id_aux (Id "range", _),_)
diff --git a/src/ast_util.mli b/src/ast_util.mli
index 6e22d173..8a07b83f 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -119,6 +119,7 @@ end
 
 val nexp_frees : nexp -> KidSet.t
 val nexp_identical : nexp -> nexp -> bool
+val is_nexp_constant : nexp -> bool
 
 val is_number : typ -> bool
 val is_vector_typ : typ -> bool
diff --git a/src/pretty_print_common.ml b/src/pretty_print_common.ml
index a627328d..52c3753a 100644
--- a/src/pretty_print_common.ml
+++ b/src/pretty_print_common.ml
@@ -215,7 +215,7 @@ let doc_typ, doc_atomic_typ, doc_nexp =
       group (parens (nexp ne))
 
   (* expose doc_typ, doc_atomic_typ and doc_nexp *)
-  in typ, atomic_typ, nexp
+  in typ, atomic_typ, atomic_nexp_typ
 
 let pp_format_id (Id_aux(i,_)) =
   match i with
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 2619cc51..1f557e74 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -300,10 +300,6 @@ and contains_bitvector_typ_arg (Typ_arg_aux (targ, _)) = match targ with
   | Typ_arg_typ t -> contains_bitvector_typ t
   | _ -> false
 
-let const_nexp (Nexp_aux (nexp,_)) = match nexp with
-  | Nexp_constant _ -> true
-  | _ -> false
-
 (* Check for variables in types that would be pretty-printed.
    In particular, in case of vector types, only the element type and the
    length argument are checked for variables, and the latter only if it is
@@ -318,11 +314,11 @@ let rec contains_t_pp_var (Typ_aux (t,a) as typ) = match t with
   | Typ_app (c,targs) ->
       if is_bitvector_typ typ then
         let (_,length,_,_) = vector_typ_args_of typ in
-        not (const_nexp ((*normalize_nexp*) length))
+        not (is_nexp_constant ((*normalize_nexp*) length))
       else List.exists contains_t_arg_pp_var targs
 and contains_t_arg_pp_var (Typ_arg_aux (targ, _)) = match targ with
   | Typ_arg_typ t -> contains_t_pp_var t
-  | Typ_arg_nexp nexp -> not (const_nexp ((*normalize_nexp*) nexp))
+  | Typ_arg_nexp nexp -> not (is_nexp_constant ((*normalize_nexp*) nexp))
   | _ -> false
 
 let prefix_recordtype = true
@@ -1260,33 +1256,30 @@ let doc_dec_lem (DEC_aux (reg,(l,annot))) =
   match reg with
   | DEC_reg(typ,id) ->
      (match typ with
-      | Typ_aux (Typ_app (r, [Typ_arg_aux (Typ_arg_typ rt, _)]), _)
-        when string_of_id r = "register" && is_vector_typ rt ->
-        let env = env_of_annot (l,annot) in
-        let (start, size, order, etyp) = vector_typ_args_of (Env.base_typ_of env rt) in
-        (match is_bit_typ (Env.base_typ_of env etyp), start, size with
-          | true, Nexp_aux (Nexp_constant start, _), Nexp_aux (Nexp_constant size, _) ->
-             let o = if is_order_inc order then "true" else "false" in
-             (doc_op equals)
-             (string "let" ^^ space ^^ doc_id_lem id)
-             (string "Register" ^^ space ^^
-                align (separate space [string_lit(doc_id_lem id);
-                                       doc_int (size);
-                                       doc_int (start);
-                                       string o;
-                                       string "[]"]))
-                ^/^ hardline
-           | _ ->
-              let (Id_aux (Id name,_)) = id in
-              failwith ("can't deal with register " ^ name))
-      | Typ_aux (Typ_app(r, [Typ_arg_aux (Typ_arg_typ (Typ_aux (Typ_id idt, _)), _)]), _)
-        when string_of_id r = "register" ->
-         separate space [string "let";doc_id_lem id;equals;
-                         string "build_" ^^ string (string_of_id idt);string_lit (doc_id_lem id)] ^/^ hardline
       | Typ_aux (Typ_id idt, _) ->
-         separate space [string "let";doc_id_lem id;equals;
-                         string "build_" ^^ string (string_of_id idt);string_lit (doc_id_lem id)] ^/^ hardline
-      |_-> empty)
+        separate space [string "let";doc_id_lem id;equals;
+                        string "build_" ^^ string (string_of_id idt);string_lit (doc_id_lem id)] ^/^ hardline
+      | _ ->
+        let rt = (match annot with
+          | Some (env, _, _) -> Env.base_typ_of env typ
+          | _ -> typ) in
+        if is_vector_typ rt then
+          let (start, size, order, etyp) = vector_typ_args_of rt in
+          if is_bit_typ etyp && is_nexp_constant start && is_nexp_constant size then
+            let o = if is_order_inc order then "true" else "false" in
+            (doc_op equals)
+            (string "let" ^^ space ^^ doc_id_lem id)
+            (string "Register" ^^ space ^^
+               align (separate space [string_lit(doc_id_lem id);
+                                      doc_nexp (size);
+                                      doc_nexp (start);
+                                      string o;
+                                      string "[]"]))
+               ^/^ hardline
+          else raise (Reporting_basic.err_unreachable l
+            ("can't deal with register type " ^ string_of_typ typ))
+        else raise (Reporting_basic.err_unreachable l
+          ("can't deal with register type " ^ string_of_typ typ)))
   | DEC_alias(id,alspec) -> empty
   | DEC_typ_alias(typ,id,alspec) -> empty
 
diff --git a/src/pretty_print_ocaml.ml b/src/pretty_print_ocaml.ml
index 66252d94..fc02f568 100644
--- a/src/pretty_print_ocaml.ml
+++ b/src/pretty_print_ocaml.ml
@@ -741,28 +741,23 @@ let doc_dec_ocaml (DEC_aux (reg,(l,annot))) =
   | DEC_reg(typ,id) ->
     if is_vector_typ typ then
       let (start, size, order, itemt) = vector_typ_args_of typ in
-    (* (match annot with
-     | Base((_,t),_,_,_,_,_) ->
-       (match t.t with
-        | Tapp("register", [TA_typ {t= Tapp("vector", [TA_nexp start; TA_nexp size; TA_ord order; TA_typ itemt])}]) 
-        | Tapp("register", [TA_typ {t= Tabbrev(_,{t=Tapp("vector", [TA_nexp start; TA_nexp size; TA_ord order; TA_typ itemt])})}]) -> *)
-          (match is_bit_typ itemt, start, size with
-           | true, Nexp_aux (Nexp_constant start, _), Nexp_aux (Nexp_constant size, _) ->
-             let o = if is_order_inc order then string "true" else string "false" in
-             separate space [string "let";
-                             doc_id_ocaml id;
-                             equals;
-                             string "Vregister";
-                             parens (separate comma [separate space [string "ref";
-                                                                     parens (separate space
-                                                                               [string "Array.make";
-                                                                                doc_int size;
-                                                                                string "Vzero";])];
-                                                     doc_int start;
-                                                     o;
-                                                     string_lit (doc_id id);
-                                                     brackets empty])]
-           | _ -> empty)
+      if is_bit_typ itemt && is_nexp_constant start && is_nexp_constant size then
+        let o = if is_order_inc order then string "true" else string "false" in
+        separate space [string "let";
+                        doc_id_ocaml id;
+                        equals;
+                        string "Vregister";
+                        parens (separate comma [separate space [string "ref";
+                                                                parens (separate space
+                                                                          [string "Array.make";
+                                                                           doc_nexp size;
+                                                                           string "Vzero";])];
+                                                doc_nexp start;
+                                                o;
+                                                string_lit (doc_id id);
+                                                brackets empty])]
+      else raise (Reporting_basic.err_unreachable l
+        ("can't deal with register type " ^ string_of_typ typ))
     else
       (match typ with
         | Typ_aux (Typ_id idt, _) ->
@@ -773,7 +768,8 @@ let doc_dec_ocaml (DEC_aux (reg,(l,annot))) =
                           equals;
                           doc_id_ocaml idt;
                           string "None"]
-        |_-> failwith "type was not handled in register declaration")
+        |_-> raise (Reporting_basic.err_unreachable l
+          ("can't deal with register type " ^ string_of_typ typ)))
      (* | _ -> failwith "annot was not Base") *)
     | DEC_alias(id,alspec) -> empty (*
         doc_op equals (string "register alias" ^^ space ^^ doc_id id) (doc_alias alspec) *)
diff --git a/src/rewriter.ml b/src/rewriter.ml
index b1f54c5f..33fe3c25 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -2767,7 +2767,7 @@ let rec rewrite_var_updates ((E_aux (expaux,((l,_) as annot))) as exp) =
          | _ ->
            raise (Reporting_basic.err_unreachable l
              "assignment without effects annotation") in
-       if not (List.exists (function BE_aux (BE_lset,_) -> true | _ -> false) effs) then
+       if effectful exp then
          Same_vars (E_aux (E_assign (lexp,vexp),annot))
        else 
          (match lexp with
diff --git a/src/type_check.ml b/src/type_check.ml
index eb6eab09..2a23eb93 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -2008,10 +2008,11 @@ and bind_assignment env (LEXP_aux (lexp_aux, _) as lexp) (E_aux (_, (l, ())) as
             end
        in
        let regtyp, inferred_flexp, is_register = infer_flexp flexp in
-       let eff = if is_register then mk_effect [BE_wreg] else no_effect in
        typ_debug ("REGTYP: " ^ string_of_typ regtyp ^ " / " ^ string_of_typ (Env.expand_synonyms env regtyp));
        match Env.expand_synonyms env regtyp with
+       | Typ_aux (Typ_app (Id_aux (Id "register", _), [Typ_arg_aux (Typ_arg_typ (Typ_aux (Typ_id regtyp_id, _)), _)]), _)
        | Typ_aux (Typ_id regtyp_id, _) when Env.is_regtyp regtyp_id env ->
+          let eff = mk_effect [BE_wreg] in
           let base, top, ranges = Env.get_regtyp regtyp_id env in
           let range, _ =
             try List.find (fun (_, id) -> Id.compare id field = 0) ranges with
@@ -2027,6 +2028,7 @@ and bind_assignment env (LEXP_aux (lexp_aux, _) as lexp) (E_aux (_, (l, ())) as
           let checked_exp = crule check_exp env exp vec_typ in
           annot_assign (annot_lexp (LEXP_field (annot_lexp_effect inferred_flexp regtyp eff, field)) vec_typ) checked_exp, env
        | Typ_aux (Typ_id rectyp_id, _) | Typ_aux (Typ_app (rectyp_id, _), _) when Env.is_record rectyp_id env ->
+          let eff = if is_register then mk_effect [BE_wreg] else no_effect in
           let (typq, Typ_aux (Typ_fn (rectyp_q, field_typ, _), _)) = Env.get_accessor rectyp_id field env in
           let unifiers = try unify l env rectyp_q regtyp with Unification_error (l, m) -> typ_error l ("Unification error: " ^ m) in
           let field_typ' = subst_unifiers unifiers field_typ in
@@ -2180,6 +2182,7 @@ and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
        let inferred_exp = irule infer_exp env exp in
        match Env.expand_synonyms env (typ_of inferred_exp) with
        (* Accessing a (bit) field of a register *)
+       | Typ_aux (Typ_app (Id_aux (Id "register", _), [Typ_arg_aux (Typ_arg_typ (Typ_aux (Typ_id regtyp, _)), _)]), _)
        | Typ_aux (Typ_id regtyp, _) when Env.is_regtyp regtyp env ->
           let base, top, ranges = Env.get_regtyp regtyp env in
           let range, _ =