Translate: Stackless Scala With Free Monads

Scala 编译器的尾调用优化 (Tail Call Elimination, TCE) 仅限于自递归方法 (Self-Recursive Methods) 而不是所有的尾调用. 这很容易使得那些由小函数组合而成的函数造成堆栈溢出. 拥有一个通用的尾调用优化方法非常有益于 Scala 语言, 尤其是有益于 Scala 的函数式编程. 针对那些本不支持尾调用优化的语言, Trampolining 是一个可以用于解决尾调用优化问题的常用方法. 本文引入 Trampolining 技术, 以将其推广为一个对任何方法调用(甚至是非尾部调用)的优化技术. 该技术能够完全的消除 Scala 程序对调用堆栈的使用.

调用堆栈是虚拟机的有限资源, 稍微有点经验的 Scala 程序员应该都遇见过一些看上去很正常的函数用尽了调用堆栈后使得程序崩溃并抛出 StackOverflowError 的错误. 考虑一个实用的例子, 在遍历一个 List 的同时维护一些状态. 我们采用 State 数据结构来表征一个简单的状态机的变换.

case class State[S, +A](runS: S => (A, S)) {
  def map[B](f: A => B) = State[S, B] { s => {
    val (a, s1) = runS(s)
    (f(a), s1)
  }}

  def flatMap[B](f: A => State[S, B]) = State[S, B] { s => {
    val (a, s1) = runS(s)
    f(a).runS(s1)
  }}
}

函数 runS 以一个类型为 S 的状态作为输入, 返回一个类型为 A 的值以及一个新的状态. 而 map 和 flatMap 方法允许通过 for-comprehension 以一种命令式的代码方式将状态变换组合在一起, 如:

def getState[S]: State[S, S] = State(s => (s, s))

def setState[S](s: S): State[S, Unit] = State(_ => ((), s))

def pureState[S, A](a: A): State[S, A] = State(s => (a, s))

注意到 pureState 和 flatMap 两个方法使得 State 成为一个 Monad.

作为一个简单的例子, 我们用 State 来实现对一个 List 中的元素编号. 这并不见得是 State 的一个常见用例, 仅仅因为这能够简单的诠释上文所说的堆栈溢出问题.

def zipIndex[A](as: List[A]): List[(Int, A)] = as.foldLeft(pureState[Int, List[(Int, A)]](Nil)) { (acc, a) =>
  for {
    xs <- acc
    n <- getState[Int]
    _ <- setState[Int](n + 1)
  } yield (n, a) :: xs
}.runS(0)._1.reverse

译者注: 对于 State 的 flatMap 方法, 每次调用都会在堆中产生一个新的 State.class 类型实例及其 runS 方法的函数类型 State.1.class 的实例. (A Look at How Scala Compiles to Java)

我们使用一个 foldLeft 来强调对 list 的遍历是一个尾递归的调用. 这个函数以一个空的 list 为起点, 不断在列表的头部添加元素, 所以与原 list 顺序相反. 而 state 是一个从零开始每步递增的整数, 所有的 state 的生成都发生在反转结果之前.

但是如果用一个元素数量超过虚拟机调用堆栈大小的 list 作为参数, State.flatMap 方法会抛出 StackOverflowError 的错误. 发生错误的原因是最终的 state 本质上是一个大函数, 由与 list 大小成正比的多个函数组合而成. 尽管它让我们觉得这是一个离散的一系列步骤, 但每个步骤使用了一种无法另编译器优化的方式去调用下一个步骤.

这种情况极大的限制了 Scala 的函数式编程特性. 而这篇论文就是来探讨解决方案的:

• 第三部分将探讨熟知的 Trampolining 技术. 在一个 Trampolined 的程序中, 我们不再让每个步骤直接去调用另一个步骤, 而采用一种 Trampoline 控制循环来暂缓下一步骤的发生. 这使得我们可以用堆去代替调用栈.
• 之后(在第四部分), 我们会将这种方法扩展为一个 Monad, 使得所有的调用都可以被转化为可依次消除的尾递归调用.
• 在实现这种 monad 时, 如果实现不当, 依然会导致在一些场景堆栈溢出的错误. 在 4.3 节, 我们将探讨这些场景以及如何消除这些问题.
• Trampolined 程序可以是多层嵌套的来提供 cooperative coroutines, 我们将在第五部分探讨.
• 在第六部分, 我们将 Trampolines 推广为 Free Monad, 一种极为通用的递归数据结构. 通过分析操作这种结构的部分函数, 利用 Trampoline 的方法来解决 Free Monad 的堆栈问题.

