benchmark

src/Init/Data/Array/InsertionSort.lean

@@ -5,24 +5,91 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Init.Data.Array.Basic
+import Init.Data.Nat.Fold
+import Init.Data.Vector.Lemmas
 
-@[inline] def Array.insertionSort (a : Array α) (lt : α → α → Bool := by exact (· < ·)) : Array α :=
-  traverse a 0 a.size
-where
-  @[specialize] traverse (a : Array α) (i : Nat) (fuel : Nat) : Array α :=
-    match fuel with
-    | 0      => a
-    | fuel+1 =>
-      if h : i < a.size then
-        traverse (swapLoop a i h) (i+1) fuel
-      else
-        a
-  @[specialize] swapLoop (a : Array α) (j : Nat) (h : j < a.size) : Array α :=
-    match (generalizing := false) he:j with -- using `generalizing` because we don't want to refine the type of `h`
-    | 0    => a
-    | j'+1 =>
-      have h' : j' < a.size := by subst j; exact Nat.lt_trans (Nat.lt_succ_self _) h
-      if lt a[j] a[j'] then
-        swapLoop (a.swap j j') j' (by rw [size_swap]; assumption; done)
-      else
-        a
+namespace Vector
+
+/-- Swap the `i`-th element repeatedly to the left, while the element to its left is not `lt` it. -/
+@[specialize, inline] def swapLeftWhileLT {n} (a : Vector α n) (i : Nat) (h : i < n)
+    (lt : α → α → Bool := by exact (· < ·)) : Vector α n :=
+  match h' : i with
+  | 0 => a
+  | i'+1 =>
+    if lt a[i] a[i'] then
+      swapLeftWhileLT (a.swap i' i) i' (by omega) lt
+    else
+      a
+
+end Vector
+
+open Vector
+namespace Array
+
+/-- Sort an array in place using insertion sort. -/
+@[inline] def insertionSort (a : Array α) (lt : α → α → Bool := by exact (· < ·)) : Array α :=
+  a.size.fold (init := ⟨a, rfl⟩) (fun i h acc => swapLeftWhileLT acc i h lt) |>.toArray
+
+/-- Insert an element into an array, after the last element which is not `lt` the inserted element. -/
+def orderedInsert (a : Array α) (x : α) (lt : α → α → Bool := by exact (· < ·)) : Array α :=
+ swapLeftWhileLT ⟨a.push x, rfl⟩ a.size (by simp) lt |>.toArray
+
+end Array
+
+/-! ### Verification -/
+
+namespace Vector
+
+theorem swapLeftWhileLT_push {n} (a : Vector α n) (x : α) (j : Nat) (h : j < n) :
+    swapLeftWhileLT (a.push x) j (by omega) lt = (swapLeftWhileLT a j h lt).push x := by
+  induction j generalizing a with
+  | zero => simp [swapLeftWhileLT]
+  | succ j ih =>
+    simp [swapLeftWhileLT]
+    split <;> rename_i h
+    · rw [Vector.getElem_push_lt (by omega), Vector.getElem_push_lt (by omega)] at h
+      rw [← Vector.push_swap, ih, if_pos h]
+    · rw [Vector.getElem_push_lt (by omega), Vector.getElem_push_lt (by omega)] at h
+      rw [if_neg h]
+
+theorem swapLeftWhileLT_cast {n m} (a : Vector α n) (j : Nat) (h : j < n) (h' : n = m) :
+    swapLeftWhileLT (a.cast h') j (by omega) lt = (swapLeftWhileLT a j h lt).cast h' := by
+  subst h'
+  simp
+
+end Vector
+
+namespace Array
+
+@[simp] theorem size_insertionSort (a : Array α) : (a.insertionSort lt).size = a.size := by
+  simp [insertionSort]
+
+private theorem insertionSort_push' (a : Array α) (x : α) :
+    (a.push x).insertionSort lt =
+      (swapLeftWhileLT ⟨(a.insertionSort lt).push x, rfl⟩ a.size (by simp) lt).toArray := by
+  rw [insertionSort, Nat.fold_congr (size_push a x), Nat.fold]
+  have : (a.size.fold (fun i h acc => swapLeftWhileLT acc i (by simp; omega) lt) ⟨a.push x, rfl⟩) =
+      ((a.size.fold (fun i h acc => swapLeftWhileLT acc i h lt) ⟨a, rfl⟩).push x).cast (by simp) := by
+    rw [Vector.eq_cast_iff]
+    simp only [Nat.fold_eq_finRange_foldl]
+    rw [← List.foldl_hom (fun a => (Vector.push x a)) _ (fun v ⟨i, h⟩ => swapLeftWhileLT v i (by omega) lt)]
+    rw [Vector.push_mk]
+    rw [← List.foldl_hom (Vector.cast _) _ (fun v ⟨i, h⟩ => swapLeftWhileLT v i (by omega) lt)]
+    · simp
+    · intro v i
+      simp only
+      rw [swapLeftWhileLT_cast]
+    · simp [swapLeftWhileLT_push]
+  rw [this]
+  simp only [Nat.lt_add_one, swapLeftWhileLT_cast, Vector.toArray_cast]
+  unfold insertionSort
+  simp only [Vector.push]
+  congr
+  all_goals simp
+
+theorem insertionSort_push (a : Array α) (x : α) :
+    (a.push x).insertionSort lt = (a.insertionSort lt).orderedInsert x lt := by
+  rw [insertionSort_push', orderedInsert]
+  simp
+
+end Array

src/Init/Data/Array/Lemmas.lean

@@ -595,6 +595,20 @@ theorem getElem_set (a : Array α) (i : Nat) (h' : i < a.size) (v : α) (j : Nat
     (ne : i ≠ j) : (a.set i v)[j]? = a[j]? := by
   by_cases h : j < a.size <;> simp [getElem?_lt, getElem?_ge, Nat.ge_of_not_lt, ne, h]
 
+theorem push_set (a : Array α) (x y : α) {i : Nat} {hi} :
+    (a.set i x).push y = (a.push y).set i x (by simp; omega):= by
+  ext j h₁ h₂
+  · simp
+  · if h' : j = a.size then
+      rw [getElem_push, getElem_set_ne, dif_neg]
+      all_goals simp_all <;> omega
+    else
+      rw [getElem_push_lt, getElem_set, getElem_set]
+      split
+      · rfl
+      · rw [getElem_push_lt]
+      simp_all; omega
+
 /-! # setIfInBounds -/
 
 @[simp] theorem set!_is_setIfInBounds : @set! = @setIfInBounds := rfl
@@ -1718,6 +1732,11 @@ theorem swap_comm (a : Array α) {i j : Nat} {hi hj} : a.swap i j hi hj = a.swap
     · split <;> simp_all
     · split <;> simp_all
 
+theorem push_swap (a : Array α) (x : α) {i j : Nat} {hi hj} :
+    (a.swap i j hi hj).push x = (a.push x).swap i j (by simp; omega) (by simp; omega) := by
+  rw [swap_def, swap_def]
+  simp [push_set, getElem_push_lt, hi, hj]
+
 /-! ### eraseIdx -/
 
 theorem eraseIdx_eq_eraseIdxIfInBounds {a : Array α} {i : Nat} (h : i < a.size) :

src/Init/Data/Vector/Lemmas.lean

@@ -64,6 +64,9 @@ protected theorem ext {a b : Vector α n} (h : (i : Nat) → (_ : i < n) → a[i
 @[simp] theorem pop_mk {data : Array α} {size : data.size = n} :
     (Vector.mk data size).pop = Vector.mk data.pop (by simp [size]) := rfl
 
+@[simp] theorem swap_mk {data : Array α} {size : data.size = n} {i j : Nat} {hi hj} :
+    (Vector.mk data size).swap i j hi hj = Vector.mk (data.swap i j) (by simp_all) := rfl
+
 @[simp] theorem getElem_push_last {v : Vector α n} {x : α} : (v.push x)[n] = x := by
   rcases v with ⟨data, rfl⟩
   simp
@@ -93,6 +96,41 @@ defeq issues in the implicit size argument.
     subst h
     simp [pop, back, back!, ← Array.eq_push_pop_back!_of_size_ne_zero]
 
+theorem push_swap (a : Vector α n) (x : α) {i j : Nat} {hi hj} :
+    (a.swap i j hi hj).push x = (a.push x).swap i j := by
+  cases a
+  simp [Array.push_swap]
+
+/-! ### cast -/
+
+@[simp] theorem cast_mk {n m} (a : Array α) (w : a.size = n) (h : n = m) :
+    (Vector.mk a w).cast h = ⟨a, h ▸ w⟩ := by
+  simp [Vector.cast]
+
+@[simp] theorem cast_refl {n} (a : Vector α n) : a.cast rfl = a := by
+  cases a
+  simp
+
+@[simp] theorem toArray_cast {n m} (a : Vector α n) (h : n = m) :
+    (a.cast h).toArray = a.toArray := by
+  subst h
+  simp
+
+theorem cast_inj {n m} (a : Vector α n) (b : Vector α n) (h : n = m) :
+    a.cast h = b.cast h ↔ a = b := by
+  cases h
+  simp
+
+theorem cast_eq_iff {n m} (a : Vector α n) (b : Vector α m) (h : n = m) :
+    a.cast h = b ↔ a = b.cast h.symm := by
+  cases h
+  simp
+
+theorem eq_cast_iff {n m} (a : Vector α n) (b : Vector α m) (h : m = n) :
+    a = b.cast h ↔ a.cast h.symm = b := by
+  cases h
+  simp
+
 /-! ### Decidable quantifiers. -/
 
 theorem forall_zero_iff {P : Vector α 0 → Prop} :

tests/bench/insertionSort/.github/workflows/lean_action_ci.yml

@@ -0,0 +1,14 @@
+name: Lean Action CI
+
+on:
+  push:
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: leanprover/lean-action@v1

tests/bench/insertionSort/.gitignore

@@ -0,0 +1 @@
+/.lake

tests/bench/insertionSort/Main.lean

@@ -0,0 +1,15 @@
+set_option linter.unusedVariables false
+
+abbrev Elem := UInt32
+
+def badRand (seed : Elem) : Elem :=
+seed * 1664525 + 1013904223
+
+def mkRandomArray : Nat → Elem → Array Elem → Array Elem
+| 0,   seed, as => as
+| i+1, seed, as => mkRandomArray i (badRand seed) (as.push seed)
+
+def main (args : List String) : IO UInt32 := do
+  let a := mkRandomArray 20000 0 (Array.mkEmpty 20000)
+  IO.println (a.insertionSort (· < ·)).size
+  return 0

tests/bench/insertionSort/README.md

@@ -0,0 +1 @@
+# insertionSort

tests/bench/insertionSort/lakefile.toml

@@ -0,0 +1,7 @@
+name = "insertionSort"
+version = "0.1.0"
+defaultTargets = ["insertionSort"]
+
+[[lean_exe]]
+name = "insertionSort"
+root = "Main"

tests/bench/insertionSort/lean-toolchain

@@ -0,0 +1 @@
+lean4