(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Vector.ma: Fixed length polymorphic vectors, and routine operations on     *)
(*            them.                                                           *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

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

include "ASM/Util.ma".

include "arithmetics/nat.ma".

include "utilities/extranat.ma".
include "utilities/oldlib/eq.ma".

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* The datatype.                                                              *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

inductive Vector (A: Type[0]): nat → Type[0] ≝
  VEmpty: Vector A O
| VCons: ∀n: nat. A → Vector A n → Vector A (S n).

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Syntax.                                                                    *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

notation "hvbox(hd break ::: tl)"
  right associative with precedence 52
  for @{ 'vcons $hd $tl }.

notation "[[ list0 x sep ; ]]"
  non associative with precedence 90
  for ${fold right @'vnil rec acc @{'vcons $x $acc}}.

interpretation "Vector vnil" 'vnil = (VEmpty ?).
interpretation "Vector vcons" 'vcons hd tl = (VCons ? ? hd tl).

notation "hvbox(l break !!! break n)"
  non associative with precedence 90
  for @{ 'get_index_v $l $n }.

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Lookup.                                                                    *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

let rec get_index_v (A: Type[0]) (n: nat)
                   (v: Vector A n) (m: nat) (lt: m < n) on m: A ≝
  (match m with
    [ O ⇒
      match v return λx.λ_. O < x → A with
        [ VEmpty ⇒ λabsd1: O < O. ?
        | VCons p hd tl ⇒ λprf1: O < S p. hd
        ]
    | S o ⇒
      (match v return λx.λ_. S o < x → A with
        [ VEmpty ⇒ λprf: S o < O. ?
        | VCons p hd tl ⇒ λprf: S o < S p. get_index_v A p tl o ?
        ])
    ]) lt.
    [ cases (not_le_Sn_O O)
      normalize in absd1
      # H
      cases (H absd1)
    | cases (not_le_Sn_O (S o))
      normalize in prf
      # H
      cases (H prf)
    | normalize
      normalize in prf
      @ le_S_S_to_le
      assumption
    ]
qed.

definition get_index' ≝
  λA: Type[0].
  λn, m: nat.
  λb: Vector A (S (n + m)).
    get_index_v A (S (n + m)) b n ?.
  normalize
  //
qed.

let rec get_index_weak_v (A: Type[0]) (n: nat)
                         (v: Vector A n) (m: nat) on m ≝
  match m with
    [ O ⇒
      match v with
        [ VEmpty ⇒ None A
        | VCons p hd tl ⇒ Some A hd
        ]
    | S o ⇒
      match v with
        [ VEmpty ⇒ None A
        | VCons p hd tl ⇒ get_index_weak_v A p tl o
        ]
    ].
    
interpretation "Vector get_index" 'get_index_v v n = (get_index_v ? ? v n).

let rec set_index (A: Type[0]) (n: nat) (v: Vector A n) (m: nat) (a: A) (lt: m < n) on m: Vector A n ≝
  (match m with
    [ O ⇒
      match v return λx.λ_. O < x → Vector A x with
        [ VEmpty ⇒ λabsd1: O < O. [[ ]]
        | VCons p hd tl ⇒ λprf1: O < S p. (a ::: tl)
        ]
    | S o ⇒
      (match v return λx.λ_. S o < x → Vector A x with
        [ VEmpty ⇒ λprf: S o < O. [[ ]]
        | VCons p hd tl ⇒ λprf: S o < S p. hd ::: (set_index A p tl o a ?)
        ])
    ]) lt.
    normalize in prf ⊢ %;
    /2/;
qed.

let rec set_index_weak (A: Type[0]) (n: nat)
                       (v: Vector A n) (m: nat) (a: A) on m ≝
  match m with
    [ O ⇒
      match v with
        [ VEmpty ⇒ None (Vector A n)
        | VCons o hd tl ⇒ Some (Vector A n) (? (VCons A o a tl))
        ]
    | S o ⇒
      match v with
        [ VEmpty ⇒ None (Vector A n)
        | VCons p hd tl ⇒
            let settail ≝ set_index_weak A p tl o a in
              match settail with
                [ None ⇒ None (Vector A n)
                | Some j ⇒ Some (Vector A n) (? (VCons A p hd j))
                ]
        ]
    ].
    //.
qed.

let rec drop (A: Type[0]) (n: nat)
             (v: Vector A n) (m: nat) on m ≝
  match m with
    [ O ⇒ Some (Vector A n) v
    | S o ⇒
      match v with
        [ VEmpty ⇒ None (Vector A n)
        | VCons p hd tl ⇒ ? (drop A p tl o)
        ]
    ].
    //.
qed.

definition head' : ∀A:Type[0]. ∀n:nat. Vector A (S n) → A ≝
λA,n,v. match v return λx.λ_. match x with [ O ⇒ True | _ ⇒ A ] with
[ VEmpty ⇒ I | VCons _ hd _ ⇒ hd ].

definition tail : ∀A:Type[0]. ∀n:nat. Vector A (S n) → Vector A n ≝
λA,n,v. match v return λx.λ_. match x with [ O ⇒ True | S m ⇒ Vector A m ] with
[ VEmpty ⇒ I | VCons m hd tl ⇒ tl ].

let rec split' (A: Type[0]) (m, n: nat) on m: Vector A (plus m n) → (Vector A m) × (Vector A n) ≝
 match m return λm. Vector A (plus m n) → (Vector A m) × (Vector A n) with
  [ O ⇒ λv. 〈[[ ]], v〉
  | S m' ⇒ λv. let 〈l,r〉 ≝ split' A m' n (tail ?? v) in 〈head' ?? v:::l, r〉
  ].
(* Prevent undesirable unfolding. *)
let rec split (A: Type[0]) (m, n: nat) (v:Vector A (plus m n)) on v : (Vector A m) × (Vector A n) ≝
 split' A m n v.


definition head: ∀A: Type[0]. ∀n: nat. Vector A (S n) → A × (Vector A n) ≝
  λA: Type[0].
  λn: nat.
  λv: Vector A (S n).
  match v return λl. λ_: Vector A l. l = S n → A × (Vector A n) with
  [ VEmpty ⇒ λK. ⊥
  | VCons o he tl ⇒ λK. 〈he, (tl⌈Vector A o ↦ Vector A n⌉)〉
  ] (? : S ? = S ?).
  //
  [ destruct
  | lapply (injective_S … K)
    //
  ]
qed.

definition from_singl: ∀A:Type[0]. Vector A (S O) → A ≝
 λA: Type[0].
 λv: Vector A (S 0).
   fst … (head … v).
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Folds and builds.                                                          *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
    
let rec fold_right (A: Type[0]) (B: Type[0]) (n: nat)
                    (f: A → B → B) (x: B) (v: Vector A n) on v ≝
  match v with
    [ VEmpty ⇒ x
    | VCons n hd tl ⇒ f hd (fold_right A B n f x tl)
    ].

let rec fold_right_i (A: Type[0]) (B: nat → Type[0]) (n: nat)
                     (f: ∀n. A → B n → B (S n)) (x: B 0) (v: Vector A n) on v ≝
  match v with
    [ VEmpty ⇒ x
    | VCons n hd tl ⇒ f ? hd (fold_right_i A B n f x tl)
    ].

let rec fold_right2_i (A: Type[0]) (B: Type[0]) (C: nat → Type[0])
                      (f: ∀N. A → B → C N → C (S N)) (c: C O) (n: nat)
                      (v: Vector A n) (q: Vector B n) on v : C n ≝
  (match v return λx.λ_. x = n → C n with
    [ VEmpty ⇒
      match q return λx.λ_. O = x → C x with
        [ VEmpty ⇒ λprf: O = O. c
        | VCons o hd tl ⇒ λabsd. ⊥
        ]
    | VCons o hd tl ⇒ 
      match q return λx.λ_. S o = x → C x with
        [ VEmpty ⇒ λabsd: S o = O. ⊥
        | VCons p hd' tl' ⇒ λprf: S o = S p.
           (f ? hd hd' (fold_right2_i A B C f c ? tl (tl'⌈Vector B p ↦ Vector B o⌉)))⌈C (S o) ↦ C (S p)⌉
        ]
    ]) (refl ? n).
  [1,2:
    destruct
  |3,4:
    lapply (injective_S … prf)
    //
  ]
