open Array

type 'a t = 
    { length : int 
    ; offset : int 
    ; mutable data : 'a node }

and 'a node = 
  | Array of 'a array
  | Diff of int * 'a * 'a t
  | Swap of int * int * 'a t
  | Same of 'a t

let rec cache a = 
  match a.data with 
    | Array _       -> ()
    | Same b ->
        cache b ;
        begin match b.data with 
          | Array arr ->
              b.data <- Same a ;
              a.data <- Array arr
          | _ -> assert false
        end
    | Swap (i,j,b) ->
        cache b ;
        begin match b.data with 
          | Array arr ->
              let x = arr.(i) in arr.(i) <- arr.(j) ; arr.(j) <- x ;
              b.data <- Swap (i,j,a) ;
              a.data <- Array arr
          | _ -> assert false
        end
    | Diff (i,x,b) ->
        cache b ;
        begin match b.data with 
          | Array arr    ->
              let x' = arr.(i) in arr.(i) <- x ;
              b.data <- Diff (i,x',a) ;
              a.data <- Array arr
          | _ -> assert false
        end

let impl a =
  match a.data with 
    | Array arr -> arr
    | _         -> 
        cache a;
        match a.data with 
          | Array arr -> arr
          | _         -> assert false

let array a = 
  { length = length a 
  ; offset = 0 
  ; data   = Array a }

let set      a i x = 
  { length = a.length 
  ; offset = a.offset 
  ; data   = Diff (i+a.offset,x,a) }

let sub a o l = 
  { length = l 
  ; offset = a.offset + o 
  ; data   = Same a }

let swap a i j =
  { length = a.length
  ; offset = a.offset
  ; data   = Swap (i+a.offset,j+a.offset,a) }

let get      a i   = (impl a).(i + a.offset)
let length   a     = a.length
let of_array a     = array (copy a)
let to_array a     = Array.sub (impl a) a.offset a.length
let make     n x   = array (make n x)
let init     n f   = array (init n f)

let map      f a   = 
  array (Array.init a.length (fun i -> f (get a i)))

let mapi     f a   =
  array (Array.init a.length (fun i -> f i (get a i)))

let copy     a     =
  if a.offset = 0 && Array.length (impl a) = a.length then a
  else array (Array.sub (impl a) a.offset a.length)

let iter     f a   = 
  for i = 0 to a.length - 1 do f (get a i) done 

let iteri    f a   = 
  for i = 0 to a.length - 1 do f i (get a i) done

let rec fill     a i l x =
  if l = 0 then a else fill (set a i x) (i+1) (l-1) x

let rec blit     a i b j l = 
  if l = 0 then a else blit (set a i (get b j)) (i+1) b (j+1) (l-1)

let fold_left    f x a = 
  let accum = ref x in
  for i = 0 to a.length - 1 do accum := f !accum (get a i) done ;
  !accum

let fold_right   f a x =
  let accum = ref x in
  for i = a.length - 1 downto 0 do accum := f (get a i) !accum done ;
  !accum