Add support for inline val-spec declaration for mappings

This means that a mapping which formerly had to be pre-declared like

val name : a <-> b
...
mapping name {
  x <-> y,
  ...
}

can now be shortened to

mapping name : a <-> b {
  x <-> y,
  ...
}

6 files changed, 79 insertions, 20 deletions

diff --git a/src/initial_check.ml b/src/initial_check.ml
index e1dd906b..9545ce44 100644
--- a/src/initial_check.ml
+++ b/src/initial_check.ml
@@ -754,6 +754,15 @@ let to_ast_tannot_opt (k_env:kind Envmap.t) (def_ord:order) (Parse_ast.Typ_annot
     let typq,k_env,k_local = to_ast_typquant k_env tq in
     Typ_annot_opt_aux(Typ_annot_opt_some(typq,to_ast_typ k_env def_ord typ),l),k_env,k_local
 
+let to_ast_typschm_opt (k_env:kind Envmap.t) (def_ord:order) (Parse_ast.TypSchm_opt_aux(aux,l)) : tannot_opt * kind Envmap.t * kind Envmap.t =
+  match aux with
+  | Parse_ast.TypSchm_opt_none ->
+     Typ_annot_opt_aux (Typ_annot_opt_none, l), k_env, Envmap.empty
+  | Parse_ast.TypSchm_opt_some (Parse_ast.TypSchm_aux (Parse_ast.TypSchm_ts (tq, typ), l)) ->
+     let typq, k_env, k_local = to_ast_typquant k_env tq in
+     Typ_annot_opt_aux (Typ_annot_opt_some (typq, to_ast_typ k_env def_ord typ), l), k_env, k_local
+
+
 let to_ast_effects_opt (k_env : kind Envmap.t) (Parse_ast.Effect_opt_aux(e,l)) : effect_opt =
   match e with
   | Parse_ast.Effect_opt_pure -> Effect_opt_aux(Effect_opt_pure,l)
@@ -807,8 +816,9 @@ let to_ast_mapcl (names,k_env,def_ord) (Parse_ast.MCL_aux(mapcl, l)) =
 
 let to_ast_mapdef (names,k_env,def_ord) (Parse_ast.MD_aux(md,l):Parse_ast.mapdef) : (unit mapdef) envs_out =
   match md with
-  | Parse_ast.MD_mapping(id, mapcls) ->
-     MD_aux(MD_mapping(to_ast_id id, List.map (to_ast_mapcl (names,k_env,def_ord)) mapcls), (l,())), (names,k_env,def_ord)
+  | Parse_ast.MD_mapping(id, typschm_opt, mapcls) ->
+     let tannot_opt, k_env, _ = to_ast_typschm_opt k_env def_ord typschm_opt in
+     MD_aux(MD_mapping(to_ast_id id, tannot_opt, List.map (to_ast_mapcl (names,k_env,def_ord)) mapcls), (l,())), (names,k_env,def_ord)
 
 type def_progress =
     No_def
@@ -904,11 +914,11 @@ let to_ast_def (names, k_env, def_ord) partial_defs def : def_progress envs_out
           | None -> let partial_def = ref ((DEF_fundef(FD_aux(FD_function(rec_opt,unit,effects_opt,[]),(l,())))),false) in
                     (No_def,envs),((id,(partial_def,k_local))::partial_defs)
           | Some(d,k) -> typ_error l "Scattered function definition header name already in use by scattered definition" (Some id) None None)
-      | Parse_ast.SD_scattered_mapping id ->
+      | Parse_ast.SD_scattered_mapping (id, tannot_opt) ->
          let id = to_ast_id id in
-         let _,_,k_local = to_ast_tannot_opt k_env def_ord (Parse_ast.Typ_annot_opt_aux (Parse_ast.Typ_annot_opt_none, Parse_ast.Unknown)) in
+         let unit, k_env ,k_local = to_ast_tannot_opt k_env def_ord tannot_opt in
          (match (def_in_progress id partial_defs) with
-          | None -> let partial_def = ref ((DEF_mapdef(MD_aux(MD_mapping(id, []), (l, ())))), false) in
+          | None -> let partial_def = ref ((DEF_mapdef(MD_aux(MD_mapping(id, unit, []), (l, ())))), false) in
                     (No_def,envs),((id,(partial_def,k_local))::partial_defs)
           | Some(d,k) -> typ_error l "Scattered mapping definition header name already in use by scattered definition" (Some id) None None)
 
