diff --git a/library/init/control/combinators.lean b/library/init/control/combinators.lean index 664617e8e1..65a97adb4c 100644 --- a/library/init/control/combinators.lean +++ b/library/init/control/combinators.lean @@ -39,6 +39,13 @@ namespace Nat @[inline] def mfor {m} [Applicative m] (n : Nat) (f : Nat → m Unit) : m Unit := mforAux f n n +@[specialize] def mforRevAux {m} [Applicative m] (f : Nat → m Unit) : Nat → m Unit +| 0 => pure () +| i+1 => f i *> mforRevAux i + +@[inline] def mforRev {m} [Applicative m] (n : Nat) (f : Nat → m Unit) : m Unit := +mforRevAux f n + @[specialize] def mfoldAux {α : Type u} {m : Type u → Type v} [Monad m] (f : Nat → α → m α) (n : Nat) : Nat → α → m α | 0, a => pure a | i+1, a => f (n-i-1) a >>= mfoldAux i @@ -46,6 +53,13 @@ mforAux f n n @[inline] def mfold {α : Type u} {m : Type u → Type v} [Monad m] (f : Nat → α → m α) (a : α) (n : Nat) : m α := mfoldAux f n n a +@[specialize] def mfoldRevAux {α : Type u} {m : Type u → Type v} [Monad m] (f : Nat → α → m α) : Nat → α → m α +| 0, a => pure a +| i+1, a => f i a >>= mfoldRevAux i + +@[inline] def mfoldRev {α : Type u} {m : Type u → Type v} [Monad m] (f : Nat → α → m α) (a : α) (n : Nat) : m α := +mfoldRevAux f n a + -- TODO: enable after we have support for marking arguments that should be considered for specialization. /- @[specialize] diff --git a/library/init/data/nat/basic.lean b/library/init/data/nat/basic.lean index 7f62650385..8c7594480d 100644 --- a/library/init/data/nat/basic.lean +++ b/library/init/data/nat/basic.lean @@ -89,6 +89,13 @@ instance : HasMul Nat := @[inline] def fold {α : Type u} (f : Nat → α → α) (n : Nat) (a : α) : α := foldAux f n n a +@[specialize] def foldRevAux {α : Type u} (f : Nat → α → α) : Nat → α → α +| 0, a => a +| succ n, a => foldRevAux n (f n a) + +@[inline] def foldRev {α : Type u} (f : Nat → α → α) (n : Nat) (a : α) : α := +foldRevAux f n a + @[specialize] def anyAux (f : Nat → Bool) (s : Nat) : Nat → Bool | 0 => false | succ n => f (s - (succ n)) || anyAux n