flambda-backend: Support DLS API (single-domain only) (#1978)

Co-authored-by: Mark Shinwell <mshinwell@pm.me>

stdlib/domain.ml

@@ -16,6 +16,57 @@
 (*                                                                        *)
 (**************************************************************************)
 
+(* CR ocaml 5 runtime: domain-local-storage assumes single-domain,
+   i.e. calling split will never be necessary. *)
+
+module DLS = struct
+
+  let unique_value = Obj.repr (ref 0)
+  let state = ref (Array.make 8 unique_value)
+
+  type 'a key = int * (unit -> 'a)
+
+  let key_counter = ref 0
+
+  let new_key ?split_from_parent:_ init_orphan =
+    let idx = !key_counter in
+    key_counter := idx + 1;
+    (idx, init_orphan)
+
+  (* If necessary, grow the current domain's local state array such that [idx]
+   * is a valid index in the array. *)
+  let maybe_grow idx =
+    let st = !state in
+    let sz = Array.length st in
+    if idx < sz then st
+    else begin
+      let rec compute_new_size s =
+        if idx < s then s else compute_new_size (2 * s)
+      in
+      let new_sz = compute_new_size sz in
+      let new_st = Array.make new_sz unique_value in
+      Array.blit st 0 new_st 0 sz;
+      state := new_st;
+      new_st
+    end
+
+  let set (idx, _init) x =
+    let st = maybe_grow idx in
+    (* [Sys.opaque_identity] ensures that flambda does not look at the type of
+     * [x], which may be a [float] and conclude that the [st] is a float array.
+     * We do not want OCaml's float array optimisation kicking in here. *)
+    st.(idx) <- Obj.repr (Sys.opaque_identity x)
+
+  let get (idx, init) =
+    let st = maybe_grow idx in
+    let v = st.(idx) in
+    if v == unique_value then
+      let v' = Obj.repr (init ()) in
+      st.(idx) <- (Sys.opaque_identity v');
+      Obj.magic v'
+    else Obj.magic v
+end
+
 (* CR ocaml 5 runtime: domains not supported on 4.x
 
 module Raw = struct

stdlib/domain.mli

@@ -92,6 +92,7 @@ val recommended_domain_count : unit -> int
     simultaneously (including domains already running).
 
     The value returned is at least [1]. *)
+*)
 
 module DLS : sig
 (** Domain-local Storage *)
@@ -146,5 +147,3 @@ module DLS : sig
         to [k], which cannot be restored later. *)
 
 end
-
-*)

stdlib/dune

@@ -110,6 +110,8 @@
   complex.mli
   digest.ml
   digest.mli
+  domain.ml
+  domain.mli
   either.ml
   either.mli
   ephemeron.ml
@@ -245,6 +247,9 @@
   .stdlib.objs/byte/stdlib__Digest.cmi
   .stdlib.objs/byte/stdlib__Digest.cmt
   .stdlib.objs/byte/stdlib__Digest.cmti
+  .stdlib.objs/byte/stdlib__Domain.cmi
+  .stdlib.objs/byte/stdlib__Domain.cmt
+  .stdlib.objs/byte/stdlib__Domain.cmti
   .stdlib.objs/byte/stdlib__Either.cmi
   .stdlib.objs/byte/stdlib__Either.cmt
   .stdlib.objs/byte/stdlib__Either.cmti
@@ -466,6 +471,7 @@
   .stdlib.objs/native/stdlib__Marshal.cmx
   .stdlib.objs/native/stdlib__BytesLabels.cmx
   .stdlib.objs/native/stdlib__Digest.cmx
+  .stdlib.objs/native/stdlib__Domain.cmx
   .stdlib.objs/native/stdlib__Atomic.cmx
   .stdlib.objs/native/stdlib__Effect.cmx
   .stdlib.objs/native/stdlib__Either.cmx

stdlib/stdlib.ml

@@ -629,9 +629,9 @@ module Complex        = Complex
 module Condition      = Condition
 *)
 module Digest         = Digest
+module Domain         = Domain
 (* CR ocaml 5 runtime:
    BACKPORT
-module Domain         = Domain
 module Effect         = Effect
 *)
 module Either         = Either

stdlib/stdlib.mli

@@ -1408,13 +1408,13 @@ module Complex        = Complex
 module Condition      = Condition
 *)
 module Digest         = Digest
-(* CR ocaml 5 runtime:
-BACKPORT
 module Domain         = Domain
 [@@alert "-unstable"]
 [@@alert unstable
     "The Domain interface may change in incompatible ways in the future."
 ]
+(* CR ocaml 5 runtime:
+BACKPORT
 module Effect         = Effect
 [@@alert "-unstable"]
 [@@alert unstable