reuse Lua's coroutine

gistfile1.txt

--coroutine pool
local freelist={};

function find_free(f)
	if #freelist > 0 then
		local free = freelist[#freelist]
		table.remove(freelist)
		return {thread = free,func = f,status="create"}
	end
end

local coroutine_create = coroutine.create
local coroutine_resume = coroutine.resume
local coroutine_status = coroutine.status
local coroutine_running = coroutine.running

local running = {__mode = "kv"} --must use kv, for some co may not resume to end.

--for debug
g_all_co = {}
function trace_co()
	for i,t in ipairs(g_all_co) do
		print(i,t.alloc,t.status,coroutine_status(t.thread),t.thread)
	end
	print(string.format("run %d free %d alloc %d call %d",table.size(running),#freelist,g_alloc_co,g_call_co))
	
	print("running...")
	for i,t in pairs(running) do 
		if i == "__mode" then
		else
			print(i,t.alloc,t.status,coroutine_status(t.thread),t.thread)
		end
	end
end

g_alloc_co = 0
g_call_co = 0
function create_co(f)
	g_alloc_co = g_alloc_co + 1
	local tb = {func = f,status="create",alloc=g_alloc_co}
	tb.thread = coroutine_create(
	function (...)
		local arg = {...}
		while true do
			assert(type(arg[2]) == "function" and type(arg[1]) == "table")
			local intb = arg[1]
			local thread = coroutine_running()
			running[thread] = intb
			arg[2](unpack(arg,3))
			intb.status = "dead"
			running[thread] = nil
			table.insert(freelist,thread)
			arg = {coroutine.yield("co_end")}
		end
	end)
	return tb
end

function coroutine.create(f)
	g_call_co = g_call_co + 1
	local tb = find_free(f) or create_co(f)	
	tb.call = g_call_co
	assert(tb.status == "create")
	--table.insert(g_all_co,tb) --uncomment it for debug
	return tb
end

function _coroutine_resume(co,...)
	if type(co) == "thread" then
		return coroutine_resume(co,...)
	end
	assert(type(co) == "table",type(co))
	if co.status == "create" then
		co.status = "running"
		return coroutine_resume(co.thread,co,co.func,...)
	elseif co.status == "running" then
		return coroutine_resume(co.thread,...)
	else --dead
		error("error!!! co.status = "..co.status)
	end
end

function coroutine.resume(co,...)
	local isok,msg = _coroutine_resume(co,...)
	if not isok then
		print("coroutine.resume error!")
	end
	return isok,msg
end

function coroutine.status(co)
	if type(co) == "thread" then
		return coroutine_status(co)
	end
	assert(type(co) == "table")
	if co.status == "dead" or co.status == "create" then
		return "dead"
	end
	return coroutine_status(co.thread)
end

function coroutine.running()
	return running[coroutine_running()]
end

--help func
function co_wrap(f)
	local free = coroutine.create(f)
	return function(...)
		return coroutine.resume(free,...)
	end
end

function co_wrap1(f)
	local free = coroutine.create(f)
	return function(...)
		local isok, msg = coroutine.resume(free,...)
		return free,isok, msg
	end
end

function co_run(f,...)
	local free = coroutine.create(f)
	return coroutine.resume(free,...)
end

function co_run1(f,...)
	local free = coroutine.create(f)
	return free,coroutine.resume(free,...)
end

--test
local test_co1
function t1(p1)
	test_co1 = coroutine.running()
	print("inco t1 1",p1)
	local rt = coroutine.yield("t1 y1")
	print("rt",rt)
	assert(rt == "r1")
	print("inco t1 end",rt)
end

local test_co2
function t2(p1,p2)
	test_co2 = coroutine.running()
	print("inco t2 1",p1,p2)
	local rt = coroutine.yield("t2 y2")
	assert(rt == "r2")
	print("inco t2 2",rt)
	rt = coroutine.yield("t2 y2")
	assert(rt == "r2")
	print("inco t2 end",rt)
end

local test_co3
function t3(p1,p2)
	test_co3 = coroutine.running()
	print("inco t3 1",p1,p2)
	local rt = coroutine.yield("t3 y3")
	assert(rt == "r3")
	error("inco t3")
	rt = coroutine.yield("t3 y3")
	assert(rt == "r3")
	print("inco t3 end",rt)
end

function test()
	freelist={}
	print("#freelist",#freelist)
	assert(#freelist == 0)
	local co,ok,rt = co_run1(t1,"p1")
	print('run(t1,"p1")',ok,rt)
	ok,rt = coroutine.resume(co,"r1")
	print('coroutine.resume(test_co1,"r1")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 1)

	co,ok,rt = co_run1(t2,"p1","p2")
	assert(test_co1.thread == test_co2.thread and test_co1 ~= test_co2)
	--print("freelist",#freelist)
	print('run(t2,"p1","p2")',ok,rt)
	ok,rt = coroutine.resume(co,"r2")
	print('coroutine.resume(test_co2,"r2")',ok,rt)
	ok,rt = coroutine.resume(co,"r2")
	print('coroutine.resume(test_co2,"r2")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 1)
end

function test_wrap()
	freelist={}
	print("#freelist",#freelist)
	assert(#freelist == 0)
	local f1 = co_wrap(t1)
	local ok,rt = f1("p1")
	print('f1("p1")',ok,rt)
	ok,rt = coroutine.resume(test_co1,"r1")
	print('coroutine.resume(test_co1,"r1")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 1)

	local f2 = co_wrap(t2)
	ok,rt = f2("p1","p2")
	assert(test_co1.thread == test_co2.thread and test_co1 ~= test_co2)
	--print("freelist",#freelist)
	print('f2("p1","p2")',ok,rt)
	ok,rt = coroutine.resume(test_co2,"r2")
	print('coroutine.resume(test_co2,"r2")',ok,rt)
	ok,rt = coroutine.resume(test_co2,"r2")
	print('coroutine.resume(test_co2,"r2")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 1)	
end

function test_error()
	freelist={}
	print("#freelist",#freelist)
	assert(#freelist == 0)
	local f3 = co_wrap(t3)
	print("local f3 = wrap(t3)")
	ok,rt = f3("p1","p2")
	print('f3("p1","p2")',ok,rt)
	ok,rt = coroutine.resume(test_co3,"r3")
	print('coroutine.resume(test_co3,"r3")',ok,rt)
	ok,rt = coroutine.resume(test_co3,"r3")
	print('coroutine.resume(test_co3,"r3")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 0)	
end

function test_error_tag()
	freelist={}
	print("#freelist",#freelist)
	assert(#freelist == 0)
	local f1 = co_wrap1(t1)
	local co,ok,rt = f1("p1")
	print('f1("p1")',ok,rt)
	ok,rt = coroutine.resume(co,"r1")
	print('coroutine.resume(test_co1,"r1")',ok,rt)
	assert(#freelist == 1)
	ok,rt = coroutine.resume(co,"r1") --CO_TAG will used here!
	print('coroutine.resume(test_co1,"r1")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 1)
end

function test_error_oldnew()
	freelist={}
	print("#freelist",#freelist)
	assert(#freelist == 0)
	local f1 = co_wrap1(t1)
	local co1,ok,rt = f1("p1")
	print('f1("p1")',ok,rt)
	ok,rt = coroutine.resume(co1,"r1")
	print('coroutine.resume(test_co1,"r1")',ok,rt)
	print("#freelist",#freelist)
	assert(#freelist == 1)

	local f2 = co_wrap1(t2)
	print("local f2 = wrap(t2)")
	local co2
	co2,ok,rt = f2("p1","p2")
	assert(test_co1 == co1 and test_co2 == co2)
	assert(test_co1.thread == test_co2.thread and test_co1 ~= test_co2 and co1 ~= co2)
	print("freelist",#freelist)
	assert(#freelist == 0)	
	print('f2("p1","p2")',ok,rt)
	ok,rt = coroutine.resume(co2,"r2")
	print('coroutine.resume(test_co2,"r2")',ok,rt)

	ok,rt = coroutine.resume(co1,"r1")
	print('coroutine.resume(test_co1,"r1")',ok,rt)
	print("#freelist",#freelist)
end

--test_error_oldnew()