sized-array/SizedArray.lean

def SizedArray (α : Type u) (n : Nat) :=
  {as : List α // as.length = n}

namespace SizedArray

def inner {n : Nat} : SizedArray α n → List α :=
  λxs => xs.val

def empty : SizedArray α 0 := ⟨[], rfl⟩
def singleton (a : α) : SizedArray α 1 := ⟨[a], rfl⟩ 

def cons {n : Nat} : α → SizedArray α n → SizedArray α (n.succ) := 
  λx xs =>
    let ⟨xs, h⟩ := xs
    ⟨x :: xs, by rw [List.length_cons, h]⟩ 

def replicate (n : Nat) (x : α) : SizedArray α n :=
  match n with
  | Nat.zero => empty
  | Nat.succ n => cons x (replicate n x)

def append {a b : Nat} (as : SizedArray α a) (bs : SizedArray α b)
  : SizedArray α (a + b) := by
  let ⟨as, heqa⟩ := as
  let ⟨bs, heqb⟩ := bs

  let combined := as ++ bs
  let (len_proof : combined.length = a + b) := by
    rw [List.length_append]
    rw [heqa, heqb]

  exact ⟨combined, len_proof⟩  

def from_list (xs : List α) : SizedArray α (xs.length) := 
  match xs with
  | List.nil => empty
  | List.cons x xs => cons x (from_list xs)

@[simp]
def tail_neq_implies_neq
  (a b : α) (as bs : List α) (h : as ≠ bs)
  : a :: as ≠ b :: bs := by
  induction as with
  | nil => 
    induction bs with
    | nil => simp at h
    | cons b bs => simp
  | cons a' as' =>
    simp [*]

@[simp]
def length_neq_implies_neq :
  ∀{as bs : List α}, as.length ≠ bs.length → as ≠ bs
  | [],     [], h => absurd h (by simp)
  | [], (_::_), _ => by simp
  | (_::_), [], _ => by simp
  | (x::xs), (y::ys), h =>
    let (h₂ : xs.length ≠ ys.length) := by
      cases xs with
      | nil =>
        rw [List.length] at h
        simp at h
        rw [List.length_nil]
        exact h
      | cons x xs => 
        rw [List.length, List.length, List.length] at h
        simp at h
        exact h
    let h₃ := @length_neq_implies_neq α xs ys h₂
    tail_neq_implies_neq x y xs ys h₃ 

-- def length_g0_imp_nonempty (xs : List α) ⦃h : xs.length > 0⦄
--   : List α ≠ List.nil := sorry

-- def is_not_empty {n : Nat} {h : n > 0}
--   : inner (SizedArray α n) ≠ [] := sorry

def gt_flip {a b : Nat} (h : a > b) : (b < a) := by simp [*]
def lt_flip {a b : Nat} (h : a < b) : (b > a) := by simp [*]
def ge_flip {a b : Nat} (h : a ≥ b) : (b ≤ a) := by simp [*]
def le_flip {a b : Nat} (h : a ≤ b) : (b ≥ a) := by simp [*]

def gt_comm {a b : Nat} : a > b ↔ b < a := by simp [*]
def lt_comm {a b : Nat} : a < b ↔ b > a := by simp [*]
def ge_comm {a b : Nat} : a ≥ b ↔ b ≤ a := by simp [*]
def le_comm {a b : Nat} : a ≤ b ↔ b ≥ a := by simp [*]

def gt_implies_neq {a b : Nat} (h : a > b) : a ≠ b := by
  have hlt := gt_flip h
  have neq := Nat.ne_of_lt hlt
  exact Ne.symm neq

def lt_implies_neq {a b : Nat} (h : a < b) : a ≠ b := by 
  have hgt := lt_flip h
  have (neq : b ≠ a) := gt_implies_neq hgt
  exact Ne.symm neq

def length_gt_zero_implies_nonempty {xs : List α} (h : xs.length > 0)
  : xs ≠ [] := 
    let (hnz : List.length xs ≠ 0) := gt_implies_neq h
    let (hne_len : xs.length ≠ [].length) := by
      rw [List.length_nil]
      exact hnz

    length_neq_implies_neq hne_len

def length_eq_zero_implies_empty {xs : List α} (_ : xs.length = 0)
  : xs = [] :=
  match xs with
  | List.nil => rfl

def empty_implies_length_eq_zero {xs : List α} (h : xs = [])
  : xs.length = 0 := by
    rw [h]
    rfl

def length_zero_iff_empty {xs : List α}
  : xs.length = 0 ↔ xs = [] :=
  ⟨
    length_eq_zero_implies_empty, 
    empty_implies_length_eq_zero
  ⟩ 

def not_empty_implies_length_neq_zero {xs : List α} (h : xs ≠ [])
  : xs.length ≠ 0 := 
  match xs with
  | List.nil => by
    have (nh : [] = []) := Eq.refl ([] : List α)
    exact absurd nh h
  | List.cons head tail => by
    rw [List.length_cons]
    exact Nat.succ_ne_zero (List.length tail)

def length_gt_zero_iff_nonempty {xs : List α}
  : xs.length > 0 ↔ xs ≠ [] := 
  ⟨
    length_gt_zero_implies_nonempty,
    by
      intro h
      have hnz := not_empty_implies_length_neq_zero h
      exact Nat.zero_lt_of_ne_zero hnz
  ⟩  


@[simp]
def length_succ_implies_nonempty
  {n : Nat} {xs : List α} (h : xs.length = n.succ)
  : xs ≠ [] := by
    apply length_gt_zero_implies_nonempty
    have (hsnnz : n.succ ≠ 0) := Nat.succ_ne_zero n
    have (hsngz : n.succ > 0) := Nat.zero_lt_of_ne_zero hsnnz
    have (h :  xs.length > 0) := Eq.substr h hsngz
    exact lt_flip h

def head {n : Nat} (h : n > 0) : SizedArray α n → α :=
  λxs =>
    let ⟨xs, h₂⟩ := xs
    let (hgz : List.length xs > 0) := Eq.substr h₂ h 
    let hne := length_gt_zero_implies_nonempty hgz
    
    List.head xs hne

def succ_inj {a b : Nat} (h : a.succ = b.succ) : a = b := by
  induction a with
  | zero =>
    induction b with
    | zero => rfl
    | succ n =>
      simp at h
      simp
      rw [Nat.add_succ] at h
      simp at h
  | succ n _ =>
    simp at h
    rw [Nat.succ_eq_add_one]
    exact h

-- def gt_and_lt_implies_false {a b : Nat} (hgt : a > b) (hlt : a < b)
--   : False := sorry

def succ_inj_gt {a b : Nat} (h : a > b) : a.succ > b.succ := by
  have (h₂ : b < a) := gt_flip h
  have (h₃ : Nat.succ b < Nat.succ a) := Nat.succ_lt_succ h₂ 
  exact gt_flip h₃ 

def tail_length_lt_source_length (x : α) (xs : List α) : ((x :: xs).length > xs.length) := by
  rw [List.length_cons]
  induction (List.length xs) with
  | zero => simp
  | succ n ih =>
    apply succ_inj_gt
    exact ih

@[simp]
def gt_self_implies_false {n : Nat} (h : n > n) : False := 
  let neq := gt_implies_neq h
  let eq  := Eq.refl n
  absurd eq neq
  
def tail_list {n : Nat} (h : n > 0) (xs : SizedArray α n) 
  : List α := by
    let ⟨xs, heqn⟩ := xs
    let (hgz : List.length xs > 0) := Eq.substr heqn h 
    let (hne : xs ≠ []) := length_gt_zero_implies_nonempty hgz

    match xs with
    | List.nil => exact (by simp at hne)
    | List.cons _ xs => exact xs

def safe_tail (xs : List α) (h : xs.length > 0) : List α :=
  match xs with
  | List.nil => by
    rw [List.length_nil] at h
    exact (gt_self_implies_false h).elim
  | List.cons _ xs => xs

def safe_tail_len (xs : List α) (h : xs.length > 0)
  : (safe_tail xs h).length = xs.length.pred := by
  induction xs with
  | nil => simp at h
  | cons a as _ =>  
    rw [safe_tail]
    simp
  
def tail {n : Nat} (h : n > 0) (xs : SizedArray α n)
  : SizedArray α n.pred := by
    let ⟨xs, heqn⟩ := xs
    let (hgz : List.length xs > 0) := Eq.substr heqn h 
    let tail := safe_tail xs hgz
    let (tail_len : (safe_tail xs hgz).length = (List.length xs).pred)
      := safe_tail_len xs hgz
    rw [heqn] at tail_len

    exact ⟨tail, tail_len⟩ 

def uncons (h : n > 0) (xs : SizedArray α n) 
  : α × (SizedArray α n.pred) := 
    ⟨head h xs, tail h xs⟩ 

def map {n : Nat} (f : α → β) (xs : SizedArray α n) : SizedArray β n :=
  match n with
  | Nat.zero => empty
  | Nat.succ n => 
    let (hnz : n.succ > 0) := lt_flip (Nat.zero_lt_succ n)
    let ⟨head, tail⟩ := uncons hnz xs 
    cons (f head) (map f tail)

def reverse {n : Nat} (xs : SizedArray α n) : SizedArray α n := by
  let ⟨xs, heqn⟩ := xs
  let reversed := xs.reverse
  let h := List.length_reverse xs
  rw [heqn] at h
  exact ⟨reversed, h⟩ 
  -- match n with
  -- | Nat.zero => empty
  -- | Nat.succ n => by
  --   let ⟨xs, heqn⟩ := xs
  --   let (hne : xs ≠ []) := length_succ_implies_nonempty heqn
  --   let (hnz : xs.length ≠ 0) := not_empty_implies_length_neq_zero hne
  --   let (hgz : xs.length > 0) := Nat.zero_lt_of_ne_zero hnz
  --   let ⟨head, tail⟩ := uncons hgz ⟨xs, rfl⟩ 
    
  --   let reversed := append (reverse tail) (singleton head)
  --   rw [heqn, Nat.pred_succ, ← Nat.succ_eq_add_one] at reversed
  --   exact reversed

def matchable {n : Nat} (xs : SizedArray α n)
  : Option (α × (SizedArray α n.pred)) :=
  match n with
  | Nat.zero => none
  | Nat.succ x => by
    let (hgz : x.succ > 0) := by
      have (h : 0 < x.succ) := Nat.zero_lt_succ x
      exact lt_flip h

    let (h : α × (SizedArray α x)) := by
      have h₂ := uncons hgz xs
      simp at h₂ 
      exact h₂

    rw [Nat.pred_succ]
    exact some h

inductive Either (α β : Type u) where
| left  : α → Either α β
| right : β → Either α β

def matchable' {n : Nat} (xs : SizedArray α n) 
  : Either (SizedArray α 0) (α × (SizedArray α n.pred)) := by
  match n with
  | Nat.zero => exact Either.left empty
  | Nat.succ x => 
    apply Either.right

    rw [Nat.pred_succ]

    let (hgz : x.succ > 0) := by
      have h₂ := Nat.zero_lt_succ x
      exact lt_flip h₂ 

    exact uncons hgz xs


-- def matchable_with_proofs {n : Nat} (xs : SizedArray α n)
--   : Either ((SizedArray α 0) × n = 0) ((α × (SizedArray α n.pred)) × n > 0) := by
--   match n with
--   | zero => sorry

-- this is actually not true!
-- 1.pred = 0 ∧ 0.pred = 0 so this would prove 1 = 0 which is FALSE!
-- sure is inconvenient though
-- def pred_eq {a b : Nat} (h : a.pred = b.pred) : a = b := by
--   induction a with
--   | zero =>
--     rw [Nat.pred_zero] at h
--     induction b with
--     | zero => rfl
--     | succ n ih =>
--       rw [Nat.pred_succ] at h
-- 
--   | succ n ih => sorry

def pred_sub_assoc {a b : Nat}
  : (Nat.pred a) - b = Nat.pred (a - b) := by
    induction a with
    | zero => rw [Nat.pred_zero, Nat.zero_sub, Nat.pred_zero ]
    | succ n _ =>
      rw [← Nat.sub_succ, Nat.pred_succ]
      rw [Nat.succ_sub_succ]


def drop {from_len : Nat} (n : Nat) (xs : SizedArray α from_len) 
  : SizedArray α (from_len - n) := by
  let ⟨inner, heqn⟩ := xs
  match inner with
  | List.nil => 
    rw [List.length_nil] at heqn
    rw [← heqn]
    rw [Nat.zero_sub]
    exact empty
  | List.cons x xs' => 
    rw [List.length_cons] at heqn

    match n with 
    | Nat.zero =>
      rw [Nat.sub_zero]
      exact xs
    | Nat.succ n =>
      let (hgz : from_len > 0) := by
        let hnz := Nat.succ_ne_zero (List.length xs')
        rw [← heqn] 
        exact Nat.zero_lt_of_ne_zero hnz

      rw [Nat.sub_succ]
      rw [← pred_sub_assoc]
      exact drop n (tail hgz xs)

example : SizedArray Nat 2 := drop 3 (from_list [1, 2, 3, 4, 5])
example : SizedArray Nat 0 := drop 100 (from_list [1, 2, 3, 4, 5])

#check drop 1 (drop 2 (from_list [1,2,3,4,5,6]))

def drop_end {from_len : Nat} (n : Nat) (xs : SizedArray α from_len)
  : SizedArray α (from_len - n) := 
  let reversed := reverse xs
  let dropped := drop n reversed
  reverse dropped

#eval drop_end 1 (drop 2 (from_list [1,2,3,4,5,6]))

-- def min_comm {a b : Nat} : min a b = min b a := by
--   rw [Nat.min_def, Nat.min_def]
--   sorry

def succ_le_succ_iff_le {a b : Nat} : a.succ ≤ b.succ ↔ a ≤ b :=
  ⟨
    Nat.le_of_succ_le_succ,
    Nat.succ_le_succ
  ⟩ 

#check ite_congr

def map_ite (f : β → γ) {c : Prop} [Decidable c] {a b : β}
  : f (if c then a else b) = (if c then (f a) else (f b)) := by
  cases h : decide c with
  | false => 
    let (hcf : c = False) := eq_false_of_decide h
    simp only [hcf]
    rw [ite_false, ite_false]
  | true =>
    let (hcf : c = True) := eq_true_of_decide h
    simp only [hcf]
    rw [ite_true, ite_true]

def succ_min {a b : Nat} : (min a b).succ = min a.succ b.succ := by
  rw [Nat.min_def]
  rw [Nat.min_def]
  let (h₁ : (a ≤ b) = (a.succ ≤ b.succ)) := by rw [succ_le_succ_iff_le]
  rw [map_ite Nat.succ]
  simp only [h₁]

def zip {α β : Type u} {a b : Nat} (as : SizedArray α a) (bs : SizedArray β b)
  : SizedArray (α × β) (min a b) := 
  match a, b with
  | 0, _ => by
    rw [Nat.min_def]
    simp
    exact empty
  | x', 0 => by
    rw [Nat.min_def]
    simp
    match x' with
    | 0 => simp; exact empty
    | Nat.succ x => simp; exact empty
  | Nat.succ a, Nat.succ b => by
    let ⟨as_inner, haeqn⟩ := as
    let ⟨bs_inner, hbeqn⟩ := bs
    let hane := length_succ_implies_nonempty haeqn
    let hagz :=
      not_empty_implies_length_neq_zero hane
      |> Nat.zero_lt_of_ne_zero 
      |> lt_flip
    let hbne := length_succ_implies_nonempty hbeqn
    let hbgz :=
      not_empty_implies_length_neq_zero hbne
      |> Nat.zero_lt_of_ne_zero 
      |> lt_flip
    let hngz := lt_flip <| Eq.subst haeqn hagz
    let ⟨a_head, a_tail⟩ := uncons hngz as
    let hngz := lt_flip <| Eq.subst hbeqn hbgz
    let ⟨b_head, b_tail⟩ := uncons hngz bs

    let bleh := cons (a_head, b_head) (zip a_tail b_tail)
    rw [Nat.pred_succ, Nat.pred_succ] at bleh
    rw [succ_min] at bleh
    exact bleh

def take {from_len : Nat} (n : Nat) (xs : SizedArray α from_len)
  : SizedArray α (min n from_len) := 
  let dummy := replicate n ()
  let zipped := zip dummy xs
  map Prod.snd zipped
  -- let reversed := reverse xs
  -- let dropped := drop (from_len - n) reversed
  -- let unreversed := reverse dropped
  -- cases n with
  -- | zero => 
  --   rw [Nat.sub_zero, Nat.sub_self] at unreversed
  --   rw [Nat.min_def]
  --   simp
  --   exact empty
  -- | succ x => 
  --   let n := x.succ
  --   let (result : SizedArray α (min n from_len)) := sorry

  --   exact result
    -- sep sep sep sep
    -- cases from_len with
    -- | zero => 
    --   rw [Nat.min_def]
    --   simp
    --   exact empty
    -- | succ y => 
    --   c
    --   sorry

  -- match n with
  -- | Nat.zero => 
  --   rw [Nat.min]
  --   simp
  --   exact empty
  -- | Nat.succ x => 
  --   match from_len with
  --   | Nat.zero => 
  --     rw [Nat.min]
  --     simp
  --     exact empty
  --   | Nat.succ x' =>
  --     let ⟨inner, heqsn⟩ := xs
  --     let (hne : inner ≠ []) := length_succ_implies_nonempty heqsn
  --     let (hnz : inner.length ≠ 0) := not_empty_implies_length_neq_zero hne
  --     let (hgz : inner.length > 0) := Nat.zero_lt_of_ne_zero hnz
  --     -- let ⟨head, tail⟩ := uncons hgz ⟨inner, rfl⟩ 
  --     -- have taken := cons head (take x tail)
  --     -- simp at taken
  --     -- let tail := take x' (tail hgz ⟨inner, rfl⟩)
  --     -- simp [*] at tail
      

  --     sorry

-- example : SizedArray Nat 3 := take 3 (by simp) (from_list [1,2,3,4,5])

#eval from_list [1,2,3,4,5,6,7,8,9,10] |> take 7 |> drop 2 |> reverse |> map (· * 2) |> (tail (by simp) ·)

def left_pad' (to : Nat) {n : Nat} (with' : α) (xs : SizedArray α n)
              (h : n ≤ to)
  : SizedArray α to := by
    let pad_amt := to - n
    let (h₂ : pad_amt = to - n) := rfl

    let padding := replicate pad_amt with'
    let padded := append padding xs

    rw [h₂, Nat.sub_add_cancel] at padded

    exact padded
    exact h

def left_pad (to : Nat) {n : Nat} (with' : α) (xs : SizedArray α n)
  : SizedArray α (max to n) := by
    cases hle : (decide (to ≤ n)) with
    | true  => 
      rw [max, Nat.instMaxNat, maxOfLe]
      let hle : to ≤ n := of_decide_eq_true hle
      simp [hle]
      exact xs
    | false => 
      let hle : ¬to ≤ n := of_decide_eq_false hle

      rw [max, Nat.instMaxNat, maxOfLe]
      simp [hle]

      let hnle := hle
      rw [Nat.not_le_eq, ← Nat.succ_eq_add_one] at hnle
      let hnle := Nat.le_of_succ_le hnle
      
      exact left_pad' to with' xs hnle

#eval left_pad 10 413 (from_list [1,2,3,4,5,6,7])

def right_pad (to : Nat) {n : Nat} (with' : α) (xs : SizedArray α n)
  : SizedArray α (max to n) := xs |> reverse |> left_pad to with' |> reverse

#eval right_pad 10 413 (from_list [1,2,3,4,5,6,7])

notation xs "[" a " ... " b "]" => take (b - a) $ drop a xs
notation xs "["   " ... " b "]" => take b xs
notation xs "[" a " ..= " b "]" => take (b - (a - 1)) $ drop a xs
notation xs "["   " ..= " b "]" => take (b + 1) xs
-- ⦋⦌
-- ⹗⹘
def List.to_sized (xs : List α) := (from_list xs)

#eval (from_list [1,2,3]) [1 ... 2]
#eval (from_list [1,2,3]) [  ... 2]
#eval (from_list [1,2,3]) [1 ..= 2]
#eval (from_list [1,2,3]) [  ..= 2]

def read_exact : α → α := id

example : SizedArray UInt8 8 := Id.run do
  let mut header := replicate (1 + 8 + 2) (0 : UInt8)
  header := read_exact header
  let bla := header [1 ... 9]
  by
    apply pure
    have h : 8 = (min (9 - 1) (1 + 8 + 2 - 1)) := by simp
    rw [← h] at bla
    exact bla