Support TagMacro parameter substitution for polymorphic function types on Scala 3

izumi-reflect/izumi-reflect/src/main/scala-3/izumi/reflect/TagMacro.scala

@@ -62,6 +62,22 @@ final class TagMacro(using override val qctx: Quotes) extends InspectorBase {
       }
     }
 
+    def summonLTTIfAvailable(typeRepr: TypeRepr): Option[Expr[LightTypeTag]] = {
+      typeRepr match {
+        case _: TypeBounds =>
+          None
+        case _ =>
+          val tagTypeRepr = AppliedType(tagSymbolTypeRef, List(typeRepr))
+          Implicits.search(tagTypeRepr) match {
+            case s: ImplicitSearchSuccess =>
+              val tagExpr = s.tree.asExpr.asInstanceOf[Expr[Tag[?]]]
+              Some('{ $tagExpr.tag })
+            case _: ImplicitSearchFailure =>
+              None
+          }
+      }
+    }
+
     def summonIfNotLambdaParamOf(typeRepr: TypeRepr, lam: TypeRepr): Expr[Option[LightTypeTag]] = {
       if (isLambdaParamOf(typeRepr, lam)) {
         '{ None }
@@ -178,43 +194,98 @@ final class TagMacro(using override val qctx: Quotes) extends InspectorBase {
         val cls = closestClassOfTypeRepr(parent)
         val parentLtt = summonLTT(parent)
 
-        val (allTypeMembers, termMembers) = members.partitionMap {
-          case (s, n, tb: TypeBounds) if allPartsStrong(owners, tb) => Left(Left((n, tb)))
-          case (_, n, TypeBounds(lo, hi)) if lo == hi => Left(Right((n, hi)))
-          case (_, _, tb @ TypeBounds(lo, hi)) =>
+        members.foreach {
+          case (_, _, tb @ TypeBounds(lo, hi)) if lo != hi && !allPartsStrong(owners, tb) =>
             report.errorAndAbort(
               s"TagMacro: resolving type parameters inside type bounds is not supported, got weak types in bounds=${tb.show}, in type=$typeReprDealiased"
             )
-          case x => Right(x)
+          case _ =>
         }
-        val (strongTypeBounds, weakTypeMembers) = allTypeMembers.partitionMap(identity)
-        // FIXME: once we add resolution for method/val members too, not just type members
-        //  this struct will no longer be 'weak'. In fact we'll want to add a new constructor
-        //  instead of `refinedTag` that will be better suited to fully resolved struct tags
-        val termAndStrongTpesOnlyWeakStructLtt = {
-          val termOnlyRefinementTypeRepr = termMembers.foldRight(defn.AnyRefClass.typeRef: TypeRepr) {
-            case ((_, name, tpe), refinement) =>
-              Refinement(parent = refinement, name = name, info = tpe)
+
+        val weakTypeMemberSubstitutions = {
+          val resolvedWeakMembers = mutable.LinkedHashMap.empty[String, Expr[LightTypeTag]]
+          def tryResolveWeak(part: TypeRepr, localBinders: Set[TypeRepr]): Unit = {
+            if (ReflectionUtil.topLevelWeakType(owners, localBinders, part)) {
+              val sym = part.typeSymbol
+              if (!sym.isNoSymbol) {
+                val key = sym.fullName
+                if (!resolvedWeakMembers.contains(key)) {
+                  summonLTTIfAvailable(part).foreach { ltt =>
+                    resolvedWeakMembers += key -> ltt
+                  }
+                }
+              }
+            }
           }
-          val withStrongTpesRefinementTypeRepr = strongTypeBounds.foldRight(termOnlyRefinementTypeRepr) {
-            case ((name, tpe), refinement) =>
-              Refinement(parent = refinement, name = name, info = tpe)
+
+          def collectWeakMembers(info0: TypeRepr, localBinders: Set[TypeRepr]): Unit = {
+            val info = info0._dealiasSimplifiedFull
+            tryResolveWeak(info, localBinders)
+            info match {
+              case AppliedType(tycon, args) =>
+                collectWeakMembers(tycon, localBinders)
+                args.foreach(collectWeakMembers(_, localBinders))
+              case AndType(lhs, rhs) =>
+                collectWeakMembers(lhs, localBinders)
+                collectWeakMembers(rhs, localBinders)
+              case OrType(lhs, rhs) =>
+                collectWeakMembers(lhs, localBinders)
+                collectWeakMembers(rhs, localBinders)
+              case TypeRef(prefix, _) =>
+                collectWeakMembers(prefix, localBinders)
+              case TermRef(prefix, _) =>
+                collectWeakMembers(prefix, localBinders)
+              case ThisType(prefix) =>
+                collectWeakMembers(prefix, localBinders)
+              case TypeBounds(lo, hi) =>
+                collectWeakMembers(lo, localBinders)
+                collectWeakMembers(hi, localBinders)
+              case lam @ TypeLambda(_, _, res) =>
+                collectWeakMembers(res, localBinders + lam)
+              case poly @ PolyType(_, bounds, res) =>
+                bounds.foreach { case TypeBounds(lo, hi) =>
+                  collectWeakMembers(lo, localBinders + poly)
+                  collectWeakMembers(hi, localBinders + poly)
+                }
+                collectWeakMembers(res, localBinders + poly)
+              case MethodType(_, argTypes, res) =>
+                argTypes.foreach(collectWeakMembers(_, localBinders))
+                collectWeakMembers(res, localBinders)
+              case ByNameType(inner) =>
+                collectWeakMembers(inner, localBinders)
+              case refinementInfo: Refinement =>
+                collectWeakMembers(refinementInfo.parent, localBinders)
+                collectWeakMembers(refinementInfo.info, localBinders)
+              case _: ParamRef | _: ConstantType | _: NoPrefix =>
+              case _: TypeRepr =>
+            }
           }
-          Inspect.inspectTypeRepr(withStrongTpesRefinementTypeRepr)
+
+          members.iterator
+            .flatMap { case (_, _, info) => refinementInfoToParts(info) }
+            .foreach(collectWeakMembers(_, Set.empty))
+          resolvedWeakMembers.toList
         }
-        val resolvedTypeMemberLtts = weakTypeMembers.map {
-          case (name, tpe) => '{ (${ Expr(name) }, ${ summonLTT(tpe) }) }
+
+        val unresolvedWeakStructLtt = Inspect.inspectTypeRepr(refinement)
+        val resolvedWeakTypeMembersExpr = weakTypeMemberSubstitutions.map {
+          case (name, ltt) => '{ (${ Expr(name) }, $ltt) }
+        }
+        val resolvedStructLtt = if (resolvedWeakTypeMembersExpr.isEmpty) {
+          unresolvedWeakStructLtt
+        } else {
+          '{ Tag.resolveWeakTypeMembers(${ unresolvedWeakStructLtt }, Map(${ Varargs(resolvedWeakTypeMembersExpr) }: _*)) }
         }
-        // NB: we're resolving LTTs anew for all type members here, instead of optimizing
-        // to resolve only for 'weak' members as in Scala 2.
+
         log(
           s"""Got refinement $refinement
              |parent=$parent
-             |members=$members
-             |closestClass=$cls
-             |""".stripMargin
+              |members=$members
+              |closestClass=$cls
+              |resolvedWeakMembers=${weakTypeMemberSubstitutions.map(_._1)}
+              |""".stripMargin
         )
-        '{ Tag.refinedTag[T](${ cls }, List(${ parentLtt }), ${ termAndStrongTpesOnlyWeakStructLtt }, Map(${ Varargs(resolvedTypeMemberLtts) }: _*)) }
+        '{ Tag.refinedTag[T](${ cls }, List(${ parentLtt }), ${ resolvedStructLtt }, Map.empty) }
 
       // error: the entire type is just a proper type parameter with no type arguments
       // it cannot be resolved further

izumi-reflect/izumi-reflect/src/main/scala-3/izumi/reflect/Tags.scala

@@ -22,6 +22,7 @@ import izumi.reflect.dottyreflection.Inspect
 import izumi.reflect.macrortti.{LTag, LightTypeTag, LightTypeTagRef}
 
 import scala.annotation.implicitNotFound
+import scala.collection.mutable
 
 trait AnyTag extends Serializable {
   def tag: LightTypeTag
@@ -82,8 +83,6 @@ trait AnyTag extends Serializable {
   * Without a `TagK` constraint above, this example would fail with `no TypeTag available for MyService[F]` error
   *
   * Currently some limitations apply as to when a `Tag` will be correctly constructed:
-  *   * Type Parameters do not yet resolve in structural refinement methods, e.g. T in {{{ Tag[{ def x: T}] }}}
-  *     They do resolve in refinement type members however, e.g. {{{ Tag[ Any { type Out = T } ] }}}
   *   * TagK* does not resolve for constructors with bounded parameters, e.g. S in {{{ class Abc[S <: String]; TagK[Abc] }}}
   *     (You can still have a bound in partial application: e.g. {{{ class Abc[S <: String, A]; TagK[Abc["hi", _]] }}}
   *   * Further details at [[https://github.com/7mind/izumi/issues/374]]
@@ -165,6 +164,62 @@ object Tag {
     Tag(lubClass, LightTypeTag.refinedType(intersection, structType, additionalTypeMembers))
   }
 
+  // Internal helper used by TagMacro to substitute weak outer type members inside structural defs/vals.
+  def resolveWeakTypeMembers(structType: LightTypeTag, replacements: Map[String, LightTypeTag]): LightTypeTag = {
+    if (replacements.isEmpty) {
+      structType
+    } else {
+      val replacementRefs = replacements.iterator.map {
+        case (name, ltt) =>
+          val replacementRef = ltt.ref match {
+            case ref: LightTypeTagRef.AbstractReference => ref
+          }
+          LightTypeTagRef.SymName.SymTypeName(name) -> replacementRef
+      }.toMap
+
+      val rewriter = new izumi.reflect.macrortti.RuntimeAPI.Rewriter(replacementRefs)
+
+      def rewrite(ref: LightTypeTagRef.AbstractReference): LightTypeTagRef.AbstractReference = {
+        rewriter.replaceRefs(ref)
+      }
+
+      def rewriteName(ref: LightTypeTagRef.NameReference): LightTypeTagRef.NameReference = {
+        rewrite(ref) match {
+          case n: LightTypeTagRef.NameReference => n
+          case _ => ref
+        }
+      }
+
+      def mergeDb[T](entries: Iterator[(T, Set[T])]): Map[T, Set[T]] = {
+        val acc = mutable.HashMap.empty[T, Set[T]]
+        entries.foreach {
+          case (k, v) =>
+            val existing = acc.getOrElse(k, Set.empty)
+            acc.update(k, existing ++ v.filterNot(_ == k))
+        }
+        acc.toMap
+      }
+
+      val rewrittenRef = rewrite(structType.ref.asInstanceOf[LightTypeTagRef.AbstractReference])
+      val rewrittenBases = structType.basesdb.iterator.map {
+        case (child, parents) =>
+          rewrite(child) -> parents.map(rewrite)
+      }
+      val rewrittenInheritance = structType.idb.iterator.map {
+        case (child, parents) =>
+          rewriteName(child) -> parents.map(rewriteName)
+      }
+      val mergedBases = mergeDb(
+        rewrittenBases ++ replacements.valuesIterator.flatMap(_.basesdb.iterator)
+      )
+      val mergedInheritance = mergeDb(
+        rewrittenInheritance ++ replacements.valuesIterator.flatMap(_.idb.iterator)
+      )
+
+      LightTypeTag(rewrittenRef, mergedBases, mergedInheritance)
+    }
+  }
+
   inline implicit final def tagFromTagMacro[T <: AnyKind]: Tag[T] = ${ TagMacro.createTagExpr[T] }
 }
 

izumi-reflect/izumi-reflect/src/main/scala-3/izumi/reflect/dottyreflection/ReflectionUtil.scala

@@ -336,6 +336,17 @@ private[reflect] object ReflectionUtil {
       case ThisType(tpe) => allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, tpe)
       case NoPrefix() => true
       case TypeBounds(lo, hi) => allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, lo) && allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, hi)
+      case MethodType(_, args, res) =>
+        args.forall(allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, _)) &&
+        allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, res)
+      case poly @ PolyType(_, bounds, res) =>
+        val withLocalBinders = outerLambdas + poly
+        bounds.forall {
+          case TypeBounds(lo, hi) =>
+            allPartsStrong(shift, outerOwnerClassDefs, withLocalBinders, lo) &&
+            allPartsStrong(shift, outerOwnerClassDefs, withLocalBinders, hi)
+        } &&
+        allPartsStrong(shift, outerOwnerClassDefs, withLocalBinders, res)
       case lam @ TypeLambda(_, _, body) => allPartsStrong(shift, outerOwnerClassDefs, outerLambdas + lam, body)
       case Refinement(parent, _, tpe) => allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, tpe) && allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, parent)
       case ByNameType(tpe) => allPartsStrong(shift, outerOwnerClassDefs, outerLambdas, tpe)

izumi-reflect/izumi-reflect/src/main/scala/izumi/reflect/macrortti/RuntimeAPI.scala

@@ -96,8 +96,10 @@ object RuntimeAPI {
     res
   }
 
-  final class Rewriter(_rules: Map[SymName.LambdaParamName, AbstractReference]) {
-    private val rules: Map[SymName, AbstractReference] = _rules.toMap
+  final class Rewriter(_rules: collection.immutable.Map[_ <: SymName, AbstractReference]) {
+    private val rules: Map[SymName, AbstractReference] = _rules.iterator.map {
+      case (k, v) => (k: SymName) -> v
+    }.toMap
 
     def complete(@unused context: AppliedNamedReference, ref: AbstractReference): AbstractReference = {
       ref
@@ -108,7 +110,11 @@ object RuntimeAPI {
       reference match {
         case l: Lambda =>
           val bad = l.input.iterator.toSet
-          val fixed = new Rewriter(_rules.filterKeys(!bad.contains(_)).iterator.toMap).replaceRefs(l.output)
+          val fixedRules = rules.iterator.filterNot {
+            case (k: SymName.LambdaParamName, _) => bad.contains(k)
+            case _ => false
+          }.toMap
+          val fixed = new Rewriter(fixedRules).replaceRefs(l.output)
           l.copy(output = fixed)
 
         case o: AppliedReference =>

izumi-reflect/izumi-reflect/src/test/scala-3/izumi/reflect/test/TagTest.scala

@@ -82,6 +82,25 @@ class TagTest extends SharedTagTest with TagAssertions {
       assertSameStrict(goodCombine.tag, Tag[Int with Unit].tag)
     }
 
+    "resolve outer type parameters inside structural defs and vals" in {
+      def t1[T: Tag]: Tag[{ def x: T }] = Tag[{ def x: T }]
+      def t2[T: Tag]: Tag[{ val x: T }] = Tag[{ val x: T }]
+
+      assertSameStrict(t1[Int].tag, Tag[{ def x: Int }].tag)
+      assertSameStrict(t2[Int].tag, Tag[{ val x: Int }].tag)
+    }
+
+    "resolve outer type parameters inside polymorphic function type refinements" in {
+      trait CustomError
+
+      def polyTag[AppF[_]: TagK, InjF[_]: TagK] = Tag[[A] => AppF[A] => InjF[Either[CustomError, A]]]
+
+      val combined = polyTag[Option, List]
+      val concrete = Tag[[A] => Option[A] => List[Either[CustomError, A]]]
+
+      assertSameStrict(combined.tag, concrete.tag)
+    }
+
   }
 
 }

izumi-reflect/izumi-reflect/src/test/scala/izumi/reflect/test/SharedTagProgressionTest.scala

@@ -28,10 +28,10 @@ abstract class SharedTagProgressionTest extends AnyWordSpec with TagAssertions w
       def t1[T: Tag]: Tag[{ def x: T }] = Tag[{ def x: T }]
       def t2[T: Tag]: Tag[{ val x: T }] = Tag[{ val x: T }]
 
-      broken {
+      brokenOnScala2 {
         assertSameStrict(t1[Int].tag, Tag[{ def x: Int }].tag)
       }
-      broken {
+      brokenOnScala2 {
         assertSameStrict(t2[Int].tag, Tag[{ val x: Int }].tag)
       }
     }