Tree.lua

Tree.lua

local Tree = {}
local Tree_mt = {__index = Tree}

function Tree.New(value, parent, index)
    local self = setmetatable({}, Tree_mt)
    self.value = value
    self.parent = parent
    self.index = index or 0
    return self
end

function Tree:Set(value)
    self.value = value
end

function Tree:IsRoot()
    return self.index == 0
end

function Tree:GetRoot()
    local node = self
    while node.parent do
        node = node.parent
    end
    return node
end

function Tree:GetParent()
    return self.parent
end

function Tree:AddChild(value)
    local i = #self - 1
    self[i] = Tree.New(value, self, i)
    return self[i]
end

function Tree:GetChild(index)
    return self[index]
end

function Tree:FindChild(value)
    if #self == 0 then
        return
    end
    
    if self.value == value then
        return self
    end
    
    for _, node in ipairs(self) do
        if node.value == value then
            return node
        end
        
        local child = node:FindChild(value)
        if child then
            return child
        end
    end
end

function Tree:IterativeChildren(callback)
    if #self == 0 then
        return
    end
    
    for _, node in ipairs(self) do
        callback(node)
        node:IterativeChildren(callback)
    end
end

return Tree