棧程式設計和函式控制流: 從 continuation 與 CPS 講到 call/cc 與協程

RzBu11023R發表於2021-08-28

原標題:尾遞迴優化 快速排序優化 CPS 變換 call/cc setjmp/longjmp coroutine 協程 棧程式設計和控制流 講解

本文為部分函數語言程式設計的擴充套件及最近接觸程式語言控制流的學習和思考,主題是棧程式設計和控制流相關,涉及內容有 堆疊程式設計總結, 函式式語言的CPS變換,python 如何實現尾遞迴優化裝飾器及其思想方法的總結應用,快速排序的演算法導論寫法的一種視角/分析,C 語言setjmp/longjmp 函式的作用和實現分析,如何實現 C/C++ 下的協程庫編寫的一些思路和要點分析。雖然是大雜燴,但是主要內容都與棧和控制流相關。主要服務於不久後進行的網路程式設計專案中協程庫部分的前情提要,另外也為將來進一步學習魔法預備。


遞迴與棧

首先理解棧程式模型,其函式呼叫是依據壓棧進行的,這裡給一副圖加深印象:
RISC-V 呼叫約定
函式呼叫最後返回的時候 callee 需要從棧恢復 callee save 暫存器然後返回到 caller,caller 再恢復 caller save 暫存器,這一步就是佔用記憶體的消費。


尾遞迴

我們是否能夠把遞迴變成 \(O(1)\) 空間的呢?實際上是可以的,但是這對程式有條件,就是尾遞迴。下面來看幾個程式:

int fib(int n){                        int sum(int n) {
    if(n < 2) return 1;                    if(n<1) return 0;
    return fib(n-1) + fib(n-2);            return sum(n-1) + n;
}                                      }

int sum2(int n, int acc){              int fact(int n, int acc){
    if(n<1) return acc;                    if(n<1) return acc;
    return sum2(n-1, acc + n);             return fact(n-1, acc * n);
}                                      }

這裡一共有四個程式,我們很容易理解只有下面的兩個可以 \(O(1)\) 返回,這是因為中間的運算是無必要的,函式呼叫者不需要在遞迴呼叫返回後進行其他運算再返回,這樣我們彙編層面不需要壓棧儲存暫存器從而直接 jmp 到函式呼叫就行了。一直執行到最後一個函式呼叫就能 a0 作為返回值返回結束函式。上面的階層函式寫成虛擬碼就是這樣:

fact:
let register a0 as n, a1 as acc
label 1: if (n<1) a0 = acc, ret
label 2: acc = acc * n
label 3: n = n-1
label 4: goto label 1

這種就是尾遞迴和尾遞迴優化。


Python 尾遞迴優化

我們在 Python 裡可以寫一個裝飾器讓通用的尾遞迴函式可以進行優化從而去掉執行的堆疊。廢話少說,直接上程式碼:

class TailRecurseException:
    def __init__(self, args, kwargs):
        self.args = args
        self.kwargs = kwargs
def tail_call_optimized(g):
    def func(*args, **kwargs):
        f = sys._getframe()
        if f.f_back and f.f_back.f_back \
            and f.f_back.f_back.f_code == f.f_code:
            # 丟擲異常
            raise TailRecurseException(args, kwargs)
        else:
            while 1:
                try:
                    return g(*args, **kwargs)
                except TailRecurseException, e:
                    args = e.args
                    kwargs = e.kwargs
    func.__doc__ = g.__doc__
    return func
@tail_call_optimized
def factorial(n, acc=1):
    "calculate a factorial"
    if n == 0:
        return acc
    return factorial(n-1, n*acc)

你看懂了嗎?一開始看可能有點一頭霧水,當然我們不瞭解 Python 的執行堆疊的概念就無法理解。結合 decorator 的程式設計方法,我們可以做這樣的思想實驗:我們先偽裝一個函式把引數控制權給偷過來(裝飾器內自定義函式即中間人),注意裝飾器一旦裝飾,其原函式(factorial)內部遞迴將轉接到裝飾器內的自定義函式(func)中。然後我們知道是遞迴的時候就清空棧,然後重新呼叫一個原函式傳入新的引數。

