fix: allow disabling CSE inside certain subtrees

Author: AayushSabharwal

src/code.jl

@@ -703,6 +703,16 @@ end
 
 @inline newsym(::Type{T}) where T = Sym{T}(gensym("cse"))
 
+"""
+    $(TYPEDSIGNATURES)
+
+Return `true` if CSE should descend inside `sym`, which has operation `f` and
+arguments `args...`.
+"""
+function cse_inside_expr(sym, f, args...)
+    return true
+end
+
 """
 $(SIGNATURES)
 
@@ -727,10 +737,16 @@ function topological_sort(graph)
             return visited[node]
         end
         if iscall(node)
+            op = operation(node)
+            args = arguments(node)
+            if !cse_inside_expr(node, op, args...)
+                visited[node] = node
+                return node
+            end
             args = map(dfs, arguments(node))
             # use `term` instead of `maketerm` because we only care about the operation being performed
             # and not the representation. This avoids issues with `newsym` symbols not having sizes, etc.
-            new_node = term(operation(node), args...)
+            new_node = term(operation(node), args...; type = symtype(node))
             sym = newsym(symtype(new_node))
             push!(sorted_nodes, sym ← new_node)
             visited[node] = sym

test/code.jl

@@ -274,3 +274,17 @@ end
     expected = eval(toexpr(Let([x ← 1, y ← 2, z ← 3], sarr)))
     @test fn(1, 2, 3) ≈ expected
 end
+
+function foo(args...) end
+
+SymbolicUtils.Code.cse_inside_expr(sym, ::typeof(foo), args...) = false
+
+@testset "`cse_inside_expr`" begin
+    @syms x y
+    ex1 = (x^2 + y^2)
+    exfoo = term(foo, ex1; type = Real)
+    ex2 = ex1 + exfoo
+    letblock = cse(ex2)
+    ex3 = letblock.body
+    @test any(isequal(exfoo), arguments(ex3))
+end