custom similarterm for egraphs, update metatheory, fix preprocessing

Author: Alessandro

Project.toml

@@ -36,7 +36,7 @@ DocStringExtensions = "0.8"
 DynamicPolynomials = "0.3"
 IfElse = "0.1"
 LabelledArrays = "1.5"
-Metatheory = "1.1"
+Metatheory = "1.2"
 MultivariatePolynomials = "0.3"
 NaNMath = "0.3"
 Setfield = "0.7, 0.8"

src/egraph.jl

@@ -1,7 +1,7 @@
 using Metatheory.Rewriters
 
-function EGraphs.preprocess(t)
-    Chain([Postwalk(toterm), Postwalk(unflatten)])(t) 
+function EGraphs.preprocess(t::Symbolic)
+    toterm(unflatten(t)) 
 end
 
 function symbolicegraph(ex)
@@ -76,12 +76,26 @@ function costfun(n::ENodeTerm, g::EGraph, an)
     cost
 end
 
-costfun(n::ENodeLiteral, g::EGraph, an) = 1
+costfun(n::ENodeLiteral, g::EGraph, an) = 0
 
+egraph_simterm(x, head, args, symtype=nothing; metadata=nothing, exprhead=exprhead(x)) = 
+TermInterface.similarterm(typeof(x), head, args, symtype; metadata=metadata, exprhead=exprhead)
+
+
+# Custom similarterm to use in EGraphs on <:Symbolic types that treats everything as a Term 
+function egraph_simterm(x::Type{<:Term}, f, args, symtype=nothing; metadata=nothing, exprhead=:call)
+    T = symtype
+    if T === nothing
+        T = _promote_symtype(f, args)
+    end
+    res = Term{T}(f isa Symbol ? eval(f) : f, args; metadata=metadata);
+    return res
+end 
 
 function optimize(ex; params=SaturationParams(timeout=20))
     # @show ex
     g = symbolicegraph(ex)
+    params.simterm = egraph_simterm
     saturate!(g, opt_theory, params)
     return extract!(g, costfun)
 end

test/egraph.jl

@@ -1,19 +1,19 @@
 using SymbolicUtils
+using Metatheory
 
 sym_eq(x, y) = simplify(x == y)
-(≖)(x, y) = sym_eq(x,y)
 
-@syms a b c x
 @testset "Basic optimization" begin 
     ex = 2a + 2b - (a*(a + b))
     res = optimize(ex)
-    @test res ≖ (a+b)*(2-a)
+    @test sym_eq(res, (a+b)*(2-a))
 
     ex = sin(a^2)/cos(a^2)
     res = optimize(ex)
-    @test isequal(res, tan(a^2)) # sym_eq does not work
+    @test isequal(res, tan(a^2)) # sym_eq does not work 
+    # TODO report issue
 
     ex = sin(1/x * (a * b + a * c))^2 + cos((a*(b+c))/x)^2
     res = optimize(ex)
     @test res == 1
-end
+end