Deferred removal of first occurence of subexpressions

src/Nonlinear/model.jl

@@ -11,6 +11,9 @@ function MOI.empty!(model::Model)
     empty!(model.parameters)
     model.operators = OperatorRegistry()
     model.last_constraint_index = 0
+    empty!(model.cache)
+    empty!(model.cache_expressions)
+    empty!(model.cache_pending)
     return
 end
 
@@ -322,6 +325,7 @@ function evaluate(
     expr::Expression;
     evaluated_expressions = Dict{Int,Float64}(),
 )
+    MOI.Utilities.final_touch(model)
     storage = zeros(length(expr.nodes))
     adj = adjacency_matrix(expr.nodes)
     children_arr = SparseArrays.rowvals(adj)

src/Nonlinear/parse.jl

@@ -36,17 +36,187 @@ function parse_expression(::Model, ::Expression, x::Any, ::Int)
     )
 end
 
+"""
+    MOI.Utilities.final_touch(model::Model, index_map = nothing)
+
+Extract the subexpressions that were detected by [`parse_expression`](@ref)
+since the last call to `final_touch`.
+
+When a `MOI.ScalarNonlinearFunction` is parsed a second time, its first
+occurrence, which was parsed inline into some expression, has to be moved into
+a new subexpression so that both occurrences can share it. Doing this
+immediately would shift the nodes of that expression, so every node index
+stored in the cache (and in the stack of the current parse) would need to be
+updated for each extraction, which is quadratic. Instead, `parse_expression`
+reserves an [`ExpressionIndex`](@ref) for the subexpression, with an empty
+placeholder [`Expression`](@ref), and records the node to extract in
+`model.cache_pending`. This function then rewrites each affected expression
+in a single pass, and remaps the cache in a single pass.
+
+The subexpressions are invalid until this function is called. It is called by
+the [`Evaluator`](@ref) constructors and by [`evaluate`](@ref), and it is a
+no-op if there is nothing to extract.
+"""
+function MOI.Utilities.final_touch(model::Model, ::Any = nothing)
+    pending = model.cache_pending
+    if isempty(pending)
+        return
+    end
+    # Group the nodes to extract by expression, sorted by node index.
+    sort!(pending)
+    maps = Dict{Int,Tuple{Vector{Int},Vector{Int}}}()
+    first = 1
+    while first <= length(pending)
+        tape = pending[first][1]
+        last = first
+        while last < length(pending) && pending[last+1][1] == tape
+            last += 1
+        end
+        maps[tape] = _extract_subexpressions!(model, tape, pending, first, last)
+        first = last + 1
+    end
+    empty!(pending)
+    # Remap the cached nodes of the expressions that were rewritten.
+    for (key, (tape, node)) in model.cache
+        if node != 0 && haskey(maps, tape)
+            new_tape, new_node = maps[tape]
+            model.cache[key] = (new_tape[node], new_node[node])
+        end
+    end
+    return
+end
+
+"""
+    _extract_subexpressions!(model::Model, tape::Int, pending, first, last)
+
+Extract the nodes `pending[i][2]` of `model.cache_expressions[tape]` into the
+subexpressions `model.expressions[pending[i][3]]` for `i in first:last`. The
+nodes `pending[first:last]` must be sorted.
+
+The nodes are stored in depth-first order so a node is visited after its
+parent. Each node is therefore appended to the same expression as its parent,
+unless it is to be extracted, in which case its parent gets a
+`NODE_SUBEXPRESSION` node instead and the node becomes the root of its
+subexpression. Nested extractions are handled in the same way.
+
+The nodes that stay in `model.cache_expressions[tape]` can only move to a
+smaller index so we rewrite it in place.
+
+Return `(new_tape, new_node)` such that node `i` of the original expression is
+now node `new_node[i]` of `model.cache_expressions[new_tape[i]]`.
+"""
+function _extract_subexpressions!(
+    model::Model,
+    tape::Int,
+    pending::Vector{Tuple{Int,Int,Int}},
+    first::Int,
+    last::Int,
+)
+    expr = model.cache_expressions[tape]
+    n = length(expr.nodes)
+    new_tape = Vector{Int}(undef, n)
+    new_node = Vector{Int}(undef, n)
+    # Number of nodes and values kept in `expr` so far.
+    num_nodes = num_values = 0
+    for i in 1:n
+        node = expr.nodes[i]
+        if node.parent == -1
+            parent_tape, parent = tape, -1
+        else
+            parent_tape, parent = new_tape[node.parent], new_node[node.parent]
+        end
+        if first <= last && pending[first][2] == i
+            subexpr = pending[first][3]
+            first += 1
+            # Replace the node by a `NODE_SUBEXPRESSION` in its parent's
+            # expression and continue with the node as root of `subexpr`.
+            subexpr_node = Node(NODE_SUBEXPRESSION, subexpr, parent)
+            if parent_tape == tape
+                num_nodes += 1
+                expr.nodes[num_nodes] = subexpr_node
+            else
+                push!(model.cache_expressions[parent_tape].nodes, subexpr_node)
+            end
+            # `subexpr` may contain cached nodes so it needs to be in
+            # `cache_expressions`.
+            push!(model.cache_expressions, model.expressions[subexpr])
+            parent_tape, parent = length(model.cache_expressions), -1
+        end
+        new_tape[i] = parent_tape
+        index = node.index
+        if parent_tape == tape
+            if node.type == NODE_VALUE
+                num_values += 1
+                expr.values[num_values] = expr.values[index]
+                index = num_values
+            end
+            num_nodes += 1
+            expr.nodes[num_nodes] = Node(node.type, index, parent)
+            new_node[i] = num_nodes
+        else
+            target = model.cache_expressions[parent_tape]
+            if node.type == NODE_VALUE
+                push!(target.values, expr.values[index])
+                index = length(target.values)
+            end
+            push!(target.nodes, Node(node.type, index, parent))
+            new_node[i] = length(target.nodes)
+        end
+    end
+    resize!(expr.nodes, num_nodes)
+    resize!(expr.values, num_values)
+    return new_tape, new_node
+end
+
 function parse_expression(
     data::Model,
     expr::Expression,
     x::MOI.ScalarNonlinearFunction,
     parent_index::Int,
 )
     stack = Tuple{Int,Any}[(parent_index, x)]
