Create closures with only necessary environment

tools/liquidity/liquidBoundVariables.ml

@@ -19,6 +19,19 @@ let rec bv code =
                      (StringSet.union (bv ifthen) (bv ifelse))
   | Seq (x, y) -> StringSet.union (bv x) (bv y)
   | Const (ty, cst) ->  StringSet.empty
+
+  | Apply (Prim_unknown, _,
+           ({ desc = Var (name, _, [])} :: args)) ->
+    let set =
+      try
+        LiquidTypes.primitive_of_string name |> ignore;
+        StringSet.empty
+      with Not_found -> StringSet.singleton name
+    in
+    List.fold_left (fun set arg ->
+        StringSet.union set (bv arg)
+      ) set args
+
   | Apply (prim, loc, args) ->
      List.fold_left (fun set arg ->
          StringSet.union set (bv arg)
@@ -115,6 +128,23 @@ let rec bound code =
   | Const (ty, cst) ->
      mk code.desc code StringSet.empty
 
+  | Apply (Prim_unknown, loc,
+           ({ desc = Var (name, varloc, [])} :: args)) ->
+    let args = List.map bound args in
+    let bv =
+      try
+        LiquidTypes.primitive_of_string name |> ignore;
+        StringSet.empty
+      with Not_found -> StringSet.singleton name
+    in
+    let v = mk (Var (name, varloc, [])) code bv in
+    let bv =
+      List.fold_left (fun set arg ->
+        StringSet.union set arg.bv
+        ) bv args
+    in
+    mk (Apply(Prim_unknown, loc, v :: args)) code bv
+
   | Apply (prim, loc, args) ->
      let args = List.map bound args in
      let bv =

tools/liquidity/liquidCheck.ml

@@ -134,23 +134,23 @@ let find_var ?(count_used=true) env loc name =
       error loc "unbound variable %S" name
 
 (* Create environment for closure *)
-let env_for_clos env loc arg_name arg_type =
-  let free_vars = match env.clos_env with
-    | Some ce ->
-      StringMap.map
-        (fun (bname, btype, index, (cpt_in, cpt_out)) ->
-           (bname, btype, index + 1, (cpt_in, cpt_out)))
-        ce.env_vars
-    | None -> StringMap.empty
-  in
-  let _, free_vars =
-    StringMap.fold (fun n (bname, btype, count) (index, free_vars) ->
-      match btype with
-      | Tlambda _ -> (index, free_vars)
-      | _ ->
+let env_for_clos env loc bvs arg_name arg_type =
+  let _, free_vars = StringSet.fold (fun v (index, free_vars) ->
+      try
         let index = index + 1 in
-        (index, StringMap.add n (bname, btype, index, (ref 0, count)) free_vars)
-    ) env.vars (StringMap.cardinal free_vars, free_vars)
+        match env.clos_env with
+        | None ->
+          let (bname, btype, cpt_out) = StringMap.find v env.vars in
+          (index,
+           StringMap.add v (bname, btype, index, (ref 0, cpt_out)) free_vars)
+        | Some ce ->
+          let bname, btype, _, (cpt_in, cpt_out) =
+            StringMap.find v ce.env_vars in
+          (index,
+           StringMap.add v (bname, btype, index, (cpt_in, cpt_out)) free_vars)
+      with Not_found ->
+        (index, free_vars)
+    ) bvs (0, StringMap.empty)
   in
   let free_vars_l =
     StringMap.bindings free_vars
@@ -737,47 +737,46 @@ let rec loc_exp env e = match e.desc with
        assert (res_type = Tunit);
        (* let env = { env with vars = StringMap.empty } in *)
        (* let (arg_name, env, arg_count) = new_binding env arg_name arg_type in *)
-       let env, arg_name, arg_type, call_env =
-         env_for_clos env loc arg_name arg_type in
-       let body, _fail, transfer = typecheck env body in
-       if transfer then
-         error loc "no transfer in lambda";
-       let is_real_closure = match env.clos_env with
-         | None -> false
-         | Some clos_env ->
-           StringMap.exists (fun name (_, (cpt_in, _)) ->
-               !cpt_in <> 0 && name <> lambda_arg_name
-             ) clos_env.env_bindings
-       in
-       (* begin match env.clos_env with *)
-       (*   | None -> () *)
-       (*   | Some clos_env -> *)
-       (*     Format.eprintf "--- Closure %s (real:%b)---@." arg_name is_real_closure; *)
-       (*     StringMap.iter (fun name (e, (cpt_in, cpt_out)) -> *)
-       (*         Format.eprintf "%s -> %s , (%d, %d)@." *)
-       (*           name (LiquidPrinter.Liquid.string_of_code e) !cpt_in !cpt_out *)
-       (*       ) clos_env.env_bindings *)
-       (* end; *)
-       if not is_real_closure then
-         (* recreate lambda from scratch *)
+       let bvs = LiquidBoundVariables.bv exp in
+       if StringSet.is_empty bvs then
+         (* not a closure, create a real lambda *)
          let env = { env_at_lambda with vars = StringMap.empty } in
          let (new_arg_name, env, arg_count) =
            new_binding env lambda_arg_name lambda_arg_type in
          let body, _fail, transfer = typecheck env lambda_body in
+         if transfer then
+           error loc "no transfer in lambda";
          check_used env lambda_arg_name loc arg_count;
          let desc =
            Lambda (new_arg_name, lambda_arg_type, loc, body, body.ty) in
          let ty = Tlambda (lambda_arg_type, body.ty) in
          mk desc ty false, false, false
-       else begin
+       else
+         (* create closure with environment *)
+         let env, arg_name, arg_type, call_env =
+           env_for_clos env loc bvs arg_name arg_type in
+         let body, _fail, transfer = typecheck env body in
+         if transfer then
+           error loc "no transfer in closure";
+         (* begin match env.clos_env with *)
+         (*   | None -> () *)
+         (*   | Some clos_env -> *)
+         (*     Format.eprintf "--- Closure %s (real:%b)---@." arg_name is_real_closure; *)
+         (*     StringMap.iter (fun name (e, (cpt_in, cpt_out)) -> *)
+         (*         Format.eprintf "%s -> %s , (%d, %d)@." *)
+         (*           name (LiquidPrinter.Liquid.string_of_code e) !cpt_in !cpt_out *)
+         (*       ) clos_env.env_bindings *)
+         (* end; *)
          check_used_in_env env lambda_arg_name loc;
          let desc =
            Closure (arg_name, arg_type, loc, call_env, body, body.ty) in
-         let call_env_type =
-           Ttuple (List.map (fun (_, t) -> t.ty) call_env) in
+         let call_env_type = match call_env with
+           | [] -> assert false
+           | [_, t] -> t.ty
+           | _ -> Ttuple (List.map (fun (_, t) -> t.ty) call_env)
+         in
          let ty = Tclosure ((lambda_arg_type, call_env_type), body.ty) in
          mk desc ty false, false, false
-       end
 
     | Closure _ -> assert false