Make copying ScalarNonlinearFunction type-stable

src/functions.jl

@@ -353,27 +353,21 @@ end
 # copy() doesn't recursively copy the children, and deepcopy seems to have a
 # performance problem for deeply nested structs.
 function Base.copy(f::ScalarNonlinearFunction)
-    stack, result_stack = Any[f], Any[]
+    result = ScalarNonlinearFunction(f.head, similar(f.args))
+    stack = Tuple{ScalarNonlinearFunction,ScalarNonlinearFunction}[(f, result)]
     while !isempty(stack)
-        arg = pop!(stack)
-        if arg isa ScalarNonlinearFunction
-            # We need some sort of hint so that the next time we see this on the
-            # stack we evaluate it using the args in `result_stack`. One option
-            # would be a custom type. Or we can just wrap in (,) and then check
-            # for a Tuple, which isn't (currently) a valid argument.
-            push!(stack, (arg,))
-            for child in arg.args
-                push!(stack, child)
+        source, destination = pop!(stack)
+        for (i, arg) in enumerate(source.args)
+            if arg isa ScalarNonlinearFunction
+                child = ScalarNonlinearFunction(arg.head, similar(arg.args))
+                destination.args[i] = child
+                push!(stack, (arg, child))
+            else
+                destination.args[i] = copy(arg)
             end
-        elseif arg isa Tuple{<:ScalarNonlinearFunction}
-            result = only(arg)
-            args = Any[pop!(result_stack) for i in 1:length(result.args)]
-            push!(result_stack, ScalarNonlinearFunction(result.head, args))
-        else
-            push!(result_stack, copy(arg))
         end
     end
-    return only(result_stack)
+    return result
 end
 
 constant(f::ScalarNonlinearFunction, ::Type{T} = Float64) where {T} = zero(T)

test/General/test_functions.jl

@@ -575,7 +575,7 @@ function test_copy_ScalarNonlinearFunction()
         g = MOI.ScalarNonlinearFunction(:^, Any[x[i], 1])
         f2 = MOI.ScalarNonlinearFunction(:+, Any[f2, g])
     end
-    f_copy = copy(f1)
+    f_copy = @inferred copy(f1)
     @test ≈(f_copy, f2)
     f1.args[2].args[2] = 2.0  # x[1]^1 --> x[1]^2
     @test !isapprox(f_copy, f1)
@@ -596,14 +596,49 @@ function test_copy_ScalarNonlinearFunction_with_arg()
         g = f2 = Float64(i) * x[i] + Float64(i)
         f2 = MOI.ScalarNonlinearFunction(:+, Any[f2, g])
     end
-    f_copy = copy(f1)
+    f_copy = @inferred copy(f1)
     @test ≈(f_copy, f2)
     f1.args[2].constant += 1
     @test !isapprox(f_copy, f1)
     @test isapprox(f_copy, f2)
     return
 end
 
+function test_copy_ScalarNonlinearFunction_leaf_types()
+    x = MOI.VariableIndex(1)
+    affine = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(big"2", x)], big"3")
+    quadratic = MOI.ScalarQuadraticFunction(
+        [MOI.ScalarQuadraticTerm(2 // 3, x, x)],
+        MOI.ScalarAffineTerm{Rational{Int}}[],
+        1 // 2,
+    )
+    child = MOI.ScalarNonlinearFunction(:+, Any[affine, quadratic])
+    f = MOI.ScalarNonlinearFunction(
+        :+,
+        Any[child, 1.0f0, big"4.0", 1//5, big"6"],
+    )
+    g = @inferred copy(f)
+    @test g ≈ f
+    @test g.args !== f.args
+    @test g.args[1].args !== child.args
+    @test g.args[1].args[1].terms !== affine.terms
+    @test g.args[1].args[2].quadratic_terms !== quadratic.quadratic_terms
+    @test typeof(g.args[2]) == Float32
+    @test typeof(g.args[3]) == BigFloat
+    @test typeof(g.args[4]) == Rational{Int}
+    @test typeof(g.args[5]) == BigInt
+    affine.constant += 1
+    quadratic.constant += 1
+    @test g.args[1].args[1].constant == 3
+    @test g.args[1].args[2].constant == 1 // 2
+    empty_f = MOI.ScalarNonlinearFunction(:+, Any[])
+    empty_g = @inferred copy(empty_f)
+    @test empty_g.head == :+
+    @test isempty(empty_g.args)
+    @test empty_g.args !== empty_f.args
+    return
+end
+
 function test_isapprox_Number()
     x = MOI.VariableIndex(1)
     for f in Any[