StateT#zip preserves all state transformations

jnape · jnape · commit b993936ea6ae · 2019-10-13T15:31:01.000-05:00
diff --git a/src/main/java/com/jnape/palatable/lambda/monad/transformer/builtin/StateT.java b/src/main/java/com/jnape/palatable/lambda/monad/transformer/builtin/StateT.java
@@ -152,11 +152,7 @@ public <B> StateT<S, M, B> fmap(Fn1<? super A, ? extends B> fn) {
      */
     @Override
     public <B> StateT<S, M, B> zip(Applicative<Fn1<? super A, ? extends B>, StateT<S, M, ?>> appFn) {
-        return stateT(s -> runStateT(s)
-                .zip(appFn.<StateT<S, M, Fn1<? super A, ? extends B>>>coerce()
-                             .runStateT(s)
-                             .fmap(Tuple2::_1)
-                             .fmap(f -> t -> t.biMapL(f))));
+        return MonadT.super.zip(appFn).coerce();
     }
 
     /**
diff --git a/src/test/java/com/jnape/palatable/lambda/functor/builtin/StateTest.java b/src/test/java/com/jnape/palatable/lambda/functor/builtin/StateTest.java
@@ -12,8 +12,8 @@
 import testsupport.traits.FunctorLaws;
 import testsupport.traits.MonadLaws;
 import testsupport.traits.MonadReaderLaws;
-import testsupport.traits.MonadWriterLaws;
 import testsupport.traits.MonadRecLaws;
+import testsupport.traits.MonadWriterLaws;
 
 import static com.jnape.palatable.lambda.adt.Unit.UNIT;
 import static com.jnape.palatable.lambda.adt.hlist.HList.tuple;
diff --git a/src/test/java/com/jnape/palatable/lambda/monad/transformer/builtin/StateTTest.java b/src/test/java/com/jnape/palatable/lambda/monad/transformer/builtin/StateTTest.java
@@ -1,6 +1,7 @@
 package com.jnape.palatable.lambda.monad.transformer.builtin;
 
 import com.jnape.palatable.lambda.adt.Maybe;
+import com.jnape.palatable.lambda.adt.Unit;
 import com.jnape.palatable.lambda.adt.hlist.Tuple2;
 import com.jnape.palatable.lambda.functor.builtin.Identity;
 import com.jnape.palatable.traitor.annotations.TestTraits;
@@ -15,10 +16,16 @@
 import testsupport.traits.MonadRecLaws;
 import testsupport.traits.MonadWriterLaws;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import static com.jnape.palatable.lambda.adt.Maybe.just;
 import static com.jnape.palatable.lambda.adt.Unit.UNIT;
 import static com.jnape.palatable.lambda.adt.hlist.HList.tuple;
 import static com.jnape.palatable.lambda.functor.builtin.Identity.pureIdentity;
+import static com.jnape.palatable.lambda.optics.functions.Set.set;
+import static com.jnape.palatable.lambda.optics.lenses.ListLens.elementAt;
+import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
 import static testsupport.traits.Equivalence.equivalence;
 
@@ -57,11 +64,14 @@ public void mapStateT() {
 
     @Test
     public void zipping() {
-        assertEquals(new Identity<>(tuple(4, "final state: FOO")),
-                     StateT.<String, Identity<?>>modify(s -> new Identity<>(s.toUpperCase()))
-                             .discardL(StateT.gets(s -> new Identity<>(s.length())))
-                             .flatMap(x -> StateT.stateT(s -> new Identity<>(tuple(x + 1, "final state: " + s))))
-                             .<Identity<Tuple2<Integer, String>>>runStateT("foo"));
+        Tuple2<Unit, List<String>> result = StateT.<List<String>, Identity<?>>modify(
+                s -> new Identity<>(set(elementAt(s.size()), just("one"), s)))
+                .discardL(StateT.modify(s -> new Identity<>(set(elementAt(s.size()), just("two"), s))))
+                .<Identity<Tuple2<Unit, List<String>>>>runStateT(new ArrayList<>())
+                .runIdentity();
+
+        assertEquals(tuple(UNIT, asList("one", "two")),
+                     result);
     }
 
     @Test
