AndThen data type

modules/core/arrow-extras-extensions/src/main/kotlin/arrow/data/extensions/andthen.kt

@@ -0,0 +1,89 @@
+package arrow.data.extensions
+
+import arrow.Kind
+import arrow.Kind2
+import arrow.core.Either
+import arrow.data.*
+import arrow.extension
+import arrow.typeclasses.*
+
+@extension
+interface AndThenSemigroup<A, B> : Semigroup<AndThen<A, B>> {
+  fun SB(): Semigroup<B>
+
+  override fun AndThen<A, B>.combine(b: AndThen<A, B>): AndThen<A, B> = SB().run {
+    AndThen { a: A -> invoke(a).combine(b.invoke(a)) }
+  }
+
+}
+
+@extension
+interface AndThenMonoid<A, B> : Monoid<AndThen<A, B>>, AndThenSemigroup<A, B> {
+
+  fun MB(): Monoid<B>
+
+  override fun SB(): Semigroup<B> = MB()
+
+  override fun empty(): AndThen<A, B> =
+    AndThen { MB().empty() }
+
+}
+
+@extension
+interface AndThenFunctor<X> : Functor<AndThenPartialOf<X>> {
+  override fun <A, B> AndThenOf<X, A>.map(f: (A) -> B): AndThen<X, B> =
+    fix().map(f)
+}
+
+@extension
+interface AndThenApplicative<X> : Applicative<AndThenPartialOf<X>>, AndThenFunctor<X> {
+  override fun <A> just(a: A): AndThenOf<X, A> =
+    AndThen.just(a)
+
+  override fun <A, B> AndThenOf<X, A>.ap(ff: AndThenOf<X, (A) -> B>): AndThen<X, B> =
+    fix().ap(ff)
+
+  override fun <A, B> AndThenOf<X, A>.map(f: (A) -> B): AndThen<X, B> =
+    fix().map(f)
+
+}
+
+@extension
+interface AndThenMonad<X> : Monad<AndThenPartialOf<X>>, AndThenApplicative<X> {
+  override fun <A, B> AndThenOf<X, A>.flatMap(f: (A) -> AndThenOf<X, B>): AndThen<X, B> =
+    fix().flatMap(f)
+
+  override fun <A, B> tailRecM(a: A, f: (A) -> AndThenOf<X, Either<A, B>>): AndThen<X, B> =
+    AndThen.tailRecM(a, f)
+
+  override fun <A, B> AndThenOf<X, A>.map(f: (A) -> B): AndThen<X, B> =
+    fix().map(f)
+
+  override fun <A, B> AndThenOf<X, A>.ap(ff: AndThenOf<X, (A) -> B>): AndThen<X, B> =
+    fix().ap(ff)
+
+}
+
+@extension
+interface AndThenCategory : Category<ForAndThen> {
+  override fun <A> id(): AndThen<A, A> =
+    AndThen.id()
+
+  override fun <A, B, C> AndThenOf<B, C>.compose(arr: Kind2<ForAndThen, A, B>): AndThen<A, C> =
+    fix().compose(arr::invoke)
+
+}
+
+@extension
+interface AndThenContravariant<O> : Contravariant<Conested<ForAndThen, O>> {
+
+  override fun <A, B> Kind<Conested<ForAndThen, O>, A>.contramap(f: (B) -> A): Kind<Conested<ForAndThen, O>, B> =
+    counnest().fix().contramap(f).conest()
+
+}
+
+@extension
+interface AndThenProfunctor : Profunctor<ForAndThen> {
+  override fun <A, B, C, D> AndThenOf<A, B>.dimap(fl: (C) -> A, fr: (B) -> D): AndThen<C, D> =
+    fix().andThen(fr).compose(fl)
+}

modules/core/arrow-extras/src/main/kotlin/arrow/data/AndThen.kt

