(* Fibonacci function with memoisation *)

module FibMemo

  use option.Option
  use int.Int
  use int.Fibonacci
  use ref.Ref
  use map.Map

  type table = map int (option int)

  predicate inv (t : table) =
    forall x y : int. t[x] = Some y -> y = fib x

  val table : ref table

  val add (x:int) (y:int) : unit
    writes  { table }
    ensures { !table = (old !table)[x <- Some y] }

  exception Not_found

  val find (x:int) : int
    ensures { !table[x] = Some result }
    raises  { Not_found -> !table[x] = None }

  let rec fibo n =
    requires { 0 <= n /\ inv !table }
    ensures  { result = fib n /\ inv !table }
    variant  { 2*n }
    if n <= 1 then
      n
    else
      memo_fibo (n-1) + memo_fibo (n-2)

  with memo_fibo n =
    requires { 0 <= n /\ inv !table }
    ensures  { result = fib n /\ inv !table }
    variant  { 2*n+1 }
    try  find n
    with Not_found -> let fn = fibo n in add n fn; fn end

end
