Permit n-argument operators

CHANGELOG.md

@@ -0,0 +1,38 @@
+# CHANGELOG
+
+## 2.0.0
+
+- Nodes can now have arbitrary numbers of children, not just binary trees. This was designed in such a way to have identical performance to the previous version. Essentially, the `Node{T,D}` type is a wrapper around a tuple of `D` children.
+    - Note that `Node{T}` automatically converts to `Node{T,2}` in many contexts, for backwards compatibility.
+- Node type signature changed from `Node{T}` to `Node{T,D}`. Usually `D` is 2.
+- Direct property access `.l` and `.r` is automatically forwarded to `get_child(tree, 1)` and `get_child(tree, 2)`,
+    but it is recommended to use the generic accessors instead.
+- Similarly, `.l = child` should be replaced with `set_child!(tree, child, 1)` and similar for `.r`.
+- All internal code migrated to use generic accessors.
+
+### Backwards Compatibility
+
+- Existing `.l` and `.r` access continues to work without warnings
+
+### Breaking Changes
+
+The only breaking change is if:
+
+1. You have any types that are subtyped to `<:AbstractExpressionNode{T}` or `<:AbstractNode{T}`. These should now be subtyped to `<:AbstractExpressionNode{T,2}` or `<:AbstractNode{T,2}`. You may also allow a `D` parameter in case you want to support higher-arity trees.
+2. You assume a tree has type, e.g., `=== Node{T}`, rather than `<: Node{T}`. So any methods dispatched to `::Type{Node{T}}` will also break. (To be safe you should always use a form `<: Node{T}` in case of future type changes - in any library.)
+3. You assume `tree.degree <= 2` in conditional logic, and your code interacts with a tree that is _not_ a binary tree. For example, the following pattern was common before this change:
+
+    ```julia
+    if tree.degree == 0
+        #= operations on leaf node =#
+    elseif tree.degree == 1
+        #= operations on unary node =#
+    else
+        # BAD: ASSUMED TO BE BINARY
+        #= operations on binary node, using `.l` and `.r` only =#
+    end
+    ```
+
+    This will obviously break if you pass a tree that is not binary, such as `tree::Node{T,3}`.
+    - To fix this, you can use the `get_children` function to get the children of the tree as a tuple of `D` children, and then index up to `tree.degree`. 
+    - Inside DynamicExpressions, we commonly use `Base.Cartesian.@nif` to generate code for different degrees, to avoid any unstable types.

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "1.10.2"
+version = "2.0.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/api.md

@@ -60,7 +60,7 @@ When using these node constructors, types will automatically be promoted.
 You can convert the type of a node using `convert`:
 
 ```@docs
-convert(::Type{AbstractExpressionNode{T1}}, tree::AbstractExpressionNode{T2}) where {T1, T2}
+convert(::Type{N1}, tree::N2) where {T1,T2,D1,D2,N1<:AbstractExpressionNode{T1,D1},N2<:AbstractExpressionNode{T2,D2}}
 ```
 
 You can set a `tree` (in-place) with `set_node!`:
@@ -75,6 +75,42 @@ You can create a copy of a node with `copy_node`:
 copy_node
 ```
 
+## Generic Node Accessors
+
+For working with nodes of arbitrary arity:
+
+```@docs
+get_child
+set_child!
+get_children
+set_children!
+```
+
+Examples:
+
+```julia
+# Define operators including ternary
+my_ternary(x, y, z) = x + y * z
+operators = OperatorEnum(((sin,), (+, *), (my_ternary,)))  # (unary, binary, ternary)
+
+tree = Node{Float64,3}(; op=1, children=(Node{Float64,3}(; val=1.0), Node{Float64,3}(; val=2.0)))
+new_child = Node{Float64,3}(; val=3.0)
+
+left_child = get_child(tree, 1)
+right_child = get_child(tree, 2)
+
+set_child!(tree, new_child, 1)
+
+children = get_children(tree)
+left, right = get_children(tree, Val(2))  # type stable
+
+# Transform to ternary operation
+child1, child2, child3 = Node{Float64,3}(; val=4.0), Node{Float64,3}(; val=5.0), Node{Float64,3}(; val=6.0)
+set_children!(tree, (child1, child2, child3))
+tree.op = 1  # my_ternary
+tree.degree = 3
+```
+
 ## Graph Nodes
 
 You can describe an equation as a *graph* rather than a tree
@@ -109,7 +145,7 @@ This means that we only need to change it once
 to have changes propagate across the expression:
 
 ```julia
-julia> y.r.val *= 0.9
+julia> get_child(y, 2).val *= 0.9
 1.35
 
 julia> z

docs/src/utils.md

