belt_Array.ml

(***********************************************************************)
(*                                                                     *)
(*                           OCaml                                     *)
(*                                                                     *)
(*            Xavier Leroy, projet Cristal, INRIA Rocquencourt         *)
(*                                                                     *)
(*  Copyright 1996 Institut National de Recherche en Informatique et   *)
(*  en Automatique.  All rights reserved.  This file is distributed    *)
(*  under the terms of the GNU Library General Public License, with    *)
(*  the special exception on linking described in file ../LICENSE.     *)
(*                                                                     *)
(***********************************************************************)

(* Array operations *)
type 'a t = 'a array

external length : 'a t -> int = "%array_length"

external size : 'a t -> int = "%array_length"

external getUnsafe : 'a t -> int -> 'a = "%array_unsafe_get"

external setUnsafe : 'a t -> int -> 'a -> unit = "%array_unsafe_set"

external ( .!() ) : 'a t -> int -> 'a = "%array_unsafe_get"

external ( .!()<- ) : 'a t -> int -> 'a -> unit = "%array_unsafe_set"

external getUndefined : 'a t -> int -> 'a Js.undefined = "%array_unsafe_get"

(* external get: 'a t -> int -> 'a = "%array_safe_get" *)
let get arr i = if i >= 0 && i < length arr then Some arr.!(i) else None

let getExn arr i =
  assert (i >= 0 && i < length arr);
  arr.!(i)

let set arr i v =
  if i >= 0 && i < length arr then (
    arr.!(i) <- v;
    true)
  else false

let setExn arr i v =
  assert (i >= 0 && i < length arr);
  arr.!(i) <- v

external truncateToLengthUnsafe : 'a t -> int -> unit = "length" [@@bs.set]

external makeUninitialized : int -> 'a Js.undefined array = "Array" [@@bs.new]

external makeUninitializedUnsafe : int -> 'a array = "Array" [@@bs.new]

external copy : 'a t -> (_[@bs.as 0]) -> 'a t = "slice" [@@bs.send]

let swapUnsafe xs i j =
  let tmp = xs.!(i) in
  xs.!(i) <- xs.!(j);
  xs.!(j) <- tmp

let shuffleInPlace xs =
  let len = length xs in
  for i = 0 to len - 1 do
    swapUnsafe xs i (Js_math.random_int i len) (* [i,len)*)
  done

let shuffle xs =
  let result = copy xs in
  shuffleInPlace result;
  (* TODO: improve*)
  result

let reverseAux xs ofs len =
  for i = 0 to (len / 2) - 1 do
    swapUnsafe xs (ofs + i) (ofs + len - i - 1)
  done

let reverseInPlace xs =
  let len = length xs in
  reverseAux xs 0 len

let reverse xs =
  let len = length xs in
  let result = makeUninitializedUnsafe len in
  for i = 0 to len - 1 do
    result.!(i) <- xs.!(len - 1 - i)
  done;
  result

let make l f =
  if l <= 0 then [||]
  else
    let res = makeUninitializedUnsafe l in
    for i = 0 to l - 1 do
      res.!(i) <- f
    done;
    res

(* See #6575. We could also check for maximum array size, but this depends
     on whether we create a float array or a regular one... *)
let makeByU l f =
  if l <= 0 then [||]
  else
    let res = makeUninitializedUnsafe l in
    for i = 0 to l - 1 do
      res.!(i) <- (f i [@bs])
    done;
    res

let makeBy l f = makeByU l (fun [@bs] a -> f a)

let makeByAndShuffleU l f =
  let u = makeByU l f in
  shuffleInPlace u;
  u

let makeByAndShuffle l f = makeByAndShuffleU l (fun [@bs] a -> f a)

let range start finish =
  let cut = finish - start in
  if cut < 0 then [||]
  else
    let arr = makeUninitializedUnsafe (cut + 1) in
    for i = 0 to cut do
      arr.!(i) <- start + i
    done;
    arr

let rangeBy start finish ~step =
  let cut = finish - start in
  if cut < 0 || step <= 0 then [||]
  else
    let nb = (cut / step) + 1 in
    let arr = makeUninitializedUnsafe nb in
    let cur = ref start in
    for i = 0 to nb - 1 do
      arr.!(i) <- cur.contents;
      cur.contents <- cur.contents + step
    done;
    arr

