Java小技巧：利用蹦床和CPS消除递归中栈溢出

在编写递归程序时候，出现栈溢出并不是什么稀奇的事，这种时候，我们有时候可以利用尾递归优化，有时候可以改为循环写法，甚至还可以调整java虚拟机参数来改变栈大小。但是并非所有的递归都可以改写为其他形式，有些程序的递归写法甚至是最优最简写法，改成其他形式就得手动维护计算的上下文，这是非常麻烦的。这个时候就得直面栈溢出问题了。本文就介绍一个简洁但不简单的工具来消除栈溢出。

示例程序

这是一段经典斐波那契数列的示例程序：

public static int fib(int n) {
    if (n < 2) {
        return 1;
    } else {
        return fib(n - 1) + fib(n - 2);
    }
}

这个程序保持这个计算过程的情况下，没法直接直接进行尾递归优化（除非改写为迭代形式）。

第一步：利用CPS改写为尾调用形式

我们知道对于迭代计算过程的递归调用之所以尾递归优化可行，是因为在进行递归调用的时候当前函数栈已经没有用了，或者说该递归调用就是当前函数调用中的最后一步。但是这里有两个调用：fib(n - 1) + fib(n - 2)，且调用后还要做一次加法，并不能进行尾递归优化。因此，我们利用CPS将其改为尾调用（tail-call）的形式：

public static int fib(int n, Function<Integer, Integer> cont) {
    if (n < 2) {
        return cont.apply(1);
    } else {
        return fib(n - 1, a -> fib(n - 2, b -> cont.apply(a + b)));
    }
}

这样每一次递归调用都是尾调用，也就是说是当前过程中的最后一步。调用的代码就是：

int num = fib(3, i -> i);

第二步：利用蹦床消除栈的累加

上面的尾调用形式仍然会出现栈溢出，因为Java并不会进行尾调用的优化。因此我们需要手动进行这部分优化，做法就是利用蹦床技术：

public static Supplier fib(int n, Function<Integer, Supplier> cont) {
    if (n < 2) {
        return () -> cont.apply(1);
    } else {
        return () -> fib(n - 1, a -> fib(n - 2, b -> () -> cont.apply(a + b)));
    }
}

调用的地方就是：

Supplier supplier = fib(9, i -> () -> i);
Object r = supplier.get();
while(r instanceof Supplier){
    r = ((Supplier<?>) r).get();
}
System.out.println("r:" + r);

到目前为止就完成了栈溢出的消除。（关于蹦床的详细内容另外再写文章介绍）

最后：编写辅助功能工具来“简化”编写

上面的代码虽然消除了栈溢出，但是包含了泛型擦除、强制类型转换，而且每次写新的算法都要编写一个while来求值，有些不安全且不方便，因此就写了个工具来稍微优化下，利用这个工具，最后的代码为：

Function<Integer, Integer> fibFun = Trampoline.defun(thisFun -> (n, cont) -> {
    if (n < 2) {
        return cont.returnBack(1);
    } else {
        return thisFun.call(n - 1, a -> thisFun.call(n - 2, b -> cont.returnBack(a + b)));
    }
});

调用的代码就是简化为了：

Integer num2 = fibFun.apply(index);

Trampoline类的实现如下：

import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;

public class Trampoline {
    public interface ProxyFun<T, R> {
        ValurOrSupplier<R> call(T t, Continuation<R> cont);
    }

    public interface Continuation<R> {
        ValurOrSupplier<R> returnBack(R r);
    }

    public static class ValurOrSupplier<T> {
        final T result;
        final Supplier<ValurOrSupplier<T>> supplier;

        final boolean isEnd;

        ValurOrSupplier(T result, Supplier<ValurOrSupplier<T>> supplier, boolean isEnd) {
            this.result = result;
            this.supplier = supplier;
            this.isEnd = isEnd;
        }

        static <T> ValurOrSupplier<T> end(T t) {
            return new ValurOrSupplier<>(t, null, true);
        }

        static <T> ValurOrSupplier<T> cont(Supplier<ValurOrSupplier<T>> supplier) {
            return new ValurOrSupplier<>(null, supplier, false);
        }
    }

    public static <T, R> Function<T, R> defun(Function<ProxyFun<T, R>, ProxyFun<T, R>> fun) {
        AtomicReference<ProxyFun<T, R>> proxyFunRef = new AtomicReference<>();
        proxyFunRef.set((t, cont) -> ValurOrSupplier.cont(() -> fun.apply(proxyFunRef.get()).call(t, r -> ValurOrSupplier.cont(() -> cont.returnBack(r)))));

        return t -> {
            ValurOrSupplier<R> result = fun.apply(proxyFunRef.get()).call(t, ValurOrSupplier::end);
            while (!result.isEnd) {
                result = result.supplier.get();
            }
            return result.result;
        };
    }
}

注意：这个只是单参数的形式，多个参数也可以按照这个思路来处理。