changes to package object and companion object handling

Author: lgmyrek

plugin/src/main/scala/io/estatico/newtype/compat/NewTypeTreeTransformer.scala

@@ -19,32 +19,51 @@ class NewTypeTreeTransformer extends UntypedTreeMap:
 
   /** Transform a list of statements, expanding @newtype/@newsubtype classes. */
   private def transformStats(stats: List[Tree])(using Context): List[Tree] =
+    val statsList = stats.toIndexedSeq
+    // Pre-scan: for each @newtype TypeDef, find a ModuleDef anywhere later in the same
+    // statement list whose term-name matches. Macro paradise on Scala 2 finds companions via the
+    // typer's symbol lookup, which works regardless of intervening `type`/`val` decls. We need to
+    // emulate that here — using only the immediately-next stat is too narrow for codebases that
+    // place companion objects after `type` aliases (a common idiom inside `package object`s).
+    val consumedCompanion = scala.collection.mutable.Set.empty[Int]
+    val companionForNewtype = scala.collection.mutable.Map.empty[Int, Int]
+    for i <- statsList.indices do
+      statsList(i) match
+        case td: TypeDef if td.rhs.isInstanceOf[Template] && findNewTypeAnnotation(td).isDefined =>
+          val target = td.name.toTermName
+          var j = i + 1
+          var found = -1
+          while j < statsList.length && found < 0 do
+            statsList(j) match
+              case md: ModuleDef if md.name == target && !consumedCompanion(j) =>
+                found = j
+              case _ =>
+            j += 1
+          if found >= 0 then
+            companionForNewtype(i) = found
+            consumedCompanion += found
+        case _ =>
+
     val result = ListBuffer.empty[Tree]
     var i = 0
-    val statsList = stats.toIndexedSeq
     while i < statsList.length do
-      val stat = statsList(i)
-      stat match
-        case td: TypeDef if td.rhs.isInstanceOf[Template] =>
-          findNewTypeAnnotation(td) match
-            case Some((annotation, isSubtype)) =>
-              val params = AnnotationParams.extract(annotation, isSubtype)
-              // Look ahead for companion object
-              val companionOpt = if i + 1 < statsList.length then
-                statsList(i + 1) match
-                  case md: ModuleDef if md.name.toTermName == td.name.toTermName => Some(md)
-                  case _ => None
-              else None
-
-              val expanded = expandNewType(td, companionOpt, params)
-              result ++= expanded
-
-              if companionOpt.isDefined then i += 1 // skip companion, we merged it
-            case None =>
-              result += transform(stat)
-        case _ =>
-          result += transform(stat)
-      i += 1
+      if consumedCompanion(i) then
+        i += 1
+      else
+        statsList(i) match
+          case td: TypeDef if td.rhs.isInstanceOf[Template] =>
+            findNewTypeAnnotation(td) match
+              case Some((annotation, isSubtype)) =>
+                val params = AnnotationParams.extract(annotation, isSubtype)
+                val companionOpt = companionForNewtype.get(i)
+                  .map(j => statsList(j).asInstanceOf[ModuleDef])
+                val expanded = expandNewType(td, companionOpt, params)
+                result ++= expanded
+              case None =>
+                result += transform(statsList(i))
+          case _ =>
+            result += transform(statsList(i))
+        i += 1
     result.toList
 
   /** Find @newtype or @newsubtype annotation on a TypeDef. */
@@ -138,11 +157,6 @@ class NewTypeTreeTransformer extends UntypedTreeMap:
          |  def derivingK[TC[_[_]]](implicit ev: TC[Repr]): TC[Type] = ev.asInstanceOf[TC[Type]]""".stripMargin
     else ""
 
-    // Merge with existing companion body
-    val existingCompanionCode = companionOpt.map { md =>
-      md.impl.body.map(plain(_)).mkString("\n  ", "\n  ", "")
-    }.getOrElse("")
-
     // Coercible implicits: for parameterized types, use def with type params
     val C = "_root_.io.estatico.newtype.Coercible"
     val coercibleImplicits = if tparams.isEmpty then
@@ -182,15 +196,35 @@ class NewTypeTreeTransformer extends UntypedTreeMap:
       |$opsCode
       |$derivingCode
       |$derivingKCode
-      |$existingCompanionCode
       |}""".stripMargin
 
     if params.debug then
       println(s"[newtype-compat] Generated source for @${if params.isSubtype then "newsubtype" else "newtype"} $nameStr:")
       source.linesIterator.foreach(line => println(s"[newtype-compat]   $line"))
 
