使用 Coroutine 實作 Iterator

,

若想在不同的資料結構中套用相同的演算法操作資料，往往會使用 iterator pattern。Iterator 的原理很簡單：不同的容器都要實作出具有相同介面的 iterator，把巡訪資料的過程封裝在 iterator 之中，使用者只需要知道如何使用這個共通的 iterator 介面，就可以操作各式各樣結構不同的容器。

傳統的 iterator 實作方法，往往是把巡訪容器時的狀態儲存在某個變數當中。比如說巡訪陣列時，會在 iterator 中儲存陣列的 index；而在巡訪 linked list 時則是儲存其中一個 node。在更複雜的場合－－比如說樹狀結構的 iterator 可能要儲存 stack，實作起來就沒那麼簡單了。然而帶有 call stack 的 coroutine 卻可以很漂亮地實做出這種複雜結構的 iterator。

簡介 Lua Function 及 Closure

由於接下來會大量使用 closure，在這邊先為不習慣 functional programming 的讀者做簡單的介紹。

Lua 中的函式屬於 first-class object，它就像數字或字串那樣，你可以把函式存在變數當中、當作參數傳進另一個函式、或是當成其它函式的回傳值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
local foo = function(n)     -- 把函式存在 foo 這個變數中
    if n%7 == 0 then        -- 檢查 n 是不是 7 的倍數
        print(n .." is dividable by 7.")
    end
end
foo(14)                     -- 印出 14 is dividable by 7.
function check_between(min, max, checker)
    for i=min, max do       -- 把 min 到 max 之間的數
        checker(i)          -- 都代入 checker 函式檢查
    end
end
check_between(10, 100, foo) -- 列出 10 到 100 之間所有 7 的倍數

把函式當作回傳值則頗為微妙，因為這個回傳的函式可以存取上一層的區域變數，比如說以下這個例子：

1
2
3
4
5
6
7
8
9
10
11
12
function make_counter()
    local c = 0
    return function()
        c = c + 1
        print("c = " .. c)
    end
end
local counter = make_counter()
counter()           -- 印出 c = 1
counter()           -- 印出 c = 2
counter()           -- 印出 c = 3

若依照一般對 C 語言的認知，區域變數在函式結束後就會消失，因此在回傳的函式中使用它似乎會造成非法記憶體存取。但在 functional programming 之中，編譯器會偵測出這類的外層變數（稱之為 upvalue），並且把它綁在回傳值上。因此在呼叫 counter() 時，c 這個變數就彷彿全域變數般會永久存在，但在 make_counter() 外部卻又看不到。

這樣的語言特性稱之為 closure，而這也是 Lua iterator 的基礎。

Lua 中的 Iterator

Iterator 需要儲存目前巡訪容器的狀態，比如說陣列的 index。C++ 或 Java 這類語言通常把這些狀態包成物件，但 Lua 則是包裝在 closure 之中：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
function array_iterator(array)
    local i = 0                 -- 封裝在 iterator 中的 index
    local n = #array            -- 取得陣列長度
    return function()
        i = i + 1               -- 指向下一個元素
        if i <= n then          -- Lua 陣列是 1 開始的
            return array[i]
        else
            return nil
        end
    end
end
it = array_iterator({1, 1, 2, 3, 5})
local e = it()
while e ~= nil do
    print(e)
    e = it()
end

Lua 中的 iterator 是個包含 upvalue 的 closure，每次呼叫時會回傳下一個元素，直到結束時回傳 nil 為止。上述的 while 迴圈看起來不是很直觀，因此 Lua 提供了以下的 syntactic sugar：

1
2
3
for e in array_iterator({1, 1, 2, 3, 5}) do
    print(e)
end

Lua 會自動把它轉化成 while 迴圈的形式。

使用 Coroutine 巡訪陣列

在上一篇 coroutine 的簡介中有提到，coroutine 可以視為「可中斷及繼續執行的函式」，同時又能用 coroutine.yield() 來傳遞資料。因此我們可以把巡訪容器這件事寫成 coroutine，以 yield 來回傳容器中的元素，這麼一來就達成了 iterator 的目標：把巡訪容器與操作元素的邏輯分離開。

使用 coroutine 改寫上述的 array_iterator() 會變成下面這樣子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
function array_iterator2(array)
    local function visit()
        for i=1, #array do
            coroutine.yield(array[i])
        end
    end
    local co = coroutine.create(visit)
    return function()
        local status, value = coroutine.resume(co)
        if status then
            return value
        else
            return nil
        end
    end
end