Scala 的编译器能够对一种叫做自调用尾递归 (Self-recursive Call) 的代码进行优化. 例如, 下面实例代码中对 list 的 左折 操作可以被编译器优化为一个使用固定栈空间的函数:

def foldl[A, B](as: List[A], b: B, f: (B, A) => B): B =
  as match {
    case Nil => b
    case x :: xs => foldl(xs, f(b, x), f)
  }

当编译器发现一个函数在尾部调用了自己, 并且 该函数不能被重载 (例如, 使用 private 或 final 关键字声明), 则该函数的自调用会被编译替换为一个简单的跳转 (Jump). 这就相当于将该递归写成一个这循环:

def foldl[A, B](as: List[A], b: B, f: (B, A) => B): B = {
  var z = b
  var az = as
  while (true) {
    az match {
      case Nil => return z
      case x :: xs =>
        z = f(b, x)
        az = xs
    }
  }
}

除了跳转比方法调用更高效这个优势外, 这种优化使得无需使用任何调用栈. 但是,用跳转替换一个普通的尾调用就没有自递归调用优化那么简单了. 目前, JVM 只允许本地跳转, 所以根本没法从一个方法尾部直接跳转至另一个方法.例如, 一个相互跳转的递归是不能被编译器优化的, 即便它的调用在方法的尾部.

def even[A](ns: List[A]): Boolean = ns match {
  case Nil => true
  case x :: xs => odd(xs)
}

def odd[A](ns: List[A]): Boolean = ns match {
  case Nil => false
  case x :: xs = even(xs)
}

这两个方法会在参数长度大于堆栈长度时引发堆栈溢出.

尽管 JVM 很可能会在未来改进这点, 但也并不是一件容易的事, 并且有可能没有想象中那么有用. 比如说, JVM 的调用模型需要把每个线程的执行状态都保存在线程的调用堆栈里. 另外, 异常处理也是通过将异常压栈后将堆栈暴露给程序进行检查的. 事实上, JVM 的安全模型也是基于检查每个调用栈帧 (Stack Frame) 的权限来实现的. 而这又与子类, 动态分配策略, 即时编译器偶合在一起, 使得 Scala 编译器的尾调用优化难以实现.

不过还好, 我们可以回避所有这些问题, 通过一种简单的数据结构来用堆空间来代替栈空间.

我们从一个非常简单的 Trampoline 数据类型来说明. 下例与 scala.util.control.TailCalls.TailRec 包中的 Trampoline 从原理上是完全相同的, 尽管实现的方式不同.

sealed trait Trampoline[+A] {
  final def runT: A = this match {
    case More(k) => k().runT
    case Done(v) => v
  }
}

case class More[+A](k: () => Trampoline[A])
    extends Trampoline[A]

case class Done[+A](v: A)
    extends Trampoline[A]

Trampoline 将计算分解为多个步骤, 每个步骤可以有两种形式. Done 包含一个用于返回的值 v, 表示后续再无其他步骤. More 则包含一个生成后续步骤的函数 k. 而 runT 方法则是一个简单的尾递归, 用于执行所有的步骤. 它被定义为 final 以便让 Scala 编译器进行优化. 这就可以解决上文提到的互相跳转的尾调用递归问题. 我们只需要机械的将所有类型为 T 的返回值替换为 Trampoline[T] 就行了. 具体如下:

def even[A](ns: List[A]): Trampoline[Boolean] = ns match {
  case Nil => Done(true)
  case x :: xs => More(() => odd(xs))
}

def odd[A](ns: List[A]): Trampoline[Boolean] = ns match {
  case Nil => Done(false)
  case x :: xs = More(() => even(xs))
}

我们从原来的直接递归调用, 变成使其返回一个表示下一个步骤的 Trampoline, 然后可以通过调用可以被优化的方法 runT 来递归执行. 这样, 无论输入的 List 有多长, 上面的代码也不会再引发堆栈溢出.

让我们看看如何使用 Trampoline 技术来解决论文一开始提出的 State 遍历 list 的 堆栈溢出问题. 首先, 我们需要改写 State 的行为来返回可供尾递归的 Trampoline 类型.

