Allow type variables to be introduced by global let bindings.

This was technically allowed previously but the rules for type
variable names in function types were too strict so it didn't work.

Also fixed a bug where Nexp_app constructors were never considered identical

and fixed a bug where top-level let bindings got annotated with the
wrong environment

2 files changed, 31 insertions, 10 deletions

diff --git a/src/rewrites.ml b/src/rewrites.ml
index 25fe9e03..0091ace0 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -3028,15 +3028,25 @@ let rewrite_check_annot =
     try
       prerr_endline ("CHECKING: " ^ string_of_exp exp ^ " : " ^ string_of_typ (typ_of exp));
       let _ = check_exp (env_of exp) (strip_exp exp) (typ_of exp) in
-      (if not (alpha_equivalent (env_of exp) (typ_of exp) (Env.expand_synonyms (env_of exp) (typ_of exp)))
-       then raise (Reporting_basic.err_typ Parse_ast.Unknown "Found synonym in annotation")
+      let typ1 = typ_of exp in
+      let typ2 = Env.expand_synonyms (env_of exp) (typ_of exp) in
+      (if not (alpha_equivalent (env_of exp) typ1 typ2)
+       then raise (Reporting_basic.err_typ Parse_ast.Unknown
+                    ("Found synonym in annotation " ^ string_of_typ typ1 ^ " vs " ^ string_of_typ typ2))
        else ());
       exp
     with
       Type_error (l, err) -> raise (Reporting_basic.err_typ l (string_of_type_error err))
   in
+  let check_pat pat =
+    prerr_endline ("CHECKING PAT: " ^ string_of_pat pat ^ " : " ^ string_of_typ (pat_typ_of pat));
+    let _, _ = bind_pat_no_guard (pat_env_of pat) (strip_pat pat) (pat_typ_of pat) in
+    pat
+  in
+
   let rewrite_exp = { id_exp_alg with e_aux = (fun (exp, annot) -> check_annot (E_aux (exp, annot))) } in
-  rewrite_defs_base { rewriters_base with rewrite_exp = (fun _ -> fold_exp rewrite_exp) }
+  rewrite_defs_base { rewriters_base with rewrite_exp = (fun _ -> fold_exp rewrite_exp);
+                                          rewrite_pat = (fun _ -> check_pat) }
 
 let rewrite_defs_check = [
   ("check_annotations", rewrite_check_annot);
diff --git a/src/type_check.ml b/src/type_check.ml
index 2ce7aebf..b4fff878 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -634,7 +634,10 @@ end = struct
 
   let freshen_kid env kid (typq, typ) =
     let fresh = fresh_kid ~kid:kid env in
-    (typquant_subst_kid kid fresh typq, typ_subst_kid kid fresh typ)
+    if KidSet.mem kid (KidSet.of_list (List.map kopt_kid (quant_kopts typq))) then
+      (typquant_subst_kid kid fresh typq, typ_subst_kid kid fresh typ)
+    else
+      (typq, typ)
 
   let freshen_bind env bind =
     List.fold_left (fun bind (kid, _) -> freshen_kid env kid bind) bind (KBindings.bindings env.typ_vars)
@@ -1283,6 +1286,8 @@ let rec nexp_identical (Nexp_aux (nexp1, _)) (Nexp_aux (nexp2, _)) =
   | Nexp_minus (n1a, n1b), Nexp_minus (n2a, n2b) -> nexp_identical n1a n2a && nexp_identical n1b n2b
   | Nexp_exp n1, Nexp_exp n2 -> nexp_identical n1 n2
   | Nexp_neg n1, Nexp_neg n2 -> nexp_identical n1 n2
+  | Nexp_app (f1, args1), Nexp_app (f2, args2) when List.length args1 = List.length args2 ->
+     Id.compare f1 f2 = 0 && List.for_all2 nexp_identical args1 args2
   | _, _ -> false
 
 let ord_identical (Ord_aux (ord1, _)) (Ord_aux (ord2, _)) =
@@ -2393,6 +2398,12 @@ and bind_pat env (P_aux (pat_aux, (l, ())) as pat) (Typ_aux (typ_aux, _) as typ)
           let env = Env.add_constraint (nc_and (nc_lteq lo (nvar kid)) (nc_lteq (nvar kid) hi)) env in
           let typed_pat, env, guards = bind_pat env pat (atom_typ (nvar kid)) in
           annot_pat (P_var (typed_pat, kid)) typ, env, guards
+       | None, Typ_aux (Typ_app (id, [Typ_arg_aux (Typ_arg_nexp n, _)]), _)
+            when Id.compare id (mk_id "atom") == 0 ->
+          let env = Env.add_typ_var kid BK_nat env in
+          let env = Env.add_constraint (nc_eq n (nvar kid)) env in
+          let typed_pat, env, guards = bind_pat env pat (atom_typ (nvar kid)) in
+          annot_pat (P_var (typed_pat, kid)) typ, env, guards
        | None, _ -> typ_error l ("Cannot bind type variable against non existential or numeric type")
      end
   | P_wild -> annot_pat P_wild typ, env, []
@@ -3395,16 +3406,16 @@ and propagate_lexp_effect_aux = function
 (* 6. Checking toplevel definitions                                       *)
 (**************************************************************************)
 
-let check_letdef env (LB_aux (letbind, (l, _))) =
+let check_letdef orig_env (LB_aux (letbind, (l, _))) =
   begin
     match letbind with
     | LB_val (P_aux (P_typ (typ_annot, pat), _), bind) ->
-       let checked_bind = crule check_exp env (strip_exp bind) typ_annot in
-       let tpat, env = bind_pat_no_guard env (strip_pat pat) typ_annot in
-       [DEF_val (LB_aux (LB_val (P_aux (P_typ (typ_annot, tpat), (l, Some (env, typ_annot, no_effect))), checked_bind), (l, None)))], env
+       let checked_bind = crule check_exp orig_env (strip_exp bind) typ_annot in
+       let tpat, env = bind_pat_no_guard orig_env (strip_pat pat) typ_annot in
+       [DEF_val (LB_aux (LB_val (P_aux (P_typ (typ_annot, tpat), (l, Some (orig_env, typ_annot, no_effect))), checked_bind), (l, None)))], env
     | LB_val (pat, bind) ->
-       let inferred_bind = irule infer_exp env (strip_exp bind) in
-       let tpat, env = bind_pat_no_guard env (strip_pat pat) (typ_of inferred_bind) in
+       let inferred_bind = irule infer_exp orig_env (strip_exp bind) in
+       let tpat, env = bind_pat_no_guard orig_env (strip_pat pat) (typ_of inferred_bind) in
        [DEF_val (LB_aux (LB_val (tpat, inferred_bind), (l, None)))], env
   end