Handle type variables generated while inferring applications in monomorphisation

Also handle any type variables from assignments and degrade gracefully
during constant propagation when unification is not possible.

2 files changed, 48 insertions, 26 deletions

diff --git a/src/constant_propagation.ml b/src/constant_propagation.ml
index 2fb19c05..89a4ba82 100644
--- a/src/constant_propagation.ml
+++ b/src/constant_propagation.ml
@@ -684,7 +684,9 @@ let const_props defs ref_vars =
          let exp_typ = typ_of exp in
          let pat_typ = typ_of_pat pat in
          let goals = KidSet.diff (tyvars_of_typ pat_typ) (tyvars_of_typ exp_typ) in
-         let unifiers = Type_check.unify l env goals pat_typ exp_typ in
+         let unifiers =
+           try Type_check.unify l env goals pat_typ exp_typ
+           with _ -> KBindings.empty in
          let is_nexp (k,a) = match a with
            | A_aux (A_nexp n,_) -> Some (k,n)
            | _ -> None
diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index 008d5e24..c011f4d0 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -1817,6 +1817,19 @@ let update_env env deps pat typ_env_pre typ_env_post =
   let var_deps = List.fold_left (fun ds v -> Bindings.add v deps ds) env.var_deps bound in
   update_env_new_kids { env with var_deps = var_deps } deps typ_env_pre typ_env_post
 
+(* A function argument may end up with fresh type variables due to coercing
+   unification (which will eventually be existentially bound in the type of
+   the function).  Here we record the dependencies for these variables. *)
+
+let add_arg_only_kids env typ_env typ deps =
+  let all_vars = tyvars_of_typ typ in
+  let check_kid kid kid_deps =
+    if KBindings.mem kid kid_deps then kid_deps
+    else KBindings.add kid deps kid_deps
+  in
+  let kid_deps = KidSet.fold check_kid all_vars env.kid_deps in
+  { env with kid_deps }
+
 let assigned_vars_exps es =
   List.fold_left (fun vs exp -> IdSet.union vs (Spec_analysis.assigned_vars exp))
     IdSet.empty es
@@ -2001,39 +2014,40 @@ let rec analyse_exp fn_id env assigns (E_aux (e,(l,annot)) as exp) =
      follow the type checker in processing them in-order to detect the automatic
      unpacking of existentials.  When we spot a new type variable (using
      update_env_new_kids) we set them to depend on the previous argument. *)
-  let non_det_args es =
+  let non_det_args es typs =
     let assigns = remove_assigns es " assigned in non-deterministic expressions" in
-    let rec aux prev_typ_env prev_deps env = function
-      | [] -> [], empty
-      | h::t ->
+    let rec aux env = function
+      | [], _ -> [], empty, env
+      | (E_aux (_,ann) as h)::t, typ::typs ->
          let typ_env = env_of h in
-         let env = update_env_new_kids env prev_deps prev_typ_env typ_env in
-         let new_deps, _, new_r = analyse_exp fn_id env assigns h in
-         let t_deps, t_r = aux typ_env new_deps env t in
-         new_deps::t_deps, merge new_r t_r
-    in
-    let deps, r = match es with
-      | [] -> [], empty
-      | h::t ->
          let new_deps, _, new_r = analyse_exp fn_id env assigns h in
-         let t_deps, t_r = aux (env_of h) new_deps env t in
-         new_deps::t_deps, merge new_r t_r
+         let env = add_arg_only_kids env typ_env typ new_deps in
+         let t_deps, t_r, t_env = aux env (t,typs) in
+         new_deps::t_deps, merge new_r t_r, t_env
     in
-    (deps, assigns, r)
+    let deps, r, env = aux env (es,typs) in
+    (deps, assigns, r, env)
   in
   let merge_deps deps = List.fold_left dmerge dempty deps in
   let deps, assigns, r =
     match e with
     | E_block es ->
