From 3ec1d1a424672754428b384587611d01b05c927e Mon Sep 17 00:00:00 2001
From: Michael Zingale <michael.zingale@stonybrook.edu>
Date: Wed, 8 May 2024 11:54:35 -0400
Subject: [PATCH] restructure to avoid function pointers

---
 integration/BackwardEuler/actual_integrator.H |  9 +++-
 integration/RKC/actual_integrator.H           |  9 +++-
 integration/VODE/actual_integrator.H          |  8 +++-
 integration/integrator_setup.H                | 46 ++++++++++++++-----
 4 files changed, 54 insertions(+), 18 deletions(-)

diff --git a/integration/BackwardEuler/actual_integrator.H b/integration/BackwardEuler/actual_integrator.H
index 8426b69a77..e53e47731f 100644
--- a/integration/BackwardEuler/actual_integrator.H
+++ b/integration/BackwardEuler/actual_integrator.H
@@ -17,8 +17,13 @@ void actual_integrator (BurnT& state, const amrex::Real dt, bool is_retry=false)
 
     constexpr int int_neqs = integrator_neqs<BurnT>();
 
-    integrator_setup<BurnT, be_t<int_neqs>>(state, dt, is_retry,
-                                            be_integrator<BurnT, be_t<int_neqs>>);
+    auto be_state = integrator_setup<BurnT, be_t<int_neqs>>(state, dt, is_retry);
+
+    auto state_save = integrator_backup(state);
+
+    auto istate = be_integrator(state, be_state);
+
+    integrator_cleanup(be_state, state, istate, state_save, dt);
 
 }
 
diff --git a/integration/RKC/actual_integrator.H b/integration/RKC/actual_integrator.H
index 8beeb0ff00..d5326ca96b 100644
--- a/integration/RKC/actual_integrator.H
+++ b/integration/RKC/actual_integrator.H
@@ -16,8 +16,13 @@ void actual_integrator (BurnT& state, amrex::Real dt, bool is_retry=false)
 {
     constexpr int int_neqs = integrator_neqs<BurnT>();
 
-    integrator_setup<BurnT, rkc_t<int_neqs>>(state, dt, is_retry,
-                                             rkc<BurnT, rkc_t<int_neqs>>);
+    auto rkc_state = integrator_setup<BurnT, rkc_t<int_neqs>>(state, dt, is_retry);
+
+    auto state_save = integrator_backup(state);
+
+    auto istate = rkc(state, rkc_state);
+
+    integrator_cleanup(rkc_state, state, istate, state_save, dt);
 
 }
 
diff --git a/integration/VODE/actual_integrator.H b/integration/VODE/actual_integrator.H
index c9f853526c..b51006d1d7 100644
--- a/integration/VODE/actual_integrator.H
+++ b/integration/VODE/actual_integrator.H
@@ -17,8 +17,12 @@ void actual_integrator (BurnT& state, amrex::Real dt, bool is_retry=false)
 
     constexpr int int_neqs = integrator_neqs<BurnT>();
 
-    integrator_setup<BurnT, dvode_t<int_neqs>>(state, dt, is_retry,
-                                               dvode<BurnT, dvode_t<int_neqs>>);
+    auto vode_state = integrator_setup<BurnT, dvode_t<int_neqs>>(state, dt, is_retry);
+    auto state_save = integrator_backup(state);
+
+    auto istate = dvode(state, vode_state);
+
+    integrator_cleanup(vode_state, state, istate, state_save, dt);
 
 }
 
diff --git a/integration/integrator_setup.H b/integration/integrator_setup.H
index 6cdaa3d704..9af9d44426 100644
--- a/integration/integrator_setup.H
+++ b/integration/integrator_setup.H
@@ -14,10 +14,17 @@
 
 #include <extern_parameters.H>
 
+struct state_backup_t {
+    amrex::Real T_in{};
+    amrex::Real e_in{};
+#ifndef AMREX_USE_GPU
+    amrex::Real xn_in[NumSpec]{};
+#endif
+};
+
 template <typename BurnT, typename IntegratorT>
 AMREX_GPU_HOST_DEVICE AMREX_INLINE
-void integrator_setup (BurnT& state, amrex::Real dt, bool is_retry,
-                       int (*integrator)(BurnT &, IntegratorT &))
+IntegratorT integrator_setup (BurnT& state, amrex::Real dt, bool is_retry)
 {
 
     IntegratorT int_state{};
@@ -80,19 +87,34 @@ void integrator_setup (BurnT& state, amrex::Real dt, bool is_retry,
 
     // Save the initial composition, temperature, and energy for our later diagnostics.
 
+    return int_state;
+}
+
+
+template <typename BurnT>
+AMREX_GPU_HOST_DEVICE AMREX_INLINE
+state_backup_t integrator_backup (const BurnT& state) {
+
+    state_backup_t state_save;
+
 #ifndef AMREX_USE_GPU
-    amrex::Real xn_in[NumSpec];
     for (int n = 0; n < NumSpec; ++n) {
-        xn_in[n] = state.xn[n];
+        state_save.xn_in[n] = state.xn[n];
     }
-    const amrex::Real T_in = state.T;
+    state_save.T_in = state.T;
 #endif
-    const amrex::Real e_in = state.e;
+    state_save.e_in = state.e;
+
+    return state_save;
+
+}
 
-    // Call the integration routine.
 
-    int istate = integrator(state, int_state);
-    state.error_code = istate;
+template <typename BurnT, typename IntegratorT>
+AMREX_GPU_HOST_DEVICE AMREX_INLINE
+void integrator_cleanup (IntegratorT& int_state, BurnT& state,
+                         int istate, const state_backup_t& state_save, amrex::Real dt)
+{
 
     // Copy the integration data back to the burn state.
 
@@ -115,7 +137,7 @@ void integrator_setup (BurnT& state, amrex::Real dt, bool is_retry,
     // to get back only the generated energy during the burn).
 
     if (integrator_rp::subtract_internal_energy) {
-        state.e -= e_in;
+        state.e -= state_save.e_in;
     }
 
     // Normalize the final abundances (except if they are number
@@ -177,9 +199,9 @@ void integrator_setup (BurnT& state, amrex::Real dt, bool is_retry,
             std::cout << "zone = (" << state.i << ", " << state.j << ", " << state.k << ")" << std::endl;
             std::cout << "time = " << int_state.t << std::endl;
             std::cout << "dt = " << std::setprecision(16) << dt << std::endl;
-            std::cout << "temp start = " << std::setprecision(16) << T_in << std::endl;
+            std::cout << "temp start = " << std::setprecision(16) << state_save.T_in << std::endl;
             std::cout << "xn start = ";
-            for (const double X : xn_in) {
+            for (const double X : state_save.xn_in) {
                 std::cout << std::setprecision(16) << X << " ";
             }
             std::cout << std::endl;