add --wait stdio (#5188)

* add --wait stdio

* fix for old ocaml (see #5174)

src/main.ml

@@ -683,17 +683,9 @@ let rec process_params create pl =
 	) in
 	loop [] pl
 
-and wait_loop boot_com host port =
-	let sock = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
-	(try Unix.setsockopt sock Unix.SO_REUSEADDR true with _ -> ());
-	(try Unix.bind sock (Unix.ADDR_INET (Unix.inet_addr_of_string host,port)) with _ -> failwith ("Couldn't wait on " ^ host ^ ":" ^ string_of_int port));
-	Unix.listen sock 10;
+and wait_loop verbose accept =
 	Sys.catch_break false;
-	let verbose = boot_com.verbose in
 	let has_parse_error = ref false in
-	if verbose then print_endline ("Waiting on " ^ host ^ ":" ^ string_of_int port);
-	let bufsize = 1024 in
-	let tmp = String.create bufsize in
 	let cache = {
 		c_haxelib = Hashtbl.create 0;
 		c_files = Hashtbl.create 0;
@@ -840,33 +832,8 @@ and wait_loop boot_com host port =
 	);
 	let run_count = ref 0 in
 	while true do
-		let sin, _ = Unix.accept sock in
+		let read, write, close = accept() in
 		let t0 = get_time() in
-		Unix.set_nonblock sin;
-		if verbose then print_endline "Client connected";
-		let b = Buffer.create 0 in
-		let rec read_loop count =
-			try
-				let r = Unix.recv sin tmp 0 bufsize [] in
-				if r = 0 then
-					failwith "Incomplete request"
-				else begin
-					if verbose then Printf.printf "Reading %d bytes\n" r;
-					Buffer.add_substring b tmp 0 r;
-					if tmp.[r-1] = '\000' then
-						Buffer.sub b 0 (Buffer.length b - 1)
-					else
-						read_loop 0
-				end
-			with Unix.Unix_error((Unix.EWOULDBLOCK|Unix.EAGAIN),_,_) ->
-				if count = 100 then
-					failwith "Aborting inactive connection"
-				else begin
-					if verbose then print_endline "Waiting for data...";
-					ignore(Unix.select [] [] [] 0.05); (* wait a bit *)
-					read_loop (count + 1);
-				end
-		in
 		let rec cache_context com =
 			if com.display = DMNone then begin
 				List.iter cache_module com.modules;
@@ -881,8 +848,8 @@ and wait_loop boot_com host port =
 			ctx.flush <- (fun() ->
 				incr compilation_step;
 				compilation_mark := !mark_loop;
-				List.iter (fun s -> ssend sin (s ^ "\n"); if verbose then print_endline ("> " ^ s)) (List.rev ctx.messages);
-				if ctx.has_error then ssend sin "\x02\n" else cache_context ctx.com;
+				List.iter (fun s -> write (s ^ "\n"); if verbose then print_endline ("> " ^ s)) (List.rev ctx.messages);
+				if ctx.has_error then write "\x02\n" else cache_context ctx.com;
 			);
 			ctx.setup <- (fun() ->
 				if verbose then begin
@@ -900,11 +867,11 @@ and wait_loop boot_com host port =
 					Hashtbl.iter (fun _ m -> if m.m_extra.m_file = file then m.m_extra.m_dirty <- true) cache.c_modules
 				end
 			);
-			ctx.com.print <- (fun str -> ssend sin ("\x01" ^ String.concat "\x01" (ExtString.String.nsplit str "\n") ^ "\n"));
+			ctx.com.print <- (fun str -> write ("\x01" ^ String.concat "\x01" (ExtString.String.nsplit str "\n") ^ "\n"));
 			ctx
 		in
 		(try
-			let s = (read_loop 0) in
+			let s = read() in
 			let hxml =
 				try
 					let idx = String.index s '\001' in
@@ -914,7 +881,6 @@ and wait_loop boot_com host port =
 					s
 			in
 			let data = parse_hxml_data hxml in
-			Unix.clear_nonblock sin;
 			if verbose then print_endline ("Processing Arguments [" ^ String.concat "," data ^ "]");
 			(try
 				Common.display_default := DMNone;
@@ -933,11 +899,11 @@ and wait_loop boot_com host port =
 				start_time := get_time();
 				process_params create data;
 				close_times();
-				if !measure_times then report_times (fun s -> ssend sin (s ^ "\n"))
+				if !measure_times then report_times (fun s -> write (s ^ "\n"))
 			with
 			| Completion str ->
 				if verbose then print_endline ("Completion Response =\n" ^ str);
-				ssend sin str
+				write str
 			| Arg.Bad msg ->
 				prerr_endline ("Error: " ^ msg);
 			);
@@ -950,10 +916,10 @@ and wait_loop boot_com host port =
 		| e ->
 			let estr = Printexc.to_string e in
 			if verbose then print_endline ("Uncaught Error : " ^ estr);
-			(try ssend sin estr with _ -> ());
+			(try write estr with _ -> ());
 			if is_debug_run() then print_endline (Printexc.get_backtrace());
 		);
-		Unix.close sin;
+		close();
 		current_stdin := None;
 		(* prevent too much fragmentation by doing some compactions every X run *)
 		incr run_count;
@@ -968,6 +934,70 @@ and wait_loop boot_com host port =
 		end else Gc.minor();
 	done
 
+and init_wait_stdio() =
+	set_binary_mode_in stdin true;
+	set_binary_mode_out stderr true;
+
+	let chin = IO.input_channel stdin in
+	let cherr = IO.output_channel stderr in
+
+	let berr = Buffer.create 0 in
+	let read = fun () ->
+		let len = IO.read_i32 chin in
+		IO.really_nread chin len
+	in
+	let write = Buffer.add_string berr in
+	let close = fun() ->
+		IO.write_i32 cherr (Buffer.length berr);
+		IO.nwrite cherr (Buffer.contents berr);
+		IO.flush cherr
+	in
+	fun() ->
+		Buffer.clear berr;
+		read, write, close
+
+and init_wait_socket verbose host port =
+	let sock = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
+	(try Unix.setsockopt sock Unix.SO_REUSEADDR true with _ -> ());
+	(try Unix.bind sock (Unix.ADDR_INET (Unix.inet_addr_of_string host,port)) with _ -> failwith ("Couldn't wait on " ^ host ^ ":" ^ string_of_int port));
+	if verbose then print_endline ("Waiting on " ^ host ^ ":" ^ string_of_int port);
+	Unix.listen sock 10;
+	let bufsize = 1024 in
+	let tmp = String.create bufsize in
+	let accept() = (
+		let sin, _ = Unix.accept sock in
+		Unix.set_nonblock sin;
+		if verbose then print_endline "Client connected";
+		let b = Buffer.create 0 in
+		let rec read_loop count =
+			try
+				let r = Unix.recv sin tmp 0 bufsize [] in
+				if r = 0 then
+					failwith "Incomplete request"
+				else begin
+					if verbose then Printf.printf "Reading %d bytes\n" r;
+					Buffer.add_substring b tmp 0 r;
+					if tmp.[r-1] = '\000' then
+						Buffer.sub b 0 (Buffer.length b - 1)
+					else
+						read_loop 0
+				end
+			with Unix.Unix_error((Unix.EWOULDBLOCK|Unix.EAGAIN),_,_) ->
+				if count = 100 then
+					failwith "Aborting inactive connection"
+				else begin
+					if verbose then print_endline "Waiting for data...";
+					ignore(Unix.select [] [] [] 0.05); (* wait a bit *)
+					read_loop (count + 1);
+				end
+		in
+		let read = fun() -> (let s = read_loop 0 in Unix.clear_nonblock sin; s) in
+		let write = ssend sin in
+		let close() = Unix.close sin in
+		read, write, close
+	) in
+	accept
+
 and do_connect host port args =
 	let sock = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
 	(try Unix.connect sock (Unix.ADDR_INET (Unix.inet_addr_of_string host,port)) with _ -> failwith ("Couldn't connect on " ^ host ^ ":" ^ string_of_int port));
@@ -1338,8 +1368,15 @@ try
 			evals := s :: !evals;
 		), " : evaluates argument as Haxe module code");
 		("--wait", Arg.String (fun hp ->
-			let host, port = (try ExtString.String.split hp ":" with _ -> "127.0.0.1", hp) in
-			wait_loop com host (try int_of_string port with _ -> raise (Arg.Bad "Invalid port"))
+			let accept = match hp with
+				| "stdio" ->
+					init_wait_stdio()
+				| _ ->
+					let host, port = (try ExtString.String.split hp ":" with _ -> "127.0.0.1", hp) in
+					let port = try int_of_string port with _ -> raise (Arg.Bad "Invalid port") in
+					init_wait_socket com.verbose host port
+			in
+			wait_loop com.verbose accept
 		),"<[host:]port> : wait on the given port for commands to run)");
 		("--connect",Arg.String (fun _ ->
 			assert false