Unit Searcher

UnitSearcher.lua

-- Mapping of all root unit tokens in the game.
local UNIT_TOKENS = {
    ["player"] = true,
    ["target"] = true,
    ["mouseover"] = true,
    ["focus"] = true,
    ["pet"] = true,
    ["npc"] = true,
    ["vehicle"] = true,
};

-- Something-something hardcoded constants yadda yadda.
for i = 1, 40 do
    if i <= 3 then
        UNIT_TOKENS["arena" .. i] = true;
    end

    if i <= 4 then
        UNIT_TOKENS["party" .. i] = true;
        UNIT_TOKENS["partypet" .. i] = true;
    end

    if i <= 5 then
        UNIT_TOKENS["boss" .. i] = true;
    end

    UNIT_TOKENS["raid" .. i] = true;
    UNIT_TOKENS["raidpet" .. i] = true;
    UNIT_TOKENS["nameplate" .. i] = true;
end

-- Returns true if the given unit token is a "unittarget" type token.
local function IsDescendantUnitToken(unitToken)
    return not UNIT_TOKENS[unitToken];
end

-- Returns the descendant target unit token for a given parent unit token.
local function GetDescendantUnitToken(unitToken)
    return unitToken .. "target"; -- TODO: Could cache these.
end

-- Returns the root ancestor unit token from a descendant unit token.
local function GetRootAncestorUnitToken(unitToken)
    -- We'll cut 6 characters off the end of the token so long as the third
    -- character from the end of it is a "g". As the "target" unit exists
    -- as a valid root, we need to keep a minimum of 6 characters.
    local finish = #unitToken;
    while finish > 6 and unitToken:byte(finish - 2, finish - 2) == 0x67 do
        finish = finish - 6;
    end

    return unitToken:sub(1, finish); -- TODO: Could cache these.
end

-- Returns the next unit token from a mapping that exists ingame. If the
-- starting token has any descendant target units that exist, those will be
-- returned first as a priority.
local function NextExistingUnitToken(unitTokens, unitToken)
    -- If the current unit token is valid we should test its descendants
    -- before moving onto the next root token.
    if unitToken and UnitExists(unitToken) then
        local descendantToken = GetDescendantUnitToken(unitToken);
        if UnitExists(descendantToken) then
            return descendantToken;
        elseif IsDescendantUnitToken(unitToken) then
            unitToken = GetRootAncestorUnitToken(unitToken);
        end
    end

    -- Advance root unit tokens until we find something.
    repeat
        unitToken = next(unitTokens, unitToken);
        if UnitExists(unitToken) then
            return unitToken;
        end
    until not unitToken;
end

-- Returns an iterator for accessing all existing units from the given
-- mapping of unit tokens.
local function EnumerateExistingUnitTokens(unitTokens)
    return NextExistingUnitToken, unitTokens or UNIT_TOKENS, nil;
end

-- Returns a table of all existing units as a mapping of GUIDs to unit tokens.
local function GetExistingUnits()
    local units = {};

    -- unitToken is the existing unit token we're currently testing,
    -- rootToken is its root ancestor.
    local unitToken = NextExistingUnitToken(UNIT_TOKENS, nil);
    local rootToken = unitToken;
    while unitToken do
        local unitGUID = UnitGUID(unitToken);
        if not units[unitGUID] then
            -- This is our first time seeing the GUID; store it and descend.
            units[unitGUID] = unitToken;

            unitToken = NextExistingUnitToken(UNIT_TOKENS, unitToken);
            rootToken = UNIT_TOKENS[unitToken] and unitToken or rootToken;
        else
            -- We've seen this GUID before, skip to the next root token.
            unitToken = next(UNIT_TOKENS, rootToken);
            rootToken = unitToken;
        end
    end

    return units;
end

-- Usage:
print("Existing units:");
for unitToken in EnumerateExistingUnitTokens() do
    print(unitToken);
end

print("Unit Mapping:");
local units = GetExistingUnits();
for unitGUID, unitToken in pairs(units) do
    print(unitGUID, unitToken)
end