[eval] rework capture variable handling (#8017)

src/context/common.ml

@@ -379,6 +379,7 @@ let get_config com =
 			pf_pad_nulls = true;
 			pf_uses_utf16 = false;
 			pf_supports_threads = true;
+			pf_capture_policy = CPWrapRef;
 		}
 
 let memory_marker = [|Unix.time()|]

src/macro/eval/evalContext.ml

@@ -90,7 +90,7 @@ type env = {
 	env_locals : value array;
 	(* The reference to the environment's captured local variables. Indices are determined during compile-time,
 	   or can be obtained through `env_info.capture_infos`. *)
-	env_captures : value ref array;
+	env_captures : value array;
 	(* Map of extra variables added while debugging. Keys are hashed variable names. *)
 	mutable env_extra_locals : value IntMap.t;
 	(* The parent of the current environment, if exists. *)
@@ -394,7 +394,7 @@ let push_environment ctx info num_locals num_captures =
 	let captures = if num_captures = 0 then
 		empty_array
 	else
-		Array.make num_captures (ref vnull)
+		Array.make num_captures vnull
 	in
 	let env = {
 		env_info = info;

src/macro/eval/evalDebugMisc.ml

@@ -124,7 +124,7 @@ let get_variable env capture_infos scopes name env =
 	with Not_found ->
 		let slot = get_capture_slot_by_name capture_infos name in
 		let value = try env.env_captures.(slot) with _ -> raise Not_found in
-		!value
+		value
 
 (* Expr to value *)
 

src/macro/eval/evalDebugSocket.ml

@@ -200,7 +200,7 @@ let output_scopes ctx env =
 
 let output_capture_vars infos env =
 	let vars = Hashtbl.fold (fun slot vi acc ->
-		let value = !(env.env_captures.(slot)) in
+		let value = (env.env_captures.(slot)) in
 		(var_to_json vi.vi_name value (Some vi) env) :: acc
 	) infos [] in
 	JArray vars
@@ -660,7 +660,7 @@ let handler =
 			| CaptureScope(infos,env) ->
 				let value = get_value env in
 				let slot = get_capture_slot_by_name infos name in
-				env.env_captures.(slot) := value;
+				env.env_captures.(slot) <- value;
 				var_to_json "" value None env
 			| DebugScope _ | StackFrame _ ->
 				JNull

src/macro/eval/evalEmitter.ml

@@ -80,7 +80,7 @@ let emit_local_declaration i exec env =
 	vnull
 
 let emit_capture_declaration i exec env =
-	env.env_captures.(i) <- ref (exec env);
+	env.env_captures.(i) <- exec env;
 	vnull
 
 let emit_const v _ = v
@@ -368,7 +368,7 @@ let emit_call exec execs p env =
 
 let emit_local_read i env = env.env_locals.(i)
 
-let emit_capture_read i env = !(env.env_captures.(i))
+let emit_capture_read i env = env.env_captures.(i)
 
 let emit_array_length_read exec p env = vint (as_array p (exec env)).alength
 
@@ -444,7 +444,7 @@ let emit_local_write slot exec env =
 
 let emit_capture_write slot exec env =
 	let v = exec env in
-	env.env_captures.(slot) := v;
+	env.env_captures.(slot) <- v;
 	v
 
 let emit_proto_field_write proto i exec2 env =
@@ -522,10 +522,10 @@ let emit_local_read_write slot exec fop prefix env =
 	if prefix then v else v1
 
 let emit_capture_read_write slot exec fop prefix env =
-	let v1 = !(env.env_captures.(slot)) in
+	let v1 = (env.env_captures.(slot)) in
 	let v2 = exec env in
 	let v = fop v1 v2 in
-	env.env_captures.(slot) := v;
+	env.env_captures.(slot) <- v;
 	if prefix then v else v1
 
 let emit_proto_field_read_write proto i exec2 fop prefix env =
@@ -702,42 +702,78 @@ let emit_neg exec p env = match exec env with
 
 (* Function *)
 
-let handle_capture_arguments exec varaccs env =
-	List.iter (fun (slot,i) ->
-		env.env_captures.(i) <- ref env.env_locals.(slot)
-	) varaccs;
-	exec env
+type env_creation = {
+	ec_info : env_info;
+	ec_num_locals : int;
+	ec_num_captures : int;
+}
 
-let get_normal_env ctx info num_locals num_captures _ =
-	push_environment ctx info num_locals num_captures
+let execute_set_local i env v =
+	env.env_locals.(i) <- v
 
-let get_closure_env ctx info num_locals num_captures refs =
-	let env = push_environment ctx info num_locals num_captures in
-	Array.iteri (fun i vr -> env.env_captures.(i) <- vr) refs;
-	env
+let execute_set_capture i env v =
+	env.env_captures.(i) <- v
+
+let process_arguments fl vl env =
+	let rec loop fl vl = match fl,vl with
+		| f :: fl,v :: vl ->
+			f env v;
+			loop fl vl
+		| f :: fl,[] ->
+			f env vnull;
+			loop fl []
+		| [],[] ->
+			()
+		| _ ->
+			exc_string "Something went wrong"
+	in
+	loop fl vl
+[@@inline]
 
-let execute_set_local env i v =
-	env.env_locals.(i) <- v
+let emit_function_ret ctx eci refs exec fl vl =
+	let env = push_environment ctx eci.ec_info eci.ec_num_locals eci.ec_num_captures in
+	Array.iter (fun (i,vr) -> env.env_captures.(i) <- vr) refs;
+	process_arguments fl vl env;
+	let v = try exec env with Return v -> v in
+	pop_environment ctx env;
+	v
 
-let emit_function_ret ctx get_env refs exec vl =
-	let env = get_env refs in
-	List.iteri (execute_set_local env) vl;
+let create_function_noret ctx eci exec fl vl =
+	let env = push_environment ctx eci.ec_info eci.ec_num_locals eci.ec_num_captures in
+	process_arguments fl vl env;
+	let v = exec env in
+	pop_environment ctx env;
+	v
+
+let create_function ctx eci exec fl vl =
+	let env = push_environment ctx eci.ec_info eci.ec_num_locals eci.ec_num_captures in
+	process_arguments fl vl env;
 	let v = try exec env with Return v -> v in
 	pop_environment ctx env;
 	v
 
-let emit_function_noret ctx get_env refs exec vl =
-	let env = get_env refs in
-	List.iteri (execute_set_local env) vl;
+let create_closure_noret ctx eci refs exec fl vl =
+	let env = push_environment ctx eci.ec_info eci.ec_num_locals eci.ec_num_captures in
+	Array.iter (fun (i,vr) -> env.env_captures.(i) <- vr) refs;
+	process_arguments fl vl env;
 	let v = exec env in
 	pop_environment ctx env;
 	v
 
-let create_function ctx get_env hasret refs exec =
-	if hasret || ctx.debug.support_debugger then (emit_function_ret ctx get_env refs exec)
-	else (emit_function_noret ctx get_env refs exec)
+let create_closure refs ctx eci exec fl vl =
+	let env = push_environment ctx eci.ec_info eci.ec_num_locals eci.ec_num_captures in
+	Array.iter (fun (i,vr) -> env.env_captures.(i) <- vr) refs;
+	process_arguments fl vl env;
+	let v = try exec env with Return v -> v in
+	pop_environment ctx env;
+	v
 
-let emit_closure ctx num_captures get_env hasret exec env =
-	let refs = Array.sub env.env_captures 0 num_captures in
-	let f = create_function ctx get_env hasret refs exec in
+let emit_closure ctx mapping eci hasret exec fl env =
+	let refs = Array.map (fun (i,slot) -> i,emit_capture_read slot env) mapping in
+	let create = match hasret,eci.ec_num_captures with
+		| true,0 -> create_function
+		| false,0 -> create_function_noret
+		| _ -> create_closure refs
+	in
+	let f = create ctx eci exec fl in
 	vstatic_function f

src/macro/eval/evalJit.ml

@@ -62,7 +62,7 @@ open EvalJitContext
 let rec op_assign ctx jit e1 e2 = match e1.eexpr with
 	| TLocal var ->
 		let exec = jit_expr jit false e2 in
-		if var.v_capture then emit_capture_write (get_capture_slot jit var.v_id) exec
+		if var.v_capture then emit_capture_write (get_capture_slot jit var) exec
 		else emit_local_write (get_slot jit var.v_id e1.epos) exec
 	| TField(ef,fa) ->
 		let name = hash (field_name fa) in
@@ -111,7 +111,7 @@ let rec op_assign ctx jit e1 e2 = match e1.eexpr with
 and op_assign_op jit op e1 e2 prefix = match e1.eexpr with
 	| TLocal var ->
 		let exec = jit_expr jit false e2 in
-		if var.v_capture then emit_capture_read_write (get_capture_slot jit var.v_id) exec op prefix
+		if var.v_capture then emit_capture_read_write (get_capture_slot jit var) exec op prefix
 		else emit_local_read_write (get_slot jit var.v_id e1.epos) exec op prefix
 	| TField(ef,fa) ->
 		let name = hash (field_name fa) in
@@ -221,14 +221,20 @@ and jit_expr jit return e =
 		emit_type_expr proto
 	| TFunction tf ->
 		let jit_closure = EvalJitContext.create ctx in
-		jit_closure.captures <- jit.captures;
-		jit_closure.capture_infos <- jit.capture_infos;
 		jit.num_closures <- jit.num_closures + 1;
-		let exec = jit_tfunction jit_closure true e.epos tf in
-		let num_captures = Hashtbl.length jit.captures in
+		let fl,exec = jit_tfunction jit_closure true e.epos tf in
 		let hasret = jit_closure.has_nonfinal_return in
-		let get_env = get_env jit_closure false tf.tf_expr.epos.pfile (EKLocalFunction jit.num_closures) in
-		emit_closure ctx num_captures get_env hasret exec
+		let eci = get_env_creation jit_closure false tf.tf_expr.epos.pfile (EKLocalFunction jit.num_closures) in
+		let captures = Hashtbl.fold (fun vid (i,declared) acc -> (i,vid,declared) :: acc) jit_closure.captures [] in
+		let captures = List.sort (fun (i1,_,_) (i2,_,_) -> Pervasives.compare i1 i2) captures in
+		(* Check if the out-of-scope var is in the outer scope because otherwise we have to promote outwards. *)
+		List.iter (fun var -> ignore(get_capture_slot jit var)) jit_closure.captures_outside_scope;
+		let captures = ExtList.List.filter_map (fun (i,vid,declared) ->
+			if declared then None
+			else Some (i,fst (try Hashtbl.find jit.captures vid with Not_found -> Error.error "Something went wrong" e.epos))
+		) captures in
+		let mapping = Array.of_list captures in
+		emit_closure ctx mapping eci hasret exec fl
 	(* branching *)
 	| TIf(e1,e2,eo) ->
 		let exec_cond = jit_expr jit false e1 in
@@ -511,7 +517,7 @@ and jit_expr jit return e =
 		end
 	(* read *)
 	| TLocal var ->
-		if var.v_capture then emit_capture_read (get_capture_slot jit var.v_id)
+		if var.v_capture then emit_capture_read (get_capture_slot jit var)
 		else emit_local_read (get_slot jit var.v_id e.epos)
 	| TField(e1,fa) ->
 		let name = hash (field_name fa) in
@@ -651,10 +657,13 @@ and jit_tfunction jit static pos tf =
 	push_scope jit pos;
 	(* Declare `this` (if not static) and function arguments as local variables. *)
 	if not static then ignore(declare_local_this jit);
-	let varaccs = ExtList.List.filter_map (fun (var,_) ->
-		let slot = add_local jit var in
-		if var.v_capture then Some (slot,add_capture jit var) else None
+	let fl = List.map (fun (var,_) ->
+		let varacc = declare_local jit var in
+		match varacc with
+		| Env slot -> execute_set_capture slot
+		| Local slot -> execute_set_local slot
 	) tf.tf_args in
+	let fl = if static then fl else (execute_set_local 0) :: fl in
 	(* Add conditionals for default values. *)
 	let e = List.fold_left (fun e (v,cto) -> match cto with
 		| None -> e
@@ -672,33 +681,25 @@ and jit_tfunction jit static pos tf =
 	in
 	(* Jit the function expression. *)
 	let exec = jit_expr jit true e in
-	(* Deal with captured arguments, if necessary. *)
-	let exec = match varaccs with
-		| [] -> exec
-		| _ -> handle_capture_arguments exec varaccs
-	in
 	pop_scope jit;
-	exec
+	fl,exec
 
-and get_env jit static file info =
-	let ctx = jit.ctx in
-	let num_locals = jit.max_num_locals in
-	let num_captures = Hashtbl.length jit.captures in
-	let info = create_env_info static file info jit.capture_infos in
-	match info.kind with
-	| EKLocalFunction _ -> get_closure_env ctx info num_locals num_captures
-	| _ -> get_normal_env ctx info num_locals num_captures
+and get_env_creation jit static file info = {
+	ec_info = create_env_info static file info jit.capture_infos;
+	ec_num_locals = jit.max_num_locals;
+	ec_num_captures = Hashtbl.length jit.captures;
+}
 
 (* Creates a [EvalValue.vfunc] of function [tf], which can be [static] or not. *)
 let jit_tfunction ctx key_type key_field tf static pos =
 	let t = Timer.timer [(if ctx.is_macro then "macro" else "interp");"jit"] in
 	(* Create a new JitContext with an initial scope *)
 	let jit = EvalJitContext.create ctx in
-	let exec = jit_tfunction jit static pos tf in
+	let fl,exec = jit_tfunction jit static pos tf in
 	(* Create the [vfunc] instance depending on the number of arguments. *)
 	let hasret = jit.has_nonfinal_return in
-	let get_env = get_env jit static tf.tf_expr.epos.pfile (EKMethod(key_type,key_field)) in
-	let f = create_function ctx get_env hasret empty_array exec in
+	let eci = get_env_creation jit static tf.tf_expr.epos.pfile (EKMethod(key_type,key_field)) in
+	let f = if hasret then create_function ctx eci exec fl else create_function_noret ctx eci exec fl in
 	t();
 	f
 

src/macro/eval/evalJitContext.ml

@@ -17,7 +17,7 @@ type t = {
 	(* The scope stack. *)
 	mutable scopes : scope list;
 	(* The captured variables declared in this context. Maps variable IDs to capture slots. *)
-	mutable captures : (int,int) Hashtbl.t;
+	mutable captures : (int,int * bool) Hashtbl.t;
 	(* The current number of locals. *)
 	mutable num_locals : int;
 	(* The maximum number of locals. *)
@@ -28,6 +28,8 @@ type t = {
 	mutable has_nonfinal_return : bool;
 	(* The name of capture variables. Maps local slots to variable names. Only filled in debug mode. *)
 	mutable capture_infos : (int,var_info) Hashtbl.t;
+	(* Variables which are accessed but not declared in this scope. *)
+	mutable captures_outside_scope : tvar list;
 }
 
 let var_info_of_var var =
@@ -47,6 +49,7 @@ let create ctx = {
 	num_closures = 0;
 	has_nonfinal_return = false;
 	capture_infos = Hashtbl.create 0;
+	captures_outside_scope = []
 }
 
 (* Returns the number of locals in [scope]. *)
@@ -78,9 +81,9 @@ let increase_num_locals jit =
 	if jit.num_locals > jit.max_num_locals then jit.max_num_locals <- jit.num_locals
 
 (* Adds capture variable [var] to context [jit]. *)
-let add_capture jit var =
+let add_capture jit var declared =
 	let i = Hashtbl.length jit.captures in
-	Hashtbl.add jit.captures var.v_id i;
+	Hashtbl.add jit.captures var.v_id (i,declared);
 	if jit.ctx.debug.support_debugger then begin
 		Hashtbl.replace jit.capture_infos i (var_info_of_var var)
 	end;
@@ -108,7 +111,7 @@ let add_local jit var = match jit.scopes with
 	Returns either [Env slot] if the variable is captured or [Local slot] otherwise.
 *)
 let declare_local jit var =
-	if var.v_capture then Env (add_capture jit var)
+	if var.v_capture then Env (add_capture jit var true)
 	else Local (add_local jit var)
 
 (*
@@ -119,7 +122,7 @@ let declare_local jit var =
 *)
 let declare_arg jit var =
 	let varacc = add_local jit var in
-	if var.v_capture then add_capture jit var,Some varacc else varacc,None
+	if var.v_capture then add_capture jit var true,Some varacc else varacc,None
 
 (* Declares a variable for `this` in context [jit]. *)
 let declare_local_this jit = match jit.scopes with
@@ -151,5 +154,8 @@ let get_slot jit vid p =
 	with Not_found -> EvalMisc.throw_string "Unbound variable" p
 
 (* Gets the slot of captured variable id [vid] in context [jit]. *)
-let get_capture_slot jit vid =
-	Hashtbl.find jit.captures vid
+let get_capture_slot jit var =
+	try fst (Hashtbl.find jit.captures var.v_id)
+	with Not_found ->
+		jit.captures_outside_scope <- var :: jit.captures_outside_scope;
+		add_capture jit var false