refactor: additional tricks to try to improve type stability

Author: MilesCranmer

src/base.jl

@@ -121,8 +121,8 @@ struct TreeMapreducer{
 end
 
 @generated function call_mapreducer(
-    mapreducer::TreeMapreducer{D,ID}, tree::AbstractNode
-) where {D,ID}
+    mapreducer::TreeMapreducer{D,ID,F1,F2,G,H}, tree::AbstractNode
+) where {D,ID,F1,F2,G,H}
     quote
         key = ID <: Dict ? objectid(tree) : nothing
         if ID <: Dict && haskey(mapreducer.id_map, key)
@@ -353,7 +353,7 @@ end
 Collect all nodes in a tree into a flat array in depth-first order.
 """
 function collect(tree::AbstractNode; break_sharing::Val{BS}=Val(false)) where {BS}
-    return filter(Returns(true), tree; break_sharing=Val(BS))
+    return filter(_ -> true, tree; break_sharing=Val(BS))
 end
 Base.IteratorSize(::Type{<:AbstractNode}) = Base.HasLength()
 

test/Project.toml

@@ -29,5 +29,5 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 Aqua = "0.7"
 
 [preferences.DynamicExpressions]
-dispatch_doctor_mode = "error"
 dispatch_doctor_codegen_level = "min"
+dispatch_doctor_mode = "error"

test/test_extra_node_fields.jl

@@ -27,7 +27,9 @@ function DynamicExpressions.with_type_parameters(
 ) where {T,N<:FrozenNode}
     return FrozenNode{T,max_degree(N)}
 end
-function DynamicExpressions.with_max_degree(::Type{N}, ::Val{D}) where {T,N<:FrozenNode{T}}
+function DynamicExpressions.with_max_degree(
+    ::Type{N}, ::Val{D}
+) where {T,N<:FrozenNode{T},D}
     return FrozenNode{T,D}
 end
 function DynamicExpressions.leaf_copy(t::FrozenNode{T}) where {T}