source: src/ASM/BitVectorTrie.ma @ 847

include "basics/types.ma".

include "utilities/option.ma".
include "ASM/BitVector.ma".

inductive BitVectorTrie (A: Type[0]): nat → Type[0] ≝
  Leaf: A → BitVectorTrie A O
| Node: ∀n: nat. BitVectorTrie A n → BitVectorTrie A n → BitVectorTrie A (S n)
| Stub: ∀n: nat. BitVectorTrie A n.

axiom fold:
 ∀A, B: Type[0].
 ∀n: nat.
 ∀f: BitVector n → A → B → B.
 ∀t: BitVectorTrie A n.
 ∀b: B.
   B.

let rec lookup_opt (A: Type[0]) (n: nat)
                (b: BitVector n) (t: BitVectorTrie A n) on t
       : option A ≝
 (match t return λx.λ_. BitVector x → option A with
  [ Leaf l ⇒ λ_.Some ? l
  | Node h l r ⇒ λb. lookup_opt A ? (tail … b) (if head' … b then r else l)
  | Stub _ ⇒ λ_.None ?
  ]) b.

let rec lookup (A: Type[0]) (n: nat)
                (b: BitVector n) (t: BitVectorTrie A n) (a: A) on b
       : A ≝
  (match b return λx.λ_. x = n → A with
    [ VEmpty ⇒ 
      (match t return λx.λ_. O = x → A with
        [ Leaf l ⇒ λ_.l
        | Node h l r ⇒ λK.⊥
        | Stub s ⇒ λ_.a
        ])
    | VCons o hd tl ⇒
      match t return λx.λ_. (S o) = x → A with
        [ Leaf l ⇒ λK.⊥
        | Node h l r ⇒
           match hd with
             [ true ⇒ λK. lookup A h (tl⌈o ↦ h⌉) r a
             | false ⇒ λK. lookup A h (tl⌈o ↦ h⌉) l a
             ]
        | Stub s ⇒ λ_. a]
    ]) (refl ? n).
  [1,2:
    destruct
  |*:
    @ injective_S
    //
  ]
qed.

let rec prepare_trie_for_insertion (A: Type[0]) (n: nat) (b: BitVector n) (a:A) on b : BitVectorTrie A n ≝
   match b with
    [ VEmpty ⇒ Leaf A a
    | VCons o hd tl ⇒
      match hd with
        [ true ⇒  Node A o (Stub A o) (prepare_trie_for_insertion A o tl a)
        | false ⇒ Node A o (prepare_trie_for_insertion A o tl a) (Stub A o)
        ]
    ].

let rec insert (A: Type[0]) (n: nat) (b: BitVector n) (a: A) on b: BitVectorTrie A n → BitVectorTrie A n ≝
  (match b with
    [ VEmpty ⇒ λ_. Leaf A a
    | VCons o hd tl ⇒ λt.
          match t return λy.λ_. S o = y → BitVectorTrie A (S o) with
            [ Leaf l ⇒ λprf.⊥
            | Node p l r ⇒ λprf.
               match hd with
                [ true ⇒  Node A o (l⌈p ↦ o⌉) (insert A o tl a (r⌈p ↦ o⌉))
                | false ⇒ Node A o (insert A o tl a (l⌈p ↦ o⌉)) (r⌈p ↦ o⌉)
                ]
            | Stub p ⇒ λprf. (prepare_trie_for_insertion A ? (hd:::tl) a)
            ] (refl ? (S o))
    ]).
  [ destruct
  |*:
    @ injective_S
    //
  ]
qed.

let rec update (A: Type[0]) (n: nat) (b: BitVector n) (a: A) on b: BitVectorTrie A n → option (BitVectorTrie A n) ≝
  (match b with
    [ VEmpty ⇒ λt. match t return λy.λ_. O = y → option (BitVectorTrie A O) with
                   [ Leaf _ ⇒ λ_. Some ? (Leaf A a)
                   | Stub _ ⇒ λ_. None ?
                   | Node _ _ _ ⇒ λprf. ⊥
                   ] (refl ? O)
    | VCons o hd tl ⇒ λt.
          match t return λy.λ_. S o = y → option (BitVectorTrie A (S o)) with
            [ Leaf l ⇒ λprf.⊥
            | Node p l r ⇒ λprf.
               match hd with
                [ true ⇒  option_map ?? (λv. Node A o (l⌈p ↦ o⌉) v) (update A o tl a (r⌈p ↦ o⌉))
                | false ⇒ option_map ?? (λv. Node A o v (r⌈p ↦ o⌉)) (update A o tl a (l⌈p ↦ o⌉))
                ]
            | Stub p ⇒ λprf. None ?
            ] (refl ? (S o))
    ]).
  [ 1,2: destruct
  |*:
    @ injective_S @sym_eq @prf
  ]
qed.

let rec merge (A: Type[0]) (n: nat) (b: BitVectorTrie A n) on b: BitVectorTrie A n → BitVectorTrie A n ≝
  match b return λx. λ_. BitVectorTrie A x → BitVectorTrie A x with
  [ Stub _ ⇒ λc. c
  | Leaf l ⇒ λc. match c with [ Leaf a ⇒ Leaf ? a | _ ⇒ Leaf ? l ]
  | Node p l r ⇒
    λc.
    (match c return λx. λ_. x = (S p) → BitVectorTrie A (S p) with
    [ Node p' l' r' ⇒ λprf. Node ? ? (merge ?? l (l'⌈p' ↦ p⌉)) (merge ?? r (r'⌈p' ↦ p⌉))
    | Stub _ ⇒ λprf. Node ? p l r
    | Leaf _ ⇒ λabsd. ?
    ] (refl ? (S p)))
  ].
  [1:
      destruct(absd)
  |2,3:
      @ injective_S
        assumption
  ]
qed.