Adding IterateT#foldCut for folding with early termination

jnape · jnape · commit 6abbd2a4d7ac · 2020-07-30T14:32:21.000-05:00
diff --git a/src/main/java/com/jnape/palatable/lambda/monad/transformer/builtin/IterateT.java b/src/main/java/com/jnape/palatable/lambda/monad/transformer/builtin/IterateT.java
@@ -28,7 +28,6 @@
 import static com.jnape.palatable.lambda.adt.choice.Choice2.a;
 import static com.jnape.palatable.lambda.adt.choice.Choice2.b;
 import static com.jnape.palatable.lambda.adt.hlist.HList.tuple;
-import static com.jnape.palatable.lambda.functions.builtin.fn1.Constantly.constantly;
 import static com.jnape.palatable.lambda.functions.builtin.fn2.Into.into;
 import static com.jnape.palatable.lambda.functions.builtin.fn2.Tupler2.tupler;
 import static com.jnape.palatable.lambda.functions.builtin.fn3.FoldLeft.foldLeft;
@@ -138,11 +137,30 @@ public IterateT<M, A> concat(IterateT<M, A> other) {
      */
     public <B, MB extends MonadRec<B, M>> MB fold(Fn2<? super B, ? super A, ? extends MonadRec<B, M>> fn,
                                                   MonadRec<B, M> acc) {
+        return foldCut((b, a) -> fn.apply(b, a).fmap(RecursiveResult::recurse), acc);
+    }
+
+    /**
+     * Monolithically fold the spine of this {@link IterateT} (with the possibility of early termination) by
+     * {@link MonadRec#trampolineM(Fn1) trampolining} the underlying effects (for iterative folding, use
+     * {@link IterateT#trampolineM(Fn1) trampolineM} directly).
+     *
+     * @param fn   the folding function
+     * @param acc  the starting accumulation effect
+     * @param <B>  the accumulation type
+     * @param <MB> the witnessed target result type
+     * @return the folded effect result
+     */
+    public <B, MB extends MonadRec<B, M>> MB foldCut(
+            Fn2<? super B, ? super A, ? extends MonadRec<RecursiveResult<B, B>, M>> fn,
+            MonadRec<B, M> acc) {
         return acc.fmap(tupler(this))
                 .trampolineM(into((as, b) -> maybeT(as.runIterateT())
-                        .flatMap(into((a, aas) -> maybeT(fn.apply(b, a).fmap(tupler(aas)).fmap(Maybe::just))))
+                        .flatMap(into((a, aas) -> maybeT(fn.apply(b, a).fmap(Maybe::just)).fmap(tupler(aas))))
                         .runMaybeT()
-                        .fmap(maybeRecur -> maybeRecur.match(constantly(terminate(b)), RecursiveResult::recurse))))
+                        .fmap(maybeR -> maybeR.match(
+                                __ -> terminate(b),
+                                into((rest, rr) -> rr.biMapL(tupler(rest)))))))
                 .coerce();
     }
 
diff --git a/src/test/java/com/jnape/palatable/lambda/monad/transformer/builtin/IterateTTest.java b/src/test/java/com/jnape/palatable/lambda/monad/transformer/builtin/IterateTTest.java
@@ -30,6 +30,8 @@
 import static com.jnape.palatable.lambda.adt.hlist.HList.tuple;
 import static com.jnape.palatable.lambda.functions.builtin.fn2.LTE.lte;
 import static com.jnape.palatable.lambda.functions.builtin.fn3.Times.times;
+import static com.jnape.palatable.lambda.functions.recursion.RecursiveResult.recurse;
+import static com.jnape.palatable.lambda.functions.recursion.RecursiveResult.terminate;
 import static com.jnape.palatable.lambda.functor.builtin.Identity.pureIdentity;
 import static com.jnape.palatable.lambda.functor.builtin.Lazy.lazy;
 import static com.jnape.palatable.lambda.functor.builtin.Writer.listen;
@@ -41,6 +43,7 @@
 import static com.jnape.palatable.lambda.monad.transformer.builtin.IterateT.singleton;
 import static com.jnape.palatable.lambda.monad.transformer.builtin.IterateT.unfold;
 import static com.jnape.palatable.lambda.monoid.builtin.AddAll.addAll;
+import static com.jnape.palatable.lambda.monoid.builtin.Join.join;
 import static com.jnape.palatable.traitor.framework.Subjects.subjects;
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
@@ -152,6 +155,18 @@ public void fold() {
                              .runWriter(addAll(ArrayList::new)));
     }
 
+    @Test
+    public void foldCut() {
+        assertEquals(tuple(3, "012"),
+                     IterateT.of(writer(tuple(1, "1")),
+                                 writer(tuple(2, "2")),
+                                 writer(tuple(3, "3")))
+                             .<Integer, Writer<String, Integer>>foldCut(
+                                     (x, y) -> listen(y == 2 ? terminate(x + y) : recurse(x + y)),
+                                     writer(tuple(0, "0")))
+                             .runWriter(join()));
+    }
+
     @Test
     public void zipUsesCartesianProduct() {
         assertThat(IterateT.of(new Identity<>(1), new Identity<>(2), new Identity<>(3))
