Repository: 7mind/izumi
Pull request #1814 (2022-11-10 19:24:10)

distage: Fix bug in `PlanVerifier` when weak set references to a named key were wrongly treated as strong references


Patch:
diff --git a/distage/distage-core/src/main/scala/izumi/distage/planning/solver/GraphPreparations.scala b/distage/distage-core/src/main/scala/izumi/distage/planning/solver/GraphPreparations.scala
index c8700a1e7..f143b476a 100644
--- a/distage/distage-core/src/main/scala/izumi/distage/planning/solver/GraphPreparations.scala
+++ b/distage/distage-core/src/main/scala/izumi/distage/planning/solver/GraphPreparations.scala
@@ -9,6 +9,7 @@ import izumi.distage.model.planning.AxisPoint
 import izumi.distage.model.reflection.DIKey
 import izumi.distage.planning.BindingTranslator
 import izumi.distage.planning.solver.SemigraphSolver.SemiEdgeSeq
+import izumi.fundamentals.collections.MutableMultiMap
 import izumi.fundamentals.graphs.WeakEdge
 import izumi.fundamentals.graphs.struct.IncidenceMatrix
 import izumi.fundamentals.graphs.tools.gc.Tracer
@@ -23,18 +24,11 @@ class GraphPreparations(
   import scala.collection.compat._
 
   def findWeakSetMembers(
-    sets: Map[Annotated[DIKey], Node[DIKey, InstantiationOp]],
-    matrix: SemiEdgeSeq[Annotated[DIKey], DIKey, InstantiationOp],
+    setOps: Map[Annotated[DIKey], Node[DIKey, InstantiationOp]],
+    execOpIndex: MutableMultiMap[DIKey, InstantiationOp],
     roots: Set[DIKey],
   ): Set[WeakEdge[DIKey]] = {
-    import izumi.fundamentals.collections.IzCollections._
-
-    val indexed = matrix.links.map {
-      case (successor, node) =>
-        (successor.key, node.meta)
-    }.toMultimapMut
-
-    val setDefs = sets
+    val setDefs = setOps
       .collect {
         case (target, Node(_, s: CreateSet)) =>
           (target, s.members)
@@ -46,14 +40,27 @@ class GraphPreparations(
           .diff(roots)
           .flatMap {
             member =>
-              indexed.get(member).toSeq.flatten.collect {
-                case ExecutableOp.WiringOp.ReferenceKey(_, Wiring.SingletonWiring.Reference(_, referenced, true), _) =>
-                  WeakEdge(referenced, member)
-              }
+              getSetElementWeakEdges(execOpIndex, member)
           }
     }.toSet
   }
 
+  def executableOpIndex(matrix: SemiEdgeSeq[Annotated[DIKey], DIKey, InstantiationOp]): MutableMultiMap[DIKey, InstantiationOp] = {
+    import izumi.fundamentals.collections.IzCollections._
+
+    matrix.links.map {
+      case (successor, node) =>
+        (successor.key, node.meta)
+    }.toMultimapMut
+  }
+
+  def getSetElementWeakEdges(execOpIndex: MutableMultiMap[DIKey, InstantiationOp], member: DIKey): Iterator[WeakEdge[DIKey]] = {
+    execOpIndex.getOrElse(member, Nil).iterator.collect {
+      case ExecutableOp.WiringOp.ReferenceKey(_, Wiring.SingletonWiring.Reference(_, referenced, true), _) =>
+        WeakEdge(referenced, member)
+    }
+  }
+
   def getRoots(input: Roots, allOps: Seq[(Annotated[DIKey], InstantiationOp)]): Set[DIKey] = {
     val effective = input match {
       case Roots.Of(roots) =>
diff --git a/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanSolver.scala b/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanSolver.scala
index e249b7a49..8f0951f01 100644
--- a/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanSolver.scala
+++ b/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanSolver.scala
@@ -97,7 +97,7 @@ object PlanSolver {
         preps.getRoots(input.roots, allOps)
 
       val weakSetMembers: Set[WeakEdge[DIKey]] =
-        preps.findWeakSetMembers(sets, matrix, roots)
+        preps.findWeakSetMembers(sets, preps.executableOpIndex(matrix), roots)
 
       Right(Problem(activations, matrix, roots, weakSetMembers))
     }
diff --git a/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanVerifier.scala b/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanVerifier.scala
index 8cf4fa536..daf3c08a7 100644
--- a/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanVerifier.scala
+++ b/distage/distage-core/src/main/scala/izumi/distage/planning/solver/PlanVerifier.scala
@@ -11,14 +11,13 @@ import izumi.distage.model.planning.{ActivationChoices, AxisPoint}
 import izumi.distage.model.reflection.DIKey.SetElementKey
 import izumi.distage.model.reflection.{DIKey, SafeType}
 import izumi.distage.planning.BindingTranslator
-import izumi.distage.planning.solver.PlanVerifier.PlanIssue._
+import izumi.distage.planning.solver.PlanVerifier.PlanIssue.*
 import izumi.distage.planning.solver.PlanVerifier.{PlanIssue, PlanVerifierResult}
 import izumi.distage.planning.solver.SemigraphSolver.SemiEdgeSeq
-import izumi.functional.IzEither._
-import izumi.fundamentals.collections.ImmutableMultiMap
-import izumi.fundamentals.collections.IzCollections._
+import izumi.functional.IzEither.*
+import izumi.fundamentals.collections.{ImmutableMultiMap, MutableMultiMap}
+import izumi.fundamentals.collections.IzCollections.*
 import izumi.fundamentals.collections.nonempty.{NonEmptyList, NonEmptyMap, NonEmptySet}
-import izumi.fundamentals.graphs.WeakEdge
 import izumi.fundamentals.platform.strings.IzString.toRichIterable
 
 import java.util.concurrent.TimeUnit
@@ -69,7 +68,7 @@ class PlanVerifier(
     val justMutators = mutators.map { case (k, op, _) => (k.key, (op, k.axis)) }.toMultimap
 
     val rootKeys: Set[DIKey] = preps.getRoots(roots, justOps)
-    val weakSetMembers: Set[WeakEdge[DIKey]] = preps.findWeakSetMembers(setOps.toMap, matrix, rootKeys)
+    val execOpIndex: MutableMultiMap[DIKey, InstantiationOp] = preps.executableOpIndex(matrix)
 
     val mutVisited = mutable.HashSet.empty[(DIKey, Set[AxisPoint])]
     val effectType = SafeType.getK[F]
@@ -78,7 +77,7 @@ class PlanVerifier(
     var after = before
     val issues =
       try {
-        trace(allAxis, mutVisited, matrixToTrace, weakSetMembers, justMutators, providedKeys, excludedActivations, rootKeys, effectType)
+        trace(allAxis, mutVisited, matrixToTrace, execOpIndex, justMutators, providedKeys, excludedActivations, rootKeys, effectType)
       } finally {
         after = System.currentTimeMillis()
       }
@@ -96,7 +95,7 @@ class PlanVerifier(
     allAxis: Map[String, Set[String]],
     allVisited: mutable.HashSet[(DIKey, Set[AxisPoint])],
     matrix: ImmutableMultiMap[DIKey, (InstantiationOp, Set[AxisPoint])],
-    weakSetMembers: Set[WeakEdge[DIKey]],
+    execOpIndex: MutableMultiMap[DIKey, InstantiationOp],
     justMutators: ImmutableMultiMap[DIKey, (InstantiationOp, Set[AxisPoint])],
     providedKeys: DIKey => Boolean,
     excludedActivations: Set[NonEmptySet[AxisPoint]],
@@ -179,11 +178,11 @@ class PlanVerifier(
                   } else {
                     Right(())
                   }
-                next <- checkConflicts(allAxis, opsWithMergedSets, weakSetMembers, excludedActivations, effectType)
+                next <- checkConflicts(allAxis, opsWithMergedSets, execOpIndex, excludedActivations, effectType)
               } yield {
                 allVisited.add((key, currentActivation))
 
-                val mutators = justMutators.getOrElse(key, Set.empty).iterator.filter(ac allValid _._2).flatMap(m => depsOf(weakSetMembers)(m._1)).toSeq
+                val mutators = justMutators.getOrElse(key, Set.empty).iterator.filter(ac allValid _._2).flatMap(m => depsOf(execOpIndex)(m._1)).toSeq
 
                 val goNext = next.iterator.map {
                   case (nextActivation, nextDeps) =>
@@ -249,7 +248,7 @@ class PlanVerifier(
   protected[this] def checkConflicts(
     allAxis: Map[String, Set[String]],
     withoutCurrentActivations: Set[(InstantiationOp, Set[AxisPoint], Set[AxisPoint])],
-    weakSetMembers: Set[WeakEdge[DIKey]],
+    execOpIndex: MutableMultiMap[DIKey, InstantiationOp],
     excludedActivations: Set[NonEmptySet[AxisPoint]],
     effectType: SafeType,
   ): Either[List[PlanIssue], Seq[(Set[AxisPoint], Set[DIKey])]] = {
@@ -267,7 +266,7 @@ class PlanVerifier(
       val next = withoutCurrentActivations.iterator.map {
         case (op, activations, _) =>
           // TODO: I'm not sure if it's "correct" to "activate" all the points together but it simplifies things greatly
-          val deps = depsOf(weakSetMembers)(op)
+          val deps = depsOf(execOpIndex)(op)
           val acts = op match {
             case _: ExecutableOp.CreateSet =>
               Set.empty[AxisPoint]
@@ -280,13 +279,16 @@ class PlanVerifier(
     }
   }
 
-  protected[this] final def depsOf(weakSetMembers: Set[WeakEdge[DIKey]])(op: InstantiationOp): Set[DIKey] = {
+  protected[this] final def depsOf(
+    execOpIndex: MutableMultiMap[DIKey, InstantiationOp]
+  )(op: InstantiationOp
+  ): Set[DIKey] = {
     op match {
       case cs: CreateSet =>
         // we completely ignore weak members, they don't make any difference in case they are unreachable through other paths
         val members = cs.members.filter {
           case m: SetElementKey =>
-            !weakSetMembers.contains(WeakEdge(m.reference, m))
+            preps.getSetElementWeakEdges(execOpIndex, m).isEmpty
           case _ =>
             true
         }


Test patch:
diff --git a/distage/distage-core/src/test/scala/izumi/distage/injector/PlanVerifierTest.scala b/distage/distage-core/src/test/scala/izumi/distage/injector/PlanVerifierTest.scala
index 66ed5b709..310a41c8a 100644
--- a/distage/distage-core/src/test/scala/izumi/distage/injector/PlanVerifierTest.scala
+++ b/distage/distage-core/src/test/scala/izumi/distage/injector/PlanVerifierTest.scala
@@ -15,6 +15,7 @@ import org.scalatest.exceptions.TestFailedException
 import org.scalatest.wordspec.AnyWordSpec
 
 class PlanVerifierTest extends AnyWordSpec with MkInjector {
+
   "Verifier handles resources (uniform case, Roots.Everything, https://github.com/7mind/izumi/issues/1476)" in {
     import PlanVerifierCase1._
 
@@ -602,4 +603,18 @@ class PlanVerifierTest extends AnyWordSpec with MkInjector {
     val result = PlanVerifier().verify[Identity](definition, Roots.target[Set[Fork2]], Injector.providedKeys(), Set.empty)
     assert(result.issues.isEmpty)
   }
+
+  "Verifier handles weak sets: named weak references are still weak" in {
+    import PlanVerifierCase1._
+
+    val definition = new ModuleDef {
+      many[Fork1]
+        .weak[ImplA]
+        .weak[ImplB]("missing")
+    }
+
+    val result = PlanVerifier().verify[Identity](definition, Roots.target[Set[Fork1]], Injector.providedKeys(), Set.empty)
+    assert(result.issues.isEmpty)
+  }
+
 }