refactor: clean up extensions

Author: MilesCranmer

ext/DynamicExpressionsCUDAExt.jl

@@ -51,11 +51,7 @@ function eval_tree_array(
     ## the GPU kernel, with size equal to the number of rows
     ## in the input data by the number of nodes in the tree.
     ## It has one extra row to store the constant values.
-    gworkspace = if gpu_workspace === nothing
-        similar(gcX, num_elem + 1, num_nodes)
-    else
-        gpu_workspace
-    end
+    gworkspace = @something(gpu_workspace, similar(gcX, num_elem + 1, num_nodes))
     gval = @view gworkspace[end, :]
     if _update_buffers
         copyto!(gval, val)

src/AsArray.jl

@@ -2,14 +2,17 @@ module AsArrayModule
 
 using ..NodeModule: AbstractExpressionNode, tree_mapreduce, count_nodes
 
+function as_array(
+    ::Type{I}, trees::N; buffer=nothing
+) where {T,N<:AbstractExpressionNode{T},I}
+    return as_array(I, (trees,); buffer=buffer)
+end
+
 function as_array(
     ::Type{I},
-    trees::Union{N,Tuple{N,Vararg{N}},AbstractVector{N}};
+    trees::Union{Tuple{N,Vararg{N}},AbstractVector{N}};
     buffer::Union{AbstractArray,Nothing}=nothing,
 ) where {T,N<:AbstractExpressionNode{T},I}
-    if trees isa N
-        return as_array(I, (trees,); buffer=buffer)
-    end
     each_num_nodes = (t -> count_nodes(t; break_sharing=Val(true))).(trees)
     num_nodes = sum(each_num_nodes)
 
@@ -25,7 +28,7 @@ function as_array(
     val = Array{T}(undef, num_nodes)
 
     ## Views of the same matrix:
-    buffer = buffer === nothing ? Array{I}(undef, 8, num_nodes) : buffer
+    buffer = @something(buffer, Array{I}(undef, 8, num_nodes))
     degree = @view buffer[1, :]
     feature = @view buffer[2, :]
     op = @view buffer[3, :]