let zip xs ys =
  let lenx, leny = (length xs, length ys) in
  let len = Pervasives.min lenx leny in
  let s = makeUninitializedUnsafe len in
  for i = 0 to len - 1 do
    s.!(i) <- (xs.!(i), ys.!(i))
  done;
  s

let zipByU xs ys f =
  let lenx, leny = (length xs, length ys) in
  let len = Pervasives.min lenx leny in
  let s = makeUninitializedUnsafe len in
  for i = 0 to len - 1 do
    s.!(i) <- (f xs.!(i) ys.!(i) [@bs])
  done;
  s

let zipBy xs ys f = zipByU xs ys (fun [@bs] a b -> f a b)

let concat a1 a2 =
  let l1 = length a1 in
  let l2 = length a2 in
  let a1a2 = makeUninitializedUnsafe (l1 + l2) in
  for i = 0 to l1 - 1 do
    a1a2.!(i) <- a1.!(i)
  done;
  for i = 0 to l2 - 1 do
    a1a2.!(l1 + i) <- a2.!(i)
  done;
  a1a2

let concatMany arrs =
  let lenArrs = length arrs in
  let totalLen = ref 0 in
  for i = 0 to lenArrs - 1 do
    totalLen.contents <- totalLen.contents + length arrs.!(i)
  done;
  let result = makeUninitializedUnsafe totalLen.contents in
  totalLen.contents <- 0;
  for j = 0 to lenArrs - 1 do
    let cur = arrs.!(j) in
    for k = 0 to length cur - 1 do
      result.!(totalLen.contents) <- cur.!(k);
      totalLen.contents <- totalLen.contents + 1
    done
  done;
  result

let slice a ~offset ~len =
  if len <= 0 then [||]
  else
    let lena = length a in
    let ofs = if offset < 0 then Pervasives.max (lena + offset) 0 else offset in
    let hasLen = lena - ofs in
    let copyLength = Pervasives.min hasLen len in
    if copyLength <= 0 then [||]
    else
      let result = makeUninitializedUnsafe copyLength in
      for i = 0 to copyLength - 1 do
        result.!(i) <- a.!(ofs + i)
      done;
      result



let sliceToEnd a offset =
  let lena = length a in
  let ofs = if offset < 0 then Pervasives.max (lena + offset) 0 else offset in
  let len = if lena > ofs then lena - ofs else 0 in
  let result = makeUninitializedUnsafe len in
  for i = 0 to len - 1 do
    result.!(i) <- a.!(ofs + i)
  done;
  result

let fill a ~offset ~len v =
  if len > 0 then
    let lena = length a in
    let ofs = if offset < 0 then Pervasives.max (lena + offset) 0 else offset in
    let hasLen = lena - ofs in
    let fillLength = Pervasives.min hasLen len in
    if fillLength > 0 then
      for i = ofs to ofs + fillLength - 1 do
        a.!(i) <- v
      done

let blitUnsafe ~src:a1 ~srcOffset:srcofs1 ~dst:a2 ~dstOffset:srcofs2
    ~len:blitLength =
  if srcofs2 <= srcofs1 then
    for j = 0 to blitLength - 1 do
      a2.!(j + srcofs2) <- a1.!(j + srcofs1)
    done
  else
    for j = blitLength - 1 downto 0 do
      a2.!(j + srcofs2) <- a1.!(j + srcofs1)
    done

(* We don't need check `blitLength` since when `blitLength < 0` the
   for loop will be nop
*)
let blit ~src:a1 ~srcOffset:ofs1 ~dst:a2 ~dstOffset:ofs2 ~len =
  let lena1 = length a1 in
  let lena2 = length a2 in
  let srcofs1 = if ofs1 < 0 then Pervasives.max (lena1 + ofs1) 0 else ofs1 in
  let srcofs2 = if ofs2 < 0 then Pervasives.max (lena2 + ofs2) 0 else ofs2 in
  let blitLength =
    Pervasives.min len (Pervasives.min (lena1 - srcofs1) (lena2 - srcofs2))
  in
  (* blitUnsafe a1 srcofs1 a2 srcofs2 blitLength *)
  if srcofs2 <= srcofs1 then
    for j = 0 to blitLength - 1 do
      a2.!(j + srcofs2) <- a1.!(j + srcofs1)
    done
  else
    for j = blitLength - 1 downto 0 do
      a2.!(j + srcofs2) <- a1.!(j + srcofs1)
    done

