feat: rename to allocate_container and copy_into!

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -9,6 +9,7 @@ using DispatchDoctor: @stable, @unstable
     include("OperatorEnum.jl")
     include("Node.jl")
     include("NodeUtils.jl")
+    include("NodePreallocation.jl")
     include("Strings.jl")
     include("Evaluate.jl")
     include("EvaluateDerivative.jl")
@@ -41,11 +42,11 @@ import .ValueInterfaceModule:
     GraphNode,
     Node,
     copy_node,
-    copy_node!,
     set_node!,
     tree_mapreduce,
     filter_map,
     filter_map!
+import .NodePreallocationModule: allocate_container, copy_into!
 import .NodeModule:
     constructorof,
     with_type_parameters,

src/Expression.jl

@@ -317,14 +317,6 @@ 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,
@@ -510,4 +502,21 @@ function (ex::AbstractExpression)(
     return get_tree(ex)(X, get_operators(ex, operators); kws...)
 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.
+function copy_into!(::Nothing, src::AbstractExpression)
+    return copy(src)
+end
+function allocate_container(::AbstractExpression, ::Union{Nothing,Integer}=nothing)
+    return nothing
+end
+function allocate_container(prototype::Expression, n::Union{Nothing,Integer}=nothing)
+    return (; tree=allocate_container(get_contents(prototype), n))
+end
+function copy_into!(dest::NamedTuple, src::Expression)
+    tree = copy_into!(dest.tree, get_contents(src))
+    return with_contents(src, tree)
+end
+
 end

src/Interfaces.jl

@@ -13,12 +13,12 @@ using ..NodeModule:
     default_allocator,
     with_type_parameters,
     leaf_copy,
-    leaf_copy!,
+    leaf_copy_into!,
     leaf_convert,
     leaf_hash,
     leaf_equal,
     branch_copy,
-    branch_copy!,
+    branch_copy_into!,
     branch_convert,
     branch_hash,
     branch_equal,
@@ -264,10 +264,10 @@ function _check_leaf_copy(tree::AbstractExpressionNode)
     tree.degree != 0 && return true
     return leaf_copy(tree) isa typeof(tree)
 end
-function _check_leaf_copy!(tree::AbstractExpressionNode{T}) where {T}
+function _check_leaf_copy_into!(tree::AbstractExpressionNode{T}) where {T}
     tree.degree != 0 && return true
     new_leaf = constructorof(typeof(tree))(; val=zero(T))
-    ret = leaf_copy!(new_leaf, tree)
+    ret = leaf_copy_into!(new_leaf, tree)
     return new_leaf == tree && ret === new_leaf
 end
 function _check_leaf_convert(tree::AbstractExpressionNode)
@@ -292,16 +292,16 @@ function _check_branch_copy(tree::AbstractExpressionNode)
         return branch_copy(tree, tree.l, tree.r) isa typeof(tree)
     end
 end
-function _check_branch_copy!(tree::AbstractExpressionNode{T}) where {T}
+function _check_branch_copy_into!(tree::AbstractExpressionNode{T}) where {T}
     if tree.degree == 0
         return true
     end
     new_branch = constructorof(typeof(tree))(; val=zero(T))
     if tree.degree == 1
-        ret = branch_copy!(new_branch, tree, copy(tree.l))
+        ret = branch_copy_into!(new_branch, tree, copy(tree.l))
         return new_branch == tree && ret === new_branch
     else
-        ret = branch_copy!(new_branch, tree, copy(tree.l), copy(tree.r))
+        ret = branch_copy_into!(new_branch, tree, copy(tree.l), copy(tree.r))
         return new_branch == tree && ret === new_branch
     end
 end
@@ -373,12 +373,12 @@ ni_components = (
     ),
     optional = (
         leaf_copy = "copies a leaf node" => _check_leaf_copy,
-        leaf_copy! = "copies a leaf node in-place" => _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,
         leaf_hash = "computes the hash of a leaf node" => _check_leaf_hash,
         leaf_equal = "checks equality of two leaf nodes" => _check_leaf_equal,
         branch_copy = "copies a branch node" => _check_branch_copy,
-        branch_copy! = "copies a branch node in-place" => _check_branch_copy!,
+        branch_copy_into! = "copies a branch node in-place" => _check_branch_copy_into!,
         branch_convert = "converts a branch node" => _check_branch_convert,
         branch_hash = "computes the hash of a branch node" => _check_branch_hash,
         branch_equal = "checks equality of two branch nodes" => _check_branch_equal,
@@ -419,7 +419,7 @@ ni_description = (
     [Arguments()]
 )
 @implements(
-    NodeInterface{all_ni_methods_except((:leaf_copy!, :branch_copy!))},
+    NodeInterface{all_ni_methods_except((:leaf_copy_into!, :branch_copy_into!))},
     GraphNode,
     [Arguments()]
 )

src/NodePreallocation.jl

@@ -0,0 +1,68 @@
+module NodePreallocationModule
+
+using ..NodeModule:
+    AbstractExpressionNode,
+    with_type_parameters,
+    tree_mapreduce,
+    leaf_copy,
+    branch_copy,
+    set_node!
+
+"""
+    copy_into!(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_into!(
+    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_into!(@inbounds(dest[_ref.x += 1]), leaf),
+        identity,
+        ((p, c::Vararg{Any,M}) where {M}) ->
+            branch_copy_into!(@inbounds(dest[_ref.x += 1]), p, c...),
+        src,
+        N;
+        break_sharing=Val(BS),
+    )
+end
+function leaf_copy_into!(dest::N, src::N) where {N<:AbstractExpressionNode}
+    set_node!(dest, src)
+    return dest
+end
+function branch_copy_into!(
+    dest::N, src::N, children::Vararg{N,M}
+) where {N<:AbstractExpressionNode,M}
+    dest.degree = M
+    dest.op = src.op
+    dest.l = children[1]
+    if M == 2
+        dest.r = children[2]
+    end
+    return dest
+end
+
+"""
+    allocate_container(prototype::AbstractExpressionNode, n=nothing)
+
+Preallocate an array of `n` empty nodes matching the type of `prototype`.
+If `n` is not provided, it will be computed from `length(prototype)`.
+
+A given return value of this will be passed to `copy_into!` as the first argument,
+so it should be compatible.
+"""
+function allocate_container(
+    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
+
+end

src/ParametricExpression.jl

@@ -14,13 +14,12 @@ import ..NodeModule:
     with_type_parameters,
     preserve_sharing,
     leaf_copy,
-    leaf_copy!,
     leaf_convert,
     leaf_hash,
     leaf_equal,
-    branch_copy!,
-    copy_node!,
-    preallocate_expression
+    set_node!,
+    copy_into!,
+    allocate_container
 import ..NodeUtilsModule:
     count_constant_nodes,
     index_constant_nodes,
@@ -447,16 +446,16 @@ end
         return node_type(; val=ex)
     end
 end
-function preallocate_expression(
+function allocate_container(
     prototype::ParametricExpression, n::Union{Nothing,Integer}=nothing
 )
     return (;
-        tree=preallocate_expression(get_contents(prototype), n),
+        tree=allocate_container(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))
+function copy_into!(dest::NamedTuple, src::ParametricExpression)
+    new_tree = copy_into!(dest.tree, get_contents(src))
     metadata = get_metadata(src)
     new_parameters = dest.parameters
     new_parameters .= metadata.parameters

src/base.jl

@@ -485,81 +485,6 @@ function branch_copy(t::N, children::Vararg{Any,M}) where {T,N<:AbstractExpressi
     return constructorof(N)(T; op=t.op, children)
 end
 
-# In-place versions
-
-"""
-    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; 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),
-        identity,
-        ((p, c::Vararg{Any,M}) where {M}) ->
-            branch_copy!(@inbounds(dest[_ref.x += 1]), p, c...),
-        src,
-        N;
-        break_sharing=Val(BS),
-    )
-end
-function leaf_copy!(dest::N, src::N) where {T,N<:AbstractExpressionNode{T}}
-    dest.degree = 0
-    if src.constant
-        dest.constant = true
-        dest.val = src.val
-    else
-        dest.constant = false
-        dest.feature = src.feature
-    end
-    return dest
-end
-function branch_copy!(
-    dest::N, src::N, children::Vararg{N,M}
-) where {T,N<:AbstractExpressionNode{T},M}
-    dest.degree = M
-    dest.op = src.op
-    dest.l = children[1]
-    if M == 2
-        dest.r = children[2]
-    end
-    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))
 

