include "BitVector.ma".
(*include "Compare.ma".*)
include "Bool.ma".
include "Maybe.ma".

ninductive BitVectorTrie (A: Type[0]): Nat → Type[0] ≝
  Leaf: A → BitVectorTrie A Z
| Node: ∀n: Nat. BitVectorTrie A n → BitVectorTrie A n → BitVectorTrie A (S n)
| Stub: ∀n: Nat. BitVectorTrie A n.

nlet 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
    [ Empty ⇒ 
      (match t return λx.λ_. Z = x → A with
        [ Leaf l ⇒ λ_.l
        | Node h l r ⇒ λK.⊥
        | Stub s ⇒ λ_.a
        ])
    | Cons 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⌈h ↦ o⌉) l a
             | false ⇒ λK. lookup A h (tl⌈h ↦ o⌉) r a
             ]
        | Stub s ⇒ λ_. a
        ]
    ]) (refl ? n).
    ndestruct; //.
nqed.

nlet rec prepare_trie_for_insertion (A: Type[0]) (n: Nat)
                                    (b: BitVector n) (a:A) on b
       : BitVectorTrie A n ≝
   match b with
    [ Empty ⇒ Leaf A a
    | Cons 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)
        ]
    ].

nlet rec insert (A: Type[0]) (n: Nat)
                (b: BitVector n) (a: A) on b:
                 BitVectorTrie A n → BitVectorTrie A n ≝
  (match b with
    [ Empty ⇒ λ_. Leaf A a
    | Cons 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⌈o ↦ p⌉) (insert A o tl a (r⌈o ↦ p⌉))
                | false ⇒ Node A o (insert A o tl a (l⌈o ↦ p⌉)) (r⌈o ↦ p⌉)
                ]
            | Stub p ⇒ λprf. ? (prepare_trie_for_insertion A ? b a)
            ] (refl ? (S o))
    ]).
    ##[ ndestruct;
    ##| ndestruct; @;
    ##| ndestruct; @;
    ##| ndestruct; @;
    ##| ndestruct; @;
    ##| #H; nassumption;
    ##]
nqed.

nlemma insert_lookup_stub:
  ∀A: Type[0].
  ∀n: Nat.
  ∀b: BitVector n.
  ∀t: BitVectorTrie A n.
  ∀a, c: A.
    (lookup A n b (insert A n b a (Stub A n)) a) = a.
  #A n b t a c.
  nelim b.
  //.
  #N H V H2.
  nnormalize.
  @.
nqed.

(*
nlemma test:
  ∀n: Nat.
  ∀b: BitVector n.
    length n b = n.
  #n b.
  nelim b.
  //.
  #N H V IH.
  ncases H.

nlemma insert_lookup_leaf:
  ∀A: Type[0].
  ∀n: Nat.
  ∀b: BitVector n.
  ∀a, c: A.
  ∀t: BitVectorTrie A n.
    lookup A ? b (insert A ? b a t) c = a.
  #A n b a c t.
  nelim b.
  nnormalize.
  @.
  #N H V IH.
*)