@@ -281,13 +281,15 @@ module Ctx = struct
281281 ; effect_warning : bool ref
282282 ; cps_calls : Effects .cps_calls
283283 ; deadcode_sentinal : Var .t
284+ ; mutated_vars : Code.Var.Set .t Code.Addr.Map .t
284285 }
285286
286287 let initial
287288 ~warn_on_unhandled_effect
288289 ~exported_runtime
289290 ~should_export
290291 ~deadcode_sentinal
292+ ~mutated_vars
291293 blocks
292294 live
293295 cps_calls
@@ -302,6 +304,7 @@ module Ctx = struct
302304 ; effect_warning = ref (not warn_on_unhandled_effect)
303305 ; cps_calls
304306 ; deadcode_sentinal
307+ ; mutated_vars
305308 }
306309end
307310
@@ -1383,9 +1386,73 @@ and translate_instr ctx expr_queue instr =
13831386 mutator_p
13841387 [ J. Expression_statement (J. EBin (J. Eq , Mlvalue.Array. field cx cy, cz)), loc ]
13851388
1386- and translate_instrs ctx expr_queue instr last =
1389+ and translate_instrs ( ctx : Ctx.t ) expr_queue instr last =
13871390 match instr with
13881391 | [] -> [] , expr_queue
1392+ | (Let (_ , Closure _ ), _ ) :: _ -> (
1393+ let names, mut, pcs, all, rem = collect_closures ctx instr in
1394+ match Code.Var.Set. cardinal mut with
1395+ | 0 ->
1396+ let st_rev, expr_queue =
1397+ List. fold_left all ~init: ([] , expr_queue) ~f: (fun (st_rev , expr_queue ) i ->
1398+ let l, expr_queue = translate_instr ctx expr_queue i in
1399+ List. rev_append l st_rev, expr_queue)
1400+ in
1401+ let instrs, expr_queue = translate_instrs ctx expr_queue rem last in
1402+ List. rev_append st_rev instrs, expr_queue
1403+ | _ ->
1404+ let muts =
1405+ Code.Var.Set. diff mut names
1406+ |> Code.Var.Set. elements
1407+ |> List. map ~f: (fun x -> x, Code.Var. fork x)
1408+ in
1409+ (* Rewrite blocks using well-scoped closure variables *)
1410+ let ctx =
1411+ let map =
1412+ List. fold_left muts ~init: Var.Map. empty ~f: (fun acc (x , x' ) ->
1413+ Var.Map. add x x' acc)
1414+ in
1415+ let p, _visited =
1416+ List. fold_left
1417+ pcs
1418+ ~init: (ctx.blocks, Addr.Set. empty)
1419+ ~f: (fun (blocks , visited ) pc ->
1420+ Subst. cont' (Subst. from_map map) pc blocks visited)
1421+ in
1422+ { ctx with blocks = p }
1423+ in
1424+ (* Let bind mutable variables that are part of closures *)
1425+ let let_bindings_rev, expr_queue =
1426+ let expr_queue, st_rev, l_rev =
1427+ List. fold_left
1428+ muts
1429+ ~init: (expr_queue, [] , [] )
1430+ ~f: (fun (expr_queue , st_rev , l_rev ) (v , v' ) ->
1431+ let instrs, ((_px, cx), expr_queue) =
1432+ access_queue_may_flush expr_queue v' v
1433+ in
1434+ let l_rev = (J. V v', (cx, J. N )) :: l_rev in
1435+ expr_queue, List. rev_append instrs st_rev, l_rev)
1436+ in
1437+ (J. variable_declaration ~kind: Let (List. rev l_rev), J. N ) :: st_rev, expr_queue
1438+ in
1439+ (* Mutually recursive need to be properly scoped. *)
1440+ let st_rev, expr_queue =
1441+ List. fold_left all ~init: ([] , expr_queue) ~f: (fun (st_rev , expr_queue ) i ->
1442+ let l, expr_queue = translate_instr ctx expr_queue i in
1443+ let l_rev =
1444+ List. rev_map l ~f: (fun (e , loc' ) ->
1445+ match e with
1446+ (* FIXME: This pattern is too fragile * )
1447+ | J. Variable_statement
1448+ (Var , [ DeclIdent (x, Some (J. EFun (None , dcl), _loc)) ]) ->
1449+ J. Function_declaration (x, dcl), loc'
1450+ | _ -> e, loc')
1451+ in
1452+ List. append l_rev st_rev, expr_queue)
1453+ in
1454+ let instrs, expr_queue = translate_instrs ctx expr_queue rem last in
1455+ List. rev_append let_bindings_rev (List. rev_append st_rev instrs), expr_queue)
13891456 | instr :: rem ->
13901457 let st, expr_queue = translate_instr ctx expr_queue instr in
13911458 let instrs, expr_queue = translate_instrs ctx expr_queue rem last in
@@ -1750,6 +1817,14 @@ and compile_closure ctx (pc, args) =
17501817 if debug () then Format. eprintf " }@]@;" ;
17511818 res
17521819
1820+ and collect_closures ctx l =
1821+ match l with
1822+ | ((Let (x , Closure (_ , (pc , _ ))), _loc ) as i ) :: rem ->
1823+ let names', mut', pcs', i', rem' = collect_closures ctx rem in
1824+ let mut = Code.Addr.Map. find pc ctx.Ctx. mutated_vars in
1825+ Code.Var.Set. add x names', Code.Var.Set. union mut mut', pc :: pcs', i :: i', rem'
1826+ | _ -> Code.Var.Set. empty, Code.Var.Set. empty, [] , [] , l
1827+
17531828let generate_shared_value ctx =
17541829 let strings =
17551830 ( J. variable_declaration
@@ -1808,12 +1883,14 @@ let f
18081883 let exported_runtime =
18091884 if exported_runtime then Some (Code.Var. fresh_n " runtime" , ref false ) else None
18101885 in
1886+ let mutated_vars = Freevars. f p in
18111887 let ctx =
18121888 Ctx. initial
18131889 ~warn_on_unhandled_effect
18141890 ~exported_runtime
18151891 ~should_export
18161892 ~deadcode_sentinal
1893+ ~mutated_vars
18171894 p.blocks
18181895 live_vars
18191896 cps_calls
0 commit comments