Maintain order of expression generators when calling flatten! (#4196)

Author: odow

src/nlp_expr.jl

@@ -511,7 +511,7 @@ a single n-ary operation.
 
 Nonlinear expressions created using operator overloading can be deeply nested
 and unbalanced. For example, `prod(x for i in 1:4)` creates
-`*(x, *(x, *(x, x)))` instead of the more preferable `*(x, x, x, x)`.
+`*(*(*(x, x), x, x)` instead of the more preferable `*(x, x, x, x)`.
 
 ## Example
 
@@ -536,27 +536,35 @@ function flatten!(expr::GenericNonlinearExpr{V}) where {V}
         return expr
     end
     stack = Tuple{GenericNonlinearExpr{V},Int,GenericNonlinearExpr{V}}[]
-    for i in length(expr.args):-1:1
+    for i in 1:length(expr.args)
         if _needs_flatten(expr, expr.args[i])
             push!(stack, (expr, i, expr.args[i]))
         end
     end
     while !isempty(stack)
         parent, i, arg = pop!(stack)
         if parent.head in (:+, :*) && arg.head == parent.head
-            n = length(parent.args)
-            resize!(parent.args, n + length(arg.args) - 1)
-            for j in length(arg.args):-1:1
-                parent_index = j == 1 ? i : n + j - 1
+            m, n = length(arg.args), length(parent.args)
+            # We're going to splice the `m` args of the child into
+            # `parent.args[i:i+m-1]`. To do so, extend the args to make space:
+            resize!(parent.args, n + (m - 1))
+            # and then shift the parent args along to make room:
+            for j in n:-1:(i+1)
+                parent.args[m+j-1] = parent.args[j]
+            end
+            # Now we can put each child arg into the parent.
+            for j in 1:m
+                parent_index = i + j - 1
                 if _needs_flatten(parent, arg.args[j])
+                    # The child needs flattening itself
                     push!(stack, (parent, parent_index, arg.args[j]))
                 else
                     parent.args[parent_index] = arg.args[j]
                 end
             end
         else
             parent.args[i] = arg
-            for j in length(arg.args):-1:1
+            for j in 1:length(arg.args)
                 if _needs_flatten(arg, arg.args[j])
                     push!(stack, (arg, j, arg.args[j]))
                 end

test/test_nlp_expr.jl

@@ -124,6 +124,17 @@ function test_extension_flatten_nary(
     return
 end
 
+function test_extension_flatten_nary_order_maintained(
+    ModelType = Model,
+    VariableRefType = VariableRef,
+)
+    model = ModelType()
+    @variable(model, x[1:3])
+    @test sum(log(x[i]) for i in 1:3) |> flatten! |> string ==
+          "log(x[1]) + log(x[2]) + log(x[3])"
+    return
+end
+
 function test_extension_zero_one(
     ModelType = Model,
     VariableRefType = VariableRef,
@@ -1070,11 +1081,11 @@ function test_printing_truncation()
     @variable(model, x[1:100])
     y = @expression(model, sum(sin.(x) .* 2))
     @test occursin(
-        "(sin(x[72]) * 2) + [[...40 terms omitted...]] + (sin(x[31]) * 2)",
+        "(sin(x[30]) * 2) + [[...40 terms omitted...]] + (sin(x[71]) * 2)",
         function_string(MIME("text/plain"), y),
     )
     @test occursin(
-        "{\\left({\\textsf{sin}\\left({x_{72}}\\right)} * {2}\\right) + {[[\\ldots\\text{40 terms omitted}\\ldots]]} + {\\left({\\textsf{sin}\\left({x_{31}}\\right)} * {2}\\right)}",
+        "{\\left({\\textsf{sin}\\left({x_{30}}\\right)} * {2}\\right) + {[[\\ldots\\text{40 terms omitted}\\ldots]]} + {\\left({\\textsf{sin}\\left({x_{71}}\\right)} * {2}\\right)}",
         function_string(MIME("text/latex"), y),
     )
     return