feat: upstream more List operations

src/Init/Data/List/Basic.lean

@@ -27,24 +27,32 @@ Recall that `length`, `get`, `set`, `foldl`, and `concat` have already been defi
 The operations are organized as follow:
 * Equality: `beq`, `isEqv`.
 * Lexicographic ordering: `lt`, `le`, and instances.
+* Head and tail operators: `head`, `head?`, `headD?`, `tail`, `tail?`, `tailD`.
 * Basic operations:
-  `map`, `filter`, `filterMap`, `foldr`, `append`, `join`, `pure`, `bind`, `replicate`, and `reverse`.
+  `map`, `filter`, `filterMap`, `foldr`, `append`, `join`, `pure`, `bind`, `replicate`, and
+  `reverse`.
+* Additional functions defined in terms of these: `leftpad`, `rightPad`, and `reduceOption`.
 * List membership: `isEmpty`, `elem`, `contains`, `mem` (and the `∈` notation),
   and decidability for predicates quantifying over membership in a `List`.
 * Sublists: `take`, `drop`, `takeWhile`, `dropWhile`, `partition`, `dropLast`,
-  `isPrefixOf`, `isPrefixOf?`, `isSuffixOf`, `isSuffixOf?`, `Subset`, `Sublist`, `rotateLeft` and `rotateRight`.
-* Manipulating elements: `replace`, `insert`, `erase`, `eraseP`, `eraseIdx`, `find?`, `findSome?`, and `lookup`.
+  `isPrefixOf`, `isPrefixOf?`, `isSuffixOf`, `isSuffixOf?`, `Subset`, `Sublist`,
+  `rotateLeft` and `rotateRight`.
+* Manipulating elements: `replace`, `insert`, `erase`, `eraseP`, `eraseIdx`.
+* Finding elements: `find?`, `findSome?`, `findIdx`, `indexOf`, `findIdx?`, `indexOf?`,
+ `countP`, `count`, and `lookup`.
 * Logic: `any`, `all`, `or`, and `and`.
 * Zippers: `zipWith`, `zip`, `zipWithAll`, and `unzip`.
 * Ranges and enumeration: `range`, `iota`, `enumFrom`, and `enum`.
 * Minima and maxima: `minimum?` and `maximum?`.
-* Other functions: `intersperse`, `intercalate`, `eraseDups`, `eraseReps`, `span`, `groupBy`, `removeAll`
+* Other functions: `intersperse`, `intercalate`, `eraseDups`, `eraseReps`, `span`, `groupBy`,
+  `removeAll`
   (currently these functions are mostly only used in meta code,
   and do not have API suitable for verification).
 
-Further operations are defined in `Init.Data.List.BasicAux` (because they use `Array` in their implementations), namely:
+Further operations are defined in `Init.Data.List.BasicAux`
+(because they use `Array` in their implementations), namely:
 * Variant getters: `get!`, `get?`, `getD`, `getLast`, `getLast!`, `getLast?`, and `getLastD`.
-* Head and tail: `head`, `head!`, `head?`, `headD`, `tail!`, `tail?`, and `tailD`.
+* Head and tail: `head!`, `tail!`.
 * Other operations on sublists: `partitionMap`, `rotateLeft`, and `rotateRight`.
 -/
 
@@ -315,6 +323,16 @@ def headD : (as : List α) → (fallback : α) → α
 @[simp 1100] theorem headD_nil : @headD α [] d = d := rfl
 @[simp 1100] theorem headD_cons : @headD α (a::l) d = a := rfl
 
+/-! ### tail -/
+
+/-- Get the tail of a nonempty list, or return `[]` for `[]`. -/
+def tail : List α → List α
+  | []    => []
+  | _::as => as
+
+@[simp] theorem tail_nil : @tail α [] = [] := rfl
+@[simp] theorem tail_cons : @tail α (a::as) = as := rfl
+
 /-! ### tail? -/
 
 /--
@@ -577,6 +595,28 @@ theorem replicate_succ (a : α) (n) : replicate (n+1) a = a :: replicate n a :=
   | zero => simp
   | succ n ih => simp only [ih, replicate_succ, length_cons, Nat.succ_eq_add_one]
 
