luatraverse.lua

-------------------------------------------------------------------------------

-- This module implements a function that traverses all live objects.

-- You can implement your own function to pass as a parameter of traverse

-- and give you the information you want. As an example we have implemented

-- countreferences and findallpaths

--

-- Alexandra Barros - 2006.03.15

-------------------------------------------------------------------------------

--luacheck: ignore 143/debug 113/getfenv

local List = {}

function List.new ()

	return {first = 0, last = -1}

end

function List.push (list, value)

	local last = list.last + 1

    list.last = last

    list[last] = value

end

function List.pop (list)

    local first = list.first

    if first > list.last then error("list is empty") end

    local value = list[first]

    list[first] = nil

    list.first = first + 1

    return value

end

function List.isempty (list)

	return list.first > list.last

end

local traverse, edge;

local types = {};

local get_metatable, get_environment, get_registry, get_locals, upvalues

-- select implementation of get_metatable depending on available API

if type( debug ) == "table" and

   type( debug.getmetatable ) == "function" then

  local get_mt = debug.getmetatable

  function get_metatable( val, enabled )

    if enabled then return get_mt( val ) end

  end

elseif type( getmetatable ) == "function" then

  function get_metatable( val, enabled )

    if enabled then return getmetatable( val ) end

  end

else

  function get_metatable() end

end

-- select implementation of get_environment depending on available API

if type( debug ) == "table" and

   type( debug.getfenv ) == "function" then

  local get_fe = debug.getfenv

  function get_environment( val, n, enabled )

    if enabled and n == 1 then

      local uv = get_fe( val )

      return uv, uv ~= nil

    end

    return nil, false

  end

elseif type( debug ) == "table" and

       type( debug.getuservalue ) == "function" then

  local get_uv = debug.getuservalue

  if _VERSION < "Lua 5.4" then

    local get_uv_noindex = get_uv

    function get_uv( val, n )

      if n == 1 then

        local uv = get_uv_noindex( val )

        return uv, uv ~= nil

      end

      return nil, false

    end

  end

  function get_environment( val, n, enabled )

    if enabled then

      -- getuservalue in Lua 5.2 throws on light userdata!

      local ok, res1, res2 = pcall( get_uv, val, n )

      if ok then return res1, res2 end

      return nil, false

    end

  end

elseif type( getfenv ) == "function" then

  function get_environment( val, n, enabled )

    if enabled and n == 1 and type( val ) == "function" then

      local uv = getfenv( val )

      return uv, uv ~= nil

    end

    return nil, false

  end

else

  function get_environment() return nil, false end

end

-- select implementation of get_registry

if type( debug ) == "table" and

   type( debug.getregistry ) == "function" then

  get_registry = debug.getregistry

else

  function get_registry() end

end

-- select implementation of get_locals

-- locs = {

--   {

--      name = "example_function";

--      func = <function>;

--      [1] = { "local_name", <local value> };

--      [2] = ...;

--   };

-- }

if type( debug ) == "table" and

   type( debug.getinfo ) == "function" and

   type( debug.getlocal ) == "function" then

  local getinfo, getlocal = debug.getinfo, debug.getlocal

  local function getinfo_nothread( _, func, what )

    return getinfo( func, what )

  end

  local function getlocal_nothread( _, level, loc )

    return getlocal( level, loc )

  end

  function get_locals( thread, enabled )

    if enabled then

      local locs = {}

      local start = 1

      local gi, gl = getinfo, getlocal

      if not thread then

        gi, gl = getinfo_nothread, getlocal_nothread

      end

      local info, i = gi( thread, 0, "nf" ), 0

      while info do

        local t = { name = info.name, func = info.func }

        local j, n,v = 1, gl( thread, i, 1 )

        while n ~= nil do

          t[ j ] = { n, v }

          j = j + 1

          n,v = gl( thread, i, j )

        end

        i = i + 1

        locs[ i ] = t

        -- We skip the currently-executing traverse() function

        -- and anything that is only referenced through it

        if info.func == traverse then start = i+1 end

        info = gi( thread, i, "nf" )

      end

      return locs, start

    end

  end

else

  function get_locals() end

end

-- select implementation of upvalues depending on available API

local function dummy_iter() end

