DSL: Fix .addDependency not working with resource bindings. Throw when .modify/.annotateParameter are used on non-function bindings instead of silently creating a mutator. Throw when parameter is not found in .annotateParameter instead of doing nothing. (#2255)

* DSL: Fix .addDependency not working with resource bindings. Throw when .modify/.annotateParameter are used on non-function bindings instead of silently creating a mutator. Throw when parameter is not found in .annotateParameter instead of doing nothing.

* fix build

* In `annotateParameter`, throw by default if parameter doesn't exist, move old behaviour to `annotateParameterIfExists`

Author: neko-kai

distage/distage-core-api/src/main/scala/izumi/distage/model/definition/dsl/AbstractBindingDefDSL.scala

@@ -300,28 +300,30 @@ object AbstractBindingDefDSL {
     }
 
     override protected def _modifyBy(f: Functoid[T] => Functoid[T]): ModifyTaggingDSL[T] = by(f)
+    override protected def _addDependencies(keys: Iterable[DIKey]): ModifyTaggingDSL[T] = by(_.addDependencies(keys))
   }
 
   trait AddDependencyDSL[T, Self] extends Any {
     protected def _modifyBy(f: Functoid[T] => Functoid[T]): Self
+    protected def _addDependencies(keys: Iterable[DIKey]): Self
 
-    def addDependency[B: Tag]: Self = {
-      _modifyBy(_.addDependency[B])
+    final def addDependency[B: Tag]: Self = {
+      addDependency(DIKey[B])
     }
 
-    def addDependency[B: Tag](name: Identifier): Self = {
-      _modifyBy(_.addDependency[B](name))
+    final def addDependency[B: Tag](name: Identifier): Self = {
+      addDependency(DIKey[B](name))
     }
 
-    def addDependency(key: DIKey): Self = {
-      _modifyBy(_.addDependency(key))
+    final def addDependency(key: DIKey): Self = {
+      _addDependencies(key :: Nil)
     }
 
-    def addDependencies(keys: Iterable[DIKey]): Self = {
-      _modifyBy(_.addDependencies(keys))
+    final def addDependencies(keys: Iterable[DIKey]): Self = {
+      _addDependencies(keys)
     }
 
-    def annotateParameter[P: Tag](name: Identifier): Self = {
+    final def annotateParameter[P: Tag](name: Identifier): Self = {
       _modifyBy(_.annotateParameter[P](name))
     }
   }
@@ -390,7 +392,8 @@ object AbstractBindingDefDSL {
         case _: SetId => 0
         case _: SetIdFromImplName => 1
         case _: Modify[?] => 2
-        case _: AliasTo => 3
+        case _: AddDependencies => 3
+        case _: AliasTo => 4
       }
       sortedOps.foreach {
         case SetImpl(implDef) =>
@@ -403,19 +406,33 @@ object AbstractBindingDefDSL {
         case SetIdFromImplName() =>
           b = b.withTarget(DIKey.IdKey(b.key.tpe, b.implementation.implType.tag.longNameWithPrefix.toLowerCase))
         case Modify(functoidModifier: (Functoid[t] => Functoid[u])) =>
+          b = b.withImplDef(b.implementation match {
+            case implDef: ImplDef.ProviderImpl =>
+              applyFunctoidModifier(implDef, functoidModifier)
+            case ImplDef.ResourceImpl(implType, effectHKTypeCtor, resourceImpl: ImplDef.ProviderImpl) =>
+              ImplDef.ResourceImpl(implType, effectHKTypeCtor, applyFunctoidModifier(resourceImpl, functoidModifier))
+            case ImplDef.EffectImpl(implType, effectHKTypeCtor, effectImpl: ImplDef.ProviderImpl) =>
+              ImplDef.EffectImpl(implType, effectHKTypeCtor, applyFunctoidModifier(effectImpl, functoidModifier))
+            case _ =>
+              throw new InvalidFunctoidModifier(
+                s"""Cannot apply Functoid modifier $functoidModifier to binding $b - Functoid is inaccessible in binding implementation. Expected `ImplDef.ProviderImpl`, but got `ImplDef.${b.implementation.productPrefix}`
+                   |  Please use a separate mutator binding `modify[T].by { <your-modifier> }` instead. (${initial.origin})""".stripMargin
+              )
+          })
+        case AddDependencies(dependencies) =>
+          val functoidModifier = (_: Functoid[Any]).addDependencies(dependencies)
           b.implementation match {
-            case ImplDef.ProviderImpl(implType, function) =>
-              val newProvider = functoidModifier(Functoid(function)).get
-              if (newProvider.ret <:< implType) {
-                b = b.withImplDef(ImplDef.ProviderImpl(implType, newProvider))
-              } else {
-                throw new InvalidFunctoidModifier(
-                  s"Cannot apply invalid Functoid modifier $functoidModifier, new return type `${newProvider.ret}` is not a subtype of the old return type `${function.ret}` (${initial.origin})"
-                )
-              }
+            case providerImpl: ImplDef.ProviderImpl =>
+              b = b.withImplDef(applyFunctoidModifier(providerImpl, functoidModifier))
+            case ImplDef.ResourceImpl(implType, effectHKTypeCtor, resourceImpl: ImplDef.ProviderImpl) =>
+              b = b.withImplDef(ImplDef.ResourceImpl(implType, effectHKTypeCtor, applyFunctoidModifier(resourceImpl, functoidModifier)))
+            case ImplDef.EffectImpl(implType, effectHKTypeCtor, effectImpl: ImplDef.ProviderImpl) =>
+              b = b.withImplDef(ImplDef.EffectImpl(implType, effectHKTypeCtor, applyFunctoidModifier(effectImpl, functoidModifier)))
             case _ =>
               // add an independent mutator instead of modifying the original functoid, if no original functoid is available
-              val newProvider = functoidModifier(Functoid.identityKey[t](b.key)).get
+              // this is ok for `addDependencies` because we don't need to access/modify arguments of the original functoid,
+              // which might be necessary for a general functoid modifier such as `annotateParameter`.
+              val newProvider = Functoid.identityKey[Any](b.key).addDependencies(dependencies).get
               val newRef = SingletonBinding(b.key, ImplDef.ProviderImpl(newProvider.ret, newProvider), Set.empty, b.origin, isMutator = true)
               refs = newRef :: refs
           }
@@ -436,6 +453,18 @@ object AbstractBindingDefDSL {
       ops += op
       this
     }
+
+    private def applyFunctoidModifier[A, B](implDef: ImplDef.ProviderImpl, functoidModifier: Functoid[A] => Functoid[B]): ImplDef.ProviderImpl = {
+      val ImplDef.ProviderImpl(implType, function) = implDef
+      val newProvider = functoidModifier(Functoid(function)).get
+      if (newProvider.ret <:< implType) {
+        ImplDef.ProviderImpl(implType, newProvider)
+      } else {
+        throw new InvalidFunctoidModifier(
+          s"Cannot apply invalid Functoid modifier $functoidModifier, new return type `${newProvider.ret}` is not a subtype of the old return type `${function.ret}` (${initial.origin})"
+        )
+      }
+    }
   }
 
   final class SetRef(initial: EmptySetBinding[DIKey.TypeKey]) extends BindingRef {
@@ -558,6 +587,7 @@ object AbstractBindingDefDSL {
     final case class SetId(id: Identifier) extends SingletonInstruction
     final case class SetIdFromImplName() extends SingletonInstruction
     final case class Modify[T](functoidModifier: Functoid[T] => Functoid[T]) extends SingletonInstruction
+    final case class AddDependencies(dependencies: Iterable[DIKey]) extends SingletonInstruction
     final case class AliasTo(key: DIKey.BasicKey, pos: SourceFilePosition) extends SingletonInstruction
   }
 

distage/distage-core-api/src/main/scala/izumi/distage/model/definition/dsl/ModuleDefDSL.scala

@@ -791,6 +791,10 @@ object ModuleDefDSL {
 
     override protected def _modifyBy(f: Functoid[T] => Functoid[T]): Self = modifyBy(f)
 
+    override protected def _addDependencies(keys: Iterable[DIKey]): Self = {
+      addOp(AddDependencies(keys))(toSame)
+    }
+
   }
 
   final class SetDSL[T](

distage/distage-core/src/main/scala/izumi/distage/modules/support/AnyBIOSupportModule.scala

@@ -30,7 +30,7 @@ class AnyBIOSupportModule[F[+_, +_]: TagKK](implicit t: TagK[F[Throwable, _]], t
 
   make[QuasiIORunner2[F]]
     .from[QuasiIORunner.BIOImpl[F]]
-    .annotateParameter[ExecutionContext]("cpu") // scala.js
+    .modifyBy(_.annotateParameterIfExists[ExecutionContext]("cpu")) // scala.js
 
   make[QuasiIO2[F]]
     .aliased[QuasiPrimitives2[F]]

distage/distage-core/src/test/scala/izumi/distage/dsl/DSLTest.scala

@@ -9,6 +9,7 @@ import izumi.distage.model.definition.Binding.{SetElementBinding, SingletonBindi
 import izumi.distage.model.definition.StandardAxis.{Mode, Repo}
 import izumi.distage.model.definition.dsl.IncludesDSL.TagMergePolicy
 import izumi.distage.model.definition.{Binding, BindingOrigin, BindingTag, Bindings, ImplDef, Lifecycle, Module, ModuleBase}
+import izumi.distage.model.exceptions.dsl.{InvalidFunctoidModifier, ParameterNotFoundForAnnotation}
 import izumi.distage.model.planning.PlanIssue
 import izumi.fundamentals.platform.functional.Identity
 import izumi.fundamentals.platform.language.SourceFilePosition
@@ -806,7 +807,67 @@ class DSLTest extends AnyWordSpec with MkInjector with should.Matchers {
       assert(verification.verificationFailed)
       assert(verification.issues.get.forall(_.isInstanceOf[PlanIssue.MissingImport]))
       val imports = verification.issues.get.toSet.collect { case i: PlanIssue.MissingImport => (i.dependee, i.key) }
-      assert(imports == Set(DIKey[Int] -> DIKey[String], DIKey[Unit]("x") -> DIKey[String]))
+      assert(
+        imports == Set(
+          DIKey[Int] -> DIKey[String],
+          DIKey[Unit]("x") -> DIKey[String],
+        )
+      )
+      // two modifier bindings were added
+      assert(definition.bindings.size == (3 + 2))
+    }
+
+    "addDependency supports adding dependencies for .fromResource/.fromEffect bindings" in {
+      val definition = new ModuleDef {
+        make[Int].fromResource(Lifecycle.pure(5)).addDependency[String]
+        make[Long].fromResource(() => Lifecycle.pure(5L)).addDependency[String]
+        make[Short].fromEffect[Identity, Short](() => 5: Identity[Short]).addDependency[String]
+      }
+
+      val verification = PlanVerifier().verify[Identity](definition, Roots.Everything, Set.empty, Set.empty)
+      assert(verification.verificationFailed)
+      assert(verification.issues.get.forall(_.isInstanceOf[PlanIssue.MissingImport]))
+      val imports = verification.issues.get.toSet.collect { case i: PlanIssue.MissingImport => (i.dependee, i.key) }
+      assert(
+        imports == Set(
+          DIKey[Int] -> DIKey[String],
+          DIKey.ResourceKey(DIKey[Long], SafeType.get[Lifecycle[Identity, Long]]) -> DIKey[String],
+          DIKey.EffectKey(DIKey[Short], SafeType.get[Short]) -> DIKey[String],
+        )
+      )
+      // only one modifier binding was added for .fromValue-like make[Int] binding
+      assert(definition.bindings.size == (3 + 1))
+    }
+
+    "modify & annotateParameter do not support adding dependencies for .fromValue and .using bindings" in {
+      intercept[InvalidFunctoidModifier](new ModuleDef {
+        make[Int]
+          .fromValue(5)
+          .modifyBy(_.addDependency[String])
+      }.bindings)
+      intercept[InvalidFunctoidModifier](new ModuleDef {
+        make[Unit].fromValue(())
+        make[Unit].named("x").using[Unit].annotateParameter[String]("special")
+      }.bindings)
+    }
+
+    "annotateParameter throws when annotating non-existent parameters" in {
+      val okDef = new ModuleDef {
+        make[Int]
+          .from((_: Long).toInt)
+          .annotateParameter[Long]("special")
+      }
+      assert(okDef.bindings != null)
+
+      val badDef = new ModuleDef {
+        make[Int]
+          .from((_: Long).toInt)
+          .annotateParameter[String]("special")
+      }
+      val err = intercept[ParameterNotFoundForAnnotation](badDef.bindings)
+      assert(err.getMessage contains "Long")
+      assert(err.getMessage contains "special")
+      assert(err.getMessage contains "String")
     }
 
   }

distage/distage-framework-docker/src/main/scala/izumi/distage/docker/DockerContainer.scala

@@ -94,21 +94,17 @@ object DockerContainer {
       selfTag: distage.Tag[DockerContainer[T]],
       mutateModule: ModuleDefDSL#MutationContext,
     ): Functoid[ContainerResource[F, T]] = {
-      addContainerDependency[containerDecl.Tag]
-      self
-        .addDependency[DockerContainer[containerDecl.Tag]]
-        .annotateParameter[Set[DockerContainer[Any]]](DependencyTag.get[containerDecl.Tag])
+      addContainerToDependenciesSet[containerDecl.Tag]
+      self.addDependency[DockerContainer[containerDecl.Tag]]
     }
 
     def dependOnContainer[T2](
       implicit tag: distage.Tag[DockerContainer[T2]],
       selfTag: distage.Tag[DockerContainer[T]],
       mutateModule: ModuleDefDSL#MutationContext,
     ): Functoid[ContainerResource[F, T]] = {
-      addContainerDependency[T2]
-      self
-        .addDependency[DockerContainer[T2]]
-        .annotateParameter[Set[DockerContainer[Any]]](DependencyTag.get[T2])
+      addContainerToDependenciesSet[T2]
+      self.addDependency[DockerContainer[T2]]
     }
 
     /**
@@ -183,7 +179,7 @@ object DockerContainer {
       }
     }
 
-    private def addContainerDependency[T2](
+    private def addContainerToDependenciesSet[T2](
       implicit tag: distage.Tag[DockerContainer[T2]],
       selfTag: distage.Tag[DockerContainer[T]],
       mutateModule: ModuleDefDSL#MutationContext,

distage/distage-framework/.jvm/src/main/scala/izumi/distage/framework/services/ResourceRewriter.scala

@@ -55,6 +55,7 @@ class ResourceRewriter(
 
   private def rewrite[TGT](tgt: SafeType, resourceType: SafeType)(convert: TGT => Lifecycle[Identity, TGT])(b: Binding): Seq[Binding] = {
     b match {
+      case b if b.isMutator => Seq(b) // do not rewrite mutators
       case implBinding: Binding.ImplBinding =>
         implBinding match {
           case binding: Binding.SingletonBinding[?] =>

fundamentals/fundamentals-functoid/src/main/scala/izumi/distage/model/exceptions/dsl/ParameterNotFoundForAnnotation.scala

@@ -0,0 +1,3 @@
+package izumi.distage.model.exceptions.dsl
+
+class ParameterNotFoundForAnnotation(message: String) extends RuntimeException(message)

fundamentals/fundamentals-functoid/src/main/scala/izumi/distage/model/providers/AbstractFunctoid.scala

@@ -1,8 +1,10 @@
 package izumi.distage.model.providers
 
 import izumi.distage.model.definition.Identifier
+import izumi.distage.model.exceptions.dsl.ParameterNotFoundForAnnotation
 import izumi.distage.model.reflection.{DIKey, Provider, SafeType}
 import izumi.fundamentals.platform.language.Quirks.Discarder
+import izumi.fundamentals.preamble.toRichIterable
 import izumi.reflect.Tag
 
 import scala.annotation.unchecked.uncheckedVariance
@@ -65,21 +67,47 @@ trait AbstractFunctoid[+A, Ftoid[+X] <: AbstractFunctoid[X, Ftoid]] {
     * Add an `@Id` annotation to an unannotated parameter `P`, e.g.
     * for .annotateParameter[P]("my-id"), transform lambda `(p: P) => x(p)`
     * into `(p: P @Id("my-id")) => x(p)`
+    *
+    * @throws ParameterNotFoundForAnnotation if there's no unannotated parameter `p: P` in functoid
     */
   def annotateParameter[P: Tag](name: Identifier): Ftoid[A] = {
+    val newFn = annotateParameterIfExists[P](name)
+    if (newFn.get.parameters == this.get.parameters) {
+      throw new ParameterNotFoundForAnnotation(
+        s"""Could not annotate parameter `${DIKey[P]}` with annotation `${name.idContract.repr(name.id)}`:
+           |Parameter `${DIKey[P]}` not found.
+           |Found other parameters:${this.get.parameters.map(_.key).niceList()}
+           |In Functoid: `$this`
+           |Use `annotateParameterIfExists` if the parameter missing is not an error in your circumstance""".stripMargin
+      )
+    } else {
+      newFn
+    }
+  }
+
+  /**
+    * Add an `@Id` annotation to an unannotated parameter `P`, e.g.
+    * for .annotateParameter[P]("my-id"), transform lambda `(p: P) => x(p)`
+    * into `(p: P @Id("my-id")) => x(p)`
+    *
+    * Does nothing if there's no unannotated parameter `p: P` in functoid
+    */
+  def annotateParameterIfExists[P: Tag](name: Identifier): Ftoid[A] = {
     val paramTpe = SafeType.get[P]
     annotateParameterWhen(name) {
       case DIKey.TypeKey(tpe, _) => tpe == paramTpe
       case _: DIKey.IdKey[?] => false
     }
   }
+
   /** Add an `@Id(name)` annotation to all unannotated parameters */
   def annotateAllParameters(name: Identifier): Ftoid[A] = {
     annotateParameterWhen(name) {
       case _: DIKey.TypeKey => true
       case _: DIKey.IdKey[?] => false
     }
   }
+
   /** Add an `@Id(name)` annotation to all parameters matching `predicate` */
   def annotateParameterWhen(name: Identifier)(predicate: DIKey.BasicKey => Boolean): Ftoid[A] = {
     val newProvider = this.get.replaceKeys {