🎓 A1Learner
🔍
0 of 52 lessons completed

Lua Iterators over Trees

A tree iterator is an iterator that walks a nested table depth-first, visiting a node before its children. It keeps an explicit stack of nodes that still need visiting.

A stack iterator visits a small tree

The iterator starts with the root on the stack. It visits the root, then its children in order.

local function make_tree_iter(root)
  local stack = {root}
  return function()
    while #stack > 0 do
      local node = table.remove(stack)
      for i = #node.children, 1, -1 do
        stack[#stack + 1] = node.children[i]
      end
      return node.value
    end
    return nil
  end
end
local tree = {
  value = "root",
  children = {
    {value = "left", children = {}},
    {value = "right", children = {}}
  }
}
for v in make_tree_iter(tree) do
  print(v)
end

Depth-first finishes one branch before the next

Child a is walked fully before sibling b starts. That order proves the walk is depth-first.

local function make_tree_iter(root)
  local stack = {root}
  return function()
    while #stack > 0 do
      local node = table.remove(stack)
      for i = #node.children, 1, -1 do
        stack[#stack + 1] = node.children[i]
      end
      return node.value
    end
    return nil
  end
end
local tree = {
  value = "root",
  children = {
    {
      value = "a",
      children = {
        {value = "a1", children = {}},
        {value = "a2", children = {}}
      }
    },
    {
      value = "b",
      children = {
        {value = "b1", children = {}}
      }
    }
  }
}
for v in make_tree_iter(tree) do
  print(v)
end

The stack lets the walk pause and resume

Manual calls take the first two values. The for loop then resumes from the same stack and prints the rest.

local function make_tree_iter(root)
  local stack = {root}
  return function()
    while #stack > 0 do
      local node = table.remove(stack)
      for i = #node.children, 1, -1 do
        stack[#stack + 1] = node.children[i]
      end
      return node.value
    end
    return nil
  end
end
local tree = {
  value = "root",
  children = {
    {
      value = "a",
      children = {
        {value = "a1", children = {}},
        {value = "a2", children = {}}
      }
    },
    {
      value = "b",
      children = {
        {value = "b1", children = {}}
      }
    }
  }
}
local it = make_tree_iter(tree)
print(it())
print(it())
for v in it do
  print(v)
end

Example

local function make_tree_iter(root)
  local stack = {root}
  return function()
    while #stack > 0 do
      local node = table.remove(stack)
      for i = #node.children, 1, -1 do
        stack[#stack + 1] = node.children[i]
      end
      return node.value
    end
    return nil
  end
end
local tree = {
  value = "root",
  children = {
    {
      value = "docs",
      children = {
        {value = "readme", children = {}},
        {value = "guide", children = {}}
      }
    },
    {
      value = "src",
      children = {
        {value = "main", children = {}}
      }
    }
  }
}
local count = 0
for v in make_tree_iter(tree) do
  print(v)
  count = count + 1
end
print(count)

Key points

  • Node: a table that holds a value and a list of child nodes.
  • Depth-first order: a walk that visits a node, then walks each child fully before moving to the next sibling.
  • Stack: a list used as a pile where new nodes go on top and the next node comes off the top.
  • Explicit stack: a stack held in a local variable lets the iterator pause between calls and resume where it stopped.
  • Leaf node: a node with an empty children list, so the walk moves back to its siblings after it.

In this example, the iterator holds its stack in a local variable and visits the tree depth-first. It prints root, then docs and both of its children, then src and its child. The final count of 6 shows that every node was visited once.

Note

pushing children in reverse order makes them come off the stack in their original order.

Try it Yourself

Edit the code, then press Run.

Running is simulated

Exercise

A tree iterator visits root, then child a with children a1 and a2, then child b. What order does depth-first print?