具體實現涉及一些小技巧,比較魔幻,下面講解:

  • 裝飾器啟動後,我們呼叫 factorial 100,此時實際呼叫的是 func 100,然後 func 100 會呼叫 g 100
  • g 99 進入原來的 factorial 99,然後 return func 98(注意這裡遞迴呼叫是被掉包了)
  • func 98 將再次進入 func,此時的堆疊是: func 100 -> g 100 -> func 99
  • func 發現發生了遞迴,於是其丟擲一個異常,並且把這個最新遞迴呼叫的引數給異常帶走了
  • 一旦異常發生,棧將直接清空直到回到第一次呼叫產生的地方,那就是第一個 funcwhile 1 中的抓異常的地方
  • 然後我們完成了清空棧的目的,然後再呼叫 g 99 ,就等價於尾遞迴控制權直接給下一任返回

感覺這裡是不是好像 DNS 查詢裡的迭代查詢?就本來建立了一條查詢鏈的 local 向伺服器 A 發起 “查詢” 呼叫,伺服器 A 向 B 發起又一次 “查詢” 呼叫,直到返回後再返回給 local。這樣效率不好而且涉及佔用伺服器資源,連線必須等待。改進則變成了 local 向 A 查詢,A 說你去向 B 查詢,A 又向 B 查詢。這裡的 A 只是計算出 B,即上文中的 “引數”,在這裡,函式呼叫是 “查詢” 函式,而引數則是伺服器的名字。


快排遞迴優化版

(本節和後續主題內容無關)對巨量資料進行快速排序真正用起來是會爆棧的,但是我們清楚的認識到快速排序的程式碼並不算尾遞迴,所以很明顯他無法進行尾遞迴優化,但是其實我們還是可以優化一部分暫存器的儲存和恢復的操作的,那就是進行剪枝,當然這種寫法很難說不算尾遞迴優化 idea 的啟發。演算法導論永遠的神。

    public int patition(int[]nums, int l, int r){
        int pivot = nums[r-1];
        while(l<r-1){
            while(l<r-1&&nums[l]<=pivot)l++;
            if(l<r-1) nums[--r] = nums[l];
            while(l<r-1&&nums[r-1]>pivot) r--;
            if(l<r-1)nums[l++] = nums[r-1];
        }
        nums[l] = pivot;
        return l;
    }
    public void qs(int[]nums, int l, int r){
        int sep;
        while(l<r){
            sep = patition(nums, l, r);
            qs(nums, l, sep);
            l = sep + 1;
        }
    }
    public void qs_v2(int[]nums, int l, int r){
        int sep;
        while(l<r){
            sep = patition(nums, l, r);
            qs_v2(nums, sep+1, r);
            r = sep;
        }
    }

分析我就不做了,很容易理解的本來是對左右子樹分別遞迴下去,但是由於我們知道樹的左邊界或者右邊界是固定的,很容易想到消除一部分遞迴呼叫,即左樹,左樹的左樹,左樹的左樹的左樹都在 while 迴圈裡面搞定了,即減少了一半的暫存器壓棧出棧,最後結果就是節約一點點空間。說這是一種尾遞迴優化是因為他的思想和尾遞迴優化是一致的。
優化節約空間


尾呼叫

但是其實尾遞迴優化是沒有意義的,這是因為我們很容易就能把尾遞迴程式寫成迴圈的形式,像 Java (Java 8)這種純 OO 語言編譯器本身就不提供尾遞迴優化,因為程式設計師總是能在希望提高效能時手動完成尾遞迴到迴圈的轉換,而 C/C++ 的編譯器如 gcc/g++ 就提供該優化,具體實現方案則正如上文提到的那樣是不進行暫存器的儲存和恢復直接 jmp

