classifying and applying substitutive invariant subcomp

FStarLang · aseemr · Nov 22, 2022 · Oct 7, 2022 · Oct 10, 2022 · Oct 10, 2022
commit 8ae7d09e7b54b2d4c3ddfc5a76a9e1ce71d306a8
diff --git a/src/typechecker/FStar.TypeChecker.Rel.fst b/src/typechecker/FStar.TypeChecker.Rel.fst
@@ -1801,6 +1801,7 @@ let simplify_guard env g = match g.guard_f with
 // returns the (subcomp guard, new sub problems, worklist)
 //
 let apply_substitutive_indexed_subcomp (env:Env.env)
+  (k:S.indexed_effect_combinator_kind)
   (bs:binders)
   (subcomp_c:comp)
   (ct1:comp_typ) (ct2:comp_typ)
@@ -1854,11 +1855,24 @@ let apply_substitutive_indexed_subcomp (env:Env.env)
     subst@f_substs in
 
   // add substitutions for the g computation
-  let bs, subst =
-    let g_bs, bs = List.splitAt (List.length args2) bs in
-    let g_substs = List.map2 (fun g_b (arg, _) -> NT (g_b.binder_bv, arg)) g_bs args2 in
-    bs,
-    subst@g_substs in
+  let bs, subst, f_g_args_eq_sub_probs, wl =
+    if Substitutive_combinator? k
+    then begin
+      let g_bs, bs = List.splitAt (List.length args2) bs in
+      let g_substs = List.map2 (fun g_b (arg, _) -> NT (g_b.binder_bv, arg)) g_bs args2 in
+      bs,
+      subst@g_substs,
+      [],
+      wl
+    end
+    else if Substitutive_invariant_combinator? k
+    then begin
+      let probs, wl = List.fold_left2 (fun (ps, wl) (t1, _) (t2, _) ->
+        let p, wl = sub_prob wl t1 EQ t2 "substitutive inv subcomp args" in
+        ps@[p], wl) ([], wl) args1 args2 in
+      bs, subst, probs, wl
+    end
+    else failwith "Impossible (rel.apply_substitutive_indexed_subcomp unexpected k" in
 
   // peel off the f:repr a is binder from bs
   let bs = List.splitAt (List.length bs - 1) bs |> fst in
@@ -1886,7 +1900,7 @@ let apply_substitutive_indexed_subcomp (env:Env.env)
     Env.pure_precondition_for_trivial_post env u subcomp_ct.result_typ wp Range.dummyRange in
 
   fml,
-  eff_params_sub_probs,
+  eff_params_sub_probs@f_g_args_eq_sub_probs,
   wl
 
 //
@@ -4314,8 +4328,7 @@ and solve_c (env:Env.env) (problem:problem comp) (wl:worklist) : solution =
             let fml, sub_probs, wl =
               if kind = Ad_hoc_combinator
               then apply_ad_hoc_indexed_subcomp env bs subcomp_c c1 c2 sub_prob wl subcomp_name r
-              else let Substitutive_combinator l = kind in
-                   apply_substitutive_indexed_subcomp env bs subcomp_c c1 c2 sub_prob
+              else apply_substitutive_indexed_subcomp env kind bs subcomp_c c1 c2 sub_prob
                    num_eff_params
                    wl
                    subcomp_name r in

diff --git a/src/typechecker/FStar.TypeChecker.TcEffect.fst b/src/typechecker/FStar.TypeChecker.TcEffect.fst
@@ -512,7 +512,7 @@ let subcomp_combinator_kind (env:env)
   (k:typ)
   (num_effect_params:int)
 
-  : option (list indexed_effect_binder_kind) =
+  : option S.indexed_effect_combinator_kind =
 
   // the idea is same as that of bind
   //   we will check that each binder in k has expected type,
@@ -551,27 +551,6 @@ let subcomp_combinator_kind (env:env)
 
     (f_bs, f_bs_kinds, rest_bs) |> Some in
 
-  let? g_bs, g_bs_kinds, rest_bs =
-    let g_sig_bs =
-      let _, sig = Env.inst_tscheme_with n_sig_ts [U_name u] in
-      sig |> U.arrow_formals
-          |> fst
-          |> (fun (a::bs) ->
-             let sig_bs, bs = List.splitAt num_effect_params bs in
-             let ss = List.fold_left2 (fun ss sig_b b ->
-               ss@[NT (sig_b.binder_bv, b.binder_bv |> S.bv_to_name)]
-             ) [NT (a.binder_bv, a_b.binder_bv |> S.bv_to_name)] sig_bs eff_params_bs in
-             bs |> SS.subst_binders ss) in
-
-    let? g_bs, rest_bs =
-      if List.length rest_bs < List.length g_sig_bs
-      then None
-      else List.splitAt (List.length g_sig_bs) rest_bs |> Some in
-
-    let? g_bs_kinds = eq_binders env g_sig_bs g_bs in
-
-    (g_bs, g_bs_kinds, rest_bs) |> Some in
-
   // peel off the f:repr a is binder
   let? rest_bs, f_b =
     if List.length rest_bs >= 1
