perf: reuse environment entries across process spawns

bench/micro/dune_bench/env_bench.ml

@@ -0,0 +1,33 @@
+open Stdune
+
+let make_env size ~temp_dir =
+  let env =
+    List.init size ~f:(fun i ->
+      let prefix = if i mod 2 = 0 then 'A' else 'Z' in
+      Printf.sprintf "%c_DUNE_BENCH_%04d=value-%04d" prefix i i)
+    |> Array.of_list
+    |> Env.of_unix
+  in
+  if temp_dir
+  then Env.add env ~var:Env.Var.temp_dir ~value:"original-temp-directory"
+  else env
+;;
+
+let render_environment_override ~temp_dir variables =
+  let env = make_env variables ~temp_dir in
+  fun () ->
+    Env.to_unix_with_override env ~var:Env.Var.temp_dir ~value:"dune-temp-directory"
+    |> ignore
+;;
+
+let%bench_fun
+    ("render environment override: present" [@indexed variables = [ 10; 100; 1_000 ]])
+  =
+  render_environment_override ~temp_dir:true variables
+;;
+
+let%bench_fun
+    ("render environment override: absent" [@indexed variables = [ 10; 100; 1_000 ]])
+  =
+  render_environment_override ~temp_dir:false variables
+;;

otherlibs/stdune/src/env.ml

@@ -52,6 +52,15 @@ let binding var value =
   Bytes.unsafe_to_string result
 ;;
 
+let unix_entry_has_var_caseless entry var =
+  let var_length = String.length var in
+  let rec loop i =
+    i = var_length
+    || (Char.lowercase_ascii entry.[i] = Char.lowercase_ascii var.[i] && loop (i + 1))
+  in
+  String.length entry > var_length && entry.[var_length] = '=' && loop 0
+;;
+
 let to_unix t =
   match t.unix with
   | Some v -> v
@@ -61,6 +70,48 @@ let to_unix t =
     res
 ;;
 
