feat: more List.Sublist theorems (#5048)

src/Init/Data/List/Count.lean

@@ -81,6 +81,10 @@ theorem Sublist.countP_le (s : l₁ <+ l₂) : countP p l₁ ≤ countP p l₂ :
   simp only [countP_eq_length_filter]
   apply s.filter _ |>.length_le
 
+theorem IsPrefix.countP_le (s : l₁ <+: l₂) : countP p l₁ ≤ countP p l₂ := s.sublist.countP_le _
+theorem IsSuffix.countP_le (s : l₁ <:+ l₂) : countP p l₁ ≤ countP p l₂ := s.sublist.countP_le _
+theorem IsInfix.countP_le (s : l₁ <:+: l₂) : countP p l₁ ≤ countP p l₂ := s.sublist.countP_le _
+
 theorem countP_filter (l : List α) :
     countP p (filter q l) = countP (fun a => p a ∧ q a) l := by
   simp only [countP_eq_length_filter, filter_filter]
@@ -140,6 +144,10 @@ theorem count_le_length (a : α) (l : List α) : count a l ≤ l.length := count
 
 theorem Sublist.count_le (h : l₁ <+ l₂) (a : α) : count a l₁ ≤ count a l₂ := h.countP_le _
 
+theorem IsPrefix.count_le (h : l₁ <+: l₂) (a : α) : count a l₁ ≤ count a l₂ := h.sublist.count_le _
+theorem IsSuffix.count_le (h : l₁ <:+ l₂) (a : α) : count a l₁ ≤ count a l₂ := h.sublist.count_le _
+theorem IsInfix.count_le (h : l₁ <:+: l₂) (a : α) : count a l₁ ≤ count a l₂ := h.sublist.count_le _
+
 theorem count_le_count_cons (a b : α) (l : List α) : count a l ≤ count a (b :: l) :=
   (sublist_cons_self _ _).count_le _
 

src/Init/Data/List/Lemmas.lean

@@ -1463,6 +1463,24 @@ theorem append_eq_append_iff {a b c d : List α} :
   | nil => simp_all
   | cons a as ih => cases c <;> simp [eq_comm, and_assoc, ih, and_or_left]
 
+theorem append_inj_of_length_left {a b c d : List α}
+    (h : a ++ b = c ++ d) (hl : length a = length c) : a = c ∧ b = d := by
+  rcases append_eq_append_iff.mp h with (⟨a', rfl, rfl⟩ | ⟨c', rfl, rfl⟩)
+  · simp only [length_append] at hl
+    have : a'.length = 0 := (Nat.add_left_cancel hl).symm
+    simp_all
+  · simp only [length_append] at hl
+    have : c'.length = 0 := (Nat.add_left_cancel hl.symm).symm
+    simp_all
+
+theorem append_inj_of_length_right {a b c d : List α}
+    (h : a ++ b = c ++ d) (hl : length b = length d) : a = c ∧ b = d := by
+  have : length a = length c :=  by
+    replace h := congrArg length h
+    simp only [length_append, hl] at h
+    exact Nat.add_right_cancel h
+  exact append_inj_of_length_left h this
+
 @[simp] theorem mem_append {a : α} {s t : List α} : a ∈ s ++ t ↔ a ∈ s ∨ a ∈ t := by
   induction s <;> simp_all [or_assoc]
 
@@ -1515,9 +1533,55 @@ theorem set_append {s t : List α} :
 @[simp] theorem foldr_append (f : α → β → β) (b) (l l' : List α) :
     (l ++ l').foldr f b = l.foldr f (l'.foldr f b) := by simp [foldr_eq_foldrM]
 
+theorem filterMap_eq_append (f : α → Option β) :
+    filterMap f l = L₁ ++ L₂ ↔ ∃ l₁ l₂, l = l₁ ++ l₂ ∧ filterMap f l₁ = L₁ ∧ filterMap f l₂ = L₂ := by
+  constructor
+  · induction l generalizing L₁ with
+    | nil =>
+      simp only [filterMap_nil, nil_eq_append, and_imp]
+      rintro rfl rfl
+      exact ⟨[], [], by simp⟩
+    | cons x l ih =>
+      simp only [filterMap_cons]
+      split
+      · intro h
+        obtain ⟨l₁, l₂, rfl, rfl, rfl⟩ := ih h
+        refine ⟨x :: l₁, l₂, ?_⟩
+        simp_all
+      · rename_i b w
+        intro h
+        rcases cons_eq_append.mp h with (⟨rfl, rfl⟩ | ⟨L₁, ⟨rfl, h⟩⟩)
+        · refine ⟨[], x :: l, ?_⟩
+          simp [filterMap_cons, w]
+        · obtain ⟨l₁, l₂, rfl, rfl, rfl⟩ := ih ‹_›
+          refine ⟨x :: l₁, l₂, ?_⟩
+          simp [filterMap_cons, w]
+  · rintro ⟨l₁, l₂, rfl, rfl, rfl⟩
+    simp
+
+theorem append_eq_filterMap (f : α → Option β) :
+    L₁ ++ L₂ = filterMap f l ↔ ∃ l₁ l₂, l = l₁ ++ l₂ ∧ filterMap f l₁ = L₁ ∧ filterMap f l₂ = L₂ := by
+  rw [eq_comm, filterMap_eq_append]
+
+theorem filter_eq_append (p : α → Bool) :
+    filter p l = L₁ ++ L₂ ↔ ∃ l₁ l₂, l = l₁ ++ l₂ ∧ filter p l₁ = L₁ ∧ filter p l₂ = L₂ := by
+  rw [← filterMap_eq_filter, filterMap_eq_append]
+
+theorem append_eq_filter (p : α → Bool) :
+    L₁ ++ L₂ = filter p l ↔ ∃ l₁ l₂, l = l₁ ++ l₂ ∧ filter p l₁ = L₁ ∧ filter p l₂ = L₂ := by
+  rw [eq_comm, filter_eq_append]
+
 @[simp] theorem map_append (f : α → β) : ∀ l₁ l₂, map f (l₁ ++ l₂) = map f l₁ ++ map f l₂ := by
   intro l₁; induction l₁ <;> intros <;> simp_all
 
+theorem map_eq_append (f : α → β) :
+    map f l = L₁ ++ L₂ ↔ ∃ l₁ l₂, l = l₁ ++ l₂ ∧ map f l₁ = L₁ ∧ map f l₂ = L₂ := by
+  rw [← filterMap_eq_map, filterMap_eq_append]
+
+theorem append_eq_map (f : α → β) :
+    L₁ ++ L₂ = map f l ↔ ∃ l₁ l₂, l = l₁ ++ l₂ ∧ map f l₁ = L₁ ∧ map f l₂ = L₂ := by
+  rw [eq_comm, map_eq_append]
+
 /-! ### concat
 
 Note that `concat_eq_append` is a `@[simp]` lemma, so `concat` should usually not appear in goals.
@@ -1571,6 +1635,8 @@ theorem eq_nil_or_concat : ∀ l : List α, l = [] ∨ ∃ L b, l = concat L b
   | cons =>
     simp [join, length_append, *]
 
+theorem join_singleton (l : List α) : [l].join = l := by simp
+
 @[simp] theorem mem_join : ∀ {L : List (List α)}, a ∈ L.join ↔ ∃ l, l ∈ L ∧ a ∈ l
   | [] => by simp
   | b :: l => by simp [mem_join, or_and_right, exists_or]
@@ -1618,6 +1684,20 @@ theorem foldr_join (f : α → β → β) (b : β) (L : List (List α)) :
     filter p (join L) = join (map (filter p) L) := by
   induction L <;> simp [*, filter_append]
 
+@[simp]
+theorem join_filter_not_isEmpty  :
+    ∀ {L : List (List α)}, join (L.filter fun l => !l.isEmpty) = L.join
+  | [] => rfl
+  | [] :: L
+  | (a :: l) :: L => by
+      simp [join_filter_not_isEmpty (L := L)]
+
+@[simp]
+theorem join_filter_ne_nil [DecidablePred fun l : List α => l ≠ []] {L : List (List α)} :
+    join (L.filter fun l => l ≠ []) = L.join := by
+  simp only [ne_eq, ← isEmpty_iff, Bool.not_eq_true, Bool.decide_eq_false,
+    join_filter_not_isEmpty]
+
 @[simp] theorem join_append (L₁ L₂ : List (List α)) : join (L₁ ++ L₂) = join L₁ ++ join L₂ := by
   induction L₁ <;> simp_all
 
@@ -1627,6 +1707,22 @@ theorem join_concat (L : List (List α)) (l : List α) : join (L ++ [l]) = join
 theorem join_join {L : List (List (List α))} : join (join L) = join (map join L) := by
   induction L <;> simp_all
 
+/-- Two lists of sublists are equal iff their joins coincide, as well as the lengths of the
+sublists. -/
+theorem eq_iff_join_eq : ∀ (L L' : List (List α)),
+    L = L' ↔ L.join = L'.join ∧ map length L = map length L'
+  | _, [] => by simp_all
+  | [], x' :: L' => by simp_all
+  | x :: L, x' :: L' => by
+    simp
+    rw [eq_iff_join_eq]
+    constructor
+    · rintro ⟨rfl, h₁, h₂⟩
+      simp_all
+    · rintro ⟨h₁, h₂, h₃⟩
+      obtain ⟨rfl, h⟩ := append_inj_of_length_left h₁ h₂
+      exact ⟨rfl, h, h₃⟩
+
 /-! ### bind -/
 
 theorem bind_def (l : List α) (f : α → List β) : l.bind f = join (map f l) := by rfl
@@ -1643,6 +1739,11 @@ theorem exists_of_mem_bind {b : β} {l : List α} {f : α → List β} :
 theorem mem_bind_of_mem {b : β} {l : List α} {f : α → List β} {a} (al : a ∈ l) (h : b ∈ f a) :
     b ∈ l.bind f := mem_bind.2 ⟨a, al, h⟩
 
+@[simp]
+theorem bind_eq_nil {l : List α} {f : α → List β} : List.bind l f = [] ↔ ∀ x ∈ l, f x = [] :=
+  join_eq_nil_iff.trans <| by
+    simp only [mem_map, forall_exists_index, and_imp, forall_apply_eq_imp_iff₂]
+
 theorem forall_mem_bind {p : β → Prop} {l : List α} {f : α → List β} :
     (∀ (x) (_ : x ∈ l.bind f), p x) ↔ ∀ (a) (_ : a ∈ l) (b) (_ : b ∈ f a), p b := by
   simp only [mem_bind, forall_exists_index, and_imp]
@@ -1798,6 +1899,14 @@ theorem map_const' (l : List α) (b : β) : map (fun _ => b) l = replicate l.len
     simp only [mem_append, mem_replicate, ne_eq]
     rintro (⟨-, rfl⟩ | ⟨_, rfl⟩) <;> rfl
 
+theorem append_eq_replicate {l₁ l₂ : List α} {a : α} :
+    l₁ ++ l₂ = replicate n a ↔
+      l₁.length + l₂.length = n ∧ l₁ = replicate l₁.length a ∧ l₂ = replicate l₂.length a := by
+  simp only [eq_replicate, length_append, mem_append, true_and, and_congr_right_iff]
+  exact fun _ =>
+    { mp := fun h => ⟨fun b m => h b (Or.inl m), fun b m => h b (Or.inr m)⟩,
+      mpr := fun h b x => Or.casesOn x (fun m => h.left b m) fun m => h.right b m }
+
 @[simp] theorem map_replicate : (replicate n a).map f = replicate n (f a) := by
   ext1 n
   simp only [getElem?_map, getElem?_replicate]

src/Init/Data/List/Sublist.lean

@@ -790,6 +790,86 @@ theorem isPrefix_iff : l₁ <+: l₂ ↔ ∀ i (h : i < l₁.length), l₂[i]? =
 
 -- See `Init.Data.List.Nat.Sublist` for `isSuffix_iff` and `ifInfix_iff`.
 
+theorem isPrefix_filterMap_iff {β} (f : α → Option β) {l₁ : List α} {l₂ : List β} :
+    l₂ <+: filterMap f l₁ ↔ ∃ l, l <+: l₁ ∧ l₂ = filterMap f l := by
+  simp only [IsPrefix, append_eq_filterMap]
+  constructor
+  · rintro ⟨_, l₁, l₂, rfl, rfl, rfl⟩
+    exact ⟨l₁, ⟨l₂, rfl⟩, rfl⟩
+  · rintro ⟨l₁, ⟨l₂, rfl⟩, rfl⟩
+    exact ⟨_, l₁, l₂, rfl, rfl, rfl⟩
+
+theorem isSuffix_filterMap_iff {β} (f : α → Option β) {l₁ : List α} {l₂ : List β} :
+    l₂ <:+ filterMap f l₁ ↔ ∃ l, l <:+ l₁ ∧ l₂ = filterMap f l := by
+  simp only [IsSuffix, append_eq_filterMap]
+  constructor
+  · rintro ⟨_, l₁, l₂, rfl, rfl, rfl⟩
+    exact ⟨l₂, ⟨l₁, rfl⟩, rfl⟩
+  · rintro ⟨l₁, ⟨l₂, rfl⟩, rfl⟩
+    exact ⟨_, l₂, l₁, rfl, rfl, rfl⟩
+
+theorem isInfix_filterMap_iff {β} (f : α → Option β) {l₁ : List α} {l₂ : List β} :
+    l₂ <:+: filterMap f l₁ ↔ ∃ l, l <:+: l₁ ∧ l₂ = filterMap f l := by
+  simp only [IsInfix, append_eq_filterMap, filterMap_eq_append]
+  constructor
+  · rintro ⟨_, _, _, l₁, rfl, ⟨⟨l₂, l₃, rfl, rfl, rfl⟩, rfl⟩⟩
+    exact ⟨l₃, ⟨l₂, l₁, rfl⟩, rfl⟩
+  · rintro ⟨l₃, ⟨l₂, l₁, rfl⟩, rfl⟩
+    exact ⟨_, _, _, l₁, rfl, ⟨⟨l₂, l₃, rfl, rfl, rfl⟩, rfl⟩⟩
+
+theorem isPrefix_filter_iff {p : α → Bool} {l₁ l₂ : List α} :
+    l₂ <+: l₁.filter p ↔ ∃ l, l <+: l₁ ∧ l₂ = l.filter p := by
+  rw [← filterMap_eq_filter, isPrefix_filterMap_iff]
+
+theorem isSuffix_filter_iff {p : α → Bool} {l₁ l₂ : List α} :
+    l₂ <:+ l₁.filter p ↔ ∃ l, l <:+ l₁ ∧ l₂ = l.filter p := by
+  rw [← filterMap_eq_filter, isSuffix_filterMap_iff]
+
+theorem isInfix_filter_iff {p : α → Bool} {l₁ l₂ : List α} :
+    l₂ <:+: l₁.filter p ↔ ∃ l, l <:+: l₁ ∧ l₂ = l.filter p := by
+  rw [← filterMap_eq_filter, isInfix_filterMap_iff]
+
+theorem isPrefix_map_iff {β} (f : α → β) {l₁ : List α} {l₂ : List β} :
+    l₂ <+: l₁.map f ↔ ∃ l, l <+: l₁ ∧ l₂ = l.map f := by
+  rw [← filterMap_eq_map, isPrefix_filterMap_iff]
+
+theorem isSuffix_map_iff {β} (f : α → β) {l₁ : List α} {l₂ : List β} :
+    l₂ <:+ l₁.map f ↔ ∃ l, l <:+ l₁ ∧ l₂ = l.map f := by
+  rw [← filterMap_eq_map, isSuffix_filterMap_iff]
+
+theorem isInfix_map_iff {β} (f : α → β) {l₁ : List α} {l₂ : List β} :
+    l₂ <:+: l₁.map f ↔ ∃ l, l <:+: l₁ ∧ l₂ = l.map f := by
+  rw [← filterMap_eq_map, isInfix_filterMap_iff]
+
+theorem isPrefix_replicate_iff {n} {a : α} {l : List α} :
+    l <+: List.replicate n a ↔ l.length ≤ n ∧ l = List.replicate l.length a := by
+  rw [IsPrefix]
+  simp only [append_eq_replicate]
+  constructor
+  · rintro ⟨_, rfl, _, _⟩
+    exact ⟨le_add_right .., ‹_›⟩
+  · rintro ⟨h, w⟩
+    refine ⟨replicate (n - l.length) a, ?_, ?_⟩
+    · simpa using add_sub_of_le h
+    · simpa using w
+
+theorem isSuffix_replicate_iff {n} {a : α} {l : List α} :
+    l <:+ List.replicate n a ↔ l.length ≤ n ∧ l = List.replicate l.length a := by
+  rw [← reverse_prefix, reverse_replicate, isPrefix_replicate_iff]
+  simp [reverse_eq_iff]
+
+theorem isInfix_replicate_iff {n} {a : α} {l : List α} :
+    l <:+: List.replicate n a ↔ l.length ≤ n ∧ l = List.replicate l.length a := by
+  rw [IsInfix]
+  simp only [append_eq_replicate, length_append]
+  constructor
+  · rintro ⟨_, _, rfl, ⟨-, _, _⟩, _⟩
+    exact ⟨le_add_right_of_le (le_add_left ..), ‹_›⟩
+  · rintro ⟨h, w⟩
+    refine ⟨replicate (n - l.length) a, [], ?_, ?_⟩
+    · simpa using Nat.sub_add_cancel h
+    · simpa using w
+
 theorem infix_of_mem_join : ∀ {L : List (List α)}, l ∈ L → l <:+: join L
   | l' :: _, h =>
     match h with