Inline Traits & Specialized Traits

[starswap]

Final year master's project at EPFL. Introduces specialization for Scala 3 as an experimental feature, via:

• Inline Traits (as built by @timotheeandres, with some modifications and fixes)
• Specialized Traits (as originally described in Specialized traits design draft #21950 also with some modifications and fixes)

The latest specification for each of these is included in docs accompanying the PR. To see a demo of the generated code and speed gain offered, try the benchmarks in tools/benchmarks/specializedtraits

How much have you relied on LLM-based tools in this contribution?

Moderately, for debugging and generating some of the benchmark examples.

How was the solution tested?

New automated tests

[starswap]

Check: Is this ok? I'm not sure if we can do this because maybe it is actually necessary to read the contents of the body annotation in the context of the body annotation. It seems to pass tests but I'm not sure if it's ok. Also this code was originally copied from the inline methods case when inline traits got built so it is possible not all of it is needed for inline traits.

[jchyb]

This is incredible work, with a staggering amount of corner cases considered and tested (and well documented!). Thank you so much! Honestly, this is mostly mergeable already, but I did come up with some questions/concerns.

This is a huge amount of changes, and I was only assigned this last week (although I was familiar with the earlier proposals for specializations in Scala 3). All of this is to say, I might not have all of the context behind the project here, so please don't hesitate to correct me. I also still saw some leftover TODOs, and I tried to reply to most of them (generally, I think it's best if we keep those to the minimum when we merge).

I should mention also that for now I only skimmed over the TypeErasure.scala and the Inlines.scala parts, so I might have more suggestions for those later on.

One more NIT is that the code here, especially with the comments, can be very wide, which can make it hard to read depending on the device. Currently in the repo we don't have a standardized formatting unfortunately, but there were plans at some point to have scalafmt enabled. So for especially wide code like

case Specialization(spec) if spec.isFullySpecializedToTopClassesOrNothing => specializations.addErasedImplementation(spec) // We never inline into anonymous class instances (avoids cycles in inline trait inlining), so 
                                                                                                                           // all anonymous class instances must have a non-anonymous class final representation as an $impl$ class.

I think it's best to be proactive and uninline the comment:

// We never inline into anonymous class instances (avoids cycles in inline trait inlining), so
// all anonymous class instances must have a non-anonymous class final representation as an $impl$ class.
case Specialization(spec) if spec.isFullySpecializedToTopClassesOrNothing => specializations.addErasedImplementation(spec)

since I don't really know if scalafmt would be able to work with this (it might do something way uglier!). Most of the code is fine of course, I'm only speaking here mostly about a few examples, most of them in DesugarSpecialisedTraits

