Add unittests and docs for AbstractNode

Author: MilesCranmer

docs/src/types.md

@@ -83,3 +83,9 @@ You can create a copy of a node with `copy_node`:
 ```@docs
 copy_node(tree::Node)
 ```
+
+There is also an abstract type `AbstractNode` which is a supertype of `Node`:
+
+```@docs
+AbstractNode
+```

src/Equation.jl

@@ -13,8 +13,13 @@ Abstract type for binary trees. Must have the following fields:
 - `degree::Integer`: Degree of the node. Either 0, 1, or 2. If 1,
     then `l` needs to be defined as the left child. If 2,
     then `r` also needs to be defined as the right child.
-- `l::AbstractNode`: Left child of the current node.
-- `r::AbstractNode`: Right child of the current node.
+- `l::AbstractNode`: Left child of the current node. Should only be
+    defined if `degree >= 1`; otherwise, leave it undefined (see the
+    the constructors of `Node{T}` for an example).
+    Don't use `nothing` to represent an undefined value
+    as it will incur a large performance penalty.
+- `r::AbstractNode`: Right child of the current node. Should only
+    be defined if `degree == 2`.
 """
 abstract type AbstractNode end
 

test/test_custom_node_type.jl

@@ -0,0 +1,37 @@
+using DynamicExpressions
+using Test
+
+mutable struct MyCustomNode{A,B} <: AbstractNode
+    degree::Int
+    val1::A
+    val2::B
+    l::MyCustomNode{A,B}
+    r::MyCustomNode{A,B}
+
+    MyCustomNode(val1, val2) = new{typeof(val1),typeof(val2)}(0, val1, val2)
+    MyCustomNode(val1, val2, l) = new{typeof(val1),typeof(val2)}(1, val1, val2, l)
+    MyCustomNode(val1, val2, l, r) = new{typeof(val1),typeof(val2)}(2, val1, val2, l, r)
+end
+
+node1 = MyCustomNode(1.0, 2)
+
+@test typeof(node1) == MyCustomNode{Float64,Int}
+@test node1.degree == 0
+@test count_depth(node1) == 1
+@test count_nodes(node1) == 1
+
+node2 = MyCustomNode(1.5, 3, node1)
+
+@test typeof(node2) == MyCustomNode{Float64,Int}
+@test node2.degree == 1
+@test node2.l.degree == 0
+@test count_depth(node2) == 2
+@test count_nodes(node2) == 2
+
+node2 = MyCustomNode(1.5, 3, node1, node1)
+
+@test count_depth(node2) == 2
+@test count_nodes(node2) == 3
+@test sum(t -> t.val1, node2) == 1.5 + 1.0 + 1.0
+@test sum(t -> t.val2, node2) == 3 + 2 + 2
+@test count(t -> t.degree == 0, node2) == 2

test/unittest.jl

@@ -91,3 +91,7 @@ end
 @safetestset "Test helpers break upon redefining" begin
     include("test_safe_helpers.jl")
 end
+
+@safetestset "Test custom node type" begin
+    include("test_custom_node_type.jl")
+end