Fix speed regression due to type union

Author: MilesCranmer

src/Equation.jl

@@ -82,7 +82,7 @@ function Base.convert(
     get!(id_map, tree) do
         if tree.degree == 0
             if tree.constant
-                Node(0, tree.constant, convert(T1, tree.val))
+                Node(0, tree.constant, convert(T1, (tree.val::T2)))
             else
                 Node(T1, 0, tree.constant, nothing, tree.feature)
             end
@@ -180,7 +180,7 @@ function set_node!(tree::Node{T}, new_tree::Node{T}) where {T}
     if new_tree.degree == 0
         tree.constant = new_tree.constant
         if new_tree.constant
-            tree.val = new_tree.val
+            tree.val = new_tree.val::T
         else
             tree.feature = new_tree.feature
         end
@@ -215,7 +215,7 @@ end
 function copy_node_break_topology(tree::Node{T})::Node{T} where {T}
     if tree.degree == 0
         if tree.constant
-            Node(; val=copy(tree.val))
+            Node(; val=copy(tree.val::T))
         else
             Node(T; feature=copy(tree.feature))
         end
@@ -248,7 +248,7 @@ function copy_node_with_topology(
     get!(id_map, tree) do
         if tree.degree == 0
             if tree.constant
-                Node(; val=copy(tree.val))
+                Node(; val=copy(tree.val::T))
             else
                 Node(T; feature=copy(tree.feature))
             end
@@ -355,11 +355,11 @@ function print_tree(
     return println(string_tree(tree, operators; varMap=varMap))
 end
 
-function Base.hash(tree::Node)::UInt
+function Base.hash(tree::Node{T})::UInt where {T}
     if tree.degree == 0
         if tree.constant
             # tree.val used.
-            return hash((0, tree.val))
+            return hash((0, tree.val::T))
         else
             # tree.feature used.
             return hash((1, tree.feature))

src/EvaluateEquation.jl

@@ -156,7 +156,7 @@ function deg0_eval(
 )::Tuple{AbstractVector{T},Bool} where {T<:Real}
     n = size(cX, 2)
     if tree.constant
-        return (fill(tree.val, n), true)
+        return (fill(tree.val::T, n), true)
     else
         return (cX[tree.feature, :], true)
     end
@@ -173,8 +173,8 @@ function deg1_l2_ll0_lr0_eval(
     op = operators.unaops[op_idx]
     op_l = operators.binops[op_l_idx]
     if tree.l.l.constant && tree.l.r.constant
-        val_ll = tree.l.l.val
-        val_lr = tree.l.r.val
+        val_ll = tree.l.l.val::T
+        val_lr = tree.l.r.val::T
         @return_on_check val_ll T n
         @return_on_check val_lr T n
         x_l = op_l(val_ll, val_lr)::T
@@ -183,7 +183,7 @@ function deg1_l2_ll0_lr0_eval(
         @return_on_check x T n
         return (fill(x, n), true)
     elseif tree.l.l.constant
-        val_ll = tree.l.l.val
+        val_ll = tree.l.l.val::T
         @return_on_check val_ll T n
         feature_lr = tree.l.r.feature
         cumulator = Array{T,1}(undef, n)
@@ -195,7 +195,7 @@ function deg1_l2_ll0_lr0_eval(
         return (cumulator, true)
     elseif tree.l.r.constant
         feature_ll = tree.l.l.feature
-        val_lr = tree.l.r.val
+        val_lr = tree.l.r.val::T
         @return_on_check val_lr T n
         cumulator = Array{T,1}(undef, n)
         @inbounds @simd for j in 1:n
@@ -229,7 +229,7 @@ function deg1_l1_ll0_eval(
     op = operators.unaops[op_idx]
     op_l = operators.unaops[op_l_idx]
     if tree.l.l.constant
-        val_ll = tree.l.l.val
+        val_ll = tree.l.l.val::T
         @return_on_check val_ll T n
         x_l = op_l(val_ll)::T
         @return_on_check x_l T n
@@ -254,16 +254,16 @@ function deg2_l0_r0_eval(
     n = size(cX, 2)
     op = operators.binops[op_idx]
     if tree.l.constant && tree.r.constant
-        val_l = tree.l.val
+        val_l = tree.l.val::T
         @return_on_check val_l T n
-        val_r = tree.r.val
+        val_r = tree.r.val::T
         @return_on_check val_r T n
         x = op(val_l, val_r)::T
         @return_on_check x T n
         return (fill(x, n), true)
     elseif tree.l.constant
         cumulator = Array{T,1}(undef, n)
-        val_l = tree.l.val
+        val_l = tree.l.val::T
         @return_on_check val_l T n
         feature_r = tree.r.feature
         @inbounds @simd for j in 1:n
@@ -273,7 +273,7 @@ function deg2_l0_r0_eval(
     elseif tree.r.constant
         cumulator = Array{T,1}(undef, n)
         feature_l = tree.l.feature
-        val_r = tree.r.val
+        val_r = tree.r.val::T
         @return_on_check val_r T n
         @inbounds @simd for j in 1:n
             x = op(cX[feature_l, j], val_r)::T
@@ -300,7 +300,7 @@ function deg2_l0_eval(
     @return_on_nonfinite_array cumulator T n
     op = operators.binops[op_idx]
     if tree.l.constant
-        val = tree.l.val
+        val = tree.l.val::T
         @return_on_check val T n
         @inbounds @simd for j in 1:n
             x = op(val, cumulator[j])::T
@@ -325,7 +325,7 @@ function deg2_r0_eval(
     @return_on_nonfinite_array cumulator T n
     op = operators.binops[op_idx]
     if tree.r.constant
-        val = tree.r.val
+        val = tree.r.val::T
         @return_on_check val T n
         @inbounds @simd for j in 1:n
             x = op(cumulator[j], val)::T
@@ -361,7 +361,7 @@ function _eval_constant_tree(
 end
 
 @inline function deg0_eval_constant(tree::Node{T})::Tuple{T,Bool} where {T<:Real}
-    return tree.val, true
+    return tree.val::T, true
 end
 
 function deg1_eval_constant(
@@ -397,7 +397,7 @@ function differentiable_eval_tree_array(
     n = size(cX, 2)
     if tree.degree == 0
         if tree.constant
-            return (ones(T, n) .* convert(T, tree.val), true)
+            return (ones(T, n) .* (tree.val::T), true)
         else
             return (cX[tree.feature, :], true)
         end
@@ -476,7 +476,7 @@ function eval_tree_array(
 ) where {T,N}
     if tree.degree == 0
         if tree.constant
-            return tree.val, true
+            return (tree.val::T), true
         else
             if N == 1
                 return cX[tree.feature], true