feat: complete node interface for n-arity

Author: MilesCranmer

src/Interfaces.jl

@@ -12,6 +12,7 @@ using ..NodeModule:
     constructorof,
     default_allocator,
     with_type_parameters,
+    children,
     leaf_copy,
     leaf_convert,
     leaf_hash,
@@ -248,8 +249,8 @@ function _check_eltype(tree::AbstractExpressionNode{T}) where {T}
 end
 function _check_with_type_parameters(tree::AbstractExpressionNode{T}) where {T}
     N = typeof(tree)
-    NT = with_type_parameters(Base.typename(N).wrapper, eltype(tree))
-    return NT == typeof(tree)
+    Nf16 = with_type_parameters(N, Float16)
+    return Nf16 <: AbstractExpressionNode{Float16}
 end
 function _check_default_allocator(tree::AbstractExpressionNode)
     N = Base.typename(typeof(tree)).wrapper
@@ -299,35 +300,21 @@ function _check_leaf_equal(tree::AbstractExpressionNode)
     return leaf_equal(tree, copy(tree))
 end
 function _check_branch_copy(tree::AbstractExpressionNode)
-    if tree.degree == 0
-        return true
-    elseif tree.degree == 1
-        return branch_copy(tree, tree.l) isa typeof(tree)
-    else
-        return branch_copy(tree, tree.l, tree.r) isa typeof(tree)
-    end
+    tree.degree == 0 && return true
+    return branch_copy(tree, children(tree, Val(tree.degree))...) isa typeof(tree)
 end
 function _check_branch_copy_into!(tree::AbstractExpressionNode{T}) where {T}
-    if tree.degree == 0
-        return true
-    end
+    tree.degree == 0 && return true
     new_branch = constructorof(typeof(tree))(; val=zero(T))
-    if tree.degree == 1
-        ret = branch_copy_into!(new_branch, tree, copy(tree.l))
-        return new_branch == tree && ret === new_branch
-    else
-        ret = branch_copy_into!(new_branch, tree, copy(tree.l), copy(tree.r))
-        return new_branch == tree && ret === new_branch
-    end
+    ret = branch_copy_into!(
+        new_branch, tree, map(copy, children(tree, Val(tree.degree)))...
+    )
+    return new_branch == tree && ret === new_branch
 end
 function _check_branch_convert(tree::AbstractExpressionNode)
-    if tree.degree == 0
-        return true
-    elseif tree.degree == 1
-        return branch_convert(typeof(tree), tree, tree.l) isa typeof(tree)
-    else
-        return branch_convert(typeof(tree), tree, tree.l, tree.r) isa typeof(tree)
-    end
+    tree.degree == 0 && return true
+    return branch_convert(typeof(tree), tree, children(tree, Val(tree.degree))...) isa
+           typeof(tree)
 end
 function _check_branch_hash(tree::AbstractExpressionNode)
     tree.degree == 0 && return true

src/Node.jl

@@ -225,7 +225,7 @@ end
 @inline children(node::AbstractNode) = node.children
 @inline function children(node::AbstractNode, ::Val{n}) where {n}
     cs = children(node)
-    return ntuple(i -> cs[i], Val(n))
+    return ntuple(i -> cs[i], Val(Int(n)))
 end
 
 ################################################################################

test/test_node_interface.jl

@@ -33,3 +33,27 @@
         ],
     )
 end
+
+@testitem "Node interface on n-arity nodes" begin
+    using DynamicExpressions
+    using DynamicExpressions: NodeInterface
+    using Interfaces: Interfaces
+
+    for D in (3, 4, 5)
+        x = [Node{Float64,D}(; feature=i) for i in 1:3]
+        operator_tuple = ((sin, cos, exp), (+, *, /, -), (fma, clamp), (max, min), ())
+        operators = OperatorEnum(operator_tuple[1:D])
+        DynamicExpressions.OperatorEnumConstructionModule.empty_all_globals!()
+        let tree = Node{Float64,D}(; op=2, children=(x[1], x[2]))  # *
+            if D > 2
+                fma_idx = 1
+                tree = Node{Float64,D}(; op=fma_idx, children=(tree, x[1], x[2]))  # fma
+            end
+            if D > 3
+                idx_max = 1
+                tree = Node{Float64,D}(; op=idx_max, children=(tree, x[1], x[2], x[3]))  # max
+            end
+            @test Interfaces.test(NodeInterface, Node, tree)
+        end
+    end
+end