feat: add preallocation for abstract structured expression

Author: MilesCranmer

src/Expression.jl

@@ -19,6 +19,7 @@ import ..NodeUtilsModule:
     count_scalar_constants,
     get_scalar_constants,
     set_scalar_constants!
+import ..NodePreallocationModule: copy_into!, allocate_container
 import ..EvaluateModule: eval_tree_array, differentiable_eval_tree_array
 import ..EvaluateDerivativeModule: eval_grad_tree_array
 import ..EvaluationHelpersModule: _grad_evaluator

src/StructuredExpression.jl

@@ -6,12 +6,14 @@ using ..ExpressionModule: AbstractExpression, Metadata, node_type
 using ..ChainRulesModule: NodeTangent
 
 import ..NodeModule: constructorof
+import ..NodePreallocationModule: copy_into!, allocate_container
 import ..ExpressionModule:
     get_contents,
     get_metadata,
     get_tree,
     get_operators,
     get_variable_names,
+    with_contents,
     Metadata,
     _copy,
     _data,
@@ -164,4 +166,16 @@ function set_scalar_constants!(e::AbstractStructuredExpression, constants, refs)
     return e
 end
 
+function allocate_container(
+    e::AbstractStructuredExpression, n::Union{Nothing,Integer}=nothing
+)
+    ts = get_contents(e)
+    return (; trees=NamedTuple{keys(ts)}(map(t -> allocate_container(t, n), values(ts))))
+end
+function copy_into!(dest::NamedTuple, src::AbstractStructuredExpression)
+    ts = get_contents(src)
+    new_contents = NamedTuple{keys(ts)}(map(copy_into!, values(dest.trees), values(ts)))
+    return with_contents(src, new_contents)
+end
+
 end