let forEachU a f =
  for i = 0 to length a - 1 do
    (f a.!(i) [@bs])
  done

let forEach a f = forEachU a (fun [@bs] a -> f a)

let mapU a f =
  let l = length a in
  let r = makeUninitializedUnsafe l in
  for i = 0 to l - 1 do
    r.!(i) <- (f a.!(i) [@bs])
  done;
  r

let map a f = mapU a (fun [@bs] a -> f a)

let flatMapU a f = concatMany (mapU a f)

let flatMap a f = flatMapU a (fun [@bs] a -> f a)

let getByU a p =
  let l = length a in
  let i = ref 0 in
  let r = ref None in
  while r.contents = None && i.contents < l do
    let v = a.!(i.contents) in
    if p v [@bs] then r.contents <- Some v;
    i.contents <- i.contents + 1
  done;
  r.contents

let getBy a p = getByU a (fun [@bs] a -> p a)

let getIndexByU a p =
  let l = length a in
  let i = ref 0 in
  let r = ref None in
  while r.contents = None && i.contents < l do
    let v = a.!(i.contents) in
    if p v [@bs] then r.contents <- Some i.contents;
    i.contents <- i.contents + 1
  done;
  r.contents

let getIndexBy a p = getIndexByU a (fun [@bs] a -> p a)

let keepU a f =
  let l = length a in
  let r = makeUninitializedUnsafe l in
  let j = ref 0 in
  for i = 0 to l - 1 do
    let v = a.!(i) in
    if f v [@bs] then (
      r.!(j.contents) <- v;
      j.contents <- j.contents + 1)
  done;
  truncateToLengthUnsafe r j.contents;
  r

let keep a f = keepU a (fun [@bs] a -> f a)

let keepWithIndexU a f =
  let l = length a in
  let r = makeUninitializedUnsafe l in
  let j = ref 0 in
  for i = 0 to l - 1 do
    let v = a.!(i) in
    if f v i [@bs] then (
      r.!(j.contents) <- v;
      j.contents <- j.contents + 1)
  done;
  truncateToLengthUnsafe r j.contents;
  r

let keepWithIndex a f = keepWithIndexU a (fun [@bs] a i -> f a i)

let keepMapU a f =
  let l = length a in
  let r = makeUninitializedUnsafe l in
  let j = ref 0 in
  for i = 0 to l - 1 do
    let v = a.!(i) in
    match f v [@bs] with
    | None -> ()
    | Some v ->
        r.!(j.contents) <- v;
        j.contents <- j.contents + 1
  done;
  truncateToLengthUnsafe r j.contents;
  r

let keepMap a f = keepMapU a (fun [@bs] a -> f a)

let forEachWithIndexU a f =
  for i = 0 to length a - 1 do
    (f i a.!(i) [@bs])
  done

let forEachWithIndex a f = forEachWithIndexU a (fun [@bs] a b -> f a b)

let mapWithIndexU a f =
  let l = length a in
  let r = makeUninitializedUnsafe l in
  for i = 0 to l - 1 do
    r.!(i) <- (f i a.!(i) [@bs])
  done;
  r

let mapWithIndex a f = mapWithIndexU a (fun [@bs] a b -> f a b)

let reduceU a x f =
  let r = ref x in
  for i = 0 to length a - 1 do
    r.contents <- (f r.contents a.!(i) [@bs])
  done;
  r.contents

let reduce a x f = reduceU a x (fun [@bs] a b -> f a b)

let reduceReverseU a x f =
  let r = ref x in
  for i = length a - 1 downto 0 do
    r.contents <- (f r.contents a.!(i) [@bs])
  done;
  r.contents

let reduceReverse a x f = reduceReverseU a x (fun [@bs] a b -> f a b)

let reduceReverse2U a b x f =
  let r = ref x in
  let len = Pervasives.min (length a) (length b) in
  for i = len - 1 downto 0 do
    r.contents <- (f r.contents a.!(i) b.!(i) [@bs])
  done;
  r.contents

