Merge branch 'tree-map' into constant-optimization

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -21,15 +21,7 @@ end
 
 using Reexport: @reexport
 @reexport import .EquationModule:
-    Node,
-    string_tree,
-    print_tree,
-    copy_node,
-    set_node!,
-    map,
-    tree_mapreduce,
-    any,
-    filter_and_map
+    Node, string_tree, print_tree, copy_node, set_node!, tree_mapreduce, filter_map
 @reexport import .EquationUtilsModule:
     count_nodes,
     count_constants,

src/Equation.jl

@@ -1,7 +1,7 @@
 module EquationModule
 
 import ..OperatorEnumModule: AbstractOperatorEnum
-import ..UtilsModule: @generate_idmap, @use_idmap
+import ..UtilsModule: @memoize_on, @with_memoize
 
 const DEFAULT_NODE_TYPE = Float32
 
@@ -62,53 +62,7 @@ mutable struct Node{T}
 end
 ################################################################################
 
-include("tree_map.jl")
-
-"""
-    convert(::Type{Node{T1}}, n::Node{T2}) where {T1,T2}
-
-Convert a `Node{T2}` to a `Node{T1}`.
-This will recursively convert all children nodes to `Node{T1}`,
-using `convert(T1, tree.val)` at constant nodes.
-
-# Arguments
-- `::Type{Node{T1}}`: Type to convert to.
-- `tree::Node{T2}`: Node to convert.
-"""
-function Base.convert(
-    ::Type{Node{T1}}, tree::Node{T2}; preserve_sharing::Bool=false
-) where {T1,T2}
-    if T1 == T2
-        return tree
-    end
-    if preserve_sharing
-        @use_idmap(_convert(Node{T1}, tree), IdDict{Node{T2},Node{T1}}())
-    else
-        _convert(Node{T1}, tree)
-    end
-end
-
-@generate_idmap tree function _convert(::Type{Node{T1}}, tree::Node{T2}) where {T1,T2}
-    if tree.degree == 0
-        if tree.constant
-            val = tree.val::T2
-            if !(T2 <: T1)
-                # e.g., we don't want to convert Float32 to Union{Float32,Vector{Float32}}!
-                val = convert(T1, val)
-            end
-            Node(T1, 0, tree.constant, val)
-        else
-            Node(T1, 0, tree.constant, nothing, tree.feature)
-        end
-    elseif tree.degree == 1
-        l = _convert(Node{T1}, tree.l)
-        Node(1, tree.constant, nothing, tree.feature, tree.op, l)
-    else
-        l = _convert(Node{T1}, tree.l)
-        r = _convert(Node{T1}, tree.r)
-        Node(2, tree.constant, nothing, tree.feature, tree.op, l, r)
-    end
-end
+include("base.jl")
 
 """
     Node([::Type{T}]; val=nothing, feature::Int=nothing) where {T}
@@ -226,45 +180,6 @@ function set_node!(tree::Node{T}, new_tree::Node{T}) where {T}
     return nothing
 end
 
-"""
-    copy_node(tree::Node; preserve_sharing::Bool=false)
-
-Copy a node, recursively copying all children nodes.
-This is more efficient than the built-in copy.
-With `preserve_sharing=true`, this will also
-preserve linkage between a node and
-multiple parents, whereas without, this would create
-duplicate child node copies.
-
-id_map is a map from `objectid(tree)` to `copy(tree)`.
-We check against the map before making a new copy; otherwise
-we can simply reference the existing copy.
-[Thanks to Ted Hopp.](https://stackoverflow.com/questions/49285475/how-to-copy-a-full-non-binary-tree-including-loops)
-
-Note that this will *not* preserve loops in graphs.
-"""
-function copy_node(tree::Node{T}; preserve_sharing::Bool=false)::Node{T} where {T}
-    if preserve_sharing
-        @use_idmap(_copy_node(tree), IdDict{Node{T},Node{T}}())
-    else
-        _copy_node(tree)
-    end
-end
-
-@generate_idmap tree function _copy_node(tree::Node{T})::Node{T} where {T}
-    if tree.degree == 0
-        if tree.constant
-            Node(; val=copy(tree.val::T))
-        else
-            Node(T; feature=copy(tree.feature))
-        end
-    elseif tree.degree == 1
-        Node(copy(tree.op), _copy_node(tree.l))
-    else
-        Node(copy(tree.op), _copy_node(tree.l), _copy_node(tree.r))
-    end
-end
-
 const OP_NAMES = Dict(
     "safe_log" => "log",
     "safe_log2" => "log2",
@@ -365,51 +280,4 @@ function print_tree(
     return println(string_tree(tree, operators; varMap=varMap))
 end
 
-function Base.hash(tree::Node{T})::UInt where {T}
-    if tree.degree == 0
-        if tree.constant
-            # tree.val used.
-            return hash((0, tree.val::T))
-        else
-            # tree.feature used.
-            return hash((1, tree.feature))
-        end
-    elseif tree.degree == 1
-        return hash((1, tree.op, hash(tree.l)))
-    else
-        return hash((2, tree.op, hash(tree.l), hash(tree.r)))
-    end
-end
-
-function is_equal(a::Node{T}, b::Node{T})::Bool where {T}
-    if a.degree == 0
-        b.degree != 0 && return false
-        if a.constant
-            !(b.constant) && return false
-            return a.val::T == b.val::T
-        else
-            b.constant && return false
-            return a.feature == b.feature
-        end
-    elseif a.degree == 1
-        b.degree != 1 && return false
-        a.op != b.op && return false
-        return is_equal(a.l, b.l)
-    else
-        b.degree != 2 && return false
-        a.op != b.op && return false
-        return is_equal(a.l, b.l) && is_equal(a.r, b.r)
-    end
-end
-
-function Base.:(==)(a::Node{T}, b::Node{T})::Bool where {T}
-    return is_equal(a, b)
-end
-
-function Base.:(==)(a::Node{T1}, b::Node{T2})::Bool where {T1,T2}
-    T = promote_type(T1, T2)
-    # TODO: Should also have preserve_sharing check...
-    return is_equal(convert(Node{T}, a), convert(Node{T}, b))
-end
-
 end

src/EquationUtils.jl

@@ -1,7 +1,7 @@
 module EquationUtilsModule
 
 import Compat: Returns
-import ..EquationModule: Node, copy_node, tree_mapreduce, any, filter_and_map
+import ..EquationModule: Node, copy_node, tree_mapreduce, any, filter_map
 
 """
     count_nodes(tree::Node{T})::Int where {T}
@@ -64,7 +64,7 @@ The function `set_constants!` sets them in the same order,
 given the output of this function.
 """
 function get_constants(tree::Node{T}) where {T}
-    return filter_and_map(is_node_constant, t -> (t.val::T), tree; result_type=T)
+    return filter_map(is_node_constant, t -> (t.val::T), tree, T)
 end
 
 """

src/Utils.jl

@@ -81,7 +81,7 @@ isgood(x) = true
 isbad(x) = !isgood(x)
 
 """
-    @generate_idmap tree function my_function_on_tree(tree::Node)
+    @memoize_on tree function my_function_on_tree(tree::Node)
         ...
     end
 
@@ -91,14 +91,14 @@ IdDict()), it will use use the `id_map` to avoid recomputing the same value
 for the same node in a tree. Use this to automatically create functions that
 work with trees that have shared child nodes.
 """
-macro generate_idmap(tree, def)
-    idmap_def = _generate_idmap(tree, def)
+macro memoize_on(tree, def)
+    idmap_def = _memoize_on(tree, def)
     return quote
         $(esc(def)) # The normal function
         $(esc(idmap_def)) # The function with an id_map argument
     end
 end
-function _generate_idmap(tree::Symbol, def::Expr)
+function _memoize_on(tree::Symbol, def::Expr)
     sdef = splitdef(def)
 
     # Add an id_map argument
@@ -127,13 +127,13 @@ function _generate_idmap(tree::Symbol, def::Expr)
 end
 
 """
-    @use_idmap(call, id_map)
+    @with_memoize(call, id_map)
 
 This simple macro simply puts the `id_map`
-into the call, to be consistent with the `@generate_idmap` macro.
+into the call, to be consistent with the `@memoize_on` macro.
 
 ```
-@use_idmap(_copy_node(tree), IdDict{Any,Any}())
+@with_memoize(_copy_node(tree), IdDict{Any,Any}())
 ````
 
 is converted to 
@@ -143,7 +143,7 @@ _copy_node(tree, IdDict{Any,Any}())
 ```
 
 """
-macro use_idmap(def, id_map)
+macro with_memoize(def, id_map)
     idmap_def = _add_idmap_to_call(def, id_map)
     return quote
         $(esc(idmap_def))