@@ -601,36 +580,63 @@ let subcomp_combinator_kind (env:env)
     then Some ()
     else None in
 
-  // check subcomp return type is expected
-  let? _ret_t_ok_ =
+  let check_ret_t (f_or_g_bs:binders) : option unit =
     let expected_t =
       match n_repr_ts with
       | Some repr_ts ->
         let _, t = Env.inst_tscheme_with repr_ts [U_name u] in
         S.mk_Tm_app t
           ((a_b.binder_bv |> S.bv_to_name |> S.as_arg)::
-           (List.map (fun {binder_bv=b} -> b |> S.bv_to_name |> S.as_arg) (eff_params_bs@g_bs)))
+           (List.map (fun {binder_bv=b} -> b |> S.bv_to_name |> S.as_arg) (eff_params_bs@f_or_g_bs)))
           Range.dummyRange
       | None ->
         U.arrow [S.null_binder S.t_unit]
           (mk_Comp ({
              comp_univs = [U_name u];
              effect_name = n_eff_name;
              result_typ = a_b.binder_bv |> S.bv_to_name;
-             effect_args = (eff_params_bs@g_bs) |> List.map (fun b -> b.binder_bv |> S.bv_to_name |> S.as_arg);
+             effect_args = (eff_params_bs@f_or_g_bs) |> List.map (fun b -> b.binder_bv |> S.bv_to_name |> S.as_arg);
              flags = []})) in
     if U.eq_tm (U.comp_result k_c) expected_t = U.Equal
     then Some ()
     else None in
 
-  // rest of the binders are ad-hoc
-  let rest_kinds = List.map (fun _ -> Ad_hoc_binder) rest_bs in
-
-  Some ([Type_binder]        @
-        eff_params_bs_kinds  @
-        f_bs_kinds           @
-        g_bs_kinds@rest_kinds@
-        [Repr_binder])
+  if Some? (check_ret_t f_bs)
+  then Some Substitutive_invariant_combinator
+  else begin
+    let? g_bs, g_bs_kinds, rest_bs =
+      let g_sig_bs =
+        let _, sig = Env.inst_tscheme_with n_sig_ts [U_name u] in
+        sig |> U.arrow_formals
+            |> fst
+            |> (fun (a::bs) ->
+               let sig_bs, bs = List.splitAt num_effect_params bs in
+               let ss = List.fold_left2 (fun ss sig_b b ->
+                 ss@[NT (sig_b.binder_bv, b.binder_bv |> S.bv_to_name)]
+               ) [NT (a.binder_bv, a_b.binder_bv |> S.bv_to_name)] sig_bs eff_params_bs in
+               bs |> SS.subst_binders ss) in
+
+      let? g_bs, rest_bs =
+        if List.length rest_bs < List.length g_sig_bs
+        then None
+        else List.splitAt (List.length g_sig_bs) rest_bs |> Some in
+
+      let? g_bs_kinds = eq_binders env g_sig_bs g_bs in
+
+      (g_bs, g_bs_kinds, rest_bs) |> Some in
+
+    // check subcomp return type is expected
+    let? _ret_t_ok_ = check_ret_t g_bs in
+
+    // rest of the binders are ad-hoc
+    let rest_kinds = List.map (fun _ -> Ad_hoc_binder) rest_bs in
+
+    Some (([Type_binder]        @
+           eff_params_bs_kinds  @
+           f_bs_kinds           @
+           g_bs_kinds@rest_kinds@
+           [Repr_binder]) |> Substitutive_combinator)
+  end
 
 //
 // Validate indexed effect subcomp (including polymonadic subcomp) shape
@@ -723,19 +729,19 @@ let validate_indexed_effect_subcomp_shape (env:env)
 
   let k = k |> N.remove_uvar_solutions env |> SS.compress in
 
-  let lopt = subcomp_combinator_kind env m_eff_name n_eff_name
+  let kopt = subcomp_combinator_kind env m_eff_name n_eff_name
     m_sig_ts n_sig_ts
     m_repr_ts n_repr_ts
     u
     k
     num_effect_params in
 
   let kind =
-    match lopt with
+    match kopt with
     | None ->
       log_ad_hoc_combinator_warning subcomp_name r;
       Ad_hoc_combinator
-    | Some l -> Substitutive_combinator l in
+    | Some k -> k in
 
   if Env.debug env <| Options.Other "LayeredEffectsTc"
   then BU.print2 "Subcomp %s has %s kind\n" subcomp_name