Skip to content

Commit

Permalink
[kernel] Loop optimizations (#1061)
Browse files Browse the repository at this point in the history 
<!--
PRs require an approval from any of the core contributors, other than
the PR author.

Include this header if applicable:
Fixes #issue1, #issue2, ...
-->

### Problem
<!--
Explain here the context, and why you're making this change. What is the
problem you're trying to solve?
-->
I'm taking a look at profiling sessions for Kyo's benchmarks and noticed
that `Loop` is getting the `Safepoint` thread local on each iteration
because of the `run` call being in a default parameter. The compiler
generates a method without the `Safepoint` evidence.

### Solution
<!--
Describe your solution. Focus on helping reviewers understand your
technical approach and implementation decisions.
-->
Restructure methods to avoid the default parameter.
  • Loading branch information
@fwbrasil
fwbrasil authored Jan 31, 2025
1 parent 617ca5d commit 724e9fe
Show file tree
Hide file tree
Showing 2 changed files with 142 additions and 53 deletions.
140 changes: 87 additions & 53 deletions kyo-kernel/shared/src/main/scala/kyo/kernel/Loop.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -253,18 +253,18 @@ object Loop:
safepoint: Safepoint
): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(i1: A)(v: Outcome[A, O] < S = run(i1))(using Safepoint): O < S =
@tailrec def loop(v: Outcome[A, O] < S)(using Safepoint): O < S =
v match
case next: Continue[A] @unchecked =>
loop(next._1)()
loop(run(next._1))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome[A, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(i1)(kyo(v, context))
loop(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(input)()
loop(Loop.continue(input))
end apply

/** Executes a loop with two state values.
Expand All @@ -287,18 +287,18 @@ object Loop:
safepoint: Safepoint
): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(i1: A, i2: B)(v: Outcome2[A, B, O] < S = run(i1, i2))(using Safepoint): O < S =
@tailrec def loop(v: Outcome2[A, B, O] < S)(using Safepoint): O < S =
v match
case next: Continue2[A, B] @unchecked =>
loop(next._1, next._2)()
loop(run(next._1, next._2))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome2[A, B, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(i1, i2)(kyo(v, context))
loop(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(input1, input2)()
loop(Loop.continue(input1, input2))
end apply

/** Executes a loop with three state values.
Expand All @@ -321,18 +321,18 @@ object Loop:
inline run: Safepoint ?=> (A, B, C) => Outcome3[A, B, C, O] < S
)(using inline _frame: Frame, safepoint: Safepoint): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(i1: A, i2: B, i3: C)(v: Outcome3[A, B, C, O] < S = run(i1, i2, i3))(using Safepoint): O < S =
@tailrec def loop(v: Outcome3[A, B, C, O] < S)(using Safepoint): O < S =
v match
case next: Continue3[A, B, C] @unchecked =>
loop(next._1, next._2, next._3)()
loop(run(next._1, next._2, next._3))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome3[A, B, C, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(i1, i2, i3)(kyo(v, context))
loop(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(input1, input2, input3)()
loop(Loop.continue(input1, input2, input3))
end apply

/** Executes a loop with four state values.
Expand All @@ -357,18 +357,18 @@ object Loop:
inline run: Safepoint ?=> (A, B, C, D) => Outcome4[A, B, C, D, O] < S
)(using inline _frame: Frame, safepoint: Safepoint): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(i1: A, i2: B, i3: C, i4: D)(v: Outcome4[A, B, C, D, O] < S = run(i1, i2, i3, i4))(using Safepoint): O < S =
@tailrec def loop(v: Outcome4[A, B, C, D, O] < S)(using Safepoint): O < S =
v match
case next: Continue4[A, B, C, D] @unchecked =>
loop(next._1, next._2, next._3, next._4)()
loop(run(next._1, next._2, next._3, next._4))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome4[A, B, C, D, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(i1, i2, i3, i4)(kyo(v, context))
loop(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(input1, input2, input3, input4)()
loop(Loop.continue(input1, input2, input3, input4))
end apply

/** Executes an indexed loop without state values.
Expand All @@ -387,18 +387,18 @@ object Loop:
safepoint: Safepoint
): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(idx: Int)(v: Outcome[Unit, O] < S = run(idx))(using Safepoint): O < S =
@tailrec def loop(idx: Int)(v: Outcome[Unit, O] < S)(using Safepoint): O < S =
v match
case next: Continue[Unit] @unchecked =>
loop(idx + 1)()
loop(idx + 1)(run(idx))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome[Unit, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(idx)(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(0)()
loop(0)(Loop.continue)
end indexed

/** Executes an indexed loop with a single state value.
Expand All @@ -418,18 +418,18 @@ object Loop:
safepoint: Safepoint
): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(idx: Int, i1: A)(v: Outcome[A, O] < S = run(idx, i1))(using Safepoint): O < S =
@tailrec def loop(idx: Int)(v: Outcome[A, O] < S)(using Safepoint): O < S =
v match
case next: Continue[A] @unchecked =>
loop(idx + 1, next._1)()
loop(idx + 1)(run(idx, next._1))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome[A, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(idx, i1)(kyo(v, context))
loop(idx)(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(0, input)()
loop(0)(Loop.continue(input))
end indexed

/** Executes an indexed loop with two state values.
Expand All @@ -450,18 +450,18 @@ object Loop:
inline run: Safepoint ?=> (Int, A, B) => Outcome2[A, B, O] < S
)(using inline _frame: Frame, safepoint: Safepoint): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(idx: Int, i1: A, i2: B)(v: Outcome2[A, B, O] < S = run(idx, i1, i2))(using Safepoint): O < S =
@tailrec def loop(idx: Int)(v: Outcome2[A, B, O] < S)(using Safepoint): O < S =
v match
case next: Continue2[A, B] @unchecked =>
loop(idx + 1, next._1, next._2)()
loop(idx + 1)(run(idx, next._1, next._2))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome2[A, B, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(idx, i1, i2)(kyo(v, context))
loop(idx)(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(0, input1, input2)()
loop(0)(Loop.continue(input1, input2))
end indexed

/** Executes an indexed loop with three state values.
Expand All @@ -484,18 +484,18 @@ object Loop:
inline run: Safepoint ?=> (Int, A, B, C) => Outcome3[A, B, C, O] < S
)(using inline _frame: Frame, safepoint: Safepoint): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(idx: Int, i1: A, i2: B, i3: C)(v: Outcome3[A, B, C, O] < S = run(idx, i1, i2, i3))(using Safepoint): O < S =
@tailrec def loop(idx: Int)(v: Outcome3[A, B, C, O] < S)(using Safepoint): O < S =
v match
case next: Continue3[A, B, C] @unchecked =>
loop(idx + 1, next._1, next._2, next._3)()
loop(idx + 1)(run(idx, next._1, next._2, next._3))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome3[A, B, C, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(idx, i1, i2, i3)(kyo(v, context))
loop(idx)(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(0, input1, input2, input3)()
loop(0)(Loop.continue(input1, input2, input3))
end indexed

/** Executes an indexed loop with four state values.
Expand All @@ -520,19 +520,18 @@ object Loop:
inline run: Safepoint ?=> (Int, A, B, C, D) => Outcome4[A, B, C, D, O] < S
)(using inline _frame: Frame, safepoint: Safepoint): O < S =
@nowarn("msg=anonymous")
@tailrec def loop(idx: Int, i1: A, i2: B, i3: C, i4: D)(v: Outcome4[A, B, C, D, O] < S =
run(idx, i1, i2, i3, i4))(using Safepoint): O < S =
@tailrec def loop(idx: Int)(v: Outcome4[A, B, C, D, O] < S)(using Safepoint): O < S =
v match
case next: Continue4[A, B, C, D] @unchecked =>
loop(idx + 1, next._1, next._2, next._3, next._4)()
loop(idx + 1)(run(idx, next._1, next._2, next._3, next._4))
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome4[A, B, C, D, O], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, O, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(idx, i1, i2, i3, i4)(kyo(v, context))
loop(idx)(kyo(v, context))
case res =>
res.asInstanceOf[O]
loop(0, input1, input2, input3, input4)()
loop(0)(Loop.continue(input1, input2, input3, input4))
end indexed

/** Executes a loop that continues until explicitly completed.
Expand All @@ -547,18 +546,18 @@ object Loop:
*/
inline def foreach[S](inline run: Safepoint ?=> Outcome[Unit, Unit] < S)(using inline _frame: Frame, safepoint: Safepoint): Unit < S =
@nowarn("msg=anonymous")
@tailrec def loop(v: Outcome[Unit, Unit] < S = run)(using Safepoint): Unit < S =
@tailrec def loop(v: Outcome[Unit, Unit] < S)(using Safepoint): Unit < S =
v match
case next: Continue[Unit] @unchecked =>
loop()
loop(run)
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome[Unit, Unit], S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, Unit, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(kyo(v, context))
case res =>
()
loop()
loop(Loop.continue)
end foreach

/** Repeats an operation a specified number of times.
Expand All @@ -575,21 +574,20 @@ object Loop:
*/
inline def repeat[S](n: Int)(inline run: Safepoint ?=> Unit < S)(using inline _frame: Frame, safepoint: Safepoint): Unit < S =
@nowarn("msg=anonymous")
@tailrec def loop(i: Int)(v: Outcome[Unit, Unit] < S = run)(using Safepoint): Unit < S =
if i == n then ()
@tailrec def loop(i: Int)(v: Unit < S)(using Safepoint): Unit < S =
if i > n then ()
else
v match
case kyo: KyoSuspend[IX, OX, EX, Any, Outcome[Unit, Unit], S] @unchecked =>
case kyo: KyoSuspend[IX, OX, EX, Any, Unit, S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, Unit, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(i)(kyo(v, context))
end new
case _ =>
loop(i + 1)()
loop(i + 1)(run)
end if
end loop
loop(0)()
loop(0)(())
end repeat

/** Executes a loop indefinitely until explicitly terminated.
Expand All @@ -602,14 +600,50 @@ object Loop:
* @return
* Nothing, as this loop runs forever unless interrupted
*/
inline def forever[S](inline run: Safepoint ?=> Unit < S)(using Frame, Safepoint): Unit < S =
def _loop(): Unit < S = loop()
@tailrec def loop(): Unit < S =
run match
case kyo: Kyo[Unit, S] @unchecked =>
kyo.andThen(_loop())
inline def forever[S](inline run: Safepoint ?=> Unit < S)(using inline _frame: Frame, safepoint: Safepoint): Unit < S =
@nowarn("msg=anonymous")
@tailrec def loop(v: Unit < S)(using Safepoint): Unit < S =
v match
case kyo: KyoSuspend[IX, OX, EX, Any, Unit, S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, Unit, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(kyo(v, context))
case _ =>
loop()
loop()
loop(run)
end loop
loop(())
end forever

/** Executes an operation repeatedly while a condition remains true.
*
* @param condition
* The condition to check before each iteration. The loop continues while this is true.
* @param run
* The operation to execute in each iteration
* @return
* Unit after the loop completes
*/
inline def whileTrue[S](inline condition: Safepoint ?=> Boolean < S)(inline run: Unit < S)(
using
inline _frame: Frame,
safepoint: Safepoint
): Unit < S =
@nowarn("msg=anonymous")
def loop(v: Unit < S)(using Safepoint): Unit < S =
condition.map {
case true =>
v match
case kyo: KyoSuspend[IX, OX, EX, Any, Unit, S] @unchecked =>
new KyoContinue[IX, OX, EX, Any, Unit, S](kyo):
def frame = _frame
def apply(v: OX[Any], context: Context)(using Safepoint) =
loop(kyo(v, context))
case _ =>
loop(run)
case false => ()
}
end loop
loop(())
end whileTrue
end Loop
55 changes: 55 additions & 0 deletions kyo-kernel/shared/src/test/scala/kyo/kernel/LoopTest.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -674,4 +674,59 @@ class LoopTest extends Test:
typeCheckFailure("implicitly[Flat[Loop.Outcome4[Int, String, Boolean, Double, Int < Any]]]")(error)
}
}

"whileTrue" - {
"continues while condition is true" in {
var counter = 0
Loop.whileTrue(counter < 5) {
counter += 1
}.eval
assert(counter == 5)
}

"does not execute when condition is initially false" in {
var executed = false
Loop.whileTrue(false) {
executed = true
}.eval
assert(!executed)
}

"with suspended condition" in {
var counter = 0
val result = Loop.whileTrue(Effect.defer(counter < 3)) {
counter += 1
}
result.eval
assert(counter == 3)
}

"with suspended body" in {
var counter = 0
val result = Loop.whileTrue(counter < 3) {
Effect.defer(counter += 1)
}
result.eval
assert(counter == 3)
}

"stack safety" in {
var counter = 0
val largeNumber = 100000
Loop.whileTrue(counter < largeNumber) {
counter += 1
}.eval
assert(counter == largeNumber)
}

"stack safety with suspended operations" in {
var counter = 0
val largeNumber = 10000
val result = Loop.whileTrue(Effect.defer(counter < largeNumber)) {
Effect.defer(counter += 1)
}
result.eval
assert(counter == largeNumber)
}
}
end LoopTest

0 comments on commit 724e9fe

Please sign in to comment.