(******************************************************************************)
(*                                                                            *)
(*                                  Monolith                                  *)
(*                                                                            *)
(*                              François Pottier                              *)
(*                                                                            *)
(*  Copyright Inria. All rights reserved. This file is distributed under the  *)
(*  terms of the GNU Lesser General Public License as published by the Free   *)
(*  Software Foundation, either version 3 of the License, or (at your         *)
(*  option) any later version, as described in the file LICENSE.              *)
(*                                                                            *)
(******************************************************************************)

open Misc

(* Variables. *)

type level =
  | Level of int [@@unboxed]
    (* a de Bruijn level: 0 is the most ancient variable *)

type var =
  level

let level (Level x) =
  x

(* For efficiency, an environment is represented as an array, indexed by de
   Bruijn levels. Environment lookup and environment extension are both
   constant-time operations. The array is resized when necessary; this is
   expected to be infrequent. *)

(* For efficiency, we also maintain an integer array whose size is the maximum
   size of the environment. This array is used as auxiliary storage in the
   implementation of the operation [var], which chooses an element of the
   environment that satisfies a predicate [p]. Outside of this operation, it
   is unused. *)

type 'v env = {
  mutable data : 'v array;
  mutable n : int;
  mutable storage: int array;
  dummy : 'v;
}

let empty capacity dummy =
  assert (0 < capacity);
  let data = Array.make capacity dummy
  and n = 0
  and storage = Array.make capacity 0 in
  { data; n; storage; dummy }

let resize env =
  let data = env.data in
  let capacity = Array.length data in
  let capacity' = 2 * capacity in
  env.data <- Array.make capacity' env.dummy;
  Array.blit data 0 env.data 0 capacity;
  env.storage <- Array.make capacity' 0

let clear env =
  env.n <- 0

let length env =
  env.n

let limit env =
  Level (length env)

let lookup env (Level x) =
  env.data.(x)

let rec bind env v =
  let { data; n; _ } = env in
  if n = Array.length data then begin
    resize env;
    assert (n < Array.length env.data);
    bind env v
  end
  else begin
    Array.set data n v;
    env.n <- n + 1
  end

let foreach env f =
  let data, n = env.data, env.n in
  assert (n <= Array.length data);
  for x = 0 to n - 1 do
    f (Level x) data.(x)
  done

let choose env f =
  let data, n, storage = env.data, env.n, env.storage in
  (* Construct an auxiliary array of the indices of the values that
     satisfy [f]. This information is stored in the array [storage],
     so we save the cost of allocating and initializing an array. *)
  let k = ref 0 in
  for i = 0 to n - 1 do
    let x = Level i in
    match f x data.(i) with
    | None ->
        ()
    | Some _w ->
        (* We cannot store [w], as we have no storage for it. *)
        (* We record the fact that [i] is a good index. *)
        storage.(postincrement k) <- i
  done;
  (* Pick an index among our [k] candidates. *)
  let i = storage.(Gen.int !k ()) in
  let x = Level i in
  (* Invoke [f] again so as to obtain [w]. *)
  match f x data.(i) with
  | Some w ->
      w
  | None ->
      assert false