if type( debug ) == "table" and

   type( debug.getupvalue ) == "function" then

  local get_up, uv_iter = debug.getupvalue

  if _VERSION == "Lua 5.1" then

    function uv_iter( state )

      local name, uv = get_up( state.value, state.n )

      state.n = state.n + 1

      return name, uv, nil

    end

  else -- Lua 5.2 (and later) mixes upvalues and environments

    local get_upid

    if type( debug.upvalueid ) == "function" then

      get_upid = debug.upvalueid

    end

    function uv_iter( state )

      local name, uv = get_up( state.value, state.n )

      state.n = state.n + 1

      if name == "_ENV" and not state.show_env then

        return uv_iter( state )

      end

      local id = nil

      if get_upid ~= nil and name ~= nil then

        id = get_upid( state.value, state.n - 1 )

      end

      return name, uv, id

    end

  end

  function upvalues( val, enabled, show_env )

    if enabled then

      return uv_iter, { value = val, n = 1, show_env = show_env }

    else

      return dummy_iter

    end

  end

else

  function upvalues()

    return dummy_iter

  end

end

local function handle_metatable(env, obj)

	local mtable = get_metatable(obj)

	if mtable ~= nil then

		edge(env, obj, mtable, "ismetatable", nil)

	end

end

local function handle_environment(env, obj)

	local n = 0

	repeat

		n = n + 1

		local fenv, has_env = get_environment(obj, n)

		if has_env and fenv ~= nil then

			edge(env, obj, fenv, "environment", nil)

		end

	until not has_env

end

local function handle_locals(env, obj)

	local stack = get_locals(obj);

	if not stack then return; end

	for _, frame in ipairs(stack) do

		for local_entry in ipairs(frame) do

			local local_name, local_value = local_entry[1], local_entry[2];

			edge(env, obj, local_name, "isname", nil);

			edge(env, obj, local_value, "local", local_name);

		end

	end

end

-- Main function

-- 'funcs' is a table that contains a funcation for every lua type and also the

-- function edge edge (traverseedge).

function traverse(funcs, ignoreobjs, seenobjs)

	-- The keys of the marked table are the objects (for example, table: 00442330).

	-- The value of each key is true if the object has been found and false

	-- otherwise.

	local env = {marked = seenobjs or {}, list=List.new(), funcs=funcs}

	if ignoreobjs then

		for i=1, #ignoreobjs do

			env.marked[ignoreobjs[i]] = true

		end

	end

	env.marked["traverse"] = true

	env.marked[traverse] = true

	-- marks and inserts on the list

	edge(env, nil, "_G", "isname", nil)

	edge(env, nil, _G, "value", "_G")

	-- traverse the registry

	local registry = get_registry();

	if registry ~= nil then

		edge(env, nil, registry, "value", "REGISTRY");

	end

	-- traverses the active thread

	-- inserts the local variables

	-- interates over the function on the stack, starting from the one that

	-- called traverse

	handle_locals(env, nil);

	while not List.isempty(env.list) do

		local obj = List.pop(env.list)

		local t = type(obj)

		types[t](env, obj)

	end

end

function types.table(env, obj)

	local f = env.funcs.table

	if f then f(obj) end

	for key, value in pairs(obj) do

		edge(env, obj, key, "key", nil)

		edge(env, obj, value, "value", key)

	end

	handle_metatable(env, obj)

end

function types.string(env, obj)

	local f = env.funcs.string

	if f then f(obj) end

end

function types.userdata(env, obj)

	local f = env.funcs.userdata

	if f then f(obj) end

	handle_metatable(env, obj)

	handle_environment(env, obj)

end

types["function"] = function (env, obj)

	local f = env.funcs.func

	if f then f(obj) end

	handle_environment(env, obj)

	-- handle upvalues

	for name, value in upvalues(obj, true, true) do

		edge(env, obj, name, "isname", nil)

		edge(env, obj, value, "upvalue", name)

	end

end

function types.thread(env, t)

	local f = env.funcs.thread

	if f then f(t) end

	handle_environment(env, t)

	handle_locals(env, t);

end

-- 'how' is a string that identifies the content of 'to' and 'value':

-- 		if 'how' is "key", then 'to' is a key and 'name' is nil.

-- 		if 'how' is "value", then 'to' is an object and 'name' is the name of the

--		key.

function edge(env, from, to, how, name)

	local t = type(to)

	if to and (t~="boolean") and (t~="number") and (t~="new") then

		-- If the destination object has not been found yet

		if not env.marked[to] then

			env.marked[to] = true

			List.push(env.list, to) -- puts on the list to be traversed

		end

		local f = env.funcs.edge

		if f then f(from, to, how, name) end

	end

end

return {

	traverse = traverse;

};