add type annotation patterns to mpats

6 files changed, 20 insertions, 2 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index d571c916..a766b846 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -463,6 +463,7 @@ and map_mpat_annot_aux f = function
   | MP_vector mpats -> MP_vector (List.map (map_mpat_annot f) mpats)
   | MP_cons (mpat1, mpat2) -> MP_cons (map_mpat_annot f mpat1, map_mpat_annot f mpat2)
   | MP_string_append (mpat1, mpat2) -> MP_string_append (map_mpat_annot f mpat1, map_mpat_annot f mpat2)
+  | MP_typ (mpat, typ) -> MP_typ (map_mpat_annot f mpat, typ)
 
 and map_fpat_annot f (FP_aux (FP_Fpat (id, pat), annot)) = FP_aux (FP_Fpat (id, map_pat_annot f pat), f annot)
 and map_mfpat_annot f (MFP_aux (MFP_mpat (id, mpat), annot)) = MFP_aux (MFP_mpat (id, map_mpat_annot f mpat), f annot)
@@ -499,6 +500,7 @@ let def_loc = function
   | DEF_kind (KD_aux (_, (l, _)))
   | DEF_type (TD_aux (_, (l, _)))
   | DEF_fundef (FD_aux (_, (l, _)))
+  | DEF_mapdef (MD_aux (_, (l, _)))
   | DEF_val (LB_aux (_, (l, _)))
   | DEF_spec (VS_aux (_, (l, _)))
   | DEF_default (DT_aux (_, l))
@@ -752,7 +754,8 @@ and string_of_mpat (MP_aux (pat, l)) =
   | MP_vector_concat pats -> string_of_list " : " string_of_mpat pats
   | MP_vector pats -> "[" ^ string_of_list ", " string_of_mpat pats ^ "]"
   | MP_string_append (pat1, pat2) -> string_of_mpat pat1 ^ " ^^ " ^ string_of_mpat pat2
-  | _ -> "PAT"
+  | MP_typ (mpat, typ) -> "(" ^ string_of_mpat mpat ^ " : " ^ string_of_typ typ ^ ")"
+  | _ -> "MPAT"
 
 and string_of_lexp (LEXP_aux (lexp, _)) =
   match lexp with
@@ -976,6 +979,7 @@ let rec tyvars_of_typ (Typ_aux (t,_)) =
   | Typ_id _ -> KidSet.empty
   | Typ_var kid -> KidSet.singleton kid
   | Typ_fn (t1,t2,_) -> KidSet.union (tyvars_of_typ t1) (tyvars_of_typ t2)
+  | Typ_bidir (t1, t2) -> KidSet.union (tyvars_of_typ t1) (tyvars_of_typ t2)
   | Typ_tup ts ->
      List.fold_left (fun s t -> KidSet.union s (tyvars_of_typ t))
        KidSet.empty ts
@@ -1020,7 +1024,7 @@ let rec undefined_of_typ mwords l annot (Typ_aux (typ_aux, _) as typ) =
         initial_check.ml. i.e. the rewriter should only encounter this
         case when re-writing those functions. *)
      wrap (E_id (prepend_id "typ_" (id_of_kid kid))) typ
-  | Typ_fn _ | Typ_exist _ -> assert false (* Typ_exist should be re-written *)
+  | Typ_bidir _ | Typ_fn _ | Typ_exist _ -> assert false (* Typ_exist should be re-written *)
 and undefined_of_typ_args mwords l annot (Typ_arg_aux (typ_arg_aux, _) as typ_arg) =
   match typ_arg_aux with
   | Typ_arg_nexp n -> [E_aux (E_sizeof n, (l, annot (atom_typ n)))]
diff --git a/src/initial_check.ml b/src/initial_check.ml
index b766daa1..3f7a7052 100644
--- a/src/initial_check.ml
+++ b/src/initial_check.ml
@@ -782,6 +782,7 @@ let rec to_ast_mpat k_env def_ord (Parse_ast.MP_aux(mpat,l)) =
     | Parse_ast.MP_list(mpats) -> MP_list(List.map (to_ast_mpat k_env def_ord) mpats)
     | Parse_ast.MP_cons(pat1, pat2) -> MP_cons (to_ast_mpat k_env def_ord pat1, to_ast_mpat k_env def_ord pat2)
     | Parse_ast.MP_string_append (pat1, pat2) -> MP_string_append (to_ast_mpat k_env def_ord pat1, to_ast_mpat k_env def_ord pat2)
