(* Example: affine arrays *)

module type AF_ARRAY = sig
  type 'a array : A

  val new : int → 'a → 'a array
  val set : 'a array → int ⊸ 'a ⊸ 'a array
  val get : 'a array → int ⊸ 'a * 'a array

  val size : 'a array → int * 'a array
end

#load "libarray"
module A = Array

module AfArray : AF_ARRAY = struct
  type 'a array = 'a A.array

  let new = A.new
  let set (a: 'a array) (ix: int) (v: 'a) =
    A.set a ix v; a
  let get (a: 'a array) (ix: int) =
    (A.get a ix, a)

  let size (a: 'a array) = (A.size a, a)
end

let deposit (a: int AfArray.array) (ix: int) (amount: int) =
  let (balance, a) = AfArray.get a ix in
    AfArray.set a ix (balance + amount)

(*** Some definitions used by the next example. ***)

 (* Placing them here ensures that "make test" will catch if they
  * stop typing, since the example where they are used demonstrates
  * a type error. *)

module A = AfArray

(* Swap the values at the given array indices *)
let swapIndices a i j =
  let (ai, a) = A.get a i in
  let (aj, a) = A.get a j in
    A.set (A.set a i aj) j ai

(* Fisher-Yates shuffle *)
let inPlaceShuffle a =
  let rec loop i a =
    if i == 0
      then a
      else let j = random_int () % (i + 1) in
             loop (i - 1) (swapIndices a i j) in
  let (n, a) = A.size a in
    loop (n - 1) a

(* Quicksort *)
let inPlaceSort a =
  let rec quicksort start limit a =
    if limit > start
      then let (pivot, a) = A.get a limit in
           let rec loop i j a =
             if i < limit
               then let (ai, a) = A.get a i in
                      if ai ≤ pivot
                        then loop (i + 1) (j + 1) (swapIndices a i j)
                        else loop (i + 1) j a
               else (j, a) in
           let (j, a) = loop start start a in
           let a      = swapIndices a j limit in
           let a      = quicksort start (j - 1) a in
             quicksort (j + 1) limit a
      else a in
  let (n, a) = A.size a in
    quicksort 0 (n - 1) a

(* For testing: *)
let listToArray (x ∷ xs) =
  let n = length xs + 1 in
  let rec loop i xs a =
    match xs with
    | []     → a
    | x ∷ xs → loop (i + 1) xs (A.set a i x)
   in loop 1 xs (A.new n x)

let arrayToList a =
  let (n, a) = A.size a in
  let rec loop i xs a =
        if i < 0
          then (xs, a)
          else let (ai, a) = A.get a i in
                 loop (i - 1) (ai ∷ xs) a
   in loop (n - 1) [] a

module Tests = struct
  let unsorted  = [4, 1, 0, 3, 2]
  let sorted    = [0, 1, 2, 3, 4]
  let sorted'   = fst(arrayToList(inPlaceSort(listToArray(unsorted))))
  let () = if sorted == sorted'
             then ()
             else failwith "test failed: inPlaceSort (1)"

  let sorted' =
        fst(arrayToList(inPlaceSort(inPlaceShuffle(listToArray(sorted)))))
  let () = if sorted == sorted'
             then ()
             else failwith "test failed: inPlaceSort (2)"
end