但是我們知道 CS 的發展歷程很多時候一個 idea 並不總是在本地方有用的,尾遞迴程式設計師本身可以寫成迴圈,我們遵循唯物主義的教導從特殊到一般分析,對於尾遞迴的更一般的情況尾呼叫來說,程式設計師是無法在不改變 Coupling 情況下優化的,這時候這個尾呼叫優化不進行暫存器的儲存和恢復就大顯身手。然而 Java 出於某些堆疊計數依賴的原因並不提供,希望編寫高效能()的 Java 程式時留意。

所以為什麼有尾呼叫的需求,正常不是我們返回值然後呼叫新的函式不就行了嗎?這其實很好理解,比如有一個應用——非同步程式設計,我們在 C# 和 javascript 程式設計裡面已經見的多了,那就是大名鼎鼎的非同步程式設計,此時這個尾呼叫實際上是一個回撥函式,這樣實際上尾呼叫會延遲在另一個核心裡執行,我並不需要等待他的返回就行執行其他的程式流。接下來就要具體講解這種尾呼叫的程式設計風格


閉包與柯里化

前面 Python 的曲線尾遞迴優化的 中間插一腿 的思路能幫助我們理解 Continuation Passing Style 程式設計。雖然這個東西理論上好像很複雜,搞魔法的可能可以用邏輯學/形式邏輯里程式等價於證明的 Curry-howard Correspondence 理論相關的東西來講。我這裡比較低階,只做一些簡單的筆記,我們先從講解函式式語言裡容易理解的 閉包(closure) 概念開始。

函數語言程式設計裡面函式是一等公民,除了引數就就沒有變數了,因為變數會存在 state ,這個東西會導致函式呼叫有副作用,所以要實現迴圈都用遞迴來實現的,引數傳值。

我們來看一段 Common Lisp 程式碼:

(define f (lambda (x) (lambda (t) (+ x t))))

這裡的 (lambda (t) (+ x t)) 就是一個閉包函式,他能夠讀取到 (lambda (x) ...) 函式的內部變數 x 的值並且用來做計算,但是這個變數對閉包函式來說是其外部的。

閉包的概念即這個函式把他的外部給包了起來,這麼說,當定義的時候,閉包函式的外部是某個函式,而當他作為高階函式返回給別人的時候,他就出到另外的環境了,函式的 Scope 變了但是他還能訪問到之前的那個外部環境,這就是因為他是把之前定義的地方的外部狀態給封閉打包了。

如果還是不理解閉的包是外部的包,請想象成你想辦法把房子打包瞭然後吃下去了,當你到別的地方的時候,你又還是住在你自己的房子裡活動。

當然 OOP 裡的類這種封裝實際上也有閉包,我們可以認為成員變數的 properties 就是內部變數,而 selector 就是閉包函式能夠給外部人一種手段訪問內部變數,當然這裡的成員變數是 mutable 的,FP 裡的是 immutable 的。

