Expand macro to simplify debugging

Author: MilesCranmer

src/tree_map.jl

@@ -97,12 +97,26 @@ function tree_mapreduce(
     preserve_sharing::Bool=false,
     result_type::Type{RT}=Nothing,
 ) where {T,N<:Node{T},F1<:Function,F2<:Function,G<:Function,RT}
-    preserve_sharing && return @with_memoization(
-        _tree_mapreduce(f_leaf, f_branch, op, tree), IdDict{N,RT}()
-    )
+    preserve_sharing && return _tree_mapreduce(f_leaf, f_branch, op, tree, IdDict{N,RT}())
     return _tree_mapreduce(f_leaf, f_branch, op, tree)
 end
-@memoize_on tree function _tree_mapreduce(
+function _tree_mapreduce(
+    f_leaf::F1, f_branch::F2, op::G, tree::Node, id_map
+) where {F1<:Function,F2<:Function,G<:Function}
+    get!(id_map, tree) do
+        if tree.degree == 0
+            return @inline(f_leaf(tree))
+        elseif tree.degree == 1
+            l = _tree_mapreduce(f_leaf, f_branch, op, tree.l, id_map)
+            return @inline(op(@inline(f_branch(tree)), l))
+        else
+            l = _tree_mapreduce(f_leaf, f_branch, op, tree.l, id_map)
+            r = _tree_mapreduce(f_leaf, f_branch, op, tree.r, id_map)
+            return @inline(op(@inline(f_branch(tree)), l, r))
+        end
+    end
+end
+function _tree_mapreduce(
     f_leaf::F1, f_branch::F2, op::G, tree::Node
 ) where {F1<:Function,F2<:Function,G<:Function}
     if tree.degree == 0