Lazy For() Comprehension Generator

vavr/generator/Generator.scala

@@ -743,6 +743,151 @@ def generateMainClasses(): Unit = {
         """
       }
 
+      def genLazyFor(im: ImportManager, packageName: String, className: String): String = {
+        xs"""
+          ${
+          monadicTypesFor
+            .filterNot(_ == "Iterable")
+            .gen(mtype => (2 to N).gen(i => {
+              val forClassName = s"ForLazy$i$mtype"
+              val isComplex = monadicTypesThatNeedParameter.contains(mtype)
+              val parameterInset = if (isComplex) "L, " else ""
+              val generics = parameterInset + (1 to i).gen(j => s"T$j")(", ")
+
+              val params = (1 to i).gen { j =>
+                if (j == 1)
+                  s"$mtype<${parameterInset}T1> ts1"
+                else {
+                  val inputTypes = (1 until j).gen(k => s"? super T$k")(", ")
+                  s"Function${j - 1}<$inputTypes, $mtype<${parameterInset}T$j>> ts$j"
+                }
+              }(", ")
+
+              val ctorArgs = (1 to i).gen(j => s"ts$j")(", ")
+
+              xs"""
+                    /$javadoc
+                     * Creates a lazy {@code For}-comprehension over ${i.numerus(mtype)}.
+                     *
+                     * <p>The first argument ({@code ts1}) is the initial ${mtype}. Each subsequent
+                     * argument ({@code ts2} .. {@code ts$i}) is a function that receives all values
+                     * bound so far and returns the next ${mtype}. This method only constructs the
+                     * lazy comprehension; underlying effects are evaluated when {@code yield(...)}
+                     * is invoked.</p>
+                     *
+                     ${(0 to i).gen(j => if (j == 0) "*" else s"* @param ts$j the ${j.ordinal} ${mtype}")("\n")}
+                     ${if (isComplex) s"* @param <L> the common left-hand type of all ${mtype}s\n" else ""}
+                     ${(1 to i).gen(j => s"* @param <T$j> the component type of the ${j.ordinal} ${mtype}")("\n")}
+                     * @return a new {@code $forClassName} builder of arity $i
+                     * @throws NullPointerException if any argument is {@code null}
+                     */
+                    public static <$generics> $forClassName<$generics> For($params) {
+                        ${(1 to i).gen(j => xs"""$Objects.requireNonNull(ts$j, "ts$j is null");""")("\n")}
+                        return new $forClassName<>($ctorArgs);
+                    }
+                  """
+            })("\n\n"))("\n\n")
+        }
+
+          ${
+          monadicTypesFor
+            .filterNot(_ == "Iterable")
+            .gen(mtype => (2 to N).gen(i => {
+              val rtype = mtype
+              val forClassName = s"ForLazy$i$mtype"
+              val parameterInset = if (monadicTypesThatNeedParameter.contains(mtype)) "L, " else ""
+              val generics = parameterInset + (1 to i).gen(j => s"T$j")(", ")
+
+              val fields = (1 to i).gen { j =>
+                if (j == 1)
+                  s"private final $mtype<${parameterInset}T1> ts1;"
+                else {
+                  val inputTypes = (1 until j).gen(k => s"? super T$k")(", ")
+                  s"private final Function${j - 1}<$inputTypes, $mtype<${parameterInset}T$j>> ts$j;"
+                }
+              }("\n")
+
+              val ctorParams = (1 to i).gen { j =>
+                if (j == 1)
+                  s"$mtype<${parameterInset}T1> ts1"
+                else {
+                  val inputTypes = (1 until j).gen(k => s"? super T$k")(", ")
+                  s"Function${j - 1}<$inputTypes, $mtype<${parameterInset}T$j>> ts$j"
+                }
+              }(", ")
+
+              val assignments = (1 to i).gen(j => s"this.ts$j = ts$j;")("\n")
+
+              val yieldBody = {
+                def nestedLambda(j: Int): String = {
+                  val base = " " * 23
+                  val indent = "  " * (j + 1)
+                  if (j == i) {
+                    val argsList = (1 to i).map(k => s"t$k").mkString(", ")
+                    val inputArgs = (1 until i).map(k => s"t$k").mkString(", ")
+                    s"ts$i.apply($inputArgs).map(t$i -> f.apply($argsList))"
+                  } else if (j == 1) {
+                    s"ts1.flatMap(t1 -> {\n" +
+                      s"${base}${indent}  return ${nestedLambda(j + 1)};\n" +
+                      s"${base}${indent}})"
+                  } else {
+                    val inputArgs = (1 until j).map(k => s"t$k").mkString(", ")
+                    s"ts$j.apply($inputArgs).flatMap(t$j -> {\n" +
+                      s"${base}${indent}  return ${nestedLambda(j + 1)};\n" +
+                      s"${base}${indent}})"
+                  }
+                }
+
+                nestedLambda(1)
+              }
+
+              xs"""
+                  /$javadoc
+                   * A lazily evaluated {@code For}-comprehension with ${i.numerus(mtype)}.
+                   *
+                   * <p>Constructed via {@code For(...)} and evaluated by calling {@code yield(...)}.
+                   * Construction is side-effect free; underlying ${i.plural(mtype)} are traversed
+                   * only when {@code yield(...)} is invoked.</p>
+                   *
+                   ${if (monadicTypesThatNeedParameter.contains(mtype)) s"* @param <L> the common left-hand type of all ${mtype}s\n" else ""}
+                   ${(1 to i).gen(j => s"* @param <T$j> the component type of the ${j.ordinal} ${mtype}")("\n")}
+                   */
+                  public static class $forClassName<$generics> {
+
+                      $fields
+
+                      private $forClassName($ctorParams) {
+                          $assignments
+                      }
+
+                      /$javadoc
+                       * Produces results by mapping the Cartesian product of all bound values.
+                       *
+                       * <p>Evaluation is lazy and delegated to the underlying ${i.plural(mtype)} by
+                       * composing {@code flatMap} and {@code map} chains.</p>
+                       *
+                       * @param f a function mapping a tuple of bound values to a result
+                       * @param <R> the element type of the resulting {@code $rtype}
+                       * @return an {@code $rtype} containing mapped results
+                       * @throws NullPointerException if {@code f} is {@code null}
+                       */
+                      public <R> $rtype<${parameterInset}R> yield(${
+                if (i == 2)
+                  s"BiFunction<? super T1, ? super T2, ? extends R>"
+                else
+                  s"Function$i<${(1 to i).map(j => s"? super T$j").mkString(", ")}, ? extends R>"
+              } f) {
+                          $Objects.requireNonNull(f, "f is null");
+                          return $yieldBody;
+                      }
+                  }
+                """
+            })("\n\n"))("\n\n")
+        }
+        """
+      }
+
+
       def genFor(im: ImportManager, packageName: String, className: String): String = {
         xs"""
           //
@@ -1440,6 +1585,8 @@ def generateMainClasses(): Unit = {
 
             ${genFor(im, packageName, className)}
 
+            ${genLazyFor(im, packageName, className)}
+
             ${genMatch(im, packageName, className)}
         }
       """
@@ -2975,8 +3122,27 @@ def generateTestClasses(): Unit = {
                       ).yield(${(i > 1).gen("(")}${(1 to i).gen(j => s"i$j")(", ")}${(i > 1).gen(")")} -> ${(1 to i).gen(j => s"i$j")(" + ")});
                       $assertThat(result.get()).isEqualTo(${(1 to i).sum});
                   }
+
+                  @$test
+                  public void shouldIterateLazyFor$mtype$i() {
+                      final $mtype<${parameterInset}Integer> result = For(
+                          ${(1 to i).gen(j => if (j == 1) {
+                              s"$mtype.${builderName}($j)"
+                          } else {
+                              val args = (1 until j).map(k => s"r$k").mkString(", ")
+                              val argsUsed = (1 until j).map(k => s"r$k").mkString(" + ")
+                              s"($args) -> $mtype.${builderName}(${argsUsed} + $j)"
+                          } )(",\n")}
+                      ).yield(${(i > 1).gen("(")}${(1 to i).gen(j => s"i$j")(", ")}${(i > 1).gen(")")} -> ${(1 to i).gen(j => s"i$j")(" + ")});
+
+                      // Each step builds on the sum of all previous results plus its index
+                      // This forms a sequence rₙ = 2ⁿ - 1, and the yield sums all rᵢ.
+                      // Hence total = Σ(2ⁱ - 1) for i = 1..n = (2ⁿ⁺¹ - 2) - n
+                      assertThat(result.get()).isEqualTo((1 << ($i + 1)) - 2 - $i);
+                  }
                 """})("\n\n"))("\n\n")}
 
+
                 ${monadicFunctionTypesFor.gen(mtype => (1 to N).gen(i => { xs"""
                   @$test
                   public void shouldIterateFor$mtype$i() {