perf: improve iterator/range benchmarks, use shortcut instances for Int ranges (#10197)

This PR is the result of analyzing the elaborator performance regression
introduced by #10005. It makes the `workspaceSymboldNewRanges` and
`iterators` benchmarks less noisy. It also replaces some range-related
instances for `Nat` with shortcuts to the general-purpose instances.
This is a trade-off between the ergonomics and the synthesis cost of
having general-purpose instances.

src/Init/Data/Range/Polymorphic/Nat.lean

@@ -6,8 +6,11 @@ Authors: Paul Reichert
 module
 
 prelude
-public import Init.Data.Nat.Lemmas
-public import Init.Data.Range.Polymorphic.Basic
+import Init.Data.Nat.Lemmas
+public import Init.Data.Nat.Order
+public import Init.Data.Range.Polymorphic.Instances
+public import Init.Data.Order.Classes
+import Init.Data.Order.Lemmas
 
 public section
 
@@ -20,6 +23,10 @@ instance : UpwardEnumerable Nat where
 instance : Least? Nat where
   least? := some 0
 
+instance : LawfulUpwardEnumerableLeast? Nat where
+  eq_succMany?_least? a := by
+    simpa [Least?.least?] using ⟨a, by simp [UpwardEnumerable.succMany?]⟩
+
 instance : LawfulUpwardEnumerableLE Nat where
   le_iff a b := by
     constructor
@@ -30,98 +37,29 @@ instance : LawfulUpwardEnumerableLE Nat where
       rw [← hn]
       exact Nat.le_add_right _ _
 
-instance : LawfulUpwardEnumerableLT Nat where
-  lt_iff a b := by
-    constructor
-    · intro h
-      refine ⟨b - a - 1, ?_⟩
-      simp [UpwardEnumerable.succMany?]
-      rw [Nat.sub_add_cancel, Nat.add_sub_cancel']
-      · exact Nat.le_of_lt h
-      · rwa [Nat.lt_iff_add_one_le, ← Nat.le_sub_iff_add_le'] at h
-        exact Nat.le_trans (Nat.le_succ _) h
-    · rintro ⟨n, hn⟩
-      simp only [UpwardEnumerable.succMany?, Option.some.injEq] at hn
-      rw [← hn]
-      apply Nat.lt_add_of_pos_right
-      apply Nat.zero_lt_succ
-
 instance : LawfulUpwardEnumerable Nat where
   succMany?_zero := by simp [UpwardEnumerable.succMany?]
   succMany?_succ := by simp [UpwardEnumerable.succMany?, UpwardEnumerable.succ?, Nat.add_assoc]
   ne_of_lt a b hlt := by
-    rw [← LawfulUpwardEnumerableLT.lt_iff] at hlt
-    exact Nat.ne_of_lt hlt
+    have hn := hlt.choose_spec
+    simp only [UpwardEnumerable.succMany?, Option.some.injEq] at hn
+    omega
 
-instance : LawfulUpwardEnumerableLowerBound .closed Nat where
-  isSatisfied_iff a l := by
-    simp [← LawfulUpwardEnumerableLE.le_iff, BoundedUpwardEnumerable.init?,
-      SupportsLowerBound.IsSatisfied]
-
-instance : LawfulUpwardEnumerableUpperBound .closed Nat where
-  isSatisfied_of_le u a b hub hab := by
-    rw [← LawfulUpwardEnumerableLE.le_iff] at hab
-    exact Nat.le_trans hab hub
-
-instance : LawfulUpwardEnumerableLowerBound .open Nat where
-  isSatisfied_iff a l := by
-    simp [← LawfulUpwardEnumerableLE.le_iff, BoundedUpwardEnumerable.init?,
-      SupportsLowerBound.IsSatisfied, UpwardEnumerable.succ?, Nat.lt_iff_add_one_le]
-
-instance : LawfulUpwardEnumerableUpperBound .open Nat where
-  isSatisfied_of_le u a b hub hab := by
-    rw [← LawfulUpwardEnumerableLE.le_iff] at hab
-    exact Nat.lt_of_le_of_lt hab hub
-
-instance : LawfulUpwardEnumerableLowerBound .unbounded Nat where
-  isSatisfied_iff a l := by
-    simp [← LawfulUpwardEnumerableLE.le_iff, BoundedUpwardEnumerable.init?,
-      SupportsLowerBound.IsSatisfied, Least?.least?]
-
-instance : LawfulUpwardEnumerableUpperBound .unbounded Nat where
-  isSatisfied_of_le _ _ _ _ _ := .intro
-
-instance : LinearlyUpwardEnumerable Nat where
-  eq_of_succ?_eq a b := by simp [UpwardEnumerable.succ?]
+instance : LawfulUpwardEnumerableLT Nat := inferInstance
+instance : LawfulUpwardEnumerableLowerBound .closed Nat := inferInstance
+instance : LawfulUpwardEnumerableUpperBound .closed Nat := inferInstance
+instance : LawfulUpwardEnumerableLowerBound .open Nat := inferInstance
+instance : LawfulUpwardEnumerableUpperBound .open Nat := inferInstance
+instance : LawfulUpwardEnumerableLowerBound .unbounded Nat := inferInstance
+instance : LawfulUpwardEnumerableUpperBound .unbounded Nat := inferInstance
 
 instance : InfinitelyUpwardEnumerable Nat where
   isSome_succ? a := by simp [UpwardEnumerable.succ?]
 
-private def rangeRev (k : Nat) :=
-  match k with
-  | 0 => []
-  | k + 1 => k :: rangeRev k
-
-private theorem mem_rangeRev {k l : Nat} (h : l < k) : l ∈ rangeRev k := by
-  induction k
-  case zero => cases h
-  case succ k ih =>
-    rw [rangeRev]
-    by_cases hl : l = k
-    · simp [hl]
-    · apply List.mem_cons_of_mem
-      exact ih (Nat.lt_of_le_of_ne (Nat.le_of_lt_succ h) hl)
-
-@[no_expose]
-instance : HasFiniteRanges .closed Nat where
-  finite init u := by
-    refine ⟨u - init + 1, ?_⟩
-    simp only [UpwardEnumerable.succMany?, SupportsUpperBound.IsSatisfied, Nat.not_le,
-      Option.elim_some]
-    omega
-
-@[no_expose]
-instance : HasFiniteRanges .open Nat where
-  finite init u := by
-    refine ⟨u - init, ?_⟩
-    simp only [UpwardEnumerable.succMany?, SupportsUpperBound.IsSatisfied, Option.elim_some]
-    omega
-
 instance : RangeSize .closed Nat where
   size bound a := bound + 1 - a
 
-instance : RangeSize .open Nat where
-  size bound a := bound - a
+instance : RangeSize .open Nat := .openOfClosed
 
 instance : LawfulRangeSize .closed Nat where
   size_eq_zero_of_not_isSatisfied upperBound init hu := by
@@ -135,17 +73,16 @@ instance : LawfulRangeSize .closed Nat where
       Option.some.injEq] at hu h ⊢
     omega
 
-instance : LawfulRangeSize .open Nat where
-  size_eq_zero_of_not_isSatisfied upperBound init hu := by
-    simp only [SupportsUpperBound.IsSatisfied, RangeSize.size] at hu ⊢
-    omega
-  size_eq_one_of_succ?_eq_none upperBound init hu h := by
-    simp only [UpwardEnumerable.succ?] at h
-    cases h
-  size_eq_succ_of_succ?_eq_some upperBound init hu h := by
-    simp only [SupportsUpperBound.IsSatisfied, RangeSize.size, UpwardEnumerable.succ?,
-      Option.some.injEq] at hu h ⊢
-    omega
+instance : LawfulRangeSize .open Nat := inferInstance
+instance : HasFiniteRanges .closed Nat := inferInstance
+instance : HasFiniteRanges .open Nat := inferInstance
+instance : LinearlyUpwardEnumerable Nat := by
+  exact instLinearlyUpwardEnumerableOfTotalLeOfLawfulUpwardEnumerableOfLawfulUpwardEnumerableLE
+
+/-!
+The following instances are used for the implementation of array slices a.k.a. `Subarray`.
+See also `Init.Data.Slice.Array`.
+-/
 
 instance : ClosedOpenIntersection ⟨.open, .open⟩ Nat where
   intersection r s := PRange.mk (max (r.lower + 1) s.lower) (min r.upper s.upper)

tests/bench/iterators.lean

@@ -1,5 +1,18 @@
 import Std.Data.Iterators
 
+/-
+This benchmark measures the performance of iterators. The file starts with various function
+declarations. The functions are then called from `main`.
+
+The benchmark is run in three settings.
+
+* The runtime of the compiled program is measured in `iterators (compiled)`.
+* The time taken to interpret the script, including running `main`, is measured in
+  `iterators (interpreted)`.
+* The time taken to interpret the script, without running `main`, is measured in
+  `interators (elab)`.
+-/
+
 /- definitions -/
 
 def sum₁ (xs : Array Nat) : Nat :=
@@ -69,38 +82,31 @@ def longChainOfCombinators (xs : Array Nat) : Nat :=
     |>.takeWhile (fun x => x < 5000000)
     |>.fold (init := 0) (· + ·)
 
-/- evaluations -/
-
 def xs : Array Nat := (*...100000).iter.toArray
 
 def l : List Nat := (*...100000).iter.toList
 
-#eval sum₁ xs
+/- evaluations -/
 
-#eval sum₂ xs
+@[noinline]
+def run' (f : Unit → α) : IO α := do
+  return f ()
 
-#eval isolatedMap xs
+notation "run " t => (fun _ => ()) <$> run' fun _ => t
 
-#eval isolatedFilterMap xs
-
-#eval isolatedTake xs 1000000
-
-#eval isolatedDrop xs 100000
-
-#eval isolatedTakeWhile xs
-
-#eval isolatedDropWhile xs
-
-#eval isolatedZip xs xs
-
-#eval isolatedSteppedRange 1000000
-
-#eval longChainOfCombinators xs
-
-#eval (*...1000000).iter.fold (init := 0) (· + ·)
-
-#eval primes 3000
-
-#eval printEveryNth l 10000
-
-#eval printEveryNthSliceBased xs 10000
+def main : IO Unit := do
+  run sum₁ xs
+  run sum₂ xs
+  run isolatedMap xs
+  run isolatedFilterMap xs
+  run isolatedTake xs 1000000
+  run isolatedDrop xs 1000000
+  run isolatedTakeWhile xs
+  run isolatedDropWhile xs
+  run isolatedZip xs xs
+  run isolatedSteppedRange 1000000
+  run longChainOfCombinators xs
+  run (*...1000000).iter.fold (init := 0) (· + ·)
+  run primes 3000
+  printEveryNth l 10000
+  printEveryNthSliceBased xs 10000

tests/bench/speedcenter.exec.velcom.yaml

@@ -568,7 +568,21 @@
     cmd: lean riscv-ast.lean
     max_runs: 2
 - attributes:
-    description: iterators
+    description: iterators (compiled)
+    tags: [fast]
+  run_config:
+    <<: *time
+    cmd: ./iterators.lean.out
+  build_config:
+    cmd: ./compile.sh iterators.lean
+- attributes:
+    description: iterators (interpreted)
+    tags: [fast]
+  run_config:
+    <<: *time
+    cmd: lean --run iterators.lean
+- attributes:
+    description: iterators (elab)
     tags: [fast]
   run_config:
     <<: *time
@@ -578,7 +592,9 @@
     tags: [fast]
   run_config:
     <<: *time
-    cmd: lean workspaceSymbolsNewRanges.lean
+    cmd: ./workspaceSymbolsNewRanges.lean.out
+  build_config:
+    cmd: ./compile.sh workspaceSymbolsNewRanges.lean
 - attributes:
     description: hashmap.lean
     tags: [fast]