feat: require user to pass ArrayBuffer object explicitly

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -73,6 +73,7 @@ import .StringsModule: get_op_name
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!
 @reexport import .EvaluateModule:
     eval_tree_array, differentiable_eval_tree_array, EvalOptions
+import .EvaluateModule: ArrayBuffer
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
 @reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!

src/Evaluate.jl

@@ -70,7 +70,7 @@ function get_feature_array(buffer::ArrayBuffer, X::AbstractMatrix, feature::Inte
 end
 
 """
-    EvalOptions{T,B,E}
+    EvalOptions
 
 This holds options for expression evaluation, such as evaluation backend.
 
@@ -86,8 +86,10 @@ This holds options for expression evaluation, such as evaluation backend.
     the entire buffer. This early exit is performed to avoid wasting compute cycles.
     Setting `Val{false}` will continue the computation as usual and thus result in
     `NaN`s only in the elements that actually have `NaN`s.
+- `buffer::Union{ArrayBuffer,Nothing}`: If not `nothing`, use this buffer for evaluation.
+    This should be an instance of `ArrayBuffer` which has an `array` field and an
+    `index` field used to iterate which buffer slot to use.
 """
-
 struct EvalOptions{T,B,E,BUF<:Union{ArrayBuffer,Nothing}}
     turbo::Val{T}
     bumper::Val{B}
@@ -99,24 +101,17 @@ end
     turbo::Union{Bool,Val}=Val(false),
     bumper::Union{Bool,Val}=Val(false),
     early_exit::Union{Bool,Val}=Val(true),
-    buffer::Union{AbstractMatrix,Nothing}=nothing,
-    buffer_ref::Union{Base.RefValue{<:Integer},Nothing}=nothing,
+    buffer::Union{ArrayBuffer,Nothing}=nothing,
 )
     v_turbo = _to_bool_val(turbo)
     v_bumper = _to_bool_val(bumper)
     v_early_exit = _to_bool_val(early_exit)
 
     if v_bumper isa Val{true}
-        @assert buffer === nothing && buffer_ref === nothing
-    end
-
-    array_buffer = if buffer === nothing
-        nothing
-    else
-        ArrayBuffer(buffer, buffer_ref)
+        @assert buffer === nothing
     end
 
-    return EvalOptions(v_turbo, v_bumper, v_early_exit, array_buffer)
+    return EvalOptions(v_turbo, v_bumper, v_early_exit, buffer)
 end
 
 @unstable @inline _to_bool_val(x::Bool) = x ? Val(true) : Val(false)

test/test_buffered_evaluation.jl

@@ -1,5 +1,6 @@
 @testitem "Buffer creation and validation" begin
     using DynamicExpressions
+    using DynamicExpressions: ArrayBuffer
 
     # Test data setup
     X = rand(2, 10)  # 2 features, 10 samples
@@ -11,13 +12,13 @@
     # Basic buffer creation - buffer shape should match (num_leafs, num_samples)
     buffer = zeros(5, size(X, 2))  # 5 leaves should be enough for our test tree
     buffer_ref = Ref(0)
-    eval_options = EvalOptions(; buffer=buffer, buffer_ref=buffer_ref)
+    eval_options = EvalOptions(; buffer=ArrayBuffer(buffer, buffer_ref))
     @test eval_options.buffer.array === buffer
     @test eval_options.buffer.index === buffer_ref
 
     # Test buffer is not allowed with bumper
     @test_throws AssertionError EvalOptions(;
-        bumper=true, buffer=buffer, buffer_ref=buffer_ref
+        bumper=true, buffer=ArrayBuffer(buffer, buffer_ref)
     )
 
     # Basic evaluation should work
@@ -28,6 +29,7 @@ end
 
 @testitem "Buffer correctness" begin
     using DynamicExpressions
+    using DynamicExpressions: ArrayBuffer
 
     X = rand(2, 10)
     operators = OperatorEnum(; binary_operators=[+, *], unary_operators=[sin])
@@ -45,7 +47,7 @@ end
         # Evaluation with buffer
         buffer = zeros(5, size(X, 2))
         buffer_ref = Ref(0)
-        eval_options = EvalOptions(; buffer=buffer, buffer_ref=buffer_ref)
+        eval_options = EvalOptions(; buffer=ArrayBuffer(buffer, buffer_ref))
         result2, ok2 = eval_tree_array(tree, X, operators; eval_options)
 
         # Results should be identical
@@ -56,7 +58,7 @@ end
 
 @testitem "Buffer index management" begin
     using DynamicExpressions
-
+    using DynamicExpressions: ArrayBuffer
     X = rand(2, 10)
     operators = OperatorEnum(; binary_operators=[+, *], unary_operators=[sin])
 
@@ -70,7 +72,7 @@ end
     # This tree needs more buffer space due to intermediate computations
     buffer = zeros(10, size(X, 2))
     buffer_ref = Ref(0)
-    eval_options = EvalOptions(; buffer=buffer, buffer_ref=buffer_ref)
+    eval_options = EvalOptions(; buffer=ArrayBuffer(buffer, buffer_ref))
 
     # Index should start at 1
     @test buffer_ref[] == 0
@@ -93,6 +95,7 @@ end
 
 @testitem "Buffer error handling" begin
     using DynamicExpressions
+    using DynamicExpressions: ArrayBuffer
 
     X = rand(2, 10)
     operators = OperatorEnum(; binary_operators=[+, /, *], unary_operators=[sin])
@@ -106,7 +109,7 @@ end
 
     buffer = zeros(5, size(X, 2))
     buffer_ref = Ref(0)
-    eval_options = EvalOptions(; buffer=buffer, buffer_ref=buffer_ref)
+    eval_options = EvalOptions(; buffer=ArrayBuffer(buffer, buffer_ref))
 
     # Test with early_exit=true
     result1, ok1 = eval_tree_array(tree, X, operators; eval_options)
@@ -115,6 +118,7 @@ end
 
 @testitem "Random tree buffer evaluation" begin
     using DynamicExpressions
+    using DynamicExpressions: ArrayBuffer
     using Random
     using LoopVectorization
     include("tree_gen_utils.jl")
@@ -139,7 +143,7 @@ end
         # Buffer evaluation
         buffer = Array{Float64}(undef, 2n_nodes, size(X, 2))
         buffer_ref = Ref(rand(1:10))  # Random starting index (will be reset)
-        eval_options = EvalOptions(; turbo, buffer, buffer_ref)
+        eval_options = EvalOptions(; turbo, buffer=ArrayBuffer(buffer, buffer_ref))
         result2, ok2 = eval_tree_array(tree, X, operators; eval_options)
 
         # Results should be identical