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

include "Nat.ma".
include "List.ma".

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

ninductive Vector (A: Type[0]): Nat → Type[0] ≝
  VEmpty: Vector A Z
| 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.                                                                    *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

nlet rec get_index_v (A: Type[0]) (n: Nat)
                   (v: Vector A n) (m: Nat) (lt: m < n) on m: A ≝
  (match m with
    [ Z ⇒
      match v return λx.λ_. Z < x → A with
        [ VEmpty ⇒ λabsd1: Z < Z. ?
        | VCons p hd tl ⇒ λprf1: Z < S p. hd
        ]
    | S o ⇒
      (match v return λx.λ_. S o < x → A with
        [ VEmpty ⇒ λprf: S o < Z. ?
        | VCons p hd tl ⇒ λprf: S o < S p. get_index_v A p tl o ?
        ])
    ]) lt.
    ##[ ncases (nothing_less_than_Z Z); #K; ncases (K absd1)
    ##| ncases (nothing_less_than_Z (S o)); #K; ncases (K prf)
    ##| napply succ_less_than_injective; nassumption
    ##]
nqed.

ndefinition get_index' ≝
  λA: Type[0].
  λn, m: Nat.
  λb: Vector A (S (n + m)).
    get_index_v A (S (n + m)) b n ?.
  nnormalize;
  napply less_than_or_equal_monotone;
  napply less_than_or_equal_plus;
nqed.

nlet rec get_index_weak_v (A: Type[0]) (n: Nat)
                          (v: Vector A n) (m: Nat) on m ≝
  match m with
    [ Z ⇒
      match v with
        [ VEmpty ⇒ Nothing A
        | VCons p hd tl ⇒ Just A hd
        ]
    | S o ⇒
      match v with
        [ VEmpty ⇒ Nothing 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).

nlet 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
    [ Z ⇒
      match v return λx.λ_. Z < x → Vector A x with
        [ VEmpty ⇒ λabsd1: Z < Z. [[ ]]
        | VCons p hd tl ⇒ λprf1: Z < S p. (a ::: tl)
        ]
    | S o ⇒
      (match v return λx.λ_. S o < x → Vector A x with
        [ VEmpty ⇒ λprf: S o < Z. [[ ]]
        | VCons p hd tl ⇒ λprf: S o < S p. hd ::: (set_index A p tl o a ?)
        ])
    ]) lt.
    napply succ_less_than_injective.
    nassumption.
nqed.
    
nlet rec set_index_weak (A: Type[0]) (n: Nat)
                        (v: Vector A n) (m: Nat) (a: A) on m ≝
  match m with
    [ Z ⇒
      match v with
        [ VEmpty ⇒ Nothing (Vector A n)
        | VCons o hd tl ⇒ Just (Vector A n) (? (VCons A o a tl))
        ]
    | S o ⇒
      match v with
        [ VEmpty ⇒ Nothing (Vector A n)
        | VCons p hd tl ⇒
            let settail ≝ set_index_weak A p tl o a in
              match settail with
                [ Nothing ⇒ Nothing (Vector A n)
                | Just j ⇒ Just (Vector A n) (? (VCons A p hd j))
                ]
        ]
    ].
    //.
nqed.

nlet rec drop (A: Type[0]) (n: Nat)
              (v: Vector A n) (m: Nat) on m ≝
  match m with
    [ Z ⇒ Just (Vector A n) v
    | S o ⇒
      match v with
        [ VEmpty ⇒ Nothing (Vector A n)
        | VCons p hd tl ⇒ ? (drop A p tl o)
        ]
    ].
    //.
nqed.

nlet rec split (A: Type[0]) (m,n: Nat) on m
             : Vector A (m+n) → (Vector A m) × (Vector A n)
≝
 match m return λm. Vector A (m+n) → (Vector A m) × (Vector A n) with
  [ Z ⇒ λv.〈[[ ]], v〉
  | S m' ⇒ λv.
     match v return λl.λ_:Vector A l.l = S (m' + n) → (Vector A (S m')) × (Vector A n) with
      [ VEmpty ⇒ λK.⊥
      | VCons o he tl ⇒ λK.
         match split A m' n (tl⌈Vector A o↦Vector A (m'+n)⌉) with
          [ mk_Cartesian v1 v2 ⇒ 〈he:::v1, v2〉 ]] (?: (S (m' + n)) = S (m' + n))].
// [ ndestruct | nlapply (S_inj … K); //]
nqed.

ndefinition head: ∀A:Type[0]. ∀n. Vector A (S n) → A × (Vector A n)
≝ λA,n,v.
 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 ?).
