`Thunk.asFunction0` is broken in Scala 3 nightly

[armanbilge]

This was identified in scala/scala3#14395 (comment), apparently due to scala/scala3#14295.

[armanbilge]

An optimization is applied: If the argument e to a cbn parameter is already
of type () ?=> T and is a pure expression, we avoid (2) and (3), i.e. we
pass e directly instead of () ?=> e.apply().

This reads to me a lot like the optimization that Scala 2 has and that Scala 3 was previously lacking as identified in #2225. But I don't understand why our trick broke.

[djspiewak]

Yeah I'm perplexed as well. I think this merits some more experimentation. We should try to figure out what's going on so we can give them feedback before they release. At a minimum, the fact that this is detectably different between versions of Scala 3 is frustrating, since it means that our users on one or the other version will have a significantly degraded experience.

[armanbilge]

benchmarks/Jmh/run -wi 10 -i 10 -f 2 -t 1 DeepBindBenchmark.delay --jvmArgs -Dcats.effect.tracing.mode=none

3.1.3-RC1-bin-20220207-74c2954-NIGHTLY

[info] Benchmark                (size)   Mode  Cnt     Score     Error  Units
[info] DeepBindBenchmark.delay   10000  thrpt   20  4058.603 ± 105.582  ops/s

3.1.1

[info] Benchmark                (size)   Mode  Cnt     Score     Error  Units
[info] DeepBindBenchmark.delay   10000  thrpt   20  3712.925 ± 332.421  ops/s

Not sure if we can really read anything into it.

[armanbilge]

At a minimum, the fact that this is detectably different between versions of Scala 3 is frustrating, since it means that our users on one or the other version will have a significantly degraded experience.

Yep, this is what got my attention too. Is there no way to check the Scala 3 version at runtime and use that to choose an implementation?

[armanbilge]

I tried restoring the original Java-based workaround from #2226, but the test still fails.

cats-effect/core/jvm/src/main/java/cats/effect/ThunkImpl.java

Lines 24 to 29 in 6d59eeb

	class ThunkImpl extends Thunk {
	@Override
	public <A> scala.Function0<A> asFunction0(scala.Function0<A> thunk) {
	return thunk;
	}
	}

So it must be that the thunk is being wrapped earlier than it used to be?

[djspiewak]

I think we may have to do some bytecode dumpster diving to figure this one out.

[armanbilge]

Decompiled Java is about as much as I'll tolerate 😤

"Thunk.asFunction0" should {
  "return the same function" in {
    var i = 0
    val f = () => i += 1
    Thunk.asFunction0(f()) eq f
  }
}

3.0.2

 private static final boolean $init$$$anonfun$2$$anonfun$1() {
    IntRef i = IntRef.create(0);
    Function0 f = () -> {
       int var1 = i.elem + 1;
       i.elem = var1;
    };
    return cats.effect.Thunk..MODULE$.asFunction0(f) == f;
 }

3.1.3-RC1-bin-20220207-74c2954-NIGHTLY

 private static final void $init$$$anonfun$1$$anonfun$1$$anonfun$1(final Function0 f$1) {
    f$1.apply$mcV$sp();
 }

 private static final boolean $init$$$anonfun$1$$anonfun$1() {
    IntRef i = IntRef.create(0);
    Function0 f = () -> {
       int var1 = i.elem + 1;
       i.elem = var1;
    };
    return cats.effect.Thunk..MODULE$.asFunction0(() -> {
       $init$$$anonfun$1$$anonfun$1$$anonfun$1(f);
       return BoxedUnit.UNIT;
    }) == f;
 }

[djspiewak]

Uh. That looks like a regression honestly. Unless im misreading the above, it would seem that the change is causing unnecessary wrapping of the thunk.

[vasilmkd]

This should be enough to file an issue upstream, no?

[djspiewak]

