include "basics/types.ma".
include "basics/bool.ma".
include "basics/lists/list.ma".

include "arithmetics/nat.ma".

include "utilities/adt/equal.ma".
include "utilities/adt/ordering.ma".

include "ASM/Util.ma".

inductive table (dom_type: Type[0]) (rng_type: Type[0]): Type[0] ≝
  | empty: table dom_type rng_type
  | node : nat → dom_type → rng_type → table dom_type rng_type → table dom_type rng_type → table dom_type rng_type.

definition tbl_height ≝
  λdom_type : Type[0].
  λrng_type : Type[0].
  λthe_table: table dom_type rng_type.
  match the_table with
  [ empty ⇒ 0
  | node sz key val l r ⇒ sz
  ].

let rec tbl_size
  (dom_type : Type[0]) (rng_type : Type[0]) (the_table: table dom_type rng_type)
    on the_table: nat ≝
  match the_table with
  [ empty ⇒ 0
  | node sz key val l r ⇒ S (tbl_size … l) + (tbl_size … r)
  ].
  
definition tbl_empty ≝
  λdom_type: Type[0].
  λrng_type: Type[0].
    empty dom_type rng_type.

definition tbl_is_empty ≝
  λdom_type: Type[0].
  λrng_type: Type[0].
  λthe_table: table dom_type rng_type.
  match the_table with
  [ empty ⇒ true
  | _     ⇒ false
  ].

definition tbl_is_empty_p ≝
  λdom_type: Type[0].
  λrng_type: Type[0].
  λthe_table: table dom_type rng_type.
  match the_table with
  [ empty ⇒ True
  | _     ⇒ False
  ].

let rec tbl_to_list
  (dom_type: Type[0]) (rng_type: Type[0]) (the_table: table dom_type rng_type)
    on the_table: list (dom_type × rng_type) ≝
  match the_table with
  [ empty ⇒ [ ]
  | node sz key val l r ⇒ 〈key, val〉 :: tbl_to_list … l @ tbl_to_list … r
  ].

let rec tbl_min_binding
  (key_type: Type[0]) (rng_type: Type[0]) (the_table: table key_type rng_type)
    (proof: ¬ (tbl_is_empty_p … the_table))
      on the_table: key_type × rng_type ≝
  match the_table return λx. ¬ (tbl_is_empty_p … x) → key_type × rng_type with
  [ empty               ⇒ λabsurd. ?
  | node sz key val l r ⇒ λproof.
    match l return λx. x = l → key_type × rng_type with
    [ empty ⇒ λproof. 〈key, val〉
    | _     ⇒ λproof. tbl_min_binding key_type rng_type l ?
    ] (refl … l)
  ] proof.
  [1: normalize in absurd;
      elim absurd;
      #ABSURD
      @⊥ @ABSURD @I 
  |2: <proof
      normalize
      /2/
  ]
qed.

definition tbl_create ≝
  λkey_type: Type[0].
  λrng_type: Type[0].
  λleft    : table key_type rng_type.
  λkey     : key_type.
  λval     : rng_type.
  λright   : table key_type rng_type.
  let hl ≝ tbl_height … left in
  let hr ≝ tbl_height … right in
  let height ≝ if gtb hl hr then hl + 1 else hr + 1 in
    node key_type rng_type height key val left right.
      
