解密自動CPS變換

博客園已經停止更新，后續更新都在我的 Github Pages 博客中。

1 前言

我最一開始聽到 CPS 變換這個詞是在王垠的博客里 (請求不要噴我)，就是那篇他第一次宣傳他的40行代碼的文章。我當時什么都看不懂，所以沒太注意，不過我也正在學程序語言方面的東西，不久我就在 EOPL 和 The Little Schemer 里面又見到了 CPS。我有點不服氣，知道了 CPS 不過就是這么個東西，於是我也開始想自己重造王垠40行代碼，然后我很驚訝地也花了剛好一個星期寫了出來，而且還基本上跟他的一模一樣......

(注意，我不是引戰)

畢竟可能每個人有自己的思考方式，我在這里只是分享一下我自己的思路，我寫出這個 CPS 程序的經歷。當然為了顯得我稍微強一點的樣子，我把中間許多非常蠢的錯誤都省略了。其實我也不知道我怎么就把這段代碼寫出來了。

我寫出這個程序以后，又去看了那篇經典論文 Representing Control，這里有一個更友好一點的版本 How to compile with continuations，發現我的思路和他的完全不同，我反倒覺得我的思路很清楚，他的思路我要繞個彎才能看懂，雖然實質上是一樣的。

這篇文章不涉及什么是 CPS 變換，CPS 變換入門請參考 The Little Schemer，我直接就開始寫 cps 函數了，我們的目標就是王垠的那40行代碼，我也把我自己的程序的變量名什么的都改成了跟那段代碼一樣的格式，方便對照，當然也有一些地方不一樣，懶得改了。

我采用 Racket 語言，就是因為用起來方便一些。就這樣了。

2 簡單的CPS變換

其實一個正確的CPS變換程序，只要學過一點點寫解釋器的人都會寫，所以我就不細講了，只是提供一個回憶，你如果要看下去的話最好是這個會自己寫。我們先只考慮λ-calculus的語法，就只有3個分支，變量，λ，函數調用。先把最簡單，沒有經過優化的程序寫出來。

主要的函數是cps1，它有兩個參數expr和ctx， expr就是一個s-expression， ctx是當前的context，是一個symbol或'(λ...)，比如(cps1'(fx)'(λ(x)x))的值為 '(fx(λ(x)x))。所以(cpsexpr)函數就是(cps1expr'id)或(cps1expr'(λ(x)x))，我暫且用前者。

先把用到的輔助函數都列出來

(define (cps expr)

  (define (atom? x)

    (not (pair? x)))

  (define n -1)

  (define (fv)

    (set! n (add1 n))

    (string->symbol (string-append "v" (number->string n))))

  (define (cps1 expr ctx)

    ....)

  (cps1 expr 'id))

下面就開始寫 cps1

(define (cps1 expr ctx)

  (match expr

    [(? atom? expr) ....]

    [`(λ(,x) ,body) ....]

    [`(,rator ,rand) ....]))

第一項，當 expr為一個變量(比如 ' x) 或一個原始類型(比如 123)時，我們就直接把它返回，比方說，

(cps1 'x 'k) ==> '(k x)

(cps1 'x '(λ(x) ....))

==> '((λ(x) ....) x)

因此，

[(? atom? expr) `(,ctx ,expr)]

對於一個λ表達式，我們同樣直接把它返回，只是先遞歸進函數體內，把函數體進行CPS變換。每個λ的continuation都為 ' k，比如，

(cps1 '(λ(x) x) 'id)

==> '(id (λ(x k) (k x)))

因此，

[`(λ(,x) ,body)

`(,ctx (λ(,x k) ,(cps1 body 'k)))]

處理函數的部分稍微復雜一些，CPS輸出應該先分別計算函數和參數，然后調用，舉個例子，在我們最終要完成的代碼里，應該大致是這樣的，

(cps1 '((f a) (g b)) 'id)

==>

'(f a

(λ(v0) (g b

(λ(v1) (v0 v1 id)))))

v0是 ( f a )的結果， v1是 ( g b )的結果，然后調用。

我們先要實現的會產生一些多余代碼，比如，

(cps1 '(f x) 'id)

==> '((λ(v0)

        ((λ(v1) (v0 v1 id))

         x))

      f)

<=> '(let ([v0 f])

       (let ([v1 x])

         (v0 v1 id)))

可以看出來，我們先分別CPS了 f和 x，然后調用 ( vf vx ctx )，其中 vf和 vx 是我們前面定義的 ( fv )產生的，這樣保證了每次都不一樣。 (說實話 ( gensym )也是可以的)

[`(,rator ,rand)

(define v-rator (fv))

(define v-rand  (fv))

(cps1 rator

       `(λ (,v-rator)

          ,(cps1 rand

                 `(λ (,v-rand)

                    (,v-rator ,v-rand ,ctx)))))]

於是就全部搞定了。

於是我就興高采烈地試了一下這個程序，

(cps '((f a) (g b)))

==>

'((λ (v4)

   ((λ (v5)

      (v4

       v5

       (λ (v0)

         ((λ (v2)

            ((λ (v3)

               (v2

                v3

                (λ (v1) (v0 v1 id))))

             b))

          g))))

    a))

f)

