From 463cc1a81f67ff57729001b3e7747bb68132e277 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Mon, 16 Sep 2024 14:25:08 -0400
Subject: [PATCH] Fix some JET failures

---
 src/MatrixFields/field_matrix_solver.jl | 14 +-------------
 src/MatrixFields/field_name_set.jl      |  5 +++--
 2 files changed, 4 insertions(+), 15 deletions(-)

diff --git a/src/MatrixFields/field_matrix_solver.jl b/src/MatrixFields/field_matrix_solver.jl
index ccad25690c..1ade13250a 100644
--- a/src/MatrixFields/field_matrix_solver.jl
+++ b/src/MatrixFields/field_matrix_solver.jl
@@ -247,17 +247,6 @@ function check_field_matrix_solver(::BlockDiagonalSolve, _, A, b)
     end
 end
 
-# TODO: we can remove the uniform_vertical_levels
-# limitation while still using static shared memory
-# once Nv is in the type space.
-function uniform_vertical_levels(x, names)
-    _, _, _, Nv1, _ = size(Fields.field_values(x[first(names)]))
-    return all(Base.tail(names)) do name
-        _, _, _, Nv, _ = size(Fields.field_values(x[name]))
-        Nv == Nv1
-    end
-end
-
 NVTX.@annotate function run_field_matrix_solver!(
     ::BlockDiagonalSolve,
     cache,
@@ -267,8 +256,7 @@ NVTX.@annotate function run_field_matrix_solver!(
 )
     names = matrix_row_keys(keys(A))
     if length(names) == 1 ||
-       all(name -> A[name, name] isa UniformScaling, names.values) ||
-       !uniform_vertical_levels(x, names.values)
+       all(name -> A[name, name] isa UniformScaling, names.values)
         foreach(names) do name
             single_field_solve!(cache[name], x[name], A[name, name], b[name])
         end
diff --git a/src/MatrixFields/field_name_set.jl b/src/MatrixFields/field_name_set.jl
index 880746a420..8556174e69 100644
--- a/src/MatrixFields/field_name_set.jl
+++ b/src/MatrixFields/field_name_set.jl
@@ -276,10 +276,11 @@ end
 values_string(values) =
     length(values) == 2 ? join(values, " and ") : join(values, ", ", ", and ")
 
-combine_eltypes(T1, T2) =
-    T1 == T2 ? T1 :
+@noinline combine_eltypes(::T1, ::T2) where {T1, T2} =
     errror("Mismatched FieldNameSets: Cannot combine a $T1 with a $T2")
 
+@inline combine_eltypes(::Type{T}, ::Type{T}) where {T} = T
+
 combine_name_trees(::Nothing, ::Nothing) = nothing
 combine_name_trees(name_tree1, ::Nothing) = name_tree1
 combine_name_trees(::Nothing, name_tree2) = name_tree2