如何在Lua中编写支持for循环的二叉树迭代器？

需求说明 #

现有基于回调函数的二叉树前序遍历实现，希望改造为支持for node in tree:visit() do ... end形式的迭代器，需满足：

• 支持任意深度的增量遍历，禁止预生成完整节点列表后用pairs遍历
• 实现方案简洁，避免过度复杂

原回调式遍历代码如下：

function klass(s,    t) t={ako=s}; t.__index=t; return t end

NODE=klass"NODE"
function NODE.new(x,l,r) return setmetatable( {x=x, left=l, right=r},NODE) end

function NODE:visit(fun,  lvl)
  lvl=lvl or 0
  fun(lvl,self)
  for _,kid in pairs{self.left, self.right} do 
    if kid then kid:visit(fun, lvl+1) end end end

--- 测试生成树 -----------------------------
function eg(x,stop,     node)
  node = NODE.new(x)
  if x < stop/2 then
    node.left = eg(2*x,stop)
    node.right = eg(2*x+1,stop) end
  return node end 

tree=eg(1,32)
tree:visit( function(lvl,node) print(('|.. '):rep(lvl)..node.x) end )

解决方案：基于协程实现迭代器 #

利用Lua的**协程（coroutine）**特性，将递归遍历逻辑封装到协程中，通过yield逐步返回节点，实现增量遍历。修改NODE:visit方法，使其返回符合for循环要求的迭代器函数：

function klass(s,    t) t={ako=s}; t.__index=t; return t end

NODE=klass"NODE"
function NODE.new(x,l,r) return setmetatable( {x=x, left=l, right=r},NODE) end

-- 改造为迭代器版本的visit方法
function NODE:visit()
  -- 定义递归遍历的协程函数
  local function traverse(node, lvl)
    lvl = lvl or 0
    coroutine.yield(lvl, node)  -- 暂停协程，返回当前节点和层级
    for _, kid in pairs{node.left, node.right} do
      if kid then traverse(kid, lvl + 1) end
    end
  end

  -- 返回迭代器函数：每次调用恢复协程，获取下一个节点
  return coroutine.wrap(function()
    traverse(self, 0)
  end)
end

--- 测试生成树 -----------------------------
function eg(x,stop,     node)
  node = NODE.new(x)
  if x < stop/2 then
    node.left = eg(2*x,stop)
    node.right = eg(2*x+1,stop) end
  return node end 

tree=eg(1,32)

-- 使用for循环遍历
for lvl, node in tree:visit() do
  print(('|.. '):rep(lvl)..node.x)
end

代码说明 #

1. 协程封装遍历逻辑：将原回调函数的fun(lvl,self)替换为coroutine.yield(lvl, node)，让遍历过程可以暂停并返回当前节点。
2. 迭代器函数生成：通过coroutine.wrap创建协程包装器，这个包装器就是for循环需要的迭代器——每次调用它会恢复协程执行，直到下一个yield或遍历结束。
3. 增量遍历特性：遍历过程中不会一次性生成所有节点，每次for循环迭代才会计算并返回下一个节点，完美支持任意深度的树结构。

运行效果 #

执行上述代码后，输出结果与原回调式遍历完全一致，实现了预期的for循环遍历形式。

内容的提问来源于stack exchange，提问作者Tim Menzies