feat: add preallocation utilities for expression

Author: MilesCranmer

src/Expression.jl

@@ -317,6 +317,14 @@ function extract_gradient(
     return extract_gradient(gradient.tree, get_tree(ex))
 end
 
+function preallocate_expression(prototype::Expression, n::Union{Nothing,Integer}=nothing)
+    return (; tree=preallocate_expression(DE.get_contents(prototype), n))
+end
+function DE.copy_node!(dest::NamedTuple, src::Expression)
+    tree = DE.copy_node!(dest.tree, DE.get_contents(src))
+    return DE.with_contents(src, tree)
+end
+
 """
     string_tree(
         ex::AbstractExpression,

src/ParametricExpression.jl

@@ -4,7 +4,8 @@ using DispatchDoctor: @stable, @unstable
 using ChainRulesCore: ChainRulesCore as CRC, NoTangent, @thunk
 
 using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
-using ..NodeModule: AbstractExpressionNode, Node, tree_mapreduce
+using ..NodeModule:
+    AbstractExpressionNode, Node, tree_mapreduce, with_contents, with_metadata
 using ..ExpressionModule: AbstractExpression, Metadata
 using ..ChainRulesModule: NodeTangent
 
@@ -17,7 +18,9 @@ import ..NodeModule:
     leaf_convert,
     leaf_hash,
     leaf_equal,
-    branch_copy!
+    branch_copy!,
+    copy_node!,
+    preallocate_expression
 import ..NodeUtilsModule:
     count_constant_nodes,
     index_constant_nodes,
@@ -444,6 +447,28 @@ end
         return node_type(; val=ex)
     end
 end
+function preallocate_expression(
+    prototype::ParametricExpression, n::Union{Nothing,Integer}=nothing
+)
+    return (;
+        tree=preallocate_expression(get_contents(prototype), n),
+        parameters=similar(get_metadata(prototype).parameters),
+    )
+end
+function copy_node!(dest::NamedTuple, src::ParametricExpression)
+    new_tree = copy_node!(dest.tree, get_contents(src))
+    metadata = get_metadata(src)
+    new_parameters = dest.parameters
+    new_parameters .= metadata.parameters
+    new_metadata = Metadata((;
+        operators=metadata.operators,
+        variable_names=metadata.variable_names,
+        parameters=new_parameters,
+        parameter_names=metadata.parameter_names,
+    ))
+    # TODO: Better interface for this^
+    return with_metadata(with_contents(src, new_tree), new_metadata)
+end
 ###############################################################################
 
 end

src/base.jl

@@ -488,23 +488,25 @@ end
 # In-place versions
 
 """
-    copy_node!(dest::AbstractArray{N}, src::N; break_sharing::Val{BS}=Val(false)) where {BS,N<:AbstractExpressionNode}
+    copy_node!(dest::AbstractArray{N}, src::N) where {BS,N<:AbstractExpressionNode}
 
 Copy a node, recursively copying all children nodes, in-place to an
 array of pre-allocated nodes. This should result in no extra allocations.
 """
 function copy_node!(
-    dest::AbstractArray{N},
-    src::N;
-    break_sharing::Val{BS}=Val(false),
-    ref::Base.RefValue{<:Integer}=Ref(0),
-) where {BS,N<:AbstractExpressionNode}
-    ref.x = 0
+    dest::AbstractArray{N}, src::N; ref::Union{Nothing,Base.RefValue{<:Integer}}=nothing
+) where {N<:AbstractExpressionNode}
+    _ref = if ref === nothing
+        Ref(0)
+    else
+        ref.x = 0
+        ref
+    end
     return tree_mapreduce(
-        leaf -> leaf_copy!(@inbounds(dest[ref.x += 1]), leaf),
+        leaf -> leaf_copy!(@inbounds(dest[_ref.x += 1]), leaf),
         identity,
         ((p, c::Vararg{Any,M}) where {M}) ->
-            branch_copy!(@inbounds(dest[ref.x += 1]), p, c...),
+            branch_copy!(@inbounds(dest[_ref.x += 1]), p, c...),
         src,
         N;
         break_sharing=Val(BS),
@@ -533,6 +535,31 @@ function branch_copy!(
     return dest
 end
 
+"""
+    preallocate_expression(prototype::AbstractExpressionNode, n=nothing)
+
+Preallocate an array of empty nodes matching the type of `prototype`. If `n` is provided, use that length, otherwise use `length(prototype)`.
+
+A given return value of this will be passed to `copy_node!` as the first argument,
+so it should be compatible.
+"""
+function preallocate_expression(
+    prototype::N, n::Union{Nothing,Integer}=nothing
+) where {T,N<:AbstractExpressionNode{T}}
+    num_nodes = @something(n, length(prototype))
+    return N[with_type_parameters(N, T)() for _ in 1:num_nodes]
+end
+
+function copy_node!(::Nothing, src::AbstractExpression)
+    return copy(src)
+end
+function preallocate_expression(::AbstractExpression, ::Union{Nothing,Integer}=nothing)
+    return nothing
+end
+# We don't require users to overload this, as it's not part of the required interface.
+# Also, there's no way to generally do this from the required interface, so for backwards
+# compatibility, we just return nothing.
+
 """
     copy(tree::AbstractExpressionNode; break_sharing::Val=Val(false))