CreateCompanionObjects transformer

A transformer that provides a convenient way to create
companion objects.

Author: DarkDimius

src/dotty/tools/dotc/transform/CreateCompanionObjects.scala

@@ -0,0 +1,53 @@
+package dotty.tools.dotc.transform
+
+import dotty.tools.dotc.transform.TreeTransforms.{TransformerInfo, TreeTransform, TreeTransformer}
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts.Context
+import scala.collection.mutable.ListBuffer
+import dotty.tools.dotc.core.{Scopes, Flags}
+import dotty.tools.dotc.core.Symbols.NoSymbol
+import scala.annotation.tailrec
+import dotty.tools.dotc.core._
+import Symbols._
+import scala.Some
+import dotty.tools.dotc.transform.TreeTransforms.{NXTransformations, TransformerInfo, TreeTransform, TreeTransformer}
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts.Context
+import scala.collection.mutable
+import dotty.tools.dotc.core.Names.Name
+import NameOps._
+
+/** A transformer that provides a convenient way to create companion objects
+  */
+abstract class CreateCompanionObjects(group: TreeTransformer, idx: Int) extends TreeTransform(group, idx) {
+
+  import tpd._
+
+  /** Given class definition should return true if companion object creation should be enforced
+    */
+  def predicate(cls: TypeDef): Boolean
+
+  override def transformStats(trees: List[Tree])(implicit ctx: Context, info: TransformerInfo): List[tpd.Tree] = {
+    @tailrec
+    def transformStats0(trees: List[Tree], acc: ListBuffer[Tree]): List[Tree] = {
+      trees match {
+        case Nil => acc.toList
+        case (claz: TypeDef) :: stats if claz.symbol.isClass && !(claz.symbol is Flags.Module) => {
+          val moduleExists = !(claz.symbol.companionModule eq NoSymbol)
+          if (moduleExists || !predicate(claz)) transformStats0(stats, acc += claz)
+          else {
+            val moduleSymbol = ctx.newCompleteModuleSymbol(claz.symbol.owner, claz.name.toTermName, Flags.Synthetic, Flags.Synthetic, List(defn.ObjectClass.typeRef), Scopes.newScope)
+            if (moduleSymbol.owner.isClass) moduleSymbol.entered
+            val companion = tpd.ModuleDef(moduleSymbol, List(EmptyTree))
+            acc += claz
+            acc += companion
+            transformStats0(stats, acc)
+          }
+        }
+        case stat :: stats => transformStats0(stats, acc += stat)
+      }
+    }
+
+    transformStats0(trees, ListBuffer())
+  }
+}

test/test/transform/CreateCompanionObjectsTest.scala

@@ -0,0 +1,117 @@
+package test.transform
+
+
+import org.junit.{Assert, Test}
+import test.DottyTest
+import dotty.tools.dotc.core._
+import dotty.tools.dotc.ast.{tpd, Trees}
+import Contexts._
+import Flags._
+import Denotations._
+import NameOps._
+import Symbols._
+import Types._
+import Decorators._
+import Trees._
+import dotty.tools.dotc.transform.TreeTransforms.{TreeTransform, TreeTransformer}
+import dotty.tools.dotc.transform.PostTyperTransformers.PostTyperTransformer
+import dotty.tools.dotc.transform.CreateCompanionObjects
+
+
+class CreateCompanionObjectsTest extends DottyTest {
+
+  import tpd._
+
+  @Test
+  def shouldCreateNonExistingObjectsInPackage = checkCompile("frontend", "class A{} ") {
+    (tree, context) =>
+      implicit val ctx = context
+
+      val transformer = new PostTyperTransformer {
+        override def transformations = Array(new CreateCompanionObjects(_, _) {
+          override def predicate(cts: TypeDef): Boolean = true
+        })
+
+        override def name: String = "test"
+      }
+      val transformed = transformer.transform(tree).toString
+      println(transformed)
+      val classPattern = "TypeDef(Modifiers(,,List()),A,"
+      val classPos = transformed.indexOf(classPattern)
+      val moduleClassPattern = "TypeDef(Modifiers(final module <synthetic>,,List()),A$"
+      val modulePos = transformed.indexOf(moduleClassPattern)
+
+      Assert.assertTrue("should create non-existing objects in package",
+        classPos < modulePos
+      )
+  }
+
+  @Test
+  def shouldCreateNonExistingObjectsInBlock = checkCompile("frontend", "class D {def p = {class A{}; 1}} ") {
+    (tree, context) =>
+      implicit val ctx = context
+      val transformer = new PostTyperTransformer {
+        override def transformations = Array(new CreateCompanionObjects(_, _) {
+          override def predicate(cts: TypeDef): Boolean = true
+        })
+
+        override def name: String = "test"
+      }
+      val transformed = transformer.transform(tree).toString
+      val classPattern = "TypeDef(Modifiers(,,List()),A,"
+      val classPos = transformed.indexOf(classPattern)
+      val moduleClassPattern = "TypeDef(Modifiers(final module <synthetic>,,List()),A$"
+      val modulePos = transformed.indexOf(moduleClassPattern)
+
+      Assert.assertTrue("should create non-existing objects in block",
+        classPos < modulePos
+      )
+  }
+
+  @Test
+  def shouldCreateNonExistingObjectsInTemplate = checkCompile("frontend", "class D {class A{}; } ") {
+    (tree, context) =>
+      implicit val ctx = context
+      val transformer = new PostTyperTransformer {
+        override def transformations = Array(new CreateCompanionObjects(_, _) {
+          override def predicate(cts: TypeDef): Boolean = true
+        })
+
+        override def name: String = "test"
+      }
+      val transformed = transformer.transform(tree).toString
+      val classPattern = "TypeDef(Modifiers(,,List()),A,"
+      val classPos = transformed.indexOf(classPattern)
+      val moduleClassPattern = "TypeDef(Modifiers(final module <synthetic>,,List()),A$"
+      val modulePos = transformed.indexOf(moduleClassPattern)
+
+      Assert.assertTrue("should create non-existing objects in template",
+        classPos < modulePos
+      )
+  }
+
+  @Test
+  def shouldCreateOnlyIfAskedFor = checkCompile("frontend", "class DONT {class CREATE{}; } ") {
+    (tree, context) =>
+      implicit val ctx = context
+      val transformer = new PostTyperTransformer {
+        override def transformations = Array(new CreateCompanionObjects(_, _) {
+          override def predicate(cts: TypeDef): Boolean = cts.name.toString.contains("CREATE")
+        })
+
+        override def name: String = "test"
+      }
+      val transformed = transformer.transform(tree).toString
+      val classPattern = "TypeDef(Modifiers(,,List()),A,"
+      val classPos = transformed.indexOf(classPattern)
+      val moduleClassPattern = "TypeDef(Modifiers(final module <synthetic>,,List()),CREATE$"
+      val modulePos = transformed.indexOf(moduleClassPattern)
+
+      val notCreatedModulePattern = "TypeDef(Modifiers(final module <synthetic>,,List()),DONT"
+      val notCreatedPos = transformed.indexOf(notCreatedModulePattern)
+
+      Assert.assertTrue("should create non-existing objects in template",
+        classPos < modulePos && (notCreatedPos < 0)
+      )
+  }
+}