Curry-ing 就是一種多引數函式的閉包改寫方法,λx.λy.x+y 即是 (lambda (x y) (+ x y) 的閉包改寫。接下來我們會講 CPS 和 CPS 變換,這兩者的關係就和 閉包與柯里化 的關係差不多。


回撥

閉包函式能夠把內部變數通過某些控制後暴露給外部,那麼接下來講回撥。Callback 實際上很常用,比如非同步常常就是用 callback 實現的,callback 本質上是一種 CPS 模式。

算了不講了,回撥很容易理解,就是非同步程式設計裡面另一個執行緒去執行的尾呼叫而已。

如果程式有連續的太多的非同步呼叫,程式碼裡面就會引發回撥地獄,這一點不好看,就像 Lisp 那種括號語言,偷到的程式碼最後一頁肯定全是括號。


CPS 和 CPS 變換

CPS 是一種 style,式如其名,就是函式的呼叫引數裡面有一個 continuation,而程式不會直接返回而是隨著這個 continuation 執行下去,所以說這是一個控制流上的 style。

CPS 變換則是完成這種風格轉換的變換,當然我們可以通過人工來變換,但是實際上熟悉 FP 中 everything is data 的概念,我們順理成章地可以通過操縱程式資料用程式完成 CPS 變換。

為了緩解枯燥 () 這裡插播一個魔法知識點:

「CPST 就是 Gödel–Gentzen 變換的 Curry–Howard 像而已,這有什麼難理解的?」 CPS 變換有什麼作用? - Belleve的回答

這句話的理解可能得搞魔法的那群人才看得懂,這裡 Curry-Howard Correspondence 是一個理論說命題的證明和程式是同構的,然後 Gödel–Gentzen 變換 不過是邏輯學裡的一個定理:對於任意的經典邏輯下的證明,我們可以把它轉換為一個直覺主義邏輯下證明而不損失任何證明能力,對於經典邏輯和自覺邏輯的區別請有興趣的同學自行學習魔法。所以 CPS 變換的一個程式其實是哥德爾-根岑變換這個定理的證明通過 Curry-Howard 同像理論在程式集中的一個對應。

回來講 CPS,前面我們都是從尾呼叫的角度講的,這樣其實這個時候對 CPS 的印象已經很清楚了,應該就是一個函式指標的回撥而已唄。下面再從非尾呼叫的函式的角度來看一下 CPS 的樣子。

對於尾遞迴本身是一個 \(O(n)\) 空間的遞迴程式,可以優化到 \(O(1)\) 並且還能避免資料過大導致的 Stack Overflow 問題,但是對於普通的遞迴來說,\(O(n)\) 的空間消耗的必須的,因為他一定要儲存狀態回退/回溯。 如果能把遞迴轉換成尾呼叫的程式設計風格,即CPS 變換也就能優化一個棧溢位的問題,把棧空間消耗放到堆上去而已。

當然我們其實知道人工用 資料結構 + 迴圈 完全模擬遞迴呼叫也可以達到目的,比如下面的快速排序程式(頭條校招面試題):

    public void qsort(int[]nums, int l, int r){
        Deque<int[]> s = new ArrayDeque<int[]>();
        s.addLast(new int[]{l, r});
        int count = 0;
        while(!s.isEmpty()){
            int[] temp = s.pollFirst();
            if(temp[0]>=temp[1])continue;
            l = temp[0]; r = temp[1];
            int sep = patition(nums, l, r); //經典分割槽程式碼略
            s.addLast(new int[]{l, sep});
            s.addLast(new int[]{sep+1, r});
        }
    }

樹形遞迴其實有點超綱了(其實只是為了複習一下題目),我們還是看回比較正常的遞迴吧:

 int sum(int n) {               int sum2(int n, int acc){     
     if(n<1) return 0;              if(n<1) return acc;       
     return sum(n-1) + n;           return sum2(n-1, acc + n);
 }                              }                             

sum 來說需要先計算出 sum(n-1) 才能計算出 sum,而 sum2 這種也不是 CPS 風格(這是依賴特定程式結構而人工才能改寫的),下面給出 CPS 風格的程式設計 (請回顧將函式結果作為引數持續執行下去):

#如果理不清請多看幾遍:
def sum_cps(n, c):
  if n == 0:
    return c(0)
  else:
    return sum_cps(n-1, lambda x: c(n + x))
sum_cps(10000, lambda x:x)

我們可以看見,這種情況下的 CPS 變換(儘管是人工) 是一種奇技淫巧,因為他不是用 棧/佇列 等資料結構來實現呼叫的模擬,而是用閉包來實現的,我們複習 FP 裡面閉包是可以通過引數實現區域性變數(成員物件放堆上)的(這裡涉及 FP 語言的編譯器和直譯器,就和我們當時用編寫 scheme 直譯器那種,值得注意的是一般還涉及 GC),所以可以用尾遞迴形式配合編譯器把爆棧問題解決掉。

實際的執行好像就不算回溯的版了,而是一個類似自底向上的鏈式求解過程:sum 100 -> sum 99 ->... 而關鍵的棧部分或者說 acc 部分已經通過閉包的形式存在了這個 lambda 匿名函式中去了!

結論是 CPS 程式設計本質是基於閉包的無棧程式設計(當然其直譯器如何工作則另說)。

當然程式上的 CPS 變換 太過於複雜,我目前還是先略過,附上論文地址供將來學習魔法的時候非同步回撥學習。。。。Representing Control: A study of the CPS transformation 以及一個 PL 課程:PL


setjmp 與 longjmp

那麼如果所有的函式都 CPS 變換後,就能用簡單的方法實現一個名為call/cc的在 FP 中實現控制流的函式,這個函式本身是無法用 FP 語言定義的,為了能夠理解 call/cc 與 CPS 變換的關係,我們想要知道 call/cc 幹什麼,在那之前我們先學習 C 語言中的一個簡化版的 call/cc

對這個的學習也是 C 語言異常控制的一個思路。但是 jmp 系列只能儲存暫存器通過恢復PC來實現跳轉,共用一個棧,所以協程還是老實用 ucontext 好了(jmp 系列也能實現,到時候看一下),不過 CPS 的 style 對於實現 協程 yield 不是很友好嗎,畢竟協程就是為了使用者態實現 sequential 執行的迷你執行緒。

#include <setjmp.h>
int setjmp(jmp_buf env);
// 返回值:若直接呼叫則返回0,若從 longjmp 呼叫返回則返回非0值
void longjmp(jmp_buf env, int val);

jmp 系列函式內容特別簡單,甚至能讓人推出他的彙編實現。當然這裡還有一些要點值得注意的,第一點,我們已經說了他是共享棧的,所以千萬不能讓某個函式 return 之後再跳過來,到時候由於 sp 指標對應的棧很有可能已經被新程式用過了,此時引用記憶體變數涉及編譯器對棧的使用,馬上會導致未定義行為

結論是如果想要實現協程,我們應當編寫 while 迴圈,或不使用棧上的變數而只使用堆的,或者人工編寫另一套棧,這就變成了無棧程式設計,反而更難受。ucontext 的結構體內容中就包含棧空間的指標,這個實現可能和核心的 signal handler 註冊的時候能指定棧的實現相關,想實現協程的可以再深入學習。不過感覺就是 6.s081 裡使用者執行緒的 lab thread 那種感覺?我覺得很有道理,畢竟 xv6 裡面 user proc 是不會分散到 multicore 上執行的,所以這個 lab thread 本質就是協程啊哈哈哈(對於協程的理解本文是採用構建不併行的最小化程式目的理解的,然後協程分為無棧協程和有棧協程),雖然目的是方便上面(user space)排程,但是協程實際上使用,也可以保留能並行的功能,但是我主要理解是依據協程是另一種意義上的基於 continuation 的控制流。

所以這裡再附送一個 lab thread 裡面的 thread 結構體程式碼如下,也許這個就是 ucontext 的簡化版吧。

lab thread


call/cc

講完 setjmplongjmp ,再來講 lisp 裡面的 call/cc。這個有點複雜,全程是 call with current continuation

continuation 的概念上文已經很熟悉了,這裡再指定在 lisp 中的概念,continuation 是當 call/cc 被呼叫時建立的當前呼叫 call/cc 上下文中的 continuation。第二點是發生什麼,具體來說就是 call/cc 打包當前的 continuation 然後傳給其引數(打包後類似與 jmp 系列中的 jmp_env),當這個 continuation 被呼叫的時候,其引數會被返回。下面看一個具體例子:

(+ 1 (call/cc (lambda (k) (k (+ 2 3)))))

可能你很難理解,我們拆開來分析吧!當前我們正處於 (call/cc ...) 的語境下,馬上可以分析出其 continuation 是:(+ 1 ...) 請留意這個 ...,因為他是我們呼叫 (continuation return-value)return-value 返回的地方(即取代 call/cc 的函式呼叫)!

然後我們再看 call/cc 的引數,這個東西將會被呼叫,即控制流轉到這個 lambda (k) (k (+ 2 3)) 函式中去,而 k 將會被 call/cc 傳 continuation 進來!其具體流程用 python 講解:

#原始碼等價於:
  f = lambda k: k(2+3)
  plus(1, call/cc(f))
#call/cc 執行時等價於:
  def continuation(x):
    plus(1, x)
    f(continuation)
#結果:
    f(continuation) 
  = continuation(2+3)
  = plus(1, 2+3)
  = 6           

在分析具體流程的時候是不是很容易有一種 CPS 的感覺?這種 CPS 的感覺實際上就是 continuation 作為了一個間接的引數實現了即我們本來要做一件事情,中間插一腳給 call/cc 呼叫的函式,之後再 continuation 回來到 call/cc 上,只不過 CPS 程式設計中吧 continuation 放到了引數裡而已。我們可以發現如果寫程式的時候我們完全用 CPS 來寫,call/cc 的實現將輕而易舉(call/cc 此時本身也用 CPS 寫了)!比如下面這段 javascript 程式碼(程式碼來自知乎):

function callcc(f, k) {
  return f(k);
}

callcc(function(x) {return x(4 * 3);}, 
       function(y) {return 1 + y;});

所以這也解釋了為什麼能實現 CPS 變換之後 call/cc 的實現就變簡單了(直接有 continuation 不用打包),因為 call/cc 就是 CPS 程式設計的不用引數版!

(define call/cc (lambda (f k) (f (lambda (v k0) (k v)) k)))

當然要完整學習 CPS 變換太長時間了了,我得留個坑,這個東西涉及 PL 的魔法,凡人避免走火入魔。


coroutine

本來感覺協程和本文的內容好像也沒有什麼關係?非也,coroutine 可以把本來的多次回撥變成一個連續的過程,協程就是 continuation 的帶排程管理器的版本,不如說是擴充套件化的 continuation。下面講解協程怎麼去實現。

第一點是,我們希望協程是 sequential 執行的而不是 parallel 的,這樣才能有效避免使用各種併發控制手段如鎖,並且因為程式本身知道所有同步資訊,能夠最大效率排列協程的執行,而不存在鎖與輪詢的浪費空轉

協程分有棧和無棧的,對於棧的處理這個我們上面講 jmp 的時候講過了。

我這裡再提一個點,如果你想實現自己的協程庫,要考慮的處理棧的處理,還有一個關鍵是對阻塞 I/O 的處理。為什麼要關心這個呢?這是因為我們用 協程 主要是網路引用下寫的,所以必然涉及到 read 和 write 等阻塞式系統呼叫,這時候整個執行緒都會阻塞,我們必須解決讓協程掛起,一種方案是通過單獨的執行緒去完成阻塞呼叫,一種是 hook 系統呼叫。hook 的原理也很簡單,我們曾經學習過程式設計師自我修養連結裝載與庫,很容易想到通過連結時的同名函式覆蓋即可實現 hook,當然涉及hook中呼叫原來函式則要使用dlsym 去查詢 so,具體的很多內容我忘記了,我們之後會重新複習連結裝載與庫這本書(當然也配合APUE裡面還是齊全的)再來編寫協程庫的部落格。

I/O複用模式(事件驅動,Linux下的select、poll 和 epoll 負責將底層 socket 的時分(理論上是包封)複用給封裝出來)下本來是非同步的,這個另說,接下來我們將會系統學習網路程式設計和 Unix 高階程式設計。(待補充...)

建議學習的協程庫開源專案:libgo


總結

...(留個坑)

相關文章