qed.
  
let rec fold_left (A: Type[0]) (B: Type[0]) (n: nat)
                    (f: A → B → A) (x: A) (v: Vector B n) on v ≝
  match v with
    [ VEmpty ⇒ x
    | VCons n hd tl ⇒ fold_left A B n f (f x hd) tl
    ].
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Maps and zips.                                                             *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

let rec map (A: Type[0]) (B: Type[0]) (n: nat)
             (f: A → B) (v: Vector A n) on v ≝
  match v with
    [ VEmpty ⇒ [[ ]]
    | VCons n hd tl ⇒ (f hd) ::: (map A B n f tl)
    ].

let rec zip_with (A: Type[0]) (B: Type[0]) (C: Type[0]) (n: nat)
             (f: A → B → C) (v: Vector A n) (q: Vector B n) on v ≝
  (match v return (λx.λr. x = n → Vector C x) with
    [ VEmpty ⇒ λ_. [[ ]]
    | VCons n hd tl ⇒
      match q return (λy.λr. S n = y → Vector C (S n)) with
        [ VEmpty ⇒ ?
        | VCons m hd' tl' ⇒
            λe: S n = S m.
              (f hd hd') ::: (zip_with A B C n f tl ?)
        ]
    ])
    (refl ? n).
  [ #e
    destruct(e);
  | lapply (injective_S … e)
    # H
    > H
    @ tl'
  ]
qed.

definition zip ≝
  λA, B: Type[0].
  λn: nat.
  λv: Vector A n.
  λq: Vector B n.
    zip_with A B (A × B) n (pair A B) v q.

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Building vectors from scratch                                              *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

let rec replicate (A: Type[0]) (n: nat) (h: A) on n ≝
  match n return λn. Vector A n with
    [ O ⇒ [[ ]]
    | S m ⇒ h ::: (replicate A m h)
    ].

(* DPM: fixme.  Weird matita bug in base case. *)
let rec append (A: Type[0]) (n: nat) (m: nat)
                (v: Vector A n) (q: Vector A m) on v ≝
  match v return (λn.λv. Vector A (n + m)) with
    [ VEmpty ⇒ (? q)
    | VCons o hd tl ⇒ hd ::: (append A o m tl q)
    ].
    # H
    assumption
qed.
    
notation "hvbox(l break @@ r)"
  right associative with precedence 47
  for @{ 'vappend $l $r }.
    
interpretation "Vector append" 'vappend v1 v2 = (append ??? v1 v2).
    
let rec scan_left (A: Type[0]) (B: Type[0]) (n: nat)
                   (f: A → B → A) (a: A) (v: Vector B n) on v ≝
  a :::
    (match v with
       [ VEmpty ⇒ VEmpty A
       | VCons o hd tl ⇒ scan_left A B o f (f a hd) tl
       ]).

let rec scan_right (A: Type[0]) (B: Type[0]) (n: nat)
                    (f: A → B → A) (b: B) (v: Vector A n) on v ≝
  match v with
    [ VEmpty ⇒ ?
    | VCons o hd tl ⇒ f hd b :: (scan_right A B o f b tl)
    ].
    //
qed.
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Other manipulations.                                                       *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

(* At some points matita will attempt to reduce reverse with a known vector,
   which reduces the equality proof for the cast.  Normalising this proof needs
   to be fast enough to keep matita usable, so use plus_n_Sm_fast. *)

let rec revapp (A: Type[0]) (n: nat) (m:nat)
                (v: Vector A n) (acc: Vector A m) on v : Vector A (n + m) ≝
  match v return λn'.λ_. Vector A (n' + m) with
    [ VEmpty ⇒ acc
    | VCons o hd tl ⇒ (revapp ??? tl (hd:::acc))⌈Vector A (o+S m) ↦ Vector A (S o + m)⌉
    ].
> plus_n_Sm_fast @refl qed.

let rec reverse (A: Type[0]) (n: nat) (v: Vector A n) on v : Vector A n ≝
  (revapp A n 0 v [[ ]])⌈Vector A (n+0) ↦ Vector A n⌉.
< plus_n_O @refl qed.

let rec pad_vector (A:Type[0]) (a:A) (n,m:nat) (v:Vector A m) on n : Vector A (n+m) ≝
match n return λn.Vector A (n+m) with
[ O ⇒ v
| S n' ⇒ a:::(pad_vector A a n' m v)
].

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Conversions to and from lists.                                             *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