+/-! ## Additional functions -/
+
+/-! ### leftpad and rightpad -/
+
+/--
+Pads `l : List α` on the left with repeated occurrences of `a : α` until it is of length `n`.
+If `l` is initially larger than `n`, just return `l`.
+-/
+def leftpad (n : Nat) (a : α) (l : List α) : List α := replicate (n - length l) a ++ l
+
+/--
+Pads `l : List α` on the right with repeated occurrences of `a : α` until it is of length `n`.
+If `l` is initially larger than `n`, just return `l`.
+-/
+def rightpad (n : Nat) (a : α) (l : List α) : List α := l ++ replicate (n - length l) a
+
+/-! ### reduceOption-/
+
+/-- Drop `none`s from a list, and replace each remaining `some a` with `a`. -/
+@[inline] def reduceOption {α} : List (Option α) → List α :=
+  List.filterMap id
+
 /-! ## List membership
 
 * `L.contains a : Bool` determines, using a `[BEq α]` instance, whether `L` contains an element `· == a`.
@@ -1080,6 +1120,8 @@ def eraseIdx : List α → Nat → List α
 @[simp] theorem eraseIdx_cons_zero : (a::as).eraseIdx 0 = as := rfl
 @[simp] theorem eraseIdx_cons_succ : (a::as).eraseIdx (i+1) = a :: as.eraseIdx i := rfl
 
+/-! Finding elements -/
+
 /-! ### find? -/
 
 /--
@@ -1117,6 +1159,46 @@ theorem findSome?_cons {f : α → Option β} :
     (a::as).findSome? f = match f a with | some b => some b | none => as.findSome? f :=
   rfl
 
+/-! ### findIdx -/
+
+/-- Returns the index of the first element satisfying `p`, or the length of the list otherwise. -/
+@[inline] def findIdx (p : α → Bool) (l : List α) : Nat := go l 0 where
+  /-- Auxiliary for `findIdx`: `findIdx.go p l n = findIdx p l + n` -/
+  @[specialize] go : List α → Nat → Nat
+  | [], n => n
+  | a :: l, n => bif p a then n else go l (n + 1)
+
+/-! ### indexOf -/
+
+/-- Returns the index of the first element equal to `a`, or the length of the list otherwise. -/
+def indexOf [BEq α] (a : α) : List α → Nat := findIdx (· == a)
+
+/-! ### findIdx? -/
+
+/-- Return the index of the first occurrence of an element satisfying `p`. -/
+def findIdx? (p : α → Bool) : List α → (start : Nat := 0) → Option Nat
+| [], _ => none
+| a :: l, i => if p a then some i else findIdx? p l (i + 1)
+
+/-! ### indexOf? -/
+
+/-- Return the index of the first occurrence of `a` in the list. -/
+@[inline] def indexOf? [BEq α] (a : α) : List α → Option Nat := findIdx? (· == a)
+
+/-! ### countP -/
+
+/-- `countP p l` is the number of elements of `l` that satisfy `p`. -/
+@[inline] def countP (p : α → Bool) (l : List α) : Nat := go l 0 where
+  /-- Auxiliary for `countP`: `countP.go p l acc = countP p l + acc`. -/
+  @[specialize] go : List α → Nat → Nat
+  | [], acc => acc
+  | x :: xs, acc => bif p x then go xs (acc + 1) else go xs acc
+
+/-! ### count -/
+
+/-- `count a l` is the number of occurrences of `a` in `l`. -/
+@[inline] def count [BEq α] (a : α) : List α → Nat := countP (· == a)
+
 /-! ### lookup -/
 
 /--
@@ -1281,6 +1363,14 @@ where
 
 @[simp] theorem range_zero : range 0 = [] := rfl
 
+/-! ### range' -/
+
+/-- `range' start len step` is the list of numbers `[start, start+step, ..., start+(len-1)*step]`.
+  It is intended mainly for proving properties of `range` and `iota`. -/
+def range' : (start len : Nat) → (step : Nat := 1) → List Nat
+  | _, 0, _ => []
+  | s, n+1, step => s :: range' (s+step) n step
+
 /-! ### iota -/
 
 /--

src/Init/Data/List/Impl.lean

@@ -193,6 +193,17 @@ theorem replicateTR_loop_eq : ∀ n, replicateTR.loop a n acc = replicate n a ++
   apply funext; intro α; apply funext; intro n; apply funext; intro a
   exact (replicateTR_loop_replicate_eq _ 0 n).symm
 
+/-! ## Additional functions -/
+
+/-! ### leftpad -/
+
+/-- Optimized version of `leftpad`. -/
+@[inline] def leftpadTR (n : Nat) (a : α) (l : List α) : List α :=
+  replicateTR.loop a (n - length l) l
+
+@[csimp] theorem leftpad_eq_leftpadTR : @leftpad = @leftpadTR := by
+  funext α n a l; simp [leftpad, leftpadTR, replicateTR_loop_eq]
+
 /-! ## Sublists -/
 
 /-! ### take -/
