lua coroutine for iterator

背景

前面的文章演示了使用闭包函数实现 状态的迭代器。

本文演示使用 coroutine来产生迭代器的例子。

coroutine迭代器例子 -- 遍历二叉树

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 }
    }
}




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

-- 計算元素總和
for e in tree_iterator(binary_tree)
do
    print(e)
end

LOG:

>lua -e "io.stdout:setvbuf 'no'" "luatest.lua"
1
3
4
5
5.5
7
7.4
9
11
>Exit code: 0

coroutine迭代器例子 -- 字符串分段

local function SegIter(instr, segSize)
    local strIndex = 1

    local function GetSegStr(str)

        local segStart = strIndex
        local segEnd = strIndex + segSize - 1
        local strseg = string.sub(str, segStart, segEnd)

        if #strseg > 0 then
            strIndex = strIndex + segSize

            -- return the current element of the iterator
            coroutine.yield(strseg)

            GetSegStr(str)
        end
    end

    local co = coroutine.create(function()
        GetSegStr(instr)
    end)

    -- iterator
    return function()
        local code, res = coroutine.resume(co)
        if not code then return nil end
        return res
    end
end


for element in SegIter("66666666666666666", 2)
do
   print(element)
end

LOG:

>lua -e "io.stdout:setvbuf 'no'" "luatest.lua"
66
66
66
66
66
66
66
66
6
>Exit code: 0

注意两个协程执行的主体函数都有递归调用的影子。 这个方法， 保证每个递归中的 yield 都能够执行， 知道最后yield执行完毕后， 函数执行完毕， 返回nil。