Test various ways to dynamically dispatch heterogeneous functions on heterogeneous objects

test_scheduling.jl

using BenchmarkTools

using FunctionWrappers: FunctionWrapper

# *** Agents

abstract type Agent
end

mutable struct Agent1 <: Agent
	x :: Int
end

mutable struct Agent2 <: Agent
	y :: Int
end

# *** "Simulation"

@noinline step1(a :: Agent1) = a.x += rand(1:2) 

@noinline step1(a :: Agent2) = a.y -= rand(1:2)

@noinline step2(a :: Agent1) = a.x += rand(1:4)

@noinline step2(a :: Agent2) = a.y -= rand(1:4)

# *** Functors

abstract type Action
end

struct Action1{AGENT} <: Action
	a :: AGENT
end

function (a::Action1{A})() where {A}
	step1(a.a)
end

struct Action2{AGENT} <: Action
	a :: AGENT
end

function (a::Action2{A})() where {A}
	step2(a.a)
end

# *** prepare agents

const na = 10_000

const arr1 = [Agent1(-i) for i=1:na]
const arr2 = [Agent2(i) for i=1:na]

function reset()
	for i in 1:length(arr1)
		arr1[i].x = -i
	end
	for i in 1:length(arr2)
		arr2[i].y = i
	end
end

# *** various ways to schedule actions

# static dispatch
function sched_static(a1, a2, n)
	agents1 = Agent1[]
	agents2 = Agent2[]
	action = Int[]

	# use some time on random selection as in the
	# other versions
	for i in 1:n
		if rand(1:2) == 1
			push!(agents1, rand(a1))
			push!(action, rand(1:2))
		else
			push!(agents2, rand(a2))
			push!(action, rand(1:2))
		end
	end

	i = 1
	for a in agents1
		action[i] == 1 ? step1(a) : step2(a)
		i += 1
	end
	for a in agents2
		action[i] == 1 ? step1(a) : step2(a)
		i += 1
	end
end


# plain generic function calls and abstract agent type
# use integer to select action
function sched_gen(a1, a2, n)
	agents = Agent[]
	action = Int[]

	for i in 1:n
		push!(agents, rand(1:2)==1 ? rand(a1) : rand(a2))
		push!(action, rand(1:2))
	end

	i = 1
	for a in agents
		action[i] == 1 ? step1(a) : step2(a)
		i += 1
	end
end


# abstract agents, generic functions, but functions stored directly
function sched_allgen(a1, a2, n)
	agents = Agent[]
	action = Function[]

	for i in 1:n
		push!(agents, rand(1:2)==1 ? rand(a1) : rand(a2))
		push!(action, rand(1:2)==1 ? step1 : step2)
	end

	i = 1
	for a in agents
		action[i](a)
		i += 1
	end
end


# use closures to store functions plus agents
function sched_plainclosure(a1, a2, n)
	action = Function[]

	for i in 1:n
		a = rand(1:2)==1 ? rand(a1) : rand(a2)
		push!(action, rand(1:2)==1 ? () -> step1(a) : () -> step2(a))
	end

	for f in action
		f()
	end
end


# closures, but explicitly typed
function sched_typedclos(a1, a2, n)
	action = Function[]

	for i in 1:n
		if rand(1:2)==1 
			aa1::Agent1 = rand(a1)
			push!(action, rand(1:2)==1 ? () -> step1(aa1) : () -> step2(aa1))
		else
			aa2::Agent2 = rand(a2)
			push!(action, rand(1:2)==1 ? () -> step1(aa2) : () -> step2(aa2))
		end
	end

	for f in action
		f()
	end
end


# manually dispatch on agent type and command tag
function sched_manual(a1, a2, n)
	agents = Agent[]
	action = Int[]

	for i in 1:n
		push!(agents, rand(1:2)==1 ? rand(a1) : rand(a2))
		push!(action, rand(1:2))
	end

	i = 1
	for a in agents
		if isa(a, Agent1)
			action[i] == 1 ? step1(a::Agent1) : step2(a::Agent1)
			i += 1
		elseif isa(a, Agent2)
			action[i] == 1 ? step1(a::Agent2) : step2(a::Agent2)
			i += 1
		else
			action[i] == 1 ? step1(a) : step2(a)
			i += 1
		end
	end
end


# store functions and agents in functor objects,
# *should* be identical to typed closures
function sched_functors(a1, a2, n)
	actions = Action[]

	for i in 1:n
		push!(actions, rand(1:2)==1 ? 
			(rand(1:2)==1 ? Action1{Agent1}(rand(a1)) : Action2{Agent1}(rand(a1))) :
			(rand(1:2)==1 ? Action1{Agent2}(rand(a2)) : Action2{Agent2}(rand(a2))))
	end

	for a in actions
		a()
	end
end

# *** a few vsrsions using FunctionWrapper

const FW = FunctionWrapper{Int, Tuple{}}

# FunctionWrapper + Functor
function sched_fw(a1, a2, n)
	actions = FW[]

	for i in 1:n
		push!(actions, rand(1:2)==1 ? 
			FW(rand(1:2)==1 ? Action1{Agent1}(rand(a1)) : Action2{Agent1}(rand(a1))) :
			FW(rand(1:2)==1 ? Action1{Agent2}(rand(a2)) : Action2{Agent2}(rand(a2))))
	end

	for a in actions
		a()
	end
end


# FunctionWrapper + closure
function sched_fwc(a1, a2, n)
	actions = FW[]

	for i in 1:n
		if rand(1:2) == 1
			ag1::Agent1 = rand(a1)
			push!(actions, rand(1:2) == 1 ? 
				FW(() -> step1(ag1::Agent1)) :
				FW(() -> step2(ag1::Agent1)))
		else
			ag2::Agent2 = rand(a2)
			push!(actions, rand(1:2) == 1 ? 
				FW(() -> step1(ag2::Agent2)) :
				FW(() -> step2(ag2::Agent2)))
		end
	end

	for a in actions
		a()
	end
end

# *** run everything

const ns = 100_000

sched_static(arr1, arr2, 1)
sched_gen(arr1, arr2, 1)
sched_allgen(arr1, arr2, 1)
sched_plainclosure(arr1, arr2, 1)
sched_typedclos(arr1, arr2, 1)
sched_manual(arr1, arr2, 1)
sched_functors(arr1, arr2, 1)
sched_fw(arr1, arr2, 1)
sched_fwc(arr1, arr2, 1)

reset()

println("static dispatch:")
@btime sched_static(arr1, arr2, ns)
reset()
println("generic function, command array:")
@btime sched_gen(arr1, arr2, ns)
reset()
println("generic function:")
@btime sched_allgen(arr1, arr2, ns)
reset()
println("plain closures:")
@btime sched_plainclosure(arr1, arr2, ns)
reset()
println("typed closures:")
@btime sched_typedclos(arr1, arr2, ns)
reset()
println("manual type check:")
@btime sched_manual(arr1, arr2, ns)
reset()
println("array of functors:")
@btime sched_functors(arr1, arr2, ns)
reset()
println("FunctionWrapper with functor:")
@btime sched_fw(arr1, arr2, ns)
reset()
println("FunctionWrapper with closure:")
@btime sched_fwc(arr1, arr2, ns)
reset()