include "basics/types.ma".
include "ASM/String.ma".
include "ASM/Arithmetic.ma".
include "common/Errors.ma".

(* identifiers and their generators are tagged to differentiate them, and to
   provide extra type checking. *)

(* in common/PreIdentifiers.ma, via Errors.ma.
inductive identifier (tag:String) : Type[0] ≝
  an_identifier : Word → identifier tag.
*)

record universe (tag:String) : Type[0] ≝
{
  next_identifier : Word;
  counter_overflow: bool
}.

definition new_universe : ∀tag:String. universe tag ≝
  λtag. mk_universe tag (zero ?) false.

(* Fresh identifier generation uses delayed overflow checking.  To make sure
   that the identifiers really were fresh, use the check_universe_ok function
   below afterwards. *)
definition fresh : ∀tag:String. universe tag → identifier tag × (universe tag) ≝
  λtag.
  λuniv: universe tag.
  let 〈gen, carries〉 ≝ add_with_carries ? (next_identifier ? univ) (zero ?) true in
    if get_index_v … carries 0 ? then
      〈an_identifier tag (next_identifier ? univ), mk_universe tag gen true〉
    else
      〈an_identifier tag (next_identifier ? univ), mk_universe tag gen false〉.
  //
qed.

axiom TooManyIdentifiers : String.

definition check_universe_ok : ∀tag:String. universe tag → res unit ≝
  λtag, univ.
  if counter_overflow ? univ
  then Error ? (msg TooManyIdentifiers)
  else OK ? it.

definition eq_identifier : ∀t. identifier t → identifier t → bool ≝
  λt,l,r.
  match l with
  [ an_identifier l' ⇒
    match r with
    [ an_identifier r' ⇒
      eq_bv ? l' r'
    ]
  ].

lemma eq_identifier_elim : ∀P:bool → Type[0]. ∀t,x,y.
  (x = y → P true) → (x ≠ y → P false) →
  P (eq_identifier t x y).
