Source file iarray.ml

(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*             Xavier Leroy, projet Cristal, INRIA Rocquencourt           *)
(*               Antal Spector-Zabusky, Jane Street, New York             *)
(*                                                                        *)
(*   Copyright 1996 Institut National de Recherche en Informatique et     *)
(*     en Automatique.                                                    *)
(*   Copyright 2023 Jane Street Group LLC                                 *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

open! Stdlib

(* NOTE: If you update this file, please also update iarrayLabels.ml; from the
   declaration of [type +'a t = 'a iarray] on down, they're the same.  This is a
   temporary state of affairs, but for now, please copy things! *)

(* In this file, we use four different implementation strategies:

     1. Reusing [external]s for mutable arrays.  (E.g., [get].)
     2. Copying implementations from [array.ml], which in this new context read
        from immutable arrays.  (E.g., [iter].)
     3. As (2), but they construct a mutable array, which we unsafely
        reinterpret as an immutable array at the very end (E.g., [map].)
     4. (Only for sorting.) Copying the immutable array and passing it to a
        mutating function.  (E.g., [sort].)

   The first and third strategies are safe because mutable and immutable arrays
   have the same runtime representation, and we only apply them to functions
   that satisfy the following three properties:

     A. They do not mutate their array inputs;
     B. They do not hold on to their array inputs; and
     C. They return a fresh array if they return an array.

   We do not expose other primitives (e.g., [unsafe_set]) or functions (e.g.,
   [fill]).

   We choose between the four strategies as follows:

     1. We use [external]s if there is a corresponding [external].
     2. Functions that only read arrays can have their implementations safely
        copied.
     3. Functions that create an immutable array have to do so by creating a
        mutable array and freezing it, since immutable arrays are, well,
        immutable.  For efficiency, we freeze them unsafely.
     4. Sorting is implemented in-place and this allows to to reuse it.
*)

[@@@ocaml.flambda_o3]

(* An alias for the type of immutable arrays. *)
type (+'a : any mod separable) t = 'a iarray

(* Array operations *)

external length : ('a : any mod separable). local_ 'a iarray -> int @@ portable
  = "%array_length"
[@@layout_poly]
external get :
  ('a : any mod separable). ('a iarray[@local_opt]) -> int -> ('a[@local_opt])
  @@ portable = "%array_safe_get"
[@@layout_poly]
external unsafe_get :
  ('a : any mod separable). ('a iarray[@local_opt]) -> int -> ('a[@local_opt])
  @@ portable = "%array_unsafe_get"
[@@layout_poly]
external concat : ('a : any mod separable). 'a iarray list -> 'a iarray
  @@ portable = "caml_array_concat"

external append_prim :
  ('a : any mod separable). 'a iarray -> 'a iarray -> 'a iarray @@ portable
  = "caml_array_append"
external unsafe_sub :
  ('a : any mod separable). 'a iarray -> int -> int -> 'a iarray @@ portable
  = "caml_array_sub"
external unsafe_of_array : ('a : any mod separable). 'a array -> 'a iarray
  @@ portable = "%array_to_iarray"
external unsafe_to_array : ('a : any mod separable). 'a iarray -> 'a array
  @@ portable = "%array_of_iarray"

(* Used only to reimplement [init] *)
external unsafe_set_mutable :
  ('a : any mod separable). 'a array -> int -> 'a -> unit @@ portable
  = "%array_unsafe_set"
[@@layout_poly]

(* We can't use immutable array literals in this file, since we don't want to
   require the stdlib to be compiled with extensions, so instead of [[::]] we
   use [unsafe_of_(local_)array [||]] below. *)

(* The implementation is copied from [Array] so that [f] can be [local_] *)
let init l (local_ f) =
  if l = 0 then unsafe_of_array [||] else
  if l < 0 then invalid_arg "Iarray.init"
  (* See #6575. We could also check for maximum array size, but this depends
     on whether we create a float array or a regular one... *)
  else
   let res = Array.make l (f 0) in
   for i = 1 to pred l do
     unsafe_set_mutable res i (f i)
   done;
   unsafe_of_array res

let append a1 a2 =
  if length a1 = 0 then a2 (* Safe because they're immutable *)
  else if length a2 = 0 then a1
  else append_prim a1 a2

let sub a ~pos ~len =
  if pos < 0 || len < 0 || pos > length a - len
  then invalid_arg "Iarray.sub"
  else unsafe_sub a pos len

let iter f a =
  for i = 0 to length a - 1 do f(unsafe_get a i) done

let iter2 f a b =
  if length a <> length b then
    invalid_arg "Iarray.iter2: arrays must have the same length"
  else
    for i = 0 to length a - 1 do f (unsafe_get a i) (unsafe_get b i) done

let map f a =
  let l = length a in
  if l = 0 then unsafe_of_array [||] else begin
    let r = Array.make l (f(unsafe_get a 0)) in
    for i = 1 to l - 1 do
      Array.unsafe_set r i (f(unsafe_get a i))
    done;
    unsafe_of_array r
  end

let map2 f a b =
  let la = length a in
  let lb = length b in
  if la <> lb then
    invalid_arg "Iarray.map2: arrays must have the same length"
  else begin
    if la = 0 then unsafe_of_array [||] else begin
      let r = Array.make la (f (unsafe_get a 0) (unsafe_get b 0)) in
      for i = 1 to la - 1 do
        Array.unsafe_set r i (f (unsafe_get a i) (unsafe_get b i))
      done;
      unsafe_of_array r
    end
  end

let iteri f a =
  for i = 0 to length a - 1 do f i (unsafe_get a i) done

let mapi f a =
  let l = length a in
  if l = 0 then unsafe_of_array [||] else begin
    let r = Array.make l (f 0 (unsafe_get a 0)) in
    for i = 1 to l - 1 do
      Array.unsafe_set r i (f i (unsafe_get a i))
    done;
    unsafe_of_array r
  end

let to_list a =
  let rec tolist i res =
    if i < 0 then res else tolist (i - 1) (unsafe_get a i :: res) in
  tolist (length a - 1) []

let of_list l = unsafe_of_array (Array.of_list l)

(* This shouldn't violate the forward-pointers restriction because the list
   elements already exist *)
let to_array ia = Array.copy (unsafe_to_array ia)

let of_array ma = unsafe_of_array (Array.copy ma)

let fold_left f x a =
  let r = ref x in
  for i = 0 to length a - 1 do
    r := f !r (unsafe_get a i)
  done;
  !r

let fold_left_map f acc input_array =
  let len = length input_array in
  if len = 0 then (acc, unsafe_of_array [||]) else begin
    let acc, elt = f acc (unsafe_get input_array 0) in
    let output_array = Array.make len elt in
    let acc = ref acc in
    for i = 1 to len - 1 do
      let acc', elt = f !acc (unsafe_get input_array i) in
      acc := acc';
      Array.unsafe_set output_array i elt;
    done;
    !acc, unsafe_of_array output_array
  end

let fold_right f a x =
  let r = ref x in
  for i = length a - 1 downto 0 do
    r := f (unsafe_get a i) !r
  done;
  !r

(* CR aspectorzabusky: Why do I need this?  Shouldn't mode-crossing handle doing
   this? *)
let[@inline always] globalize_bool : local_ bool -> bool = fun b -> b

let exists p a =
  let n = length a in
  let rec loop i = exclave_
    if i = n then false
    else if p (unsafe_get a i) then true
    else loop (succ i) in
  globalize_bool (loop 0)

let for_all p a =
  let n = length a in
  let rec loop i = exclave_
    if i = n then true
    else if p (unsafe_get a i) then loop (succ i)
    else false in
  globalize_bool (loop 0)

let for_all2 p l1 l2 =
  let n1 = length l1
  and n2 = length l2 in
  if n1 <> n2 then invalid_arg "Iarray.for_all2"
  else let rec loop i = exclave_
    if i = n1 then true
    else if p (unsafe_get l1 i) (unsafe_get l2 i) then loop (succ i)
    else false in
  globalize_bool (loop 0)

let exists2 p l1 l2 =
  let n1 = length l1
  and n2 = length l2 in
  if n1 <> n2 then invalid_arg "Iarray.exists2"
  else let rec loop i = exclave_
    if i = n1 then false
    else if p (unsafe_get l1 i) (unsafe_get l2 i) then true
    else loop (succ i) in
  globalize_bool (loop 0)

let mem x a =
  let n = length a in
  let rec loop i = exclave_
    if i = n then false
    else if compare (unsafe_get a i) x = 0 then true
    else loop (succ i) in
  globalize_bool (loop 0)

let memq x a =
  let n = length a in
  let rec loop i = exclave_
    if i = n then false
    else if x == (unsafe_get a i) then true
    else loop (succ i) in
  globalize_bool (loop 0)

let find_opt p a =
  let n = length a in
  let rec loop i =
    if i = n then None
    else
      let x = unsafe_get a i in
      if p x then Some x
      else loop (succ i)
  in
  loop 0 [@nontail]

let find_map f a =
  let n = length a in
  let rec loop i =
    if i = n then None
    else
      match f (unsafe_get a i) with
      | None -> loop (succ i)
      | Some _ as r -> r
  in
  loop 0 [@nontail]

let split x =
  if x = unsafe_of_array [||]
  then unsafe_of_array [||], unsafe_of_array [||]
  else begin
    let a0, b0 = unsafe_get x 0 in
    let n = length x in
    let a = Array.make n a0 in
    let b = Array.make n b0 in
    for i = 1 to n - 1 do
      let ai, bi = unsafe_get x i in
      Array.unsafe_set a i ai;
      Array.unsafe_set b i bi
    done;
    unsafe_of_array a, unsafe_of_array b
  end

let combine a b =
  let na = length a in
  let nb = length b in
  if na <> nb then invalid_arg "Iarray.combine";
  let r = if na = 0 then [||]
  else begin
    let x = Array.make na (unsafe_get a 0, unsafe_get b 0) in
    for i = 1 to na - 1 do
      Array.unsafe_set x i (unsafe_get a i, unsafe_get b i)
    done;
    x
  end in
  unsafe_of_array r

(* Must be fully applied due to the value restriction *)
let lift_sort sorter cmp iarr =
  let arr = to_array iarr in
  sorter cmp arr;
  unsafe_of_array arr

let sort cmp iarr = lift_sort Array.sort cmp iarr
let stable_sort cmp iarr = lift_sort Array.stable_sort cmp iarr
let fast_sort cmp iarr = lift_sort Array.fast_sort cmp iarr

let to_seq a =
  let rec aux i () =
    if i < length a
    then
      let x = unsafe_get a i in
      Seq.Cons (x, aux (i+1))
    else Seq.Nil
  in
  aux 0

let to_seqi a =
  let rec aux i () =
    if i < length a
    then
      let x = unsafe_get a i in
      Seq.Cons ((i,x), aux (i+1))
    else Seq.Nil
  in
  aux 0

let of_seq i = unsafe_of_array (Array.of_seq i)