
(** {1 Finite Maps} *)

(** {2 Polymorphic maps to be used in spec/ghost only} *)

module Fmap

  use int.Int
  use map.Map
  use set.Fset as S

  type fmap 'k 'v = abstract {
    contents: 'k -> 'v;
      domain: S.fset 'k;
  }
  meta coercion function contents

  predicate (==) (m1 m2: fmap 'k 'v) =
    S.(==) m1.domain m2.domain /\
    forall k. S.mem k m1.domain -> m1[k] = m2[k]

  axiom extensionality:
    forall m1 m2: fmap 'k 'v. m1 == m2 -> m1 = m2

  predicate mem (k: 'k) (m: fmap 'k 'v) =
    S.mem k m.domain

  predicate mapsto (k: 'k) (v: 'v) (m: fmap 'k 'v) =
    mem k m /\ m[k] = v

  lemma mem_mapsto:
    forall k: 'k, m: fmap 'k 'v. mem k m -> mapsto k m[k] m

  predicate is_empty (m: fmap 'k 'v) =
    S.is_empty m.domain

  function mk (d: S.fset 'k) (m: 'k -> 'v) : fmap 'k 'v

  axiom mk_domain:
    forall d: S.fset 'k, m: 'k -> 'v. domain (mk d m) = d

  axiom mk_contents:
    forall d: S.fset 'k, m: 'k -> 'v, k: 'k.
    S.mem k d -> (mk d m)[k] = m[k]

  constant empty: fmap 'k 'v

  axiom is_empty_empty: is_empty (empty: fmap 'k 'v)

  function add (k: 'k) (v: 'v) (m: fmap 'k 'v) : fmap 'k 'v

  function ([<-]) (m: fmap 'k 'v) (k: 'k) (v: 'v) : fmap 'k 'v =
    add k v m

  (*FIXME? (add k v m).contents = m.contents[k <- v] *)

  axiom add_contents_k:
    forall k v, m: fmap 'k 'v. (add k v m)[k] = v

  axiom add_contents_other:
    forall k v, m: fmap 'k 'v, k1. mem k1 m -> k1 <> k -> (add k v m)[k1] = m[k1]

  axiom add_domain:
    forall k v, m: fmap 'k 'v. (add k v m).domain = S.add k m.domain

  (* FIXME? find_opt (k: 'k) (m: fmap 'k 'v) : option 'v *)

  function find (k: 'k) (m: fmap 'k 'v) : 'v

  axiom find_def:
    forall k, m: fmap 'k 'v. mem k m -> find k m = m[k]

  function remove (k: 'k) (m: fmap 'k 'v) : fmap 'k 'v

  axiom remove_contents:
    forall k, m: fmap 'k 'v, k1. mem k1 m -> k1 <> k -> (remove k m)[k1] = m[k1]

  axiom remove_domain:
    forall k, m: fmap 'k 'v. (remove k m).domain = S.remove k m.domain

  function size (m: fmap 'k 'v) : int =
    S.cardinal m.domain

end

(** {2 Finite monomorphic maps to be used in programs only}

A program function `eq` deciding equality on the `key` type must be provided when cloned.
*)

(** {3 Applicative maps} *)

module MapApp

  use int.Int
  use map.Map
  use export Fmap

  type key

  (* we enforce type `key` to have a decidable equality
     by requiring the following function *)
  val eq (x y: key) : bool
    ensures { result <-> x = y }

  type t 'v = abstract {
    to_fmap: fmap key 'v;
  }
  meta coercion function to_fmap

  val create () : t 'v
    ensures { result.to_fmap = empty }

  val mem (k: key) (m: t 'v) : bool
    ensures { result <-> mem k m }

  val is_empty (m: t 'v) : bool
    ensures { result <-> is_empty m }

  val add (k: key) (v: 'v) (m: t 'v) : t 'v
    ensures { result = add k v m }

  val find (k: key) (m: t 'v) : 'v
    requires { mem k m }
    ensures  { result = m[k] }
    ensures  { result = find k m }

  use ocaml.Exceptions

  val find_exn (k: key) (m: t 'v) : 'v
    ensures { S.mem k m.domain }
    ensures { result = m[k] }
    raises  { Not_found ->  not (S.mem k m.domain) }

  val remove (k: key) (m: t 'v) : t 'v
    ensures { result = remove k m }

  val size (m: t 'v) : int
    ensures { result = size m }

end

(** {3 Applicative maps of integers} *)

module MapAppInt

  use int.Int

  clone export MapApp with type key = int, val eq = Int.(=), axiom .

end

(** {3 Imperative maps} *)

module MapImp

  use int.Int
  use map.Map
  use export Fmap

  type key

  val eq (x y: key) : bool
    ensures { result <-> x = y }

  type t 'v = abstract {
    mutable to_fmap: fmap key 'v;
  }
  meta coercion function to_fmap

  val create () : t 'v
    ensures { result.to_fmap = empty }

  val mem (k: key) (m: t 'v) : bool
    ensures { result <-> mem k m }

  val is_empty (m: t 'v) : bool
    ensures { result <-> is_empty m }

  val add (k: key) (v: 'v) (m: t 'v) : unit
    writes  { m }
    ensures { m = add k v (old m) }

  val find (k: key) (m: t 'v) : 'v
    requires { mem k m }
    ensures  { result = m[k] }
    ensures  { result = find k m }

  use ocaml.Exceptions

  val find_exn (k: key) (m: t 'v) : 'v
    ensures { S.mem k m.domain }
    ensures { result = m[k] }
    raises  { Not_found ->  not (S.mem k m.domain) }

  val remove (k: key) (m: t 'v) : unit
    writes  { m }
    ensures { m = remove k (old m) }

  val size (m: t 'v) : int
    ensures { result = size m }

  val clear (m: t 'v) : unit
    writes  { m }
    ensures { m = empty }

end

(** {3 Imperative maps of integers} *)

module MapImpInt

  use int.Int

  clone export MapImp with type key = int, val eq = Int.(=), axiom .

end