-    // Parse the generated source code into untyped trees
-    parseStats(source)
+    // Parse the generated source code into untyped trees, then splice the companion's body
+    // structurally (rather than via tree.show round-trip — which emits the `module` soft keyword
+    // for nested objects and produces source that won't reparse).
+    val stats = parseStats(source)
+    companionOpt match
+      case None => stats
+      case Some(companion) =>
+        val companionTermName = name.toTermName
+        stats.map {
+          case md: ModuleDef if md.name == companionTermName =>
+            // Splice the companion's full Template into the synthesized one. We need to carry
+            // over `parents` (e.g. `extends WithGenericSalesforceId[...]`), `derived` (any
+            // `derives` clauses), and `self` from the original — the synthesized object only
+            // provides default-empty values for those. Body members from both are concatenated.
+            val mergedImpl = cpy.Template(md.impl)(
+              parents = if companion.impl.parents.nonEmpty then companion.impl.parents else md.impl.parents,
+              derived = if companion.impl.derived.nonEmpty then companion.impl.derived else md.impl.derived,
+              self    = if companion.impl.self != md.impl.self then companion.impl.self else md.impl.self,
+              body    = md.impl.body ++ companion.impl.body
+            )
+            untpd.ModuleDef(md.name, mergedImpl).withMods(md.mods).withSpan(md.span)
+          case other => other
+        }
 
   private def buildOpsCodeFromMethods(
     nameStr: String, reprNameStr: String, reprTypeStr: String,
@@ -285,6 +319,9 @@ class NewTypeTreeTransformer extends UntypedTreeMap:
         self = transformSub(md.impl.self),
         body = transformStats(md.impl.body)
       )
-      untpd.ModuleDef(md.name, newImpl).withSpan(md.span)
+      // Preserve the original modifiers (incl. the package-object marker) — without `withMods`,
+      // a `package object` would be reconstructed as a regular `object`, breaking sub-packages
+      // and parent-package-object scope inheritance.
+      untpd.ModuleDef(md.name, newImpl).withMods(md.mods).withSpan(md.span)
     case _ =>
       super.transform(tree)

tests/src/test/scala/io/estatico/newtype/macros/NewTypeCompanionAfterTypeAliasesTest.scala

@@ -0,0 +1,40 @@
+package io.estatico.newtype.macros
+
+import io.estatico.newtype.ops._
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+/*
+ * Regression: a `@newtype` declaration followed by `type` aliases, and then its companion object.
+ *
+ * On Scala 2 this is fine because macro paradise looks up the companion symbol via the typer
+ * regardless of intervening declarations. The original Scala 3 plugin only inspected the
+ * immediately-next stat (`statsList(i + 1)`) for a companion, so this idiom (common inside
+ * `package object` blocks where helper aliases are colocated with the newtype) caused the
+ * companion to be missed entirely. The synthesized companion and the original then coexisted as
+ * two `object Foo` declarations, producing `is already defined` errors.
+ */
+class NewTypeCompanionAfterTypeAliasesTest extends AnyFlatSpec with Matchers {
+
+  import NewTypeCompanionAfterTypeAliasesTest._
+
+  behavior of "@newtype companion separated from declaration by type aliases"
+
+  it should "still merge the companion that follows after intervening `type` decls" in {
+    val id = OrderId("ord-1")
+    id.coerce[String] shouldBe "ord-1"
+    OrderId.fromString("ord-2").coerce[String] shouldBe "ord-2"
+    OrderId.Default.coerce[String] shouldBe "ord-default"
+  }
+}
+
+object NewTypeCompanionAfterTypeAliasesTest {
+  @newtype final case class OrderId(value: String)
+  // Intervening type aliases — must not prevent the companion below from being recognized.
+  type OrderIdHelper = String
+  type OrderIdAnotherHelper = OrderIdHelper
+  object OrderId {
+    def fromString(s: String): OrderId = OrderId(s)
+    val Default: OrderId = OrderId("ord-default")
+  }
+}

tests/src/test/scala/io/estatico/newtype/macros/NewTypeCompanionExtendsTraitTest.scala

