feat: more extensible expression interface

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -66,7 +66,8 @@ import .NodeModule:
 @reexport import .EvaluationHelpersModule
 @reexport import .ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node
 @reexport import .RandomModule: NodeSampler
-@reexport import .ExpressionModule: AbstractExpression, Expression, with_tree
+@reexport import .ExpressionModule:
+    AbstractExpression, Expression, with_tree, with_metadata, get_contents, get_metadata
 import .ExpressionModule:
     get_tree, get_operators, get_variable_names, Metadata, default_node_type, node_type
 @reexport import .ParseModule: @parse_expression, parse_expression

src/Expression.jl

@@ -140,16 +140,46 @@ end
 function set_constants!(ex::AbstractExpression{T}, constants, refs) where {T}
     return error("`set_constants!` function must be implemented for $(typeof(ex)) types.")
 end
+function get_contents(ex::AbstractExpression)
+    return error("`get_contents` function must be implemented for $(typeof(ex)) types.")
+end
+function get_metadata(ex::AbstractExpression)
+    return error("`get_metadata` function must be implemented for $(typeof(ex)) types.")
+end
 ########################################################
 
 """
     with_tree(ex::AbstractExpression, tree::AbstractExpressionNode)
+    with_tree(ex::AbstractExpression, tree::AbstractExpression)
 
 Create a new expression based on `ex` but with a different `tree`
 """
+function with_tree(ex::AbstractExpression, tree::AbstractExpression)
+    return with_tree(ex, get_contents(tree))
+end
 function with_tree(ex::AbstractExpression, tree)
-    return constructorof(typeof(ex))(tree, ex.metadata)
+    return constructorof(typeof(ex))(tree, get_metadata(ex))
 end
+function get_contents(ex::Expression)
+    return ex.tree
+end
+
+"""
+    with_metadata(ex::AbstractExpression, metadata)
+    with_metadata(ex::AbstractExpression; metadata...)
+
+Create a new expression based on `ex` but with a different `metadata`.
+"""
+function with_metadata(ex::AbstractExpression; metadata...)
+    return with_metadata(get_contents(ex), Metadata(metadata))
+end
+function with_metadata(ex::AbstractExpression, metadata::Metadata)
+    return constructorof(typeof(ex))(get_contents(ex), metadata)
+end
+function get_metadata(ex::Expression)
+    return ex.metadata
+end
+
 function preserve_sharing(::Union{E,Type{E}}) where {T,N,E<:AbstractExpression{T,N}}
     return preserve_sharing(N)
 end

src/Interfaces.jl

@@ -43,7 +43,10 @@ using ..ExpressionModule:
     get_tree,
     get_operators,
     get_variable_names,
+    get_contents,
+    get_metadata,
     with_tree,
+    with_metadata,
     default_node_type
 using ..ParametricExpressionModule: ParametricExpression, ParametricNode
 
@@ -52,6 +55,14 @@ using ..ParametricExpressionModule: ParametricExpression, ParametricNode
 ###############################################################################
 
 ## mandatory
+function _check_get_contents(ex::AbstractExpression)
+    new_ex = with_tree(ex, get_contents(ex))
+    return new_ex == ex && new_ex isa typeof(ex)
+end
+function _check_get_metadata(ex::AbstractExpression)
+    new_ex = with_metadata(ex, get_metadata(ex))
+    return new_ex == ex && new_ex isa typeof(ex)
+end
 function _check_get_tree(ex::AbstractExpression{T,N}) where {T,N}
     return get_tree(ex) isa N
 end
@@ -67,6 +78,15 @@ function _check_copy(ex::AbstractExpression)
     # TODO: Could include checks for aliasing here
     return preserves
 end
+function _check_with_tree(ex::AbstractExpression)
+    new_ex = with_tree(ex, get_contents(ex))
+    new_ex2 = with_tree(ex, ex)
+    return new_ex == ex && new_ex isa typeof(ex) && new_ex2 == ex && new_ex2 isa typeof(ex)
+end
+function _check_with_metadata(ex::AbstractExpression)
+    new_ex = with_metadata(ex, get_metadata(ex))
+    return new_ex == ex && new_ex isa typeof(ex)
+end
 
 ## optional
 function _check_count_nodes(ex::AbstractExpression)
@@ -116,10 +136,14 @@ end
 #! format: off
 ei_components = (
     mandatory = (
+        get_contents = "extracts the runtime contents of an expression" => _check_get_contents,
+        get_metadata = "extracts the runtime metadata of an expression" => _check_get_metadata,
         get_tree = "extracts the expression tree from [`AbstractExpression`](@ref)" => _check_get_tree,
         get_operators = "returns the operators used in the expression (or pass `operators` explicitly to override)" => _check_get_operators,
         get_variable_names = "returns the variable names used in the expression (or pass `variable_names` explicitly to override)" => _check_get_variable_names,
         copy = "returns a copy of the expression" => _check_copy,
+        with_tree = "returns the expression with different tree" => _check_with_tree,
+        with_metadata = "returns the expression with different metadata" => _check_with_metadata,
     ),
     optional = (
         count_nodes = "counts the number of nodes in the expression tree" => _check_count_nodes,

src/ParametricExpression.jl

@@ -19,7 +19,13 @@ import ..EvaluateModule: eval_tree_array
 import ..EvaluateDerivativeModule: eval_grad_tree_array
 import ..EvaluationHelpersModule: _grad_evaluator
 import ..ExpressionModule:
-    get_tree, get_operators, get_variable_names, max_feature, default_node_type
+    get_contents,
+    get_metadata,
+    get_tree,
+    get_operators,
+    get_variable_names,
+    max_feature,
+    default_node_type
 import ..ParseModule: parse_leaf
 
 """A type of expression node that also stores a parameter index"""
@@ -127,9 +133,9 @@ end
 ###############################################################################
 # Abstract expression interface ###############################################
 ###############################################################################
-function get_tree(ex::ParametricExpression)
-    return ex.tree
-end
+get_contents(ex::ParametricExpression) = ex.tree
+get_metadata(ex::ParametricExpression) = ex.metadata
+get_tree(ex::ParametricExpression) = ex.tree
 function get_operators(ex::ParametricExpression, operators=nothing)
     return operators === nothing ? ex.metadata.operators : operators
 end

test/test_multi_expression.jl

@@ -10,6 +10,13 @@
         trees::TREES
         metadata::Metadata{D}
 
+        function MultiScalarExpression(trees::NamedTuple, metadata::Metadata{D}) where {D}
+            example_tree = first(values(trees))
+            N = typeof(example_tree)
+            T = eltype(example_tree)
+            return new{T,N,typeof(trees),D}(trees, metadata)
+        end
+
         """
         Create a multi-expression expression type.
 
@@ -65,6 +72,12 @@
     end
 
     tree_factory(f::F, trees) where {F} = f(; trees...)
+    function DE.get_contents(ex::MultiScalarExpression)
+        return ex.trees
+    end
+    function DE.get_metadata(ex::MultiScalarExpression)
+        return ex.metadata
+    end
     function DE.get_tree(ex::MultiScalarExpression{T,N}) where {T,N}
         fused_expression = parse_expression(
             tree_factory(ex.metadata.tree_factory, ex.trees)::Expr;