@@ -366,6 +377,26 @@ def unzipTR (l : List (α × β)) : List α × List β :=
 
 /-! ## Ranges and enumeration -/
 
+/-! ### range' -/
+
+/-- Optimized version of `range'`. -/
+@[inline] def range'TR (s n : Nat) (step : Nat := 1) : List Nat := go n (s + step * n) [] where
+  /-- Auxiliary for `range'TR`: `range'TR.go n e = [e-n, ..., e-1] ++ acc`. -/
+  go : Nat → Nat → List Nat → List Nat
+  | 0, _, acc => acc
+  | n+1, e, acc => go n (e-step) ((e-step) :: acc)
+
+@[csimp] theorem range'_eq_range'TR : @range' = @range'TR := by
+  funext s n step
+  let rec go (s) : ∀ n m,
+    range'TR.go step n (s + step * n) (range' (s + step * n) m step) = range' s (n + m) step
+  | 0, m => by simp [range'TR.go]
+  | n+1, m => by
+    simp [range'TR.go]
+    rw [Nat.mul_succ, ← Nat.add_assoc, Nat.add_sub_cancel, Nat.add_right_comm n]
+    exact go s n (m + 1)
+  exact (go s n 0).symm
+
 /-! ### iota -/
 
 /-- Tail-recursive version of `List.iota`. -/

src/Init/Data/List/Lemmas.lean

@@ -2577,6 +2577,10 @@ theorem cons_subset_cons {l₁ l₂ : List α} (a : α) (s : l₁ ⊆ l₂) : a
 theorem Subset.map {l₁ l₂ : List α} (f : α → β) (H : l₁ ⊆ l₂) : map f l₁ ⊆ map f l₂ :=
   fun x => by simp only [mem_map]; exact .imp fun a => .imp_left (@H _)
 
+@[deprecated (since := "2024-07-28")]
+theorem map_subset {l₁ l₂ : List α} {f : α → β} (h : l₁ ⊆ l₂) : map f l₁ ⊆ map f l₂ :=
+  Subset.map f h
+
 theorem Subset.filter {l₁ l₂ : List α} (p : α → Bool) (H : l₁ ⊆ l₂) : filter p l₁ ⊆ filter p l₂ :=
   fun x => by simp_all [mem_filter, subset_def.1 H]
 

tests/lean/run/starsAndBars.lean

@@ -1,6 +1,6 @@
-@[simp] def List.count [DecidableEq α] : List α → α → Nat
+@[simp] def List.count' [DecidableEq α] : List α → α → Nat
   | [], _    => 0
-  | a::as, b => if a = b then as.count b + 1 else as.count b
+  | a::as, b => if a = b then as.count' b + 1 else as.count' b
 
 inductive StarsAndBars : Nat → Nat → Type where
   | nil  : StarsAndBars 0 0
@@ -11,7 +11,7 @@ namespace StarsAndBars
 
 namespace Ex1
 
-def toList : StarsAndBars s b → { bs : List Bool // bs.count false = s ∧ bs.count true = b }
+def toList : StarsAndBars s b → { bs : List Bool // bs.count' false = s ∧ bs.count' true = b }
   | nil    => ⟨[], by simp⟩
   | star h => ⟨false :: toList h, by simp [(toList h).2]⟩
   | bar  h => ⟨true :: toList h, by simp [(toList h).2]⟩
@@ -27,13 +27,13 @@ end Ex1
 /- Same example with explicit proofs -/
 namespace Ex2
 
-theorem toList_star_proof (h : bs.count false = s ∧ bs.count true = b) : (false :: bs).count false = s.succ ∧ (false :: bs).count true = b := by
+theorem toList_star_proof (h : bs.count' false = s ∧ bs.count' true = b) : (false :: bs).count' false = s.succ ∧ (false :: bs).count' true = b := by
   simp [*]
 
-theorem toList_bar_proof (h : bs.count false = s ∧ bs.count true = b) : (true :: bs).count false = s ∧ (true :: bs).count true = b.succ := by
+theorem toList_bar_proof (h : bs.count' false = s ∧ bs.count' true = b) : (true :: bs).count' false = s ∧ (true :: bs).count' true = b.succ := by
   simp [*]
 
-def toList : StarsAndBars s b → { bs : List Bool // bs.count false = s ∧ bs.count true = b }
+def toList : StarsAndBars s b → { bs : List Bool // bs.count' false = s ∧ bs.count' true = b }
   | nil    => ⟨[], by simp⟩
   | star h => ⟨false :: toList h, toList_star_proof (toList h).property⟩
   | bar  h => ⟨true :: toList h,  toList_bar_proof (toList h).property⟩