perf: improve kernel reduction of List.sublists (leanprover-communi…

…ty#301)

Std/Data/List/Basic.lean

@@ -889,10 +889,20 @@ sublists [1, 2, 3] = [[], [1], [2], [1, 2], [3], [1, 3], [2, 3], [1, 2, 3]]
 ```
 -/
 def sublists (l : List α) : List (List α) :=
+  l.foldr (fun a acc => acc.bind fun x => [x, a :: x]) [[]]
+
+/-- A version of `List.sublists` that has faster runtime performance but worse kernel performance -/
+def sublistsFast (l : List α) : List (List α) :=
   let f a arr := arr.foldl (init := Array.mkEmpty (arr.size * 2))
     fun r l => (r.push l).push (a :: l)
   (l.foldr f #[[]]).toList
 
+-- The fact that this transformation is safe is proved in mathlib4 as `sublists_eq_sublistsFast`.
+-- Using a `csimp` lemma here is impractical as we are missing a lot of lemmas about lists.
+-- TODO(std4#307): upstream the necessary results about `sublists` and put the `csimp` lemma in
+-- `Std/Data/List/Lemmas.lean`.
+attribute [implemented_by sublistsFast] sublists
+
 section Forall₂
 
 variable {r : α → β → Prop} {p : γ → δ → Prop}

test/list_sublists.lean

@@ -0,0 +1,10 @@
+import Std.Data.List.Basic
+
+-- this times out with `sublistsFast`
+set_option maxRecDepth 561 in
+example : [1, 2, 3].sublists.sublists.length = 256 := rfl
+
+-- TODO(std4#307): until we have the `csimp` lemma in std, this is a sanity check that these two
+-- are equal.
+example : ([] : List Nat).sublists = [].sublistsFast := rfl
+example : [1, 2, 3].sublists = [1, 2, 3].sublistsFast := rfl