refactor: more generic set_node! implementation

Author: MilesCranmer

src/Node.jl

@@ -420,17 +420,25 @@ Set every field of `tree` equal to the corresponding field of `new_tree`.
 function set_node!(tree::AbstractExpressionNode, new_tree::AbstractExpressionNode)
     tree.degree = new_tree.degree
     if new_tree.degree == 0
-        tree.constant = new_tree.constant
-        if new_tree.constant
-            tree.val = new_tree.val::eltype(new_tree)
-        else
-            tree.feature = new_tree.feature
-        end
+        set_leaf!(tree, new_tree)
     else
-        tree.op = new_tree.op
-        set_children!(tree, get_children(new_tree))
+        set_branch!(tree, new_tree)
     end
     return nothing
 end
+function set_leaf!(tree::AbstractExpressionNode, new_leaf::AbstractExpressionNode)
+    tree.constant = new_leaf.constant
+    if new_leaf.constant
+        tree.val = new_leaf.val::eltype(new_leaf)
+    else
+        tree.feature = new_leaf.feature
+    end
+    return nothing
+end
+function set_branch!(tree::AbstractExpressionNode, new_branch::AbstractExpressionNode)
+    tree.op = new_branch.op
+    set_children!(tree, get_children(new_branch))
+    return nothing
+end
 
 end

src/ParametricExpression.jl

@@ -21,7 +21,7 @@ import ..NodeModule:
     leaf_convert,
     leaf_hash,
     leaf_equal,
-    set_node!,
+    set_leaf!,
     @make_accessors
 import ..NodePreallocationModule: copy_into!, allocate_container
 import ..NodeUtilsModule:
@@ -145,24 +145,18 @@ function leaf_copy(t::ParametricNode{T}) where {T}
         return n
     end
 end
-function set_node!(tree::ParametricNode, new_tree::ParametricNode)
-    tree.degree = (deg = new_tree.degree)
-    if deg == 0
-        if new_tree.constant
-            tree.constant = true
-            tree.val = new_tree.val
-        elseif !new_tree.is_parameter
-            tree.constant = false
-            tree.is_parameter = false
-            tree.feature = new_tree.feature
-        else
-            tree.constant = false
-            tree.is_parameter = true
-            tree.parameter = new_tree.parameter
-        end
+function set_leaf!(tree::ParametricNode, new_leaf::ParametricNode)
+    if new_leaf.constant
+        tree.constant = true
+        tree.val = new_leaf.val
+    elseif !new_leaf.is_parameter
+        tree.constant = false
+        tree.is_parameter = false
+        tree.feature = new_leaf.feature
     else
-        tree.op = new_tree.op
-        set_children!(tree, get_children(new_tree))
+        tree.constant = false
+        tree.is_parameter = true
+        tree.parameter = new_leaf.parameter
     end
     return nothing
 end