无脑 CPS 变换

CPS 变换真是一个让人又爱又恨的技术，在 PL 学术界也是引发无数讨论的话题。按我对 scheme 语言这么多年的接触，也是经历肯定，否定，再肯定，再否定...无休止的过程，不过每次理解都会更深一点，倒是好的。老实说，是好是坏，还是看它带来多少好处，以及有多少缺点。比如，反方观点：

• CPS 过后的代码，简直没法读。
• CPS 变换引入了一大堆多余的代码，后期又需要额外添加步骤，将代码给优化掉，凭空增加了复杂性。
• CPS 之后的代码永不返回，直接转成 C 会爆栈。
• CPS 引入了很多闭包，会影响性能。

正方观点：

• CPS 之后得到的 IR 更容易实现一些优化过程
• CPS 完全将控制流暴露出来，使 call/cc 功能非常容易实现
• CPS 之后调用变成变成带参数的跳转，简化了函数调用协议

先看一下如何实现 CPS 变换。网上能找到不少教程，包括王垠的 40 行代码。不过，那段代码太精妙了，以至于我没法知道它背后的推导过程（推导过程才是对理解至关重要的）。所以还是从最基础的方式开始，大致的变换规则如下：

(func cps
      lit cc => [cc lit] where (or (variable? lit) (const? lit))
      ['if a b c] cc => (let va (gensym 'r)
                          (cps a ['lambda [va]
                                   ['if va
                                        (cps b cc)
                                        (cps c cc)]]))
      ['lambda args body] cc => (let k (gensym 'k)
                                  [cc ['lambda (cons k args)
                                        (cps body k)]])
      [f x] cc => (let f0 (gensym 'f)
                    (cps f ['lambda [f0]
                             (cps x ['lambda [x]
                                      [f0 cc x]])])))

这样子变换出来的结果问题比较多，比如说我们变换这段代码：

