From b5f5ef864faf2809337473a4ce3ceb039bd07d16 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Wed, 4 Oct 2023 13:07:19 -0700
Subject: [PATCH 1/2] Simplify dispatch

---
 perf/flame.jl             |  2 +-
 src/solvers/imex_ark.jl   | 15 ++++++---------
 src/solvers/imex_ssprk.jl |  5 ++---
 3 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/perf/flame.jl b/perf/flame.jl
index 31597769..625e5753 100644
--- a/perf/flame.jl
+++ b/perf/flame.jl
@@ -35,7 +35,7 @@ end
 p = @allocated do_work!(integrator, cache)
 using Test
 @testset "Allocations" begin
-    @test_broken p == 0
+    @test p == 0
 end
 
 import ProfileCanvas
diff --git a/src/solvers/imex_ark.jl b/src/solvers/imex_ark.jl
index 8a3e8f3c..d4b45dd7 100644
--- a/src/solvers/imex_ark.jl
+++ b/src/solvers/imex_ark.jl
@@ -45,12 +45,10 @@ function init_cache(prob::DiffEqBase.AbstractODEProblem, alg::IMEXAlgorithm{Unco
     return IMEXARKCache(U, T_lim, T_exp, T_imp, temp, γ, newtons_method_cache)
 end
 
-step_u!(integrator, cache::IMEXARKCache) = step_u!(integrator, cache, integrator.sol.prob.f, integrator.alg.name)
-
-# include("hard_coded_ars343.jl")
 # generic fallback
-function step_u!(integrator, cache::IMEXARKCache, f, name)
+function step_u!(integrator, cache::IMEXARKCache)
     (; u, p, t, dt, alg) = integrator
+    (; f) = integrator.sol.prob
     (; post_explicit!, post_implicit!) = f
     (; T_lim!, T_exp!, T_imp!, lim!, dss!) = f
     (; tableau, newtons_method) = alg
@@ -107,11 +105,10 @@ function step_u!(integrator, cache::IMEXARKCache, f, name)
             @. temp = U
             post_explicit!(U, p, t_imp)
             # TODO: can/should we remove these closures?
-            implicit_equation_residual! =
-                (residual, Ui) -> begin
-                    T_imp!(residual, Ui, p, t_imp)
-                    @. residual = temp + dt * a_imp[i, i] * residual - Ui
-                end
+            implicit_equation_residual! = (residual, Ui) -> begin
+                T_imp!(residual, Ui, p, t_imp)
+                @. residual = temp + dt * a_imp[i, i] * residual - Ui
+            end
             implicit_equation_jacobian! = (jacobian, Ui) -> T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
             call_post_implicit! = Ui -> begin
                 post_implicit!(Ui, p, t_imp)
diff --git a/src/solvers/imex_ssprk.jl b/src/solvers/imex_ssprk.jl
index 5f5982ba..d9a0c6a1 100644
--- a/src/solvers/imex_ssprk.jl
+++ b/src/solvers/imex_ssprk.jl
@@ -52,10 +52,9 @@ function init_cache(prob::DiffEqBase.AbstractODEProblem, alg::IMEXAlgorithm{SSP}
     return IMEXSSPRKCache(U, U_exp, U_lim, T_lim, T_exp, T_imp, temp, β, γ, newtons_method_cache)
 end
 
-step_u!(integrator, cache::IMEXSSPRKCache) = step_u!(integrator, cache, integrator.sol.prob.f, integrator.alg.name)
-
-function step_u!(integrator, cache::IMEXSSPRKCache, f, name)
+function step_u!(integrator, cache::IMEXSSPRKCache)
     (; u, p, t, dt, alg) = integrator
+    (; f) = integrator.sol.prob
     (; post_explicit!, post_implicit!) = f
     (; T_lim!, T_exp!, T_imp!, lim!, dss!) = f
     (; tableau, newtons_method) = alg

From 1726619cb5e307e9ca675e8fecaed837462d5a85 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Wed, 4 Oct 2023 13:45:25 -0700
Subject: [PATCH 2/2] Bump patch version

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 9a573241..916c4dfb 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.12"
+version = "0.7.13"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"