@@ -0,0 +1,243 @@
+package arrow.data
+
+import arrow.core.*
+import arrow.higherkind
+
+operator fun <A, B> AndThenOf<A, B>.invoke(a: A): B = fix().invoke(a)
+
+/**
+ * [AndThen] wraps a function of shape `(A) -> B` and can be used to do function composition.
+ * It's similar to [arrow.core.andThen] and [arrow.core.compose] and can be used to build stack safe
+ * data structures that make use of lambdas. Usage is typically used for signature such as `A -> Kind<F, A>` where
+ * `F` has a [arrow.typeclasses.Monad] instance i.e. [StateT.flatMap].
+ *
+ * As you can see the usage of [AndThen] is the same as `[arrow.core.andThen] except we start our computation by
+ * wrapping our function in [AndThen].
+ *
+ * ```kotlin:ank:playground
+ * import arrow.core.andThen
+ * import arrow.data.AndThen
+ * import arrow.data.extensions.list.foldable.foldLeft
+ *
+ * fun main(args: Array<String>) {
+ *   //sampleStart
+ *   val f = (0..10000).toList()
+ *     .fold({ x: Int -> x + 1 }) { acc, _ ->
+ *       acc.andThen { it + 1 }
+ *     }
+ *
+ *   val f2 = (0..10000).toList()
+ *     .foldLeft(AndThen { x: Int -> x + 1 }) { acc, _ ->
+ *       acc.andThen { it + 1 }
+ *     }
+ *   //sampleEnd
+ *   println("f(0) = ${f(0)}, f2(0) = ${f2(0)}")
+ * }
+ * ```
+ *
+ */
+@higherkind
+sealed class AndThen<A, B> : (A) -> B, AndThenOf<A, B> {
+
+  private data class Single<A, B>(val f: (A) -> B, val index: Int) : AndThen<A, B>()
+
+  private data class Concat<A, E, B>(val left: AndThen<A, E>, val right: AndThen<E, B>) : AndThen<A, B>() {
+    override fun toString(): String = "AndThen.Concat(...)"
+  }
+
+  /**
+   * ```kotlin:ank:playground
+   * import arrow.data.AndThen
+   * import arrow.data.extensions.list.foldable.foldLeft
+   *
+   * fun main(args: Array<String>) {
+   *   //sampleStart
+   *   val f = (0..10000).toList()
+   *     .foldLeft(AndThen { x: Int -> x + 1 }) { acc, _ ->
+   *       acc.andThen { it + 1 }
+   *     }
+   *
+   *   val result = f(0)
+   *   //sampleEnd
+   *   println("result = $result")
+   * }
+   * ```
+   *
+   * @param g function to invoke after this [AndThen].
+   * @return a function wrapped in [AndThen] from [A] to [X].
+   */
+  fun <X> andThen(g: (B) -> X): AndThen<A, X> =
+    when (this) {
+      // Fusing calls up to a certain threshold, using the fusion technique implemented for `IO#map`
+      is Single -> if (index != maxStackDepthSize) Single(f andThen g, index + 1)
+      else andThenF(AndThen(g))
+      else -> andThenF(AndThen(g))
+    }
+
+  /**
+   * ```kotlin:ank:playground
+   * import arrow.data.AndThen
+   * import arrow.data.extensions.list.foldable.foldLeft
+   *
+   * fun main(args: Array<String>) {
+   *   //sampleStart
+   *   val f = (0..10000).toList().foldLeft(AndThen { i: Int -> i + 1 }) { acc, _ ->
+   *     acc.compose { it + 1 }
+   *   }
+   *
+   *   val result = f(0)
+   *   //sampleEnd
+   *   println("result = $result")
+   * }
+   * ```
+   *
+   * @param g function to invoke before this [AndThen].
+   * @return a function wrapped in [AndThen] from [C] to [B].
+   */
+  infix fun <C> compose(g: (C) -> A): AndThen<C, B> =
+    when (this) {
+      // Fusing calls up to a certain threshold, using the fusion technique implemented for `IO#map`
+      is Single -> if (index != maxStackDepthSize) Single(f compose g, index + 1)
+      else composeF(AndThen(g))
+      else -> composeF(AndThen(g))
+    }
+
+  /**
+   * Alias for [andThen]
+   *
+   * @see andThen
+   */
+  fun <C> map(f: (B) -> C): AndThen<A, C> =
+    andThen(f)
+
+  /**
+   * Alias for [andThen]
+   *
+   * @see compose
+   */
+  fun <C> contramap(f: (C) -> A): AndThen<C, B> =
+    this compose f
+
+  fun <C> flatMap(f: (B) -> AndThenOf<A, C>): AndThen<A, C> =
+    AndThen { a: A -> f(this.invoke(a)).fix().invoke(a) }
+
+  fun <C> ap(ff: AndThenOf<A, (B) -> C>): AndThen<A, C> =
+    ff.fix().flatMap { f ->
+      map(f)
+    }
+
+  /**
+   * Invoke the `[AndThen]` function
+   *
+   * ```kotlin:ank:playground
+   * import arrow.data.AndThen
+   *
+   * fun main(args: Array<String>) {
+   *   //sampleStart
+   *   val f: AndThen<Int, String> = AndThen(Int::toString)
+   *   val result = f.invoke(0)
+   *   //sampleEnd
+   *   println("result = $result")
+   * }
+   * ```
+   *
+   * @param a value to invoke function with
+   * @return result of type [B].
+   **/
+  @Suppress("UNCHECKED_CAST")
+  override fun invoke(a: A): B = loop(this as AndThen<Any?, Any?>, a)
+
+  override fun toString(): String = "AndThen(...)"
+
+  companion object {
+
+    fun <A, B> just(b: B): AndThen<A, B> =
+      AndThen { b }
+
+    fun <A> id(): AndThen<A, A> =
+      AndThen(::identity)
+
+    /**
+     * Wraps a function in [AndThen].
+     *
+     * ```kotlin:ank:playground
+     * import arrow.data.AndThen
+     *
+     * fun main(args: Array<String>) {
+     *   //sampleStart
+     *   val f = AndThen { x: Int -> x + 1 }
+     *   val result = f(0)
+     *   //sampleEnd
+     *   println("result = $result")
+     * }
+     * ```
+     *
+     * @param f the function to wrap
+     * @return wrapped function [f].
+     */
+    operator fun <A, B> invoke(f: (A) -> B): AndThen<A, B> = when (f) {
+      is AndThen<A, B> -> f
+      else -> Single(f, 0)
+    }
+
+    fun <I, A, B> tailRecM(a: A, f: (A) -> AndThenOf<I, Either<A, B>>): AndThen<I, B> =
+      AndThen { t: I -> step(a, t, f) }
+
+    private tailrec fun <I, A, B> step(a: A, t: I, fn: (A) -> AndThenOf<I, Either<A, B>>): B {
+      val af = fn(a)(t)
+      return when (af) {
+        is Either.Right -> af.b
+        is Either.Left -> step(af.a, t, fn)
+      }
+    }
+
+    /**
+     * Establishes the maximum stack depth when fusing `andThen` or `compose` calls.
+     *
+     * The default is `128`, from which we substract one as an
+     * optimization. This default has been reached like this:
+     *
+     *  - according to official docs, the default stack size on 32-bits
+     *    Windows and Linux was 320 KB, whereas for 64-bits it is 1024 KB
+     *  - according to measurements chaining `Function1` references uses
+     *    approximately 32 bytes of stack space on a 64 bits system;
+     *    this could be lower if "compressed oops" is activated
+     *  - therefore a "map fusion" that goes 128 in stack depth can use
+     *    about 4 KB of stack space
+     */
+    private const val maxStackDepthSize = 127
+  }
+
+  private fun <X> andThenF(right: AndThen<B, X>): AndThen<A, X> = Concat(this, right)
+  private fun <X> composeF(right: AndThen<X, A>): AndThen<X, B> = Concat(right, this)
+
+  @Suppress("UNCHECKED_CAST")
+  private tailrec fun loop(self: AndThen<Any?, Any?>, current: Any?): B = when (self) {
+    is Single -> self.f(current) as B
+    is Concat<*, *, *> -> {
+      when (val oldLeft = self.left) {
+        is Single<*, *> -> {
+          val left = oldLeft as Single<Any?, Any?>
+          val newSelf = self.right as AndThen<Any?, Any?>
+          loop(newSelf, left.f(current))
+        }
+        is Concat<*, *, *> -> loop(
+          rotateAccumulate(self.left as AndThen<Any?, Any?>, self.right as AndThen<Any?, Any?>),
+          current
+        )
+      }
+    }
+  }
+
+  @Suppress("UNCHECKED_CAST")
+  private tailrec fun rotateAccumulate(
+    left: AndThen<Any?, Any?>,
+    right: AndThen<Any?, Any?>): AndThen<Any?, Any?> = when (left) {
+    is Concat<*, *, *> -> rotateAccumulate(
+      left.left as AndThen<Any?, Any?>,
+      (left.right as AndThen<Any?, Any?>).andThenF(right)
+    )
+    is Single<*, *> -> left.andThenF(right)
+  }
+
+}