Let's dig into the raw bytecode first, but it is very suggestive of a bug upstream. It honestly looks like the compiler is doing the opposite of what Martin's change was intended to achieve.

[armanbilge]

Putting aside the performance regression (which FTR my benchmarks were unable to pick up), I realized that actually this change doesn't break tracing. The relevant test for that continues to pass.

cats-effect/tests/shared/src/test/scala/cats/effect/tracing/TracingSpec.scala

Lines 35 to 42 in 4f01f5f

	"generate unique traces" in {
	val a = IO(println("foo"))
	val b = IO(println("bar"))
	(a, b) match {
	case (IO.Delay(_, eventA), IO.Delay(_, eventB)) => eventA ne eventB
	case _ => false
	}
	}

Even more interesting, actually you could argue the accuracy of cached tracing is improved. This test started to fail, but actually it encodes a less-than-desirable behavior, which is that two independent IOs that happen to call the same method end up sharing the same trace due to caching.

cats-effect/tests/shared/src/test/scala/cats/effect/tracing/TracingSpec.scala

Lines 25 to 32 in 4f01f5f

	"generate identical traces" in {
	val f = () => println("foo")
	val a = IO(f())
	val b = IO(f())
	(a, b) match {
	case (IO.Delay(_, eventA), IO.Delay(_, eventB)) => eventA eq eventB
	case _ => false
	}

Because each of those thunks now generates a unique wrapping lambda, they get unique traces.

[armanbilge]

One more observation: what has changed in the compiler is actually the original 2.13->3 regression that caused us to introduce Thunk.asFunction0 in the first place.

If I restore the code to how it used to be:

   def delay[A](thunk: => A): IO[A] = {
-    val fn = Thunk.asFunction0(thunk)
+    val fn = () => thunk
     Delay(fn, Tracing.calculateTracingEvent(fn))
   }

These are the test results:

3.1.3-RC1-bin-20220207-74c2954-NIGHTLY

[info] TracingSpec
[info] IO.delay should
[error]   x generate identical traces
[error]    false (file:1)
[info]   + generate unique traces
[info] Async.delay should
[error]   x generate identical traces
[error]    false (file:1)
[info]   + generate unique traces
[info] Total for specification TracingSpec

3.0.1

[info] TracingSpec
[info] IO.delay should
[info]   + generate identical traces
[error]   x generate unique traces
[error]    false (file:1)
[info] Async.delay should
[info]   + generate identical traces
[error]   x generate unique traces
[error]    false (file:1)
[info] Total for specification TracingSpec

Essentially, there's been a sort of flip-flop.

In the code:

   def delay[A](thunk: => A): IO[A] = {
     val fn = () => thunk
     Delay(fn, Tracing.calculateTracingEvent(fn))
   }

In 3.0.2, the () => thunk was creating a wrapper. Because this wrapper always had the same class, it was causing all thunks to share the same cache key.

In the nightly, the () => thunk is optimized to the Function0 backing the thunk, just like it works in Scala 2.

However, it's the opposite story at the callsite:

var i = 0
val f = () => i += 1
IO.delay(f())

In 3.0.2, the thunk passed to IO.delay was f. This is also how it works in Scala 2.

In the nightly, the f is being wrapped in another Function0 and this is passed in as the thunk. Since the wrapping occurs at the callsite, it is now out of our control. However, it does have the benefit of making every use of IO.delay unique from the perspective of trace caching.

[djspiewak]

In the nightly, the f is being wrapped in another Function0 and this is passed in as the thunk

This is the bit that I think constitutes a regression. Definitely the improvement of the "thunk to () => A" behavior is very very welcome.

Realistically, we're probably going to have to upgrade to the latest Scala 3 in order to force call sites into this particular pattern, otherwise tracing breaks.

[armanbilge]

Realistically, we're probably going to have to upgrade to the latest Scala 3 in order to force call sites into this particular pattern, otherwise tracing break

Why would tracing break?