............

不過算作是個很好的開頭吧。

3 最簡CPS輸出

其實下面才開始真正的任務，上面一節只是因為，市面上的程序都是分好幾個函數，我要把它們合在一起。

上面的程序的問題就在於，當ctx是'(λ(v)...v...)，而且expr是一個'x之類時，輸出應該為 '...x...而不是'((λ(v)...v...)x)，照λ-calculus的術語說就是產生了一個beta-redex。

我們來觀察一下現在我們的CPS程序的3個分支產生的ctx

case 1:如果是一個atom，就產生`(,ctx,expr)，這時ctx在函數的位置。

case 2:如果是λ表達式，ctx也在函數的位置。

case 3:但如果是函數調用，這時ctx在參數的位置((vfvxctx))

很顯然，在參數位置時ctx是不可能被化簡的，因為結果必須是 (vfvxk/id)或(vfvx(λ(v?)???)) 的形式。而在函數位置時是有可能化簡的，當它是λ函數的時候。

所以 ctx在函數位置(case 1,2)時應該和case 3統一一下。就是說，為了化簡，我們把 ` ( λ ( v? ) ??? ) 的 quasiquote直接去掉，改成一個函數 ( λ ( v? ) ` ??? )，調用它就相當於直接把函數體里面的 v?替換掉了，比如

;; 原來的輸出

'((λ(v0) (f v0))

x)

;; 現在變成

((λ(v) `(f ,v))

'x)

==> '(f x)

第二，如果ctx是'k/id，就改成(λ(out)`(k/id,out))

因為總共就兩種情況:ctx在函數位置和參數位置。我們不妨把cps1函數的ctx參數改成兩個，一個叫ctx-f在函數位置，一個叫ctx-a在參數位置。

;; ctx-f : symbol -> s-exp ;; ctx-a : s-exp

(define (cps1 expr ctx-f ctx-a)

(match expr

....))

