perf: fold Task.get (Task.pure x) to just x (#12446)

This PR adds a simplification rule for `Task.get (Task.pure x) = x` into
the LCNF simplifier. This
ensures that we avoid touching the runtime for a `Task` that instantly
gets destructed anyways.

src/Lean/Compiler/LCNF/Simp/ConstantFold.lean

@@ -434,6 +434,18 @@ def stringFolders : List (Name × Folder) := [
   (``String.push, Folder.mkBinary String.push)
 ]
 
+def foldTaskGet (args : Array (Arg .pure)) : FolderM (Option (LetValue .pure)) := do
+  let #[_, .fvar taskFVar] := args | return none
+  let some (.const ``Task.pure _ #[_, val]) ← findLetValue? (pu := .pure) taskFVar | return none
+  match val with
+  | .erased => return some .erased
+  | .fvar fvarId => return some (.fvar fvarId #[])
+  | _ => return none
+
+def taskFolders : List (Name × Folder) := [
+  (``Task.get, foldTaskGet)
+]
+
 /--
 Apply all known folders to `decl`.
 -/
@@ -457,7 +469,12 @@ private def getFolder (declName : Name) : CoreM Folder := do
   ofExcept <| getFolderCore (← getEnv) (← getOptions) declName
 
 def builtinFolders : SMap Name Folder :=
-  (arithmeticFolders ++ relationFolders ++ conversionFolders ++ higherOrderLiteralFolders ++ stringFolders).foldl (init := {}) fun s (declName, folder) =>
+  (arithmeticFolders
+    ++ relationFolders
+    ++ conversionFolders
+    ++ higherOrderLiteralFolders
+    ++ stringFolders
+    ++ taskFolders).foldl (init := {}) fun s (declName, folder) =>
     s.insert declName folder
 
 structure FolderOleanEntry where

tests/lean/run/constant_fold_task_get_pure.lean

@@ -0,0 +1,12 @@
+/-! This test asserts that we detect the pattern of calling `get` on a `Task.pure` and remove
+  `Task.pure` in order to avoid interaction with the runtime. -/
+
+/--
+trace: [Compiler.saveMono] size: 0
+    def test n : Nat :=
+      return n
+-/
+#guard_msgs in
+set_option trace.Compiler.saveMono true in
+def test (n : Nat) : Nat :=
+  (Task.pure n).get