fix: fix inference of substitute

Author: AayushSabharwal

benchmark/benchmarks.jl

@@ -58,15 +58,21 @@ let r = @rule(~x => ~x), rs = RuleSet([r]),
 
     # we use `fold = false` since otherwise it dynamic dispatches to `sin`/`cos` whenever
     # both arguments in the contained addition are substituted.
-    overhead["substitute"]["a"] = @benchmarkable substitute(subs_expr, $(Dict(a=>1)); fold = false) setup=begin
+
+    fold = @static if hasmethod(SymbolicUtils.Substituter{true}, Tuple{Dict, Function})
+        Val{false}()
+    else
+        false
+    end
+    overhead["substitute"]["a"] = @static @benchmarkable substitute(subs_expr, $(Dict(a=>1)); fold) setup=begin
         subs_expr = (sin(a+b) + cos(b+c)) * (sin(b+c) + cos(c+a)) * (sin(c+a) + cos(a+b))
     end
 
-    overhead["substitute"]["a,b"] = @benchmarkable substitute(subs_expr, $(Dict(a=>1, b=>2)); fold = false) setup=begin
+    overhead["substitute"]["a,b"] = @benchmarkable substitute(subs_expr, $(Dict(a=>1, b=>2)); fold) setup=begin
         subs_expr = (sin(a+b) + cos(b+c)) * (sin(b+c) + cos(c+a)) * (sin(c+a) + cos(a+b))
     end
 
-    overhead["substitute"]["a,b,c"] = @benchmarkable substitute(subs_expr, $(Dict(a=>1, b=>2, c=>3)); fold = false) setup=begin
+    overhead["substitute"]["a,b,c"] = @benchmarkable substitute(subs_expr, $(Dict(a=>1, b=>2, c=>3)); fold) setup=begin
         subs_expr = (sin(a+b) + cos(b+c)) * (sin(b+c) + cos(c+a)) * (sin(c+a) + cos(a+b))
     end
 

src/SymbolicUtils.jl

@@ -155,6 +155,8 @@ include("substitute.jl")
 include("code.jl")
 
 PrecompileTools.@setup_workload begin
+    fold1 = Val{false}()
+    fold2 = Val{true}()
     PrecompileTools.@compile_workload begin
         @syms x y f(t) q[1:5]
         Sym{SymReal}(:a; type = Real, shape = ShapeVecT())
@@ -173,8 +175,15 @@ PrecompileTools.@setup_workload begin
         show(devnull, x ^ 2 + y * x + y / 3x)
         expand((x + y) ^ 2)
         simplify(x ^ (1//2) + (sin(x) ^ 2 + cos(x) ^ 2) + 2(x + y) - x - y)
-        substitute(x + 2y + sin(x), Dict(x => y); fold = false)
-        substitute(x + 2y + sin(x), Dict(x => 1); fold = true)
+        ex = x + 2y + sin(x)
+        rules1 = Dict(x => y)
+        rules2 = Dict(x => 1)
+        substitute(ex, rules1)
+        substitute(ex, rules1; fold = fold1)
+        substitute(ex, rules2; fold = fold1)
+        substitute(ex, rules2)
+        substitute(ex, rules1; fold = fold2)
+        substitute(ex, rules2; fold = fold2)
         q[1]
         q'q
     end

src/substitute.jl

@@ -3,12 +3,22 @@ struct Substituter{Fold, D <: AbstractDict, F}
     filter::F
 end
 
+@inline function Substituter{Fold}(d::AbstractDict, filter::F) where {Fold, F}
+    Substituter{Fold, typeof(d), F}(d, filter)
+end
+@inline function Substituter{Fold}(d::Pair, filter::F) where {Fold, F}
+    Substituter{Fold}(Dict(d), filter)
+end
+@inline function Substituter{Fold}(d::AbstractArray{<:Pair}, filter::F) where {Fold, F}
+    Substituter{Fold}(Dict(d), filter)
+end
+
 function (s::Substituter)(ex)
     return get(s.dict, ex, ex)
 end
 
 function (s::Substituter)(ex::AbstractArray)
-    map(s, ex)
+    [s(x) for x in ex]
 end
 
 function (s::Substituter)(ex::SparseMatrixCSC)
@@ -90,22 +100,22 @@ end
 end
 
 """
-    substitute(expr, dict; fold=true)
+    substitute(expr, dict; fold=Val(true))
 
 substitute any subexpression that matches a key in `dict` with
-the corresponding value. If `fold=false`,
+the corresponding value. If `fold=Val(false)`,
 expressions which can be evaluated won't be evaluated.
 
 ```julia
-julia> substitute(1+sqrt(y), Dict(y => 2), fold=true)
+julia> substitute(1+sqrt(y), Dict(y => 2), fold=Val(true))
 2.414213562373095
-julia> substitute(1+sqrt(y), Dict(y => 2), fold=false)
+julia> substitute(1+sqrt(y), Dict(y => 2), fold=Val(false))
 1 + sqrt(2)
 ```
 """
-@inline function substitute(expr, dict; fold=true, filterer=default_substitute_filter)
-    isempty(dict) && !fold && return expr
-    return Substituter{fold, typeof(dict), typeof(filterer)}(dict, filterer)(expr)
+@inline function substitute(expr, dict; fold::Val{Fold}=Val{true}(), filterer=default_substitute_filter) where {Fold}
+    isempty(dict) && !Fold && return expr
+    return Substituter{Fold}(dict, filterer)(expr)
 end
 
 """
@@ -265,7 +275,7 @@ function _reduce_eliminated_idxs(expr::BasicSymbolic{T}, output_idx::OutIdxT{T},
         for (idx, ii) in zip(iidxs, collapsed)
             subrules[ii] = idx
         end
-        return substitute(new_expr, subrules; fold = false)
+        return substitute(new_expr, subrules; fold = Val{false}())
     end
 end
 @cache function reduce_eliminated_idxs_1(expr::BasicSymbolic{SymReal}, output_idx::OutIdxT{SymReal}, ranges::RangesT{SymReal}, reduce)::BasicSymbolic{SymReal}
@@ -335,9 +345,9 @@ function scalarize(x::BasicSymbolic{T}, ::Val{toplevel} = Val{false}()) where {T
                     subrules[ii] = idxs[i]
                 end
                 if toplevel
-                    substitute(new_expr, subrules; fold = true)
+                    substitute(new_expr, subrules; fold = Val{true}())
                 else
-                    scalarize(substitute(new_expr, subrules; fold = true))
+                    scalarize(substitute(new_expr, subrules; fold = Val{true}()))
                 end
             end
         end

src/types.jl

@@ -3155,10 +3155,10 @@ Base.@propagate_inbounds function _getindex(arr::BasicSymbolic{T}, idxs::Union{B
             end
             if isempty(new_output_idx)
                 new_expr = reduce_eliminated_idxs(expr, output_idx, ranges, reduce)
-                result = substitute(new_expr, subrules; fold = false, filterer = !isarrayop)
+                result = substitute(new_expr, subrules; fold = Val{false}(), filterer = !isarrayop)
                 return result
             else
-                new_expr = substitute(expr, subrules; fold = false, filterer = !isarrayop)
+                new_expr = substitute(expr, subrules; fold = Val{false}(), filterer = !isarrayop)
                 if term !== nothing
                     term = getindex(term, idxs...)
                 end