@@ -15,7 +15,6 @@ mapreduce(f::F, op::G, tree::AbstractNode; return_type, f_on_shared, break_shari
 any(f::F, tree::AbstractNode) where {F<:Function}
 all(f::F, tree::AbstractNode) where {F<:Function}
 map(f::F, tree::AbstractNode, result_type::Type{RT}=Nothing; break_sharing::Val=Val(false)) where {F<:Function,RT}
-convert(::Type{<:AbstractExpressionNode{T1}}, n::AbstractExpressionNode{T2}) where {T1,T2}
 hash(tree::AbstractExpressionNode{T}, h::UInt; break_sharing::Val=Val(false)) where {T}
 ```
 

ext/DynamicExpressionsBumperExt.jl

@@ -5,8 +5,7 @@ using DynamicExpressions:
     OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array, EvalOptions
 using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
-import DynamicExpressions.ExtensionInterfaceModule:
-    bumper_eval_tree_array, bumper_kern1!, bumper_kern2!
+import DynamicExpressions.ExtensionInterfaceModule: bumper_eval_tree_array, bumper_kern!
 
 function bumper_eval_tree_array(
     tree::AbstractExpressionNode{T},
@@ -37,8 +36,7 @@ function bumper_eval_tree_array(
             branch_node -> branch_node,
             # In the evaluation kernel, we combine the branch nodes
             # with the arrays created by the leaf nodes:
-            ((args::Vararg{Any,M}) where {M}) ->
-                dispatch_kerns!(operators, args..., eval_options),
+            KernelDispatcher(operators, eval_options),
             tree;
             break_sharing=Val(true),
         )
@@ -49,63 +47,44 @@ function bumper_eval_tree_array(
     return (result, all_ok[])
 end
 
-function dispatch_kerns!(
-    operators, branch_node, cumulator, eval_options::EvalOptions{<:Any,true,early_exit}
-) where {early_exit}
-    cumulator.ok || return cumulator
-
-    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, eval_options)
-    return ResultOk(out, early_exit ? is_valid_array(out) : true)
-end
-function dispatch_kerns!(
-    operators,
-    branch_node,
-    cumulator1,
-    cumulator2,
-    eval_options::EvalOptions{<:Any,true,early_exit},
-) where {early_exit}
-    cumulator1.ok || return cumulator1
-    cumulator2.ok || return cumulator2
-
-    out = dispatch_kern2!(
-        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, eval_options
-    )
-    return ResultOk(out, early_exit ? is_valid_array(out) : true)
+struct KernelDispatcher{O<:OperatorEnum,E<:EvalOptions{<:Any,true,<:Any}} <: Function
+    operators::O
+    eval_options::E
 end
 
-@generated function dispatch_kern1!(unaops, op_idx, cumulator, eval_options::EvalOptions)
-    nuna = counttuple(unaops)
+@generated function (kd::KernelDispatcher{<:Any,<:EvalOptions{<:Any,true,early_exit}})(
+    branch_node, inputs::Vararg{Any,degree}
+) where {degree,early_exit}
     quote
-        Base.@nif(
-            $nuna,
-            i -> i == op_idx,
-            i -> let op = unaops[i]
-                return bumper_kern1!(op, cumulator, eval_options)
-            end,
-        )
+        Base.Cartesian.@nexprs($degree, i -> inputs[i].ok || return inputs[i])
+        cumulators = Base.Cartesian.@ntuple($degree, i -> inputs[i].x)
+        out = dispatch_kerns!(kd.operators, branch_node, cumulators, kd.eval_options)
+        return ResultOk(out, early_exit ? is_valid_array(out) : true)
     end
 end
-@generated function dispatch_kern2!(
-    binops, op_idx, cumulator1, cumulator2, eval_options::EvalOptions
-)
-    nbin = counttuple(binops)
+@generated function dispatch_kerns!(
+    operators::OperatorEnum{OPS},
+    branch_node,
+    cumulators::Tuple{Vararg{Any,degree}},
+    eval_options::EvalOptions,
+) where {OPS,degree}
+    nops = length(OPS.types[degree].types)
     quote
-        Base.@nif(
-            $nbin,
+        op_idx = branch_node.op
+        Base.Cartesian.@nif(
+            $nops,
             i -> i == op_idx,
-            i -> let op = binops[i]
-                return bumper_kern2!(op, cumulator1, cumulator2, eval_options)
-            end,
+            i -> bumper_kern!(operators[$degree][i], cumulators, eval_options)
         )
     end
 end
-function bumper_kern1!(op::F, cumulator, ::EvalOptions{false,true}) where {F}
-    @. cumulator = op(cumulator)
-    return cumulator
-end
-function bumper_kern2!(op::F, cumulator1, cumulator2, ::EvalOptions{false,true}) where {F}
-    @. cumulator1 = op(cumulator1, cumulator2)
-    return cumulator1
+
+function bumper_kern!(
+    op::F, cumulators::Tuple{Vararg{Any,degree}}, ::EvalOptions{false,true,early_exit}
+) where {F,degree,early_exit}
+    cumulator_1 = first(cumulators)
+    @. cumulator_1 = op(cumulators...)
+    return cumulator_1
 end
 
 end