+    expr_index = 0
     while !isempty(stack)
         parent_node, arg = pop!(stack)
         if arg isa MOI.ScalarNonlinearFunction
-            _parse_without_recursion_inner(stack, data, expr, arg, parent_node)
+            if haskey(data.cache, arg)
+                # `arg` was already parsed, we reuse it as a subexpression.
+                _expr_index, _node = data.cache[arg]
+                if _node == 0
+                    # It was already extracted into a subexpression.
+                    subexpr = ExpressionIndex(_expr_index)
+                else
+                    # It was parsed inline at node `_node` of
+                    # `data.cache_expressions[_expr_index]`. This is its second
+                    # occurrence so it needs to be extracted into a
+                    # subexpression. We reserve its index and defer the
+                    # extraction to `MOI.Utilities.final_touch` so that no
+                    # node is moved in the meantime.
+                    push!(data.expressions, Expression())
+                    subexpr = ExpressionIndex(length(data.expressions))
+                    push!(
+                        data.cache_pending,
+                        (_expr_index, _node, subexpr.value),
+                    )
+                    data.cache[arg] = (subexpr.value, 0)
+                end
+                parse_expression(data, expr, subexpr, parent_node)
+            else
+                _parse_without_recursion_inner(
+                    stack,
+                    data,
+                    expr,
+                    arg,
+                    parent_node,
+                )
+                # Cache the node of `arg` in `expr` so that we can extract it
+                # if we encounter it again.
+                if expr_index == 0
+                    push!(data.cache_expressions, expr)
+                    expr_index = length(data.cache_expressions)
+                end
+                data.cache[arg] = (expr_index, length(expr.nodes))
+            end
         else
             # We can use recursion here, because ScalarNonlinearFunction only
             # occur in other ScalarNonlinearFunction.
@@ -82,7 +252,7 @@ function _parse_without_recursion_inner(stack, data, expr, x, parent)
     parent = length(expr.nodes)
     # Args need to be pushed onto the stack in reverse because the stack is a
     # first-in last-out datastructure.
-    for arg in reverse(x.args)
+    for arg in Iterators.Reverse(x.args)
         push!(stack, (parent, arg))
     end
     return

src/Nonlinear/types.jl

@@ -76,9 +76,10 @@ tree.
 struct Expression
     nodes::Vector{Node}
     values::Vector{Float64}
-    Expression() = new(Node[], Float64[])
 end
 
+Expression() = Expression(Node[], Float64[])
+
 function Base.:(==)(x::Expression, y::Expression)
     return x.nodes == y.nodes && x.values == y.values
 end
@@ -165,6 +166,19 @@ mutable struct Model
     operators::OperatorRegistry
     # This is a private field, used only to increment the ConstraintIndex.
     last_constraint_index::Int64
+    # Cache values are (expression index, node index). A zero node index
+    # identifies an extracted subexpression in `expressions`; otherwise the
+    # expression index refers to `cache_expressions`. The node indices refer
+    # to the expressions as they are before the pending extractions are done
+    # by `MOI.Utilities.final_touch`.
+    # We use a separate vector `cache_expressions` instead of storing
+    # a `Tuple{Expression,Int}` in the `Dict` as creating such tuple
+    # would trigger a heap allocation for each expression
+    cache::Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}}
+    cache_expressions::Vector{Expression}
+    # Extractions deferred to `MOI.Utilities.final_touch`. Each element is
+    # (index in `cache_expressions`, node index, index in `expressions`).
+    cache_pending::Vector{Tuple{Int,Int,Int}}
     function Model()
         return new(
             nothing,
@@ -173,6 +187,9 @@ mutable struct Model
             Float64[],
             OperatorRegistry(),
             0,
+            Dict{MOI.ScalarNonlinearFunction,Tuple{Int,Int}}(),
+            Expression[],
+            Tuple{Int,Int,Int}[],
         )
     end
 end
@@ -272,6 +289,7 @@ supported feature is `:ExprGraph`.
 struct ExprGraphOnly <: AbstractAutomaticDifferentiation end
 
 function Evaluator(model::Model, ::ExprGraphOnly, ::Vector{MOI.VariableIndex})
+    MOI.Utilities.final_touch(model)
     return Evaluator(model)
 end
 
@@ -289,6 +307,7 @@ function Evaluator(
     ::SparseReverseMode,
     ordered_variables::Vector{MOI.VariableIndex},
 )
+    MOI.Utilities.final_touch(model)
     return Evaluator(model, ReverseAD.NLPEvaluator(model, ordered_variables))
 end
 
@@ -305,5 +324,6 @@ function Evaluator(
     ::SymbolicMode,
     ordered_variables::Vector{MOI.VariableIndex},
 )
+    MOI.Utilities.final_touch(model)
     return Evaluator(model, SymbolicAD.Evaluator(model, ordered_variables))
 end