diff --git a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean index 426251c74c..e09f7376af 100644 --- a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean +++ b/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 diff --git a/tests/lean/run/constant_fold_task_get_pure.lean b/tests/lean/run/constant_fold_task_get_pure.lean new file mode 100644 index 0000000000..a1acb56105 --- /dev/null +++ b/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