let rec list_of_vector (A: Type[0]) (n: nat)
                        (v: Vector A n) on v ≝
  match v return λn.λv. list A with
    [ VEmpty ⇒ []
    | VCons o hd tl ⇒ hd :: (list_of_vector A o tl)
    ].

let rec vector_of_list (A: Type[0]) (l: list A) on l ≝
  match l return λl. Vector A (length A l) with
    [ nil ⇒ ?
    | cons hd tl ⇒ hd ::: (vector_of_list A tl)
    ].
    normalize
    @ VEmpty
qed.

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)   
(* Rotates and shifts.                                                        *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
    
let rec rotate_left (A: Type[0]) (n: nat)
                     (m: nat) (v: Vector A n) on m: Vector A n ≝
  match m with
    [ O ⇒ v
    | S o ⇒
        match v with
          [ VEmpty ⇒ [[ ]]
          | VCons p hd tl ⇒
             rotate_left A (S p) o ((append A p ? tl [[hd]])⌈Vector A (p + S O) ↦ Vector A (S p)⌉)
          ]
    ].
  //
qed.

definition rotate_right ≝
  λA: Type[0].
  λn, m: nat.
  λv: Vector A n.
    reverse A n (rotate_left A n m (reverse A n v)).

definition shift_left_1 ≝
  λA: Type[0].
  λn: nat.
  λv: Vector A (S n).
  λa: A.
   match v return λy.λ_. y = S n → Vector A y with
     [ VEmpty ⇒ λH.⊥
     | VCons o hd tl ⇒ λH.reverse … (a::: reverse … tl)
     ] (refl ? (S n)).
 destruct.
qed.


(* XXX this is horrible - but useful to ensure that we can normalise in the proof assistant. *)
definition switch_bv_plus : ∀A:Type[0]. ∀n,m. Vector A (n+m) → Vector A (m+n) ≝
λA,n,m. match commutative_plus_faster n m return λx.λ_.Vector A (n+m) → Vector A x with [ refl ⇒ λi.i ].

definition shift_right_1 ≝
  λA: Type[0].
  λn: nat.
  λv: Vector A (S n).
  λa: A.
    let 〈v',dropped〉 ≝ split ? n 1 (switch_bv_plus ? 1 n v) in a:::v'.
(*    reverse … (shift_left_1 … (reverse … v) a).*)

definition shift_left : ∀A:Type[0]. ∀n,m:nat. Vector A n → A → Vector A n ≝
  λA: Type[0].
  λn, m: nat.
    match nat_compare n m return λx,y.λ_. Vector A x → A → Vector A x with
    [ nat_lt _ _ ⇒ λv,a. replicate … a
    | nat_eq _   ⇒ λv,a. replicate … a
    | nat_gt d m ⇒ λv,a. let 〈v0,v'〉 ≝ split … v in switch_bv_plus … (v' @@ (replicate … a))
    ].