-       let rec aux assigns = function
+       let rec aux env assigns = function
          | [] -> (dempty, assigns, empty)
          | [e] -> analyse_exp fn_id env assigns e
+         (* There's also a lone assignment case below where no env update is needed *)
+         | E_aux (E_assign (lexp,e1),ann)::e2::es ->
+            let d1,assigns,r1 = analyse_exp fn_id env assigns e1 in
+            let assigns,r2 = analyse_lexp fn_id env assigns d1 lexp in
+            let env = update_env_new_kids env d1 (env_of_annot ann) (env_of e2) in
+            let d3, assigns, r3 = aux env assigns (e2::es) in
+            (d3, assigns, merge (merge r1 r2) r3)
          | e::es ->
             let _, assigns, r' = analyse_exp fn_id env assigns e in
-            let d, assigns, r = aux assigns es in
+            let d, assigns, r = aux env assigns es in
             d, assigns, merge r r'
        in
-       aux assigns es
+       aux env assigns es
     | E_nondet es ->
        let _, assigns, r = non_det es in
        (dempty, assigns, r)
@@ -2060,21 +2074,26 @@ let rec analyse_exp fn_id env assigns (E_aux (e,(l,annot)) as exp) =
     | E_lit _ -> (dempty,assigns,empty)
     | E_cast (_,e) -> analyse_exp fn_id env assigns e
     | E_app (id,args) ->
-       let deps, assigns, r = non_det_args args in
        let typ_env = env_of_annot (l,annot) in
-       let (_,fn_typ) = Env.get_val_spec id typ_env in
-       let fn_effect = match fn_typ with
-         | Typ_aux (Typ_fn (_,_,eff),_) -> eff
-         | _ -> Effect_aux (Effect_set [],Unknown)
+       let (_,fn_typ) = Env.get_val_spec_orig id typ_env in
+       let kid_inst = instantiation_of exp in
+       let kid_inst = KBindings.fold (fun kid -> KBindings.add (orig_kid kid)) kid_inst KBindings.empty in
+       let fn_typ = subst_unifiers kid_inst fn_typ in
+       let arg_typs, fn_effect = match fn_typ with
+         | Typ_aux (Typ_fn (args,_,eff),_) -> args,eff
+         | _ -> [], Effect_aux (Effect_set [],Unknown)
        in
+       (* We have to use the types from the val_spec here so that we can track
+          any type variables that are generated by the coercing unification that
+          the type checker applies after inferring the type of an argument, and
+          that only appear in the unifiers. *)
+       let deps, assigns, r, env = non_det_args args arg_typs in
        let eff_dep = match fn_effect with
          | Effect_aux (Effect_set ([] | [BE_aux (BE_undef,_)]),_) -> dempty
          | _ -> Unknown (l, "Effects from function application")
        in
-       let kid_inst = instantiation_of exp in
        let kid_inst = KBindings.map (simplify_size_typ_arg env typ_env) kid_inst in
        (* Change kids in instantiation to the canonical ones from the type signature *)
-       let kid_inst = KBindings.fold (fun kid -> KBindings.add (orig_kid kid)) kid_inst KBindings.empty in
        let kid_deps = KBindings.map (deps_of_typ_arg l fn_id env deps) kid_inst in
        let rdep,r' =
          if Id.compare fn_id id == 0 then
@@ -2163,6 +2182,7 @@ let rec analyse_exp fn_id env assigns (E_aux (e,(l,annot)) as exp) =
        let env = update_env env d1 pat (env_of_annot (l,annot)) (env_of e2) in
        let d2,assigns,r2 = analyse_exp fn_id env assigns e2 in
        (d2,assigns,merge r1 r2)
+    (* There's a more general assignment case above to update env inside a block. *)
     | E_assign (lexp,e1) ->
        let d1,assigns,r1 = analyse_exp fn_id env assigns e1 in
        let assigns,r2 = analyse_lexp fn_id env assigns d1 lexp in