let reduceReverse2 a b x f = reduceReverse2U a b x (fun [@bs] a b c -> f a b c)

let reduceWithIndexU a x f =
  let r = ref x in
  for i = 0 to length a - 1 do
    r.contents <- (f r.contents a.!(i) i [@bs])
  done;
  r.contents

let reduceWithIndex a x f = reduceWithIndexU a x (fun [@bs] a b c -> f a b c)

let rec everyAux arr i b len =
  if i = len then true
  else if b arr.!(i) [@bs] then everyAux arr (i + 1) b len
  else false

let rec someAux arr i b len =
  if i = len then false
  else if b arr.!(i) [@bs] then true
  else someAux arr (i + 1) b len

let everyU arr b =
  let len = length arr in
  everyAux arr 0 b len

let every arr f = everyU arr (fun [@bs] b -> f b)

let someU arr b =
  let len = length arr in
  someAux arr 0 b len

let some arr f = someU arr (fun [@bs] b -> f b)

let rec everyAux2 arr1 arr2 i b len =
  if i = len then true
  else if b arr1.!(i) arr2.!(i) [@bs] then everyAux2 arr1 arr2 (i + 1) b len
  else false

let rec someAux2 arr1 arr2 i b len =
  if i = len then false
  else if b arr1.!(i) arr2.!(i) [@bs] then true
  else someAux2 arr1 arr2 (i + 1) b len

let every2U a b p = everyAux2 a b 0 p (Pervasives.min (length a) (length b))

let every2 a b p = every2U a b (fun [@bs] a b -> p a b)

let some2U a b p = someAux2 a b 0 p (Pervasives.min (length a) (length b))

let some2 a b p = some2U a b (fun [@bs] a b -> p a b)

let eqU a b p =
  let lena = length a in
  let lenb = length b in
  if lena = lenb then everyAux2 a b 0 p lena else false

let eq a b p = eqU a b (fun [@bs] a b -> p a b)

let rec everyCmpAux2 arr1 arr2 i b len =
  if i = len then 0
  else
    let c = (b arr1.!(i) arr2.!(i) [@bs]) in
    if c = 0 then everyCmpAux2 arr1 arr2 (i + 1) b len else c

let cmpU a b p =
  let lena = length a in
  let lenb = length b in
  if lena > lenb then 1
  else if lena < lenb then -1
  else everyCmpAux2 a b 0 p lena

let cmp a b p = cmpU a b (fun [@bs] a b -> p a b)

let partitionU a f =
  let l = length a in
  let i = ref 0 in
  let j = ref 0 in
  let a1 = makeUninitializedUnsafe l in
  let a2 = makeUninitializedUnsafe l in
  for ii = 0 to l - 1 do
    let v = a.!(ii) in
    if f v [@bs] then (
      a1.!(i.contents) <- v;
      i.contents <- i.contents + 1)
    else (
      a2.!(j.contents) <- v;
      j.contents <- j.contents + 1)
  done;
  truncateToLengthUnsafe a1 i.contents;
  truncateToLengthUnsafe a2 j.contents;
  (a1, a2)

let partition a f = partitionU a (fun [@bs] x -> f x)

let unzip a =
  let l = length a in
  let a1 = makeUninitializedUnsafe l in
  let a2 = makeUninitializedUnsafe l in
  for i = 0 to l - 1 do
    let v1, v2 = a.!(i) in
    a1.!(i) <- v1;
    a2.!(i) <- v2
  done;
  (a1, a2)

let joinWithU a sep toString =
  match length a with
  | 0 -> ""
  | l ->
      let lastIndex = l - 1 in
      let rec aux i res =
        if i = lastIndex then res ^ (toString a.!(i) [@bs])
        else aux (i + 1) (res ^ (toString a.!(i) [@bs]) ^ sep)
      in
      aux 0 ""

let joinWith a sep toString = joinWithU a sep (fun [@bs] x -> toString x)

let initU n f =
  let v = makeUninitializedUnsafe n in
  for i = 0 to n - 1 do
    v.!(i) <- (f i [@bs])
  done;
  v

let init n f = initU n (fun [@bs] i -> f i)

external push : 'a t -> 'a -> unit = "push" [@@send]