+let to_unix_with_override t ~var ~value =
+  match get t var with
+  | None ->
+    if Sys.win32
+    then of_map (Map.set t.vars var value) |> to_unix
+    else binding var value :: to_unix t
+  | Some old_value when String.equal old_value value -> to_unix t
+  | Some old_value ->
+    let new_entry = binding var value in
+    let replace =
+      if Sys.win32
+      then (
+        let rec loop = function
+          | [] -> None
+          | entry :: rest ->
+            if unix_entry_has_var_caseless entry var
+            then Some (new_entry :: rest)
+            else (
+              match loop rest with
+              | None -> None
+              | Some rest -> Some (entry :: rest))
+        in
+        loop)
+      else (
+        let old_entry = binding var old_value in
+        let rec loop = function
+          | [] -> None
+          | entry :: rest ->
+            if String.equal entry old_entry
+            then Some (new_entry :: rest)
+            else (
+              match loop rest with
+              | None -> None
+              | Some rest -> Some (entry :: rest))
+        in
+        loop)
+    in
+    (match replace (to_unix t) with
+     | Some result -> result
+     | None -> of_map (Map.set t.vars var value) |> to_unix)
+;;
+
 let of_unix arr =
   Array.to_list arr
   |> List.map ~f:(fun s ->

otherlibs/stdune/src/env.mli

@@ -27,6 +27,11 @@ val vars : t -> Var.Set.t
 val initial : t
 
 val to_unix : t -> string list
+
+(** [to_unix_with_override t ~var ~value] serializes [t] with [var] set to
+    [value]. It reuses the cached serialization of variables that are unchanged. *)
+val to_unix_with_override : t -> var:Var.t -> value:string -> string list
+
 val of_unix : string array -> t
 val get : t -> Var.t -> string option
 

otherlibs/stdune/test/dune

@@ -3,6 +3,7 @@
  (modules
   :standard
   \
+  env_tests
   path_external_build_tests
   path_tests
   proc_linux_tests
@@ -24,6 +25,14 @@
  (preprocess
   (pps ppx_expect)))
 
+(library
+ (name stdune_env_tests)
+ (modules env_tests)
+ (inline_tests)
+ (libraries stdune ppx_inline_test.config)
+ (preprocess
+  (pps ppx_inline_test)))
+
 (library
  (name stdune_path_tests)
  (modules path_tests)
@@ -88,4 +97,5 @@
 (alias
  (name runtest-windows)
  (deps
+  (alias runtest-stdune_env_tests)
   (alias runtest-stdune_external_build_tests)))

otherlibs/stdune/test/env_tests.ml

@@ -0,0 +1,47 @@
+open Stdune
+
+let test_override env ~var ~value =
+  let expected = Env.add env ~var ~value in
+  let actual_unix = Env.to_unix_with_override env ~var ~value in
+  let actual = actual_unix |> Array.of_list |> Env.of_unix in
+  if Sys.win32
+  then (
+    let expected_unix = Env.to_unix expected in
+    if not (List.equal String.equal expected_unix actual_unix)
+    then
+      Code_error.raise
+        "Env.to_unix_with_override changed the serialized environment"
+        [ "expected", Dyn.list Dyn.string expected_unix
+        ; "actual", Dyn.list Dyn.string actual_unix
+        ];
+    let matching_entries =
+      List.filter actual_unix ~f:(fun entry ->
+        match String.lsplit2 entry ~on:'=' with
+        | None -> false
+        | Some (entry_var, _) -> Ordering.is_eq (Env.Var.compare entry_var var))
+    in
+    if List.length matching_entries <> 1
+    then
+      Code_error.raise
+        "Env.to_unix_with_override returned duplicate variables"
+        [ "var", Dyn.string var; "environment", Dyn.list Dyn.string actual_unix ]);
+  if not (Env.equal expected actual)
+  then
+    Code_error.raise
+      "Env.to_unix_with_override returned the wrong environment"
+      [ "expected", Env.to_dyn expected; "actual", Env.to_dyn actual ]
+;;
+
+let%test_unit "to_unix_with_override" =
+  let env = Env.of_unix [| "A=one"; "B=two"; "C=three" |] in
+  test_override env ~var:"A" ~value:"one";
+  test_override env ~var:"B" ~value:"changed";
+  test_override env ~var:"D" ~value:"new";
+  test_override env ~var:"C" ~value:"value=with=equals";
+  if Sys.win32
+  then (
+    let mixed_case = Env.of_unix [| "A=one"; "Temp=first" |] in
+    test_override mixed_case ~var:"TEMP" ~value:"changed";
+    let duplicate_case = Env.of_unix [| "A=one"; "Temp=first"; "TEMP=second" |] in
+    test_override duplicate_case ~var:"temp" ~value:"changed")
+;;

src/dune_engine/dtemp.ml

@@ -7,9 +7,9 @@ let file ~prefix ~suffix =
   Temp.temp_in_dir File ~dir:(Lazy.force temp_dir) ~suffix ~prefix
 ;;
 
-let add_to_env env =
+let to_unix env =
   let value = Lazy.force temp_dir_value in
-  Env.add env ~var:Env.Var.temp_dir ~value
+  Env.to_unix_with_override env ~var:Env.Var.temp_dir ~value
 ;;
 
 let destroy = Temp.destroy

src/dune_engine/dtemp.mli

@@ -5,9 +5,8 @@ open Import
 (** This returns a build path, but we don't rely on that *)
 val file : prefix:string -> suffix:string -> Path.t
 
-(** Add the temp env var to the environment passed or return the initial
-    environment with the temp var added. *)
-val add_to_env : Env.t -> Env.t
+(** Serialize an environment with Dune's temp directory set. *)
+val to_unix : Env.t -> string list
 
 (** Destroy the temporary file or directory *)
 val destroy : Temp.what -> Path.t -> unit

src/dune_engine/process.ml

@@ -982,10 +982,7 @@ let spawn
     Time.now ()
   in
   let pid =
-    let env =
-      let env = Dtemp.add_to_env env in
-      Env.to_unix env |> Spawn.Env.of_list
-    in
+    let env = Dtemp.to_unix env |> Spawn.Env.of_list in
     let stdout = Io.fd stdout |> Fd.unsafe_to_unix_file_descr in
     let stderr = Io.fd stderr |> Fd.unsafe_to_unix_file_descr in
     let stdin = Io.fd stdin |> Fd.unsafe_to_unix_file_descr in