Fix Generalized Closures (#1453)

src/haz3lcore/dynamics/DHExp.re

@@ -106,7 +106,7 @@ let rec strip_casts =
 let assign_name_if_none = (t, name) => {
   let (term, rewrap) = unwrap(t);
   switch (term) {
-  | Fun(arg, ty, body, None) => Fun(arg, ty, body, name) |> rewrap
+  | Fun(arg, body, typ, None) => Fun(arg, body, typ, name) |> rewrap
   | TypFun(utpat, body, None) => TypFun(utpat, body, name) |> rewrap
   | _ => t
   };

src/haz3lcore/dynamics/EvalCtx.re

@@ -8,7 +8,7 @@ type term =
   | Seq2(DHExp.t, t)
   | Let1(Pat.t, t, DHExp.t)
   | Let2(Pat.t, DHExp.t, t)
-  | Fun(Pat.t, t, option(ClosureEnvironment.t), option(Var.t))
+  | Fun(Pat.t, t, option(Typ.t), option(Var.t))
   | FixF(Pat.t, t, option(ClosureEnvironment.t))
   | TypAp(t, Typ.t)
   | Ap1(Operators.ap_direction, t, DHExp.t)
@@ -124,9 +124,9 @@ let rec compose = (ctx: t, d: DHExp.t): DHExp.t => {
     | Let2(dp, d1, ctx) =>
       let d = compose(ctx, d);
       Let(dp, d1, d) |> wrap;
-    | Fun(dp, ctx, env, v) =>
+    | Fun(dp, ctx, typ, v) =>
       let d = compose(ctx, d);
-      Fun(dp, d, env, v) |> wrap;
+      Fun(dp, d, typ, v) |> wrap;
     | FixF(v, ctx, env) =>
       let d = compose(ctx, d);
       FixF(v, d, env) |> wrap;

src/haz3lcore/dynamics/Evaluator.re

@@ -98,9 +98,9 @@ module EvaluatorEVMode: {
 
 module Eval = Transition(EvaluatorEVMode);
 
-let rec evaluate = (state, env, d) => {
+let rec evaluate = (~in_closure=?, state, env, d) => {
   open Trampoline.Syntax;
-  let.trampoline u = Eval.transition(evaluate, state, env, d);
+  let.trampoline u = Eval.transition(evaluate, ~in_closure?, state, env, d);
   switch (u) {
   | (Final, x) => (EvaluatorEVMode.Final, x) |> Trampoline.return
   | (Uneval, x) => Trampoline.Next(() => evaluate(state, env, x))

src/haz3lcore/dynamics/EvaluatorStep.re

@@ -93,10 +93,10 @@ let rec matches =
       | Let2(d1, d2, ctx) =>
         let+ ctx = matches(env, flt, ctx, exp, act, idx);
         Let2(d1, d2, ctx) |> rewrap;
-      | Fun(dp, ctx, env', name) =>
-        let+ ctx =
-          matches(Option.value(~default=env, env'), flt, ctx, exp, act, idx);
-        Fun(dp, ctx, env', name) |> rewrap;
+      | Fun(dp, ctx, ty, name) =>
+        // TODO: Should this env include the bound variables?
+        let+ ctx = matches(env, flt, ctx, exp, act, idx);
+        Fun(dp, ctx, ty, name) |> rewrap;
       | FixF(name, ctx, env') =>
         let+ ctx =
           matches(Option.value(~default=env, env'), flt, ctx, exp, act, idx);
@@ -301,9 +301,9 @@ module Decompose = {
   };
 
   module Decomp = Transition(DecomposeEVMode);
-  let rec decompose = (state, env, exp) => {
+  let rec decompose = (~in_closure=?, state, env, exp) => {
     switch (exp) {
-    | _ => Decomp.transition(decompose, state, env, exp)
+    | _ => Decomp.transition(decompose, ~in_closure?, state, env, exp)
     };
   };
 };
@@ -343,8 +343,14 @@ module TakeStep = {
 
   module TakeStepEV = Transition(TakeStepEVMode);
 
-  let take_step = (state, env, d) =>
-    TakeStepEV.transition((_, _, _) => None, state, env, d)
+  let take_step = (~in_closure=?, state, env, d) =>
+    TakeStepEV.transition(
+      (~in_closure as _=?, _, _, _) => None,
+      ~in_closure?,
+      state,
+      env,
+      d,
+    )
     |> Option.map(DHExp.repair_ids);
 };
 

src/haz3lcore/dynamics/FilterMatcher.re

@@ -228,14 +228,7 @@ let rec matches_exp =
     | (TypFun(pat1, d1, s1), TypFun(pat2, d2, s2)) =>
       s1 == s2 && matches_utpat(pat1, pat2) && matches_exp(d1, d2)
     | (TypFun(_), _) => false
-
-    | (Fun(dp1, d1, Some(denv), _), Fun(fp1, f1, Some(fenv), _)) =>
-      matches_fun(~denv, dp1, d1, ~fenv, fp1, f1)
-    | (Fun(dp1, d1, Some(denv), _), Fun(fp1, f1, None, _)) =>
-      matches_fun(~denv, dp1, d1, ~fenv, fp1, f1)
-    | (Fun(dp1, d1, None, _), Fun(fp1, f1, Some(fenv), _)) =>
-      matches_fun(~denv, dp1, d1, ~fenv, fp1, f1)
-    | (Fun(dp1, d1, None, _), Fun(fp1, f1, None, _)) =>
+    | (Fun(dp1, d1, _, _), Fun(fp1, f1, _, _)) =>
       matches_fun(~denv, dp1, d1, ~fenv, fp1, f1)
     | (Fun(_), _) => false
 

src/haz3lcore/dynamics/Substitution.re

@@ -65,16 +65,13 @@ let rec subst_var = (m, d1: DHExp.t, x: Var.t, d2: DHExp.t): DHExp.t => {
         subst_var(m, d1, x, d3);
       };
     FixF(y, d3, env') |> rewrap;
-  | Fun(dp, d3, env, s) =>
-    /* Function closure shouldn't appear during substitution
-       (which only is called from elaboration currently) */
-    let env' = Option.map(subst_var_env(m, d1, x), env);
+  | Fun(dp, d3, ty, s) =>
     if (binds_var(m, x, dp)) {
-      Fun(dp, d3, env', s) |> rewrap;
+      Fun(dp, d3, ty, s) |> rewrap;
     } else {
       let d3 = subst_var(m, d1, x, d3);
-      Fun(dp, d3, env', s) |> rewrap;
-    };
+      Fun(dp, d3, ty, s) |> rewrap;
+    }
   | TypFun(tpat, d3, s) =>
     TypFun(tpat, subst_var(m, d1, x, d3), s) |> rewrap
   | Closure(env, d3) =>

src/haz3lcore/dynamics/Transition.re

@@ -23,6 +23,10 @@ open PatternMatch;
     are all enumerated by the <i> field, so i=0 indicates that it is the first
     sub-expression, i=1 the second etc.
 
+    If there are any sub-expressions that are not requirements, and therefore not
+    guaranteed to be run, you should add a `let.wrap_closure () = env` to ensure that
+    the closure isn't lost if the expression is indet.
+
     Finally, we have the Step construct that defines the actual step. Note "Step"s
     should be used if and only if they change the expression. If they do not change
     the expression, use `Constructor` or `Indet`.
@@ -47,7 +51,7 @@ type step_kind =
   | VarLookup
   | Seq
   | LetBind
-  | FunClosure
+  | WrapClosure
   | FixUnwrap
   | FixClosure
   | UpdateTest
@@ -140,15 +144,51 @@ module Transition = (EV: EV_MODE) => {
     | Matches(env') => r(evaluate_extend_env(env', env))
     };
 
+  /* Helper function to wrap a closure around an expression. Required for functions, but also for
+     things like if-then-else expressions where the scrutinee is indet, and for hole closures */
+  let wrap_closure_when_done = (~in_closure, expr, env, r: rule) =>
+    switch (in_closure, r) {
+    | (_, Step(_)) => r
+    | (None, Constructor | Indet | Value) =>
+      Step({
+        expr: Closure(env, expr) |> fresh,
+        state_update,
+        kind: WrapClosure,
+        is_value: false,
+      })
+    | (Some(f), Constructor | Indet | Value) =>
+      f();
+      r;
+    };
+
   /* Note[Matt]: For IDs, I'm currently using a fresh id
      if anything about the current node changes, if only its
      children change, we use rewrap */
 
-  let transition = (req, state, env, d): 'a => {
+  let transition =
+      (
+        req:
+          (
+            ~in_closure: unit => unit=?,
+            state,
+            ClosureEnvironment.t,
+            DHExp.t
+          ) =>
+          'a,
+        ~in_closure=?,
+        state,
+        env,
+        d,
+      )
+      : 'a => {
     // Split DHExp into term and id information
     let (term, rewrap) = DHExp.unwrap(d);
     let wrap_ctx = (term): EvalCtx.t => Term({term, ids: [rep_id(d)]});
 
+    let (let.wrap_closure) = (env, f: unit => rule) => {
+      wrap_closure_when_done(~in_closure, d, env, f());
+    };
+
     // Transition rules
     switch (term) {
     | Var(x) =>
@@ -167,7 +207,9 @@ module Transition = (EV: EV_MODE) => {
           kind: VarLookup,
           is_value,
         });
-      | None => Indet
+      | None =>
+        let.wrap_closure _ = env;
+        Indet;
       };
     | Seq(d1, d2) =>
       let. _ = otherwise(env, d1 => Seq(d1, d2) |> rewrap)
@@ -178,6 +220,7 @@ module Transition = (EV: EV_MODE) => {
       let. _ = otherwise(env, d1 => Let(dp, d1, d2) |> rewrap)
       and. d1' =
         req_final(req(state, env), d1 => Let1(dp, d1, d2) |> wrap_ctx, d1);
+      let.wrap_closure _ = env;
       let.match env' = (env, matches(dp, d1'));
       Step({
         expr: Closure(env', d2) |> fresh,
@@ -186,17 +229,10 @@ module Transition = (EV: EV_MODE) => {
         is_value: false,
       });
     | TypFun(_)
-    | Fun(_, _, Some(_), _) =>
-      let. _ = otherwise(env, d);
-      Constructor;
-    | Fun(p, d1, None, v) =>
+    | Fun(_, _, _, _) =>
       let. _ = otherwise(env, d);
-      Step({
-        expr: Fun(p, d1, Some(env), v) |> rewrap,
-        state_update,
-        kind: FunClosure,
-        is_value: true,
-      });
+      let.wrap_closure _ = env;
+      Value;
     | FixF(dp, d1, None) =>
       let. _ = otherwise(env, FixF(dp, d1, None) |> rewrap);
       Step({
@@ -415,6 +451,7 @@ module Transition = (EV: EV_MODE) => {
       let. _ = otherwise(env, c => If(c, d1, d2) |> rewrap)
       and. c' =
         req_final(req(state, env), c => If1(c, d1, d2) |> wrap_ctx, c);
+      let.wrap_closure _ = env;
       let-unbox b = (Bool, c');
       Step({
         expr: {
@@ -466,6 +503,7 @@ module Transition = (EV: EV_MODE) => {
           d1 => BinOp1(Bool(And), d1, d2) |> wrap_ctx,
           d1,
         );
+      let.wrap_closure _ = env;
       let-unbox b1 = (Bool, d1');
       Step({
         expr: b1 ? d2 : Bool(false) |> fresh,
@@ -481,6 +519,7 @@ module Transition = (EV: EV_MODE) => {
           d1 => BinOp1(Bool(Or), d1, d2) |> wrap_ctx,
           d1,
         );
+      let.wrap_closure _ = env;
       let-unbox b1 = (Bool, d1');
       Step({
         expr: b1 ? Bool(true) |> fresh : d2,
@@ -682,26 +721,40 @@ module Transition = (EV: EV_MODE) => {
           kind: CaseApply,
           is_value: false,
         })
-      | None => Indet
+      | None =>
+        let.wrap_closure _ = env;
+        Indet;
       };
     | Closure(env', d) =>
+      // HACK [Matt] This ref is a hack to ensure that we don't get into an infinite loop
+      // where we keep deleting and re-adding closures around forms that need closures
+      // e.g. functions.
+      let needs_closure = ref(false);
+      let in_closure = () => needs_closure := true;
       let. _ = otherwise(env, d => Closure(env', d) |> rewrap)
       and. d' =
-        req_final(req(state, env'), d1 => Closure(env', d1) |> wrap_ctx, d);
-      Step({expr: d', state_update, kind: CompleteClosure, is_value: true});
+        req_final(
+          req(~in_closure, state, env'),
+          d1 => Closure(env', d1) |> wrap_ctx,
+          d,
+        );
+      if (needs_closure^) {
+        Constructor;
+      } else {
+        Step({expr: d', state_update, kind: CompleteClosure, is_value: true});
+      };
     | MultiHole(_) =>
       let. _ = otherwise(env, d);
-      // and. _ =
-      //   req_all_final(
-      //     req(state, env),
-      //     (d1, ds) => MultiHole(d1, ds) |> wrap_ctx,
-      //     ds,
-      //   );
+      let.wrap_closure _ = env;
       Indet;
     | EmptyHole
-    | Invalid(_)
+    | Invalid(_) =>
+      let. _ = otherwise(env, d);
+      // let.wrap_closure _ = env;  // uncomment for hole closures
+      Indet;
     | DynamicErrorHole(_) =>
       let. _ = otherwise(env, d);
+      let.wrap_closure _ = env;
       Indet;
     | Cast(d, t1, t2) =>
       let. _ = otherwise(env, d => Cast(d, t1, t2) |> rewrap)
@@ -767,7 +820,7 @@ let should_hide_step_kind = (~settings: CoreSettings.Evaluation.t) =>
   | CompleteClosure
   | CompleteFilter
   | BuiltinWrap
-  | FunClosure
+  | WrapClosure
   | FixClosure
   | RemoveParens => true;
 
@@ -806,7 +859,7 @@ let stepper_justification: step_kind => string =
   | FixClosure => "fixpoint closure"
   | CompleteFilter => "complete filter"
   | CompleteClosure => "complete closure"
-  | FunClosure => "function closure"
+  | WrapClosure => "wrap closure"
   | RemoveTypeAlias => "define type"
   | RemoveParens => "remove parentheses"
   | UnOp(Meta(Unquote)) => failwith("INVALID STEP");

src/haz3lcore/dynamics/TypeAssignment.re

@@ -152,13 +152,13 @@ and typ_of_dhexp = (ctx: Ctx.t, m: Statics.Map.t, dh: DHExp.t): option(Typ.t) =>
       };
     let* ctx = dhpat_extend_ctx(dhp, ty_p, ctx);
     typ_of_dhexp(ctx, m, d);
-  | Fun(dhp, d, env, _) =>
-    let* ty_p = dhpat_synthesize(dhp, ctx);
-    let* ctx =
-      switch (env) {
-      | None => Some(ctx)
-      | Some(env) => env_extend_ctx(env, m, ctx)
+  | Fun(dhp, d, ty, _) =>
+    let* ty_p =
+      switch (ty) {
+      | None => dhpat_synthesize(dhp, ctx)
+      | Some(t) => Some(t)
       };
+
     let* ctx = dhpat_extend_ctx(dhp, ty_p, ctx);
     let* ty2 = typ_of_dhexp(ctx, m, d);
     Some(Arrow(ty_p, ty2) |> Typ.temp);

src/haz3lcore/dynamics/Unboxing.re

@@ -150,7 +150,8 @@ let rec unbox: type a. (unbox_request(a), DHExp.t) => unboxed(a) =
 
     /* Function-like things can look like the following when values */
     | (Fun, Constructor(name, _)) => Matches(Constructor(name)) // Perhaps we should check if the constructor actually is a function?
-    | (Fun, Fun(dp, d3, Some(env'), _)) => Matches(FunEnv(dp, d3, env'))
+    | (Fun, Closure(env', {term: Fun(dp, d3, _, _), _})) =>
+      Matches(FunEnv(dp, d3, env'))
     | (
         Fun,
         Cast(
@@ -191,7 +192,6 @@ let rec unbox: type a. (unbox_request(a), DHExp.t) => unboxed(a) =
         BuiltinFun(_) |
         Deferral(_) |
         DeferredAp(_) |
-        Fun(_, _, _, Some(_)) |
         ListLit(_) |
         Tuple(_) |
         Cast(_) |
@@ -221,7 +221,7 @@ let rec unbox: type a. (unbox_request(a), DHExp.t) => unboxed(a) =
         Invalid(_) | Undefined | EmptyHole | MultiHole(_) | DynamicErrorHole(_) |
         Var(_) |
         Let(_) |
-        Fun(_, _, _, None) |
+        Fun(_, _, _, _) |
         FixF(_) |
         TyAlias(_) |
         Ap(_) |

src/haz3lcore/dynamics/ValueChecker.re

@@ -51,15 +51,15 @@ module ValueCheckerEVMode: {
 
 module CV = Transition(ValueCheckerEVMode);
 
-let rec check_value = (state, env, d) =>
-  CV.transition(check_value, state, env, d);
+let rec check_value = (~in_closure=?, state, env, d) =>
+  CV.transition(check_value, ~in_closure?, state, env, d);
 
-let rec check_value_mod_ctx = ((), env, d) =>
+let rec check_value_mod_ctx = (~in_closure=?, (), env, d) =>
   switch (DHExp.term_of(d)) {
   | Var(x) =>
     switch (ClosureEnvironment.lookup(env, x)) {
-    | Some(v) => check_value_mod_ctx((), env, v)
-    | None => CV.transition(check_value_mod_ctx, (), env, d)
+    | Some(v) => check_value_mod_ctx(~in_closure?, (), env, v)
+    | None => CV.transition(check_value_mod_ctx, ~in_closure?, (), env, d)
     }
-  | _ => CV.transition(check_value_mod_ctx, (), env, d)
+  | _ => CV.transition(check_value_mod_ctx, ~in_closure?, (), env, d)
   };