feat: remove loop closures that capture mutable variables

Changes.md

@@ -12,6 +12,8 @@ Unreleased
   ([#1050](https://github.com/melange-re/melange/pull/1050))
 - core: emit `let` instead of `var` in compiled JS
   ([#1019](https://github.com/melange-re/melange/pull/1019))
+- core: in compiled JS, stop generating closures in loops that capture mutable
+  variables ([#1020](https://github.com/melange-re/melange/pull/1020))
 
 3.0.0 2024-01-28
 ---------------

jscomp/core/j.mli

@@ -270,15 +270,14 @@ and statement_desc =
   (* Function declaration and Variable declaration  *)
   | Exp of expression
   | If of expression * block * block
-  | While of label option * expression * block * Js_closure.t
+  | While of label option * expression * block
     (* check if it contains loop mutable values, happens in nested loop *)
   | ForRange of
       for_ident_expression option
       * finish_ident_expression
       * for_ident
       * for_direction
       * block
-      * Js_closure.t
   | Continue of label
   | Break (* only used when inline a fucntion *)
   | Return of expression

jscomp/core/js_dump.ml

@@ -207,13 +207,6 @@ let exp_need_paren (e : J.expression) =
   | Typeof _ | Number _ | Js_not _ | Bool _ | New _ ->
       false
 
-let comma_idents (cxt : cxt) ls = iter_lst cxt ls ident comma
-
-let pp_paren_params (cxt : cxt) (lexical : Ident.t list) : unit =
-  string cxt L.lparen;
-  let (_ : cxt) = comma_idents cxt lexical in
-  string cxt L.rparen
-
 (* Print as underscore for unused vars, may not be
     needed in the future *)
 (* let ipp_ident cxt id (un_used : bool) =
@@ -438,70 +431,29 @@ and pp_function ~return_unit ~is_method cxt ~fn_state (l : Ident.t list)
           space cxt;
           brace_vgroup cxt 1 (fun _ -> function_body ~return_unit cxt b)
       in
-      let lexical : Ident.Set.t = Js_fun_env.get_lexical_scope env in
-      let enclose lexical =
-        let handle lexical =
-          if Ident.Set.is_empty lexical then (
-            match fn_state with
-            | Is_return ->
-                return_sp cxt;
-                string cxt L.function_;
-                space cxt;
-                param_body ()
-            | No_name { single_arg } ->
-                (* see # 1692, add a paren for annoymous function for safety  *)
-                cond_paren_group cxt (not single_arg) 1 (fun _ ->
-                    string cxt L.function_;
-                    space cxt;
-                    param_body ())
-            | Name_non_top x ->
-                ignore (pp_var_assign inner_cxt x : cxt);
-                string cxt L.function_;
-                space cxt;
-                param_body ();
-                semi cxt
-            | Name_top x ->
-                string cxt L.function_;
-                space cxt;
-                ignore (ident inner_cxt x : cxt);
-                param_body ())
-          else
-            (* print our closure as
-               {[(function(x,y){ return function(..){...}} (x,y))]}
-               Maybe changed to `let` in the future
-            *)
-            let lexical = Ident.Set.elements lexical in
-            (match fn_state with
-            | Is_return -> return_sp cxt
-            | No_name _ -> ()
-            | Name_non_top name | Name_top name ->
-                ignore (pp_var_assign inner_cxt name : cxt));
-            string cxt L.lparen;
-            string cxt L.function_;
-            pp_paren_params inner_cxt lexical;
-            brace_vgroup cxt 0 (fun _ ->
-                return_sp cxt;
-                string cxt L.function_;
-                space cxt;
-                (match fn_state with
-                | Is_return | No_name _ -> ()
-                | Name_non_top x | Name_top x -> ignore (ident inner_cxt x));
-                param_body ());
-            pp_paren_params inner_cxt lexical;
-            string cxt L.rparen;
-            match fn_state with
-            | Is_return | No_name _ -> () (* expression *)
-            | _ -> semi cxt (* has binding, a statement *)
-        in
-        handle
-          (match fn_state with
-          | (Name_top name | Name_non_top name) when Ident.Set.mem lexical name
-            ->
-              (*TODO: when calculating lexical we should not include itself *)
-              Ident.Set.remove lexical name
-          | _ -> lexical)
-      in
-      enclose lexical;
+      (match fn_state with
+      | Is_return ->
+          return_sp cxt;
+          string cxt L.function_;
+          space cxt;
+          param_body ()
+      | No_name { single_arg } ->
+          (* see # 1692, add a paren for annoymous function for safety  *)
+          cond_paren_group cxt (not single_arg) 1 (fun _ ->
+              string cxt L.function_;
+              space cxt;
+              param_body ())
+      | Name_non_top x ->
+          ignore (pp_var_assign inner_cxt x : cxt);
+          string cxt L.function_;
+          space cxt;
+          param_body ();
+          semi cxt
+      | Name_top x ->
+          string cxt L.function_;
+          space cxt;
+          ignore (ident inner_cxt x : cxt);
+          param_body ());
       outer_cxt
 
 (* Assume the cond would not change the context,
@@ -1050,7 +1002,7 @@ and statement_desc top cxt (s : J.statement_desc) : cxt =
           string cxt L.else_;
           space cxt;
           brace_block cxt s2)
-  | While (label, e, s, _env) ->
+  | While (label, e, s) ->
       (*  FIXME: print scope as well *)
       (match label with
       | Some i ->
@@ -1076,7 +1028,7 @@ and statement_desc top cxt (s : J.statement_desc) : cxt =
       let cxt = brace_block cxt s in
       semi cxt;
       cxt
-  | ForRange (for_ident_expression, finish, id, direction, s, env) ->
+  | ForRange (for_ident_expression, finish, id, direction, s) ->
       let action cxt =
         vgroup cxt 0 (fun _ ->
             let cxt =
@@ -1143,24 +1095,7 @@ and statement_desc top cxt (s : J.statement_desc) : cxt =
             in
             brace_block cxt s)
       in
-      let lexical = Js_closure.get_lexical_scope env in
-      if Ident.Set.is_empty lexical then action cxt
-      else
-        (* unlike function,
-           [print for loop] has side effect,
-           we should take it out
-        *)
-        let inner_cxt = merge_scope cxt lexical in
-        let lexical = Ident.Set.elements lexical in
-        vgroup cxt 0 (fun _ ->
-            string cxt L.lparen;
-            string cxt L.function_;
-            pp_paren_params inner_cxt lexical;
-            let cxt = brace_vgroup cxt 0 (fun _ -> action inner_cxt) in
-            pp_paren_params inner_cxt lexical;
-            string cxt L.rparen;
-            semi cxt;
-            cxt)
+      action cxt
   | Continue s ->
       continue cxt s;
       cxt (* newline cxt;  #2642 *)

jscomp/core/js_fun_env.ml

@@ -46,7 +46,6 @@ type immutable_mask =
 
 type t = {
   mutable unbounded : Ident.Set.t;
-  mutable bound_loop_mutable_values : Ident.Set.t;
   used_mask : bool array;
   immutable_mask : immutable_mask;
 }
@@ -60,7 +59,6 @@ let make ?immutable_mask n =
       (match immutable_mask with
       | Some x -> Immutable_mask x
       | None -> All_immutable_and_no_tail_call);
-    bound_loop_mutable_values = Ident.Set.empty;
   }
 
 let no_tailcall x =
@@ -91,13 +89,3 @@ let set_unbounded env v =
   (* if Ident.Set.is_empty env.bound then *)
   env.unbounded <- v
 (* else assert false *)
-
-let set_lexical_scope env bound_loop_mutable_values =
-  env.bound_loop_mutable_values <- bound_loop_mutable_values
-
-let get_lexical_scope env = env.bound_loop_mutable_values
-
-(* TODO:  can be refined if it
-    only enclose toplevel variables
-*)
-(* let is_empty t = Ident.Set.is_empty t.unbounded *)

jscomp/core/js_fun_env.mli

@@ -32,8 +32,6 @@ type t
 val make : ?immutable_mask:bool array -> int -> t
 val no_tailcall : t -> bool list
 val set_unbounded : t -> Ident.Set.t -> unit
-val set_lexical_scope : t -> Ident.Set.t -> unit
-val get_lexical_scope : t -> Ident.Set.t
 val mark_unused : t -> int -> unit
 val get_unused : t -> int -> bool
 val get_mutable_params : Ident.t list -> t -> Ident.t list

jscomp/core/js_pass_scope.ml

@@ -178,10 +178,6 @@ let record_scope_pass =
                due to the recursive thing
             *)
             Js_fun_env.set_unbounded env closured_idents';
-            let lexical_scopes =
-              Ident.Set.(inter closured_idents' state.loop_mutable_values)
-            in
-            Js_fun_env.set_lexical_scope env lexical_scopes;
             (* tailcall , note that these varibles are used in another pass *)
             {
               state with
@@ -247,7 +243,7 @@ let record_scope_pass =
     statement =
       (fun self state x ->
         match x.statement_desc with
-        | ForRange (_, _, loop_id, _, _, a_env) ->
+        | ForRange (_, _, loop_id, _, _) ->
             (* TODO: simplify definition of For *)
             let {
               defined_idents = defined_idents';
@@ -280,8 +276,6 @@ let record_scope_pass =
                   (diff closured_idents' defined_idents')
                   state.loop_mutable_values)
             in
-            let () = Js_closure.set_lexical_scope a_env lexical_scope in
-            (* set scope *)
             {
               state with
               used_idents = Ident.Set.union state.used_idents used_idents';
@@ -299,7 +293,7 @@ let record_scope_pass =
               closured_idents =
                 Ident.Set.union state.closured_idents lexical_scope;
             }
-        | While (_label, pred, body, _env) ->
+        | While (_label, pred, body) ->
             with_in_loop
               (self.block self
                  (with_in_loop (self.expression self state pred) true)

jscomp/core/js_stmt_make.ml

@@ -318,17 +318,14 @@ let if_ ?comment ?declaration ?else_ (e : J.expression) (then_ : J.block) : t =
 let assign ?comment id e : t =
   { statement_desc = J.Exp (E.assign (E.var id) e); comment }
 
-let while_ ?comment ?label ?env (e : E.t) (st : J.block) : t =
-  let env = match env with None -> Js_closure.empty () | Some x -> x in
-  { statement_desc = While (label, e, st, env); comment }
+let while_ ?comment ?label (e : E.t) (st : J.block) : t =
+  { statement_desc = While (label, e, st); comment }
 
-let for_ ?comment ?env for_ident_expression finish_ident_expression id direction
+let for_ ?comment for_ident_expression finish_ident_expression id direction
     (b : J.block) : t =
-  let env = match env with None -> Js_closure.empty () | Some x -> x in
   {
     statement_desc =
-      ForRange
-        (for_ident_expression, finish_ident_expression, id, direction, b, env);
+      ForRange (for_ident_expression, finish_ident_expression, id, direction, b);
     comment;
   }
 

jscomp/core/js_stmt_make.mli

@@ -130,17 +130,10 @@ val assign : ?comment:string -> J.ident -> J.expression -> t
    J.ident ->
    t *)
 
-val while_ :
-  ?comment:string ->
-  ?label:J.label ->
-  ?env:Js_closure.t ->
-  J.expression ->
-  J.block ->
-  t
+val while_ : ?comment:string -> ?label:J.label -> J.expression -> J.block -> t
 
 val for_ :
   ?comment:string ->
-  ?env:Js_closure.t ->
   J.for_ident_expression option ->
   J.finish_ident_expression ->
   J.for_ident ->