// [ ndestruct | nlapply (S_inj … K); //]
nqed.

ndefinition from_singl: ∀A:Type[0]. Vector A (S Z) → A ≝
 λA,v. first … (head … v).
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Folds and builds.                                                          *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
    
nlet 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)
    ].

nlet 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 Z) (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.λ_. Z = x → C x with
        [ VEmpty ⇒ λprf: Z = Z. c
        | VCons o hd tl ⇒ λabsd. ⊥
        ]
    | VCons o hd tl ⇒ 
      match q return λx.λ_. S o = x → C x with
        [ VEmpty ⇒ λabsd: S o = Z. ⊥
        | 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: ndestruct | ##3,4: nlapply (S_inj … prf); // ]
nqed.
  
nlet 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 ⇒ f (fold_left A B n f x tl) hd
    ].
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Maps and zips.                                                             *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

nlet 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)
    ].

nlet 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;
          ndestruct(e);
          ##
        | nlapply (S_inj … e); #H; nrewrite > H;
          napply tl'
          ##
        ]
nqed.

ndefinition zip ≝
  λA, B: Type[0].
  λn: Nat.
  λv: Vector A n.
  λq: Vector B n.
    zip_with A B (Cartesian A B) n (mk_Cartesian A B) v q.

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

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

nlet 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)
    ].
    
notation "hvbox(l break @@ r)"
  right associative with precedence 47
  for @{ 'vappend $l $r }.
    
interpretation "Vector append" 'vappend v1 v2 = (append ??? v1 v2).
    
nlet 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
       ]).

nlet 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)
    ].
    //.
nqed.
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Other manipulations.                                                       *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

nlet rec reverse (A: Type[0]) (n: Nat)
                 (v: Vector A n) on v ≝
  match v return (λm.λv. Vector A m) with
    [ VEmpty ⇒ [[ ]]
    | VCons o hd tl ⇒ ? (append A o ? (reverse A o tl) [[hd]])
    ].
    nrewrite < (succ_plus ? ?).
    nrewrite > (plus_zero ?).
    //.
nqed.

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

nlet 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)
    ].

nlet rec vector_of_list (A: Type[0]) (l: List A) on l ≝
  match l return λl. Vector A (length A l) with
    [ Empty ⇒ [[ ]]
    | Cons hd tl ⇒ hd ::: (vector_of_list A tl)
    ].

(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)   
(* Rotates and shifts.                                                        *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
    
nlet rec rotate_left (A: Type[0]) (n: Nat)
                     (m: Nat) (v: Vector A n) on m: Vector A n ≝
  match m with
    [ Z ⇒ 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 Z) ↦ Vector A (S p)⌉)
          ]
    ].
 //.
nqed.

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

ndefinition 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)).
 ndestruct.
nqed.

ndefinition shift_right_1 ≝
  λA: Type[0].
  λn: Nat.
  λv: Vector A (S n).
  λa: A.
    reverse … (shift_left_1 … (reverse … v) a).
    
ndefinition shift_left ≝
  λA: Type[0].
  λn, m: Nat.
  λv: Vector A (S n).
  λa: A.
    iterate … (λx. shift_left_1 … x a) v m.
    
ndefinition 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.                                                        *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)

nlet 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.λ_. n = x → Bool with
    [ VEmpty ⇒
      match c return λx.λ_. x = Z → Bool with
        [ VEmpty ⇒ λ_. true
        | VCons p hd tl ⇒ λabsd.⊥
        ]
    | VCons o hd tl ⇒
        match c return λx.λ_. x = S o → Bool with
          [ VEmpty ⇒ λabsd.⊥
          | VCons p hd' tl' ⇒
            λprf.
              if (f hd hd') then
                (eq_v A o f tl (tl'⌈Vector A p ↦ Vector A o⌉))
              else
                false
          ]
    ]) (refl ? n).
    ##[##1,2: ndestruct
      | nlapply (S_inj … prf); #X; nrewrite < X; @]
nqed.

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

ndefinition mem ≝
 λA: Type[0].
 λeq_a : A → A → Bool.
 λn: Nat.
 λl: Vector A n.
 λx: A.
  fold_right … (λy,v. inclusive_disjunction (eq_a x y) v) false l.

ndefinition 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. conjunction (mem ? f ? q x) v) true v.
    
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)
(* Lemmas.                                                                    *)
(* -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= *)    
    
nlemma 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.
  nelim v.
  nnormalize.
  @.
  #N H V H2.
  nnormalize.
  nrewrite > H2.
  @.
nqed.