test/test_copy_inplace.jl

@@ -1,6 +1,6 @@
-@testitem "copy_node! - random trees" begin
+@testitem "copy_into! - random trees" begin
     using DynamicExpressions
-    using DynamicExpressions: copy_node!
+    using DynamicExpressions: copy_into!
     include("tree_gen_utils.jl")
 
     operators = OperatorEnum(; binary_operators=[+, *, /], unary_operators=[sin, cos])
@@ -15,7 +15,7 @@
         orig_nodes = dest_array[(n_nodes + 1):end]  # Save reference to unused nodes
 
         ref = Ref(0)
-        result = copy_node!(dest_array, tree; ref)
+        result = copy_into!(dest_array, tree; ref)
 
         @test ref[] == n_nodes  # Increment once per node
 
@@ -35,17 +35,17 @@
     end
 end
 
-@testitem "copy_node! - leaf nodes" begin
+@testitem "copy_into! - leaf nodes" begin
     using DynamicExpressions
-    using DynamicExpressions: copy_node!
+    using DynamicExpressions: copy_into!
 
     leaf_constant = Node{Float64}(; val=1.0)
     leaf_feature = Node{Float64}(; feature=1)
 
     for leaf in [leaf_constant, leaf_feature]
         dest_array = [Node{Float64}() for _ in 1:1]
         ref = Ref(0)
-        result = copy_node!(dest_array, leaf; ref=ref)
+        result = copy_into!(dest_array, leaf; ref=ref)
         @test ref[] == 1
         @test result == leaf
         @test result === dest_array[1]