feat: Array.swap_perm (#6272)

This PR introduces the basic theory of permutations of `Array`s and
proves `Array.swap_perm`.

The API falls well short of what is available for `List` at this point.

src/Init/Data/Array.lean

@@ -20,3 +20,4 @@ import Init.Data.Array.MapIdx
 import Init.Data.Array.Set
 import Init.Data.Array.Monadic
 import Init.Data.Array.FinRange
+import Init.Data.Array.Perm

src/Init/Data/Array/Perm.lean

@@ -0,0 +1,65 @@
+/-
+Copyright (c) 2024 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+prelude
+import Init.Data.List.Nat.Perm
+import Init.Data.Array.Lemmas
+
+namespace Array
+
+open List
+
+/--
+`Perm as bs` asserts that `as` and `bs` are permutations of each other.
+
+This is a wrapper around `List.Perm`, and for now has much less API.
+For more complicated verification, use `perm_iff_toList_perm` and the `List` API.
+-/
+def Perm (as bs : Array α) : Prop :=
+  as.toList ~ bs.toList
+
+@[inherit_doc] scoped infixl:50 " ~ " => Perm
+
+theorem perm_iff_toList_perm {as bs : Array α} : as ~ bs ↔ as.toList ~ bs.toList := Iff.rfl
+
+@[simp] theorem perm_toArray (as bs : List α) : as.toArray ~ bs.toArray ↔ as ~ bs := by
+  simp [perm_iff_toList_perm]
+
+@[simp, refl] protected theorem Perm.refl (l : Array α) : l ~ l := by
+  cases l
+  simp
+
+protected theorem Perm.rfl {l : List α} : l ~ l := .refl _
+
+theorem Perm.of_eq {l₁ l₂ : Array α} (h : l₁ = l₂) : l₁ ~ l₂ := h ▸ .rfl
+
+protected theorem Perm.symm {l₁ l₂ : Array α} (h : l₁ ~ l₂) : l₂ ~ l₁ := by
+  cases l₁; cases l₂
+  simp only [perm_toArray] at h
+  simpa using h.symm
+
+protected theorem Perm.trans {l₁ l₂ l₃ : Array α} (h₁ : l₁ ~ l₂) (h₂ : l₂ ~ l₃) : l₁ ~ l₃ := by
+  cases l₁; cases l₂; cases l₃
+  simp only [perm_toArray] at h₁ h₂
+  simpa using h₁.trans h₂
+
+instance : Trans (Perm (α := α)) (Perm (α := α)) (Perm (α := α)) where
+  trans h₁ h₂ := Perm.trans h₁ h₂
+
+theorem perm_comm {l₁ l₂ : Array α} : l₁ ~ l₂ ↔ l₂ ~ l₁ := ⟨Perm.symm, Perm.symm⟩
+
+theorem Perm.push (x y : α) {l₁ l₂ : Array α} (p : l₁ ~ l₂) :
+    (l₁.push x).push y ~ (l₂.push y).push x := by
+  cases l₁; cases l₂
+  simp only [perm_toArray] at p
+  simp only [push_toArray, List.append_assoc, singleton_append, perm_toArray]
+  exact p.append (Perm.swap' _ _ Perm.nil)
+
+theorem swap_perm {as : Array α} {i j : Nat} (h₁ : i < as.size) (h₂ : j < as.size) :
+    as.swap i j ~ as := by
+  simp only [swap, perm_iff_toList_perm, toList_set]
+  apply set_set_perm
+
+end Array

src/Init/Data/List/Lemmas.lean

@@ -589,6 +589,14 @@ theorem getElem?_set' {l : List α} {i j : Nat} {a : α} :
   · simp only [getElem?_set_self', Option.map_eq_map, ↓reduceIte, *]
   · simp only [ne_eq, not_false_eq_true, getElem?_set_ne, ↓reduceIte, *]
 
+@[simp] theorem set_getElem_self {as : List α} {i : Nat} (h : i < as.length) :
+    as.set i as[i] = as := by
+  apply ext_getElem
+  · simp
+  · intro n h₁ h₂
+    rw [getElem_set]
+    split <;> simp_all
+
 theorem set_eq_of_length_le {l : List α} {n : Nat} (h : l.length ≤ n) {a : α} :
     l.set n a = l := by
   induction l generalizing n with

src/Init/Data/List/Nat.lean

@@ -15,3 +15,4 @@ import Init.Data.List.Nat.Find
 import Init.Data.List.Nat.BEq
 import Init.Data.List.Nat.Modify
 import Init.Data.List.Nat.InsertIdx
+import Init.Data.List.Nat.Perm

src/Init/Data/List/Nat/Perm.lean

@@ -0,0 +1,54 @@
+/-
+Copyright (c) 2024 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+prelude
+import Init.Data.List.Nat.TakeDrop
+import Init.Data.List.Perm
+
+namespace List
+
+/-- Helper lemma for `set_set_perm`-/
+private theorem set_set_perm' {as : List α} {i j : Nat} (h₁ : i < as.length) (h₂ : i + j < as.length)
+    (hj : 0 < j) :
+    (as.set i as[i + j]).set (i + j) as[i] ~ as := by
+  have : as =
+      as.take i ++ as[i] :: (as.take (i + j)).drop (i + 1) ++ as[i + j] :: as.drop (i + j + 1) := by
+    simp only [getElem_cons_drop, append_assoc, cons_append]
+    rw [← drop_append_of_le_length]
+    · simp
+    · simp; omega
+  conv => lhs; congr; congr; rw [this]
+  conv => rhs; rw [this]
+  rw [set_append_left _ _ (by simp; omega)]
+  rw [set_append_right _ _ (by simp; omega)]
+  rw [set_append_right _ _ (by simp; omega)]
+  simp only [length_append, length_take, length_set, length_cons, length_drop]
+  rw [(show i - min i as.length = 0 by omega)]
+  rw [(show i + j - (min i as.length + (min (i + j) as.length - (i + 1) + 1)) = 0 by omega)]
+  simp only [set_cons_zero]
+  simp only [append_assoc]
+  apply Perm.append_left
+  apply cons_append_cons_perm
+
+theorem set_set_perm {as : List α} {i j : Nat} (h₁ : i < as.length) (h₂ : j < as.length) :
+    (as.set i as[j]).set j as[i] ~ as := by
+  if h₃ : i = j then
+    simp [h₃]
+  else
+    if h₃ : i < j then
+      let j' := j - i
+      have t : j = i + j' := by omega
+      generalize j' = j' at t
+      subst t
+      exact set_set_perm' _ _ (by omega)
+    else
+      rw [set_comm _ _ _ (by omega)]
+      let i' := i - j
+      have t : i = j + i' := by omega
+      generalize i' = i' at t
+      subst t
+      apply set_set_perm' _ _ (by omega)
+
+end List

src/Init/Data/List/Nat/TakeDrop.lean

@@ -345,7 +345,7 @@ theorem drop_append {l₁ l₂ : List α} (i : Nat) : drop (l₁.length + i) (l
   rw [drop_append_eq_append_drop, drop_eq_nil_of_le] <;>
     simp [Nat.add_sub_cancel_left, Nat.le_add_right]
 
-theorem set_eq_take_append_cons_drop {l : List α} {n : Nat} {a : α} :
+theorem set_eq_take_append_cons_drop (l : List α) (n : Nat) (a : α) :
     l.set n a = if n < l.length then l.take n ++ a :: l.drop (n + 1) else l := by
   split <;> rename_i h
   · ext1 m

src/Init/Data/List/Perm.lean

@@ -39,6 +39,9 @@ protected theorem Perm.symm {l₁ l₂ : List α} (h : l₁ ~ l₂) : l₂ ~ l
   | swap => exact swap ..
   | trans _ _ ih₁ ih₂ => exact trans ih₂ ih₁
 
+instance : Trans (Perm (α := α)) (Perm (α := α)) (Perm (α := α)) where
+  trans h₁ h₂ := Perm.trans h₁ h₂
+
 theorem perm_comm {l₁ l₂ : List α} : l₁ ~ l₂ ↔ l₂ ~ l₁ := ⟨Perm.symm, Perm.symm⟩
 
 theorem Perm.swap' (x y : α) {l₁ l₂ : List α} (p : l₁ ~ l₂) : y :: x :: l₁ ~ x :: y :: l₂ :=
@@ -102,7 +105,7 @@ theorem perm_append_comm : ∀ {l₁ l₂ : List α}, l₁ ++ l₂ ~ l₂ ++ l
   | _ :: _, _ => (perm_append_comm.cons _).trans perm_middle.symm
 
 theorem perm_append_comm_assoc (l₁ l₂ l₃ : List α) :
-    Perm (l₁ ++ (l₂ ++ l₃)) (l₂ ++ (l₁ ++ l₃)) := by
+    (l₁ ++ (l₂ ++ l₃)) ~ (l₂ ++ (l₁ ++ l₃)) := by
   simpa only [List.append_assoc] using perm_append_comm.append_right _
 
 theorem concat_perm (l : List α) (a : α) : concat l a ~ a :: l := by simp
@@ -133,7 +136,7 @@ theorem Perm.nil_eq {l : List α} (p : [] ~ l) : [] = l := p.symm.eq_nil.symm
 
 theorem not_perm_nil_cons (x : α) (l : List α) : ¬[] ~ x :: l := (nomatch ·.symm.eq_nil)
 
-theorem not_perm_cons_nil {l : List α} {a : α} : ¬(Perm (a::l) []) :=
+theorem not_perm_cons_nil {l : List α} {a : α} : ¬((a::l) ~ []) :=
   fun h => by simpa using h.length_eq
 
 theorem Perm.isEmpty_eq {l l' : List α} (h : Perm l l') : l.isEmpty = l'.isEmpty := by
@@ -478,6 +481,15 @@ theorem Perm.flatten {l₁ l₂ : List (List α)} (h : l₁ ~ l₂) : l₁.flatt
 
 @[deprecated Perm.flatten (since := "2024-10-14")] abbrev Perm.join := @Perm.flatten
 
+theorem cons_append_cons_perm {a b : α} {as bs : List α} :
+    a :: as ++ b :: bs ~ b :: as ++ a :: bs := by
+  suffices [[a], as, [b], bs].flatten ~ [[b], as, [a], bs].flatten by simpa
+  apply Perm.flatten
+  calc
+    [[a], as, [b], bs] ~ [as, [a], [b], bs] := Perm.swap as [a] _
+    _ ~ [as, [b], [a], bs] := Perm.cons _ (Perm.swap [b] [a] _)
+    _ ~ [[b], as, [a], bs] := Perm.swap [b] as _
+
 theorem Perm.flatMap_right {l₁ l₂ : List α} (f : α → List β) (p : l₁ ~ l₂) : l₁.flatMap f ~ l₂.flatMap f :=
   (p.map _).flatten
 

src/Init/Data/Vector/Basic.lean

@@ -21,6 +21,9 @@ deriving Repr, DecidableEq
 
 attribute [simp] Vector.size_toArray
 
+/-- Convert `xs : Array α` to `Vector α xs.size`. -/
+abbrev Array.toVector (xs : Array α) : Vector α xs.size := .mk xs rfl
+
 namespace Vector
 
 /-- Syntax for `Vector α n` -/