diff --git a/core/src/main/scala-3/spire/syntax/FastForSyntax.scala b/core/src/main/scala-3/spire/syntax/FastForSyntax.scala index 7a007a735..d61ee2eeb 100644 --- a/core/src/main/scala-3/spire/syntax/FastForSyntax.scala +++ b/core/src/main/scala-3/spire/syntax/FastForSyntax.scala @@ -27,7 +27,7 @@ trait FastForSyntax: case NumericRange[Long] => Long inline def fastFor[A](inline init: A)(inline test: A => Boolean, inline next: A => A)(inline body: A => Unit): Unit = - fastForInline(init, test, next, body) + ${ fastForImpl('init, 'test, 'next, 'body) } inline def fastForRange[R <: RangeLike](inline r: R)(inline body: RangeElem[R] => Unit): Unit = ${ fastForRangeMacroGen('r, 'body) } diff --git a/core/src/main/scala-3/spire/syntax/macros/cforMacros.scala b/core/src/main/scala-3/spire/syntax/macros/cforMacros.scala index 4de6fd0d5..c40cc6946 100644 --- a/core/src/main/scala-3/spire/syntax/macros/cforMacros.scala +++ b/core/src/main/scala-3/spire/syntax/macros/cforMacros.scala @@ -15,157 +15,125 @@ package spire.syntax.macros -import quoted._ -import collection.immutable.NumericRange +import scala.quoted.* +import scala.collection.immutable.NumericRange +import scala.PartialFunction.cond import spire.syntax.fastFor.{RangeElem, RangeLike} -inline def fastForInline[R](init: R, test: R => Boolean, next: R => R, body: R => Unit): Unit = - var index = init - while test(index) do - body(index) - index = next(index) - -def fastForRangeMacroGen[R <: RangeLike: Type](r: Expr[R], body: Expr[RangeElem[R] => Unit])(using - quotes: Quotes +def fastForImpl[R: Type](init: Expr[R], test: Expr[R => Boolean], next: Expr[R => R], body: Expr[R => Unit])(using + Quotes ): Expr[Unit] = - import quotes._ - import quotes.reflect._ - - type RangeL = NumericRange[Long] + import quotes.reflect.* - (r, body) match - case '{ $r: Range } -> '{ $body: (Int => Unit) } => fastForRangeMacro(r, body) - case '{ $r: NumericRange[Long] } -> '{ $body: (Long => Unit) } => fastForRangeMacroLong(r, body) - case '{ $r } -> _ => report.error(s"Ineligible Range type ", r); '{} + def code(testRef: Expr[R => Boolean], nextRef: Expr[R => R], bodyRef: Expr[R => Unit]): Expr[Unit] = '{ + var index = $init + while $testRef(index) do + $bodyRef(index) + index = $nextRef(index) + } -end fastForRangeMacroGen + letFunc("test", test)(t => letFunc("next", next)(n => letFunc("body", body)(b => code(t, n, b)))) +end fastForImpl -def fastForRangeMacroLong(r: Expr[NumericRange[Long]], body: Expr[Long => Unit])(using quotes: Quotes): Expr[Unit] = - import quotes._ +def fastForRangeMacroGen[R <: RangeLike: Type](r: Expr[R], body: Expr[RangeElem[R] => Unit])(using + quotes: Quotes +): Expr[Unit] = import quotes.reflect.* - def strideUpUntil(fromExpr: Expr[Long], untilExpr: Expr[Long], stride: Expr[Long]): Expr[Unit] = '{ - var index = $fromExpr - val limit = $untilExpr - val body0 = $body - while index < limit do - ${ Expr.betaReduce(body) }(index) - index += $stride - } + r match + case '{ $r: Range } => RangeForImpl.ofInt(r, body.asExprOf[Int => Unit]) + case '{ $r: NumericRange[Long] } => RangeForImpl.ofLong(r, body.asExprOf[Long => Unit]) + case '{ $r } => report.error(s"Ineligible Range type ", r); '{} - def strideUpTo(fromExpr: Expr[Long], untilExpr: Expr[Long], stride: Expr[Long]): Expr[Unit] = '{ - var index = $fromExpr - val end = $untilExpr - while index <= end do - ${ Expr.betaReduce(body) }(index) - index += $stride - } +end fastForRangeMacroGen - def strideDownTo(fromExpr: Expr[Long], untilExpr: Expr[Long], stride: Expr[Long]): Expr[Unit] = '{ - var index = $fromExpr - val end = $untilExpr - while index >= end do - ${ Expr.betaReduce(body) }(index) - index -= $stride - } +private object RangeForImpl: + type Code[T] = Expr[T => Unit] => Expr[Unit] + type Test[T] = (Expr[T], Expr[T]) => Expr[Boolean] + + def ofInt(r: Expr[Range], body: Expr[Int => Unit])(using Quotes): Expr[Unit] = + val code: Code[Int] = r match + case '{ ($i: Int) to $j } => loopCode(i, j, 1, (x, y) => '{ $x <= $y }) + case '{ ($i: Int) to $j by ${ Expr(k) } } if k > 0 => loopCode(i, j, k, (x, y) => '{ $x <= $y }) + case '{ ($i: Int) to $j by ${ Expr(k) } } if k < 0 => loopCode(i, j, k, (x, y) => '{ $x >= $y }) + case '{ ($i: Int) to $j by ${ Expr(k) } } if k == 0 => zeroStride(r) + case '{ ($i: Int) until $j } => loopCode(i, j, 1, (x, y) => '{ $x < $y }) + case '{ ($i: Int) until $j by ${ Expr(k) } } if k > 0 => loopCode(i, j, k, (x, y) => '{ $x < $y }) + case '{ ($i: Int) until $j by ${ Expr(k) } } if k < 0 => loopCode(i, j, k, (x, y) => '{ $x > $y }) + case '{ ($i: Int) until $j by ${ Expr(k) } } if k == 0 => zeroStride(r) + case _ => deOpt(r, '{ $r.foreach($body) }) + + letFunc("body", body)(code) + end ofInt + + def ofLong(r: Expr[NumericRange[Long]], body: Expr[Long => Unit])(using quotes: Quotes): Expr[Unit] = + val code: Code[Long] = r match + case '{ ($i: Long) to $j } => loopCode(i, j, 1L, (x, y) => '{ $x <= $y }) + case '{ ($i: Long) to $j by ${ Expr(k) } } if k > 0 => loopCode(i, j, k, (x, y) => '{ $x <= $y }) + case '{ ($i: Long) to $j by ${ Expr(k) } } if k < 0 => loopCode(i, j, k, (x, y) => '{ $x >= $y }) + case '{ ($i: Long) to $j by ${ Expr(k) } } if k == 0 => zeroStride(r) + case '{ ($i: Long) until $j } => loopCode(i, j, 1L, (x, y) => '{ $x < $y }) + case '{ ($i: Long) until $j by ${ Expr(k) } } if k > 0 => loopCode(i, j, k, (x, y) => '{ $x < $y }) + case '{ ($i: Long) until $j by ${ Expr(k) } } if k < 0 => loopCode(i, j, k, (x, y) => '{ $x > $y }) + case '{ ($i: Long) until $j by ${ Expr(k) } } if k == 0 => zeroStride(r) + case _ => deOpt(r, '{ $r.foreach($body) }) + + letFunc("body", body)(code) + + end ofLong + + def loopCode[T: Type: ToExpr: CanLoop](i: Expr[T], j: Expr[T], s: T, test: Test[T])(using Quotes): Code[T] = + body => + '{ + var index = $i + val limit = $j + while ${ test('index, 'limit) } do + $body(index) + index = ${ 'index.stepBy(Expr(s)) } + } - def strideDownUntil(fromExpr: Expr[Long], untilExpr: Expr[Long], stride: Expr[Long]): Expr[Unit] = '{ - var index = $fromExpr - val limit = $untilExpr - while index > limit do - ${ Expr.betaReduce(body) }(index) - index -= $stride - } + def zeroStride[T, R](orig: Expr[R])(using Quotes): Code[T] = _ => + import quotes.reflect.* + report.error("zero stride", orig) + '{} - r match - case '{ ($i: Long) until $j } => strideUpUntil(i, j, Expr(1L)) - case '{ ($i: Long) to $j } => strideUpTo(i, j, Expr(1L)) - case '{ ($i: Long) until $j by $step } => - step.asTerm match { - case Literal(LongConstant(k)) if k > 0 => strideUpUntil(i, j, Expr(k)) - case Literal(LongConstant(k)) if k < 0 => strideDownUntil(i, j, Expr(-k)) - case Literal(LongConstant(k)) if k == 0 => report.error("zero stride", step); '{} - case _ => - report.warning(s"defaulting to foreach, can not optimise non-constant step", step) - '{ val b = $body; $r.foreach(b) } - } - case '{ ($i: Long) to $j by $step } => - step.asTerm match { - case Literal(LongConstant(k)) if k > 0 => strideUpTo(i, j, Expr(k)) - case Literal(LongConstant(k)) if k < 0 => strideDownTo(i, j, Expr(-k)) - case Literal(LongConstant(k)) if k == 0 => report.error("zero stride", step); '{} - case _ => - report.warning(s"defaulting to foreach, can not optimise non-constant step", step) - '{ val b = $body; $r.foreach(b) } - } + def deOpt[T, R](orig: Expr[R], foreach: Expr[Unit])(using Quotes): Code[T] = _ => + import quotes.reflect.* + report.warning(s"defaulting to foreach, can not optimise range expression", orig) + foreach - case _ => - report.warning(s"defaulting to foreach, can not optimise range expression", r) - '{ val b = $body; $r.foreach(b) } + trait CanLoop[T]: + extension (x: Expr[T]) def stepBy(y: Expr[T])(using Quotes): Expr[T] -end fastForRangeMacroLong + object CanLoop: + given CanLoop[Int] with + extension (x: Expr[Int]) def stepBy(y: Expr[Int])(using Quotes): Expr[Int] = '{ $x + $y } -def fastForRangeMacro(r: Expr[Range], body: Expr[Int => Unit])(using quotes: Quotes): Expr[Unit] = - import quotes._ - import quotes.reflect._ + given CanLoop[Long] with + extension (x: Expr[Long]) def stepBy(y: Expr[Long])(using Quotes): Expr[Long] = '{ $x + $y } - def strideUpUntil(fromExpr: Expr[Int], untilExpr: Expr[Int], stride: Expr[Int]): Expr[Unit] = '{ - var index = $fromExpr - val limit = $untilExpr - while (index < limit) { - ${ Expr.betaReduce(body) }(index) - index += $stride - } - } +end RangeForImpl - def strideUpTo(fromExpr: Expr[Int], untilExpr: Expr[Int], stride: Expr[Int]): Expr[Unit] = '{ - var index = $fromExpr - val end = $untilExpr - while index <= end do - ${ Expr.betaReduce(body) }(index) - index += $stride - } +/** + * Equivalent to `'{ val name: A => B = $rhs; ${in('name)} }`, except when `rhs` is a function literal, then equivalent + * to `in(rhs)`. + * + * This allows inlined function arguments to perform side-effects only once before their first evaluation, while still + * avoiding the creation of closures for function literal arguments. + */ +private def letFunc[A, B, C](using Quotes)(name: String, rhs: Expr[A => B])(in: Expr[A => B] => Expr[C]): Expr[C] = + import quotes.reflect.* - def strideDownTo(fromExpr: Expr[Int], untilExpr: Expr[Int], stride: Expr[Int]): Expr[Unit] = '{ - var index = $fromExpr - val end = $untilExpr - while index >= end do - ${ Expr.betaReduce(body) }(index) - index -= $stride - } + extension (t: Term) def unsafeAsExpr[A] = t.asExpr.asInstanceOf[Expr[A]] // cast without `quoted.Type[A]` - def strideDownUntil(fromExpr: Expr[Int], untilExpr: Expr[Int], stride: Expr[Int]): Expr[Unit] = '{ - var index = $fromExpr - val limit = $untilExpr - while index > limit do - ${ Expr.betaReduce(body) }(index) - index -= $stride + def isFunctionLiteral[A, B](f: Expr[A => B]): Boolean = cond(f.asTerm.underlyingArgument) { case Lambda(_, _) => + true } - r match - case '{ ($i: Int) until $j } => strideUpUntil(i, j, Expr(1)) - case '{ ($i: Int) to $j } => strideUpTo(i, j, Expr(1)) - case '{ ($i: Int) until $j by $step } => - step.asTerm match { - case Literal(IntConstant(k)) if k > 0 => strideUpUntil(i, j, Expr(k)) - case Literal(IntConstant(k)) if k < 0 => strideDownUntil(i, j, Expr(-k)) - case Literal(IntConstant(k)) if k == 0 => report.error("zero stride", step); '{} - case _ => - report.warning(s"defaulting to foreach, can not optimise non-constant step", step) - '{ val b = $body; $r.foreach(b) } - } - case '{ ($i: Int) to $j by $step } => - step.asTerm match { - case Literal(IntConstant(k)) if k > 0 => strideUpTo(i, j, Expr(k)) - case Literal(IntConstant(k)) if k < 0 => strideDownTo(i, j, Expr(-k)) - case Literal(IntConstant(k)) if k == 0 => report.error("zero stride", step); '{} - case _ => - report.warning(s"defaulting to foreach, can not optimise non-constant step", step) - '{ val b = $body; $r.foreach(b) } - } - case _ => - report.warning(s"defaulting to foreach, can not optimise range expression", r) - '{ val b = $body; $r.foreach(b) } + def let[A, B](name: String, rhs: Expr[A])(in: Expr[A] => Expr[B])(using Quotes): Expr[B] = + // Equivalent to `'{ val name = $rhs; ${in('name)} }` + ValDef.let(Symbol.spliceOwner, name, rhs.asTerm)(ref => in(ref.unsafeAsExpr[A]).asTerm).unsafeAsExpr[B] -end fastForRangeMacro + if isFunctionLiteral(rhs) then in(Expr.betaReduce(rhs)) + else let(name, rhs)(in) diff --git a/tests/shared/src/test/scala/spire/syntax/FastForSuite.scala b/tests/shared/src/test/scala-3/scala/spire/syntax/FastForSuite.scala similarity index 96% rename from tests/shared/src/test/scala/spire/syntax/FastForSuite.scala rename to tests/shared/src/test/scala-3/scala/spire/syntax/FastForSuite.scala index 5457fc9bd..64e9e8c91 100644 --- a/tests/shared/src/test/scala/spire/syntax/FastForSuite.scala +++ b/tests/shared/src/test/scala-3/scala/spire/syntax/FastForSuite.scala @@ -99,15 +99,16 @@ class FastForSuite extends munit.FunSuite { assertEquals(b.toList, List(0, 1, 2)) } - // This test distinguishes fastFor from cfor - test("doesn't capture value in closure") { + test("capture value in closure") { // same behavior as cfor val b1 = collection.mutable.ArrayBuffer.empty[() => Int] fastFor(0)(_ < 3, _ + 1) { x => b1 += (() => x) } val b2 = collection.mutable.ArrayBuffer[() => Int]() - (0 until 3).foreach { x => - b2 += (() => x) + var i = 0 + while (i < 3) { + b2 += (() => i) + i += 1 } assertEquals(b1.map(_.apply()).toList, b2.map(_.apply()).toList) }