(*    iterate … (λx. shift_left_1 … x a) v m.*)
    
definition shift_right ≝
  λA: Type[0].
  λn, m: nat.
  λv: Vector A (S n).
  λa: A.
    iterate … (λx. shift_right_1 … x a) v m.

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Decidable equality.                                                        *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

let rec eq_v (A: Type[0]) (n: nat) (f: A → A → bool) (b: Vector A n) (c: Vector A n) on b : bool ≝
  (match b return λx.λ_. Vector A x → bool with
   [ VEmpty ⇒ λc.
       match c return λx.λ_. match x return λ_.Type[0] with [ O ⇒ bool | _ ⇒ True ] with
       [ VEmpty ⇒ true
       | VCons p hd tl ⇒ I
       ]
   | VCons m hd tl ⇒ λc. andb (f hd (head' A m c)) (eq_v A m f tl (tail A m c))
   ]
  ) c.

lemma vector_inv_n: ∀A,n. ∀P:Vector A n → Type[0]. ∀v:Vector A n.
  match n return λn'. (Vector A n' → Type[0]) → Vector A n' → Type[0] with
  [ O ⇒ λP.λv.P [[ ]] → P v
  | S m ⇒ λP.λv.(∀h,t. P (VCons A m h t)) → P v
  ] P v.
#A #n #P #v generalize in match P cases v normalize //
qed.

lemma eq_v_elim: ∀P:bool → Type[0]. ∀A,f.
  (∀Q:bool → Type[0]. ∀a,b. (a = b → Q true) → (a ≠ b → Q false) → Q (f a b)) →
  ∀n,x,y.
  (x = y → P true) →
  (x ≠ y → P false) →
  P (eq_v A n f x y).
#P #A #f #f_elim #n #x elim x
[ #y @(vector_inv_n … y)
  normalize /2/
| #m #h #t #IH #y @(vector_inv_n … y)
  #h' #t' #Ht #Hf whd in ⊢ (?%)
  @(f_elim ? h h') #Eh
  [ @IH [ #Et @Ht >Eh >Et @refl | #NEt @Hf % #E' destruct (E') elim NEt /2/ ]
  | @Hf % #E' destruct (E') elim Eh /2/
  ]
] qed.

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Subvectors.                                                                *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *) 

definition mem ≝
 λA: Type[0].
 λeq_a : A → A → bool.
 λn: nat.
 λl: Vector A n.
 λx: A.
  fold_right … (λy,v. (eq_a x y) ∨ v) false l.


definition subvector_with ≝
  λA: Type[0].
  λn: nat.
  λm: nat.
  λf: A → A → bool.
  λv: Vector A n.
  λq: Vector A m.
    fold_right ? ? ? (λx, v. (mem ? f ? q x) ∧ v) true v.
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Lemmas.                                                                    *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)    
    
lemma map_fusion:
  ∀A, B, C: Type[0].
  ∀n: nat.
  ∀v: Vector A n.
  ∀f: A → B.
  ∀g: B → C.
    map B C n g (map A B n f v) = map A C n (λx. g (f x)) v.
  #A #B #C #n #v #f #g
  elim v
  [ normalize
    %
  | #N #H #V #H2
    normalize
    > H2
    %
  ]
qed.