feat: create in-place copy operator

Author: MilesCranmer

src/Interfaces.jl

@@ -13,10 +13,12 @@ using ..NodeModule:
     default_allocator,
     with_type_parameters,
     leaf_copy,
+    leaf_copy!,
     leaf_convert,
     leaf_hash,
     leaf_equal,
     branch_copy,
+    branch_copy!,
     branch_convert,
     branch_hash,
     branch_equal,
@@ -262,6 +264,12 @@ 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}
+    tree.degree != 0 && return true
+    new_leaf = constructorof(typeof(tree))(; val=zero(T))
+    ret = leaf_copy!(new_leaf, tree)
+    return new_leaf == tree && ret === new_leaf
+end
 function _check_leaf_convert(tree::AbstractExpressionNode)
     tree.degree != 0 && return true
     return leaf_convert(typeof(tree), tree) isa typeof(tree) &&
@@ -284,6 +292,19 @@ 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}
+    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))
+        return new_branch == tree && ret === new_branch
+    else
+        ret = branch_copy!(new_branch, tree, copy(tree.l), copy(tree.r))
+        return new_branch == tree && ret === new_branch
+    end
+end
 function _check_branch_convert(tree::AbstractExpressionNode)
     if tree.degree == 0
         return true
@@ -352,10 +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_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_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,
@@ -396,7 +419,7 @@ ni_description = (
     [Arguments()]
 )
 @implements(
-    NodeInterface{all_ni_methods_except(())},
+    NodeInterface{all_ni_methods_except((:leaf_copy!, :branch_copy!))},
     GraphNode,
     [Arguments()]
 )

src/ParametricExpression.jl

@@ -13,9 +13,11 @@ import ..NodeModule:
     with_type_parameters,
     preserve_sharing,
     leaf_copy,
+    leaf_copy!,
     leaf_convert,
     leaf_hash,
-    leaf_equal
+    leaf_equal,
+    branch_copy!
 import ..NodeUtilsModule:
     count_constant_nodes,
     index_constant_nodes,
@@ -122,6 +124,22 @@ function leaf_copy(t::ParametricNode{T}) where {T}
         return n
     end
 end
+function leaf_copy!(dest::N, src::N) where {T,N<:ParametricNode{T}}
+    dest.degree = 0
+    if src.constant
+        dest.constant = true
+        dest.val = src.val
+    elseif !src.is_parameter
+        dest.constant = false
+        dest.is_parameter = false
+        dest.feature = src.feature
+    else
+        dest.constant = false
+        dest.is_parameter = true
+        dest.parameter = src.parameter
+    end
+    return dest
+end
 function leaf_convert(::Type{N}, t::ParametricNode) where {T,N<:ParametricNode{T}}
     if t.constant
         return constructorof(N)(T; val=convert(T, t.val))

src/base.jl

@@ -485,6 +485,53 @@ 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; break_sharing::Val{BS}=Val(false)) 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}
+    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
+
 """
     copy(tree::AbstractExpressionNode; break_sharing::Val=Val(false))