From 5d92e5c9e14994fdcec83b235f9d187e1c979556 Mon Sep 17 00:00:00 2001
From: Adriano Meligrana <68152031+Tortar@users.noreply.github.com>
Date: Mon, 12 Aug 2024 15:16:34 +0200
Subject: [PATCH] Avoid using a Dict for index retrieval with multiagent macro
 (#1066)

* Avoid using a Dict for index retrieval with multiagent macro

* Update Project.toml

* Update model_event_queue.jl
---
 Project.toml                  |  2 +-
 src/core/model_event_queue.jl | 22 +++++++++++++---------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/Project.toml b/Project.toml
index 88c0e4e7cb..844f2722e9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Agents"
 uuid = "46ada45e-f475-11e8-01d0-f70cc89e6671"
 authors = ["George Datseris", "Tim DuBois", "Aayush Sabharwal", "Ali Vahdati", "Adriano Meligrana"]
-version = "6.1.2"
+version = "6.1.3"
 
 [deps]
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
diff --git a/src/core/model_event_queue.jl b/src/core/model_event_queue.jl
index baab7a794c..0ef6ae302a 100644
--- a/src/core/model_event_queue.jl
+++ b/src/core/model_event_queue.jl
@@ -38,7 +38,7 @@ struct EventQueueABM{
     S<:SpaceType,
     A<:AbstractAgent,
     C<:ContainerType{A},
-    P,E,R<:AbstractRNG,TF,ET,PT,FPT,TI,Q} <: AgentBasedModel{S}
+    P,E,R<:AbstractRNG,TF,ET,PT,FPT,Q} <: AgentBasedModel{S}
     # core ABM stuff
     agents::C
     space::S
@@ -49,7 +49,6 @@ struct EventQueueABM{
     # Specific to event queue
     events::E
     type_func::TF
-    type_to_idx::TI
     idx_events_each_type::ET
     propensities_each_type::PT
     idx_func_propensities_each_type::FPT
@@ -164,8 +163,13 @@ function EventQueueABM(
     queue = BinaryHeap(Base.By(last), Pair{Tuple{Int, Int}, Float64}[])
 
     agent_types = is_sumtype(A) ? values(allvariants(A)) : union_types(A)
-    type_func = is_sumtype(A) ? variantof : typeof
-    type_to_idx = Dict(t => i for (i, t) in enumerate(agent_types))
+
+    if is_sumtype(A) 
+        type_func = DynamicSumTypes.variant_idx 
+    else 
+        type_to_idx = Dict(t => i for (i, t) in enumerate(agent_types))
+        type_func = (agent) -> type_to_idx[typeof(agent)]
+    end
 
     # precompute a vector mapping the agent type index to a
     # vectors of indices, each vector corresponding
@@ -196,15 +200,15 @@ function EventQueueABM(
     # because we use the index of `type_to_idx` to access them.
 
     # construct the type
-    E,TF,ET,PT,FPT,TI,Q = typeof.((
+    E,TF,ET,PT,FPT,Q = typeof.((
         events, type_func, idx_events_each_type, propensities_each_type,
-        idx_func_propensities_each_type, type_to_idx, queue
+        idx_func_propensities_each_type, queue
     ))
-    return EventQueueABM{S,A,C,P,E,R,TF,ET,PT,FPT,TI,Q}(
+    return EventQueueABM{S,A,C,P,E,R,TF,ET,PT,FPT,Q}(
 
         agents, space, properties, rng, Ref(0), Ref(0.0),
 
-        events, type_func, type_to_idx, idx_events_each_type,
+        events, type_func, idx_events_each_type,
         propensities_each_type, idx_func_propensities_each_type,
         queue, autogenerate_on_add, autogenerate_after_action,
     )
@@ -237,7 +241,7 @@ current time of the `model`.
 function add_event!(agent, model) # TODO: Study type stability of this function
     events = abmevents(model)
     # Here, we retrieve the applicable events for the agent and corresponding info
-    idx = getfield(model, :type_to_idx)[getfield(model, :type_func)(agent)]
+    idx = getfield(model, :type_func)(agent)
     events_type = getfield(model, :idx_events_each_type)[idx]
     propensities_type = getfield(model, :propensities_each_type)[idx]
     idx_func_propensities_type = getfield(model, :idx_func_propensities_each_type)[idx]