(do (set 'square (lambda (n) (* n n)))
    (+ (square 3) 1))

会得到这样的结果：

((lambda (#arg4273)
   ((lambda (_)
      ((lambda (#f4236)
         (#f4236 (lambda (#arg4217)
                   ((lambda (x) (halt x)) (+ #arg4217 1))) 3))
       square))
    (set (quote square) #arg4273)))
 (lambda (#k4287 n)
   (#k4287 (* n n))))

这个可读性真够差的，我们可以手动简化。首先是不必要的 lambda，如果一个 lambda 是这种形式：

((lambda (x) body) v)

我们都可以改写成 let：

(let x v body)

上面的代码简化得到：

(let #arg4273 (lambda (#k4287 n)
                (#k4287 (* n n)))
     (let _ (set (quote square) #arg4273)
       (let #f4236 square
            (#f4236 (lambda (#arg4217)
                      (let x (+ #arg4217 1)
                        (halt x))) 3))))

其实改成 let 形式以后，并没有引入太多额外的 lambda。

另外我发现，let #f4236 square 这个改名跟本没啥意义。于是我们可以引入一个规则：

如果函数调用的参数本身是一个符号，并不需要对它这样改写：

      [f x] cc => (let f0 (gensym 'f)
                    (cps f ['lambda [f0]
                             (cps x ['lambda [x]
                                      [f0 cc x]])]))

可以直接用 f 就行了：

      [f x] cc => (cps x ['lambda [x0]
                            [f cc x0]])

同理，对参数也可以这样改写，省掉不必要引入的变量，也就是说如果 x 已经是最简形式，其实等价于：

      [f x] cc => [f cc x]

另外，如果一个参数只出在 lambda 的 body 中出现一次，我们可以用实参替换掉 body 内的形参，并不会导致结果膨胀：

(let x (+ #arg4217 1)
                   (halt x))

直接变成

(halt (+ #arg4217 1))

最终上面的例子变成，还算能读：

(let _ (set (quote square)
            (lambda (#k4287 n)
              (#k4287 (* n n))))
  (square (lambda (#arg4217)
            (halt (+ #arg4217 1))) 3))

这些都是我自己在手动改写简化代码过程中，总结的一些规律。后来偶然读到一篇 paper，居然跟这个做法不谋而合。这篇 paper 叫 《No-Brainer CPS Conversion》(所以博客题目也叫无脑变换)，它先没有追求所谓的 one-pass，在最原始的 CPS 变换之后，引入了一个"无脑"优化阶段，使用的规则如下：

下面这种形式，如果参数 y 是简单形式(variable,const 这类)，则可以执行 beta-替换

((lambda (x) .. x .. x .. x) y) => .. y .. y .. y

如果参数 x 在 body 里面只有一个位置引用，则可以执行 beta-替换

((lambda (x) .. x ..) triv) => .. triv ..

如果参数在 body 里面根本没使用，则可以返回 body

((lambda (x) body) triv) => body

eta 变换

(lambda (x k) (triv x k)) => triv

经过“无脑”变换之后，得到的结果还算是比较简洁的。比如 fact 的例子：

(do (set 'fact (lambda (n)
                 (if (= n 0)
                     1
                     (* n (fact (- n 1))))))
    (fact 5))

比较 naive 的变换：

((lambda (#arg4364)
   ((lambda (_)
      ((lambda (#f4327)
         (#f4327 (lambda (x)
                   (halt x)) 5))
       fact))
    (set (quote fact) #arg4364)))
 (lambda (#k4378 n)
   ((lambda (#r4387)
      (if #r4387 (#k4378 1)
          ((lambda (#f4428)
             ((lambda (#arg4434)
                (#f4428 (lambda (#arg4414)
                          (#k4378 (* n #arg4414))) #arg4434))
              (- n 1)))
           fact))) (= n 0))))
 

经过无脑优化之后：

(let _ (set (quote fact)
            (lambda (#k4874 n)
              (if (= n 0)
                  (#k4874 1)
                  (fact (lambda (#arg4919)
                          (#k4874 (* n #arg4919)))
                        (- n 1)))))
  (fact (lambda (x)
          (halt x)) 5))
 

好啦，做到这里的时候，我发现一个严重问题！CPS 变换引入的有部分 lambda 是怎么都消除不掉的！比如这里的传给 fact 的 cc，它是

(lambda (#arg4919)
        (#k4874 (* n #arg4919)))

这个会有 free variable，经过闭包变换之后，会导致有创建闭包的开销。

我发现有些情况变换出来有，有些情况变换出来不会引入额外的闭包。总结出来，非尾调用会引入额外的闭包。

原来是这样子的，由于 CPS 不带返回了，原本的栈桢就消失了。但是对于非尾调用，函数总要返回到当时的上下文里面呀，没有栈了怎么办？其实，经过 CPS 变换是把原本栈的东西，拍成 active frame 了，也就是原来可以用栈存的东西，现在都跑到 cont 里面去了，变成了 free variable！只要支持任意跳转(也就是 call/cc)，这个是不可能优化掉的。

后面发现有些 paper 优化这些东西，比如 《Callee-save registers in continuation-passing style》，在我看来，这个事情本质上解法都很丑陋。也就是经过 CPS 变换，在非尾调用的场景，原来可以放到栈里面的东西，现在都要变成 cont 里面的自由变量，导致闭包创建。这个开销本质上不可以消除。即使想优化这里也很绕。

所以我的结论是，还是不做 CPS 了。

再扯几句无关的，理论上我在寻求一个比较好的 cora 编译到 C 的方案，中间会引入一个虚拟机设计，原本是做了 CPS 变换的。但是想想这一步，不一定值得。cora 需要支持的 feature 包括：

• 闭包
• TCO
• Curry

而 call/cc 和 exception，暂时还不属于必须的 feature。如果用 CPS 实现，可以获得 TCO，call/cc，exception，但是后两者并非必须，所以我在重新考虑这个是否值得。这个取舍过程就要评估 CPS 本身的优缺点了。

这个评估过程中，首先看看能将 CPS 优化到何种程度，然后找到了一种无脑优化方式，还算可接受。再然后，又发现有些几乎完全无法优化掉的开销引入，于是就变得厌恶了。

如果不做 CPS，TCO 该怎么实现呢？对于尾调用可以用 trampoline，对于非尾调用直接用 C 语言的函数调用协议应该就行了。

============ 2018.08.28 更新 =============

如果简单的不做 CPS 变换之后，GC 变得很难处理了。

前面的方案我们将每个变成当寄存器用，在代码生成的时候就是无限的寄存器。实际上，哪些寄存器在用，哪些寄存器没有使用了，这个信息被以自由变量的形式，带到闭包里面了。

如果不做 CPS 变换，非尾调用时，这些变量(虚拟寄存器)其实存到了宿主语言的栈里面。GC 就不好 trace root 在哪里了，因为宿主环境的栈也属于需要 trace 范围了。Cheney on the M.T.A 那个可能是不得已的做法。

(+ n (f g h))

如果用之前的代码生成方式，在调用 f 的时候 n 其实是存到宿主语言的栈里面了。

补救措施可以不使用宿主的栈，而是在 VM 里面设计栈。一种实现方式是改回基于栈的虚拟机。另一种方式，要追踪虚拟寄存器和生命期，看起来比较难搞。又开始怀念 CPS 了...