let rec tbl_balance
  (key_type: Type[0]) (rng_type: Type[0]) (ord: key_type → key_type → order)
    (left: table key_type rng_type) (key: key_type) (val: rng_type)
      (right: table key_type rng_type) (left_proof: ¬ (tbl_is_empty_p … left))
        (right_proof: ¬ (tbl_is_empty_p … right))
          on left: table key_type rng_type ≝
  match gtb (tbl_height … left) ((tbl_height … right) + 2) with
  [ true  ⇒
    match left return λx. ¬ (tbl_is_empty_p … x) → table key_type rng_type with
    [ empty              ⇒ λabsurd. ?
    | node _ lv ld ll lr ⇒ λnon_empty_proof.
      if gtb (tbl_height … ll) (tbl_height … lr) then
        tbl_create … ll lv ld (tbl_create … lr key val right)
      else
        match lr return λx. ¬ (tbl_is_empty_p … x) → table key_type rng_type with
        [ empty ⇒ λabsurd. ?
        | node _ lrv lrd lrl lrr ⇒ λnon_empty_proof.
            tbl_create … (tbl_create … ll lv ld lrl) lrv lrd (tbl_create … lrr key val right)
        ] ?
    ] left_proof
  | false ⇒ ?
  ].
  [1: normalize in absurd;
      elim absurd
      #ABSURD
      @⊥ @ABSURD @I
  |3: normalize in absurd;
      elim absurd
      #ABSURD
      @⊥ @ABSURD @I
  |*: cases daemon
  ]
qed.
  
let rec tbl_remove_min_binding
  (key_type: Type[0]) (rng_type: Type[0]) (ord: key_type → key_type → order)
    (the_table: table key_type rng_type) (proof: ¬ (tbl_is_empty_p … the_table))
      on the_table: table key_type rng_type ≝
  match the_table return λx. ¬ (tbl_is_empty_p … x) → table key_type rng_type with
  [ empty               ⇒ λabsurd. ?
  | node sz key val l r ⇒ λproof.
    match l return λx. x = l → table key_type rng_type with
    [ empty ⇒ λproof. r
    | _     ⇒ λproof. tbl_balance … ord (tbl_remove_min_binding … ord l ?) key val r
    ] (refl … l)
  ] proof.
  [1: normalize in absurd;
      elim absurd;
      #ABSURD
      @⊥ @ABSURD @I
  |2: <proof
      normalize
      @nmk //
  ]
qed.
    
let rec tbl_insert
  (key_type: Type[0]) (rng_type: Type[0]) (ord: key_type → key_type → order)
    (key: key_type) (val: rng_type) (the_table: table key_type rng_type)
      on the_table: table key_type rng_type ≝
  match the_table with
  [ empty                 ⇒ node … 1 key val (empty …) (empty …)
  | node sz key' val' l r ⇒
    match ord key key' with
    [ order_lt ⇒ tbl_balance … ord (tbl_insert … ord key val l) key' val' r
    | order_eq ⇒ node … sz key val l r
    | order_gt ⇒ tbl_balance … ord l key' val' (tbl_insert … ord key val r)
    ]
  ].
  
let rec tbl_merge
  (dom_type: Type[0]) (rng_type: Type[0]) (ord: dom_type → dom_type → order)
    (left: table dom_type rng_type) (right: table dom_type rng_type)
      on left: table dom_type rng_type ≝
  match left with
  [ empty ⇒ right
  | _     ⇒
    match right return λx. x = right → table dom_type rng_type with
    [ empty               ⇒ λid_proof. left
    | node sz key val l r ⇒ λid_proof.
      let 〈x, d〉 ≝ tbl_min_binding dom_type rng_type right ? in
        tbl_balance … ord left x d (tbl_remove_min_binding … ord right ?)
    ] (refl … right)
  ].
  <id_proof
  normalize /2/
qed.

let rec tbl_remove
  (key_type: Type[0]) (rng_type: Type[0]) (ord: key_type → key_type → order)
    (key: key_type) (the_table: table key_type rng_type)
      on the_table: table key_type rng_type ≝
  match the_table with
  [ empty                 ⇒ empty …
  | node sz key' val' l r ⇒
    match ord key key' with
    [ order_lt ⇒ tbl_balance … ord (tbl_remove … ord key l) key' val' r
    | order_eq ⇒ tbl_merge … ord l r
    | order_gt ⇒ tbl_balance … ord l key' val' (tbl_remove … ord key r)
    ]
  ].
  
