Fix some bugs that surfaced in the ASL export

- Bitvector pattern rewriting had stopped working due to a line of code being
  lost in some merge.
- Fix a bug in early return rewriting that caused returns getting pulled out of
  if-statements to disappear.
- There were some variable name clashes with keywords because doc_lem_id was
  not always called.
- Ast_util.is_number failed to check for "int" and "nat" built-in types,
  causing pattern matching on natural number literals to fail.

3 files changed, 20 insertions, 10 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index e6217526..daaf5725 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -655,6 +655,8 @@ let rec lexp_to_exp (LEXP_aux (lexp_aux, annot) as le) =
 
 let rec is_number (Typ_aux (t,_)) =
   match t with
+  | Typ_id (Id_aux (Id "int", _))
+  | Typ_id (Id_aux (Id "nat", _))
   | Typ_app (Id_aux (Id "range", _),_)
   | Typ_app (Id_aux (Id "implicit", _),_)
   | Typ_app (Id_aux (Id "atom", _),_) -> true
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 5f520dfa..5ca0c1bc 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -74,6 +74,7 @@ let fix_id remove_tick name = match name with
   | "try"
   | "match"
   | "with"
+  | "check"
   | "field"
   | "LT"
   | "GT"
@@ -344,7 +345,7 @@ let rec doc_pat_lem sequential mwords apat_needed (P_aux (p,(l,annot)) as pa) =
      begin match id with
      | Id_aux (Id "None",_) -> string "Nothing" (* workaround temporary issue *)
      | _ -> doc_id_lem id end
-  | P_var(p,kid) -> parens (separate space [doc_pat_lem sequential mwords true p; string "as"; doc_var_lem kid])
+  | P_var(p,kid) -> doc_pat_lem sequential mwords true p
   | P_as(p,id) -> parens (separate space [doc_pat_lem sequential mwords true p; string "as"; doc_id_lem id])
   | P_typ(typ,p) ->
     let doc_p = doc_pat_lem sequential mwords true p in
@@ -521,12 +522,12 @@ let doc_exp_lem, doc_let_lem =
           let [id;indices;body;e5] = args in
           let varspp = match e5 with
             | E_aux (E_tuple vars,_) ->
-               let vars = List.map (fun (E_aux (E_id (Id_aux (Id name,_)),_)) -> string name) vars in
+               let vars = List.map (fun (E_aux (E_id id,_)) -> doc_id_lem id) vars in
                begin match vars with
                | [v] -> v
                | _ -> parens (separate comma vars) end
-            | E_aux (E_id (Id_aux (Id name,_)),_) ->
-               string name
+            | E_aux (E_id id,_) ->
+               doc_id_lem id
             | E_aux (E_lit (L_aux (L_unit,_)),_) ->
                string "_" in
           parens (
@@ -541,12 +542,12 @@ let doc_exp_lem, doc_let_lem =
           let [is_while;cond;body;e5] = args in
           let varspp = match e5 with
             | E_aux (E_tuple vars,_) ->
-               let vars = List.map (fun (E_aux (E_id (Id_aux (Id name,_)),_)) -> string name) vars in
+               let vars = List.map (fun (E_aux (E_id id,_)) -> doc_id_lem id) vars in
                begin match vars with
                | [v] -> v
                | _ -> parens (separate comma vars) end
-            | E_aux (E_id (Id_aux (Id name,_)),_) ->
-               string name
+            | E_aux (E_id id,_) ->
+               doc_id_lem id
             | E_aux (E_lit (L_aux (L_unit,_)),_) ->
                string "_" in
           parens (
diff --git a/src/rewriter.ml b/src/rewriter.ml
index 2028aa9c..bcc4e60a 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -1854,7 +1854,8 @@ let remove_bitvector_pat pat =
           let t = Env.base_typ_of env (typ_of_annot annot) in
           let (l,_) = annot in
           match pat, is_bitvector_typ t, contained_in_p_as with
-          | P_vector _, true, false
+          | P_vector _, true, false ->
+            P_aux (P_as (P_aux (pat,annot),fresh_id "b__" l), annot)
           | _ -> P_aux (pat,annot)
         )
     ; fP_aux = (fun (fpat,annot) -> FP_aux (fpat,annot))
@@ -2332,7 +2333,9 @@ let rewrite_defs_early_return =
     | _ -> E_block es in
 
   let e_if (e1, e2, e3) =
-    if is_return e2 && is_return e3 then E_if (e1, get_return e2, get_return e3)
+    if is_return e2 && is_return e3 then
+      let (E_aux (_, annot)) = e2 in
+      E_return (E_aux (E_if (e1, get_return e2, get_return e3), annot))
     else E_if (e1, e2, e3) in
 
   let e_case (e, pes) =
@@ -2341,8 +2344,12 @@ let rewrite_defs_early_return =
     let get_return_pexp (Pat_aux (pexp, a)) = match pexp with
     | Pat_exp (p, e) -> Pat_aux (Pat_exp (p, get_return e), a)
     | Pat_when (p, g, e) -> Pat_aux (Pat_when (p, g, get_return e), a) in
+    let annot = match pes with
+    | Pat_aux (Pat_exp (_, E_aux (_, annot)), _) :: _ -> annot
+    | Pat_aux (Pat_when (_, _, E_aux (_, annot)), _) :: _ -> annot
+    | [] -> (Parse_ast.Unknown, None) in
     if List.for_all is_return_pexp pes
-    then E_return (E_aux (E_case (e, List.map get_return_pexp pes), (Parse_ast.Unknown, None)))
+    then E_return (E_aux (E_case (e, List.map get_return_pexp pes), annot))
     else E_case (e, pes) in
 
   let e_aux (exp, (l, annot)) =