(* This is an implementation of Okasaki's binary random access lists (Okasaki,
   "Purely Functional Data Structures", 10.1.2. *)

(* A random access list represents a sequence of elements. *)

(* A sequence either is empty, or has an even number of elements, or has an
   odd number of elements. These three cases are reflected in the definition
   that follows. *)

(* When the sequence is non-empty, we have zero or one element at the root,
   as appropriate. The remainder of the sequence has even length, so we can
   organize it as a sequence of pairs, whose type is [('a * 'a) seq]. *)

(* This means that a sequence of [n] elements is organized as a list of
   [Zero] and [One] digits which corresponds to the number [n] in binary
   notation (with the least significant digit at the head of the list).
   The length of this list is, obviously, logarithmic in [n]. *)

(* Since [Zero] and [One] are used only when the sequence is non-empty,
   [Zero] is never followed with [Nil]. In other words, there are no
   trailing zeros. *)

type 'a seq =
| Nil
| Zero of     ('a * 'a) seq
| One of 'a * ('a * 'a) seq

(* An example. The number eleven is written 1101 in binary notation (least
   significant digit first). So, a sequence of eleven elements is represented
   by a data structure of the form [One (_, One (_, Zero (One (_, Nil))))].
   For example, here is the sequence of integers 0 .. 10. *)

let (_ : int seq) =
  One (0,
  One ((1, 2), 
  Zero (
  One ((((3, 4), (5, 6)), ((7, 8), (9, 10))),
  Nil
  ))))

(* The empty sequence. *)

let empty =
  Nil

(* A test for emptiness. *)

let is_empty xs =
  match xs with
  | Nil -> true
  | Zero (_) | One (_, _) -> false

(* Counting the number of elements in a sequence. *)

(* If [length] is called with a sequence of type ['a seq], it may call itself,
   recursively, with a sequence of type [('a * 'a) seq]. Thus, polymorphic
   recursion is used here. *)

(* In OCaml, when polymorphic recursion is used, this must be explicitly
   declared. The type of [length] must be given and must be explicitly
   universally quantified -- note the [ 'a . ] prefix, which means
   "for all 'a". *)

let rec length : 'a . 'a seq -> int =
  fun xs ->
    match xs with
    | Nil         ->                 0
    | Zero xs     ->     2 * length xs
    | One (_, xs) -> 1 + 2 * length xs

(* Consing -- inserting an element in front of a sequence. *)

let rec cons : 'a . 'a -> 'a seq -> 'a seq =
  fun x ys ->
    match ys with
    | Nil ->
        One (x, Nil)
    | Zero ys ->
        One (x, ys)
    | One (y, ys) ->
        Zero (cons (x, y) ys)

(* Unconsing -- extracting and removing the element found in front
   of a sequence. *)

let rec uncons : 'a . 'a seq -> ('a * 'a seq) option =
  fun xs ->
    match xs with
    | Nil ->
        None
    | One (x, Nil) ->
        Some (x, Nil)
    | One (x, ys) ->
        Some (x, Zero ys)
    | Zero ys ->
        match uncons ys with
        | Some ((x, y), ys) ->
            Some (x, One (y, ys))
        | None ->
            assert false (* cannot happen; no trailing zeros *)

(* Random access: accessing the [i]-th element for reading. *)

(* [i] must be comprised between [0] included and [length xs] excluded. *)

let rec get : 'a . int -> 'a seq -> 'a =
  fun i xs ->
    match xs with
    | Nil ->
        assert false (* cannot happen; [i] is within bounds *)
    | One (x, xs) ->
        if i = 0 then
          x
        else
          get (i - 1) (Zero xs)
    | Zero xs ->
        let (x0, x1) = get (i / 2) xs in
        if i mod 2 = 0 then x0 else x1

(* Random access: accessing the [i]-th element for updating. *)

(* [i] must be comprised between [0] included and [length xs] excluded. *)

(* The user-supplied function [f] has type ['a -> 'a]. It represents the
   update that we wish to apply to the [i]-th element. This formulation,
   suggested by Okasaki, is well-suited to a recursive formulation of
   [fupdate]. The complexity remains O(log n) because we have O(log n) nested
   recursive calls to [update] *and* we build a closure of length O(log n). *)

let rec fupdate : 'a . int -> ('a -> 'a) -> 'a seq -> 'a seq =
  fun i f xs ->
    match xs with
    | Nil ->
        assert false (* cannot happen; [i] is within bounds *)
    | One (x, xs) ->
        if i = 0 then
          One (f x, xs)
        else
          cons x (fupdate (i - 1) f (Zero xs))
    | Zero xs ->
        let f' =
          if i mod 2 = 0 then
            fun (x0, x1) -> (f x0, x1)
          else
            fun (x0, x1) -> (x0, f x1)
        in
        Zero (fupdate (i / 2) f' xs)

let update i y xs =
  fupdate i (fun _ -> y) xs