現在這個 iterator 中只夾帶了 co 這個 upvalue。每次在 iterator 前進到下一個元素時，它只是重覆地呼叫 coroutine.resume() 讓 coroutine 往下執行，並取得其中使用 coroutine.yield 所傳回的元素值。

看起來好像比原來的版本更複雜了？我們來試試不同行為的 iterator。

Shuffle Iterator

想像一個特別的 iterator，它同樣會巡訪整個陣列，但卻會先輸出奇數索引上的元素，再輸出偶數索引上的元素。

1
2
3
4
for e in shuffle_iterator({1, 2, 3, 4, 5, 6})
do
    print(e)    -- 輸出 1 3 5 2 4 6
end

製作這種 iterator 並不難，只是不太容易讓人理解：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
function shuffle_iterator(array)
    local i = -1
    local n = #array
    local even_mode = false
    return function()
        i = i + 2
        if i > n then
            if even_mode or n < 2 then
                return nil
            else
                i = 2
                even_mode = true
            end
        end
        return array[i]
    end
end

even_mode 是用來儲存這個 iterator 是否走進了偶數索引區的狀態。儘管這個 iterator 並不難，但若利用 coroutine 會更加直覺：

1
2
3
4
5
6
7
8
9
10
11
function shuffle_iterator2(array)
    local function visit()
        for i=1,#array,2 do
            coroutine.yield(array[i])
        end
        for i=2,#array,2 do
            coroutine.yield(array[i])
        end
    end
    return coroutine.wrap(visit)
end

我使用了 coroutine.wrap()，這個 Lua 內建函式做的事和我們之前做的一樣：把 coroutine 包裝在 closure 之中，每次呼叫這個 closure 時，都等於呼叫 coroutine.resume()。

使用 coroutine 的寫法，彷彿就只是單純使用迴圈把內容一一印出來，不同的地方只在於把 print() 改成 coroutine.yield() 而已。這也意味著只要我們寫一份把容器內容印出來的程式碼，它就可以馬上改寫成 iterator。

我們來看看更複雜的容器。

二元樹的例子

現在我們的任務是製作二元樹的 iterator，以中序的方式巡訪。Lua 中的二元樹可以用 table 來表示：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
local binary_tree = {
    data = 5,
    left = {
        data = 3,
        left =  { data = 1 },
        right = { data = 4 }
    },
    right = {
        data = 9,
        left = {
            data = 7,
            left =  { data = 5.5 },
            right = { data = 7.4}
        },
        right = { data = 11 }
    }
}
這顆樹大概長這個樣子

製作二元樹的 iterator 不算是個簡單 (trivial) 的工作，但前面提到：只要知道怎麼把內容印出來，就可以把這個過程利用 coroutine 轉換成 iterator。所以我們先用看起來最容易的方法，也就是遞迴，來把二元樹印出來：

1
2
3
4
5
6
7
8
9
10
11
function print_inorder(node)
    if node.left ~= nil then
        print_inorder(node.left)  -- 巡訪左邊的 subtree
    end
    print(node.data)              -- 印出這個 node 的資料
    if node.right ~= nil then
        print_inorder(node.right) -- 巡訪右邊的 subtree
    end
end

改成 coroutine iterator 就只是把它照抄一遍：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
function tree_iterator(root)
    local function visit_inorder(node)
        if node.left ~= nil then
            visit_inorder(node.left)
        end
        coroutine.yield(node.data)
        if node.right ~= nil then
            visit_inorder(node.right)
        end
    end
    return coroutine.wrap(
        function() visit_inorder(root) end
    )
end
-- 計算元素總和
local sum = 0
for e in tree_iterator(binary_tree)
do
    sum = sum + e
end

如果不使用 coroutine，就相當於把遞迴的演算法改寫成非遞迴的形式，而這樣的改寫通常需要 stack 的協助，而導致結果不易理解。以下是改寫成 iterator 的版本：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
function tree_iterator2(root)
    local node_stack = {}
    local function push_left_subtree(node)
        while node ~= nil do
            table.insert(node_stack, node)
            node = node.left
        end
    end
    push_left_subtree(root)
    return function()
        if #node_stack == 0 then
            return nil
        end
        local node = table.remove(node_stack)
        push_left_subtree(node.right)
        return node.data
    end
end

使用 coroutine 實作 iterator 具有程式碼簡單易懂的特性，然而這也並非沒有缺點。因為 coroutine 內部保留了 call stack，通常它會比原本的 iterator 占用更多記憶體資源。

在下一篇文章，我會以遊戲 UI 作為範例，介紹 coroutine 的另一個應用方式。

分享到 評論