@@ -918,9 +928,9 @@ let to_ast_def (names, k_env, def_ord) partial_defs def : def_progress envs_out
           | None -> typ_error l "Scattered mapping definition clause does not match any existing mapping definition headers" (Some id) None None
           | Some (d, k) ->
              (match !d with
-              | DEF_mapdef(MD_aux(MD_mapping(_,mcls),ml)),false ->
+              | DEF_mapdef(MD_aux(MD_mapping(_,tannot_opt, mcls),ml)),false ->
                  let (MCL_aux (mapcl_aux, _)) = to_ast_mapcl (names,k_env,def_ord) mapcl in
-                 d := DEF_mapdef(MD_aux(MD_mapping(id, mcls @ [MCL_aux (mapcl_aux, (l, ()))]), ml)), false;
+                 d := DEF_mapdef(MD_aux(MD_mapping(id, tannot_opt, mcls @ [MCL_aux (mapcl_aux, (l, ()))]), ml)), false;
                  (No_def,envs),partial_defs
               | _, true -> typ_error l "Scattered mapping definition clause extends ended definition" (Some id) None None
               | _ -> typ_error l "Scattered mapping definition doesn't match existing definition header" (Some id) None None))
diff --git a/src/parse_ast.ml b/src/parse_ast.ml
index 607285c7..c31d548c 100644
--- a/src/parse_ast.ml
+++ b/src/parse_ast.ml
@@ -344,6 +344,14 @@ tannot_opt_aux =  (* Optional type annotation for functions *)
    Typ_annot_opt_none
  | Typ_annot_opt_some of typquant * atyp
 
+type
+typschm_opt_aux =
+  TypSchm_opt_none
+| TypSchm_opt_some of typschm
+
+type
+typschm_opt =
+  TypSchm_opt_aux of typschm_opt_aux * l
 
 type 
 effect_opt_aux =  (* Optional effect annotation for functions *)
@@ -456,7 +464,7 @@ type mapcl =
  | MCL_aux of ( mapcl_aux) * l
 
 type mapdef_aux =  (* mapping definition (bidirectional pattern-match function) *)
- | MD_mapping of id * ( mapcl) list
+ | MD_mapping of id * typschm_opt * ( mapcl) list
 
 type mapdef =
  | MD_aux of ( mapdef_aux) * l
