Inline Traits & Specialized Traits

bench-micro/src/main/scala/dotty/tools/benchmarks/SpecializedTraitsBenchmark.scala

@@ -0,0 +1,100 @@
+// Run with: scala-cli --power --jmh bench-micro/src/main/scala/dotty/tools/benchmarks/SpecializedTraitsBenchmark.scala
+// May have to run it again / delete .scala-build and rerun if you get a class not found
+// error from scala-cli first time - the --jmh flag is still experimental.
+// Don't forget to publish the compiler first and check that the version below corresponds to the generated version,
+// as well as to kill the bloop server if you are republishing the same version as before.
+// scala-cli --power bloop exit
+
+//> using scala 3.9.0-RC1-bin-SNAPSHOT-nonbootstrapped
+//> using options -language:experimental.specializedTraits
+
+package dotty.tools.benchmarks
+
+import org.openjdk.jmh.annotations.*
+import java.util.concurrent.TimeUnit
+
+class VecManual(elems: Array[Int]):
+  private val num = summon[Numeric[Int]]
+
+  def length = elems.length
+
+  def apply(i: Int): Int = elems(i)
+
+  def scalarProduct(other: VecManual): Int =
+    require(this.length == other.length)
+    var result = num.fromInt(0)
+    for i <- 0 until length do
+      result = num.plus(result, num.times(this(i), other(i)))
+    result
+
+class VecGeneric[T: Numeric](elems: Array[T]):
+  private val num = summon[Numeric[T]]
+
+  def length = elems.length
+
+  def apply(i: Int): T = elems(i)
+
+  def scalarProduct(other: VecGeneric[T]): T =
+    require(this.length == other.length)
+    var result = num.fromInt(0)
+    for i <- 0 until length do
+      result = num.plus(result, num.times(this(i), other(i)))
+    result
+
+inline trait VecSpecialized[T: {Specialized, Numeric2}](elems: Array[T]):
+  private val num = summon[Numeric2[T]]
+
+  def length = elems.length
+
+  def apply(i: Int): T = elems(i)
+
+  def scalarProduct(other: VecSpecialized[T]): T =
+    require(this.length == other.length)
+    var result = num.fromInt(0)
+    for i <- 0 until length do
+      result = num.plus(result, num.times(this(i), other(i)))
+    result
+
+@State(Scope.Benchmark)
+class Arrays:
+  var arr1 = Array.fill(100_000_000) {math.round(math.random().floatValue * 4)}
+  var arr2 = Array.fill(100_000_000) {math.round(math.random().floatValue * 4)}
+  val target = arr1.zip(arr2).map((x, y) => x * y).fold(0)(_ + _)
+
+@State(Scope.Benchmark)
+@BenchmarkMode(Array(Mode.AverageTime))
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Warmup(iterations = 5, time = 100, timeUnit = TimeUnit.MILLISECONDS)
+@Measurement(iterations = 15, time = 100, timeUnit = TimeUnit.MILLISECONDS)
+@Fork(1)
+class VecBench:
+  @Benchmark
+  def manual(arr: Arrays) =
+    val x = VecManual(arr.arr1)
+    val y = VecManual(arr.arr2)
+    assert(x.scalarProduct(y) == arr.target)
+
+  @Benchmark
+  def generic(arr: Arrays) =
+    val x = VecGeneric[Int](arr.arr1)
+    val y = VecGeneric[Int](arr.arr2)
+    assert(x.scalarProduct(y) == arr.target)
+
+  @Benchmark
+  def specialized(arr: Arrays) =
+    val x = new VecSpecialized[Int](arr.arr1) {}
+    val y = new VecSpecialized[Int](arr.arr2) {}
+    assert(x.scalarProduct(y) == arr.target)
+
+// You can really see the impact of Specialized on the interface usage here
+// Remove Specialized and see that the generated code gets much more boxing and unboxing
+// which slows it down substantially. 
+inline trait Numeric2[T: Specialized]:
+  def fromInt(x: Int): T
+  def plus(x: T, y: T): T
+  def times(x: T, y: T): T
+
+implicit object IntIsIntegral extends Numeric2[Int]:
+  override def fromInt(x: Int): Int = x
+  override def plus(x: Int, y: Int): Int = x + y
+  override def times(x: Int, y: Int): Int = x * y

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -43,6 +43,8 @@ class Compiler {
     List(new UnrollDefinitions) ::  // Unroll annotated methods if detected in PostTyper
     List(new sjs.PrepJSInterop) ::  // Additional checks and transformations for Scala.js (Scala.js only)
     List(new SetRootTree) ::        // Set the `rootTreeOrProvider` on class symbols
+    List(new DesugarSpecializedTraits) :: // Process Specialized traits
+    List(new SpecializeInlineTraits) ::   // Inline the code of inline traits into their children
     Nil
 
   /** Phases dealing with TASTY tree pickling and unpickling */
@@ -81,7 +83,9 @@ class Compiler {
          new ForwardDepChecks,       // Check that there are no forward references to local vals
          new SpecializeApplyMethods, // Adds specialized methods to FunctionN
          new TryCatchPatterns,       // Compile cases in try/catch
-         new PatternMatcher) ::      // Compile pattern matches
+         new PatternMatcher,         // Compile pattern matches
+         new PruneSpecializedMethods,// Remove specialized methods which have already been inlined
+         new PruneInlineTraits) ::   // Remove right-hand side of definitions in inline traits
     List(new TestRecheck.Pre) ::     // Test only: run rechecker, enabled under -Yrecheck-test
     List(new TestRecheck) ::         // Test only: run rechecker, enabled under -Yrecheck-test
     List(new cc.Setup) ::            // Preparations for check captures phase, enabled under captureChecking

