aoxu
2/14/2015 - 2:59 AM

inject.lua

local FUNC_TEMP=[[
local $ARGS
return function(...)
$SOURCE
end,
function()
return {$LOCALS}
end
]]

local temp = {}
local function wrap_locals(source, level)
	level = level + 3
	local f = debug.getinfo(level,"f").func
	if f == nil then
		return
	end
	local uv = {}
	local locals = {}
	local uv_id = {}
	local local_id = {}
	local i = 1
	while true do
		local name, value = debug.getlocal(level, i)
		if name == nil then
			break
		end
		if name:byte() ~= 40 then	-- '('
			table.insert(uv, name)
			table.insert(locals, ("[%d]=%s,"):format(i,name))
			local_id[name] = value
		end
		i = i + 1
	end
	local i = 1
	while true do
		local name = debug.getupvalue(f, i)
		if name == nil then
			break
		end
		uv_id[name] = i
		table.insert(uv, name)
		i = i + 1
	end
	temp.ARGS = table.concat(uv, ",")
	temp.SOURCE = source
	temp.LOCALS = table.concat(locals)
	local full_source = FUNC_TEMP:gsub("%$(%w+)",temp)
--	print(full_source)
	local loader, err = load(full_source)
	if loader == nil then
		return nil, err
	end
	local func, update = loader()
	-- join func's upvalues
	local i = 1
	while true do
		local name = debug.getupvalue(func, i)
		if name == nil then
			break
		end
		local local_value = local_id[name]
		if local_value then
			debug.setupvalue(func, i, local_value)
		end
		local upvalue_id = uv_id[name]
		if upvalue_id then
			debug.upvaluejoin(func, i, f, upvalue_id)
		end
		i=i+1
	end
	local vararg, v = debug.getlocal(level, -1)
	if vararg then
		local vargs = { v }
		local i = 2
		while true do
			local vararg,v = debug.getlocal(level, -i)
			if vararg then
				vargs[i] = v
			else
				break
			end
			i=i+1
		end
		return func, update, table.unpack(vargs)
	else
		return func, update
	end
end

local function exec(level, func, update, ...)
	if func == nil then
		return update
	end
	local rets = table.pack(pcall(func, ...))
	if rets[1] then
		local needupdate = update()
		for k,v in pairs(needupdate) do
			debug.setlocal(level,k,v)
		end
		return table.unpack(rets, 2, rets.n)
	else
		return rets[2]
	end
end

function run (source, level)
	level = level or 0
	return exec(level+2, wrap_locals(source, level))
end

----------------------------test-----------------

local uv = 2

function f(...)
	local a,b = 1,uv
	print("_ENV ===>", _ENV)
	print("a,b ====>", a,b)
	run[[
		print "=== inject code ==="
		print("\t_ENV ===>", _ENV)
		print("\ta,b,uv ===>", a,b,uv)
		print("\t... ", ...)
		a,b = b,a
		uv = 3
		print "==== inject end ==="
	]]
	print("a,b ====>", a,b)
end

f("Hello","world")

print("uv=",uv)

--[[ output

_ENV ===>       table: 0000000000266de0
a,b ====>       1       2
=== inject code ===
_ENV ===>       table: 0000000000266de0
a,b,uv ===>     1       2       2
...     Hello   world
==== inject end ===
a,b ====>       2       1
uv=     3
]]