@@ -498,7 +506,7 @@ scattered_def_aux =  (* Function and type union definitions that can be spread a
  | SD_scattered_funcl of funcl (* scattered function definition clause *)
  | SD_scattered_variant of id * name_scm_opt * typquant (* scattered union definition header *)
  | SD_scattered_unioncl of id * type_union (* scattered union definition member *)
- | SD_scattered_mapping of id
+ | SD_scattered_mapping of id * tannot_opt
  | SD_scattered_mapcl of id * mapcl
  | SD_scattered_end of id (* scattered definition end *)
 
diff --git a/src/parser.mly b/src/parser.mly
index a46defd6..5c513e5b 100644
--- a/src/parser.mly
+++ b/src/parser.mly
@@ -93,6 +93,16 @@ let mk_exp e n m = E_aux (e, loc n m)
 let mk_lit l n m = L_aux (l, loc n m)
 let mk_lit_exp l n m = mk_exp (E_lit (mk_lit l n m)) n m
 let mk_typschm tq t n m = TypSchm_aux (TypSchm_ts (tq, t), loc n m)
+
+let mk_typschm_opt ts n m = TypSchm_opt_aux (
+                                  TypSchm_opt_some (
+                                      ts
+                                    ),
+                                  loc n m
+                                )
+
+let mk_typschm_opt_none = TypSchm_opt_aux (TypSchm_opt_none, Unknown)
+
 let mk_nc nc n m = NC_aux (nc, loc n m)
 let mk_sd s n m = SD_aux (s, loc n m)
 let mk_sd_doc s str n m = SD_aux (s, Documented (str, loc n m))
@@ -108,7 +118,7 @@ let mk_default d n m = DT_aux (d, loc n m)
 let mk_mpexp mpexp n m = MPat_aux (mpexp, loc n m)
 let mk_mpat mpat n m = MP_aux (mpat, loc n m)
 let mk_mapcl mpexp1 mpexp2 n m = MCL_aux (MCL_mapcl (mpexp1, mpexp2), loc n m)
-let mk_map id mapcls n m = MD_aux (MD_mapping (id, mapcls), loc n m)
+let mk_map id tannot mapcls n m = MD_aux (MD_mapping (id, tannot, mapcls), loc n m)
 
 let doc_vs doc (VS_aux (v, l)) = VS_aux (v, Documented (doc, l))
 
@@ -1280,9 +1290,9 @@ mapcl_list:
 
 map_def:
   | Mapping id Eq Lcurly mapcl_list Rcurly
-    { mk_map $2 $5 $startpos $endpos }
-  (* | Mapping id Colon typschm Eq Lcurly mapcl_list Rcurly
-   *   { mk_map $2 $4 $7 $startpos $endpos } *)
+    { mk_map $2 mk_typschm_opt_none $5 $startpos $endpos }
+  | Mapping id Colon typschm Eq Lcurly mapcl_list Rcurly
+    { mk_map $2 (mk_typschm_opt $4 $startpos($4) $endpos($4)) $7 $startpos $endpos }
 
 let_def:
   | Let_ letbind
@@ -1334,7 +1344,9 @@ scattered_def:
   | Function_ id
     { mk_sd (SD_scattered_function(mk_recn, mk_tannotn, mk_eannotn, $2)) $startpos $endpos }
   | Mapping id
-    { mk_sd (SD_scattered_mapping $2) $startpos $endpos }  
+    { mk_sd (SD_scattered_mapping ($2, mk_tannotn)) $startpos $endpos }
+  | Mapping id Colon funcl_typ
+    { mk_sd (SD_scattered_mapping ($2, $4)) $startpos $endpos }
 
 scattered_clause:
   | Doc Function_ Clause funcl
diff --git a/src/pretty_print_sail.ml b/src/pretty_print_sail.ml
index 6ea669f9..4b9fa6a9 100644
--- a/src/pretty_print_sail.ml
+++ b/src/pretty_print_sail.ml
@@ -476,7 +476,7 @@ let doc_mapcl (MCL_aux (MCL_mapcl (mpexp1, mpexp2), _)) =
   let right = doc_mpexp mpexp2 in
   left ^^ space ^^ string "<->" ^^ space ^^ right
 
-let doc_mapdef (MD_aux (MD_mapping (id, mapcls), _)) =
+let doc_mapdef (MD_aux (MD_mapping (id, typa, mapcls), _)) =
   match mapcls with
   | [] -> failwith "Empty mapping"
   | _ ->
diff --git a/src/rewrites.ml b/src/rewrites.ml
index e41318dd..58ff885f 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -3758,7 +3758,7 @@ let rewrite_defs_realise_mappings (Defs defs) =
       | MPat_aux (MPat_when (mpat2, _), _)->
        realise_mpexps true (append_placeholder mpexp) (mk_mpexp (MPat_pat (mk_mpat (MP_app ((mk_id "Some"), [ mk_mpat (MP_tup [mpat2; strlen]) ])))))
   in
-  let realise_mapdef (MD_aux (MD_mapping (id, mapcls), ((l, (tannot:tannot)) as annot))) =
+  let realise_mapdef (MD_aux (MD_mapping (id, _, mapcls), ((l, (tannot:tannot)) as annot))) =
     let forwards_id = mk_id (string_of_id id ^ "_forwards") in
     let forwards_matches_id = mk_id (string_of_id id ^ "_forwards_matches") in
     let backwards_id = mk_id (string_of_id id ^ "_backwards") in
diff --git a/src/type_check.ml b/src/type_check.ml
index 09b8a9c7..fbec5111 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -674,6 +674,7 @@ end = struct
     match typ_aux with
     | Typ_id _ | Typ_var _ -> typ
     | Typ_fn (arg_typ, ret_typ, effect) -> Typ_aux (Typ_fn (map_nexps f arg_typ, map_nexps f ret_typ, effect), l)
+    | Typ_bidir (typ1, typ2) -> Typ_aux (Typ_bidir (map_nexps f typ1, map_nexps f typ2), l)
     | Typ_tup typs -> Typ_aux (Typ_tup (List.map (map_nexps f) typs), l)
     | Typ_exist (kids, nc, typ) -> Typ_aux (Typ_exist (kids, nc, map_nexps f typ), l)
     | Typ_app (id, args) -> Typ_aux (Typ_app (id, List.map (map_nexps_arg f) args), l)
@@ -1507,6 +1508,13 @@ let typ_identical env typ1 typ2 =
     match typ1, typ2 with
     | Typ_id v1, Typ_id v2 -> Id.compare v1 v2 = 0
     | Typ_var kid1, Typ_var kid2 -> Kid.compare kid1 kid2 = 0
+    | Typ_fn (arg_typ1, ret_typ1, eff1), Typ_fn (arg_typ2, ret_typ2, eff2) ->
+       typ_identical' arg_typ1 arg_typ2
+       && typ_identical' ret_typ1 ret_typ2
+       && strip_effect eff1 = strip_effect eff2
+    | Typ_bidir (typ1, typ2), Typ_bidir (typ3, typ4) ->
+       typ_identical' typ1 typ3
+       && typ_identical' typ2 typ4
     | Typ_tup typs1, Typ_tup typs2 ->
        begin
          try List.for_all2 typ_identical' typs1 typs2 with
@@ -4231,21 +4239,42 @@ let check_fundef env (FD_aux (FD_function (recopt, tannotopt, effectopt, funcls)
   else typ_error l ("Effects do not match: " ^ string_of_effect declared_eff ^ " declared and " ^ string_of_effect eff ^ " found")
 
 
-let check_mapdef env (MD_aux (MD_mapping (id, mapcls), (l, _)) as md_aux) =
+let check_mapdef env (MD_aux (MD_mapping (id, tannot_opt, mapcls), (l, _)) as md_aux) =
   typ_print (lazy ("\nChecking mapping " ^ string_of_id id));
-  let quant, typ = Env.get_val_spec id env in
+  let have_val_spec, (quant, typ), env =
+    try true, Env.get_val_spec id env, env with
+    | Type_error (l, _) as err ->
+       match tannot_opt with
+       | Typ_annot_opt_aux (Typ_annot_opt_some (quant, typ), _) ->
+          false, (quant, typ), env
+       | Typ_annot_opt_aux (Typ_annot_opt_none, _) ->
+          raise err
+  in
   let vtyp1, vtyp2, vl = match typ with
     | Typ_aux (Typ_bidir (vtyp1, vtyp2), vl) -> vtyp1, vtyp2, vl
     | _ -> typ_error l "Mapping val spec was not a mapping type"
   in
+  begin match tannot_opt with
+  | Typ_annot_opt_aux (Typ_annot_opt_none, _) -> ()
+  | Typ_annot_opt_aux (Typ_annot_opt_some (annot_typq, annot_typ), l) ->
+     if typ_identical env typ annot_typ then ()
+     else typ_error l (string_of_bind (quant, typ) ^ " and " ^ string_of_bind (annot_typq, annot_typ) ^ " do not match between mapping and val spec")
+  end;
   typ_debug (lazy ("Checking mapdef " ^ string_of_id id ^ " has type " ^ string_of_bind (quant, typ)));
+  let vs_def, env =
+    if not have_val_spec then
+      [mk_val_spec env quant typ id], Env.add_val_spec id (quant, typ) env
+    else
+      [], env
+  in
   let mapcl_env = add_typquant quant env in
   let mapcls = List.map (fun mapcl -> check_mapcl mapcl_env mapcl typ) mapcls in
   let eff = List.fold_left union_effects no_effect (List.map mapcl_effect mapcls) in
+  let env = Env.define_val_spec id env in
   if equal_effects eff no_effect then
-    [DEF_mapdef (MD_aux (MD_mapping (id, mapcls), (l, None)))], env
+    vs_def @ [DEF_mapdef (MD_aux (MD_mapping (id, tannot_opt, mapcls), (l, None)))], env
   else
-    typ_error l ("Mapping not pure:" ^ string_of_effect eff ^ " found")
+    typ_error l ("Mapping not pure: " ^ string_of_effect eff ^ " found")
 
 
 (* Checking a val spec simply adds the type as a binding in the