compiler/src/dotty/tools/dotc/ast/TreeTypeMap.scala

@@ -98,7 +98,7 @@ class TreeTypeMap(
 
   override def transform(tree: Tree)(using Context): Tree = treeMap(tree) match {
     case impl @ Template(constr, _, self, _) =>
-      val tmap = withMappedSyms(localSyms(impl :: self :: Nil))
+      val tmap = withMappedSyms(localSyms(impl :: self :: Nil)) 
       cpy.Template(impl)(
           constr = tmap.transformSub(constr),
           parents = impl.parents.mapconserve(transform),
@@ -127,14 +127,15 @@ class TreeTypeMap(
           cpy.Block(blk)(stats1, expr1)
         case lit @ Literal(Constant(tpe: Type)) =>
           cpy.Literal(lit)(Constant(mapType(tpe)))
-        case ddef @ DefDef(name, paramss, tpt, _) =>
+        case ddef @ DefDef(name, paramss, tpt, _) => // Why are we not correctly mapping foo's return type? See Reached def def ...
           val (tmap1, paramss1) = transformAllParamss(paramss)
           val res = cpy.DefDef(ddef)(name, paramss1, tmap1.transform(tpt), tmap1.transform(ddef.rhs))
           res.symbol.setParamssFromDefs(paramss1)
           res.symbol.transformAnnotations {
             case ann: BodyAnnotation => ann.derivedAnnotation(transform(ann.tree))
             case ann => ann
-          }
+          } 
+          // HERE?
           res
         case tdef @ LambdaTypeTree(tparams, body) =>
           val (tmap1, tparams1) = transformDefs(tparams)

compiler/src/dotty/tools/dotc/config/Feature.scala

@@ -28,6 +28,7 @@ object Feature:
 
   val dependent = experimental("dependent")
   val erasedDefinitions = experimental("erasedDefinitions")
+  val specializedTraits = experimental("specializedTraits")
   val strictEqualityPatternMatching = experimental("strictEqualityPatternMatching")
   val symbolLiterals = deprecated("symbolLiterals")
   val saferExceptions = experimental("saferExceptions")
@@ -67,6 +68,7 @@ object Feature:
     (scala2macros, "Allow Scala 2 macros"),
     (dependent, "Allow dependent method types"),
     (erasedDefinitions, "Allow erased definitions"),
+    (specializedTraits, "Allow specialized traits"),
     (strictEqualityPatternMatching, "relaxed CanEqual checks for ADT pattern matching"),
     (symbolLiterals, "Allow symbol literals"),
     (saferExceptions, "Enable safer exceptions"),

compiler/src/dotty/tools/dotc/core/Contexts.scala

@@ -42,6 +42,8 @@ import plugins.*
 import java.nio.file.InvalidPathException
 import dotty.tools.dotc.coverage.Coverage
 import scala.annotation.tailrec
+import dotty.tools.dotc.inlines.Inlines.InlineTraitState
+import dotty.tools.dotc.transform.SpecializedTraitState
 
 object Contexts {
 
@@ -146,6 +148,8 @@ object Contexts {
     def typerState: TyperState
     def gadt: GadtConstraint = gadtState.gadt
     def gadtState: GadtState
+    def inlineTraitState: InlineTraitState
+    def specializedTraitState: SpecializedTraitState
     def searchHistory: SearchHistory
     def source: SourceFile
 
@@ -436,6 +440,8 @@ object Contexts {
       superOrThisCallContext(owner, constrCtx.scope)
         .setTyperState(typerState)
         .setGadtState(gadtState)
+        .setInlineTraitState(inlineTraitState)
+        .setSpecializedTraitState(specializedTraitState)
         .fresh
         .setScope(this.scope)
     }
@@ -603,6 +609,12 @@ object Contexts {
 
     private var _gadtState: GadtState = uninitialized
     final def gadtState: GadtState = _gadtState
+
+    private var _inlineTraitState: InlineTraitState = uninitialized
+    final def inlineTraitState: InlineTraitState = _inlineTraitState
+
+    private var _specializedTraitState: SpecializedTraitState = uninitialized
+    final def specializedTraitState: SpecializedTraitState = _specializedTraitState
 
     private var _searchHistory: SearchHistory = uninitialized
     final def searchHistory: SearchHistory = _searchHistory
@@ -628,6 +640,8 @@ object Contexts {
       _tree = origin.tree
       _scope = origin.scope
       _gadtState = origin.gadtState
+      _inlineTraitState = origin.inlineTraitState
+      _specializedTraitState = origin.specializedTraitState
       _searchHistory = origin.searchHistory
       _source = origin.source
       _moreProperties = origin.moreProperties
@@ -691,6 +705,16 @@ object Contexts {
     def setFreshGADTBounds: this.type =
       setGadtState(gadtState.fresh)
 
+    def setInlineTraitState(inlineTraitState: InlineTraitState): this.type =
+      util.Stats.record("Context.setInlineTraitState")
+      this._inlineTraitState = inlineTraitState
+      this
+
+    def setSpecializedTraitState(specializedTraitState: SpecializedTraitState): this.type =
+      util.Stats.record("Context.setSpecializedTraitState")
+      this._specializedTraitState = specializedTraitState
+      this
+
     def setSearchHistory(searchHistory: SearchHistory): this.type =
       util.Stats.record("Context.setSearchHistory")
       this._searchHistory = searchHistory
@@ -785,6 +809,8 @@ object Contexts {
           .updated(profilerLoc, Profiler.NoOp)
       c._searchHistory = new SearchRoot
       c._gadtState = GadtState(GadtConstraint.empty)
+      c._inlineTraitState = InlineTraitState()
+      c._specializedTraitState = SpecializedTraitState()
       c
   end FreshContext
 

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -236,6 +236,7 @@ class Definitions {
   @tu lazy val ScalaCollectionImmutablePackageClass: ClassSymbol = requiredPackage("scala.collection.immutable").moduleClass.asClass
   @tu lazy val ScalaMathPackageClass: ClassSymbol = requiredPackage("scala.math").moduleClass.asClass
   @tu lazy val ScalaUtilPackageClass: ClassSymbol = requiredPackage("scala.util").moduleClass.asClass
+  @tu lazy val ScalaSpecializePackageVal: TermSymbol = requiredPackage("scala.specialize")
 
   // fundamental modules
   @tu lazy val SysPackage : Symbol = requiredModule("scala.sys.package")
@@ -793,6 +794,10 @@ class Definitions {
   @tu lazy val StringAddClass    : ClassSymbol = requiredClass("scala.runtime.StringAdd")
     @tu lazy val StringAdd_+ : Symbol = StringAddClass.requiredMethod(nme.raw.PLUS)
 
+  @tu lazy val SpecializedClass : ClassSymbol = requiredClass("scala.specialize.Specialized")
+  @tu lazy val SpecializedModule: Symbol = SpecializedClass.companionModule
+    @tu lazy val SpecializedModule_apply: Symbol = SpecializedModule.requiredMethod(nme.apply)
+
   @tu lazy val StringContextClass: ClassSymbol = requiredClass("scala.StringContext")
     @tu lazy val StringContext_s  : Symbol = StringContextClass.requiredMethod(nme.s)
     @tu lazy val StringContext_raw: Symbol = StringContextClass.requiredMethod(nme.raw_)
@@ -1734,6 +1739,9 @@ class Definitions {
   private val PredefImportFns: RootRef =
     RootRef(() => ScalaPredefModule.termRef)
 
+  private val SpecializeImportFns: RootRef = // TODO: Find a solution to importing in scala package without conflict with original specialized.
+    RootRef(() => ScalaSpecializePackageVal.termRef)
+
   @tu private lazy val YimportsImportFns: List[RootRef] = ctx.settings.Yimports.value.map { name =>
     val denot =
       getModuleIfDefined(name).suchThat(_.is(Module)) `orElse`
@@ -1749,8 +1757,8 @@ class Definitions {
   @tu private lazy val ScalaRootImportFns: List[RootRef] =
     if !ctx.settings.Yimports.isDefault then YimportsImportFns
     else if ctx.settings.YnoImports.value then Nil
-    else if ctx.settings.YnoPredef.value then ScalaImportFns
-    else ScalaImportFns :+ PredefImportFns
+    else if ctx.settings.YnoPredef.value then ScalaImportFns :+ SpecializeImportFns
+    else ScalaImportFns :+ SpecializeImportFns :+ PredefImportFns
 
   @tu private lazy val JavaRootImportTypes: List[TermRef] = JavaRootImportFns.map(_.refFn())
   @tu private lazy val ScalaRootImportTypes: List[TermRef] = ScalaRootImportFns.map(_.refFn())

compiler/src/dotty/tools/dotc/core/Flags.scala

@@ -453,13 +453,13 @@ object Flags {
 
   /** Flags representing source modifiers */
   private val CommonSourceModifierFlags: FlagSet =
-    commonFlags(Private, Protected, Final, Case, Implicit, Given, Override, JavaStatic, Transparent, Erased)
+    commonFlags(Private, Protected, Final, Case, Implicit, Given, Override, JavaStatic, Transparent, Erased, Inline)
 
   val TypeSourceModifierFlags: FlagSet =
     CommonSourceModifierFlags.toTypeFlags | Abstract | Sealed | Opaque | Open | Into
 
   val TermSourceModifierFlags: FlagSet =
-    CommonSourceModifierFlags.toTermFlags | Inline | AbsOverride | Lazy | Tracked
+    CommonSourceModifierFlags.toTermFlags | AbsOverride | Lazy | Tracked
 
   /** Flags representing modifiers that can appear in trees */
   val ModifierFlags: FlagSet =
@@ -592,6 +592,7 @@ object Flags {
   val InlineOrProxy: FlagSet                 = Inline | InlineProxy                           // An inline method or inline argument proxy */
   val InlineMethod: FlagSet                  = Inline | Method
   val InlineImplicitMethod: FlagSet          = Implicit | InlineMethod
+  val InlineTrait: FlagSet                   = Inline | Trait
   val InlineParam: FlagSet                   = Inline | Param
   val InlineByNameProxy: FlagSet             = InlineProxy | Method
   val JavaEnum: FlagSet                      = JavaDefined | Enum                             // A Java enum trait

compiler/src/dotty/tools/dotc/core/NameOps.scala

@@ -77,6 +77,9 @@ object NameOps {
     def isAnonymousFunctionName: Boolean = name.startsWith(str.ANON_FUN)
     def isUnapplyName: Boolean = name == nme.unapply || name == nme.unapplySeq
     def isRightAssocOperatorName: Boolean = name.lastPart.last == ':'
+    def isSpecializedTraitInterfaceName: Boolean = name.toString.contains(str.SPECIALIZED_TRAIT_SUFFIX)
+    def isSpecializedTraitImplementationName: Boolean = name.toString.contains(str.SPECIALIZED_TRAIT_IMPL_SUFFIX)
+    def isRawSpecializedTraitImplementationName: Boolean = name.toString.endsWith(str.SPECIALIZED_TRAIT_IMPL_SUFFIX)
 
     /** Does this name match `[{letter | digit} '_'] op`?
       *

compiler/src/dotty/tools/dotc/core/Phases.scala

@@ -234,6 +234,8 @@ object Phases {
     private var mySbtExtractAPIPhase: Phase = uninitialized
     private var myPicklerPhase: Phase = uninitialized
     private var mySetRootTreePhase: Phase = uninitialized
+    private var mySpecializeInlineTraitsPhase: Phase = uninitialized
+    private var myDesugarSpecializedTraitsPhase: Phase = uninitialized
     private var myInliningPhase: Phase = uninitialized
     private var myStagingPhase: Phase = uninitialized
     private var mySplicingPhase: Phase = uninitialized
@@ -266,6 +268,8 @@ object Phases {
     final def sbtExtractAPIPhase: Phase = mySbtExtractAPIPhase
     final def picklerPhase: Phase = myPicklerPhase
     final def setRootTreePhase: Phase = mySetRootTreePhase
+    final def specializeInlineTraitsPhase: Phase = mySpecializeInlineTraitsPhase
+    final def desugarSpecializedTraitsPhase: Phase = myDesugarSpecializedTraitsPhase
     final def inliningPhase: Phase = myInliningPhase
     final def stagingPhase: Phase = myStagingPhase
     final def splicingPhase: Phase = mySplicingPhase
@@ -298,6 +302,8 @@ object Phases {
       mySbtExtractAPIPhase = phaseOfClass(classOf[sbt.ExtractAPI])
       mySetRootTreePhase = phaseOfClass(classOf[SetRootTree])
       myPicklerPhase = phaseOfClass(classOf[Pickler])
+      mySpecializeInlineTraitsPhase = phaseOfClass(classOf[SpecializeInlineTraits])
+      myDesugarSpecializedTraitsPhase = phaseOfClass(classOf[DesugarSpecializedTraits])
       myInliningPhase = phaseOfClass(classOf[Inlining])
       myStagingPhase = phaseOfClass(classOf[Staging])
       mySplicingPhase = phaseOfClass(classOf[Splicing])
@@ -560,6 +566,8 @@ object Phases {
   def sbtExtractDependenciesPhase(using Context): Phase = ctx.base.sbtExtractDependenciesPhase
   def sbtExtractAPIPhase(using Context): Phase          = ctx.base.sbtExtractAPIPhase
   def picklerPhase(using Context): Phase                = ctx.base.picklerPhase
+  def specializeInlineTraitsPhase(using Context): Phase = ctx.base.specializeInlineTraitsPhase
+  def desugarSpecializedTraitsPhase(using Context): Phase = ctx.base.desugarSpecializedTraitsPhase
   def inliningPhase(using Context): Phase               = ctx.base.inliningPhase
   def stagingPhase(using Context): Phase                = ctx.base.stagingPhase
   def splicingPhase(using Context): Phase               = ctx.base.splicingPhase

compiler/src/dotty/tools/dotc/core/StdNames.scala

@@ -26,6 +26,11 @@ object StdNames {
     inline val LOCALDUMMY_PREFIX        = "<local "       // owner of local blocks
     inline val ANON_CLASS               = "$anon"
     inline val ANON_FUN                 = "$anonfun"
+    inline val INLINE_TRAIT_INNER_CLASS_SUFFIX = "$trait"
+    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_IMPL_SUFFIX = "$impl"
+    inline val SPECIALIZED_TRAIT_TYPE_SEP = "$"
+    /* inline val SPECIALIZED_METHOD_TARGET_NAME_SUFFIX = "$spec" */
 
     inline val REPL_SESSION_LINE  = "rs$line$"
     inline val REPL_ASSIGN_SUFFIX = "$assign"