
(** Random Access Lists.
    (Okasaki, "Purely Functional Data Structures", 10.1.2.)

    The code below uses polymorphic recursion (both in the logic
    and in the programs).

    Author: Jean-Christophe Filliâtre (CNRS)
 *)

module RandomAccessList

  use int.Int
  use int.ComputerDivision
  use list.List
  use list.Length
  use list.Nth
  use option.Option

  type ral 'a =
  | Empty
  | Zero    (ral ('a, 'a))
  | One  'a (ral ('a, 'a))

  function flatten (l: list ('a, 'a)) : list 'a =
    match l with
    | Nil -> Nil
    | Cons (x, y) l1 -> Cons x (Cons y (flatten l1))
    end

  let rec lemma length_flatten (l:list ('a, 'a))
    ensures { length (flatten l) = 2 * length l }
    variant { l }
    =
    match l with
    | Cons (_,_) q -> length_flatten q
    | Nil -> ()
    end

  function elements (l: ral 'a) : list 'a =
    match l with
    | Empty    -> Nil
    | Zero l1  -> flatten (elements l1)
    | One x l1 -> Cons x (flatten (elements l1))
    end

  let rec size (l: ral 'a) : int
    variant { l }
    ensures { result = length (elements l) }
    =
    match l with
    | Empty    -> 0
    | Zero  l1 ->     2 * size l1
    | One _ l1 -> 1 + 2 * size l1
    end

  let rec cons (x: 'a) (l: ral 'a) : ral 'a
    variant { l }
    ensures { elements result = Cons x (elements l) }
    =
    match l with
    | Empty    -> One x Empty
    | Zero l1  -> One x l1
    | One y l1 -> Zero (cons (x, y) l1)
    end

  let rec lemma nth_flatten (i: int) (l: list ('a, 'a))
    requires { 0 <= i < length l }
    variant  { l }
    ensures  { match nth i l with
               | None -> false
               | Some (x0, x1) -> Some x0 = nth (2 * i)     (flatten l) /\
                                  Some x1 = nth (2 * i + 1) (flatten l) end }
    =
    match l with
    | Nil -> ()
    | Cons _ r -> if i > 0 then nth_flatten (i-1) r
    end

  let rec lookup (i: int) (l: ral 'a) : 'a
    requires { 0 <= i < length (elements l) }
    variant  { i, l }
    ensures  { nth i (elements l) = Some result }
    =
    match l with
    | Empty    -> absurd
    | One x l1 -> if i = 0 then x else lookup (i-1) (Zero l1)
    | Zero l1  -> let (x0, x1) = lookup (div i 2) l1 in
                  if mod i 2 = 0 then x0 else x1
    end

  let rec tail (l: ral 'a) : ral 'a
    requires { elements l <> Nil }
    variant  { l }
    ensures  { match elements l with
               | Nil -> false
               | Cons _ l -> elements result = l
               end }
    =
    match l with
    | Empty    -> absurd
    | One _ l1 -> Zero l1
    | Zero  l1 -> let (_, x1) = lookup 0 l1 in One x1 (tail l1)
    end

  let rec update (i: int) (y: 'a) (l: ral 'a) : ral 'a
    requires { 0 <= i < length (elements l) }
    variant  { i, l}
    ensures  { nth i (elements result) = Some y}
    ensures  { forall j. 0 <= j < length (elements l) ->
               j <> i -> nth j (elements result) = nth j (elements l) }
    ensures  { length (elements result) = length (elements l) }
    ensures  { match result, l with
               | One _ _, One _ _ | Zero _, Zero _ -> true
               | _                                 -> false
               end }
    =
    match l with
    | Empty    -> absurd
    | One x l1 -> if i = 0 then One y l1 else
                  match update (i-1) y (Zero l1) with
                  | Empty | One _ _ -> absurd
                  | Zero l1         -> One x l1
                  end
    | Zero l1  ->
        let (x0, x1) = lookup (div i 2) l1 in
        let l1' =
          update (div i 2) (if mod i 2 = 0 then (y,x1) else (x0,y)) l1 in
        assert { forall j. 0 <= j < length (elements l) -> j <> i ->
                 match nth (div j 2) (elements l1) with
                 | None -> false
                 | Some (x0,_) -> Some x0 = nth (2 * (div j 2)) (elements l)
                 end
                 && nth j (elements l) = nth j (elements (Zero l1')) };
        Zero l1'
    end

end

(** A straightforward encapsulation with a list ghost model
    (in anticipation of module refinement) *)

module RAL

  use int.Int
  use list.List
  use list.Length
  use option.Option
  use list.Nth
  use RandomAccessList

  type t 'a = { r: ral 'a; ghost l: list 'a }
    invariant { l = elements r }

  let constant empty : t 'a = { r = Empty; l = Nil }
    ensures { result.l = Nil }

  let size (t: t 'a) : int
    ensures { result = length t.l }
    =
    size t.r

  let cons (x: 'a) (s: t 'a) : t 'a
    ensures { result.l = Cons x s.l }
    =
    { r = cons x s.r; l = Cons x s.l }

  let lookup (i: int) (s: t 'a) : 'a
    requires { 0 <= i < length s.l }
    ensures { Some result = nth i s.l }
    =
    lookup i s.r

end

(** Model using sequences instead of lists *)

module RandomAccessListWithSeq

  use int.Int
  use int.ComputerDivision
  use seq.Seq

  type ral 'a =
  | Empty
  | Zero    (ral ('a, 'a))
  | One  'a (ral ('a, 'a))

  function flatten (s: seq ('a, 'a)) : seq 'a =
    create (2 * length s)
           (fun i -> let (x0, x1) = s[div i 2] in
                     if mod i 2 = 0 then x0 else x1)

  lemma cons_flatten : forall x y :'a,s:seq ('a,'a).
    let a = flatten (cons (x,y) s) in
    let b = cons x (cons y (flatten s)) in
    a = b by a == b by forall i. 0 <= i < a.length -> a[i] = b[i]
    by (i <= 1 so (cons (x,y) s)[div i 2] = (x,y))
       \/ (i >= 2 so (cons (x,y) s)[div i 2] = s[div (i-2) 2]
       )

  function elements (l: ral 'a) : seq 'a =
    match l with
    | Empty    -> empty
    | Zero l1  -> flatten (elements l1)
    | One x l1 -> cons x (flatten (elements l1))
    end

  let rec size (l: ral 'a) : int
    variant { l }
    ensures { result = length (elements l) }
    =
    match l with
    | Empty    -> 0
    | Zero  l1 ->     2 * size l1
    | One _ l1 -> 1 + 2 * size l1
    end

  let rec cons (x: 'a) (l: ral 'a) : ral 'a
    variant { l }
    ensures { elements result == cons x (elements l) }
    =
    match l with
    | Empty    -> One x Empty
    | Zero l1  -> One x l1
    | One y l1 -> Zero (cons (x, y) l1)
    end

  let rec lookup (i: int) (l: ral 'a) : 'a
    requires { 0 <= i < length (elements l) }
    variant  { i, l }
    ensures  { (elements l)[i] = result }
    =
    match l with
    | Empty    -> absurd
    | One x l1 -> if i = 0 then x else lookup (i-1) (Zero l1)
    | Zero l1  -> let (x0, x1) = lookup (div i 2) l1 in
                  if mod i 2 = 0 then x0 else x1
    end

   let rec tail (l: ral 'a) : ral 'a
    requires { 0 < length (elements l) }
    variant  { l }
    ensures  { elements result == (elements l)[1..] }
    =
    match l with
    | Empty    -> absurd
    | One _ l1 -> Zero l1
    | Zero  l1 -> let (_, x1) as p = lookup 0 l1 in
      let tl = tail l1 in
      assert { elements l1 == cons p (elements tl) };
      One x1 tl
    end

  (** update in O(log n)
      for this, we need to pass a function that will update the element
      when we find it *)

  function setf (s: seq 'a) (i: int) (f: 'a -> 'a) : seq 'a =
    set s i (f s[i])

  let function aux (i: int) (f: 'a -> 'a) : ('a, 'a) -> ('a, 'a) =
    fun z -> let (x,y) = z in if mod i 2 = 0 then (f x, y) else (x, f y)

  let rec fupdate (f: 'a -> 'a) (i: int) (l: ral 'a) : ral 'a
    requires { 0 <= i < length (elements l) }
    variant  { i, l }
    ensures  { elements result == setf (elements l) i f }
    =
    match l with
    | Empty    -> absurd
    | One x l1 -> if i = 0 then One (f x) l1 else
                  cons x (fupdate f (i-1) (Zero l1))
    | Zero l1  ->
      let ul1 = fupdate (aux i f) (div i 2) l1 in
      let res = Zero ul1 in
      assert { forall j. 0 <= j < length (elements res) ->
        (elements res)[j] = (setf (elements l) i f)[j]
        by div j 2 <> div i 2 ->
          (elements ul1)[div j 2] = (elements l1)[div j 2]
      };
      res
    end

  let function f (y: 'a) : 'a -> 'a = fun _ -> y

  let update (i: int) (y: 'a) (l: ral 'a) : ral 'a
    requires { 0 <= i < length (elements l) }
    ensures  { elements result == set (elements l) i y}
    =
    fupdate (f y) i  l

end