case class State[S, +A](runS: S => Trampoline[(A, S)])

我们现在该如何实现组合 State 的 flatMap 呢? 先试试这样:

def flatMap[B](f: A => State[S, B]) = State[S, B] { s =>
  More[(S, B)] { () =>
    val (a, s1) = runS(s).runT
    More(() => f(a).runS(s1))
  }
}

但这样写会很低效. 而且例一中的 zipIndex 仍然可能会堆栈溢出, 甚至对于更小的 list. 问题出在 runT 的调用没有出现在尾部, 所以它并不能被封装在 Trampoline 内部, 也不能用 Trampoline 来优化.

通过使用 Trampoline 化的 State monad, 对任何大小的输入 list, 现在运行之前看到的 zipIndex 方法不会在引发堆栈溢出的异常了. 这里介绍的 Trampoline 是 Scala 中清除栈帧的一种通用手段. 我们可以用它来重写任何程序而不再需要栈空间了. 考虑一个具有如下形式的程序:

val x = f()
val y = g(x)
h(y)

可以非常轻松的这样重写:

for {
  x <- f()
  y <- g(x)
  z <- h(y)
} yield z

当然, 需要定义如下的隐式转换:

implicit def step[A](a: => A): Trampoline[A] =
  More(() => Done(a))

唯一不太适合用 step 转换的是(非偶然地)自递归调用. 这非常容易察觉, 当这种情况, 我们可以显示的调用 More 构造器, 例如下面的计算第 n 个斐波那契数的递归函数:

def fib(n: Int): Trampoline[Int] =
  if(n <= 1) Done(n) else for {
    x <- More(() => fib(n-1))
    y <- More(() => fib(n-2))
  } yield x + y

由于这种转换是完全机械化的, 可能有人会写一个编译器插件或者扩展 Scala 的编译期将所有程序转换成这样.

我们已经看到如何使用 flatMap 序列化的组合 Trampoline. 但其实可以通过交叉计算来并行的组合它们:

def zip[B](b: Trampoline[B]): Trampoline[(A, B)] =
  (this.resume, b.resume) match {
    case (Right(a), Right(B)) => Done((a, b))
    case (Left(a), Left(b)) =>
      More(() => a() zip b())
    case (Left(a), Right(b)) =>
      More(() => a() zip Done(b))
    case (Right(a), Left(b)) =>
      More(() => Done(a) zip b())
  }

为了验证, 我们引入控制台打印:

val hello: Trampoline[Unit] = for {
  _ <- print("Hello, ")
  _ <- println("World!")
} yield ()

将其与自身交错执行, 然后观察结果:

(hello zip hello).runT
// Hello, Hello, World!
// World!

尽管这是单线程的并行计算, 但很容易扩展到多线程的分布式计算. 对于一组待执行的 trampoline, 任何不是 Done 的 resume 任务都可以在任何线程执行.

可以将其推广到一个完全对称的协同程序. 我们会在下一节讨论.

我们可以将 Trampoline 看作事一个能通过 Function0 挂起并之后执行的协同程序. 那么它就不再是唯一一种有此性质的类型构造器了, 如果将该行为抽象, 我们就能获得如下的类型:

sealed trait Free[S[+_], +A] {
  private case class FlatMap[S[+_], A, +B](
    a: Free[S, A],
    f: A => Free[S, B]) extends Free[S, B]
}

case class Done[S[+_], +A](a: A) extends Free[S, A]

case class More[S[+_], +A](k: S[Free[S, A]]) extends Free[S, A]

译者注: S 是 Function0 的推广. S 具有能够挂起计算的能力.

那么, Trampoline 可以如下定义:

type Trampoline[+A] = Free[Function0, A]

就像 Done 和 FlatMap 数据构造器一样, Free[S, A] 是任意协变的 Functor S 的一种 Monad. 从范畴论角度来看, 它正是该 Functor 生成的 Free Monad [4]. S 必须是一个 Functor 即必须存在 Functor[S]³:

trait Functor[F[_]] {
  def map[A, B](m: F[A])(f: A => B): F[B]
}

对 Function0 来说是显而易见的:

implicit val f0Functor = new Functor[Function0] {
  def map[A, B](a: () => A)(f: A => B) = () => f(a())
}

后续部分是作者对当前工作的总结和展望.