Simplify branching and fix type instability in setproperty!

Author: hersle

ext/LinearSolveForwardDiffExt.jl

@@ -308,31 +308,37 @@ function SciMLBase.solve!(
     )
 end
 
-# If setting A or b for DualLinearCache, put the Dual-stripped versions in the LinearCache
-function Base.setproperty!(dc::DualLinearCache, sym::Symbol, val::AbstractArray)
+# Generalized function to set A, b or u ("x") for DualLinearCache
+function setx(dc::DualLinearCache, x::Symbol, dual_x::Symbol, partials_x::Symbol, val, invalidate::Bool)
+    # Put the Dual-stripped versions in the LinearCache
+    prop = nodual_value!(getproperty(dc.linear_cache, x), val) # Update in-place
+    setproperty!(dc.linear_cache, x, prop) # Does additional invalidation logic etc.
+
+    # Update partials
+    setfield!(dc, dual_x, val)
+    partial_vals!(getfield(dc, partials_x), val) # Update in-place
+
+    # Invalidate cache (if setting A or b)
+    invalidate && setfield!(dc, :rhs_cache_valid, false)
+end
+setA!(dc::DualLinearCache, A) = setx(dc, :A, :dual_A, :partials_A, A, true)
+setb!(dc::DualLinearCache, b) = setx(dc, :b, :dual_b, :partials_b, b, true)
+setu!(dc::DualLinearCache, u) = setx(dc, :u, :dual_u, :partials_u, u, false)
+
+function Base.setproperty!(dc::DualLinearCache, sym::Symbol, val)
     # If the property is A or b, also update it in the LinearCache
-    if sym === :A || sym === :b || sym === :u
-        prop = nodual_value!(getproperty(dc.linear_cache, sym), val) # Update in-place
-        setproperty!(dc.linear_cache, sym, prop) # Does additional invalidation logic etc.
+    if sym === :A
+        setA!(dc, val)
+    elseif sym === :b
+        setb!(dc, val)
+    elseif sym === :u
+        setu!(dc, val)
     elseif hasfield(DualLinearCache, sym)
         setfield!(dc, sym, val)
     elseif hasfield(LinearSolve.LinearCache, sym)
         setproperty!(dc.linear_cache, sym, val)
     end
-
-    # Update the partials and invalidate cache if setting A or b
-    if sym === :A
-        setfield!(dc, :dual_A, val)
-        partial_vals!(getfield(dc, :partials_A), val) # Update in-place
-        setfield!(dc, :rhs_cache_valid, false)  # Invalidate cache
-    elseif sym === :b
-        setfield!(dc, :dual_b, val)
-        partial_vals!(getfield(dc, :partials_b), val) # Update in-place
-        setfield!(dc, :rhs_cache_valid, false)  # Invalidate cache
-    elseif sym === :u
-        setfield!(dc, :dual_u, val)
-        partial_vals!(getfield(dc, :partials_u), val) # Update in-place
-    end
+    nothing
 end
 
 # "Forwards" getproperty to LinearCache if necessary