fix: various issues in preallocation interface

Author: MilesCranmer

src/Interfaces.jl

@@ -13,12 +13,10 @@ using ..NodeModule:
     default_allocator,
     with_type_parameters,
     leaf_copy,
-    leaf_copy_into!,
     leaf_convert,
     leaf_hash,
     leaf_equal,
     branch_copy,
-    branch_copy_into!,
     branch_convert,
     branch_hash,
     branch_equal,
@@ -38,6 +36,8 @@ using ..NodeUtilsModule:
     has_constants,
     get_scalar_constants,
     set_scalar_constants!
+using ..NodePreallocationModule:
+    copy_into!, leaf_copy_into!, branch_copy_into!, allocate_container
 using ..StringsModule: string_tree
 using ..EvaluateModule: eval_tree_array
 using ..EvaluateDerivativeModule: eval_grad_tree_array
@@ -96,6 +96,11 @@ function _check_with_metadata(ex::AbstractExpression)
 end
 
 ## optional
+function _check_copy_into!(ex::AbstractExpression)
+    container = allocate_container(ex)
+    prealloc_ex = copy_into!(container, ex)
+    return container !== nothing && prealloc_ex == ex && prealloc_ex !== container
+end
 function _check_count_nodes(ex::AbstractExpression)
     return count_nodes(ex) isa Int64
 end
@@ -156,6 +161,7 @@ ei_components = (
         with_metadata = "returns the expression with different metadata" => _check_with_metadata,
     ),
     optional = (
+        copy_into! = "copies an expression into a preallocated container" => _check_copy_into!,
         count_nodes = "counts the number of nodes in the expression tree" => _check_count_nodes,
         count_constant_nodes = "counts the number of constant nodes in the expression tree" => _check_count_constant_nodes,
         count_depth = "calculates the depth of the expression tree" => _check_count_depth,
@@ -260,6 +266,11 @@ function _check_tree_mapreduce(tree::AbstractExpressionNode)
 end
 
 ## optional
+function _check_copy_into!(tree::AbstractExpressionNode)
+    container = allocate_container(tree)
+    prealloc_tree = copy_into!(container, tree)
+    return container !== nothing && prealloc_tree == tree && prealloc_tree !== container
+end
 function _check_leaf_copy(tree::AbstractExpressionNode)
     tree.degree != 0 && return true
     return leaf_copy(tree) isa typeof(tree)
@@ -372,6 +383,7 @@ ni_components = (
         tree_mapreduce = "applies a function across the tree" => _check_tree_mapreduce
     ),
     optional = (
+        copy_into! = "copies a node into a preallocated container" => _check_copy_into!,
         leaf_copy = "copies a leaf node" => _check_leaf_copy,
         leaf_copy_into! = "copies a leaf node in-place" => _check_leaf_copy_into!,
         leaf_convert = "converts a leaf node" => _check_leaf_convert,

src/NodePreallocation.jl

@@ -9,7 +9,7 @@ using ..NodeModule:
     set_node!
 
 """
-    copy_into!(dest::AbstractArray{N}, src::N) where {BS,N<:AbstractExpressionNode}
+    copy_into!(dest::AbstractArray{N}, src::N) where {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.
@@ -29,8 +29,7 @@ function copy_into!(
         ((p, c::Vararg{Any,M}) where {M}) ->
             branch_copy_into!(@inbounds(dest[_ref.x += 1]), p, c...),
         src,
-        N;
-        break_sharing=Val(BS),
+        N,
     )
 end
 function leaf_copy_into!(dest::N, src::N) where {N<:AbstractExpressionNode}

src/ParametricExpression.jl

@@ -4,9 +4,8 @@ using DispatchDoctor: @stable, @unstable
 using ChainRulesCore: ChainRulesCore as CRC, NoTangent, @thunk
 
 using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
-using ..NodeModule:
-    AbstractExpressionNode, Node, tree_mapreduce, with_contents, with_metadata
-using ..ExpressionModule: AbstractExpression, Metadata
+using ..NodeModule: AbstractExpressionNode, Node, tree_mapreduce
+using ..ExpressionModule: AbstractExpression, Metadata, with_contents, with_metadata
 using ..ChainRulesModule: NodeTangent
 
 import ..NodeModule:
@@ -17,9 +16,8 @@ import ..NodeModule:
     leaf_convert,
     leaf_hash,
     leaf_equal,
-    set_node!,
-    copy_into!,
-    allocate_container
+    set_node!
+import ..NodePreallocationModule: copy_into!, allocate_container
 import ..NodeUtilsModule:
     count_constant_nodes,
     index_constant_nodes,

test/test_parametric_expression.jl

@@ -26,7 +26,7 @@ end
     using Interfaces: test
 
     ex = @parse_expression(
-        x + y + p1 * p2,
+        x + y + p1 * p2 + 1.5,
         binary_operators = [+, -, *, /],
         variable_names = ["x", "y"],
         node_type = ParametricNode,