(Also sorry about the late-night notification! I don't trust GitHub nor my laptop to keep my comments saved until the morning...)

[jchyb]

This seems like an repeat of the first case case tree: TypeDef if tree.symbol.isInlineTrait (unreachable code)

[jchyb]

The nested Inlined nodes can be generally useful for reporting/checks etc., so if there is no need for it, maybe this could be removed

[jmesyou]

It indeed did not make a single difference to any of the testcases so I removed it

[jchyb]

This might be too much to change at this point (I don't know the timeline of the project), but much of the comments here are concerned about the MegaPhase and MiniPhase structure. Considering we only have List(DesugarSpecialisedTraits, SpecialiseInlineTraits) in the MegaPhase here, and both of those operate on the same trees, and are closely coupled (requiring a transformFollowing call the next phase here), I feel like it would be the simplest to just have those as one phase, avoiding the issues/concerns (reading the docs, it seems like those started as amore legitimate latter transform phases and then were moved as problems were discovered, so I understand how we arrived at this point).

[jmesyou]

I did do some experimenting with this. However, I decided to specifically leave this these two as separate to try and distinguish the work between inline traits and specialized traits. DesugarSpecialisedTraits is specifically related to the generation of specialized implementations bounded by Specialized, while SpecializeInlineTraits is related both to inline and specialized traits. Could we leave this as follow on change as it might be too complicated of a refactor? i.e. group DesugarSpecialisedTraits and SpecialiseInlineTraits as a single pass.

[jchyb]

I feel like we should split up the logic for inlining for methods and traits into different files (e.g. keep Inlines.scala for methods and add TraitInlines.scala etc.). Here we have a needsInlining check that covers both cases, but then we have to check what the actual tree is and delegate it to the correct inlining procedure anyway, which themselves are totally different. Not to mention the Inlines.scala is already too big and confusing to navigate through (if there is any reused logic between the 2 variants, I think it can stay here)

[jmesyou]

I'm going to give this a try. The current state where we check whether needs inlining is not very elegant. At first glance, it does not look like there's too much shared logic but there is going to be some minor duplication.

[jmesyou]

After some experimenting, this turned out to be a bigger than expected change. Perhaps like the MegaPhase change, we could do it as a follow-on? I can create issues to track this items, the MegaPhases one and other remaining TODOs in this PR

[jchyb]

Sure, that sounds good, although from the 2 things I would give some more priority to the Inlines.scala split, since many different contributors are bound to touch that file for things unrelated to inline traits (and it shouldn't be too difficult to split, I'd mainly want to keep InlineParentTraits, inlineParentInlineTraits and their utils out of Inlines.scala to keep things manageable).

[jchyb]

A little confused around the changes here - aren't all inlined traits inlined in DesugarSpecialisedTraits? Or is it just the specialized ones, to do it early?

[jmesyou]

DesugarInlineTraits only generates the specialized class implementation for traits that have a type parameter with a Specialized bound, while the remaining passes do the lifting of actually inlining the defs from an inline trait into their respective sites.

Using this class as an example:

//> using options -language:experimental.specializedTraits

inline trait Foo[T: Specialized](x: T):
    def foo = x

def f(b: Foo[Int]) = 37 + b.foo
 
@main def main =
    val x = new Foo[Int](42) {}
    f(x)

Without the Specialized bound, x has the following tree after erasure:

val x: Foo =
{
       final class $anon() extends Object(), Foo(Int.box(42)) {
       private val Foo$$x: Int = 42
       override def foo(): Int = this.Foo$$x
       override def foo(): Object = Int.box(this.foo())
         
new Object with Foo {...}():Foo
}

whereas with the bound, there is specialized implementation generated with inline definitions:

  class Foo$impl$scala$Int(Foo$impl$scala$Int$$x: Int) extends Object(),
    Foo$$sp$scala$Int(), Foo(Int.box(this.Foo$impl$scala$Int$$x)), 
    Foo$$sp$scala$Int {
    private val Foo$$x: Int = this.Foo$impl$scala$Int$$x
    override def foo(): Int = this.Foo$$x
    private val Foo$impl$scala$Int$$x: Int
    override def foo(): Object = Int.box(this.foo())
  }

val x: Foo$$sp$scala$Int = new Foo$impl$scala$Int(42):Foo$$sp$scala$Int

[jchyb]

I somewhat misunderstood what the TODO comment was talking about and thought for a little bit we did some additional work in the Inlining phase. Sorry about that (and thank you for explaining).

[jchyb]

Agreeing with the comment, better to move those into Specialization

[jchyb]

I think the current way is fine

[jchyb]

If the Specialization of two applications results in the same trees being generated (and that seems to be the case), then I think we can call them equal without any concerns here - I think the TODO can be removed

[jchyb]

I think we could put paramToArgList into the constructor of the trait that includes the val specializedTypeArgs and share it that way, but that does seem indeed annoying so the current repetition is ok in my opinion

[jchyb]

Can we split the additions here into another method (like handleClassDefSpecialisations)? this is pretty independent from the rest of the method and the comments assume we know the code below refers to specialization

[jmesyou]

Fixed compilation issues with error codes before addressing comments by @jchyb (:

[jchyb]

Suggested change

	recycleAnInteger(new RecyclingBin[Anyval]() {}) // RecyclingBin[AnyVal] can be interpreted as RecyclingBin[Int] due to contravariance
	recycleAnInteger(new RecyclingBin[Any]() {}) // RecyclingBin[Any] can be interpreted as RecyclingBin[Int] due to contravariance
	recycleAnInteger(new RecyclingBin[AnyVal]() {}) // RecyclingBin[AnyVal] can be interpreted as RecyclingBin[Int] due to contravariance
	recycleAnInteger(new RecyclingBin[Any]() {}) // RecyclingBin[Any] can be interpreted as RecyclingBin[Int] due to contravariance

I think the typo is causing the CI crashes here. But it also exposes a small issue with how ErrorTypes are handled, with the specType method crashing on them. ErrorTypes are types used by the compiler in places that error out, so that a phase like Typer is able to report multiple errors, not just one and they don't conform to the regular type hierarchy.

Besides the fix above, could you add an additional neg test with something like:

//> using options -experimental -language:experimental.inlineTraits

object Snippet {
  inline trait RecyclingBin[-T: Specialized]:
      def recycle(x: T) = println(s"Recycling ${x}")

  def recycleAnInteger(bin: RecyclingBin[Int]) = bin.recycle(100)
  def recycleAnInteger(bin: Int) = () // second overload forces overload resolution

  recycleAnInteger(new RecyclingBin[Anyval]() {}) // Anyval is a typo for AnyVal
}

and a fix for it (I think returning a defn.ObjectClass.typeRef for ErrorTypes would be fine)

[jmesyou]

Ah, thanks! Let me fix that.

Edit: Fixed

[jmesyou]

@jchyb I think the change is ready for another look. With the exception of splitting the logic between inline methods and inline traits into separate files and the grouping of MegaPhase into a single pass. I did a pretty big syntactical refactor of DesugarSpecializedTraits, so hopefully it is an easier read now.

[jmesyou]

New test failure with:

tests/run-macros/tasty-extractors-owners

Possibly unrelated? Fixed

[SolalPirelli]

The PR changes the .check file, and the output matches the .check that's on main, so revert those changes?

[jchyb]

Apologies for awkward review timing. The DesugarSpecialisedTraits are way easier to read now - thank you for adjusting that. I still have some suggestions for some of the TODOs - I fear if we don’t resolve them now they will continue to stay unresolved, especially those pointing to more minor concerns. There are a couple that point to bigger issues/concerns, which I feel like are actually okay too leave in here for the added context (or for the MiniPhase merger later), so I didn’t tag them. I think after this round we’ll be able to merge.

[jchyb]

Let's remove the TODO

Suggested change

	inline val SPECIALIZED_TRAIT_SUFFIX = "$$sp" // TODO: Added an extra $ to avoid name conflict with scala 2 specialization when we do isSpecializedTraitInterfaceName but maybe ideally there's a better way.
	inline val SPECIALIZED_TRAIT_SUFFIX_SCALA_3 = "$$sp" // Added an extra $ to avoid name conflict with scala 2 specialization.

[jchyb]

Recursively inlining is fine, let's remove the TODO (Especially if we plan to merge this into a single phase anyway)

Suggested change

// TODO: We might be able to fix that with transformFollowing

[jchyb]

Let's remove the TODO, current differentiation is fine

[jchyb]

Some duplication is fine, if the if check below completely lines up with another in RefChecks, we could try to extract that one into a global scope and call RefChecks.hasConflictingMembers in both places, but if not, I think we can remove the TODO.

[jchyb]

Leftover TODO from previous review round, var is fine, let's remove the TODO

[jchyb]

I think if/when caching is added, it will be more obvious what to do with this, but I think currently the name is fine.

Let's remove the TODO.

[jmesyou]

I experimented with this a bit, my intuition was that the cache should be mutable as part of the last review round but decided similarly that it will become more obvious which is the better approach when caching is actually implemented.

[jchyb]

Maybe something simple but descriptive like:

Suggested change

	// TODO: Maybe can improve the name of this field
	val specialization: List[Type] =
	// Type params mapped to specialised arguments or left alone (when Specialized isn't used)
	val allTypeParamsRemapped: List[Type] =

(just a suggestion, I mostly want to remove the TODO and stick with something)

[jchyb]

Maybe:

Suggested change

	// TODO: This gets around the fact that inline traits doesn't define inlinedMethod
	// correctly but maybe there is a better way. Ultimately inline traits are not
	// inline methods!
	// This is reused through InlineTraitAncestors for inline traits, so inlinedMethod might not exist there

[jchyb]

I suspect it could cause some issues, but can't say for sure. Ideally this would be Inlining-only (and Typer, for transparent inlines). I suspect that if something comes up its should be easy enough to change (so we can leave the TODO intact for now, I think)

[jchyb]

Let's stick to the current way and remove the TODO:

Suggested change

	// TODO: Maybe we don't want to import by default; maybe we do but we would rather not create another special case ImportFns for the importing.
	// Problem which required the special case: the new Specialized lives in scala.specialize. This is to avoid conflict with the Scala2 specialized annotation.
	// Therefore it is not imported by default with everything else that lives in the scala package.
	private val SpecializeImportFns: RootRef =
	// The new Specialized lives in scala.specialize. This is to avoid conflict with the Scala2 specialized annotation.
	// Therefore it is not imported by default with everything else that lives in the scala package. so we additionally import it here
	private val SpecializeImportFns: RootRef =

[jmesyou]

@jchyb Thanks for taking a look through the change again (: I have removed all the TODOs that you've listed in the latest push. I also reorganized and cleaned up inline-traits.md as this latest push. It previously had some internal notes appended. I have since relocated those notes into their appropriate sections. That leaves the remaining TODOs grouped into two issues that I will open:

1. Phasing of trait inlining and specialized trait generation
2. Splitting trait inlining into its own respective file TraitInlining.scala from Inlining.scala

Other than those, there are no remaining TODOs

[jmesyou]

Resolved merge conflicts

edit: looks like the test failure is a test machine failure of some kind

[jchyb]

Just noticed some more TODO's with simple fixes (or, in the case of the test ones, maybe we don't need to worry about type inference quirks). Sorry about not catching them earlier. I promise after this we can merge.

[jchyb]

Tested locally and all tests pass, it seems that we don't need to retype here

Suggested change

	override protected def inlineCopier: tpd.TreeCopier = new TypedTreeCopier() {
	// TODO: Timothée left this comment: "it feels weird... Is this correct?"
	override def Apply(tree: Tree)(fun: Tree, args: List[Tree])(using Context): Apply =
	untpd.cpy.Apply(tree)(fun, args).withTypeUnchecked(tree.tpe)
	}
	override protected def inlineCopier: tpd.TreeCopier = new TypedTreeCopier()

[jchyb]

Suggested change

	// TODO: Is this what we want? Should it be Foo?
	// need to discuss.

[jchyb]

Suggested change

	val xs: List[Double] = :+:(9.1, :+:(68.52, :+:(18.4, :+:(83.5, Nill[Double]())))) // TODO : Can we prevent the need for an explicit type here or at least make it clearer
	val xs: List[Double] = :+:(9.1, :+:(68.52, :+:(18.4, :+:(83.5, Nill[Double]()))))

[jchyb]

Suggested change

	def foreach[S](f: T => Unit): Unit = (this: List[T]) match { // TODO: Can we avoid the need to cast this to List[T] here? Should it not already be of that type?
	def foreach[S](f: T => Unit): Unit = (this: List[T]) match {

[jmesyou]

Updated + I found a few more and cleaned up those too. I probably should've done a thorough grep beforehand, sorry about that. I have compiled a list of the remaining TODOs, I think most have been covered to some extent as part of code review but all are worth keeping.

❯ git log -2 -p | grep "TODO"
+      // TODO: Show errors where they actually show up in the inline trait and not at the type of the user 
+    // TODO: Depending on how we ultimately organize the phasing
+  // TODO: I reckon to get the best miniphase style processing we can do everything except 
+  // TODO: Do we want to compress this more by adopting e.g. specializedTypeNames from scala 2? 
+          else if (tree.symbol.isTypeParam && tree.symbol.owner.isClass) tree.span // TODO is this the correct span
-            else if (tree.symbol.isTypeParam && tree.symbol.owner.isClass) tree.span // TODO is this the correct span? -- indentation change
+      // TODO: We need to stop inlining if there is a non-inline trait or class that sits between the inline trait and the current class.
+        // TODO: This inlines also some calls to inline defs that were made in the inline trait body, is that ok? 
+      // TODO: check that things are inlined properly

[jchyb]

Thank you so much! The PR looks good now. There is a weird CI issue with that doesn't seem to be related to anything here: assertion failed: wrong source set for $this.self # -1 of class dotty.tools.dotc.ast.Trees$Select, set to compiler/src/dotty/tools/dotc/core/Contexts.scala but context had compiler/src/dotty/tools/dotc/core/TypeApplications.scala - I can reproduce it locally, and bringing back the custom inlineCopier doesn't fix it, so it has to be caused by something from the rebase. I'll try to investigate more when I have time

[jmesyou]

I was able to reproduce the test failures locally on the main branch. Did a quick investigation and it looks like the problematic commit is a4e1b441306d129813fdb8c5739edac0bfe7b9cd. Prior commits have clean tests.