[Merged by Bors] - feat(Data/Fin/Tuple): Define Fin.take and initial theorems

Mathlib/Data/Fin/Tuple/Basic.lean

@@ -72,6 +72,7 @@ For a **pivot** `p : Fin (n + 1)`,
   satisfied.
 * `Fin.append a b` : append two tuples.
 * `Fin.repeat n a` : repeat a tuple `n` times.
+* `Fin.take m h` : take the first `h : m ≤ n` elements of a tuple.
 
 -/
 
@@ -1071,6 +1072,103 @@ theorem contractNth_apply_of_ne (j : Fin (n + 1)) (op : α → α → α) (g : F
 
 end ContractNth
 
+section Take
+
+variable {n : ℕ} {α : Fin n → Sort*}
+
+/-- Take the first `m` elements of an `n`-tuple where `m ≤ n`, returning an `m`-tuple. -/
+def take (v : (i : Fin n) → α i) (m : ℕ) (h : m ≤ n) : (i : Fin m) → α (castLE h i) :=
+  fun i ↦ v (castLE h i)
+
+@[simp]
+theorem take_def (v : (i : Fin n) → α i) (m : ℕ) (h : m ≤ n) (i : Fin m) :
+    (take v m h) i = v (castLE h i) := rfl
+
+@[simp]
+theorem take_zero (v : (i : Fin n) → α i) : take v 0 (Nat.zero_le n) = fun i ↦ elim0 i := by
+  ext i; exact elim0 i
+
+@[simp]
+theorem take_eq_self (v : (i : Fin n) → α i) :
+    take v n (le_refl n) = v := by ext i; simp [take]
+
+/-- Taking `m + 1` elements is equal to taking `m` elements and adding the `(m + 1)`th one. -/
+theorem take_succ_eq_snoc (v : (i : Fin n) → α i) (m : ℕ) (h : m < n) :
+    take v m.succ h = snoc (take v m (le_of_lt h)) (v ⟨m, h⟩) := by
+  ext i
+  induction m with
+  | zero =>
+    have h' : i = 0 := by
+      ext
+      simp only [Nat.succ_eq_add_one, Nat.reduceAdd, coe_fin_one]
+    subst h'
+    simp [take, snoc, castLE]
+  | succ m _ =>
+    induction i using reverseInduction with
+    | last => simp [take, snoc, castLT]; congr
+    | cast i _ => simp [snoc_cast_add]
+
+/-- `take` commutes with `update` for indices in the range of `take`. -/
+@[simp]
+theorem take_update_of_lt (v : (i : Fin n) → α i) (m : ℕ) (h : m ≤ n) (i : Fin m)
+    (x : α (castLE h i)) : take (update v (castLE h i) x) m h = update (take v m h) i x := by
+  ext j
+  by_cases h' : j = i
+  · rw [h']
+    simp only [take, update_same]
+  · have : castLE h j ≠ castLE h i := by simp [h']
+    simp only [take, update_noteq h', update_noteq this]
+
+/-- `take` is the same after `update` for indices outside the range of `take`. -/
+@[simp]
+theorem take_update_of_ge (v : (i : Fin n) → α i) (m : ℕ) (h : m ≤ n) (i : Fin n) (hi : i ≥ m)
+    (x : α i) : take (update v i x) m h = take v m h := by
+  ext j
+  have : castLE h j ≠ i := by
+    refine ne_of_val_ne ?_
+    simp only [coe_castLE]
+    exact Nat.ne_of_lt (lt_of_lt_of_le j.isLt hi)
+  simp only [take, update_noteq this]
+
+/-- Taking the first `m ≤ n` elements of an `append u v`, where `u` is a `n`-tuple, is the same as
+taking the first `m` elements of `u`. -/
+theorem take_append_left {n' m : ℕ} {α : Sort*} (u : (i : Fin n) → α) (v : (i : Fin n') → α)
+    (h : m ≤ n) : take (append u v) m (Nat.le_add_right_of_le h) = take u m h := by
+  ext i
+  have : castLE (Nat.le_add_right_of_le h) i = castAdd n' (castLE h i) := rfl
+  simp only [take, append_left, this]
+
+/-- Taking the first `n + m` elements of an `append u v`, where `v` is a `n'`-tuple and `m ≤ n'`,
+is the same as appending `u` with the first `m` elements of `v`. -/
+theorem take_append_right {n' m : ℕ} {α : Sort*} (u : (i : Fin n) → α) (v : (i : Fin n') → α)
+    (h : m ≤ n') : take (append u v) (n + m) (Nat.add_le_add_left h n) = append u (take v m h) := by
+  ext i
+  by_cases h' : i < n
+  · have : castLE (Nat.add_le_add_left h n) i = castAdd n' ⟨i.val, h'⟩ := by
+      simp only [castAdd_mk]
+      rfl
+    rw [take, this, append_left]
+    have : i = castAdd m ⟨i.val, h'⟩ := by simp only [castAdd_mk]
+    conv_rhs => rw [this, append_left]
+  · simp only [not_lt] at h'
+    let j := subNat n (cast (Nat.add_comm _ _) i) h'
+    have : i = natAdd n j := by simp [j]
+    conv_rhs => rw [this, append_right, take]
+    have : castLE (Nat.add_le_add_left h n) i = natAdd n (castLE h j) := by
+      simp_all only [natAdd_subNat_cast, j]
+      ext : 1
+      simp_all only [coe_castLE, coe_natAdd, coe_subNat, coe_cast, Nat.add_sub_cancel']
+    rw [take, this, append_right]
+
+@[simp]
+theorem take_init {α : Fin (n + 1) → Sort*} (v : (i : Fin (n + 1)) → α i) (m : ℕ) (h : m ≤ n) :
+    take (init v) m h = take v m (Nat.le_succ_of_le h) := by
+  ext i
+  simp only [take, init]
+  congr
+
+end Take
+
 /-- To show two sigma pairs of tuples agree, it to show the second elements are related via
 `Fin.cast`. -/
 theorem sigma_eq_of_eq_comp_cast {α : Type*} :