@@ -0,0 +1,82 @@
+package io.estatico.newtype.macros
+
+import io.estatico.newtype.ops._
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+/*
+ * Regression: a `@newtype` whose companion `extends` a trait (and supplies abstract members or
+ * overrides via the trait's API).
+ *
+ * Mirrors the chili pattern in `salesforce/package.scala`:
+ *
+ *     @newtype final case class ContactId(value: String Refined ContactIdRefinement)
+ *     object ContactId extends WithGenericSalesforceId[ContactIdRefinement] { override type N = ContactId }
+ *
+ * The plugin merges `companion.impl.body` into the synthesized object's body, but it does not
+ * propagate `companion.impl.parents`, `companion.impl.derived`, or `companion.impl.self`. As a
+ * result the trait extension is silently dropped — the synthesized object becomes a plain
+ * `object Foo { ... }` with the body but without the parent it was declared to extend.
+ *
+ * The visible failure is downstream: callers expecting `Foo: HasN` lose the inherited members,
+ * and `override`s in the companion body refer to symbols that no longer exist on the companion.
+ */
+class NewTypeCompanionExtendsTraitTest extends AnyFlatSpec with Matchers {
+
+  import NewTypeCompanionExtendsTraitTest._
+
+  behavior of "@newtype companion that extends a trait"
+
+  it should "preserve the parent trait so inherited members are reachable" in {
+    Foo.label shouldBe "foo-label"
+    Foo("x").coerce[String] shouldBe "x"
+    (Foo: HasLabel) shouldBe Foo
+  }
+
+  it should "preserve overrides that depend on the parent trait" in {
+    Bar.emptyString shouldBe ""
+    (Bar: HasEmptyString[BarMarker]).emptyString shouldBe ""
+  }
+
+  it should "preserve generic parent traits with type-member overrides" in {
+    // Mirrors the chili pattern: `object ContactId extends WithGenericSalesforceId[ContactIdRefinement] { override type N = ContactId }`.
+    Baz.label shouldBe "baz-refined"
+    val ev: WithRefinement[BazRefinement] = Baz
+    ev.label shouldBe "baz-refined"
+  }
+}
+
+object NewTypeCompanionExtendsTraitTest {
+
+  trait HasLabel {
+    def label: String
+  }
+
+  trait HasEmptyString[Phantom] {
+    def emptyString: String
+  }
+
+  trait WithRefinement[R] {
+    type N
+    def label: String
+  }
+
+  @newtype final case class Foo(value: String)
+  object Foo extends HasLabel {
+    val label: String = "foo-label"
+  }
+
+  final class BarMarker
+  @newtype final case class Bar(value: String)
+  object Bar extends HasEmptyString[BarMarker] {
+    val emptyString: String = ""
+  }
+
+  // BazRefinement is a phantom type used as a refinement marker, like `ContactIdRefinement` in chili.
+  final class BazRefinement
+  @newtype final case class Baz(value: String)
+  object Baz extends WithRefinement[BazRefinement] {
+    override type N = Baz
+    val label: String = "baz-refined"
+  }
+}

tests/src/test/scala/io/estatico/newtype/macros/NewTypeCompanionWithNestedObjectTest.scala

@@ -0,0 +1,47 @@
+package io.estatico.newtype.macros
+
+import io.estatico.newtype.ops._
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+/*
+ * Regression: `@newtype` with a companion whose body contains a nested `object`.
+ *
+ * The plugin merges the existing companion body by calling `tree.show` on each member of
+ * `companionOpt.impl.body`, then splicing the resulting strings into a generated source that
+ * gets re-parsed by `Parsers.Parser`. For nested `object`s, Dotty's untyped `tree.show` emits the
+ * `module` soft keyword (e.g. `module object Patterns { ... }`), which is not valid Scala source.
+ * The resulting source fails parsing inside `<newtype-generated-N>`.
+ *
+ * This is the same failure mode observed when cross-compiling chili `core-domain` to Scala 3:
+ *
+ *     [error] <newtype-generated-84>:35:10: end of statement expected but 'object' found
+ *     [error]   module object Predef {
+ *
+ * A fix likely needs to render companion body members through a printer that does not emit the
+ * `module` modifier, or rebuild the companion structurally rather than via string splicing.
+ */
+class NewTypeCompanionWithNestedObjectTest extends AnyFlatSpec with Matchers {
+
+  import NewTypeCompanionWithNestedObjectTest._
+
+  behavior of "@newtype companion containing a nested object"
+
+  it should "preserve a nested object declared in the companion" in {
+    val s = Phone("+1 555 0100")
+    s.coerce[String] shouldBe "+1 555 0100"
+    Phone.Patterns.us shouldBe "^\\+1.*"
+    Phone.fromString("+1 555 0100").coerce[String] shouldBe "+1 555 0100"
+  }
+}
+
+object NewTypeCompanionWithNestedObjectTest {
+  @newtype final case class Phone(value: String)
+  object Phone {
+    object Patterns {
+      val us: String = "^\\+1.*"
+      val intl: String = "^\\+\\d{1,3}.*"
+    }
+    def fromString(s: String): Phone = Phone(s)
+  }
+}