modules/core/arrow-extras/src/main/kotlin/arrow/data/StateT.kt

@@ -99,7 +99,7 @@ class StateT<F, S, A>(
      * @param f the modify function to apply.
      */
     fun <F, S> modify(AF: Applicative<F>, f: (S) -> S): StateT<F, S, Unit> = AF.run {
-      StateT<F, S, Unit>(just({ s ->
+      StateT(just({ s ->
         just(f(s)).map { Tuple2(it, Unit) }
       }))
     }
@@ -169,7 +169,7 @@ class StateT<F, S, A>(
   fun <B, Z> map2(MF: Monad<F>, sb: StateTOf<F, S, B>, fn: (A, B) -> Z): StateT<F, S, Z> =
     MF.run {
       invokeF(runF.map2(sb.fix().runF) { (ssa, ssb) ->
-        ssa.andThen { fsa ->
+        AndThen(ssa).andThen { fsa ->
           fsa.flatMap { (s, a) ->
             ssb(s).map { (s, b) -> Tuple2(s, fn(a, b)) }
           }
@@ -186,7 +186,7 @@ class StateT<F, S, A>(
    */
   fun <B, Z> map2Eval(MF: Monad<F>, sb: EvalOf<StateT<F, S, B>>, fn: (A, B) -> Z): Eval<StateT<F, S, Z>> = MF.run {
     runF.map2Eval(sb.fix().map { it.runF }) { (ssa, ssb) ->
-      ssa.andThen { fsa ->
+      AndThen(ssa).andThen { fsa ->
         fsa.flatMap { (s, a) ->
           ssb((s)).map { (s, b) -> Tuple2(s, fn(a, b)) }
         }
@@ -221,7 +221,7 @@ class StateT<F, S, A>(
   fun <B> flatMap(MF: Monad<F>, fas: (A) -> StateTOf<F, S, B>): StateT<F, S, B> = MF.run {
     invokeF(
       runF.map { sfsa ->
-        sfsa.andThen { fsa ->
+        AndThen(sfsa).andThen { fsa ->
           fsa.flatMap {
             fas(it.b).runM(MF, it.a)
           }
@@ -238,7 +238,7 @@ class StateT<F, S, A>(
   fun <B> flatMapF(MF: Monad<F>, faf: (A) -> Kind<F, B>): StateT<F, S, B> = MF.run {
     invokeF(
       runF.map { sfsa ->
-        sfsa.andThen { fsa ->
+        AndThen(sfsa).andThen { fsa ->
           fsa.flatMap { (s, a) ->
             faf(a).map { b -> Tuple2(s, b) }
           }