let rec tbl_find
  (dom_type: Type[0]) (rng_type: Type[0]) (ord: dom_type → dom_type → order)
    (to_find: dom_type) (the_table: table dom_type rng_type) 
      on the_table: option rng_type ≝
  match the_table with
  [ empty               ⇒ None …
  | node sz key val l r ⇒
    match ord to_find key with
    [ order_lt ⇒ tbl_find … ord to_find l
    | order_eq ⇒ Some … val
    | order_gt ⇒ tbl_find … ord to_find r
    ]
  ].

let rec tbl_update
  (dom_type: Type[0]) (rng_type: Type[0]) (ord: dom_type → dom_type → order)
    (to_find: dom_type) (update: rng_type → rng_type)
      (the_table: table dom_type rng_type)
        on the_table: table dom_type rng_type ≝
  match the_table with
  [ empty               ⇒ empty …
  | node sz key val l r ⇒
    match ord to_find key with
    [ order_lt ⇒ tbl_update … ord to_find update l
    | order_eq ⇒ node … sz key (update val) l r
    | order_gt ⇒ tbl_update … ord to_find update r
    ]
  ].

let rec tbl_forall
  (dom_type: Type[0]) (rng_type: Type[0]) (f: rng_type → bool)
    (the_table: table dom_type rng_type)
      on the_table: bool ≝
  match the_table with
  [ empty               ⇒ true
  | node sz key val l r ⇒ f val ∧ tbl_forall … f l ∧ tbl_forall … f r
  ].
  
let rec tbl_exists
  (dom_type: Type[0]) (rng_type: Type[0]) (f: rng_type → bool)
    (the_table: table dom_type rng_type)
      on the_table: bool ≝
  match the_table with
  [ empty               ⇒ false
  | node sz key val l r ⇒ f val ∨ tbl_exists … f l ∨ tbl_exists … f r
  ].
      
axiom tbl_equal:
  ∀key_type, rng_type.
    (key_type → key_type → order) → (rng_type → rng_type → order) →
      table key_type rng_type → table key_type rng_type → bool.

let rec tbl_mapi
  (dom_type: Type[0]) (rng_type: Type[0]) (a: Type[0]) (f: dom_type → rng_type → a)
    (the_table: table dom_type rng_type)
      on the_table: table dom_type a ≝
  match the_table with
  [ empty               ⇒ empty …
  | node sz key val l r ⇒ node  … sz key (f key val) (tbl_mapi … f l) (tbl_mapi … f r)
  ].

let rec tbl_map
  (dom_type: Type[0]) (rng_type: Type[0]) (a: Type[0]) (f: rng_type → a)
    (the_table: table dom_type rng_type)
      on the_table: table dom_type a ≝
  match the_table with
  [ empty               ⇒ empty …
  | node sz key val l r ⇒ node  … sz key (f val) (tbl_map … f l) (tbl_map … f r)
  ].

let rec tbl_fold
  (dom_type: Type[0]) (a: Type[0]) (b: Type[0]) (f: dom_type → a → b → b)
    (the_table: table dom_type a) (seed: b)
      on the_table: b ≝
  match the_table with
  [ empty               ⇒ seed
  | node sz key val l r ⇒ tbl_fold … f r (f key val (tbl_fold … f l seed))
  ].

definition tbl_singleton ≝
  λkey_type : Type[0].
  λrng_type : Type[0].
  λkey      : key_type.
  λrng      : rng_type.
    node key_type rng_type 1 key rng (empty …) (empty …).

definition tbl_from_list ≝
  λkey_type : Type[0].
  λrng_type : Type[0].
  λord      : key_type → key_type → order.
  λthe_list : list (key_type × rng_type).
    foldr … (λkey_rng.
      let 〈key, rng〉 ≝ key_rng in
        tbl_insert … ord key rng) (tbl_empty …) the_list.

definition tbl_filter ≝
  λdom_type : Type[0].
  λrng_type : Type[0].
  λord      : dom_type → dom_type → order.
  λf        : rng_type → bool.
  λthe_table: table dom_type rng_type.
  let as_list  ≝ tbl_to_list … the_table in
  let filtered ≝ filter … (λx. f (\snd x)) as_list in
    tbl_from_list … ord filtered.