tests/src/test/scala/io/estatico/newtype/macros/NewTypeInPackageObjectTest.scala

@@ -0,0 +1,51 @@
+package io.estatico.newtype.macros
+
+import io.estatico.newtype.ops._
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+/*
+ * Tests for `@newtype` declarations placed directly inside a `package object` body, with and
+ * without companions, plus a nested `package object` that depends on its parent's declarations.
+ *
+ * Fixtures live in:
+ *   - tests/.../pkgobj/package.scala
+ *   - tests/.../pkgobj/nested/package.scala
+ *
+ * On Scala 3, the rebuilt `ModuleDef` for a package object drops the `Package` modifier, so the
+ * package object is emitted as a regular object. The runtime tests below illustrate the user-
+ * visible consequences (nested package inaccessible, parent type alias not in scope from the
+ * nested package). They are expected to pass once the plugin preserves package-object identity.
+ */
+class NewTypeInPackageObjectTest extends AnyFlatSpec with Matchers {
+
+  behavior of "@newtype inside a package object"
+
+  it should "expand a @newtype with a companion containing typeclass instances" in {
+    val t = pkgobj.Tenant("acme")
+    t.coerce[String] shouldBe "acme"
+    pkgobj.Tenant.Default.coerce[String] shouldBe "default"
+    pkgobj.Tenant.of("foo").coerce[String] shouldBe "foo"
+    val ord = implicitly[Ordering[pkgobj.Tenant]]
+    ord.compare(pkgobj.Tenant("a"), pkgobj.Tenant("b")) should be < 0
+  }
+
+  it should "expand a @newtype whose companion only declares vals" in {
+    pkgobj.User.Anonymous.coerce[String] shouldBe "anonymous"
+  }
+
+  it should "expand a @newtype that has no companion at all" in {
+    pkgobj.JobName("nightly").coerce[String] shouldBe "nightly"
+  }
+
+  behavior of "@newtype inside a nested package object"
+
+  it should "resolve type aliases declared in the parent package object" in {
+    pkgobj.nested.SubField("x").coerce[String] shouldBe "x"
+    pkgobj.nested.SubField.Sentinel.coerce[String] shouldBe "sentinel"
+  }
+
+  it should "expand a @newtype without a companion in a nested package object" in {
+    pkgobj.nested.SubLabel("y").coerce[String] shouldBe "y"
+  }
+}

tests/src/test/scala/io/estatico/newtype/macros/pkgobj/nested/package.scala

@@ -0,0 +1,27 @@
+package io.estatico.newtype.macros.pkgobj
+
+import io.estatico.newtype.macros.newtype
+
+/*
+ * Fixture: a nested package object inside a package whose parent declares a `package object`.
+ *
+ * Mirrors the chili `package object common.crm { @newtype case class CrmObject(value: RefinedSaneString) }`
+ * pattern, where `RefinedSaneString` lives in the parent `package object common`.
+ *
+ * Two pieces are exercised together:
+ *   - the parent's package-object types must remain in scope inside this nested package, and
+ *   - `@newtype` expansion must work inside a nested package object's body too.
+ *
+ * If the plugin loses the `Package` flag on the parent's package object, both fail: the parent is
+ * emitted as a regular object, the package `pkgobj.nested` cannot exist underneath a class named
+ * `pkgobj`, and `SaneString` is no longer reachable here without a fully-qualified import.
+ */
+package object nested {
+
+  @newtype final case class SubField(value: SaneString)
+  object SubField {
+    val Sentinel: SubField = SubField("sentinel")
+  }
+
+  @newtype final case class SubLabel(value: SaneString)
+}