#P #t * #x * #y #T #F
change with (P (eq_bv ???))
@eq_bv_elim [ /2/ | * #H @F % #E destruct /2/ ]
qed.
    
definition word_of_identifier ≝
  λt.
  λl: identifier t.
  match l with    
  [ an_identifier l' ⇒ l'
  ].

definition identifier_eq : ∀tag:String. ∀x,y:identifier tag. (x=y) + (x≠y).
#tag * #x * #y lapply (refl ? (eq_bv ? x y)) cases (eq_bv ? x y) in ⊢ (???% → %)
#E [ % | %2 ]
lapply E @eq_bv_elim
[ #H #_ >H @refl | 2,3: #_ #H destruct | #H #_ % #H' destruct /2/ ]
qed.

definition identifier_of_nat : ∀tag:String. nat → identifier tag ≝
  λtag,n. an_identifier tag (bitvector_of_nat ? n).


(* Maps from identifiers to arbitrary types. *)

include "ASM/BitVectorTrie.ma".

inductive identifier_map (tag:String) (A:Type[0]) : Type[0] ≝
  an_id_map : BitVectorTrie A 16 → identifier_map tag A.
  
definition empty_map : ∀tag:String. ∀A. identifier_map tag A ≝
  λtag,A. an_id_map tag A (Stub A 16).

definition lookup : ∀tag,A. identifier_map tag A → identifier tag → option A ≝
  λtag,A,m,l. lookup_opt A 16 (match l with [ an_identifier l' ⇒ l' ])
                              (match m with [ an_id_map m' ⇒ m' ]).

(* Always adds the identifier to the map. *)
definition add : ∀tag,A. identifier_map tag A → identifier tag → A → identifier_map tag A ≝
λtag,A,m,l,a. an_id_map tag A (insert A 16 (match l with [ an_identifier l' ⇒ l' ]) a
                                           (match m with [ an_id_map m' ⇒ m' ])).

lemma lookup_add_hit : ∀tag,A,m,i,a.
  lookup tag A (add tag A m i a) i = Some ? a.
#tag #A * #m * #i #a
@lookup_opt_insert_hit
qed.

lemma lookup_add_miss : ∀tag,A,m,i,j,a.
  (notb (eq_identifier tag i j)) →
  lookup tag A (add tag A m j a) i = lookup tag A m i.
#tag #A * #m * #i * #j #a
change with (notb (eq_bv ???) → ?)
@lookup_opt_insert_miss
qed.

lemma lookup_add_oblivious : ∀tag,A,m,i,j,a.
  (lookup tag A m i ≠ None ?) →
  lookup tag A (add tag A m j a) i ≠ None ?.
#tag #A #m #i #j #a #H
lapply (lookup_add_miss … m i j a)
@eq_identifier_elim
[ #E #_ >E >lookup_add_hit % #N destruct
| #_ #H' >H' //
] qed.

(* Extract every identifier, value pair from the map. *)
definition elements : ∀tag,A. identifier_map tag A → list (identifier tag × A) ≝
λtag,A,m.
  fold ??? (λl,a,el. 〈an_identifier tag l, a〉::el)
           (match m with [ an_id_map m' ⇒ m' ]) [ ].

axiom MissingId : String.

(* Only updates an existing entry; fails with an error otherwise. *)
definition update : ∀tag,A. identifier_map tag A → identifier tag → A → res (identifier_map tag A) ≝
λtag,A,m,l,a.
  match update A 16 (match l with [ an_identifier l' ⇒ l' ]) a
                    (match m with [ an_id_map m' ⇒ m' ]) with
  [ None ⇒ Error ? ([MSG MissingId; CTX tag l]) (* missing identifier *)
  | Some m' ⇒ OK ? (an_id_map tag A m')
  ].

definition foldi:
  ∀A, B: Type[0].
  ∀tag: String.
  (identifier tag -> A -> B -> B) -> identifier_map tag A -> B -> B ≝
λA,B,tag,f,m,b.
  match m with
  [ an_id_map m' ⇒ fold A B ? (λbv. f (an_identifier ? bv)) m' b ].

(* Sets *)

inductive identifier_set (tag:String) : Type[0] ≝
  an_id_set : BitVectorTrie unit 16 → identifier_set tag.

definition empty_set : ∀tag:String. identifier_set tag ≝
λtag. an_id_set tag (Stub unit 16).

let rec add_set (tag:String) (s:identifier_set tag) (i:identifier tag) on s : identifier_set tag ≝
  an_id_set tag (insert unit 16 (match i with [ an_identifier i' ⇒ i' ])
                             it (match s with [ an_id_set s' ⇒ s' ])).

definition singleton_set : ∀tag:String. identifier tag → identifier_set tag ≝
λtag,i. add_set tag (empty_set tag) i.

let rec mem_set (tag:String) (s:identifier_set tag) (i:identifier tag) on s : bool ≝
  match lookup_opt ? 16 (match i with [ an_identifier i' ⇒ i' ])
                        (match s with [ an_id_set s' ⇒ s' ]) with
  [ None ⇒ false
  | Some _ ⇒ true
  ].

let rec union_set (tag:String) (s:identifier_set tag) (s':identifier_set tag) on s : identifier_set tag ≝
  an_id_set tag (merge unit 16 (match s with [ an_id_set s0 ⇒ s0 ])
                               (match s' with [ an_id_set s1 ⇒ s1 ])).

interpretation "identifier set union" 'union a b = (union_set ? a b).
notation "∅" non associative with precedence 90 for @{ 'empty }.
interpretation "empty identifier set" 'empty = (empty_set ?).
interpretation "singleton identifier set" 'singl a = (add_set ? (empty_set ?) a).
interpretation "identifier set membership" 'mem a b = (mem_set ? b a).

lemma union_empty_l : ∀tag.∀s:identifier_set tag. ∅ ∪ s = s.
#tag * //
qed.

lemma union_empty_r : ∀tag.∀s:identifier_set tag. s ∪ ∅ = s.
#tag * #s cases (BitVectorTrie_Sn … s)
[ * #x * #y #E >E //
| #E >E //
] qed.