(cps1 expr (λ(out) out) 'id)

(λ(out)out)就是id這個函數。因為原來的slideshow 'id可以看成是`(λ(v?)v?)，所以化簡后就變成了(λ(out)out)

接下來就是 match的分支。

[(? atom?) (ctx-f expr)]

[`(λ(,x) ,body)

(ctx-f `(λ(,x k) ,(cps1 body

(λ(out) `(k ,out))

'k)))]

這兩個應該都是好理解的，都調用了 ctx-f來化簡當前的式子。

但是在CPS函數調用時遇到了一些麻煩,因為我們發現總共有4種情況需要討論，其中有大量重復的代碼，但是不管怎么說，先把代碼寫出來才是正道 (以下代碼會需要一點耐心)

[`(,rator ,rand)

(define v-rator (fv))

(define v-rand (fv))

(cps1 rator

       (λ(out-rator) (cps1 rand

                           (λ(out-rand) `(,out-rator ,out-rand ,ctx-a))

                           `(λ(,v-rand) (,out-rator ,v-rand ,ctx-a))))

       `(λ(,v-rator) ,(cps1 rand

                            (λ(out-rand) `(,v-rator ,out-rand ,ctx-a))

                            `(λ(,v-rand) (,v-rator ,v-rand ,ctx-a)))))]

令人驚訝的是，這個 cps函數就這么完成了! 跟前面的結果對照一下就會看出明顯區別了。它不但可以處理beta-redex，還能正確處理尾遞歸。

這段代碼應該也不難理解，只是分別討論了函數和參數分別處於函數位置和參數位置的情況。大致思路就是，首先，cps1要根據rator和rand 是否為一個atom來決定如何輸出，其次，我們不願意在遞歸進去之前就判斷一次，遞歸進去之后又要matchexpr(開頭提到的那篇論文的方法就有這個問題)。所以我們把現在的狀態分成了兩個參數，也一起遞歸進去。

它有唯一一個但很好修復的缺陷，就是v-rator和v-rand 定義地太早了，所以有時候會出現vn不連續的情況，如果不嫌麻煩的話可以在每次第一次出現v-?的地方再 (let([v-?(fv)])....)，當然這個代碼看起來就...... 另外，做出了下面一道習題后也會很好修復這個缺陷。

論文里的方法大概就是這樣，只是寫成了好幾個函數

[`(,rator ,rand)

(if (atom? rator)

     (if (atom? rand)

         ....

         ....)

     (if (atom? rand)

         ....

         ....))]

其實這樣寫也是完全可以的，這時 cps1只需要傳一個參數 ctx-f就可以了，在判斷出不是 atom以后用 ' vn調用 ctx-f，把它轉換成 ctx-a，這其實更接近王垠的版本。

甚至還有一個寫法，就是利用多返回值，再返回一個布爾值表示當前的選擇，這個方法看起來會有些麻煩。就不提了。

現在舉兩個例子，看一下這個程序是怎么做的，

(cps '(f x))

==> (cps1 'f ctx-f ctx-a)

==> (ctx-f 'f)

==> (cps1 'x .... ....)where[out-rator='f]

==> `(,out-rator ,out-rand id)where[out-rator='f out-rand='x]

==> '(f x id)

但如果沒有beta-redex可以化簡，比如，

(cps '((f a) b))

==> (cps1 '(f a) ctx-f ctx-a)

==> `(f a ,ctx-a)where[ctx-a='(λ(v0) (v0 b id))]

==> '(f a (λ(v0) (v0 b id)))

現在可以來看一下這段代碼對我們有什么啟發。

所謂的continuation-passing style多用一個參數k來告訴我們要調用的函數當前的狀態是什么，就是這個函數運行完了以后應該干什么。但這里的cps函數也有一個參數ctx，它也表示一個狀態，它表示的是現在的狀態，讓更深層遞歸的函數能得知一些外部信息。

很多時候我們發現就傳一個死的數據(比如第一個版本里的ctx) 是不夠的，不但遞歸進去的函數需要這個數據，而這個數據也要隨着當前的情況而變化。在簡單的情況下我們可以傳好幾個參數，或者一個對象進去，里面的函數選擇性地使用這些數據。但是在支持高階函數的語言里很多時候方便很多，因為我們可以傳一個函數進去。

更通常的情況下，上面的 cps1可以只有一個參數 ctx，其中 ctx是這樣的，

ctx=

(λ(position)

  (cond

    [(eq? position 'f) ....]

    [(eq? position 'a) ....]))

這兩個分支分別為原來的 ctx-f和 ctx-a。

這種模式更廣泛的應用之處在於ctx的參數不是一個用來選擇的符號，而是一個連續數值或對象的時候。我一下子想不出實際的例子，有了我會補充。

習題:請擴展這個程序以支持多參數的λ和函數調用。

令我驚訝的是，支持多參數就不用分4類討論了! 因為只要分兩類討論，依次遍歷整個列表就可以，不用區分函數和參數，所以代碼反倒簡單多了。推薦做一下這個習題。

4 簡化 cps1 函數

這一節，我們把ctx-f和ctx-a合並成一個ctx

觀察所有產生的ctx-f/a參數，總結一下總共有這些:

1. λ(out)`(k,out)

2. 'k

3. λ(out)....`(....,out....)

4. `(λ(,vn)....(....,vn....))

如果要只傳一個參數的話，我們會發現，由2可以推出1, 因為我們只要給它包一個λ就可以了。由3可以推出4，如果3是ctx，4就是 `(λ(,vn),(ctxvn))

問題就在於，13是一個形式的，24是一個形式的，我們要選擇的就是只傳13還是只傳24.

我們發現，13是兩個固定的值，而24里面是有一堆省略號的，也就是說，如果采用一點類似作弊的策略，從1也可以推出2，只要判斷ctx是否等於(λ(out)`(k,out))。但是無論如何也不可能從任意的4推出3(當然你如果使用eval 的話，我就沒話說了，按理來說是可以的，你可以自己嘗試一下，成功了記得偷偷告訴我一聲。

於是，我們決定采用13型的ctx。

先寫一個轉換的函數，把一個13型的 ctx轉成24型的。就是把 ctx-f轉成 ctx-a

(define (ctx-f->a ctx)

  (if (ctx1? ctx)

      'k

      (let ([v (fv)])

        `(λ(,v) ,(ctx v)))))

但是我們怎樣知道一個 ctx是不是 ctx1?呢，就是不直接寫出 ( λ ( out ) ` ( k , out ) )，而是定義

(define ctx1 (λ(out) `(k ,out)))

(看到了嗎，這就是王垠CPS代碼里的 ctx0)

於是，

(define (ctx1? ctx) (eq? ctx ctx1))

這樣完成了之后，原來那段代碼所有的 ctx-a都不用手寫了，只要改為 ( ctx-f->a ctx-f )就可以了。因為能這樣直接轉化，所以也沒有必要傳兩個 ctx參數了，我們在需要用到 ctx-a時現轉化就可以，於是，我們最終得到了這樣的代碼。

(define (cps1 expr ctx)

  (match expr

    [(? atom? expr) (ctx expr)]

    [`(λ(,x) ,body)

     (ctx `(λ(,x k) ,(cps1 body ctx1)))]

    [`(,rator ,rand)

     (cps1 rator

           (λ (out-rator)

             (cps1 rand

                   (λ (out-rand)

                     `(,out-rator

                       ,out-rand

                       ,(ctx-f->a ctx))))))]))

是的，這就是完整的λ-calculus的最簡的CPS變換，如果對照一下王垠的CPS變換的最后幾行，你會發現我的這個版本甚至更清晰一些，因為我用 ctx-f->a這個函數避免了 ` ( , out-rator , out-rand .... ) 這樣重復的代碼，並把判斷也放進了這個輔助函數中。

你現在可以自己隨意試驗這個程序了。

下面我們對它進行一些擴展，先增加多參數的λ和函數調用，然后是原生的幾個函數(比如+,-,zero?等)，最后添加if語句。

5 多參數和原生函數

都已經到這一步了，支持多參數其實很簡單。

第一種情況不用變。

[(? atom? expr) (ctx expr)]

第二種情況只要把λ的參數改成一個列表就可以了。

[`(λ ,args ,body)

(ctx `(λ(,@args k) ,(cps1 body ctx1)))]

第三種情況甚至只要遍歷列表就可以

[_ ; else : expr = ‘(,rator . ,rands)

(let recur ([exprs expr] [acc '()])

   (if (null? exprs)

       `(,@acc ,(ctx-f->a ctx))

       (cps1 (car exprs)

             (λ(v) (recur (cdr exprs) `(,@acc ,v))))))]

說實話，最后這個遞歸...我也不知道我自己是怎么寫出來的，思路就是這樣，首先最后返回值肯定是
( cps1 ( car exprs ) ( λ ( v ) .... ) )
然后省略號要填的是，遞歸遍歷 ( cdr exprs )，所以結構必須是這樣的，

(let recur ([exprs expr])

  (cps1 (car exprs) ; when exprs is not null

        (λ (v)

          (recur (cdr exprs))

          (process-v))))

然后當 exprs最后變成 null的時候，我們需要把之前所有的 v : v1 v2 v3 ....收集起來，返回 ` ( , v1 , v2 , v3 .... , vn , ( ctx-f->a ctx ) )，因此再多一個變量 acc，用來收集這些 v。這個程序就完成了。你可以自己試驗一下確保它正確。

接下來我們要支持一些原生的函數，比如 + - * / zero?，這些函數不需要經過CPS變換，比如，

(cps '(+ x y))

;; instead of (+ x y id)

==> '(+ x y)

(cps '(+ (f x) y))

==> '(f x (λ(v0) (+ v0 y)))

(cps '(+ (* x y) z))

==> '(+ (* x y) z)

;; when used as higher order procedure

(cps '(((λ(m) +) n) ; returns +

       x y))

==> '((λ(m k)

        (k +))

      n

      (λ(v0) (v0 x y id)))

首先，

(define (trivial? x)

(memq x '(zero? add1 sub1 + - * /)))

然后，處理函數的不分大多數不變，先在最后進行一個判斷。

[_

(let recur ([exprs expr] [acc '()])

  (if (null? exprs)

      (if (trivial? (car acc))

          ....

          `(,@acc ,(ctx-f->a ctx)))

      (cps1 (car exprs)

            (λ(v) (recur (cdr exprs) `(,@acc ,v))))))]

如果 ( car acc )是原生函數，就把整體當成一個值返回，所以非常簡單，

(if (trivial? (car acc))

(ctx acc)

`(,@acc ,(ctx-f->a ctx)))

這樣就完成了。現在只差最后的難點，就是 if語句了。

6 if語句

首先， if語句的不同之處在於，我們需要對它的兩個分支做兩次CPS變換，比如，

(cps '(λ(x) (if a b (f c))))

==> '(λ(x k) (if a (k b) (f c k)))

是的，這個沒什么難的，既然都到這里了，應該隨手就可以寫出來了，

[`(if ,test ,conseq ,alt)

(cps1 test

       (λ(t) `(if ,t

                  ,(cps1 conseq ctx)

                  ,(cps1 alt ctx))))]

把這段代碼插到CPS函數調用的前面。我一開始以為就這么完成了，結果發現，這里的 ctx是會被翻倍的。

如果 ctx是 k的話，沒有關系，但是如果是一個比較復雜的式子，就是說， if語句嵌套在了不是尾遞歸的地方，比如函數的參數，或 if 的判斷語句，就出現了一些問題，

(cps '(λ(x) (f (if a b c))))

==> '(λ (x k) (if a (f b k) (f c k)))

看起來沒什么毛病，但那是因為我們的 ctx太簡單了，制造一點更復雜的，

(cps '(λ(x) (f (g (h (if a b c))))))

==> '(λ (x k)

       (if a

           (h

            b

            (λ (v0) (g v0 (λ (v1) (f v1 k)))))

           (h

            c

            (λ (v2) (g v2 (λ (v3) (f v3 k)))))))

(cps '(λ(x) (if (if a b c) d e)))

==> '(λ (x k)

       (if a

           (if b (k d) (k e))

           (if c (k d) (k e))))

一眼看過去，長得一模一樣的代碼很多，第一個例子里面， ( h b/c .... )就只有 b和 c不同，后面完全一樣，第二個例子也是這樣。

解決方法也很直接，在發現當前的 ctx不是最簡的時候，我們用一個 let 包住當前的 ctx，最終結果變成這樣，

(cps '(λ(x) (f (if a b c))))

==> '(λ(x k)

       (let ([k (λ(v0) (f v0 k))])

         (if a (k b) (k c))))

(cps '(λ(x) (if (if a b c) d e)))

==> '(λ(x k)

       (let ([k (λ(v0) (if v0 (k d) (k e)))])

         (if a (k b) (k c))))

(這看起來怎么像是ANF和CPS的混合版)

所以，也是只要加一個判斷即可，

[`(if ,test ,conseq ,alt)

(define (if-body ctx)

   (cps1 test

         (λ(t) `(if ,t

                    ,(cps1 conseq ctx)

                    ,(cps1 alt ctx)))))

(if (ctx1? ctx)

     (if-body ctx)

     `(let ([k ,(ctx-f->a ctx)])

        ,(if-body ctx1)))]

好了，現在一切都已經完成了。這個教程可以基本圓滿地結束了。

那么最后，還有一點小問題可以修復，就是在這種 ctx是 id的時候，

(cps '(if a b c))

==> '(let ([k (λ(v0) v0)])

(if a (k b) (k c)))

我們不太希望它只是 ( let ( [ k id ] ) )，而希望直接輸出 ( if a b c )，因此，我們把 id也歸為一類 ctx (這應該是第五類了？)，先插入定義

(define id (λ(x) x))

(define (id? x) (eq? x id))

然后把處理 if的代碼加上一句，

(if (or (ctx1? ctx)

        (id? ctx))

    (if-body ctx)

    .... as before)

最后，

(define (cps expr)

....

(cps1 expr id))

就可以了。

再最后，如果有人看着 ( cps ' ( f x ) ) 變成 ' ( f x ( λ ( v0 ) v0 ) ) 感覺不爽的話，可以這么改一下，

(define (ctx-f->a ctx)

  (cond

    [(ctx1? ctx) 'k]

    [(id? ctx) 'id]

    [else

     (define v (fv))

     `(λ(,v) ,(ctx v))]))

;; 話說這里用 case 語句會更舒服一點的...

這樣， ( cps ' ( f x ) ) 就是 ' ( f x id ) 啦。

7 總結

這么多代碼看下來，其實你會發現，就只有幾個關鍵點，只要想到了，其實也沒有多難。我自己想這個程序的時候，手頭上沒有電腦，我是寫在紙上的(好痛苦啊)，但是放到電腦上測試，一次性就全是對的，畢竟不是很大的工程，也沒有各種復雜的角角落落需要考慮，思路還是很簡單的。

這個程序還有升級空間，就是 begin 和 set! 語句，提示一下，東西越來越復雜的時候，可能不得不回歸到第3節中的方式，把各種 ctx 拆開，否則處理 set! 的時候會產生一堆嵌套的 begin 語句。另外，如果是Common Lisp里的那種有返回值的賦值語句，處理起來會簡單一些，因為可以簡單地看作一個表達式。

最后就隨便說說，其實這段代碼也沒有特別的高級，只是自己寫出來了，那就開心一下就好。代碼里倒是有幾個挺特別的想法值得學習。

沒得寫了，就打個廣告吧，本文章用 scribble 生成，不過用它只是因為，它是唯一一個支持racket代碼高亮的......

文筆不好請見諒，有任何錯誤或寫的不好的地方歡迎指出。

最后的最后，我把本篇文章完整的代碼貼上，

(define (cps expr)

  (define (atom? x)

    (not (pair? x)))

  (define n -1)

  (define (fv)

    (set! n (add1 n))

    (string->symbol (string-append "v" (number->string n))))

  (define ctx1 (λ(out) `(k ,out)))

  (define (ctx1? ctx) (eq? ctx ctx1))

  (define (ctx-f->a ctx)

    (if (ctx1? ctx)

        'k

        (let ([v (fv)])

          `(λ(,v) ,(ctx v)))))

  (define (trivial? x)

    (memq x '(zero? add1 sub1 + - * /)))

  (define id (λ(x) x))

  (define (id? x) (eq? x id))

  (define (cps1 expr ctx)

    (match expr

      [(? atom?) (ctx expr)]

      [`(if ,test ,conseq ,alt)

       (define (if-body ctx)

         (cps1 test

               (λ(t) `(if ,t

                          ,(cps1 conseq ctx)

                          ,(cps1 alt ctx)))))

       (if (or (ctx1? ctx)

               (id? ctx))

           (if-body ctx)

           `(let ([k ,(ctx-f->a ctx)])

              ,(if-body ctx1)))]

      [`(λ ,args ,body)

       (ctx `(λ(,@args k) ,(cps1 body ctx1)))]

      [_

       (let recur ([exprs expr] [acc '()])

         (if (null? exprs)

             (if (trivial? (car acc))

                 (ctx acc)

                 `(,@acc ,(ctx-f->a ctx)))

             (cps1 (car exprs)

                   (λ(v) (recur (cdr exprs) `(,@acc ,v))))))]))

  (cps1 expr id))