Lazy For() Comprehension

Author: pivovarit

vavr/generator/Generator.scala

@@ -743,6 +743,133 @@ 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"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 {@code For}-comprehension of ${i.numerus(mtype)}.
+                   ${(0 to i).gen(j => if (j == 0) "*" else s"* @param ts$j the ${j.ordinal} $mtype")("\n")}
+                   ${if (isComplex) s"* @param <L> left-hand type of all ${mtype}s\n" else ""}
+                   ${(1 to i).gen(j => s"* @param <T$j> component type of the ${j.ordinal} $mtype")("\n")}
+                   * @return a new {@code For}-comprehension of arity $i
+                   */
+                  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 functionType = i match {
+                case 2 => BiFunctionType
+                case _ => s"Function$i"
+              }
+              val args = (1 to i).gen(j => s"? super 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"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"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 = "        " * 3
+                  val indent = "  " * j
+                  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
+                 * For-comprehension with ${i.numerus(mtype)}.
+                 */
+                public static class $forClassName<$generics> {
+
+                    $fields
+
+                    private $forClassName($ctorParams) {
+                        $assignments
+                    }
+
+                    /$javadoc
+                     * Yields a result for elements of the cross-product of the underlying ${i.plural(mtype)}.
+                     *
+                     * @param f a function that maps an element of the cross-product to a result
+                     * @param <R> type of the resulting {@code $rtype} elements
+                     * @return an {@code $rtype} of mapped results
+                     */
+                    public <R> $rtype<${parameterInset}R> yield($functionType<$args, ? 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 +1567,8 @@ def generateMainClasses(): Unit = {
 
             ${genFor(im, packageName, className)}
 
+            ${genLazyFor(im, packageName, className)}
+
             ${genMatch(im, packageName, className)}
         }
       """