sort_corresponding_arrays.ml

(** [sort_corresponding_arrays cmp xs ys] is equivalent to
    [Stdlib.Array.sort xs] but also applies the same reordering of indices to
    the elements of [ys], ensuring that all pairs of elements that share an
    index {i before} the sort still do so afterwards. Equivalent to:

    {[
      let a3 = Array.combine xs ys in
      Array.sort (fun (x1, _) (x2, _) -> cmp x1 x2);
      Array.iteri (fun i (x, y) ->
          xs.(i) <- x;
          ys.(i) <- y)
    ]} 
    
    Implementation adapted from {!Stdlib.Array.sort}. *)
let sort_corresponding_arrays :
    type a b. (a -> a -> int) -> a array -> b array -> unit =
 fun cmp a b ->
  let exception Bottom of int in
  let open Array in
  let maxson l i =
    let i31 = i + i + i + 1 in
    let x = ref i31 in
    if i31 + 2 < l then (
      if cmp (get a i31) (get a (i31 + 1)) < 0 then x := i31 + 1;
      if cmp (get a !x) (get a (i31 + 2)) < 0 then x := i31 + 2;
      !x)
    else if i31 + 1 < l && cmp (get a i31) (get a (i31 + 1)) < 0 then i31 + 1
    else if i31 < l then i31
    else raise (Bottom i)
  in
  let rec trickledown l i e e' =
    let j = maxson l i in
    if cmp (get a j) e > 0 then (
      set a i (get a j);
      set b i (get b j);
      trickledown l j e e')
    else (
      set a i e;
      set b i e')
  in
  let trickle l i e e' =
    try trickledown l i e e'
    with Bottom i ->
      set a i e;
      set b i e'
  in
  let rec bubbledown l i =
    let j = maxson l i in
    set a i (get a j);
    set b i (get b j);
    bubbledown l j
  in
  let bubble l i = try bubbledown l i with Bottom i -> i in
  let rec trickleup i e e' =
    let father = (i - 1) / 3 in
    assert (i <> father);
    if cmp (get a father) e < 0 then (
      set a i (get a father);
      set b i (get b father);
      if father > 0 then trickleup father e e'
      else (
        set a 0 e;
        set b 0 e'))
    else (
      set a i e;
      set b i e')
  in
  let l = length a in
  for i = ((l + 1) / 3) - 1 downto 0 do
    trickle l i (get a i) (get b i)
  done;
  for i = l - 1 downto 2 do
    let e = get a i in
    let e' = get b i in
    set a i (get a 0);
    set b i (get b 0);
    trickleup (bubble i 0) e e'
  done;
  if l > 1 then (
    let e = get a 1 in
    let e' = get b 1 in
    set a 1 (get a 0);
    set b 1 (get b 0);
    set a 0 e;
    set b 0 e')