+    | Parse_ast.MP_typ (mpat, typ) -> MP_typ (to_ast_mpat k_env def_ord mpat, to_ast_typ k_env def_ord typ)
     ), (l,()))
 
 
diff --git a/src/parse_ast.ml b/src/parse_ast.ml
index 3969663a..e59dd356 100644
--- a/src/parse_ast.ml
+++ b/src/parse_ast.ml
@@ -431,6 +431,7 @@ type mpat_aux =  (* Mapping pattern. Mostly the same as normal patterns but only
  | MP_list of ( mpat) list
  | MP_cons of ( mpat) * ( mpat)
  | MP_string_append of ( mpat) * ( mpat)
+ | MP_typ of mpat * atyp
 
 and mpat =
  | MP_aux of ( mpat_aux) * l
diff --git a/src/parser.mly b/src/parser.mly
index adf99ede..d41964b7 100644
--- a/src/parser.mly
+++ b/src/parser.mly
@@ -1241,6 +1241,9 @@ atomic_mpat:
     { mk_mpat (MP_list []) $startpos $endpos }
   | LsquareBar mpat_list RsquareBar
     { mk_mpat (MP_list $2) $startpos $endpos }
+  | atomic_mpat Colon typ
+    { mk_mpat (MP_typ ($1, $3)) $startpos $endpos }
+
 
 
 mpexp:
diff --git a/src/rewrites.ml b/src/rewrites.ml
index fdfd949a..ac36dffa 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -3504,6 +3504,7 @@ let rec exp_of_mpat (MP_aux (mpat, annot)) =
   | MP_list mpats                   -> E_aux (E_list (List.map exp_of_mpat mpats), annot)
   | MP_cons (mpat1, mpat2)          -> E_aux (E_cons (exp_of_mpat mpat1, exp_of_mpat mpat2), annot)
   | MP_string_append (mpat1, mpat2) -> E_aux (E_app (mk_id "string_append", [exp_of_mpat mpat1; exp_of_mpat mpat2]), annot)
+  | MP_typ (mpat, typ)              -> E_aux (E_cast (typ, exp_of_mpat mpat), annot)
 
 and fexps_of_mfpats mfpats flag annot =
   let fexp_of_mfpat (MFP_aux (MFP_mpat (id, mpat), annot)) =
@@ -3523,6 +3524,7 @@ let rec pat_of_mpat (MP_aux (mpat, annot)) =
   | MP_list mpats                   -> P_aux (P_list (List.map pat_of_mpat mpats), annot)
   | MP_cons (mpat1, mpat2)          -> P_aux ((P_cons (pat_of_mpat mpat1, pat_of_mpat mpat2), annot))
   | MP_string_append (mpat1, mpat2) -> P_aux ((P_string_append (pat_of_mpat mpat1, pat_of_mpat mpat2), annot))
+  | MP_typ (mpat, typ)              -> P_aux (P_typ (typ, pat_of_mpat mpat), annot)
 
 and fpats_of_mfpats mfpats =
   let fpat_of_mfpat (MFP_aux (MFP_mpat (id, mpat), annot)) =
diff --git a/src/type_check.ml b/src/type_check.ml
index a0888aa7..c96c2c3f 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -3632,6 +3632,10 @@ and infer_mpat env (MP_aux (mpat_aux, (l, ())) as mpat) =
      end
   | MP_lit lit ->
      annot_mpat (MP_lit lit) (infer_lit env lit), env, []
+  | MP_typ (mpat, typ_annot) ->
+     Env.wf_typ env typ_annot;
+     let (typed_mpat, env, guards) = bind_mpat env mpat typ_annot in
+     annot_mpat (MP_typ (typed_mpat, typ_annot)) typ_annot, env, guards
   | MP_vector (mpat :: mpats) ->
      let fold_mpats (mpats, env, guards) mpat =
        let typed_mpat, env, guards' = bind_mpat env mpat bit_typ in
@@ -3965,6 +3969,9 @@ and propagate_mpat_effect_aux = function
   | MP_vector mpats ->
      let p_mpats = List.map propagate_mpat_effect mpats in
      MP_vector p_mpats, collect_effects_mpat p_mpats
+  | MP_typ (mpat, typ) ->
+     let p_mpat = propagate_mpat_effect mpat in
+     MP_typ (p_mpat, typ), effect_of_mpat mpat
   | _ -> typ_error Parse_ast.Unknown "Unimplemented: Cannot propagate effect in mpat"