From 3c6b37ed3852d3590afe9b29f9690b0824f84a09 Mon Sep 17 00:00:00 2001
From: Benjamin Chung <ckfinite@gmail.com>
Date: Tue, 12 Nov 2024 23:46:52 -0800
Subject: [PATCH] Simplify setup function construction

---
 src/systems/callbacks.jl | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/src/systems/callbacks.jl b/src/systems/callbacks.jl
index 9c8187f6b7..b9423b8f07 100644
--- a/src/systems/callbacks.jl
+++ b/src/systems/callbacks.jl
@@ -918,24 +918,21 @@ function generate_vector_rootfinding_callback(
     initialize = nothing
     if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
         initialize = handle_optional_setup_fn(
-            map(
-                (cb, fn) -> begin
-                    if (save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
-                        let save_idxs = save_idxs
-                            custom_init = fn.initialize
-                            (i) -> begin
-                                isnothing(custom_init) && custom_init(i)
-                                for idx in save_idxs
-                                    SciMLBase.save_discretes!(i, idx)
-                                end
+            map(cbs, affect_functions) do cb, fn
+                if (save_idxs = get(ic.callback_to_clocks, cb, nothing)) !== nothing
+                    let save_idxs = save_idxs
+                        custom_init = fn.initialize
+                        (i) -> begin
+                            isnothing(custom_init) && custom_init(i)
+                            for idx in save_idxs
+                                SciMLBase.save_discretes!(i, idx)
                             end
                         end
-                    else
-                        fn.initialize
                     end
-                end,
-                cbs,
-                affect_functions),
+                else
+                    fn.initialize
+                end
+            end,
             SciMLBase.INITIALIZE_DEFAULT)
 
     else